A scalable, end-to-end trainable latent-memory framework for 100M-token contexts

arXiv • Zenodo

中文

Long-term memory is essential for general intelligence, yet the full attention bottleneck constrains most LLMs' effective context length to 128K–1M. Existing attempts，hybrid linear attention, fixed-size state memory (e.g., RNNs), and external storage like RAG/agents，either suffer rapid precision decay and latency growth at extreme scales, lack end-to-end differentiability or dynamic memory maintenance, or require complex pipelines. We present Memory Sparse Attention (MSA): an end-to-end trainable, scalable sparse latent-state memory framework. Core ideas include:

• Scalable sparse attention + document-wise RoPE (parallel/global) achieving near-linear complexity in both training and inference;
• KV cache compression with a Memory Parallel inference engine to deliver 100M token throughput on 2×A800 GPUs;
• Memory Interleave for multi-round, multi-hop reasoning across scattered memory segments.

On long-context QA and NIAH (Needle-in-a-Haystack) benchmarks, MSA surpasses same-backbone RAG, best-of-breed RAG stacks, and leading long-context models. Across an unprecedented 16K→100M token range, MSA shows < 9% degradation, suggesting a practical path to decouple memory capacity from reasoning.

Scaling from 16K→100M tokens: MSA fuses top-k selection with sparse attention to remain end-to-end differentiable while allowing document decoupling at inference. On MS MARCO, MSA sustains <9% degradation and exhibits strong extrapolation. Some baseline curves end early due to their context limits.

Figure 1: MSA scalability under extreme-long contexts

• Memory-Sparse Attention (MSA): an end-to-end trainable, scalable sparse attention layer with document-wise RoPE, realizing O(L) complexity and <9% degradation from 16K→100M tokens.
• KV cache compression + Memory Parallel: tiered storage (GPU-resident routing keys, CPU content K/V), distributed scoring, and on-demand transfers to enable 100M-token inference on 2×A800.
• Memory Interleave: adaptive alternating "generative retrieval → context expansion → generation," significantly boosting multi-hop reasoning across documents.
• Comprehensive evaluation: MSA outperforms same-backbone RAG, best-of-breed RAG pipelines, and top long-context models on long-context QA and NIAH, showing superior stability and accuracy at scale.

MSA integrates retrieval and generation into a single differentiable loop. Document latent states (K/V/Kᵣ) are chunk-mean pooled for compression. A router projector computes relevance via cosine similarity (mean-pooled over heads, then token-wise max), selects Top‑k documents, then concatenates their compressed K/V with the query's local K/V for autoregressive decoding. Routing applies only to upper layers; lower layers keep independent document processing for hierarchical alignment.

• Parallel (document-wise) RoPE: Each document resets positions from 0, preventing position drift between train-short and infer-long, enabling 64k training to extrapolate to 100M.
• Global RoPE (active context): The query's starting index is offset by k (Top‑k retrieved blocks), preserving causal ordering: background → query → generation.

Figure 2: MSA layer (sparse attention + document-wise RoPE)

Figure 2: Memory-Sparse Attention layer and parallel/global RoPE

MSA uses a three-stage pipeline (Fig. 3):

1. Global Memory Encoding (offline): forward over the corpus to cache chunk-pooled (K̄, V̄, K̄ᵣ).
2. Online Routing & Context Assembly: project query to Qᵣ, match with K̄ᵣ to pick Top‑k, then load only the selected K̄/V̄ and concatenate with local context.
3. Sparse Generation: autoregress over the sparse context.

Memory Parallel shards K̄ᵣ across GPUs (query broadcast → local scoring → global reduce). Content K̄/V̄ stays in host DRAM and is asynchronously fetched when selected—balancing VRAM and throughput for 100M-token deployment.

Figure 3: Three-stage inference & Memory Interleave

Figure 3: Offline encoding → online routing → sparse generation; optional multi-round interleave for multi-hop

Setup QA: 9 datasets (MS MARCO v1, NQ, DuReader, TriviaQA(10M), NarrativeQA, PopQA, 2WikiMultiHopQA, HotpotQA, MuSiQue), memory banks 277K→10M tokens, metric: LLM judge (0–5). NIAH (RULER): 8 subtasks, 32K→1M tokens, report average accuracy. Backbone: Qwen3‑4B‑Instruct‑2507. Compare to same-backbone RAG and best-of-breed RAG stacks (KaLMv2 + large generators, optional reranker).

Summary: Average 3.760, improving over standard RAG (+16.0%), RAG+rerank (+11.5%), and HippoRAG2 (+14.8%) using their best@k; MSA leads on all but NarrativeQA within the same-backbone group.

Table 1: Same-backbone RAG vs MSA (@1/@5/@10 vs MSA @adaptive)

Summary: Against KaLMv2+Qwen3‑235B and KaLMv2+Llama‑3.3‑70B (w/ and w/o reranking), MSA achieves the best score on 4/9 datasets and an average 3.760, with relative gains of +7.2%, +5.0%, +10.7%, and +5.4% over the strongest configurations respectively. Gaps on a few datasets (e.g., MuSiQue) are largely attributable to parameter-count and intrinsic reasoning capacity.

Table 2: SOTA RAG stacks (strong retriever + large generator + optional reranker) vs MSA

Summary: MSA maintains 94.84% at 1M tokens. The unmodified backbone collapses beyond 128K (down to 24.69% @1M). Hybrid linear-attention long-context models degrade noticeably at ≥128K/256K. External-memory agents (e.g., RL‑MemoryAgent‑14B) remain stable but are weaker in absolute accuracy and show steeper decay than MSA.

Figure 4: Accuracy vs context length (higher is better)

• Training: 158.95B-token continuous pretraining with auxiliary routing loss, followed by two-stage SFT (8k→64k curriculum).
• Ablations (paper Table 4): curriculum extension, Memory Interleave, continuous pretraining, and injecting original text all contribute substantially; removing them causes 5%–37% drops depending on task.

For full details (project structure, supported benchmarks, etc.), see QUICK_START.md.

1. Install

conda create -n msa python=3.12 -y && conda activate msa
pip install -r requirements.txt
pip install flash-attn==2.7.4.post1 --no-build-isolation

2. Download model

mkdir ckpt
huggingface-cli download --resume-download EverMind-AI/MSA-4B --local-dir ckpt/MSA-4B

3. Download benchmarks

Benchmark data is hosted on EverMind-AI/MSA-RAG-BENCHMARKS and will be automatically downloaded to data/ on first run.

4. Run

# Run inference on benchmarks
bash scripts/run_benchmarks.sh eval_benchmark

# Compute LLM-based scores
bash scripts/calculate_llm_score.sh eval_benchmark

@misc{chen2026msamemorysparseattention,
      title={MSA: Memory Sparse Attention for Efficient End-to-End Memory Model Scaling to 100M Tokens},
      author={Yu Chen and Runkai Chen and Sheng Yi and Xinda Zhao and Xiaohong Li and Jianjin Zhang and Jun Sun and Chuanrui Hu and Yunyun Han and Lidong Bing and Yafeng Deng and Tianqiao Chen},
      year={2026},
      eprint={2603.23516},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2603.23516},
}

This repository and documentation page are maintained by the MSA authors. For project updates, please visit the Homepage: https://evermind.ai/