Search papers, labs, and topics across Lattice.
This paper investigates the Associative Recurrent Memory Transformer (ARMT) to extend the context length of large language models (LLMs) while maintaining constant memory scaling and improving efficiency. By constructing domain-specific long-context datasets and employing a novel training recipe that includes continued pre-training and synthetic data generation, the authors demonstrate that ARMT-augmented models can process inputs beyond their original context limits without performance degradation. Key results show that these models generalize better to out-of-distribution context lengths and require 30% fewer FLOPs while preserving baseline performance within the original context window.
ARMT-augmented models can handle inputs far beyond their original context limits while using 30% less compute, revolutionizing efficiency in long-context processing.
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.
