Elamur Research⌕
☰
home / papers / 2024 / Simple linear attention language models balance the recall-throughput tradeoff

Simple linear attention language models balance the recall-throughput tradeoff

Paper · 2024 · Simran Arora, Sabri Eyuboglu, Michael Zhang et al.

Summary

Attention-based models are good at recalling information that has already appeared in the context (recall). However, the KV-cache grows with sequence length during generation, so these models use a lot of memory and have lower throughput. The authors, from Stanford and the University at Buffalo, combined experiments on the synthetic MQAR task with a theoretical lower bound. This allowed them to show that recall depends directly on the size of the model's recurrent state. Efficient alternatives such as Mamba, H3 and RWKV have a fixed-size state, and recall is exactly where they perform worse. As a solution, the authors propose the BASED architecture. It combines linear attention with a very narrow sliding window attention. The linear attention uses a feature map that is a second-order Taylor approximation of softmax. By varying these two parameters, one can move smoothly between the quality of full attention and the compactness of recurrent models. The authors also wrote IO-aware CUDA algorithms for prefill and decoding. On models with up to 1.3b parameters, BASED matches Mamba in perplexity, performs noticeably better on tasks that require recall, and generates text much faster. Thus, cheap generation is possible without a substantial loss in the ability to work with context.

Key points

  • On real-world tasks that require recall, standard attention outperforms Mamba by 32.2 accuracy points. BASED narrows this gap and outperforms Mamba by 10.36 points on average.
  • When a 1.3b-parameter model generates 1024 tokens (batch 128, H100 GPU), BASED achieves up to 24× higher throughput than FlashAttention-2.
  • The sliding window attention window is deliberately kept small, 64–128 tokens, so that it fully utilizes the tensor cores. For comparison, this window is 4096 tokens in Mistral-7B.
  • The models were trained on the Pile and also evaluated on DNA modeling and on the LM Eval Harness. They were compared with Transformer++, Mamba, Hyena and other architectures. The longer the training, the larger the advantage over Mamba: at the 360m scale it grew from 3.9 to 9.0 points, and at 1.3b from 9.0 to 10.4.
  • The code is open: github.com/HazyResearch/based. The paper appeared at ICML 2024 as part of the Efficient Systems for Foundation Models workshop.

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

🎓Stanford UniversityStanford · United States434 works · 314,419 citations🎓University at Buffalo, State University of New YorkBuffalo · United States12 works · 318 citations

 Fields, topics and concepts

📚 Shelves 5The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 worksComputer Vision and Pattern Recognition · 2,869 works

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

Cite

Cite as Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra and Christopher Ré. 2024. Simple linear attention language models balance the recall-throughput tradeoff. arXiv:2402.18668.

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

@misc{arora2024simple,
  title = {{Simple linear attention language models balance the recall-throughput tradeoff}},
  author = {Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra and Christopher Ré},
  year = {2024},
  doi = {10.48550/arxiv.2402.18668},
  eprint = {2402.18668},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2402.18668}
}

