The Mamba in the Llama: Distilling and Accelerating Hybrid Models
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 instruction-tuned linear RNN model.
66ce962f9ca851092d935554
Junxiong Wang
admin_assigned
2024-08-28T08:14:57.265Z
63dc68bea99b2c8a7c20f1d0
/avatars/cf8ddc91415ef1f895803f4390ff1f6f.svg
Junxiong Wang
user
JunxiongWang
66ce962f9ca851092d935555
Daniele Paliotta
admin_assigned
2024-08-28T08:15:17.427Z
66799a0655941cab0f13f5fc
/avatars/998850563291242b720b2e48ce54a29a.svg
Daniele Paliotta
user
danielepaliotta
66ce962f9ca851092d935556
Avner May
admin_assigned
2024-08-28T08:15:23.275Z
6514823916db7b869acecde9
https://cdn-avatars.huggingface.co/v1/production/uploads/6514823916db7b869acecde9/XAtcK7P0PgSZLB8qpb3rj.jpeg
Avner May
user
avnermay
66ce962f9ca851092d935557
Alexander M. Rush
claimed_verified
2024-09-20T07:48:32.484Z
5e694f697127cb58dde42fe1
https://cdn-avatars.huggingface.co/v1/production/uploads/5dd96eb166059660ed1ee413/J21QWWtFOKqhJTF5qIMlD.jpeg
Sasha Rush
user
srush
66ce962f9ca851092d935558
Tri Dao
admin_assigned
2024-08-28T08:15:35.978Z
64b8a6b5cf14c2fabe98159b
/avatars/dbc009451865435bf290791beadc4723.svg
Tri Dao
user
tridao
The Mamba in the Llama: Distilling and Accelerating Hybrid Models
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 instruction-tuned linear RNN model.
Large-scale Transformer models can be distilled into linear RNNs using attention layer weights, achieving competitive performance with enhanced deployment efficiency.
Linear RNN Mamba Transformer language modeling distillation attention layers hybrid model chat benchmarks speculative decoding inference speed Llama3-8B-Instruct AlpacaEval MT-Bench