PaperlayerView exact revision on arXiv

Source-linked AI summary

The Mamba in the Llama: Distilling and Accelerating Hybrid Models

Junxiong Wang, Daniele Paliotta, Avner May, Alexander M. Rush, Tri Dao

arXiv:2408.15237v4cs.LGcs.AI

TL;DR

Long-sequence Transformer generation is costly, motivating conversion of pretrained Transformers into faster-deploying linear RNNs. The paper transfers attention projections into Mamba-based hybrid models and develops hardware-aware speculative decoding; its distilled models match teacher chat benchmarks and compare favorably with similarly sized scratch-trained Mamba models.

  • Problem

    Transformers’ quadratic sequence complexity and large KV cache slow very-long-sequence generation, while linear RNNs offer faster inference but still require substantial training resources to scale.

    Source: §1 · p. 1

  • Method

    The approach initializes Mamba linear RNN blocks from Transformer attention projections, distills them through multiple stages, and develops hardware-aware speculative decoding for Mamba and hybrid models.

    Source: §1 · p. 2

  • Results

    Distilled models trained with 20B tokens match teacher performance on standard chat benchmarks and perform on par with or better than similarly sized Mamba models trained from scratch with 1.2T or 3.5T tokens.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The results show that Transformer knowledge can be transferred to other architectures while the resulting Mamba model nears LLM chat accuracy and can be accelerated with speculative decoding.

    Source: §8 · p. 14

    Mamba’s recurrent-state rewind during speculation would require caching all prior hidden states, creating large memory overhead.

    Source: §4.1 · p. 6

Abstract

from arXiv · show

Linear RNN architectures, like Mamba, can be competitive with Transformer models in language modeling while having advantageous deployment characteristics. Given the focus on training large-scale Transformer models, we consider the challenge of converting these pretrained models for deployment. We demonstrate that it is feasible to distill large Transformers into linear RNNs by reusing the linear projection weights from attention layers with academic GPU resources. The resulting hybrid model, which incorporates a quarter of the attention layers, achieves performance comparable to the original Transformer in chat benchmarks and outperforms open-source hybrid Mamba models trained from scratch with trillions of tokens in both chat benchmarks and general benchmarks. Moreover, we introduce a hardware-aware speculative decoding algorithm that accelerates the inference speed of Mamba and hybrid models. Overall we show how, with limited computation resources, we can remove many of the original attention layers and generate from the resulting model more efficiently. Our top-performing model, distilled from Llama3-8B-Instruct, achieves a 29.61 length-controlled win rate on AlpacaEval 2 against GPT-4 and 7.35 on MT-Bench, surpassing the best 8B scale instruction-tuned linear RNN model. We also find that the distilled model has natural length extrapolation, showing almost perfect accuracy in the needle-in-a-haystack test at 20x the distillation length. Code and pre-trained checkpoints are open-sourced at https://github.com/jxiw/MambaInLlama and https://github.com/itsdaniele/speculative_mamba.

1 Introduction

The paper targets efficient generation by distilling pretrained Transformers into linear RNNs, addressing the challenge of transferring Transformer weights and adapting speculative decoding. It reports teacher-level chat performance and competitive general-benchmark results after training with 20B tokens.

Motivation

  • Linear RNNs offer 5× higher inference throughput than Transformers, whose quadratic sequence complexity and large KV cache make very long-sequence generation slow.The introduction identifies long-document reasoning and large-codebase workflows as examples where KV-cache demands matter.Source: §1 · p. 1

Weight transfer

  • Reusing attention-layer weights enables distilling a large Transformer into a hybrid linear RNN with minimal additional compute while preserving much of its generation quality.The proposed Mamba architecture initializes directly from a pretrained model’s attention block.Source: §1 · p. 2

Distillation

  • A multistage distillation pipeline combines progressive distillation, supervised fine-tuning, and directed preference optimization, outperforming vanilla distillation in perplexity and downstream evaluation.The approach mirrors a standard LLM training pipeline.Source: §1 · p. 2

Speculative decoding

  • A hardware-aware speculative sampling algorithm and fast kernel support speculative decoding on Mamba and hybrid models, reaching over 300 tokens/second for Mamba 7B.The paper also reports that speculative decoding applies effectively to its hybrid architecture.Source: §1 · p. 2

Experimental results

  • With 20B training tokens, distilled models match their teachers on standard chat benchmarks and perform on par with or better than similarly sized Mamba models trained from scratch with 1.2T or 3.5T tokens.The general-benchmark comparisons include MMLU and TruthfulQA.Source: §1 · p. 2

2 From Transformer to Mamba

The section develops a path from attention to linear RNNs, then expands that formulation using Mamba’s state-space parameterization. It applies the resulting model by transferring attention projections and replacing attention layers stepwise while retaining frozen Transformer MLP layers.