References

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

  1. Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. In Advances in Neural Information Processing Systems, 2022. find
  2. 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. International Conference on Learning Representations, 2023a.
  3. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. volume 30, 2017.
  4. Michael Zhang, Kush Bhatia, Hermann Kumbong, and Christopher Re. The hedgehog & the porcupine: Expressive linear attentions with softmax mimicry. In The Twelfth International Conference on Learning Representations, 2024. find
  5. Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209.11895, 2022.
  6. Feyza Duman Keles, Pruthuvi Mahesakya Wijewardena, and Chinmay Hegde. On the computational complexity of self-attention. In 34th International Conference on Algorithmic Learning Theory, volume 201, page 1–23, 2023. find
  7. Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, and Shruti Bhosale. Llama 2: Open foundation and fine-tuned chat models. arXiv:2307.09288, 2023.
  8. 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 Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 10630–10643, Online and Punta Ca… find
  9. 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 Mc-Donell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, An… find
  10. Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In International Conference on Machine Learning, pages 9355–9366. PMLR, 2021. find
  11. Manzil Zaheer, Guru Guruganesh, Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and et al. Big bird: Transformers for longer sequences. Proceedings of NeurIPS, 2020.
  12. Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In International conference on machine learning, pages 4055–4064. PMLR, 2018. find
  13. Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509, 2019.
  14. Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, and Christopher Ré. Scatterbrain: Unifying sparse and low-rank attention approximation. arXiv preprint arXiv:2110.15343, 2021a. — [arXiv:2110.15343](https://arxiv.org/abs/2110.15343) find
  15. Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150, 2020.
  16. Zhanpeng Zeng, Sourav Pak, Jeffrey Kline, Glenn Fung, and Vikas Sing. Multi resolution analysis (mra) for approximate self-attention. Proceedings of the 39 th International Conference on Machine Learning, 2022. find
  17. Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International conference on machine learning, pages 5156–5165. PMLR, 2020a. find
  18. Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00396, 2021. find
  19. Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models, 2023.
  20. Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009.14794, 2020.
  21. Ekin Akyürek, Bailin Wang, Yoon Kim, and Jacob Andreas. In-context language learning: Architectures and algorithms. 2024. find
  22. Krzysztof Choromanski, Haoxian Chen, Han Lin, Yuanzhe Ma, Arijit Sehanobish, Deepali Jain, Michael S Ryoo, Jake Varley, Andy Zeng, Valerii Likhosherstov, et al. Hybrid random features. arXiv preprint arXiv:2110.04367, 2021. — [arXiv:2110.04367](https://arxiv.org/abs/2110.04367) find
  23. Daniel Y. Fu, Elliot L. Epstein, Eric Nguyen, Armin W. Thomas, Michael Zhang, Tri Dao, Atri Rudra, and Christopher Ré. Simple hardware-efficient long convolutions for sequence modeling. arXiv preprint arXiv:2302.06646, 2023b. find
  24. Zhen Qin, Weixuan Sun, Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan, Lingpeng Kong, and Yiran Zhong. cosformer: Rethinking softmax in attention. arXiv preprint arXiv:2202.08791, 2022a. — [arXiv:2202.08791](https://arxiv.org/abs/2202.08791) find
  25. Yao-Hung Hubert Tsai, Shaojie Bai, Makoto Yamada, Louis-Philippe Morency, and Ruslan Salakhutdinov. Transformer dissection: a unified understanding of transformer’s attention via the lens of kernel. arXiv preprint arXiv:1908.11775, 2019. — [arXiv:1908.11775](https://arxiv.org/abs/1908.11775) find
  26. A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In Proceedings of the International Conference on Machine Learning (ICML), 2020b. find
  27. NVIDIA. Nvidia H100 tensor core GPU architecture, 2022. find
  28. A. Vyas, A. Katharopoulos, and F. Fleuret. Fast transformers with clustered attention. 2020. find
  29. Chen Zhu, Wei Ping, Chaowei Xiao, Mohammad Shoeybi, Tom Goldstein, Anima Anandkumar, and Bryan Catanzaro. Long-short transformer: Efficient transformers for language and vision. Advances in neural information processing systems, 34:17723–17736, 2021. find
  30. Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The Pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027, 2020. find
  31. Silas Alberti, Niclas Dern, Laura Thesing, and Gitta Kutyniok. Sumformer: Universal approximation for efficient transformers. arXiv preprint arXiv:2307.02301, 2023. — [arXiv:2307.02301](https://arxiv.org/abs/2307.02301) find
  32. Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding, 2023.
  33. Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. ACM Computing Surveys, 55(6):1–28, 2022.
  34. Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019. find
  35. Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nyströmformer: A nyström-based algorithm for approximating self-attention. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 14138– 14148, 2021. find
  36. Eric Wu, Kevin Wu, Roxana Daneshjou, David Ouyang, Daniel Ho, and James Zou. How medical ai devices are evaluated: limitations and recommendations from an analysis of fda approvals. Nature Medicine, 27:1–3, 04 2021. find
  37. Yifan Chen, Qi Zeng, Heng Ji, and Yun Yang. Sky-former: Remodel self-attention with gaussian kernel and nystr\\”om method. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021b. find
  38. Xiang Deng, Prashant Shiralkar, Colin Lockard, Binxuan Huang, and Huan Sun. Dom-lm: Learning generalizable representations for html documents. 2022. find
  39. Alexandre De Brebisson and Pascal Vincent. An exploration of softmax alternatives belonging to the spherical loss family. arXiv preprint arXiv:1511.05042, 2015. — [arXiv:1511.05042](https://arxiv.org/abs/1511.05042) find
  40. Simran Arora, Brandon Yang, Sabri Eyuboglu, Avanika Narayan, Andrew Hojel, Immanuel Trummer, and Christopher Ré. Language models enable simple systems for generating structured views of heterogeneous data lakes. arXiv:2304.09433, 2023b.
  41. Emmanuel Candes, Xiaodong Li, Yi Ma, and John Wright. Robust principal component analysis? arXiv:0912.3599, 2009. find
  42. Zhenhai Zhu and Rau Soricut. H-transformer-1d: Fast onedimensional hierarchical attention for sequences. n Annual Meeting of the Association for Computational Linguistics, 2021. find
  43. NVIDIA. Getting started with cuda graphs, 2019. URL https://developer.nvidia.com/blog/ cuda-graphs/. ↗
  44. Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportunities and risks of foundation models. arXiv preprint arXiv:2108.07258, 2021.
  45. Zhen Qin, Xiaodong Han, Weixuan Sun, Dongxu Li, Lingpeng Kong, Nick Barnes, and Yiran Zhong. The devil in linear transformer. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 7025–7041, Abu Dhabi, United Arab Emirates, December 2022b. Association for C…
  46. Laura Weidinger, John Mellor, Maribeth Rauh, Conor Griffin, Jonathan Uesato, Po-Sen Huang, Myra Cheng, Mia Glaese, Borja Balle, Atoosa Kasirzadeh, et al. Ethical and social risks of harm from language models. arXiv preprint arXiv:2112.04359, 2021.
  47. Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450, 2016.
  48. Nikita Kitaev, Łukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint arXiv:2001.04451, 2020.
  49. Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems, 32, 2019. find
  50. Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768, 2020. find
  51. James W Cooley and John W Tukey. An algorithm for the machine calculation of complex fourier series. Mathematics of computation, 19(90):297–301, 1965. find
  52. David W. Romero, Anna Kuzina, Erik J. Bekkers, Jakub M. Tomczak, and Mark Hoogendoorn. Ckconv: Continuous kernel convolution for sequential data. 2022. find
  53. Hanhwi Jang, Joonsung Kim, Jae-Eon Jo, Jaewon Lee, and Jangwoo Kim. Mnnfast: A fast and scalable system architecture for memory-augmented neural networks. In 2019 ACM/IEEE 46th Annual International Symposium on Computer Architecture (ISCA), pages 250–263, 2019. find
  54. Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces are as effective as structured state spaces, 2022. find
  55. Hao Liu and Pieter Abbeel. Blockwise parallel transformer for long context large models. arXiv preprint arXiv:2305.19370, 2023. find
  56. Albert Gu, Ankit Gupta, Karan Goel, and Christopher Ré. On the parameterization and initialization of diagonal state space models, 2022. find
  57. Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoc Le. Transformer quality in linear time. In International Conference on Machine Learning, pages 9099–9117. PMLR, 2022. find
  58. Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language modeling via gated state spaces, 2022. find
  59. Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus), 2023. find
  60. Xuezhe Ma, Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig, Jonathan May, and Zettlemoyer Luke. Mega: Moving average equipped gated attention. arXiv preprint arXiv:2209.10655, 2022. find
  61. Alex Wang, Yada Pruksachatkun, Nikita Nangia, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. SuperGLUE: a stickier benchmark for general-purpose language understanding systems. Curran Associates Inc., Red Hook, NY, USA, 2019.
  62. Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, et al. A mathematical framework for transformer circuits. Transformer Circuits Thread, 1, 2021. find
  63. Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring a broad discourse context, 2016.
  64. Yann N Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. In International conference on machine learning, pages 933–941. PMLR, 2017.
  65. Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence?, 2019. find
  66. Liliang Ren, Yang Liu, Shuohang Wang, Yichong Xu, Chenguang Zhu, and ChengXiang Zhai. Sparse modular activation for efficient sequence modeling. arXiv preprint arXiv:2306.11197, 2023. — [arXiv:2306.11197](https://arxiv.org/abs/2306.11197) find
  67. Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language, 2019.
  68. Hao Liu, Matei Zaharia, and Pieter Abbeel. Ring attention with blockwise transformers for near-infinite context. arXiv preprint arXiv:2310.01889, 2023.
  69. 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, 2018.
  70. Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Symposium on Operating Systems Principles, pages 611–626, 2023.
  71. Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale, 2019. find
  72. Eric Nguyen, Michael Poli, Marjan Faizi, Armin Thomas, Callum Birch-Sykes, Michael Wornow, Aman Patel, Clayton Rabideau, Stefano Massaroli, Yoshua Bengio, Stefano Ermon, Stephen A. Baccus, and Chris Ré. Hyenadna: Long-range genomic sequence modeling at single nucleotide resolution, 2023. find
  73. Daniel Y Fu, Hermann Kumbong, Eric Nguyen, and Christopher Ré. Flashfftconv: Efficient convolutions for long sequences with tensor cores. arXiv preprint arXiv:2311.05908, 2023c. find
  74. Katarina Gresova, Vlastimil Martinek, David Cechak, Petr Simecek, and Panagiotis Alexiou. Genomic benchmarks: A collection of datasets for genomic sequence classification. bioRxiv, 2022. doi: 10.1101/2022.06.08.495248. — [doi:10.1101/2022.06.08.495248](https://doi.org/10.1101/2022.06.08.495248) find
  75. Markus N Rabe and Charles Staats. Self-attention does not need o(n 2 ) memory. arXiv preprint arXiv:2112.05682, 2021. find
  76. Colin Lockard, Prashant Shiralkar, and Xin Luna Dong. OpenCeres: When open information extraction meets the semi-structured web. In Jill Burstein, Christy Doran, and Thamar Solorio, editors, Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguis…
  77. Peter Bürgisser, Michael Clausen, and Mohammad A Shokrollahi. Algebraic complexity theory, volume 315. Springer Science & Business Media, 2013. find
  78. Thathachar S Jayram, Ravi Kumar, and Dandapani Sivakumar. The one-way communication complexity of hamming distance. Theory of Computing, 4(1):129–135, 2008. find
  79. Swastik Kopparty. Topics in algorithms and complexity theory: Spring 2020. 2020. find
  80. The appendix is organized as follows: \#\# 1. Appendix A includes an extended related works discussion. 2. Appendix B includes details on the IO-aware implementation and benchmarking for B ASED . \#\# 3. Appendix C includes additional discussion of B ASED architectural details. \#\# 4. Appendix D provides… find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/arora-2024-simple-linear-attention-language-models/