Elamur Research⌕
☰
home / papers / 2024 / The Mamba in the Llama: Distilling and Accelerating Hybrid Models

The Mamba in the Llama: Distilling and Accelerating Hybrid Models

Paper · 2024 · Junxiong Wang; Daniele Paliotta; Avner May; Alexander M. Rush; and Tri Dao

Summary

Transformer LLMs decode slowly on long sequences because of quadratic attention cost and a growing KV cache. Linear RNNs such as Mamba run inference far faster, but pretraining them at scale is expensive. The authors instead convert existing pretrained chat Transformers (Zephyr-7B, Llama-3 8B) into hybrid Mamba/Mamba2 models.

Each attention head becomes a Mamba head whose input, B and C projections reuse the attention value, key and query weights. Only the A matrix and the step size ∆ are new parameters, and the MLP blocks stay frozen. Attention layers are removed in stages, first down to 50% and then to 25%. The models are then trained with distillation-based SFT and DPO on about 20B tokens, which fits academic-scale GPU budgets. The authors also design a hardware-aware speculative decoding algorithm and kernel for Mamba and hybrid models.

The distilled hybrids roughly match their teachers on chat benchmarks. They beat Mamba models trained from scratch on trillions of tokens. This suggests that knowledge in Transformers can be transferred cheaply into architectures that are faster at inference.

Key points

  • The best model, distilled from Llama3-8B-Instruct with 25% attention layers, reaches a 29.61 length-controlled win rate on AlpacaEval 2 (vs. GPT-4) and scores 7.35 on MT-Bench.
  • It matches or exceeds Mamba 7B (trained on 1.2T tokens) and NVIDIA Hybrid Mamba2 (trained on 3.5T tokens) on tasks such as MMLU and TruthfulQA.
  • Initializing from attention weights is critical. In the Zephyr 50%-attention hybrid, random initialization drops MT-Bench from 6.69 to 1.04 and raises LAMBADA perplexity from 6.20 to 55.01. Removing the Mamba blocks entirely is worse still, with perplexity of 151.98.
  • The speculative decoding method exceeds 300 tokens/s for a Mamba 7B model, and the distilled model handles needle-in-a-haystack retrieval almost perfectly at 20× the distillation length.
  • Code and checkpoints are public in the GitHub repos jxiw/MambaInLlama and itsdaniele/speculative_mamba.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Fields, topics and concepts

📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 19,503 worksArtificial Intelligence · 18,805 worksComputer Vision and Pattern Recognition · 3,621 worksHardware and Architecture · 183 works

Cite

Cite as Junxiong Wang; Daniele Paliotta; Avner May; Alexander M. Rush; and Tri Dao. 2024. The Mamba in the Llama: Distilling and Accelerating Hybrid Models. arXiv:2408.15237.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2408.15237 · arXiv 2408.15237

@misc{rush2024mamba,
  title = {{The Mamba in the Llama: Distilling and Accelerating Hybrid Models}},
  author = {Junxiong Wang; Daniele Paliotta; Avner May; Alexander M. Rush; and Tri Dao},
  year = {2024},
  doi = {10.48550/arxiv.2408.15237},
  eprint = {2408.15237},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2408.15237}
}

References

