Elamur Research⌕
☰
home / papers / 2024 / Transformers to SSMs: Distilling Quadratic Knowledge to Subquadratic Models

Transformers to SSMs: Distilling Quadratic Knowledge to Subquadratic Models

Paper · 2024 · Aviv Bick, Kevin Y. Li, Eric P. Xing et al.

Summary

Subquadratic sequence models such as Mamba run inference more cheaply than Transformers, but they have received far less pretraining compute. They therefore lag behind the best Transformer models. This paper asks whether the effort already spent on a strong Transformer can be reused to build a strong state space model (SSM). The authors frame both attention and SSMs as matrix mixers, meaning matrices that combine token representations across the sequence. Their method, MOHAWK, transfers knowledge in three increasingly coarse stages:

  1. It aligns the student's mixing matrices with the teacher's attention matrices.
  2. It matches each block's hidden-state outputs independently.
  3. It copies the remaining weights and distills end-to-end predictions.

The teacher is Phi-1.5 (1.3B parameters). The student, Phi-Mamba, is a Mamba-2 variant shaped to mirror Phi-1.5's blocks, and distillation uses the C4 corpus. Phi-Mamba needs only about 3B tokens, under 1% of what comparable models use when trained from scratch. Even so, it clearly beats earlier open non-Transformer models of similar size. Ablations show that each stage contributes to this result. The work matters because it lets efficient architectures inherit capabilities from existing pretrained Transformers instead of being trained from scratch.

Key points

  • Token budgets and results: Phi-Mamba was distilled with 3B tokens and Hybrid Phi-Mamba with 5B tokens. On Winogrande, Phi-Mamba scores 71.7% against 60.9% for pretrained Mamba-2. On ARC-C, it scores 44.1% against 33.3%.

  • Hybrid model: Hybrid Phi-Mamba keeps 4 attention layers and swaps the other 20 for Mamba-2 blocks. It averages 66.0% on selected benchmarks, close to the 67.2% of its teacher, Phi-1.5. It also matches or beats Samba on several tasks, even though Samba has three times as many attention layers.

  • Stage ablation: Running only the final end-to-end distillation stage gives clearly worse results than the full three-stage MOHAWK pipeline.

  • Mixer comparison: When every mixer goes through identical training in Stages 2 and 3, the more expressive ones approximate attention more closely and achieve higher downstream scores:

    MixerBlock output L2 distanceAverage accuracy
    Causal Toeplitz9.635.0%
    Causal low-rank7.636.4%
    SSD (Mamba-2)5.559.2%

  • Frozen components: Training only the Mamba-2 blocks, with embeddings, head and MLPs frozen, costs little. Phi-Mamba's average drops from 62.7% to 61.4%, and the hybrid's from 64.0% to 63.9%.

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 5The field of science this work belongs to (OpenAlex fields).Computer Science · 19,503 worksArtificial Intelligence · 18,805 worksArts and Humanities · 17,959 worksLanguage and Linguistics · 17,789 worksComputer Vision and Pattern Recognition · 3,621 works

Cite

Cite as Aviv Bick, Kevin Y. Li, Eric P. Xing, J. Zico Kolter and Albert Gu. 2024. Transformers to SSMs: Distilling Quadratic Knowledge to Subquadratic Models. arXiv:2408.10189.

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

@misc{bick2024transformers,
  title = {{Transformers to SSMs: Distilling Quadratic Knowledge to Subquadratic Models}},
  author = {Aviv Bick, Kevin Y. Li, Eric P. Xing, J. Zico Kolter and Albert Gu},
  year = {2024},
  doi = {10.48550/arxiv.2408.10189},
  eprint = {2408.10189},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2408.10189}
}

