AllMem: 효율적인 긴 문맥 모델링을 위한 메모리 중심의 해법

AllMem: A Memory-centric Recipe for Efficient Long-context Modeling

대규모 언어 모델(LLM)은 자기 주의(self-attention) 메커니즘에 내재된 연산 복잡도와 메모리 오버헤드로 인해 긴 시퀀스 작업에서 심각한 성능 병목 현상을 겪습니다. 이러한 문제를 해결하기 위해, 우리는 슬라이딩 윈도우 어텐션(SWA)과 비선형 TTT(Test-Time Training) 메모리 네트워크를 통합한 새롭고 효율적인 하이브리드 아키텍처인 AllMem을 소개합니다. AllMem은 캣스트로픽 포게팅(catastrophic forgetting)을 완화하면서 모델이 초장문 맥락으로 효과적으로 확장할 수 있도록 지원합니다. 이 접근 방식은 선형 메모리 모델의 전형적인 표현 제약을 극복할 뿐만 아니라, 긴 시퀀스 추론 시 연산 및 메모리 사용량을 크게 줄여줍니다. 또한, 우리는 사전 훈련된 모델의 표준 어텐션 계층을 메모리가 증강된 슬라이딩 윈도우 계층으로 교체하기 위해 메모리 효율적 미세 조정(Fine-Tuning) 전략을 구현했습니다. 이 프레임워크는 사전 훈련된 기성 LLM을 AllMem 기반 아키텍처로 효율적으로 변환할 수 있게 해줍니다. 실증적 평가에 따르면, 우리의 4k 윈도우 모델은 37k LongBench에서 전체 어텐션 대비 0.83의 미미한 하락만으로 거의 손실 없는 성능을 달성했습니다. 또한 128k 문맥의 InfiniteBench에서 8k 윈도우 모델은 전체 어텐션보다 뛰어난 성능을 보였으며, 이는 글로벌 어텐션의 막대한 비용 없이도 노이즈를 완화하고 견고한 장거리 모델링을 유지하는 파라미터화된 메모리의 효과를 입증합니다.

Large Language Models (LLMs) encounter significant performance bottlenecks in long-sequence tasks due to the computational complexity and memory overhead inherent in the self-attention mechanism. To address these challenges, we introduce \textsc{AllMem}, a novel and efficient hybrid architecture that integrates Sliding Window Attention (SWA) with non-linear Test-Time Training (TTT) memory networks. \textsc{AllMem} enables models to effectively scale to ultra-long contexts while mitigating catastrophic forgetting. This approach not only overcomes the representation constraints typical of linear memory models but also significantly reduces the computational and memory footprint during long-sequence inference. Furthermore, we implement a Memory-Efficient Fine-Tuning strategy to replace standard attention layers in pre-trained models with memory-augmented sliding window layers. This framework facilitates the efficient transformation of any off-the-shelf pre-trained LLM into an \textsc{AllMem}-based architecture. Empirical evaluations confirm that our 4k window model achieves near-lossless performance on 37k LongBench with a marginal 0.83 drop compared to full attention. Furthermore, on InfiniteBench at a 128k context, our 8k window variant outperforms full attention, which validates the effectiveness of our parameterized memory in mitigating noise and maintaining robust long-range modeling without the prohibitive costs of global attention.