MSA-4B
Anonymous release for peer review. Inference and evaluation code lives in a separate repository:
Anoy123423123/MSA-Code.
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.
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
This checkpoint cannot be loaded with AutoModel.from_pretrained alone — the architecture is
MSAForCausalLM (model_type: msa), which is implemented in the code repository rather than in
transformers. Follow the steps below.
Quick Start
Full instructions, hardware requirements, results tables and troubleshooting are in the code
repository:
Anoy123423123/MSA-Code.
The short version:
1# 1. Code
2pip install -U "huggingface_hub==0.31.4"
3huggingface-cli download Anoy123423123/MSA-Code --repo-type=model --local-dir MSA
4cd MSA
5
6# 2. Environment (transformers must be exactly 4.51.3)
7conda create -n msa python=3.12 -y && conda activate msa
8pip install -r requirements.txt
9pip install flash-attn==2.7.4.post1 --no-build-isolation
10
11# 3. These weights
12huggingface-cli download --resume-download Anoy123423123/MSA-4B --local-dir ckpt/MSA-4B
13
14# 4. Reproduce the benchmarks (data is fetched automatically on first run)
15bash scripts/run_benchmarks.sh eval_benchmark
16
17# 5. LLM-judge scores
18export OPENROUTER_API_KEY=<your key>
19bash scripts/calculate_llm_score.sh eval_benchmark
Requires 8× 80GB GPUs for the default benchmark suite; the 100M-token ms_100M setting runs on
2× A800 via Memory Parallel.
Related Repositories
License
MIT.
