GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints
Authors: Joshua Ainslie
James Lee-Thorp
Michiel de Jong
Yury Zemlyanskiy
Federico Lebron
Sumit Sanghai
Venue: EMNLP 2023 Main (2023)
Keywords: efficient nlp
multi-query attention
fast inference
TL;DR: We present GQA, a generalization of multi-query and standard multi-head attention, and a recipe for uptraining multi-head models to fast multi-query or GQA versions.

Abstract: Multi-query attention (MQA), which only uses a single key-value head, drastically speeds up decoder inference. However, MQA can lead to quality degradation, and moreover it may not be desirable to train a separate model just for faster inference. We (1) propose a recipe for uptraining existing multi-head language model checkpoints into models with MQA using 5\% of original pre-training compute, and (2) introduce grouped-query attention (GQA), a generalization of multi-query attention which uses an intermediate (more than one, less than number of query heads) number of key-value heads. We show that uptrained GQA achieves quality close to multi-head attention with comparable speed to MQA.

GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints
Joshua Ainslie; James Lee-Thorp; Michiel de Jong"! Yury Zemlyanskiy, Federico Lebrón, Sumit Sanghai
Google Research
Abstract
show that language model checkpoints with multi- head attention (MHA) can be uptrained (Komat- suzaki et al., 2022) to use MQA with a small frac- tion of original training compute. This presents a cost-effective method to obtain fast multi-query as well as high-quality MHA checkpoints.
Multi-query attention (MQA), which only uses a single key-value head, drastically speeds up decoder inference. However, MQA can lead to quality degradation, and moreover it may not be desirable to train a separate model just for faster inference. We (1) propose a recipe for uptraining existing multi-head lan- guage model checkpoints into models with MQA using 5% of original pre-training pute, and (2) introduce grouped-query atten- tion (GQA), generalization of multi-query at- tention which uses an intermediate (more than one, less than number of query heads) number of key-value heads. We show that uptrained GQA achieves quality close to multi-head at- lention with comparable speed to MQA.
Second, propose grouped-quers attention (GQA), an interpolation between multi-head and multi-query attention with single key and value heads per subgroup of query heads. We show that uptrained GQA achieves quality close to multi- head attention while being almost as fast as multi- query attention.
2
Method
2.1 Uptraining
1 Introduction
Generating multi-query model from multi-head model takes place in two steps: first, converting the checkpoint, and second, additional pre-training allow the model to adapt to new structure. Fig- ure shows the process for converting multi-head checkpoint into a multi-query checkpoint. The pro- jection matrices for key and value heads are mean pooled into single projection matrices, which we find works better than selecting single key and value head randomly initializing new key and value heads from scratch.
Autoregressive decoder inference is a severe bottle- neck for Transformer models due to the memory bandwidth overhead from loading decoder weights and all attention keys and values at every decod- ing step (Shazeer, 2019; Pope et al., 2022; de Jong et al., 2022). The memory bandwidth from loading keys and values can be sharply reduced through multi-query attention (Shazeer, 2019), which uses multiple query heads but single key and value heads.
However, multi-query attention (MQA) can lead quality degradation and training instability, and it may not be feasible to train separate models optimized for quality and inference. Moreover, while some language models already use multi- query attention, such as PaLM (Chowdhery et al., 2022), many do not, including publicly available language models such T5 (Raffel et al., 2020) and LLaMA (Touvron et al., 2023).
Key Projection K;
4,
Key Projection
Mean Pool
Key Projection Kyjo
4,
Key Projection K4
This work contains two contributions for faster inference with large language models. First,
Figure I: Overview of conversion from multi-head to multi-query attention. Key and value projection matri- ces from all heads are mean pooled into single head.
'Equal contribution.
University of Southern Califomia. Work done at Gongle
Rescarch.
The converted checkpoint then pre-trained for

---

