Research map / Architectures and efficiency / Attention and transformers
Transformers to SSMs: Distilling Quadratic Knowledge to Subquadratic 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.
Distills a pretrained Phi-1.5 Transformer into Mamba-2-based Phi-Mamba (and a hybrid with 4 attention layers) via three MOHAWK stages (matrix orientation, hidden-state alignment, weight transfer + logit distillation) using only 3-5B C4 tokens.
Problem
Subquadratic architectures such as Mamba are pretrained with far fewer resources than strong Transformers, and it is unclear how to transfer a Transformer's learned knowledge into an SSM of similar size.
Approach
View attention and SSMs as matrix mixers and distill progressively: first match mixing matrices, then block hidden states, then end-to-end logits with transferred MLP/embedding weights, yielding a Mamba-2 student distilled from Phi-1.5 with under 1% of pretraining tokens.
Key results
- Phi-Mamba-1.5B (3B C4 tokens) reaches 62.6 average across six tasks vs. 59.7 for Mamba-1-1.4B trained on 315B tokens, with Phi-1.5 teacher at 64.9.
- Hybrid Phi-Mamba-1.5B (4 attention layers, 5B tokens) reaches 66.0 average on five tasks vs. Phi-1.5's 67.2 and Samba-1.7B's 65.4.
- Ablation shows Stage 2+3 gives 63.1 vs. 59.9 for Stage 3 alone on Phi-to-Phi distillation; Stage 3 alone is less effective than the full pipeline.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Distill Phi-1.5 into Phi-Mamba using MOHAWK | Phi-Mamba-1.5B · Downstream average | accuracy: 59.7 -> 62.6 | Helped |
| Add end-to-end Stage 3 distillation after Stage 2 (ablation) | Phi-Mamba · Downstream average | accuracy: 53.3 -> 62.3 | Helped |
| Add Stage 1 matrix orientation to stages 2-3 (ablation) | Phi-Mamba · Downstream average | accuracy: 62.3 -> 62.7 | Helped |
| Use Stage 3 alone instead of the full three-stage process (ablation) | Phi-Mamba · Downstream average | accuracy: 62.7 -> 54.5 | Hurt |
| Replace the SSD mixer with a causal Toeplitz mixer (ablation) | Phi-Mamba · Downstream average | accuracy: 59.2 -> 35.0 | Hurt |
| Replace the SSD mixer with a causal low-rank mixer (ablation) | Phi-Mamba · Downstream average | accuracy: 59.2 -> 36.4 | Hurt |
| Retain four attention layers rather than one (ablation) | Hybrid Phi-Mamba · Downstream average | accuracy: 61.8 -> 64.0 | Helped |
| Retain four attention layers rather than two (ablation) | Hybrid Phi-Mamba · Downstream average | accuracy: 62.3 -> 64.0 | Helped |
| Freeze all components except Mamba-2 blocks (ablation) | Phi-Mamba · Downstream average | accuracy: 62.7 -> 61.4 | Hurt |
| Use Mamba-2 in place of attention blocks | Phi-Mamba-1.5B · LAMBADA | accuracy: 53.4 -> 50.1 | Hurt |
| Use Mamba-2 in place of attention blocks | Phi-Mamba-1.5B · HellaSwag | accuracy: 62.6 -> 60.2 | Hurt |
Surprising findings
- Freezing everything except the Mamba-2 mixers during Stage 3 causes only a small drop (Phi-Mamba 62.7 to 61.4 average).
- Phi-Mamba distilled on C4 with 3B tokens outperforms Samba, which was trained on Phi-2's dataset with more parameters and 3x attention layers, on some tasks.
Limitations
- Phi-Mamba still lags the Phi-1.5 teacher on most benchmarks (62.6 vs 64.9 average).
- Stage 3 training showed loss spikes requiring a reduced learning rate and stabilization tricks.
- Hybrid model performance is below the pure Transformer teacher and the authors suspect the distillation process for hybrids could be improved.
Experimental setup
Datasets: C4 WikiText-103 OpenWebText IMDB
Benchmarks: WinoGrande ARC-Easy ARC-Challenge PIQA HellaSwag LAMBADA
Metrics: accuracy (normalized for ARC-C and HellaSwag) Frobenius distance L2 block distance perplexity
Baselines: Mamba-1-1.4B Mamba-2-1.3B Finch-1.6B xLSTM-1.4B Eagle-1.5B Pythia-1.4B RWKV4-1.5B DeltaNet-1.3B GLA-1.3B Samba-1.7B Mamba-SWA-MLP-1.6B
Builds on: Mamba-2 Phi-1.5 knowledge distillation structured state space duality
Closest papers by meaning
- Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality (ICML 2024 · similarity 0.92)
- Bigram Subnetworks: Mapping to Next Tokens in Transformer Language Models (NeurIPS 2025 · similarity 0.89)
- The Mamba in the Llama: Distilling and Accelerating Hybrid Models (NeurIPS 2024 · similarity 0.89)
- Mamba Modulation: On the Length Generalization of Mamba Models (NeurIPS 2025 · similarity 0.89)
- Trainable Transformer in Transformer (ICML 2024 · similarity 0.88)
- Transformer Block Coupling and its Correlation with Generalization in LLMs (ICLR 2025 · similarity 0.88)
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.