Extending LLM Context via Associative Recurrent Memory

2026-07-13

AI Summary

This work addresses the limitation of large language models in efficiently handling long-sequence tasks due to fixed context lengths. The authors propose the Associative Recurrent Memory Transformer (ARMT), which incorporates a recurrent memory mechanism to substantially extend context processing capacity with only constant-level memory overhead. By integrating continual pretraining, synthetic long-context data generation, curriculum learning, and a selective memory layer integration strategy, ARMT maintains competitive performance within the original context window while reducing FLOPs by 30% and demonstrating strong generalization to inputs far exceeding the training sequence length. The study further introduces a realistic long-context evaluation benchmark, empirically validating the method’s effectiveness and practical utility.

Abstract

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.

Problem

• large language models
• context length
• long-context processing
• transformer limitations
• memory scaling

Innovation

• Associative Recurrent Memory
• Long-context LLMs
• Constant Memory Scaling
• Curriculum Learning
• Synthetic Data Generation