2.1 Relationship Between Attention and Linear RNNs

  • Removing softmax from attention yields a linear RNN form, but the naive version tracks only one scalar per hidden dimension and performs poorly because softmax is critical to attention.Prior kernel methods expand the hidden-state representation to better match softmax’s modeling capacity.Source: §2.1 · pp. 2–3

2.2 Distilling to an Expanded Linear RNN

  • Mamba expands the linear RNN hidden state through a continuous-time state-space parameterization discretized with learned sampling intervals, enabling richer functions with relatively small efficiency costs.The expanded hidden state scales by N′ over naive linear attention; Mamba’s fused kernel avoids materializing the discrete parameters.Source: §2.2 · pp. 3–4

2.3 Attention-to-Mamba Initialization and Hybrid Stepwise Training

  • The attention-initialized Mamba block feeds Transformer Q, K, and V projections into discretization, then applies the resulting linear RNN while keeping Transformer MLP layers frozen.The design initializes Mamba’s C, B, and X projections from Q, K, and V, respectively.Source: §2.3 · p. 4, Figure 1 · p. 4
  • Replacing attention layers stepwise—first keeping every two layers, then every four—was empirically the most effective strategy reported for hybrid models.The approach distills after each replacement stage.Source: §2.3 · p. 4

3 Knowledge Distillation for Aligned LMs

The approach distills alignment stages of a pretrained language model while keeping its transferred MLP layers frozen and training its Mamba layers. It applies knowledge distillation to supervised fine-tuning and preference optimization, including sequence-level DPO.

3 Knowledge Distillation for Aligned LMs

  • The method assumes most Transformer knowledge remains in transferred MLP layers, which are frozen while Mamba layers are trained for fine-tuning and alignment.This focuses distillation on the fine-tuning and alignment steps.Source: §3 · p. 5
  • Supervised fine-tuning distillation trains the student to maximize the likelihood of a response given its prompt, using conditional-generation-style training.Two approaches are described: word-level KL divergence over next-token distributions and sequence-level distillation using teacher-generated pseudo-labels in place of ground truth.Source: §3 · p. 5
  • Preference optimization uses desired and undesired response pairs to improve outputs against a reward while maintaining closeness to a reference model.The reference is typically the model after supervised fine-tuning; DPO can optimize the objective with direct gradient updates from preferred and dispreferred outputs.Source: §3 · p. 5
  • The paper applies DPO as a sequence-level distillation objective by scoring preferred and dispreferred outputs with teacher and student models, then backpropagating to the student.The authors describe this as the first use of DPO as a distillation objective.Source: §3 · p. 5

4 Speculative Decoding Algorithms For Linear RNNs

The section adapts speculative decoding to linear RNNs by using hardware-aware multi-step generation and lazy hidden-state caching. It addresses the challenges of verifying draft tokens and reports that kernel fusion enables H100 speedups where a naive implementation did not.

4.1 Challenges in RNN Speculation

  • RNN speculation is difficult because short verification chains have uncertain conflict points, while parallel kernels target very long sequences and avoid materializing intermediate states.Rewinding an RNN may require caching all hidden states, which creates large memory overhead because the RNN state is very large.Source: §4.1 · p. 6

4.2 Multi-Step Linear RNN Speculation

  • The proposed multi-step kernel verifies draft tokens without materializing most hidden states or discrete-time parameters, and can recompute the correct state after rejection from a saved snapshot.The approach assumes decoding is memory-bound, so computing multiple steps adds little overhead over single-step decoding.Source: §4.2 · pp. 6–7
  • Verification keeps one RNN hidden state cached and advances it lazily; hybrid models use Algorithm 2 for RNN layers and parallel verification for Transformer layers.A Mamba or hybrid draft model also recomputes states for tokens accepted in the previous iteration.Source: §4.2 · p. 7

4.3 Speculation Analysis and Hardware Specific Optimization

  • Kernel fusion, recomputation, and caching in fused kernels achieve H100 speedups, whereas the naive multi-call implementation has no H100 speedup.The Table 1 comparison concerns speculative decoding with pure Mamba models.Source: §4.3 · pp. 7–8, Table 1 · p. 7

5 Results

The distilled hybrids retain strong chat and general-benchmark performance, with long-context retrieval beyond the distillation length and speculative decoding gains on hybrid models. The results cover multiple teacher models, attention-layer configurations, and evaluation settings.

5.1 Experimental Setup

  • Experiments use Zephyr-7B and Llama-3 Instruct 8B teachers, with hybrid Mamba and Mamba2 models spanning 50%, 25%, 12.5%, and 0% attention.The setup also describes later Llama-3.1 70B distillations at 3B and 8B scales, with additional DPO variants.Source: §5.1 · pp. 8–9
  • Distillation uses pseudo-labeling, supervised fine-tuning, and preference alignment; each initial hybrid-model distillation takes under five days on 8x80G A100 GPUs.The first stage uses UltraChat and UltraFeedback prompts; later stages use GenQA, InfinityInstruct, OpenHermes 2.5, and DPO datasets that vary by teacher.Source: §5.1 · p. 8

