Elamur Research⌕
☰
home / papers / 2024 / An Empirical Study of Mamba-based Language Models

An Empirical Study of Mamba-based Language Models

Paper · 2024 · Roger Waleffe, Wonmin Byeon, Duncan Riach et al.

Summary

Selective state-space models (SSMs) such as Mamba could make inference cheaper than in Transformers. They avoid attention's quadratic cost and the ever-growing key-value cache. However, earlier comparisons with identical training data covered only small models. This NVIDIA-led report addresses that gap. It trains 8B-parameter Mamba, Mamba-2, Transformer, and hybrid models with the same hyperparameters on datasets of 1.1T and 3.5T tokens. All models are scored with a single shared evaluation pipeline across 35 tasks.

Pure SSMs perform well on many benchmarks. They fall behind on copying, in-context learning, and long-context reasoning, with five-shot MMLU and Phonebook lookup as clear examples. The hybrid mixes a small share of self-attention layers with Mamba-2 and MLP layers. It outperforms the Transformer on every short-context benchmark tested and roughly matches or beats it on long-context ones after extension. It is also projected to generate tokens far faster. The work provides rare controlled evidence that hybrid SSM designs are viable at scale, and the code and checkpoints are publicly available.

Key points

  • Training setup: the training data are precursors of the Nemotron-4 corpus (70% English, 15% non-English, 15% code), tokenized with a 256K SentencePiece vocabulary. Training and evaluation code is released in Megatron-LM.
  • Hybrid design: Mamba-2-Hybrid (8B) has 24 Mamba-2, 4 self-attention, and 28 MLP layers. Trained on 3.5T tokens, it beats the Transformer on all 12 short-context tasks by +2.65 points on average and by +3.5 points on five-shot MMLU. It is estimated to generate tokens up to 8× faster.
  • Pure SSM weakness: after 1.1T tokens, pure Mamba and Mamba-2 trail the Transformer by about 15 points on MMLU. More data narrows this gap, but Phonebook and long-context tasks remain difficult.
  • Long context: across 23 long-context tasks, the 16K and 32K hybrids match or exceed the Transformer on average. A 128K variant solves Phonebook perfectly beyond 150K tokens, though multi-document QA in LongBench remains weak.
  • Released weights: Mamba-2 8B and Mamba-2-Hybrid 8B are available on Hugging Face.

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

 Institutions

🎓Carnegie Mellon UniversityPittsburgh · United States598 works · 191,225 citations🎓Princeton UniversityPrinceton · United States113 works · 48,612 citations🎓University of Wisconsin–MadisonMadison · United States54 works · 49,089 citations

 Fields, topics and concepts

📚 Shelves 5The field of science this work belongs to (OpenAlex fields).Computer Science · 18,614 worksArtificial Intelligence · 17,965 worksArts and Humanities · 17,450 worksLanguage and Linguistics · 17,296 worksComputer Vision and Pattern Recognition · 3,486 works

🧭 Topics 1Research subjects and languages the work is about.Language models & pretraining · 3,894 works

Cite

Cite as Roger Waleffe, Wonmin Byeon, Duncan Riach, Brandon Norick, Vijay Korthikanti, Tri Dao, Albert Gu, Ali Hatamizadeh, Sudhakar Singh, Deepak Narayanan, Garvit Kulshreshtha, Vartika Singh, Jared Casper, Jan Kautz, Mohammad Shoeybi and Bryan Catanzaro. 2024. An Empirical Study of Mamba-based Language Models. arXiv:2406.07887.

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

@misc{waleffe2024empirical,
  title = {{An Empirical Study of Mamba-based Language Models}},
  author = {Roger Waleffe, Wonmin Byeon, Duncan Riach, Brandon Norick, Vijay Korthikanti, Tri Dao, Albert Gu, Ali Hatamizadeh, Sudhakar Singh, Deepak Narayanan, Garvit Kulshreshtha, Vartika Singh, Jared Casper, Jan Kautz, Mohammad Shoeybi and Bryan Catanzaro},
  year = {2024},
  doi = {10.48550/arxiv.2406.07887},
  eprint = {2406.07887},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2406.07887}
}

