Long-term memory is a cornerstone of human intelligence. Enabling AI to process lifetime-scale information remains a long-standing pursuit in the field. Due to the constraints of full-attention architectures, the effective context length of large language models (LLMs) is typically limited to 1M tokens. Existing approaches, such as hybrid linear attention, fixed-size memory states (e.g., RNNs), and external storage methods like RAG or agent systems, attempt to extend this limit. However, they often suffer from severe precision degradation and rapidly increasing latency as context length grows, an inability to dynamically modify memory content, or a lack of end-to-end optimization. These bottlenecks impede complex scenarios like large-corpus summarization, Digital Twins, and long-history agent reasoning, while limiting memory capacity and slowing inference. We present Memory Sparse Attention (MSA), an end-to-end trainable, efficient, and massively scalable memory model framework. Through core innovations including scalable sparse attention and document-wise RoPE, MSA achieves linear complexity in both training and inference while maintaining exceptional stability, exhibiting less than 9% degradation when scaling from 16K to 100M tokens. Furthermore, KV cache compression, combined with Memory Parallel, enables 100M-token inference on 2xA800 GPUs. We also propose Memory Interleaving to facilitate complex multi-hop reasoning across scattered memory segments. MSA significantly surpasses frontier LLMs, state-of-the-art RAG systems, and leading memory agents in long-context benchmarks. These results demonstrate that by decoupling memory capacity from reasoning, MSA provides a scalable foundation to endow general-purpose models with intrinsic, lifetime-scale memory.
Figures & tables
Figure 1: MSA integrates top k selection with sparse attention, achieving strong scalability while remaining differentiable. This design enables end-to-end training, yet allows the documents to be decoupled at inference time, thereby providing robust extrapolation capability. MSA demonstrates exceptional scalability on the MS MARCO dataset, sustaining consistent performance with less than 9% degradation across an unprecedented memory context range from 16K to 100M tokens. Some curves terminate prematurely due to context length limitations.
Method
Lifetime Memory
Precision
Compatible w/ Mainstream LLMs
Computational Complexity
Memory Management
Catastrophic Forgetting
Parameter-Based Memory
Model-Based (LoRA / CPT)
No
High
High
Training: High Inference: Low
Hard
Yes
Test-Time Training (Titans)
No
Medium
Low
Medium
Medium
Yes
External Storage-Based Memory
RAG
Yes
Medium
Medium
O(L)
Easy
Low
MemAgent
Yes
Medium
Medium
Medium
Easy
Medium
Table 1: Comparison of Long-Term Memory Methods for LLMs
Figure 2: Memory Sparse Attention layer
Figure 3: Three-Stage Inference Process with Memory Interleave
Tokens
Qwen3-4B
Qwen3-4B (RR)
Hipporag2
MSA (Ours)
Dataset
R@1
R@5
R@10
R@1
R@5
R@10
R@1
R@5
R@10
@adaptive
MS MARCO v1 [ 2 ]
7.34M
2.893
3.011
3.005
2.934
3.032
3.017
2.676
3.005
3.019
4.141
Natural Questions [ 24 ]
1.47M
3.452
3.374
3.297
3.494
3.408
3.385
3.338
3.389
3.374
3.545
DuReader [ 14 ]
277K
3.726
3.579
3.594
3.848
3.618
3.607
2.941
3.485
3.415
4.155
TriviaQA (10M) [ 20 ]
10M
4.133
4.414
4.273
4.313
4.375
4.391
4.188
4.430
4.367
4.621
NarrativeQA [ 23 ]
538K
1.611
2.567
2.860
3.638
3.492
3.536
1.959
2.628
2.655
3.395
Table 2: Comparison of MSA with same-backbone RAG baselines (Qwen3-4B) on LLM judge results (scale 0-5). Higher scores indicate better performance. "RR" denotes RAG systems that perform both retrieval and reranking.
KaLMv2 + Qwen3-235B
KaLMv2 + Qwen3-235B (RR)
KaLMv2 + Llama3.3
KaLMv2 + Llama3.3 (RR)
MSA (Ours)
Dataset
R@1
R@5
R@10
R@1
R@5
R@10
R@1
R@5
R@10
R@1
R@5
R@10
@adaptive
MS MARCO v1
2.846
3.028
3.027
2.886
3.020
2.995
2.649
2.904
2.919
2.881
2.955
2.952
4.141
Natural Questions
3.711
3.670
3.694
3.621
3.610
3.645
3.675
3.674
3.662
3.756
3.665
3.647
3.545
DuReader
4.044
3.991
3.978
3.973
3.932
3.891
4.051
3.846
3.742
3.967
3.776
3.780
4.155
TriviaQA (10M)
4.367
4.656
4.578
4.492
4.320
4.555
4.273
4.740
4.719
4.547
4.703
4.695
4.621
NarrativeQA
1.413
2.130
2.427
3.212
3.427
3.375
1.290
2.123
2.382
3.150
3.263
3.317
3.395
Table 3: Comparison of MSA with SOTA RAG systems using large-scale backbones on LLM judge results (scale 0-5). Higher scores indicate better performance. "RR" denotes RAG systems that perform both retrieval and reranking.
Figure 4: Results on the "Needle In A Haystack" (NIAH) evaluation across varying context lengths from 32k to 1M tokens.
Model Variant
Average
MS MARCO v1
Natural Questions
DuReader
HotpotQA
MSA-S2 (Full)
3.976
4.141
3.545
4.155
4.061
MSA-S1 (Full)
3.694
3.197
3.493
4.064
4.020
w/o memory interleave
3.497
3.175
3.485
4.076
3.250
w/o continual pre-training
2.537
2.267
2.448
3.144
2.289
w/o original text
2.325
2.625
2.190
2.186
2.297
Table 4: Ablation study on four QA benchmarks. We compare the full MSA-S1 model against three variants: removing multi-round interleaved retrieval ( w/o multi-round ), skipping continual pre-training with auxiliary routing supervision ( w/o pretrain ), and disabling loading of original document text after document ID generation ( w/o original text ). Additionally, we compare MSA-S2 and MSA-S1 to demonstrate the effect of curriculum learning. All scores are reported on a 0–5 quality scale.
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 9
Dataset Source (Filename)
Queries
Tokens
Task / Domain
Long-Context & Instruction Tuning
kalmfinetune_data
5,801,540
6.46B
Knowledge Augmentation
Academic & Scientific Literature
S2ORC_citations_abstracts
500,000
7.31B
Scientific Literature
S2ORC_citations_titles
500,000
7.18B
Scientific Literature
S2ORC_title_abstract
500,000
7.11B
Scientific Literature
Appendix
Table 5: Detailed statistics of the full MSA Pre-training Dataset.
Memory-augmented language models combine persistent or temporary memory with neural computation to retain, retrieve, and use information beyond the active context.
Long-context reasoning remains a critical bottleneck for large language models, as recent recurrent-memory approaches face two inherent challenges: sequential chunk-wise updates can overwrite early critical evidence with later irrelevant content, and serial inter-chunk dependencies limit parallelism and cause latency to increase with context length. To address these issues, we propose PI-Mem (Parallel-Iterative Memory), a mechanism that processes all chunks in parallel and iteratively refines a shared memory over a bounded number of turns. In each turn, PI-Mem reads all chunks in parallel conditioned on the current memory, selects new or complementary evidence from each chunk, and merges the selected evidence into a compact shared memory for the next turn. To discourage redundant turns, we optimize the workflow through reinforcement learning with an auxiliary turn-efficiency reward, enabling the model to adaptively exit once sufficient evidence has been accumulated. We evaluate PI-Mem with Qwen3.5-35B-A3B and Qwen2.5-7B on the HotpotQA benchmark across context lengths up to 3.6 million tokens and find that it outperforms the recurrent-memory baseline by +6.25 and +7.81 absolute points while achieving 6.1× and 2.1× inference speedups, respectively. These results demonstrate that PI-Mem breaks the accuracy--efficiency trade-off in long-context reasoning and provides a scalable approach to complex multi-hop question answering over extremely long documents.
Dawei Liu, Haixu Song, Shuang Cheng +9
Shanghai Jiao Tong University · Shanghai Artificial Intelligence Laboratory · Tsinghua University +3
Recurrent memory agents extend LLMs to arbitrarily long contexts by iteratively consolidating input into a fixed-size memory window. Despite their scalability, these agents exhibit a well-documented reliability problem: end-to-end performance degrades systematically as context length grows. We diagnose this failure by decomposing performance into two factors--memory capture and memory retention--and quantitatively confirm that retention is the dominant bottleneck. Retention collapses because existing designs maintain memory as a monolithic text block, forcing every update to risk overwriting previously retained content. Motivated by this diagnosis, we propose Multi-Head Recurrent Memory (MHM), a general, training-free framework that partitions memory into independent heads governed by a stage-wise select-then-update strategy. At each step, exactly one head is selected for update while the remaining heads are structurally shielded from overwriting, shifting the burden of retention from model behavior to architectural design. As a lightweight instantiation, we introduce Least-Recently-Updated MHM (MHM-LRU), which guarantees uniform head utilization with zero additional token overhead. Extensive experiments on long-context benchmarks show that MHM-LRU substantially improves both retention and end-to-end accuracy across the 100K--1M token range, where baselines degrade sharply. On RULER-HQA at 896K tokens, MHM-LRU improves the memory retention rate from less than 30% to 73.96%. These gains generalize across model families, scales, and task types, positioning architectural optimization as a practical and cost-efficient path toward reliable long-context recurrent memory.
Jiatong Li, Samuel Yeh, Sharon Li
Department of Computer Science, University of Wisconsin-Madison
Language models struggle to generalize beyond pretraining context lengths, limiting long-horizon reasoning and retrieval. Continued pretraining on long-context data can help but is expensive due to the quadratic scaling of Attention. We observe that most tokens do not require (Global) Attention over the entire sequence and can rely on local context. Based on this, we propose L2A (Learning To Attend), a layer that enables conditional (token-wise) long-range memory access by deciding when to invoke global attention. We evaluate L2A on Qwen 2.5 and Qwen 3 models, extending their effective context length from 32K to 128K tokens. L2A matches the performance of standard long-context training to within 3% while skipping Global Attention for ∼80% of tokens, outperforming prior baselines. We also design custom Triton kernels to efficiently implement this token-wise conditional Attention on GPUs, achieving up to ∼2× improvements in training throughput and time-to-first-token over FlashAttention. Moreover, L2A enables post-training pruning of highly sparse Global Attention layers, reducing KV cache memory by up to 50% with negligible performance loss. Our code is released under Apache 2.0 at https://github.com/awslabs/hybrid-model-factory/tree/main/examples/research/L2A.