MSA: Memory Sparse Attention for Efficient End-to-End Memory Model Scaling to 100M Tokens

TL;DR

MSA (Memory Sparse Attention) is an end-to-end trainable sparse-attention mechanism in which each transformer layer (in the upper half of the stack) routes the query to a Top-k subset of document-level latent KV chunks via cosine similarity against a learned routing key. Chunked mean-pooling compresses each document’s K/V into compact representations; document-wise (parallel) RoPE resets positions per document so a model trained at 64k extrapolates to 100M tokens; routing keys live in VRAM and content K/V is held in CPU DRAM and asynchronously fetched on demand, enabling 100M-token inference on 2× A800 GPUs. On long-context QA (9 datasets, memory banks 277K–10M tokens) with a Qwen3-4B backbone, MSA beats same-backbone RAG by +16% average and matches or beats best-of-breed RAG stacks (KaLMv2 + Qwen3-235B / Llama-3.3-70B + reranker) using a 4B generator; on RULER NIAH it holds 94.84% at 1M tokens where the base backbone collapses to 24.69%, with <9% degradation across the 16K → 100M range.

Key claims

Existing long-context paradigms each fail one of {capacity, precision, end-to-end trainability, catastrophic forgetting} — parameter-based memory (LoRA / Titans) has limited capacity and forgets, RAG/MemAgent is non-differentiable, KV-centric latent memory (DSA, MemGen) has prohibitive compute at extreme lengths, and linear-attention variants lossily compress into fixed-size state and degrade catastrophically [§1, §2, Table 1].
The MSA layer replaces dense self-attention with a Top-k document-routed sparse attention: a Router Q/K projector produces routing vectors that are cosine-similar-matched against chunk-mean-pooled compressed routing keys, with relevance aggregated by mean-over-heads then max-over-query-tokens, and only the selected documents’ compressed K/V are concatenated before the query’s local KV for autoregressive decoding [§3.2.1, eqs. 1–4].
Routing is applied only in the latter half of layers — early layers do not produce semantic abstractions strong enough for retrieval — but lower layers still run “Independent Document Processing” to maintain hierarchical alignment of memory states [§3.2.1].
Parallel (document-wise) RoPE resets positions to 0 per document while Global RoPE offsets the active context by k (the Top-k count), preserving causal “background → query → generation” ordering and decoupling positional semantics from total document count, which is the load-bearing trick that lets a 64K-trained model extrapolate to 100M [§3.2.2].
Training uses 158.95B-token continuous pretraining with an auxiliary InfoNCE-style contrastive routing loss on positive/negative document pairs (warm-up at lr 5e-5 with router loss dominant; anneal to 1e-5 with generative-retrieval loss dominant), then a two-stage SFT curriculum 8k → 64k [§3.3.1, §3.3.2].
Inference is three-stage — offline global memory encoding (chunk-pooled K/V/Kr cached), online routing & context assembly (Top-k selection over compressed routing keys), and sparse generation — and “Memory Interleave” iterates routing+expansion+generation for multi-hop reasoning, with the per-round retrieval count adaptively determined by the model [§3.4.1, §3.5].
The Memory Parallel engine tiers storage: routing keys K̄r distributed across GPU VRAM (≈56 GB for 100M context), content K̄/V̄ on CPU host DRAM and asynchronously DMA’d to GPU when their document is selected, with tiled scoring matmuls to bound peak memory; the full 4B model is replicated per-GPU to avoid decode-time communication [§3.4.2].
Long-context QA (Qwen3-4B backbone, 9 datasets including MS MARCO 7.34M, TriviaQA 10M, multi-hop 2WikiMultiHopQA / HotpotQA / MuSiQue): MSA averages 3.760 vs same-backbone RAG 3.242 (R@10) — +16.0% over standard RAG, +11.5% over RAG+rerank, +14.8% over HippoRAG2 [Table 1].
Best-of-breed RAG comparison (KaLMv2 retriever + Qwen3-235B or Llama-3.3-70B generators ± reranker): MSA (4B backbone) wins on 4/9 datasets and on the average (3.760 vs 3.506–3.580), with the largest wins on MS MARCO (4.141 vs ≤3.028) and 2WikiMultiHopQA (4.280 vs ≤3.583); MuSiQue and HotpotQA remain reasoning-capacity-bound where the larger generators still win [Table 2].
RULER NIAH stability (32K → 1M): MSA holds 94.84% at 1M, the unmodified Qwen3-4B backbone collapses past 128K to 24.69% at 1M, hybrid linear-attention long-context models degrade noticeably at ≥128K/256K, and external-memory agents (RL-MemoryAgent-14B) remain stable but with lower absolute accuracy and steeper decay [Fig. 4].
Ablations: removing curriculum extension, Memory Interleave, continuous pretraining, or original-text injection causes 5–37% task-dependent drops [§Implementation Notes, Table 4].

Method

