Research map / Architectures and efficiency / Attention and transformers
Can Mamba Learn How To Learn? A Comparative Study on In-Context Learning Tasks
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.
Evaluates in-context learning of Mamba and other SSMs against Transformers trained from scratch on synthetic ICL tasks, and proposes MambaFormer, a hybrid of Mamba and attention without positional encoding.
Problem
Whether attention-free state-space models such as Mamba can perform in-context learning, which has mainly been studied for Transformers, and how they compare on synthetic ICL tasks.
Approach
Train SSMs (Mamba, S4, S4-Mamba) and Transformers from scratch on a suite of synthetic ICL tasks at matched FLOPs, then introduce MambaFormer, which interleaves Mamba and attention blocks and starts with a Mamba block so no positional encoding is needed; this combines Mamba's sparse-parity ability with the Transformer's retrieval ability.
Key results
- Mamba matches Transformer on standard regression ICL tasks and succeeds at sparse parity (d=10, k=2) where Transformers stay near random guessing up to 768 dim and 24 layers.
- On vector MQAR, Mamba has test MSE around 0.15-0.86 versus about 1e-7 to 1e-4 for Transformers, while MambaFormer nearly closes the gap to the Transformer.
- On RegBench (120 epochs), MambaFormer reaches 99.8% accuracy versus 98.9% for Transformer and 87.8% for Mamba; on GINC, LSTM leads with 96.4% ICL accuracy and Mamba gets 87.1%.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add an attention layer within the Mamba blocks (ablation) | Mamba · Vector MQAR | test loss (MSE) | Helped |
| Add an initial Mamba block to the hybrid architecture (ablation) | Vanilla hybrid · Sparse parity | convergence time | Helped |
| Add an initial Mamba block without positional encoding (ablation) | Transformer · Sparse parity | parity learning efficiency | Helped |
| Use Mamba instead of Transformer | ICL models · Sparse parity | accuracy | Helped |
| Use Mamba instead of Transformer | ICL models · Vector MQAR | test loss (MSE) | Hurt |
| Use Mamba instead of Transformer | ICL models · Orthogonal-outlier regression | performance | Hurt |
| Use Mamba instead of Transformer | ICL models · Many-outlier regression | OOD accuracy | Helped |
| Replace Mamba with MambaFormer | ICL models · Vector MQAR | test loss (MSE): 8.64e-1 -> 7.30e-6 | Helped |
| Replace Mamba with MambaFormer | ICL models · GINC, 64 examples | ICL accuracy: 87.1 ± 7.8 -> 79.6 ± 3.8 | Hurt |
| Replace Mamba with MambaFormer | ICL models · RegBench, 15 epochs | accuracy: 69.4 -> 99.8 | Helped |
| Replace Mamba with MambaFormer | ICL models · RegBench, 120 epochs | accuracy: 87.8 -> 99.8 | Helped |
| Replace Adam with AdamW plus a learning-rate scheduler (ablation) | Mamba and MambaFormer · Linear regression | performance | Helped |
Surprising findings
- Transformers fail to learn sparse parity even at 768 embedding dimension and 24 layers, while Mamba and S4-Mamba solve it with a 2-layer model.
- A Mamba block placed as the first layer makes a Transformer solve sparse parity nearly as efficiently as Mamba, without positional encoding.
Limitations
- Experiments are on small synthetic tasks and small models; language ICL is only tested on synthetic formal language benchmarks (GINC, RegBench).
- Mamba underperforms Transformers on vector MQAR and decision trees, and results are sensitive to learning rate and gradient clipping.
- The paper does not explain why SSMs learn sparse parity, leaving the mechanism for future study.
Experimental setup
Datasets: GINC RegBench
Benchmarks: linear regression sparse linear regression 2NN regression decision tree orthogonal-outlier regression many-outlier regression sparse parity chain-of-thought I/O vector MQAR GINC RegBench
Metrics: squared error loss classification accuracy mean squared error perplexity greedy-decoding accuracy convergence time
Baselines: Transformer Mamba S4 S4-Mamba LSTM Transformer with PE Transformer without PE
Builds on: Mamba S4 Transformer in-context learning regression (Garg et al.) GINC benchmark RegBench
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)
- Trained Mamba Emulates Online Gradient Descent in In-Context Linear Regression (NeurIPS 2025 · similarity 0.89)
- Mamba Modulation: On the Length Generalization of Mamba Models (NeurIPS 2025 · similarity 0.89)
- Mamba Can Learn Low-Dimensional Targets In-Context via Test-Time Feature Learning (NeurIPS 2026 · similarity 0.89)
- Block-Biased Mamba for Long-Range Sequence Processing (NeurIPS 2025 · 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.