Elamur Research⌕
☰
home / papers / 2023 / Zoology: Measuring and Improving Recall in Efficient Language Models

Zoology: Measuring and Improving Recall in Efficient Language Models

Paper · 2023 · Simran Arora; Sabri Eyuboglu; Aman Timalsina; Isys Johnson; Michael Poli; James Zou; Atri Rudra; and Christopher Re

Summary

Gated-convolution language models such as H3, Hyena and RWKV are attractive because they scale sub-quadratically with sequence length. This paper asks whether they truly match Transformers. The authors trained 17 models with 5 architectures at sizes from 70M to 1.4B parameters, all on the same Pile data and setup. Attention still won by up to 2.1 perplexity points. Most of that gap came from associative recall: predicting a token by retrieving an association that appeared earlier in the same context.

Earlier synthetic recall benchmarks hid this weakness. To expose it, the authors propose multi-query associative recall (Mqar). In Mqar, many lookups happen at varied positions and draw on a large vocabulary. They also introduce BaseConv, an operator that can simulate gated-convolution architectures. Using it, they prove that solving Mqar requires model width that grows with sequence length, whereas attention needs no such growth.

The key fix is input-dependent mixing. Adding sparse attention at only a few layers, aimed at likely recall positions, nearly closes the gap while keeping costs sub-quadratic. This makes Mqar a useful design tool for efficient architectures.

Key points

  • Recall explains most of the gap. Tokens that require recall make up only 6.4% of Pile tokens, yet they account for about 82% of the perplexity gap. On these tokens, a 70M attention model beats a 1.4B Hyena model. The same gap appears at 7B scale when comparing RWKV with Llama-2.
  • Theory. BaseConv provably covers H3, Hyena, RWKV and RetNet. For these models, the dimension needed to solve Mqar scales with input length, while attention's does not.
  • Full-attention hybrid. At 360M parameters, BaseConv with only 3 attention layers beats the Transformer by 0.85 perplexity points while using about 18% fewer FLOPs.
  • Sparse-attention hybrids. Placing sparse attention on repeated-bigram positions closes 85% of the recall gap. A learned selector using k = 256 positions closes 72%. According to the abstract, input-dependent sparse hybrids recover 97.4% of the overall gap.
  • Code: github.com/HazyResearch/zoology

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 States506 works · 330,052 citations🎓Purdue University West LafayetteWest Lafayette · United States26 works · 2,076 citations🎓University at Buffalo, State University of New YorkBuffalo · United States13 works · 3,945 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 Simran Arora; Sabri Eyuboglu; Aman Timalsina; Isys Johnson; Michael Poli; James Zou; Atri Rudra; and Christopher Re. 2023. Zoology: Measuring and Improving Recall in Efficient Language Models. arXiv:2312.04927.

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

@misc{rudra2023zoology,
  title = {{Zoology: Measuring and Improving Recall in Efficient Language Models}},
  author = {Simran Arora; Sabri Eyuboglu; Aman Timalsina; Isys Johnson; Michael Poli; James Zou; Atri Rudra; and Christopher Re},
  year = {2023},
  doi = {10.48550/arxiv.2312.04927},
  eprint = {2312.04927},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2312.04927}
}

