MSA-4B

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. Code
pip install -U "huggingface_hub==0.31.4"
huggingface-cli download Anoy123423123/MSA-Code --repo-type=model --local-dir MSA
cd MSA

# 2. Environment (transformers must be exactly 4.51.3)
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

# 3. These weights
huggingface-cli download --resume-download Anoy123423123/MSA-4B --local-dir ckpt/MSA-4B

# 4. Reproduce the benchmarks (data is fetched automatically on first run)
bash scripts/run_benchmarks.sh eval_benchmark

# 5. LLM-judge scores
export OPENROUTER_API_KEY=<your key>
bash 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.