References

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

  1. Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. “GPT-4 Technical Report”. In: arXiv preprint arXiv:2303.08774 (2023).
  2. Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. “GQA: Training Generalized Multi-Query Transformer Models from Multi-head Checkpoints”. In: arXiv preprint arXiv:2305.13245 (2023).
  3. Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Ré. “Zoology: Measuring and Improving Recall in Efficient Language Models”. In: arXiv preprint arXiv:2312.04927 (2023).
  4. Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. “Layer Normalization”. In: arXiv preprint arXiv:1607.06450 (2016).
  5. Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. “Neural Machine Translation by Jointly Learning to Align and Translate”. In: arXiv preprint arXiv:1409.0473 (2014).
  6. Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li. “LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding”. In: arXiv preprint arXiv:2308.14508 (2023).
  7. 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.
  8. bloc97. “NTK-aware Scaled RoPE allows LLaMA models to have Extended (8k+) Context Size Without any Fine-tuning and Minimal Perplexity Degradation”. In: (2023). url: https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_%20scaled_rope_allows_llama_models_to_have. ↗
  9. Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. “Language Models are Few-shot Learners”. In: Advances in Neural Information Processing Systems 33 (2020), pp. 1877– 1901.
  10. 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).
  11. 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
  12. Pradeep Dasigi, Kyle Lo, Iz Beltagy, Arman Cohan, Noah A Smith, and Matt Gardner. “A Dataset of Information-Seeking Questions and Answers Anchored in Research Papers”. In: Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Languag…
  13. Soham De, Samuel L Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Muraru, Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, et al. “Griffin: Mixing Gated Linear Recurrences with Local Attention for Efficient Language Models”. In: arXiv preprint arXiv:2402.19427… find
  14. Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac’h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Ani… find
  15. Paolo Glorioso, Quentin Anthony, Yury Tokpanov, James Whittington, Jonathan Pilault, Adam Ibrahim, and Beren Millidge. “Zamba: A Compact 7B SSM Hybrid Model”. In: arXiv preprint arXiv:2405.16712 (2024). find
  16. Albert Gu and Tri Dao. “Mamba: Linear-time Sequence Modeling with Selective State Spaces”. In: arXiv preprint arXiv:2312.00752 (2023). find
  17. Albert Gu, Karan Goel, and Christopher Re. “Efficiently Modeling Long Sequences with Structured State Spaces”. In: International Conference on Learning Representations. 2021. find
  18. Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. “Measuring Massive Multitask Language Understanding”. In: International Conference on Learning Representations. 2020.
  19. Dan Hendrycks and Kevin Gimpel. “Gaussian Error Linear Units (GELUs)”. In: arXiv preprint arXiv:1606.08415 (2016). find
  20. Xanh Ho, Anh-Khoa Duong Nguyen, Saku Sugawara, and Akiko Aizawa. “Constructing A Multi-hop QA Dataset for Comprehensive Evaluation of Reasoning Steps”. In: Proceedings of the 28th International Conference on Computational Linguistics. 2020, pp. 6609–6625.
  21. Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, and Boris Ginsburg. “RULER: What’s the Real Context Size of Your Long-Context Language Models?” In: arXiv preprint arXiv:2404.06654 (2024).
  22. Samy Jelassi, David Brandfonbrener, Sham M Kakade, and Eran Malach. “Repeat After Me: Transformers are Better than State Space Models at Copying”. In: arXiv preprint arXiv:2402.01032 (2024). find
  23. Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. “Mistral 7B”. In: arXiv preprint arXiv:2310.06825 (2023). find
  24. Qiao Jin, Bhuwan Dhingra, Zhengping Liu, William W Cohen, and Xinghua Lu. “PubMedQA: A Dataset for Biomedical Research Question Answering”. In: arXiv preprint arXiv:1909.06146 (2019).
  25. Mandar Joshi, Eunsol Choi, Daniel S Weld, and Luke Zettlemoyer. “TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension”. In: Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). 2017, pp. 1601–1611.
  26. Tomáš Kočiskỳ, Jonathan Schwarz, Phil Blunsom, Chris Dyer, Karl Moritz Hermann, Gábor Melis, and Edward Grefenstette. “The NarrativeQA Reading Comprehension Challenge”. In: Transactions of the Association for Computational Linguistics 6 (2018), pp. 317–328.
  27. Vijay Korthikanti, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Mohammad Shoeybi, and Bryan Catanzaro. “Reducing Activation Recomputation in Large Transformer Models”. In: arXiv preprint arXiv:2205.05198 (2022).
  28. Taku Kudo and John Richardson. “Sentencepiece: A Simple and Language Independent Subword Tokenizer and Detokenizer for Neural Text Processing”. In: arXiv preprint arXiv:1808.06226 (2018).
  29. Guokun Lai, Qizhe Xie, Hanxiao Liu, Yiming Yang, and Eduard Hovy. “RACE: Large-scale ReAding Comprehension Dataset From Examinations”. In: arXiv preprint arXiv:1704.04683 (2017).
  30. Kenton Lee, Ming-Wei Chang, and Kristina Toutanova. “Latent Retrieval for Weakly Supervised Open Domain Question Answering”. In: arXiv preprint arXiv:1906.00300 (2019).
  31. Xin Li and Dan Roth. “Learning Question Classifiers”. In: COLING 2002: The 19th International Conference on Computational Linguistics. 2002.
  32. Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz, et al. “Jamba: A Hybrid Transformer-mamba Language Model”. In: arXiv preprint arXiv:2403.19887 (2024).
  33. Stephanie Lin, Jacob Hilton, and Owain Evans. “TruthfulQA: Measuring How Models Mimic Human Falsehoods”. In: arXiv preprint arXiv:2109.07958 (2021).
  34. Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. “Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answering”. In: arXiv preprint arXiv:1809.02789 (2018). find
  35. Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, et al. “Efficient Large-scale Language Model Training on GPU Clusters using Megatron-LM”. In: Proceedings of the Internatio… find
  36. NVIDIA. NVIDIA H100 Tensor Core GPU. https://www.nvidia.com/en-us/data-center/h100/. 2023. ↗
  37. Jongho Park, Jaeseung Park, Zheyang Xiong, Nayoung Lee, Jaewoong Cho, Samet Oymak, Kangwook Lee, and Dimitris Papailiopoulos. “Can Mamba Learn How to Learn? A Comparative Study on In-Context Learning Tasks”. In: arXiv preprint arXiv:2402.04248 (2024). find
  38. Jupinder Parmar, Shrimai Prabhumoye, Joseph Jennings, Mostofa Patwary, Sandeep Subramanian, Dan Su, Chen Zhu, Deepak Narayanan, Aastha Jhunjhunwala, Ayush Dattagupta, et al. “Nemotron-4 15B Technical Report”. In: arXiv preprint arXiv:2402.16819 (2024).
  39. Jonathan Pilault, Mahan Fathi, Orhan Firat, Chris Pal, Pierre-Luc Bacon, and Ross Goroshin. “Block-state Transformers”. In: Advances in Neural Information Processing Systems 36 (2024).
  40. Pranav Rajpurkar, Robin Jia, and Percy Liang. “Know what you don’t Know: Unanswerable Questions for SQuAD”. In: arXiv preprint arXiv:1806.03822 (2018).
  41. Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. “WinoGrande: An Adversarial Winograd Schema Challenge at Scale”. In: Communications of the ACM 64.9 (2021), pp. 99–106. find
  42. George Saon, Ankit Gupta, and Xiaodong Cui. “Diagonal State Space Augmented Transformers for Speech Recognition”. In: ICASSP 2023-2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP). IEEE. 2023, pp. 1–5. find
  43. Uri Shaham, Elad Segal, Maor Ivgi, Avia Efrat, Ori Yoran, Adi Haviv, Ankit Gupta, Wenhan Xiong, Mor Geva, Jonathan Berant, et al. “Scrolls: Standardized Comparison over Long Language Sequences”. In: arXiv preprint arXiv:2201.03533 (2022).
  44. Noam Shazeer. “GLU Variants Improve Transformer”. In: arXiv preprint arXiv:2002.05202 (2020). find
  45. Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. “Megatron-LM: Training Multi-billion Parameter Language Models using Model Parallelism”. In: arXiv preprint arXiv:1909.08053 (2019). find
  46. Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. “Roformer: Enhanced Transformer with Rotary Position Embedding”. In: Neurocomputing 568 (2024), p. 127063.
  47. Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. “Efficient Transformers: A Survey”. In: ACM Computing Surveys 55.6 (2022), pp. 1–28.
  48. Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. “Llama 2: Open Foundation and Fine-tuned Chat Models”. In: arXiv preprint arXiv:2307.09288 (2023).
  49. Harsh Trivedi, Niranjan Balasubramanian, Tushar Khot, and Ashish Sabharwal. “MuSiQue: Multihop Questions via Single-hop Question Composition”. In: Transactions of the Association for Computational Linguistics 10 (2022), pp. 539–554. — ♫ MuSiQue: Multihop Questions via Single-hop Question Composition
  50. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. “Attention is All You Need”. In: Advances in Neural Information Processing Systems 30 (2017).
  51. Wenhan Xiong, Jingyu Liu, Igor Molybog, Hejia Zhang, Prajjwal Bhargava, Rui Hou, Louis Martin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, et al. “Effective Long-context Scaling of Foundation Models”. In: arXiv preprint arXiv:2309.16039 (2023).
  52. Zhilin Yang, Peng Qi, Saizheng Zhang, Yoshua Bengio, William Cohen, Ruslan Salakhutdinov, and Christopher D Manning. “HotpotQA: A Dataset for Diverse, Explainable Multi-hop Question Answering”. In: Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing. 2018, pp. 2369…
  53. Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. “HellaSwag: Can a Machine Really Finish your Sentence?” In: arXiv preprint arXiv:1905.07830 (2019). find
  54. Biao Zhang and Rico Sennrich. “Root Mean Square Layer Normalization”. In: Advances in Neural Information Processing Systems 32 (2019). find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/waleffe-2024-empirical-study-mamba-based-language/