85 references, 41 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Arora, S., Eyuboglu, S., Timalsina, A., Johnson, I., Poli, M., Zou, J., Rudra, A., and Ré, C. (2023). Zoology: Measuring and improving recall in efficient language models. arXiv preprint arXiv:2312.04927.
  2. Arora, S., Eyuboglu, S., Zhang, M., Timalsina, A., Alberti, S., Zinsley, D., Zou, J., Rudra, A., and Ré, C. (2024). Simple linear attention language models balance the recall-throughput tradeoff. arXiv preprint arXiv:2402.18668.
  3. BAAI (2024). Infinity instruct. arXiv preprint arXiv:2406.XXXX. find
  4. Beck, M., Pöppel, K., Spanring, M., Auer, A., Prudnikova, O., Kopp, M., Klambauer, G., Brandstetter, J., and Hochreiter, S. (2024). xlstm: Extended long short-term memory. arXiv preprint arXiv:2405.04517.
  5. Bhendawade, N., Belousova, I., Fu, Q., Mason, H., Rastegari, M., and Najibi, M. (2024). Speculative streaming: Fast llm inference without auxiliary models. arXiv preprint arXiv:2402.11131. — [arXiv:2402.11131](https://arxiv.org/abs/2402.11131) find
  6. Bick, A., Li, K. Y., Xing, E. P., Kolter, J. Z., and Gu, A. (2024). Transformers to ssms: Distilling quadratic knowledge to subquadratic models. arXiv preprint arXiv:2408.10189.
  7. Bisk, Y., Zellers, R., Gao, J., Choi, Y., et al. (2020). PIQA: Reasoning about Physical Commonsense in Natural Language. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 7432–7439.
  8. Botev, A., De, S., Smith, S. L., Fernando, A., Muraru, G.-C., Haroun, R., Berrada, L., Pascanu, R., Sessa, P. G., Dadashi, R., et al. (2024). Recurrentgemma: Moving past transformers for efficient open language models. arXiv preprint arXiv:2404.07839. find
  9. Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. (2020). Language models are few-shot learners. Advances in neural information processing systems, 33:1877–1901.
  10. Cai, T., Li, Y., Geng, Z., Peng, H., Lee, J. D., Chen, D., and Dao, T. (2024). Medusa: Simple llm inference acceleration framework with multiple decoding heads. arXiv preprint arXiv:2401.10774.
  11. Chen, C., Borgeaud, S., Irving, G., Lespiau, J.-B., Sifre, L., and Jumper, J. (2023a). Accelerating Large Language Model Decoding with Speculative Sampling. find
  12. Chen, J., Qadri, R., Wen, Y., Jain, N., Kirchenbauer, J., Zhou, T., and Goldstein, T. (2024). Genqa: Generating millions of instructions from a handful of prompts. arXiv preprint arXiv:2406.10323. — [arXiv:2406.10323](https://arxiv.org/abs/2406.10323) find
  13. Chen, Z., Yang, X., Lin, J., Sun, C., Huang, J., and Chang, K. C.-C. (2023b). Cascade speculative drafting for even faster llm inference. arXiv preprint arXiv:2312.11462. — [arXiv:2312.11462](https://arxiv.org/abs/2312.11462) find
  14. Clark, P., Cowhey, I., Etzioni, O., Khot, T., Sabharwal, A., Schoenick, C., and Tafjord, O. (2018a). Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge. arXiv preprint arXiv:1803.05457.
  15. Clark, P., Cowhey, I., Etzioni, O., Khot, T., Sabharwal, A., Schoenick, C., and Tafjord, O. (2018b). Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv:1803.05457.
  16. Cobbe, K., Kosaraju, V., Bavarian, M., Chen, M., Jun, H., Kaiser, L., Plappert, M., Tworek, J., Hilton, J., Nakano, R., et al. (2021). Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168.
  17. Cui, G., Yuan, L., Ding, N., Yao, G., Zhu, W., Ni, Y., Xie, G., Liu, Z., and Sun, M. (2023). Ultrafeedback: Boosting language models with high-quality feedback. arXiv preprint arXiv:2310.01377. find
  18. Dao, T. and Gu, A. (2024). Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. arXiv preprint arXiv:2405.21060. find
  19. De, S., Smith, S. L., Fernando, A., Botev, A., Cristian-Muraru, G., Gu, A., Haroun, R., Berrada, L., Chen, Y., Srinivasan, S., et al. (2024). Griffin: Mixing gated linear recurrences with local attention for efficient language models. arXiv preprint arXiv:2402.19427. find
  20. Ding, N., Chen, Y., Xu, B., Qin, Y., Hu, S., Liu, Z., Sun, M., and Zhou, B. (2023). Enhancing chat language models by scaling high-quality instructional conversations. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pages 3029–3051.
  21. Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., et al. (2024). The llama 3 herd of models. arXiv preprint arXiv:2407.21783.
  22. Fu, D., Arora, S., Grogan, J., Johnson, I., Eyuboglu, E. S., Thomas, A., Spector, B., Poli, M., Rudra, A., and Ré, C. (2024a). Monarch mixer: A simple sub-quadratic gemm-based architecture. Advances in Neural Information Processing Systems, 36. find
  23. Fu, D. Y., Dao, T., Saab, K. K., Thomas, A. W., Rudra, A., and Re, C. (2022). Hungry hungry hippos: Towards language modeling with state space models. In The Eleventh International Conference on Learning Representations. find
  24. Fu, Y., Bailis, P., Stoica, I., and Zhang, H. (2024b). Break the sequential dependency of llm inference using lookahead decoding. arXiv preprint arXiv:2402.02057.
  25. Gao, L., Tow, J., Abbasi, B., Biderman, S., Black, S., DiPofi, A., Foster, C., Golding, L., Hsu, J., Le Noac’h, A., Li, H., McDonell, K., Muennighoff, N., Ociepa, C., Phang, J., Reynolds, L., Schoelkopf, H., Skowron, A., Sutawika, L., Tang, E., Thite, A., Wang, B., Wang, K., and Zou, A. (2023). A fr… find
  26. Gu, A. and Dao, T. (2023). Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752. find
  27. Gu, A., Goel, K., Gupta, A., and Ré, C. (2022). On the Parameterization and Initialization of Diagonal State Space Models. Advances in Neural Information Processing Systems, 35:35971–35983. find
  28. Gu, A., Goel, K., and Ré, C. (2021). Efficiently Modeling Long Sequences with Structured State Spaces. arXiv preprint arXiv:2111.00396. find
  29. Gu, A., Rozière, B., Leather, H., Solar-Lezama, A., Synnaeve, G., and Wang, S. I. (2024). Cruxeval: A benchmark for code reasoning, understanding and execution. arXiv preprint arXiv:2401.03065.
  30. Guo, M., Ainslie, J., Uthus, D., Ontanon, S., Ni, J., Sung, Y.-H., and Yang, Y. (2021). Longt5: Efficient text-to-text transformer for long sequences. arXiv preprint arXiv:2112.07916.
  31. Gupta, A., Gu, A., and Berant, J. (2022). Diagonal State Spaces are as Effective as Structured State Spaces. Advances in Neural Information Processing Systems, 35:22982–22994. find
  32. He, Z., Zhong, Z., Cai, T., Lee, J. D., and He, D. (2023). Rest: Retrieval-based speculative decoding. arXiv preprint arXiv:2311.08252.
  33. Hendrycks, D., Burns, C., Basart, S., Zou, A., Mazeika, M., Song, D., and Steinhardt, J. (2020). Measuring Massive Multitask Language Understanding. In International Conference on Learning Representations.
  34. Hendrycks, D., Burns, C., Basart, S., Zou, A., Mazeika, M., Song, D., and Steinhardt, J. (2021). Measuring massive multitask language understanding. Proceedings of the International Conference on Learning Representations (ICLR).
  35. Hinton, G., Vinyals, O., and Dean, J. (2015). Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531.
  36. Irie, K., Schlag, I., Csordás, R., and Schmidhuber, J. (2021). Going beyond linear transformers with recurrent fast weight programmers. Advances in neural information processing systems, 34:7703–7717. find
  37. Jiang, A. Q., Sablayrolles, A., Mensch, A., Bamford, C., Chaplot, D. S., Casas, D. d. l., Bressand, F., Lengyel, G., Lample, G., Saulnier, L., et al. (2023). Mistral 7b. arXiv preprint arXiv:2310.06825. find
  38. Jin, Q., Dhingra, B., Liu, Z., Cohen, W. W., and Lu, X. (2019). PubMedQA: A Dataset for Biomedical Research Question Answering. arXiv preprint arXiv:1909.06146.
  39. Kim, Y. and Rush, A. M. (2016). Sequence-level knowledge distillation. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, pages 1317–1327.
  40. Lai, G., Xie, Q., Liu, H., Yang, Y., and Hovy, E. (2017). RACE: Large-scale ReAding Comprehension Dataset From Examinations. arXiv preprint arXiv:1704.04683.
  41. Leviathan, Y., Kalman, M., and Matias, Y. (2023). Fast Inference from Transformers via Speculative Decoding. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 19274–19286. PMLR. find
  42. Li, R., Allal, L. B., Zi, Y., Muennighoff, N., Kocetkov, D., Mou, C., Marone, M., Akiki, C., Li, J., Chim, J., et al. (2023a). Starcoder: may the source be with you! arXiv preprint arXiv:2305.06161. find
  43. Li, X., Zhang, T., Dubois, Y., Taori, R., Gulrajani, I., Guestrin, C., Liang, P., and Hashimoto, T. B. (2023b). Alpacaeval: An automatic evaluator of instruction-following models. https://github.com/tatsu-lab/ alpaca_eval. ↗
  44. Lieber, O., Lenz, B., Bata, H., Cohen, G., Osin, J., Dalmedigos, I., Safahi, E., Meirom, S., Belinkov, Y., Shalev-Shwartz, S., et al. (2024). Jamba: A hybrid transformer-mamba language model. arXiv preprint arXiv:2403.19887.
  45. Lin, B. Y. (2024). ZeroEval: A Unified Framework for Evaluating Language Models. find
  46. Lin, S., Hilton, J., and Evans, O. (2021). TruthfulQA: Measuring How Models Mimic Human Falsehoods. arXiv preprint arXiv:2109.07958.
  47. Lin, S., Hilton, J., and Evans, O. (2022). Truthfulqa: Measuring how models mimic human falsehoods. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 3214–3252.
  48. Liu, X., Hu, L., Bailis, P., Stoica, I., Deng, Z., Cheung, A., and Zhang, H. (2023). Online speculative decoding. arXiv preprint arXiv:2310.07177.
  49. Massaroli, S., Poli, M., Fu, D., Kumbong, H., Parnichkun, R., Romero, D., Timalsina, A., McIntyre, Q., Chen, B., Rudra, A., et al. (2024). Laughing hyena distillery: Extracting compact recurrences from convolutions. Advances in Neural Information Processing Systems, 36. find
  50. Mehta, H., Gupta, A., Cutkosky, A., and Neyshabur, B. (2023). Long Range Language Modeling via Gated State Spaces. In The Eleventh International Conference on Learning Representations. find
  51. Meng, Y., Xia, M., and Chen, D. (2024). Simpo: Simple preference optimization with a reference-free reward. arXiv preprint arXiv:2405.14734. find
  52. Mercat, J., Vasiljevic, I., Keh, S., Arora, K., Dave, A., Gaidon, A., and Kollar, T. (2024). Linearizing large language models. arXiv preprint arXiv:2405.06640.
  53. Merity, S., Xiong, C., Bradbury, J., and Socher, R. (2016). Pointer sentinel mixture models. arXiv preprint arXiv:1609.07843.
  54. Mihaylov, T., Clark, P., Khot, T., and Sabharwal, A. (2018). Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answering. arXiv preprint arXiv:1809.02789. find
  55. Peng, B., Alcaide, E., Anthony, Q., Albalak, A., Arcadinho, S., Cao, H., Cheng, X., Chung, M., Grella, M., GV, K. K., et al. (2023a). Rwkv: Reinventing rnns for the transformer era. arXiv preprint arXiv:2305.13048. find
  56. Peng, B., Quesnelle, J., Fan, H., and Shippole, E. (2023b). Yarn: Efficient context window extension of large language models. arXiv preprint arXiv:2309.00071. find
  57. Poli, M., Massaroli, S., Nguyen, E., Fu, D. Y., Dao, T., Baccus, S., Bengio, Y., Ermon, S., and Ré, C. (2023). Hyena hierarchy: Towards larger convolutional language models. In International Conference on Machine Learning, pages 28043–28078. PMLR. find
  58. Rafailov, R., Sharma, A., Mitchell, E., Manning, C. D., Ermon, S., and Finn, C. (2024). Direct preference optimization: Your language model is secretly a reward model. Advances in Neural Information Processing Systems, 36.
  59. Ralambomihanta, T. R., Mohammadzadeh, S., Islam, M. S. N., Jabbour, W., and Liang, L. (2024). Scavenging hyena: Distilling transformers into long convolution models. arXiv preprint arXiv:2401.17574. find
  60. Ren, L., Liu, Y., Lu, Y., Shen, Y., Liang, C., and Chen, W. (2024). Samba: Simple hybrid state space models for efficient unlimited context language modeling. arXiv preprint arXiv:2406.07522.
  61. Roziere, B., Gehring, J., Gloeckle, F., Sootla, S., Gat, I., Tan, X. E., Adi, Y., Liu, J., Remez, T., Rapin, J., et al. (2023). Code llama: Open foundation models for code. arXiv preprint arXiv:2308.12950.
  62. Sakaguchi, K., Bras, R. L., Bhagavatula, C., and Choi, Y. (2021). Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64(9):99–106. find
  63. Schlag, I., Irie, K., and Schmidhuber, J. (2021). Linear transformers are secretly fast weight programmers. In International Conference on Machine Learning, pages 9355–9366. PMLR.
  64. Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. (2017). Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347. find
  65. Shaham, U., Segal, E., Ivgi, M., Efrat, A., Yoran, O., Haviv, A., Gupta, A., Xiong, W., Geva, M., Berant, J., et al. (2022). Scrolls: Standardized comparison over long language sequences. arXiv preprint arXiv:2201.03533.
  66. Shleifer, S. and Rush, A. M. (2020). Pre-trained summarization distillation. CoRR, abs/2010.13002. — [arXiv:2010.13002](https://arxiv.org/abs/2010.13002) find
  67. Spector, B. and Re, C. (2023). Accelerating llm inference with staged speculative decoding. arXiv preprint arXiv:2308.04623. — [arXiv:2308.04623](https://arxiv.org/abs/2308.04623) find
  68. Sun, Y., Dong, L., Huang, S., Ma, S., Xia, Y., Xue, J., Wang, J., and Wei, F. (2023). Retentive network: A successor to transformer for large language models. arXiv preprint arXiv:2307.08621.
  69. Tay, Y., Dehghani, M., Abnar, S., Shen, Y., Bahri, D., Pham, P., Rao, J., Yang, L., Ruder, S., and Metzler, D. (2020). Long range arena: A benchmark for efficient transformers. In International Conference on Learning Representations.
  70. Teknium (2023). Openhermes 2.5: An open dataset of synthetic data for generalist llm assistants. find
  71. Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., et al. (2023). Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971.
  72. Tunstall, L., Beeching, E., Lambert, N., Rajani, N., Rasul, K., Belkada, Y., Huang, S., von Werra, L., Fourrier, C., Habib, N., et al. (2023). Zephyr: Direct distillation of lm alignment. arXiv preprint arXiv:2310.16944. find
  73. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. (2017). Attention is all you need. Advances in neural information processing systems, 30.
  74. Waleffe, R., Byeon, W., Riach, D., Norick, B., Korthikanti, V., Dao, T., Gu, A., Hatamizadeh, A., Singh, S., Narayanan, D., et al. (2024). An empirical study of mamba-based language models. arXiv preprint arXiv:2406.07887.
  75. Wang, J., Yan, J. N., Gu, A., and Rush, A. M. (2022). Pretraining without attention. arXiv preprint arXiv:2212.10544. find
  76. Wu, Y., Dukler, Y., Trager, M., Achille, A., Xia, W., and Soatto, S. (2024). Snakes and ladders: Accelerating SSM inference with speculative decoding. In Rezagholizadeh, M., Passban, P., Samiee, S., Partovi Nia, V., Cheng, Y., Deng, Y., Liu, Q., and Chen, B., editors, Proceedings of The 4th NeurIPS… find
  77. Xia, H., Ge, T., Wang, P., Chen, S.-Q., Wei, F., and Sui, Z. (2023). Speculative Decoding: Exploiting Speculative Execution for Accelerating Seq2seq Generation. In Findings of the Association for Computational Linguistics: EMNLP 2023, pages 3909–3925, Singapore. Association for Computational Linguis…
  78. Yang, J., Jimenez, C. E., Wettig, A., Lieret, K., Yao, S., Narasimhan, K., and Press, O. (2024a). Swe-agent: Agent computer interfaces enable software engineering language models. find
  79. Yang, N., Ge, T., Wang, L., Jiao, B., Jiang, D., Yang, L., Majumder, R., and Wei, F. (2023a). Inference with reference: Lossless acceleration of large language models. arXiv preprint arXiv:2304.04487. — [arXiv:2304.04487](https://arxiv.org/abs/2304.04487) find
  80. Yang, S., Wang, B., Shen, Y., Panda, R., and Kim, Y. (2023b). Gated linear attention transformers with hardware-efficient training. arXiv preprint arXiv:2312.06635. find
  81. Yang, S., Wang, B., Zhang, Y., Shen, Y., and Kim, Y. (2024b). Parallelizing linear transformers with the delta rule over sequence length. arXiv preprint arXiv:2406.06484. find
  82. Yao, S., Zhao, J., Yu, D., Du, N., Shafran, I., Narasimhan, K., and Cao, Y. (2022). React: Synergizing reasoning and acting in language models. arXiv preprint arXiv:2210.03629.
  83. Zellers, R., Holtzman, A., Bisk, Y., Farhadi, A., and Choi, Y. (2019). Hellaswag: Can a machine really finish your sentence? arXiv preprint arXiv:1905.07830. find
  84. Zhang, M., Bhatia, K., Kumbong, H., and Ré, C. (2024). The hedgehog & the porcupine: Expressive linear attentions with softmax mimicry. arXiv preprint arXiv:2402.04347. — [arXiv:2402.04347](https://arxiv.org/abs/2402.04347) find
  85. Zheng, L., Chiang, W.-L., Sheng, Y., Zhuang, S., Wu, Z., Zhuang, Y., Lin, Z., Li, Z., Li, D., Xing, E. P., Zhang, H., Gonzalez, J. E., and Stoica, I. (2023). Judging llm-as-a-judge with mt-bench and chatbot arena.

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/wang-2024-mamba-llama-distilling-accelerating-hybrid/