References

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

  1. 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. In International Conference on Learning Representations, 2023a. 1, 1, 2, 3.1, 3.2, 4.3, 5, 5, A.3, A.3, C.2, E.1 find
  2. 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. 1, A.3, A.3 find
  3. Junxiong Wang, Jing Nathan Yan, Albert Gu, and Alexander M Rush. Pretraining without attention. arXiv preprint arXiv:2212.10544, 2022. 1, 2 find
  4. Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. arXiv preprint arXiv:2302.10866, 2023a. 1, 1, 2, 3.1, 3.2, 4.3, A.2, A.3, B, C.2, E.1, 2, H.2, 8 find
  5. 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. 1, 2
  6. Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00396, 2021. 1, 2, A.2, A.3, 3, H.2.2 find
  7. Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, Xuzheng He, Haowen Hou, Przemyslaw Kazienko, Jan Kocon, and Jiaming et al. Kong. Rwkv: Reinventing rnns for the transformer era. arXiv:2305.13048, 2023. 1, 2… find
  8. Daniel Y. Fu, Simran Arora, Jessica Grogan, Isys Johnson, Sabri Eyuboglu, Armin W. Thomas, Benjamin Spector, Michael Poli, Atri Rudra, and Christopher Ré. Monarch mixer: A simple sub-quadratic gemm-based architecture, 2023b. 1, 2, A.3 find
  9. Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv preprint arXiv:1410.5401, 2014. 1, A.1 find
  10. Jimmy Ba, Geoffrey E Hinton, Volodymyr Mnih, Joel Z Leibo, and Catalin Ionescu. Using fast weights to attend to the recent past. Advances in neural information processing systems, 29, 2016. 1, A.1, H.7.1 find
  11. 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. 1, A.1 find
  12. 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. 1, 3.2, 4.3, A.1, E.1
  13. Shahar Lutati, Itamar Zimerman, and Lior Wolf. Focus your attention (with adaptive iir filters), 2023. 1 find
  14. Michael Hahn. Theoretical limitations of self-attention in neural sequence models. In Transactions of the Association for Computational Linguistics, volume 8, 2020. 1, 6 find
  15. William Merrill, Ashish Sabharwal, and Noah A. Smith. Saturated transformers are constant-depth threshold circuits. In Transactions of the Association for Computational Linguistics, volume 10, 2022. 1, 6 find
  16. 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. 1, 6 find
  17. 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. 2, 3.1, D, D.1 find
  18. 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. 2, 3.1, C.2, G Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob…
  19. 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. 2, A.2 find
  20. Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. ACM Computing Surveys, 55(6):1–28, 2022. 2, 4
  21. Karan Goel, Albert Gu, Chris Donahue, and Christopher Ré. It’s raw! audio generation with state-space models. Proceedings of the 39 th International Conference on Machine Learning,, 2022. 2 find
  22. Michael Zhang, Khaled Saab, Michael Poli, Tri Dao, Karan Goel, and Christopher Ré. Effectively modeling time series with simple discrete state spaces. International Conference on Learning Representations, 2023. 2 find
  23. Shuangfei Zhai, Walter Talbott, Nitish Srivastava, Chen Huang, Hanlin Goh, Ruixiang Zhang, and Josh Susskind. An attention free transformer. arXiv preprint arXiv:2105.14103, 2021. 2 — [arXiv:2105.14103](https://arxiv.org/abs/2105.14103) find
  24. Alex Andonian, Quentin Anthony, Stella Biderman, Sid Black, Preetham Gali, Leo Gao, Eric Hallahan, Josh Levy-Kramer, Connor Leahy, Lucas Nestler, Kip Parker, Michael Pieler, Jason Phang, Shivanshu Purohit, Hailey Schoelkopf, Dashiell Stander, Tri Songz, Curt Tigges, Benjamin Thérien, Phil Wang, and… find
  25. Ramin Hasani, Mathias Lechner, Tsun-Huang Wang, Makram Chahine, Alexander Amini, and Daniela Rus. Liquid structural state-space models. arXiv preprint arXiv:2209.12951, 2022. 3.1, A.3 — [arXiv:2209.12951](https://arxiv.org/abs/2209.12951) find
  26. 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. 3.1, A.3, D.1.2, 3, 6, H.2, H.2.3, 12
  27. Sabri Eyuboglu, Maya Varma, Khaled Saab, Jean-Benoit Delbrouck, Christopher Lee-Messer, Jared Dunnmon, James Zou, and Christopher Ré. Domino: Discovering systematic errors with cross-modal embeddings. In International Conference on Learning Representations, March 2022. 3.1 find
  28. David J Willshaw, O Peter Buneman, and Hugh Christopher Longuet-Higgins. Non-holographic associative memory. Nature, 222(5197):960–962, 1969. 3.1, A.1 find
  29. John J Hopfield. Neural networks and physical systems with emergent collective computational abilities. Proceedings of the national academy of sciences, 79(8):2554–2558, 1982. 3.1, A.1 find
  30. P. Bürgisser, T. Lickteig, M. Clausen, and A. Shokrollahi. Algebraic Complexity Theory. Grundlehren der mathematischen Wissenschaften. Springer Berlin Heidelberg, 1996. ISBN 9783540605829. URL https://books.google.com/books?id=dYcgjfXsYk8C. 4.1, H.13 ↗
  31. Tri Dao, Nimit S Sohoni, Albert Gu, Matthew Eichhorn, Amit Blonder, Megan Leszczynski, Atri Rudra, and Christopher Ré. Kaleidoscope: An efficient, learnable representation for all structured linear maps. arXiv preprint arXiv:2012.14966, 2020. 4.1, H.3, H.4, H.14, H.17, H.4, H.19 — [arXiv:2012.14966](https://arxiv.org/abs/2012.14966) find
  32. 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, 2023c. 4.3, A.3, C.2, 10, 13 find
  33. 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. 5, 4
  34. Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509, 2019a. 5
  35. Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150, 2020. 5, D.1.2, 4
  36. Liliang Ren, Yang Liu, Shuohang Wang, Yichong Xu, Chenguang Zhu, and ChengXiang Zhai. Sparse modular activation for efficient sequence modeling. Sparse Modular Activation for Efficient Sequence Modeling, 2023. 5, A.3 find
  37. 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. 6 find
  38. Tri Dao. FlashAttention-2: Faster attention with better parallelism and work partitioning. 2023. 6
  39. 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), 2020a. URL https://arxiv.org/abs/2006.16236. 6 ↗
  40. Chris Olah. Mechanistic interpretability, variables, and the importance of interpretable bases: An informal note on some intuitions related to mechanistic interpretability, 2022. URL https://transformer-circuits.pub/2022/mech-interp-essay/index.html. A.1, A.3 ↗
  41. Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Generalization beyond overfitting on small algorithmic datasets. arXiv:2201.02177, 2022. A.1
  42. Nick Cammarata, Shan Carter, Gabriel Goh, Chris Olah, Michael Petrov, Ludwig Schubert, Chelsea Voss, Ben Egan, and Swee Kiat Lim. Thread: circuits. Distill, 5(3):e24, 2020. A.1 find
  43. Dingquan Wang and Jason Eisner. The galactic dependencies treebanks: Getting more data by synthesizing new languages. Transactions of the Association for Computational Linguistics, 4:491–505, 2016. A.1
  44. Jennifer C White and Ryan Cotterell. Examining the inductive bias of neural language models with artificial languages. arXiv preprint arXiv:2106.01044, 2021. A.1 find
  45. Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 1, context-free grammar. arXiv preprint arXiv:2305.13673, 2023. A.1 find
  46. Shauli Ravfogel, Yoav Goldberg, and Tal Linzen. Studying the inductive biases of rnns with synthetic variations of natural languages. arXiv preprint arXiv:1903.06400, 2019. A.1
  47. Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. arXiv preprint arXiv:2111.02080, 2021. A.1, G Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint arXiv:2001.04451…
  48. JA Feldman, GE Hinton, and JA Anderson. Parallel models of associative memory, 1981. A.1 find
  49. Wei Zhang and Bowen Zhou. Learning to update auto-associative memory in recurrent neural networks for improving sequence memorization. arXiv preprint arXiv:1709.06493, 2017. A.1 — [arXiv:1709.06493](https://arxiv.org/abs/1709.06493) find
  50. Stefano Massaroli, Michael Poli, Daniel Y Fu, Hermann Kumbong, David Romero, Rom Parnichukun, Aman Timalsina, Quinn McIntyre, Beidi Chen, Atri Rudra, Ce Zhang, Christopher Ré, Stefano Ermon, and Yoshua Bengio. Laughing hyena distillery: Extracting compact recurrences from convolutions. 2023. A.2 find
  51. David W. Romero, Anna Kuzina, Erik J. Bekkers, Jakub M. Tomczak, and Mark Hoogendoorn. Ckconv: Continuous kernel convolution for sequential data. 2022. A.3 find
  52. Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces are as effective as structured state spaces, 2022. A.3 find
  53. Albert Gu, Ankit Gupta, Karan Goel, and Christopher Ré. On the parameterization and initialization of diagonal state space models, 2022. A.3 find
  54. Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language modeling via gated state spaces, 2022. A.3 find
  55. Jimmy T. H. Smith, Andrew Warrington, and Scott W. Linderman. Simplified state space layers for sequence modeling, 2023. A.3 find
  56. 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. A.3 find
  57. Transformers revolutionized ai. what will replace them?, 2023. URL https://www.forbes.com/ sites/robtoews/2023/09/03/transformers-revolutionized-ai-what-will-replace-them/?sh= 6ed698269c1f. A.3 ↗
  58. Jue Wang, Wentao Zhu, Pichao Wang, Xiang Yu, Linda Liu, Mohamed Omar, and Raffay Hamid. Selective structured state-spaces for long-form video understanding. In CVPR, 2023. A.3 find
  59. Brandon Yang, Gabriel Bender, Quoc V. Le, and Jiquan Ngiam. Condconv: Conditionally parametrized convolutions for efficient inference. In 33rd Conference on Neural Information Processing Systems (NeurIPS 2019), 2020. A.3 find
  60. Chrysoula Kosma, Giannis Nikolentzos, and Michalis Vazirgiannis. Time-parameterized convolutional neural networks for irregularly sampled time series. 2023. A.3 find
  61. Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention, 2020b. A.3 find
  62. Together. Redpajama: An open source recipe to reproduce llama training dataset, 2023. URL https://github.com/togethercomputer/RedPajama-Data. D, D.2, 4 ↗
  63. Sainbayar Sukhbaatar, Edouard Grave, Piotr Bojanowski, and Armand Joulin. Adaptive attention span in transformers. Association of Computational Linguistics, 2019. D.1.2
  64. 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. Advances in neural information processing systems, 33:1877–1901, 2020. 1
  65. Anonymous. Mamba: Linear-time sequence modeling with selective state spaces. 2023a. URL https://openreview.net/forum?id=AL1fq05o7H. 3, F.1 ↗
  66. Anonymous. Gateloop: Fully data-controlled linear recurrence for sequence modeling. 2023b. URL https://openreview.net/pdf?id=02Ug9N8DCI. 3, F.1 ↗
  67. Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509, 2019b. 4
  68. Jiezhong Qiu, Hao Ma, Omer Levy, Scott Wen tau Yih, Sinong Wang, and Jie Tang. Blockwise self-attention for long document understanding. arXiv preprint arXiv:1911.02972, 2019. 4 find
  69. Michael T Heideman and C Sidney Burrus. Multiplicative complexity, convolution, and the DFT. Springer, 1988. H.1.1 find
  70. Ilya Volkovich. A guide to learning arithmetic circuits. In Conference on Learning Theory, pages 1540–1561. PMLR, 2016. H.1.1 find
  71. Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. Proceedings of the 40th International Conference on Machine Learning, 2023b. H.2.1 find
  72. Thathachar S Jayram, Ravi Kumar, and Dandapani Sivakumar. The one-way communication complexity of hamming distance. Theory of Computing, 4(1):129–135, 2008. H.23 find
  73. Ofir Press, Noah A Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv:2108.12409, 2021. H.7.2 find
  74. Selim G Akl and Henk Meijer. Parallel binary search. IEEE Transactions on Parallel & Distributed Systems, 1(02):247–250, 1990. H.7.4, H.7.4 find
  75. Thomas H Cormen, Charles E Leiserson, Ronald L Rivest, and Clifford Stein. Introduction to algorithms. MIT press, 2022. H.32 find
  76. Miklós Ajtai, János Komlós, and Endre Szemerédi. An 0 (n log n) sorting network. In Proceedings of the fifteenth annual ACM symposium on Theory of computing, pages 1–9, 1983. H.35 find
  77. Chris Chatfield. The analysis of time series: An introduction, fifth edition. 1995. H.8.1 find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/arora-2023-zoology-measuring-improving-recall-efficient/