Multi-head
Grouped-query
Multi-query
Values
Keys
Queries
Figure 2: Overview of grouped-query method. Multi-head attention has query, key, and value heads. Multi-query attention shares single key and value heads across all query heads. Grouped-query attention instead shares single key and value heads for each group of query heads, interpolating between muati-head and multi-query attention.
a small proportion a of its original training steps the same pre-trate vve-traiiing recipe.
et al., 2022); GQA removes the waste from such partitioning. Therefore, we expect GQA present particularly good trade-off for larger models. We note that GQA is not applied the encoder self-attention layers; encoder representations are computed in parallel, and memory bandwidth is therefore generally not the primary bottleneck.
2.2 Grouped-query attention
Grouped-query attention divides query heads into G groups, each of which shares a single key head and value head. GQA-G refers to grouped-quers with G groups. GQA-1, with single group and therefore single key and value head, is equivalent to MQA, while GQA-H, with groups equal number of heads, is equivalent to MHA. Figure shows comparison of coouped-query attention and multi- head/multi-quee query attention. When converting a multi-head checkpoint to GQA checkpoint, we construct each group key and value head by mean- pooling all the original heads within that group. An intermediate number of groups leads to an interpolated model that is higher quality than MQA but faster than MHA, and, as we will show, rep- resents a favorable trade-off. Going from MHA to MQA reduces key and value heads to a sin- gle key and value head, reducing the size of the key-value cache and therefore amount of data that needs to be loaded by factor of II. However, larger models generally scale the number of heads, such that multi-query attention represents a more aggressive cut in both memory bandwidth and ca- pacity. GQA lets us keep the same proportional decrease in bandwidth and capacity as model size increases.
3
Experiments
3.1 Experimental setup
Configurations All models are based on the T5.1.1 architecture (Raffel et al., 2020), im- plemented with JAX (Bradbury et al., 2018), Flax (Heek al., 2020), and Flaxformer'. For main experiments
consider T5 Large and XXL with multi-head attention, trained versions of T5 XXL with multi-query and grouped-query attention. We use the Adafactor op- timizer with the same hyperparameters and learn- ing rate schedule as T5 (Raffel et al., 2020). We apply MQA and GQA to decoder self-attention and cross-attention, but not encoder self-attention.
well as up-
Uptraining Uptrained models are initialized from public 15.1.1 checkpoints. The key and value heads are mean-pooled to the appropriate MQA or GQA A structure, and then pre-trained for a further proportion of original pre-training steps with the original pre-training setup and dataset from (Raffel et al., 2020). For 0.05, training took approxi- mately 600 TPUv3 chip-days.
Moreover, larger models suffer relatively less from memory bandwidth overhead from attention, the KV-cache scales with model dimension while model FLOPs and parameters scale with the square of model dimension. Finally, standard shard- ing for large models replicates the single key and value head by the number of model partitions (Pope
Data We evaluate on summarization datasets CNN/Daily Mail (Nallapati et al., 2016), arXiv and PubMed (Cohan et al., 2018), MediaSum (Zhu et al., 2021), and Multi-News (Fabbri et al., 2019);
1https: //gi thub. com/ wonm/oogle//laxformon

---

Model
Tinfer
Average
CNN
arXiv
PubMed
MediaSum
MultiNews
WMT
TriviaQA
Ri
Ri
Ri
R₁
Ri
BLEU
F1
MHAgLarge
0.37
46.0
42.9
44.6
46.2
35.5
46.6
27.7
78.2
MHA-XXL
1.51
47.2
43.8
45.6
47.5
36.4
46.9
28.4
81.9
MQA-XXL
0.24
46.6
43.0
45.0
46.9
36.1
46.5
28.5
81.3
GQA-8-XXI.
0.28
47.1
43.5
45.4
47.7
36.3
47.2
28.4
81.6
Table I: Inference time and average dev set performance comparison of T5 Large and XXL models with multi- head attention, and uptrained TS-XXL models with multi-query and grouped-quers attention on summarization datasets CNN/Daily Mail, arXiv, PubMed, MediaSum, and MultiNews, translation dataset WMT, and question- answering dataset TriviaQA.
translation dataset WMT 2014 English-to-German; and question answering dataset TriviaQA (Joshi al., 2017). We do not evaluate on popular clas- sification benchmarks such as GLUE (Wang et al., 2019) as autoregressive inference less applicable for those tasks.
MIIA-XXL
GQA-XXL
47
Performance
46.5
MQA-XXL
Fine-tuning For fine-tuning, we use constant learning rate of 0.001, batch size 128, and dropout rate 0.1 for tasks. CNN/Daily Mail and WMT use input length of 512 and output length 256. Other summarization datasets use input length 2048 and output length 512. Finally, TriviaQA uses input length 2048 and output length 32. We train until convergence and select the checkpoint with the highest dev performance. We greedy decoding for inference.
46
MHA-Large
0.5
1
1.3
Time per sample (ms)
Figure 3: Uptrained MQA yields a favorable trade- off compared MHA with higher quality and faster speed than MHA-Large, and GQA achieves even better performance with similar speed gains and comparable quality to MHA-XXL. Average per- formance on all tasks as function of average inference time per sample
Timing We report time per sample per TPUv4 chip, as measured by xprof (Google, 2020). For timing experiments we use 8 TPUs with the largest batch size that fits up to 32 per TPU, and paral- lelization optimized separately for each model.
La-Lerge and T5 5XXXL with multi- head attention, and 5% uptrained T5-XXI. with MQA and GQA-8 8 attention.
3.2 Main results
uate performance on a representive subsample of tasks: CNN/Daily Mail, (short-form summariza- tion), MultiNews (long-form summarization), and TriviaQA (guestion-answering)
Figure shows average performance over all datasets as function of average inference time for MHA T5-Large and T5-XXL, and uptrained MQA and GQA-8 XXL models with uptraining 0.05. We see that larger up- trained MQA model provides favorable trade- off relative to MHA models, with higher quality and faster inference than MHA-Large. Moreover, GQA achieves significant additional quality gains, achieving performance close to MHA-XXI. with speed close to MQA. Table contains full results for all datasets.
Checkpoint conversion Figure 4 compares the performance of different methods for checkpoint conversion. Mcan pooling appears to work best, followed by selecting single head and then ran- dom initialization. Intuitively, results are ordered by the degree which information is preserved from the pre- pre-trained model.
proportion
Uptraining steps Figure 5 shows how perfor- mance varies with uptraining proportion for T5 XXL. with MQA and GQA. First, we note that GQA already achieves reasonable performance ter conversion while MQA A requires uptraining to
3.3 Ablations
This section presents experiments to investigate the effect of different modeling choices. We eval-

---

Mean
Time per sample
2
MHA
GQA
First
MQA
Random
54.4 54.6 54.8 55 55.2 55.4 55.6
8
16
32
64
GQA groups
Figure 4: Performance comparison different check- point conversion methods for T5-Large uptrained to MQA with proportion a = 0.05. 'Mean' mean-pools key and value heads, 'First' selects the first head and 'Random' initializes heads from scratch.
Figure 6: Time per sample for GQA-XXL. fune- tion of the number of GQA groups with input length 2048 and output length $12. Going from (MQA) 8 groups adds modest inference overhead, with increas- ing cost to adding more groups.
be useful. Both MQA and GQA gain from 5% uptraining with diminishing returns from 10%.
cially helpful for long inputs (Pope et al., 2022; de Jong et al., 2022). Rabe (2023) independently developed GQA with public implementation.
57
Performance
A number of other methods have been proposed to reduce memory bandwidth overhead from keys and values, as well parameters. Flash atten- tion (Dao et al., 2022) structures the attention com- putation to avoid materializing the quadratic tention scores, reducing memory and speeding up training. Quantization (Dettmers et al., 2022; Fran- tar et al., 2022) reduces the size of weights and activations, including keys and values, by lowering precision. Model distillation (Hinton et al., 2015; Gou et al., 2021) instead reduces model size at given precision, using data generated from the larger model finetune the smaller model. Layer- sparse cross-attention (de Jong al., 2022) elim- inates most cross-attention layers which make up the primary expense for longer inputs. Speculative sampling (Chen et al., 2023; Leviathan et al., 2022) ameliorates the memory bandwidth bottleneck by proposing multiple tokens with smaller model which are then scored in parallel by larger model.
MHA GQA
55
MQA
54
0
0.02
0.04
0.06
0.08
0.1
Uptraining proportion
Figure 5: Performance as function uptraining pro- portion for T5 XXL models with MQA and GQA-8.
Number of groups Figure 6 demonstrates the effect of the number of GQA groups on infer- ence speed. For larger models the memory band- width overhead from the KV cache is less con- straining (Shazeer, 2019), while the reduction in key-value size sharper due to the increased num- ber of heads. As a result, increasing the number groups from MQA only results in modest slow- downs initially, with increasing cost as we move closer to MHA. We selected 8 groups as a favor- able middle ground.
Finally, the uptraining procedure we propose is inspired by Komatsuzaki et al. (2022), which uptrains standard T5 checkpoints into sparsely acti- vated Mixture-of-LExpett models.
4 Related Work
Conclusion
This work is focused on achieving better trade- off between decoder quality and inference time through reducing the memory bandwidth over- head (Williams et al., 2009) from loading keys and values. Shazeer (2019) first proposed reducing this overhead through mulli-query attention. Follow-up work showed that multi-query attention espe-
Language models expensive for inference pri- marily due the memory bandwidth overhead from loading keys and values. Multi-query atten- tion reduces this overhead at the cost of decreased model capacity and quality. We propose to convert meati-head attention models to multi-query models

---

with small fraction of original pre-training com- pute. Moreover, we introduce grouped-query atten- tion, an interpolation of multi-query and multi-head attention that achieves quality close to multi-head at comparable speed to multi-query attention.
Yin, Toju Duke, Auselm Levskaya. Sanjay Ghe- mawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fe- dus. Denny Zhou, Daphne Ippolito. David Luan, Hyeontaek Lim, Barret Zoph. Alexander Spiridonov, Ryan Sepassi. David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankara- narayana Pillai, Marie Pellat, Aitor Lewkowycz. Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Bren- nan Saeta, Mark Diaz. Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier-Hellstern, Douglas Eek, Jeff Dean, Slav Petrov, and Noah Fiedel. 2022. Palm: Scaling language modeling with pathways.
Limitations
This paper focuses on ameliorating the memory bandwidth overhead from loading keys and values. This overhead is most important when generating longer sequences, for which quality is inherently difficult evaluate. For summarization we employ Rouge score, which we know is a flawed evaluation that does not tell the whole story; for that reason, it is difficult to be certain our trade-offs are cor- rect. Due to limited computation, we also do not compare our XXL GQA model to comparitive model trained from scratch, we do not know the relative performance of uptraining vs training from scratch. Finally, we evaluate the impact of uptrain- ing and GQA only on encoder-decoded models. Recently, decoder-only models are extremely pop- ular, and since these models do not have separate self-attention and cross-attention, we expect GQA to have a stronger advantage over MQA.
Arman Cohan, Franck Demonccuut, Doo Soon Kim, Trung Bui, Seokhwan Kim, Walter Chang, and Na- zli Goharian. 2018. A discourse-awee aware attention model for abstractive summarization of long docu- ments. Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Tech- nologies, Volume (Short Papers), pages 615-621, New Orleans, Louisiana. Association for Computa- tional Linguistics.
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.
Michiel de Jong, Yury Zemlyanskiy, Joshua Ainslie, Nicholas FitzGerald, Sumit Sanghai, Fei Sha, and William Cohen. 2022. FiDO: Fusion-in-decoder op- timized for stronger performance and faster infer- ence. arXiv preprint arXiv:2212.08153.
A Acknowlegements
We thank Santiago Ontañón, Afroz Mohiuddin, William Cohen and others at Google Research for insightful advice and discussion.
Tim Detimers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. 2022. Llm.it 8-bit ma- trix multiplication for transformers at scale. CoRR, abs/2208.07339.
References
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.
Alexander R. Fabbri, Irene Li, Tianwei She, Suyi Li. and Dragomir R. Radev. 2019. Multi-news: A large- scale multi-document summarization dataset and ab- stractive hierarchical model. In Proceedings of the 57th Conference of the Association for Computa- tional Linguistics, ACL 2019, Florence, Italy, July 28- August 2, 2019, Volume I: Long Pupers, pages 1074-1084. Association for Computational Linguis- tics.
Charlic Chen, Sebastian Borgeaud, Geoffrey Irv- ing, Jean-Baptiste Lespiau, Laurent Sifre, and John Jumper. 2023. Accelerating large language model decoding with speculative sampling. CoRR, abs/2302.01318.
Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. 2022. GPTQ: accurate bost-training quantization for generative pre-trained transformers. CoRR, abs/2210.1732..
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, Vin- odkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng
Google. 2020. Profile your model with cloud tpu
tools. https://cloud, google wooogl.cco//ppuuuu wttpttt/mo//eeeeeeeeeeeeeoooo docs/ cloud-tpu-tools Accessed: 2022-11-11.
Jianping Gou, Baosheng Yu, Stephen J. Maybank, and Dacheng Tao. 2021. Knowledge distillation: A sur- vey. Int. J. Comput. Vis., (2966):188--1181)

