Research map / Architectures and efficiency / Attention and transformers
The Mamba in the Llama: Distilling and Accelerating Hybrid 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 Zephyr-7B and Llama-3 8B Instruct into hybrid Mamba/Mamba2 linear RNNs by reusing attention projection weights, followed by SFT and DPO distillation, and adds a hardware-aware multi-step speculative decoding kernel for linear RNNs.
Problem
Pretrained large Transformer LLMs are slow for long-sequence generation due to quadratic attention and large KV caches, while training linear RNNs from scratch requires enormous compute; the paper asks how to convert pretrained Transformers into efficient linear RNN/hybrid models and how to speed up
Approach
Initialize Mamba/Mamba2 blocks from the attention Q/K/V projections (mapped to C/B/x), freeze the transferred MLPs, replace a fraction of attention layers stepwise, and distill with pseudo-label KL/SeqKD followed by SFT and DPO against the teacher; plus a fused multi-step kernel that enables speculative decoding with a single cached RNN state.
Key results
- Mamba-Llama3 (50% attention) achieves 29.61 length-controlled win rate on AlpacaEval 2 and 7.35 on MT-Bench, versus 22.90 and 8.00 for the Llama-3-Instruct teacher.
- Hybrid speculative decoding on a single RTX 3090 gives about 1.8x speedup for Mamba-Zephyr (50%) and 1.88x for 25% attention with a 2-layer draft.
- Pure Mamba 2.8B speculative decoding reaches 289 tokens/s on a 3090 (2.6x speedup) and pure Mamba 7B reaches 272 tokens/s on H100 (2x speedup).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Increase retained attention layers from 25% to 50% (ablation) | Mamba-Llama3 · MT-Bench | score: 6.86 -> 7.35 | Helped |
| Increase retained attention layers from 25% to 50% (ablation) | Mamba-Llama3 · AlpacaEval 2 | LC win %: $25.85_{1.26}$ -> 29.61$_{1.31}$ | Helped |
| Add both SFT and DPO after distillation (ablation) | Hybrid Mamba (50% attention) · MT-Bench (OpenHermes 2.5 alignment ablation) | score: 5.55 -> 6.69 | Helped |
| Add SFT alone after distillation (ablation) | Hybrid Mamba (50% attention) · MT-Bench (OpenHermes 2.5 alignment ablation) | score: 5.55 -> 5.61 | No clear effect |
| Add DPO alone after distillation (ablation) | Hybrid Mamba (50% attention) · MT-Bench (OpenHermes 2.5 alignment ablation) | score: 5.55 -> 5.42 | No clear effect |
| Initialize projections from attention weights instead of random initialization (ablation) | Zephyr-Mamba (50% attention) · MT-Bench | score: 1.04 -> 6.69 | Helped |
| Initialize projections from attention weights instead of random initialization (ablation) | Hybrid Mamba (50% attention) · LAMBADA | PPL: 55.01 -> 6.20 | Helped |
| Interleave attention layers with Mamba layers (ablation) | Hybrid Mamba (25% attention) · Zephyr perplexity ablation | PPL: 2.89 -> 2.20 | Helped |
| Include Mamba blocks rather than remove them (ablation) | Hybrid Mamba (50% attention) · MT-Bench | score: 1.01 -> 6.69 | Helped |
| Reduce retained attention layers from 50% to 0% (ablation) | Mamba2-Llama3 · LM Eval AVG | score: 63.84 -> 54.74 | Hurt |
| Use a 4-layer rather than 2-layer draft model (ablation) | Mamba-Zephyr (50%) speculative decoding · OpenHermes2.5 | speedup: 1.8x -> 1.81x | Mixed |
Surprising findings
- Removing attention entirely (0%) degrades the distilled model sharply (MT-Bench 5.64 for Mamba2-Llama3), so a fraction of attention layers is needed for chat quality.
- Reusing attention projection weights for initialization is critical: without it, LAMBADA perplexity rises from 6.20 to 55.01 and MT-Bench drops to 1.04.
Limitations
- Pure linear RNN (0% attention) distillation degrades significantly on chat benchmarks.
- Speculative decoding speedups on Llama-3 hybrids are modest because the draft model's large embedding table adds overhead.
- Mamba-Zephyr is strictly subquadratic only in part since Zephyr/Mistral uses sliding window attention.
Experimental setup
Datasets: UltraChat UltraFeedback GenQA InfinityInstruct OpenHermes 2.5 SimPO preference data The Pile
Benchmarks: AlpacaEval 2 MT-Bench LM Eval Harness (WinoGrande, PIQA, HellaSwag, ARC, MMLU, OpenBookQA, TruthfulQA, PubMedQA, RACE) Open LLM Leaderboard ZeroEval (GSM8K, CRUX) Needle in a Haystack
Metrics: length-controlled win rate win rate MT-Bench score accuracy perplexity speedup tokens/second
Baselines: TRI Mamba 7B Nvidia Hybrid Mamba2 8B Falcon Mamba 7B RecurrentGemma-9B Zephyr-7B Llama-3-Instruct 8B Llama-3.1-Instruct 8B GPT-4o
Builds on: Mamba Mamba2 speculative decoding DPO Laughing Hyena
Closest papers by meaning
- Transformers to SSMs: Distilling Quadratic Knowledge to Subquadratic Models (NeurIPS 2024 · similarity 0.89)
- Mamba Modulation: On the Length Generalization of Mamba Models (NeurIPS 2025 · similarity 0.89)
- Routing Mamba: Scaling State Space Models with Mixture-of-Experts Projection (NeurIPS 2025 · similarity 0.89)
- LongMamba: Enhancing Mamba's Long-Context Capabilities via Training-Free Receptive Field Enlargement (ICLR 2025 · similarity 0.89)
- DeciMamba: Exploring the Length Extrapolation Potential of Mamba (ICLR 2025 · similarity 0.88)
- Cobra: Extending Mamba to Multi-Modal Large Language Model for Efficient Inference (AAAI 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.