References

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

  1. Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xLSTM: Extended Long Short-Term Memory. 2024. arXiv: 2405.04517 [cs.LG].
  2. Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. “PIQA: Reasoning about Physical Commonsense in Natural Language”. In: Proceedings of the AAAI conference on Artificial Intelligence. Vol. 34. 05. 2020, pp. 7432–7439.
  3. Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter,… — Language Models are Few-Shot Learners
  4. Shiv Chandrasekaran, Patrick Dewilde, Ming Gu, T Pals, and Alle-Jan van der Veen. “Fast stable solver for sequentially semi-separable linear systems of equations”. In: International Conference on High-Performance Computing. Springer. 2002, pp. 545–554. find
  5. Yen-Chun Chen, Zhe Gan, Yu Cheng, Jingzhou Liu, and Jingjing Liu. Distilling Knowledge Learned in BERT for Text Generation. 2020. arXiv: 1911.03829 [cs.CL]. find
  6. Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. “Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge”. In: arXiv preprint arXiv:1803.05457 (2018).
  7. Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. 2022. arXiv: 2205.14135 [cs.LG]. find
  8. Tri Dao and Albert Gu. “Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality”. In: International Conference on Machine Learning (ICML). 2024. find
  9. P. Dewilde and A. J. van der Veen. “On the Hankel-norm approximation of upper-triangular operators and matrices”. In: Integral Equations and Operator Theory 17.1 (Mar. 1, 1993), pp. 1–45. doi: 10.1007/BF01322544. url: https://doi.org/10.1007/BF01322544. — [doi:10.1007/bf01322544](https://doi.org/10.1007/bf01322544) ↗
  10. Patrick Dewilde and Alle-Jan Veen. Time-Varying Systems and Computations. 1st ed. Springer Book Archive. Springer Science+Business Media Dordrecht 1998. Springer New York, NY, 1998, pp. XIV, 460. isbn: 978-1-4757-2817-0. doi: https://doi.org/10.1007/978-1-4757-2817-0. url: https://doi.org/10.1007/97… — [doi:10.1007/978-1-4757-2817-0](https://doi.org/10.1007/978-1-4757-2817-0) ↗
  11. Daniel Y. Fu, Tri Dao, Khaled K. Saab, Armin W. Thomas, Atri Rudra, and Christopher Ré. Hungry Hungry Hippos: Towards Language Modeling with State Space Models. 2023. arXiv: 2212.14052 [cs.LG]. find
  12. Albert Gu. “Modeling Sequences with Structured State Spaces”. PhD thesis. Stanford University, 2023. find
  13. Albert Gu and Tri Dao. Mamba: Linear-Time Sequence Modeling with Selective State Spaces. 2023. arXiv: 2312.00752 [cs.LG]. find
  14. Albert Gu, Karan Goel, and Christopher Ré. Efficiently Modeling Long Sequences with Structured State Spaces. 2022. arXiv: 2111.00396 [cs.LG]. find
  15. Serkan Gugercin and Athanasios C. Antoulas. “A Survey of Model Reduction by Balanced Truncation and Some New Results”. In: International Journal of Control 77.8 (2004), pp. 748–766. doi: 10.1080/00207170410001713448. eprint: https://doi.org/10.1080/00207170410001713448. url: https://doi.org/10.1080/… — [doi:10.1080/00207170410001713448](https://doi.org/10.1080/00207170410001713448) ↗
  16. Suriya Gunasekar, Yi Zhang, Jyoti Aneja, Caio César Teodoro Mendes, Allie Del Giorno, Sivakanth Gopi, Mojan Javaheripi, Piero Kauffmann, Gustavo de Rosa, Olli Saarikivi, Adil Salim, Shital Shah, Harkirat Singh Behl, Xin Wang, Sébastien Bubeck, Ronen Eldan, Adam Tauman Kalai, Yin Tat Lee, and Yuanzhi… find
  17. Sangchul Hahn and Heeyoul Choi. Self-Knowledge Distillation in Natural Language Processing. 2019. arXiv: 1908. 01851 [cs.CL]. find
  18. Md. Akmal Haidar and Mehdi Rezagholizadeh. TextKD-GAN: Text Generation using KnowledgeDistillation and Generative Adversarial Networks. 2019. arXiv: 1905.01976 [cs.CL]. — [arXiv:1905.01976](https://arxiv.org/abs/1905.01976) find
  19. Ali Hatamizadeh and Jan Kautz. MambaVision: A Hybrid Mamba-Transformer Vision Backbone. 2024. arXiv: 2407. 08083 [cs.CV]. url: https://arxiv.org/abs/2407.08083. — [arXiv:2407.08083](https://arxiv.org/abs/2407.08083) ↗
  20. Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the Knowledge in a Neural Network. 2015. arXiv: 1503.02531 [stat.ML].
  21. Minghao Hu, Yuxing Peng, Furu Wei, Zhen Huang, Dongsheng Li, Nan Yang, and Ming Zhou. Attention-Guided Answer Distillation for Machine Reading Comprehension. 2018. arXiv: 1808.07644 [cs.CL]. find
  22. Sukjun Hwang, Aakash Lahoti, Tri Dao, and Albert Gu. Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers. 2024. arXiv: 2407.09941 [cs.LG]. url: https://arxiv.org/abs/2407.09941. — [arXiv:2407.09941](https://arxiv.org/abs/2407.09941) ↗
  23. Ananya Harsh Jha, Dirk Groeneveld, Emma Strubell, and Iz Beltagy. “Large Language Model Distillation Doesn’t Need a Teacher”. In: arXiv preprint arXiv:2305.14864 (2023). — [arXiv:2305.14864](https://arxiv.org/abs/2305.14864) find
  24. Jungo Kasai, Hao Peng, Yizhe Zhang, Dani Yogatama, Gabriel Ilharco, Nikolaos Pappas, Yi Mao, Weizhu Chen, and Noah A Smith. “Finetuning pretrained transformers into rnns”. In: arXiv preprint arXiv:2103.13076 (2021). find
  25. Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention. 2020. arXiv: 2006.16236 [cs.LG]. find
  26. Chen Liang, Haoming Jiang, Zheng Li, Xianfeng Tang, Bin Yin, and Tuo Zhao. “Homodistil: Homotopic task-agnostic distillation of pre-trained transformers”. In: arXiv preprint arXiv:2302.09632 (2023). — [arXiv:2302.09632](https://arxiv.org/abs/2302.09632) find
  27. Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz, Omri Abend, Raz Alon, Tomer Asida, Amir Bergman, Roman Glozman, Michael Gokhman, Avashalom Manevich, Nir Ratner, Noam Rozen, Erez Shwartz, Mor Zusman,… — Jamba: A Hybrid Transformer-Mamba Language Model
  28. Hao Liu, Matei Zaharia, and Pieter Abbeel. Ring Attention with Blockwise Transformers for Near-Infinite Context. 2023. arXiv: 2310.01889 [cs.CL].
  29. Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long Range Language Modeling via Gated State Spaces. 2022. arXiv: 2206.13947 [cs.LG]. — [arXiv:2206.13947](https://arxiv.org/abs/2206.13947) find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/bick-2024-transformers-ssms-distilling-quadratic-knowledge/