5.2 Evaluation on Chat Benchmarks

  • The 50%-attention hybrid matches its teacher on MT-Bench and slightly exceeds it on AlpacaEval, while 25% and 12.5% variants remain competitive.The pure Mamba model degrades significantly; the distilled hybrid also outperforms Falcon Mamba, trained from scratch with more than 5T tokens.Source: §5.2 · p. 9, Table 2 · p. 9

5.3 Evaluation on General Benchmarks

  • Llama-3-distilled hybrid Mamba and Mamba2 outperform scratch-trained TRI Mamba and Nvidia Mamba on zero-shot LM Eval; at 25% attention, they remain competitive.On the Open LLM Leaderboard, distilled hybrids match the best open-source linear RNN models and outperform corresponding instruct models on GSM8K and CRUX.Source: §5.3 · p. 10, Table 3 · p. 10, Table 4 · p. 10

5.4 Evaluation on Long Context Tasks

  • After distillation on 2k-length examples, the 3B hybrids achieve perfect retrieval accuracy up to 10k, and the 8B hybrids up to 16k.Mamba-Llama3.1-8B shows good results up to 38k; the 3B result is better than Llama-3.2-3B-Instruct.Source: §5.4 · p. 11, Figure 4 · p. 11

5.5 Hybrid speculative decoding

  • Over 1.8x speedups on Zephyr hybrids are reported for both 50% and 25% models; Llama-hybrid speedups are more modest.Four-layer drafts achieve higher acceptance but add overhead; the Llama draft is larger because of its embedding table.Source: §5.5 · pp. 11–12, Table 5 · p. 12

6 Analysis

The analysis tests distillation quality, alignment choices, initialization, and the role of Mamba blocks. It reports advantages for attention-weight initialization, combined SFT and DPO, and retaining Mamba layers.

Comparison with other distillation approaches

  • Perplexity degradation is 1.03 with 50% attention, 1.09 with 25%, 1.22 with 6.35%, and 3.36 with no attention; removing more layers worsens results.The comparison with Distill Hyena is qualified because its dataset and model size differ.Source: §6 · p. 12, Table 6 · p. 12

Does distilling from preferences help?

  • SFT plus DPO yields the best alignment score, while either SFT or DPO alone yields little improvement.The comparison uses Zephyr as teacher, OpenHermes 2.5 for SFT, and UltraFeedback for DPO.Source: §6 · p. 12, Table 6 · p. 12

Pseudo Label Distillation Ablations

  • Transformer weight initialization is crucial: without it, perplexity significantly worsens for both pure Mamba and hybrid models.Freezing MLP layers can help the student focus on token interactions and better mimic attention; progressive distillation and attention/Mamba interleaving offer smaller benefits.Source: §6 · pp. 12–13, Table 7 · p. 13

Attention Initialization

  • Reusing attention-layer linear projections performs significantly better than random initialization across all evaluated benchmarks.The comparison is for Zephyr-Mamba with 50% attention.Source: §6 · p. 13, Table 8 · p. 13

Necessity of Linear RNN

  • Removing Mamba blocks significantly worsens performance relative to the model with Mamba, indicating gains are not solely attributable to remaining attention layers.Source: §6 · p. 13, Table 9 · p. 13

7 Related Work

Related work covers attention-free sequence models, attempts to distill Transformers into linear RNN-style architectures, and speculative decoding methods. The cited approaches differ in their architectures and decoding strategies.

Attention-free models

  • Attention-free work includes S4 and variants, gated SSMs, Mamba, other sub-quadratic architectures, and hybrid architectures.The cited discussion notes Mamba's claim that static dynamics do not capture input-specific context selection and reports Mamba outperforming Transformers across scales.Source: §7 · pp. 13–14

Distillation from Transformers

  • Prior Transformer-to-linear-RNN distillation includes transferring Hyena's long convolution into a state-space representation and progressively distilling small Transformers into Hyena models.Source: §7 · p. 14

Speculative Decoding

  • Speculative decoding uses a draft model to propose tokens for verification; prior work includes rejection sampling, tree-structured candidates, and trained or training-free drafts.Source: §7 · p. 14

8 Conclusion

The paper combines distillation and speculative decoding to maintain LLM abilities while increasing decoding speed. Its results suggest Transformer knowledge can transfer to other architectures, enabling inference profiles beyond attention optimization.

8 Conclusion

  • The full model nears LLM chat accuracy and is accelerated at inference with hardware-aware speculative decoding.The conclusion describes speculative decoding as significantly speeding up the Mamba model.Source: §8 · p. 14
  • Distillation on supervised instructions and preferences improves the model’s ability with relatively little compute.Source: §8 · p. 14
  • The results suggest Transformer knowledge can transfer to other architectures, opening potential to customize LLM inference profiles beyond optimizing attention.Source: §8 · p. 14
Loading 2408.15237v4…