MSA decomposes the standard attention path into a memory-side encoding and a query-side retrieval. Each document is independently forward-passed through the backbone; in every upper-half layer, the standard K/V projections produce per-head K/V matrices and an additional Router K Projector (parameters Wr) produces a routing key Kr. K, V, and Kr are then chunked (fixed-length spans) and mean-pooled into compressed matrices K̄, V̄, K̄r. At inference, the query’s hidden state similarly produces Q, K, V (local) plus a Router Q projection Qr; for each document chunk a per-head cosine relevance is computed against K̄r, mean-pooled across heads, max-pooled across query tokens, and the top-k documents are selected (max over chunks gives the document score). Selected documents contribute their compressed K̄, V̄ — concatenated before the query’s local KV in sequence order — to the layer’s attention, and decoding proceeds autoregressively against this sparsity-aware context.

Two RoPEs are applied: document-wise (parallel) RoPE for the retrieved compressed KVs, where each document’s positions start at 0, eliminating the train-short / infer-long position drift; and Global RoPE for the active context, starting at position k so the query sees retrieved documents as causally prior. Routing is restricted to upper-half layers; in lower layers each document is forward-passed independently with no cross-document attention. Training adds a contrastive InfoNCE loss over positive/negative document pairs against the relevance scores (paper eq. 5), summed with the standard generative loss; the schedule warms up the router with high routing-loss weight and large lr, then anneals to favor generation.

Serving 100M-token contexts uses Memory Parallel: routing keys K̄r are sharded across multiple GPUs (each scoring its shard, then a global reduce picks Top-k); content K̄/V̄ stays on CPU and is asynchronously fetched only for selected documents; the (relatively small) 4B model is replicated per-GPU to avoid decode-time all-reduces. Memory Interleave iterates Stage 2 + Stage 3 — the model generates document IDs ending in a special delimiter, the system retrieves those documents, appends their text to the query, and the loop continues until the model switches from emitting IDs to emitting an answer — supporting variable-depth multi-hop chains.

Results

Long-context QA, same backbone (Qwen3-4B), 9 datasets: MSA average 3.760 vs standard RAG R@10 3.242 (+16.0%), Qwen3-4B+reranker R@10 3.372 (+11.5%), HippoRAG2 R@10 3.275 (+14.8%); MSA wins on every dataset except NarrativeQA where the reranker variant is highest [Table 1].
Long-context QA vs best-of-breed RAG (large generators): MSA 3.760 average vs KaLMv2+Qwen3-235B 3.506 (+7.2%), Qwen3-235B+RR 3.580 (+5.0%), KaLMv2+Llama-3.3-70B 3.396 (+10.7%), Llama-3.3-70B+RR 3.568 (+5.4%); MSA wins MS MARCO (4.141), DuReader (4.155), PopQA (3.433), 2WikiMultiHopQA (4.280); loses on TriviaQA, MuSiQue, NarrativeQA, HotpotQA where larger generators help [Table 2].
RULER NIAH at 1M tokens: MSA 94.84%, base Qwen3-4B 24.69%, hybrid-linear long-context models degrade significantly past 128K; <9% MSA degradation across 16K → 100M [Fig. 4, §abstract].
Hardware footprint: 100M-token inference on 2× NVIDIA A800 GPUs (160 GB aggregate); compressed K̄r alone is ≈56 GB at 100M context for 8 heads × 128 dim × 18 layers × BF16 with pool kernel = 64 [§3.4.2].
Sample/training cost: 158.95B-token continuous pretraining + two-stage SFT (8k → 64k curriculum); ablations report 5–37% task-dependent drops when any of {curriculum extension, Memory Interleave, continuous pretraining, original-text injection} is removed [§Implementation Notes, Table 4].

Why it’s interesting

MSA is a useful third datapoint between Inference-Time Hyper-Scaling with KV Cache Compression (DMS — retrofitted per-head KV eviction at 8× compression on a 1K-step budget, no architectural change) and the Long context as weights candidate cluster (End-to-End Test-Time Training for Long Context, Doc-to-LoRA: Learning to Instantly Internalize Contexts — moving context into weights). MSA takes the third path: keep the KV cache, but make retrieval over it differentiable and chunk-compressed, with positional encoding designed specifically for train-short/infer-long extrapolation. The “decouple memory from reasoning” framing matches the agent-memory taxonomy in Rethinking Memory Mechanisms of Foundation Agents in the Second Half: A Survey — MSA is a latent-state external-memory hybrid under that taxonomy, getting RAG-style capacity scaling without RAG’s discrete retrieval bottleneck. The strongest result is on multi-hop QA (2WikiMultiHopQA: 4.280 vs ≤3.583 best baseline), which is where existing RAG pipelines hurt most — Memory Interleave’s iterative retrieval is the closest mechanism the team has filed to a differentiable “deep research loop”. The 100M-token-on-2×A800 deployment story is interesting independently of the algorithmic claims because it makes the engineering envelope concrete: routing keys must live on GPU (56 GB at 100M), content KV asynchronously paged from CPU DRAM.