MSA-4B

Highlights

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

Model Overview

This model is based on Qwen3-4B-Instruct-2507 with Memory Sparse Attention (MSA).

• Number of Parameters: 4.0B
• Number of Layers: 36
• Number of MSA Layers: 18
• Number of Attention Heads (GQA): 32 for Q and 8 for KV
• Based on Qwen/Qwen3-4B-Instruct-2507

For more details, including benchmark evaluation, hardware requirements, and inference performance, please refer to our paper and GitHub.

Performance

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).

Table 1: MSA vs same-backbone RAG (Qwen3‑4B)

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)

Table 2: MSA vs best-of-breed RAG (large backbones)

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

Quick Start

1. Download code from our Github and install dependencies

git clone https://github.com/EverMind-AI/MSA

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
pip install -U huggingface_hub
export HF_ENDPOINT=https://hf-mirror.com
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, based on the benchmarks specified in scripts/run_benchmarks.sh. No manual download is needed.

4. Run inference on benchmarks

bash scripts/run_benchmarks.sh eval_benchmark

5. Compute LLM-based scores

bash scripts/calculate_llm_score.sh eval_benchmark

Supported Benchmarks

Citation

If you find our work helpful, feel free to give us a cite.

@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},
}

Acknowledgments

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