The Mamba in the Llama: Distilling and Accelerating Hybrid Models

• URL: http://arxiv.org/abs/2408.15237v3
• Date: Wed, 08 Jan 2025 20:34:02 GMT
• Title: The Mamba in the Llama: Distilling and Accelerating Hybrid Models
• Authors: Junxiong Wang, Daniele Paliotta, Avner May, Alexander M. Rush, Tri Dao,
• Abstract summary: We show how to distill large Transformers into linear RNNs by reusing the linear projection weights from attention layers with academic GPU resources.<n>The resulting hybrid model achieves performance comparable to the original Transformer in chat benchmarks.<n>We also introduce a hardware-aware speculative decoding algorithm that accelerates the inference speed of Mamba and hybrid models.
• Score: 76.64055251296548
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• M1: Towards Scalable Test-Time Compute with Mamba Reasoning Models [72.75501495786297]
  We introduce a novel hybrid linear RNN reasoning model, M1, built on the Mamba architecture. Experimental results show that M1 not only outperforms previous linear RNN models but also matches the performance of state-of-the-art DeepSeek R1 distilled reasoning models.
  arXiv  Detail & Related papers  (2025-04-14T17:38:25Z)
• Thinking Slow, Fast: Scaling Inference Compute with Distilled Reasoners [72.37408197157453]
  Recent advancements have demonstrated that the performance of large language models (LLMs) can be significantly enhanced by scaling computational resources at test time. This raises a fundamental question: can models with lower complexity leverage their superior generation throughput to outperform similarly sized Transformers for a fixed computational budget? To address this question and overcome the lack of strong subquadratic reasoners, we distill pure and hybrid Mamba models from pretrained Transformers.
  arXiv  Detail & Related papers  (2025-02-27T18:08:16Z)
• MaskMamba: A Hybrid Mamba-Transformer Model for Masked Image Generation [63.73137438677585]
  MaskMamba is a novel hybrid model that combines Mamba and Transformer architectures. It achieves a remarkable $54.44%$ improvement in inference speed at a resolution of $2048times 2048$ over Transformer.
  arXiv  Detail & Related papers  (2024-09-30T04:28:55Z)
• Transformers to SSMs: Distilling Quadratic Knowledge to Subquadratic Models [92.36510016591782]
  We present a method that is able to distill a pretrained Transformer architecture into alternative architectures such as state space models (SSMs) Our method, called MOHAWK, is able to distill a Mamba-2 variant based on the Phi-1.5 architecture using only 3B tokens and a hybrid version (Hybrid Phi-Mamba) using 5B tokens. Despite using less than 1% of the training data typically used to train models from scratch, Phi-Mamba boasts substantially stronger performance compared to all past open-source non-Transformer models.
  arXiv  Detail & Related papers  (2024-08-19T17:48:11Z)
• How Effective are State Space Models for Machine Translation? [19.509486069758495]
  Transformers are the current architecture of choice for NLP, but their attention layers do not scale well to long contexts. Recent works propose to replace attention with linear recurrent layers. It remains unclear whether these models are competitive with transformers in machine translation.
  arXiv  Detail & Related papers  (2024-07-07T20:21:49Z)
• An Empirical Study of Mamba-based Language Models [69.74383762508805]
  Selective state-space models (SSMs) like Mamba overcome some shortcomings of Transformers. We present a direct comparison between 8B-context Mamba, Mamba-2, and Transformer models trained on the same datasets. We find that the 8B Mamba-2-Hybrid exceeds the 8B Transformer on all 12 standard tasks.
  arXiv  Detail & Related papers  (2024-06-12T05:25:15Z)
• Parallelizing Linear Transformers with the Delta Rule over Sequence Length [49.88826673324244]
  This work describes a hardware-efficient algorithm for training linear transformers with the delta rule. We train a 1.3B model for 100B tokens and find that it outperforms recent linear-time baselines.
  arXiv  Detail & Related papers  (2024-06-10T17:24:42Z)
• Demystify Mamba in Vision: A Linear Attention Perspective [72.93213667713493]
  Mamba is an effective state space model with linear computation complexity. We show that Mamba shares surprising similarities with linear attention Transformer. We propose a Mamba-Like Linear Attention (MLLA) model by incorporating the merits of these two key designs into linear attention.
  arXiv  Detail & Related papers  (2024-05-26T15:31:09Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.