Research map / Architectures and efficiency / Attention and transformers
Mamba Modulation: On the Length Generalization of Mamba Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Analyzes why pre-trained Mamba models fail at long contexts via the spectrum of the state transition matrix A, and proposes training-free-style calibrated scaling of A (instead of Delta) per layer to improve length generalization.
Problem
Mamba and Mamba2 language models degrade sharply when evaluated on contexts longer than their pre-training length; prior fixes focus on discretization steps Delta, lacking a principled explanation.
Approach
Connects the asymptotic state norm behavior as input length grows to the eigenvalue spectrum of the transition matrix A, and proposes spectrum scaling of A (per-layer scaling factors calibrated with zeroth-order optimization) as more effective than scaling Delta.
Key results
- On PG19, calibrated scaling of A for mamba2-1.3b keeps perplexity between 3.28 and 4.72 from 2K to 64K, versus 1479.45 for the base model at 64K.
- On LongBench with mamba2-2.7b, calibrated scaling of A raises the average score from 8.75 (base) to 14.98, versus 13.60 for MambaExtend.
- Calibrating A consistently gives lower perplexity than calibrating Delta, e.g., up to 1000x lower for mamba2-2.7b at long sequences.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Calibrate per-layer scaling factors for A instead of Δt | mamba2-2.7b · LongBench | Average: 13.60 -> 14.98 | Helped |
| Calibrate per-layer scaling factors for A instead of leaving the model unmodified | mamba2-2.7b · LongBench | Average: 8.75 -> 14.98 | Helped |
| Calibrate per-layer scaling factors for A | mamba-1.4b · PG19 64K | perplexity: 444.09 -> 19.13 | Helped |
| Calibrate per-layer scaling factors for A instead of Δt | mamba-1.4b · PG19 64K | perplexity: 1648.0 -> 19.13 | Helped |
| Use one constant A scaling factor across layers instead of calibrated per-layer factors (ablation) | mamba-1.4b · PG19 64K | perplexity: 19.13 -> 114.44 | Hurt |
| Calibrate per-layer scaling factors for A | mamba-2.8b · PG19 64K | perplexity: 508.88 -> 40.00 | Helped |
| Calibrate per-layer scaling factors for A instead of Δt | mamba-2.8b · PG19 64K | perplexity: 60.00 -> 40.00 | Helped |
| Use one constant A scaling factor across layers instead of calibrated per-layer factors (ablation) | mamba-2.8b · PG19 64K | perplexity: 40.00 -> 355.77 | Hurt |
| Calibrate per-layer scaling factors for A | mamba2-1.3b · PG19 64K | perplexity: 1479.45 -> 4.72 | Helped |
| Calibrate per-layer scaling factors for A instead of Δt | mamba2-1.3b · PG19 64K | perplexity: 27.50 -> 4.72 | Helped |
| Use one constant A scaling factor across layers instead of calibrated per-layer factors (ablation) | mamba2-1.3b · PG19 64K | perplexity: 4.72 -> 13.22 | Hurt |
| Calibrate per-layer scaling factors for A rather than Δt | Mamba and Mamba2 models · Passkey Retrieval | solved length/depth pairs | Mixed |
Surprising findings
- Constant (uncalibrated) scaling of A is generally ineffective for mamba-1.4b/2.8b and mamba2-1.3b on PG19.
- Some calibrated scaling factors greater than 1 still improve generalization, which the authors argue is not explained by prior Delta-based analysis.
Limitations
- Focused only on Mamba-style models; methodology needs adaptation for other state-transition architectures.
- Lack of instruction-tuned Mamba models limits the experiments that can be run.
- Constant scaling factors need tuning and do not work across all layers.
Experimental setup
Datasets: ProofPile PG19 GovReport Pile WikiText-103 LongBench
Benchmarks: ProofPile PG19 GovReport Passkey Retrieval LongBench
Metrics: perplexity passkey retrieval accuracy LongBench F1/accuracy/edit similarity
Baselines: MambaExtend DeciMamba LongMamba Constant Scaling A Base Model
Builds on: Mamba Mamba2 MambaExtend SPSA
Closest papers by meaning
- LongMamba: Enhancing Mamba's Long-Context Capabilities via Training-Free Receptive Field Enlargement (ICLR 2025 · similarity 0.91)
- Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality (ICML 2024 · similarity 0.90)
- Block-Biased Mamba for Long-Range Sequence Processing (NeurIPS 2025 · similarity 0.90)
- DeciMamba: Exploring the Length Extrapolation Potential of Mamba (ICLR 2025 · similarity 0.89)
- Can Mamba Learn How To Learn? A Comparative Study on In-Context Learning Tasks (ICML 2024 · similarity 0.89)
- Transformers to SSMs: Distilling Quadratic Knowledge to Subquadratic Models (NeurIPS 2024 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.