연관 회상 메모리를 활용한 LLM 컨텍스트 확장

Extending LLM Context via Associative Recurrent Memory

대규모 언어 모델(LLM)의 컨텍스트 길이를 늘리는 것은 다양한 실제 응용 분야에서 매우 중요하지만, 기존 트랜스포머 구조는 계산 복잡도가 2차 함수로 증가하고 메모리 사용량이 선형적으로 증가하는 제약이 있습니다. 본 연구에서는 연관 회상 메모리 트랜스포머(ARMT)를 LLM의 장기 컨텍스트 처리 능력 향상, 일정한 메모리 확장, 그리고 더 나은 효율성을 위한 실용적인 접근 방식으로 탐구합니다. 본 연구는 세 가지 주요 기여를 합니다. 첫째, 실제 워크로드를 평가하기 위해 설계된 두 개의 도메인 특화 장기 컨텍스트 데이터셋을 구축했으며, 특히 좁은 영역에서의 미세 조정 시나리오에 중점을 두었습니다. 둘째, 지속적인 사전 학습, 합성 장기 컨텍스트 데이터 생성, 커리큘럼 학습, 그리고 선택적으로 연관 메모리를 특정 모델 레이어에 통합하는 포괄적인 ARMT 기반 컨텍스트 확장 학습 방법을 제안합니다. 셋째, 광범위한 실험 연구를 통해 ARMT를 적용한 모델들이 (i) 원래 컨텍스트 제한을 훨씬 넘어 입력 데이터를 잘 처리하면서도 기존 기준 성능을 유지하고, (ii) 분포 외부의 다양한 컨텍스트 길이에 대해 더 효과적으로 일반화되며, (iii) 원래 컨텍스트 창 내에서 기준 성능을 유지하면서 30% 적은 FLOPs를 사용한다는 것을 보여줍니다.

Extending the context length of large language models (LLMs) is critical for many real-world applications, yet standard transformers remain constrained by quadratic compute and linear memory scaling. In this work, we investigate the Associative Recurrent Memory Transformer (ARMT) as a practical approach for enabling long-context processing in LLMs, constant memory scaling, and better efficiency. We make three main contributions. First, we construct two domain-specific long-context datasets designed to evaluate realistic workloads, focusing on narrow-domain fine-tuning scenarios. Second, we propose a comprehensive training recipe for ARMT-based context extension, combining continued pre-training, synthetic long-context data generation, curriculum learning, and selective integration of associative memory into chosen model layers. Third, we present an extensive experimental study demonstrating that ARMT-augmented models: (i) process inputs well beyond their original context limits without degrading performance relative to in-limit baselines; (ii) generalize more effectively to out-of-distribution context lengths; and (iii) need 30% less FLOPs while preserving baseline performance within the original context window.