---

Jonathan Heek, Anselm Levskaya, Avital Oliver, Mar- vin Ritter, Bertrand Rondepierre, Andreas Steiner, and Marc van Zee. 2020. Flax: A neural network library and ecosystem for JAX.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman, 2019. GLUE; A multi-task benchmark and analysis plat- form for natural language understanding. In 7th International Conference on Learning Representa- lions, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net.
Geoffrey E. Hinton, Oriol Vinyals, and Jeffrey Dean. 2015. Distilling the knowledge neural network. CoRR, 0ass/5533.02331
Samuel Williams. Andrew Waterman, and David A. Patterson. 2009. Roofline: an insightful visual per- formance model for multicore architectures. Com- mun. ACM, 52(4):65-76.
Mandar Joshi, Eunsol Choi, Daniel S. Weld, and Luke Zettlemoyer. 2017. Triviaqa: large scale distantly supervised challenge dataset for reading comprehen- sion. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics, Van- couver, Canada. Association for Computational Lin- guistics.
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 Chap- ter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLI 2021, Online, June 6-11, 2021. pages 5927-5934, Associ- ation for Computational Linguistics.
Aran Komatsuzaki, Joan Puigcerver, James Lee-Thorp. Carlos Riguelme Ruiz, Basil Mustafa, Joshua Ainslie, Yi Tay. Mostafa Dehghani, and Neil Houlsby, 2022. Sparse upcycling: Training mixture- of-experts from dense checkpoints.
Yaniv Leviathan, Matan Kalman, and Yossi Matias. 2022. Fast inference from transformers via specu- lative decoding. CoRR, abs/2211.17199.
Ramesh Nallapati, Bowen Zhou, Cícero Nogueira dos Santos, Çaglar Gülçchre, 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-12, 2016, pages 280-290. ACL.
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 in- ference. arXiv preprint arXiv:2211.05702
Markus Rabe. 2023.
t/tp com Memory-efficient attention. blob/mainn/falxfraee /coppo/////////////////// Eitention/memory httpsmooooofffffffee afficcient_attentionn Accessed: 2023-05-23.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lec, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. 2020. Exploring the limits transfer learning with a unified text-to-text trans- former. J. Mach. Learn. Res., 21:140:1-140%
Noam Shazeer. 2019.
Fast transformer decoding: One write-head is all you need. arXiv preprint arxiv:1911.02150.
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 mod- els.

