Source-linked AI summary
MSA: Memory Sparse Attention for Efficient End-to-End Memory Model Scaling to 100M Tokens
Yu Chen, Runkai Chen, Sheng Yi, Xinda Zhao, Xiaohong Li, Jianjin Zhang, Jun Sun, Chuanrui Hu, Yunyun Han, Lidong Bing, Yafeng Deng, Tianqiao Chen
TL;DR
Long-term LLM memory remains limited by context capacity and tradeoffs among precision, efficiency, dynamic memory management, and end-to-end optimization. MSA uses scalable sparse attention with document-wise RoPE, KV-cache compression, Memory Parallel, and Memory Interleave to extend intrinsic memory. It maintains less than 9% degradation from 16K to 100M tokens and outperforms leading LLMs, RAG systems, and memory agents on long-context benchmarks.
Problem
Existing memory approaches do not jointly provide scalable, high-fidelity, architecturally compatible, and end-to-end trainable lifelong memory.
Method
MSA combines differentiable sparse attention, document-wise RoPE, KV-cache compression, Memory Parallel, and Memory Interleave for lifetime-scale intrinsic memory.
Results
Less than 9% degradation occurs from 16K to 100M tokens, while MSA surpasses frontier LLMs, advanced RAG systems, and memory agents on long-context benchmarks.
Takeaways & Limitations
MSA provides a scalable framework for end-to-end modeling of lifetime-scale contexts while preserving retrieval fidelity and reasoning depth.
Takeaways & Limitations
MSA remains limited when tasks require strong, tightly coupled dependencies across multiple documents, where intrinsic memory may not preserve accurate structural alignment.
Abstract
from arXiv · showhide
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.
1 Introduction
Long-term LLM memory must scale toward hundreds of millions of tokens while preserving precise, trainable retrieval, but existing paradigms trade off capacity, efficiency, or end-to-end optimization. MSA addresses this with differentiable sparse attention and related memory mechanisms, maintaining less than 9% degradation from 16K to 100M tokens and outperforming established alternatives.
- Motivation: Models must process hundreds of millions of tokens to approach estimated human-scale lifelong memory, far beyond contemporary effective context lengths of 128K to 1M tokens.The paper connects this requirement to applications including Digital Twins and long-term multi-agent histories.
- Motivation: An effective lifelong-memory system requires architectural compatibility, scalable low-overhead computation, precise end-to-end retrieval and storage, manageable memory, and resistance to catastrophic forgetting.These desiderata define the paper’s target for long-term memory systems.
- Limitations of Existing Approaches: Existing memory paradigms separately address subsets of these criteria, leaving high-fidelity methods capacity-limited and scalable methods less reliable.The paper identifies limited scalability of high-fidelity memory and the lack of a fully differentiable, jointly optimized memory pipeline as fundamental limitations.
- MSA: MSA combines differentiable top-k selection, sparse attention, document-wise RoPE, KV-cache compression, Memory Parallel, and Memory Interleave for scalable intrinsic memory.These components target efficient scaling, extrapolation, deployment, and multi-hop reasoning across distributed memory segments.
- Results: Less than 9% degradation occurs when scaling from 16K to 100M tokens, while MSA outperforms frontier LLMs, RAG systems, and memory agents on long-context evaluations.The reported comparisons include long-context QA and Needle-In-A-Haystack benchmarks.
2 Related Work
Related work organizes long-term LLM memory into parameter-based, external-storage, and latent-state paradigms. These approaches differ in where memory is stored and how it is retrieved or compressed.
- Parameter-Based Memory: Parameter-based memory internalizes external information through parameter updates or specialized architectures, integrating knowledge deeply into the model.Examples include Continuous Pre-training, LoRA, and test-time-trained architectures.
- External Storage-Based Memory: External storage-based memory retrieves relevant memories from large databases, with RAG using vector similarity to select textual chunks.Later methods address noise and precision limitations in dense retrieval.
- Latent State-Based Memory: Latent state-based memory manages memory in the model’s internal latent states, including KV-pair representations that may still rely on external retrieval.Memory³ is described as an example constrained by the size of active KV caches.
- Linear Attention: Linear-attention memory compresses preceding context into compact recurrent hidden states instead of explicitly accessing previous key-value pairs.RWKV and DeltaNet exemplify this recurrence-based approach.
3 Memory Sparse Attention
MSA integrates sparse retrieval and answer generation into an end-to-end trainable latent-memory architecture. It compresses document chunks, routes queries to top-k documents, and uses document-wise positional encoding to support scaling across memory sizes.
- MSA jointly optimizes sparse memory retrieval and answer generation in a single latent-memory architecture.
- Sparse Attention Mechanism: Chunk-wise mean pooling compresses document Key, Value, and routing-key states into latent representations, reducing memory footprint and retrieval complexity.
- Sparse Attention Mechanism: MSA computes cosine-similarity relevance scores between routing queries and compressed routing keys, then selects the Top-k documents.
- Sparse Attention Mechanism: The selected compressed Key and Value matrices are concatenated with the query’s local cache to form the sparse context for autoregressive generation.
- Sparse Attention Mechanism: MSA applies routing selectively to the latter half of model layers, while lower layers use local context without attending to compressed memory KVs.
- Parallel and Global RoPE: Independent document-wise RoPE decouples positional semantics from the total document count, while Global RoPE offsets active-context positions by retrieved documents.
3.3 Training
MSA training combines generative retrieval, auxiliary routing supervision, and staged optimization. A two-stage SFT curriculum first builds instruction following and reasoning, then extends context length and improves data quality.
- Continuous Pre-training: Continuous pre-training on 158.95 billion deduplicated tokens trains the model to autoregressively generate unique IDs for relevant documents.
- Continuous Pre-training: The auxiliary loss supervises layer-wise Top-k routing decisions so relevant document chunks are separated from irrelevant ones in latent routing space.
- Continuous Pre-training: The optimization schedule uses L = 0.1LLM + Laux during warm-up, then L = LLM + 0.1Laux after the learning rate is annealed to 6e-6.
- SFT Curriculum: The first SFT stage uses 8k-token contexts to establish instruction-following and reasoning capabilities within a standard context window.
- SFT Curriculum: The second SFT stage cleans the data and extends memory context from 8k to 64k tokens to improve robustness when extrapolating to massive memory banks.
3.4 Inference
MSA inference uses offline document encoding, online routing and context assembly, and sparse autoregressive generation. Memory Parallel and tiered storage distribute routing and content states to support 100M-token inference on two A800 GPUs.
- Three-Stage Inference Process: The three-stage inference pipeline performs offline global memory encoding, online routing and context assembly, then sparse generation.
- Global Memory Encoding: Offline encoding computes K, V, and routing-key matrices for each document, compresses them by mean pooling, and caches the compact representations.
- Routing and Context Assembly: Online routing matches the query’s routing query against cached routing keys, selects Top-k documents, and loads only their compressed K and V matrices.
- Sparse Generation: Sparse generation applies standard attention between the active token’s query and the selected compressed KV pairs plus the local query cache.
- Memory Parallel: 100 million tokens can be processed with 2 NVIDIA A800 GPUs using Memory Parallel and a tiered strategy that distributes routing keys and CPU-offloads content KVs.
- Memory Parallel: 169GB of compressed KV and routing-key cache is estimated for 100M tokens, exceeding the 160GB aggregate capacity of a standard 2×A800 node under monolithic storage.
3.5 Memory Interleave
Memory Interleave iteratively alternates retrieval and context expansion to support multi-hop reasoning across scattered memory segments.
- Memory Interleave: Memory Interleave repeats routing and context assembly with sparse generation instead of relying on single-shot retrieval.
- Memory Interleave: Retrieved documents become part of the next query, enabling iterative generative retrieval and context expansion for complex multi-hop questions.
- Memory Interleave: During training, each multi-hop retrieval chain is divided into single-step samples based on either the query or existing document context.
4 Experiment
MSA is evaluated on diverse QA and NIAH tasks against same-backbone and best-of-breed RAG systems, long-context models, and memory agents. Results show strong QA gains, stable retrieval through 1M tokens, and measurable benefits from curriculum learning and key architectural components.
- Evaluation setup: MSA is evaluated on nine QA benchmarks spanning single-hop, multihop, and long-context scenarios, with memory banks from 277K to 10M tokens.The evaluation includes standard RAG systems with fixed retrieval depths k = {1, 5, 10}.
- QA results: MSA achieves state-of-the-art performance on eight of nine datasets against same-backbone RAG systems, with average improvements of 16.0%, 11.5%, and 14.8%.The comparisons are against standard RAG, RAG with reranking, and HippoRAG2, respectively.
- QA results: MSA secures top performance on four of nine best-of-breed RAG datasets and achieves a competitive average score of 3.760.Relative improvements over the strongest listed configurations are 7.2%, 5.0%, 10.7%, and 5.4%, respectively.
- NIAH results: 94.84% accuracy at 1M tokens demonstrates stable NIAH retrieval for MSA across 32k-to-1M-token contexts.The unmodified Qwen3-4B-Instruct backbone falls to 48.16% at 256k and 24.69% at 1M tokens.
- NIAH results: 80.78% accuracy at 1M tokens is reached by Qwen3-Next-80B-A3B after substantial decay beyond 256k tokens, while Qwen2.5-14B-1M and Qwen3-30B-A3B degrade earlier.Qwen2.5-14B-1M drops to 89.97% at 256k tokens, and Qwen3-30B-A3B drops to 79.13% at 128k tokens.
- Ablation study: 37.1% average performance decline follows disabling original-text integration, compared with 31.3% without continual pre-training, 7.6% from the second curriculum stage, and 5.3% without memory interleave.The largest task-specific drops are 46.2% on DuReader without original text, 43.1% on HotpotQA without continual pre-training, 29.5% on MS MARCO from curriculum removal, and 19.2% on HotpotQA without memory interleave.
5 Analysis
MSA is evaluated as a long-memory system that combines linear scaling with resistance to context degradation. Its training and inference pipelines separate document processing, sparse routing, and generation so that memory-size dependence remains near-linear and preprocessing can be amortized across queries.
- Information Fidelity: MSA targets both O(L) computational efficiency and minimal context degradation as noisy memory grows.The analysis evaluates QA quality from 16K to 100M tokens using MS MARCO and an LLM judge.
- Computational Efficiency: O(L) complexity in training enables MSA to scale with memory size while preserving a sparse-attention pipeline.The dominant O(LG) document-processing term becomes linear in L when L ≫ G.
- Computational Efficiency: Sparse routing costs O(ML/P), while sparse generation depends on query length and fixed hyperparameters rather than total memory size L.Routing scores are computed over L/P pooled chunks, and generation attends only to the query plus k compressed documents.
- Inference Pipeline: O(LG) offline preprocessing is performed once per memory-bank version, unlike conventional attention’s O(L^2) prefill for every query.The cached compressed representations support repeated query handling from the same memory bank.
- Inference Pipeline: O(ML/P) dominates per-query inference and scales linearly with memory size, while autoregressive generation remains independent of L.Generation operates over a sparse context of size M + kG/P for answer length T.
6 Conclusion
MSA combines sparse attention, document-wise RoPE, KV-cache compression, Memory Parallel, and Memory Interleave for lifetime-scale memory processing. The framework reportedly preserves retrieval fidelity and reasoning depth while maintaining high accuracy at 100M-token contexts.
- Conclusion: MSA extends end-to-end modeling to lifetime-scale contexts through sparse attention, document-wise RoPE, and KV-cache compression.Memory Parallel supports fast 100M-token processing, while Memory Interleave addresses reasoning across distributed memory segments.
- Conclusion: MSA surpasses mainstream state-of-the-art general-purpose LLMs on long-context QA and Needle-in-a-Haystack benchmarks.The conclusion attributes the comparison to preserved retrieval fidelity and reasoning depth.
- Conclusion: High accuracy is maintained as effective context scales to 100M tokens, with KV-cache compression reducing memory footprint and latency.The conclusion presents minimal performance degradation under extreme context lengths.
7 Limitations
MSA remains limited on tasks requiring strong, tightly coupled dependencies across multiple documents. Intrinsic memory alone can struggle to preserve structural alignment when evidence is distributed and highly interlinked.
- Scope Boundary: MSA struggles to maintain accurate structural alignment when evidence is distributed and highly interlinked across sources.The limitation concerns tasks with strong, tightly coupled dependencies across multiple documents.
- Scope Boundary: Memory Interleave is a promising mitigation, but its effectiveness depends on more efficient and principled designs preserving inter-document relationships.The proposed direction is not presented as a resolved solution.
A Prompts
The prompt evaluates predicted answers against real answers using a single 0–5 score. The rubric distinguishes exactness, completeness, relevance, redundancy, incompleteness, factual error, and unrelated or nonsensical output.
- Scoring Prompt: 0–5 scoring is based on predicted-answer accuracy, completeness, and relevance to the real answer.The output is restricted to a single number.
- Scoring Rubric: A score of 5 requires an exactly matching, correct answer, while wording differences do not affect factual accuracy.The rubric treats exact correctness as the highest score.
- Scoring Rubric: A score of 4 indicates complete, error-free core information with a small amount of non-critical redundancy.Redundant content is tolerated when it does not affect the answer’s core information.
- Scoring Rubric: A score of 3 denotes correct core information that is slightly incomplete or imprecise.The answer may differ from the real answer in some aspects without containing errors.
- Scoring Rubric: Scores of 2 and 1 represent partial relevance with significant omission or factual inaccuracy, respectively.Score 2 concerns deviation from the query’s core topic; score 1 retains basic relevance but is incorrect.
- Scoring Rubric: A score of 0 marks an answer as unrelated, gibberish, or pure hallucination without logical connection to the real answer.This is the rubric’s lowest category.
B Pre-training Data Composition
The MSA pre-training corpus combines broad-domain data with robust retrieval coverage at large scale. Its distribution is balanced by capping oversized non-KALM datasets while preserving all KALM instruction data.
- 158.95 billion tokens across 17.9 million queries form the diverse MSA pre-training corpus.The corpus is designed to support both robust retrieval capabilities and broad general knowledge.
- The corpus spans domains ranging from scientific literature to general community Q&A.
- Table 5 provides detailed statistics for the full MSA pre-training dataset.
- Datasets outside the KALM suite exceeding 0.5 million queries are downsampled to a maximum of 0.5 million.This rule is used to maintain a balanced data distribution.
- KALM instruction data is retained in its entirety rather than subjected to the non-KALM downsampling cap.