Elamur Research
☰
home / papers / 2023 / GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints

GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints

Written for the Research library (claude-opus-5-5) from the paper text. Original — aclanthology.org/2023.emnlp-main.298.pdf

Summary

During autoregressive decoding, a Transformer model loads the attention keys and values from memory at every step. As a result, memory bandwidth becomes the bottleneck. Multi-query attention (MQA) keeps only a single key-value head and noticeably speeds up generation, but it degrades quality. Moreover, training a separate model solely for the sake of speed is expensive. The authors from Google Research propose two solutions. The first is uptraining: an existing MHA checkpoint is converted by averaging (mean pooling) the key and value projection matrices. The model is then further trained on a small fraction of the original pre-training, about 5%. The second solution is grouped-query attention (GQA). The query heads are divided into G groups, and each group has one shared key head and one value head. This is an intermediate option between MQA (G=1) and MHA (G=H). The method was evaluated on T5.1.1 Large and XXL on summarization, translation and question answering tasks. GQA-8-XXL is almost on par with MHA-XXL in quality, while approaching MQA in speed. This makes it possible to cheaply speed up already trained large models without retraining them from scratch.

Key points

  • Speed and quality (average across all tasks): MHA-XXL — 1.51 ms per sample, 47.2 points. MQA-XXL — 0.24 ms, 46.6 points. GQA-8-XXL — 0.28 ms, 47.1 points. For comparison, MHA-Large — 0.37 ms and 46.0 points, meaning the uptrained MQA-XXL is both faster and higher-quality than it.
  • Cost of uptraining: with α = 0.05, uptraining T5 XXL took about 600 TPUv3 chip-days. Just 5% already yields a noticeable gain, while returns diminish beyond 10%. GQA performs reasonably even right after conversion, whereas MQA is of little use without uptraining.
  • Conversion method: mean pooling of the heads works best. Selecting the first head is worse, and random initialization is the worst. The more information from the pre-trained model is preserved, the better the result.
  • Evaluation data: CNN/Daily Mail, arXiv, PubMed, MediaSum, Multi-News (ROUGE-1 metric), WMT 2014 English-to-German (BLEU) and TriviaQA (F1). GQA and MQA were applied only to decoder self-attention and cross-attention. 8 groups were chosen as a trade-off, since the slowdown grows as the number of groups increases.
  • Limitations: only encoder-decoder models were evaluated. There is no comparison with a GQA model trained from scratch. The authors themselves acknowledge that ROUGE is an imperfect metric. For decoder-only models, they expect GQA to have an even greater advantage over MQA.

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: AI & machine learning

Cited by

584 papers of our library cite it (of the 140673 whose bibliographies we read); 21 of them have a page here, 563 not yet.

Cite

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.18653/v1/2023.emnlp-main.298 · arXiv 2305.13245

@inproceedings{ainslie2023gqa,
  title = {{GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints}},
  author = {Ainslie, Joshua and Lee-Thorp, James and Jong, Michiel de and Zemlyanskiy, Yury and Lebron, Federico and Sanghai, Sumit},
  booktitle = {Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing},
  year = {2023},
  doi = {10.18653/v1/2023.emnlp-main.298},
  eprint = {2305.13245},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.18653/v1/2023.emnlp-main.298}
}

References

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

  1. James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. 2018. JAX: composable transformations of Python+NumPy programs. find
  2. Charlie Chen, Sebastian Borgeaud, Geoffrey Irving, Jean-Baptiste Lespiau, Laurent Sifre, and John Jumper. 2023. Accelerating large language model decoding with speculative sampling. CoRR, abs/2302.01318. find
  3. Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emi… find our card
  4. Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. 2022. Flashattention: Fast and memory-efficient exact attention with io-awareness. CoRR, abs/2205.14135. find
  5. Michiel de Jong, Yury Zemlyanskiy, Joshua Ainslie, Nicholas FitzGerald, Sumit Sanghai, Fei Sha, and William Cohen. 2022. FiDO: Fusion-in-decoder optimized for stronger performance and faster inference. arXiv preprint arXiv:2212.08153. find
  6. Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. 2022. Llm.int8(): 8-bit matrix multiplication for transformers at scale. CoRR, abs/2208.07339. find
  7. Alexander R. Fabbri, Irene Li, Tianwei She, Suyi Li, and Dragomir R. Radev. 2019. Multi-news: A large-scale multi-document summarization dataset and abstractive hierarchical model. In Proceedings of the 57th Conference of the Association for Computational Linguistics, ACL 2019, Florence, Italy, July… find our card
  8. Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. 2022. GPTQ: accurate post-training quantization for generative pre-trained transformers. CoRR, abs/2210.17323. find
  9. Google. 2020. Profile your model with cloud tpu tools. https://cloud.google.com/tpu/docs/ cloud-tpu-tools. Accessed: 2022-11-11. ↗
  10. Jianping Gou, Baosheng Yu, Stephen J. Maybank, and Dacheng Tao. 2021. Knowledge distillation: A survey. Int. J. Comput. Vis., 129(6):1789–1819. find
  11. Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Marc van Zee. 2020. Flax: A neural network library and ecosystem for JAX. find
  12. Geoffrey E. Hinton, Oriol Vinyals, and Jeffrey Dean. 2015. Distilling the knowledge in a neural network. CoRR, abs/1503.02531. find our card
  13. Mandar Joshi, Eunsol Choi, Daniel S. Weld, and Luke Zettlemoyer. 2017. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics, Vancouver, Canada. Association for Computationa… find our card
  14. Aran Komatsuzaki, Joan Puigcerver, James Lee-Thorp, Carlos Riquelme Ruiz, Basil Mustafa, Joshua Ainslie, Yi Tay, Mostafa Dehghani, and Neil Houlsby. 2022. Sparse upcycling: Training mixture-of-experts from dense checkpoints. find
  15. Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. 2019. GLUE: A multi-task benchmark and analysis platform for natural language understanding. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenRevie… find our card
  16. Samuel Williams, Andrew Waterman, and David A. Patterson. 2009. Roofline: an insightful visual performance model for multicore architectures. Commun. ACM, 52(4):65–76. find
  17. Chenguang Zhu, Yang Liu, Jie Mei, and Michael Zeng. 2021. Mediasum: A large-scale media interview dataset for dialogue summarization. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2021, Onl… find
  18. Yaniv Leviathan, Matan Kalman, and Yossi Matias. 2022. Fast inference from transformers via speculative decoding. CoRR, abs/2211.17192. find
  19. Ramesh Nallapati, Bowen Zhou, Cícero Nogueira dos Santos, Çaglar Gülçehre, and Bing Xiang. 2016. Abstractive text summarization using sequence-to-sequence rnns and beyond. In Proceedings of the 20th SIGNLL Conference on Computational Natural Language Learning, CoNLL 2016, Berlin, Germany, August 11-… find
  20. Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Anselm Levskaya, Jonathan Heek, Kefan Xiao, Shivani Agrawal, and Jeff Dean. 2022. Efficiently scaling transformer inference. arXiv preprint arXiv:2211.05102. find
  21. Markus Rabe. 2023. Memory-efficient attention. https://github.com/google/flaxformer/ blob/main/flaxformer/components/ attention/memory\_efficient\_attention.py. Accessed: 2023-05-23. ↗
  22. Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. 2020. Exploring the limits of transfer learning with a unified text-to-text transformer. J. Mach. Learn. Res., 21:140:1–140:67. find
  23. Noam Shazeer. 2019. Fast transformer decoding: One write-head is all you need. arXiv preprint arXiv:1911.02150. find
  24. Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. 2023. Llama: Open and efficient foundation language models. find