---

A Training Stability
We find that multi multi-query attention can lead to train- ing instability during fine-tuning, in particular com- bined with long input tasks. We trained multiple TS-Large models with multi-query attention from scratch, In each case, pre-training suffered from frequent loss spikes and the final models diverged immediately when fine-tuning on long-input tasks, Uptrained multi quelly attention models are more stable but still display high variance, so for multi- query models on unstable tasks we report average performance over three fine-tuning runs. Uptrained grouped-query attention models, however, appear to be stable, so we did not investigate futher on the root causes of multi-query instability.


Reviews:
[{"review_id": "l8FQssaIe3", "reviewer": "Reviewer_5bUv", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 2, "soundness": 2, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their comments, and address their points below.\n\n**Reason 1:** We indeed focus on inference speed. Inference and training operate in different compute regimes; training is bottlenecked by FLOPs, while inference is bottlenecked by memory bandwidth (https://arxiv.org/abs/2211.05102). Multi-query and grouped-query attention improve inference speed by reducing memory bandwidth overhead, and therefore are not applicable for training. However, inference speed is a critical property for modern LLMs, since LLMs can receive billions of user requests per day while only being trained infrequently.\n\n**Reason 2:** We agree that there likely exist methods that achieve better trade-offs. Some examples of possible methods include pruning the least valuable KV heads, using more KV heads for important query heads, or conditionally loading some KV heads for only some tokens. However, we feel that GQA represents a significant step forward over MQA, and will independently be of value to the community (as recently seen in the LLAMA2 paper).\n\n**Reason 3:** Indeed further exploration of different evals would be valuable. We note that we evaluate tasks with a variety of metrics, including WMT (BLEU) and TriviaQA (F1). Moreover, the recent LLAMA2 work has demonstrated similar results with different eval methods, providing further evidence that our results are robust.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The uptraining method proposed in this paper can effectively upgrade the existing MHA model to the MQA model, requiring only 5% of the computing power of the original pre-training, reducing the cost of model conversion.\n2. The proposed GQA method achieves a high inference speed while maintaining the model quality, and provides an effective solution for improving the inference speed of large language models.\n3. The experimental design is reasonable, and the results fully prove the effectiveness of the proposed method.", "reasons_to_reject": "1. This article mainly focuses on the improvement of inference speed, but there is less discussion on the optimization of the model training phase.\n2. Although the GQA method has achieved a good balance between quality and speed, there is still a certain decline in model capacity and quality. In the future, we can further explore how to improve model quality while maintaining speed.\n3. If the restriction analysis points out that the evaluation method (Rouge score) is flawed, the current experiment may not fully reflect the performance of the model, and further exploration is necessary in this work.", "questions_to_authors": "", "paper_topic_and_main_contributions": "This paper mainly studies how to improve the inference speed of large language models while maintaining the quality of the model. The author proposes an uptraining method to upgrade the existing multi-head model checkpoints to a multi-query model, and introduces a method called grouped-query attention (GQA) as a compromise between MQA and MHA . The experimental results show that the uptrained GQA is close to MHA in quality and close to MQA in operation speed.", "excitement": 3, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "GYdKARC4IO", "reviewer": "Reviewer_KbGZ", "summary": "", "questions": "1. How does applying GQA and MQA on the encoder with uptraining  affect performance?", "limitations": "", "rating": 4, "confidence": 2, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their questions and comments.\n\n**Reason 1:**\nAutoregressive decoding speed is typically bottlenecked by memory bandwidth (https://arxiv.org/abs/2211.05102) rather than FLOPs, as all model weights as well as the KV cache have to be loaded to SRAM for every token (see the original multi-query paper for a discussion). Multi-query and grouped-query attention increase speed not by decreasing the FLOPs, but rather by reducing the size of the KV cache and therefore memory bandwidth overhead.\n\n**Question 1:**\nThis is a subtle point that we will clarify in the paper. The encoder representation can be computed in parallel, not token by token, and this computation is FLOPs bound, not memory bandwidth bound. Because MQA and GQA function by reducing memory bandwidth overhead, neither method is applicable to speeding up the encoder.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The paper proposes that up-training and grouping of query-headed attention while maintaining fast inference of query-headed attention and similar performance of multi-headed attention. This is useful for time-constrained tasks.\n2. The paper is easy to read,\n3. The paper performs ablations on uptraining steps, number of groups", "reasons_to_reject": "1. The authors uses time per sample to measure model efficiency. It may have been more useful to also report FLOPS to better compare with other methods.", "questions_to_authors": "1. How does applying GQA and MQA on the encoder with uptraining  affect performance?", "paper_topic_and_main_contributions": "The paper proposes uptraining for multi-query attention and transform multi-query attention to grouped-query attention where query heads are divided into  groups, each of which shares a single key head and value head. This leads to improvements in performance while maintaining the fast inference time.", "excitement": 4, "reproducibility": 3, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "Ac270FxSeh", "reviewer": "Reviewer_B4s9", "summary": "", "questions": "1. What is the result of other versions such as 4, 16, 32 as well as GQA-8 in Table 1?\n\n2. How does the performance of the model change when the alpha value exceeds 0.1?\n\n3. Why did you apply MQA and GQA only to the decoder and not to the encoder?\n\n4. What is the minimum number of groups to stabilize training with GQA?", "limitations": "", "rating": 3, "confidence": 4, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their extensive questions and comments.\n\n**Reason 1:**\n\na. This is a question also raised by KbGZ, and reflects an important point about encoder-decoder inference, which we will clarify in the paper. The encoder representation can be computed in parallel, not token by token, and this computation is FLOPs bound, not memory bandwidth bound. Because MQA and GQA function by reducing memory bandwidth overhead, neither method is applicable to speeding up the encoder.\n\nb. Thanks, the mentioned details are indeed missing, and we will add them. The optimizer is Adafactor, with the same hyperparameters and inverse sqrt learning rate schedule as T5 for uptraining, and a constant 1e-3 learning rate for fine-tuning. The uptraining dataset is C4 for alpha% of 1T tokens, where alpha=5% took approximately 600 TPUv3 chip-days (see the original T5 paper for more details on the recipe and computational cost). \n\n**Reason 2:**\n\na. We apologize for a point of confusion here. The instability is not in uptraining but rather the resulting fine-tuning. Moreover, it is not the case that increasing the uptraining ratio leads to worse results (in fact it leads to better results and more stable fine-tuning). Since the instability occurs primarily in the baseline method, and our GQA procedure solves the problem, we decided not to spend further resources on exploring the exact parameters where it occurs. We will clarify in the paper.\n\nb. This is a very valid point. Unfortunately, while profiling is cheap, fine-tuning large long-input models is very resource intensive, and we did not have the ability to evaluate performance for all values of the numbers of queries per group. We argue that GQA-8 is both very close to MQA in speed, and close to MHA in quality, and therefore must be close to optimal. Since then, the LLAMA2 work has also shown this to be a reasonable value.\n\n\n**Reason 3**:\nDue to limited computational resources, we did not experiment extensively with more than 8 query heads per KV head. In practice, more than 8 query heads per KV head does not lead to significant speedup and we do not expect this to be a common hyperparameter choice. \n\n\n**Question 1:** We direct the reviewer to our response to reason 2.b\n\n**Question 2:** Since improvement from alpha=0.05 to 0.1 is already modest, and since uptraining for longer is very costly, we did not explore higher values of alpha.\n\n**Question 3:** We direct the reviewer to our response to reason 1.a\n\n**Question 4:** We direct the reviewer to our response to reason 3.a", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. Proposing a relatively simple yet effective methodology to reduce the computational overhead of LM.\n2. Including experimental results with various parameter configurations for the suggested methodology.\n3.  well-structured and easy to read paper.", "reasons_to_reject": "1. lacking of detail about experiment design and settings.\n\n1-a) Why  MQA and GQA are applied only to the decoder, while the encoder was left unaffected during the experiments.\n1-b) Missing detail about the optimizer, the total number of epochs, and the time taken to train the entire dataset.\n\n2. lacking of sufficient experimental results to support the claims made by the authors.\n\n2-a) The authors mentioned that the model's training is unstable. Therefore, increasing the uptraining ratio (alpha) may lead to worse results. However, in Figure 5, it can be observed that the model's performance consistently improves as the maximum alpha ratio increases up to 0.1. There is a lack of experiment results regarding the point at which the performance starts to decline as the ratio further increases.\n\n2-b) The authors only demonstrated in the paper that as the number of queries per group increases in GQA, the time also increases, as shown in Figure 6. However, it remains unclear how the performance actually changes with the increase in the number of queries per group.\n\n\n\n3. The lack of detail about the claims.\n\n3-a) GQA has superior training stability compared to MQA. However, there was no mention of the details of the correlation between group size and stability (how much group size guarantees stability) are not mentioned.", "questions_to_authors": "1. What is the result of other versions such as 4, 16, 32 as well as GQA-8 in Table 1?\n\n2. How does the performance of the model change when the alpha value exceeds 0.1?\n\n3. Why did you apply MQA and GQA only to the decoder and not to the encoder?\n\n4. What is the minimum number of groups to stabilize training with GQA?", "paper_topic_and_main_contributions": "The authors proposed MQA (Multi Query Attention) and GQA (Grouped-query attention) to overcome the memory bandwidth issue caused by the loading of multiple keys and values in the self-attention mechanism of the LM model.\nMulti-query attention utilizes only a single key and value, resulting in faster computation speed but lower model capacity and quality. To address this limitation, the authors introduced GQA, which uses several grouped keys and values, achieving similar speed to MQA while maintaining comparable quality to the Multi-head approach.", "excitement": 3, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: There is a consensus that this work is sound and exciting. The reviewers highlight that the method and paper are easy to understand and yield a straightforward inference improvement at negligible performance degradation. While some concerns have been raised regarding the evaluation in terms of ROUGE, the authors provide enough evidence of strong performance across other metrics.
Decision: Accept-Main