MSA: Memory Sparse Attention for Efficient End-to-End Memory Model Scaling to 100M Tokens
Organizations: Evermind · Shanda Group · Peking University
Abstract
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
| 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 | Easy | Low | |
| MemAgent | Yes | Medium | Medium | Medium | Easy | Medium |
| 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 |
| 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 |
| 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 |
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
| 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 |
Memory-Augmented Language Models
Memory-augmented language models combine persistent or temporary memory with neural computation to retain, retrieve, and use information beyond the active context.
See all 291 papers