Source-linked AI summary
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
TL;DR
Decoder inference is limited by memory-bandwidth costs, while MQA can reduce quality and may require adapting models that were not trained with it. The paper uptrains existing MHA checkpoints and introduces GQA, whose intermediate key-value sharing achieves near-MHA quality at near-MQA speed. The authors report that 5% uptraining is beneficial, with diminishing returns at 10%.
Problem
Decoder inference incurs substantial memory-bandwidth overhead, while MQA can degrade quality and many existing language models use MHA instead.
Method
The paper uptrains existing MHA checkpoints for MQA and introduces GQA, which shares one key-value pair across each group of query heads.
Results
Uptrained GQA achieves quality close to MHA with speed close to MQA, and both MQA and GQA benefit from 5% uptraining with diminishing returns from 10%.
Takeaways & Limitations
GQA provides an intermediate attention design that trades key-value memory bandwidth against model quality between MHA and MQA.
Takeaways & Limitations
The evaluation covers only encoder-decoder models and does not compare the XXL GQA model with a model trained from scratch.
Abstract
from arXiv · showhide
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.
1 Introduction
The paper addresses decoder inference bottlenecks and the quality costs of multi-query attention by uptraining existing multi-head checkpoints and introducing grouped-query attention. Uptrained GQA remains close to multi-head quality while approaching multi-query speed.
- Decoder inference is bottlenecked by memory bandwidth required to load decoder weights and all attention keys and values at each step.
- MQA reduces this overhead with single key and value heads but can degrade quality and cause training instability.
- Existing multi-head checkpoints can be uptrained into MQA models using a small fraction of the original training compute.
- GQA uses single key and value heads for subgroups of query heads, interpolating between multi-head and multi-query attention.
- Uptrained GQA achieves quality close to multi-head attention while being almost as fast as multi-query attention.
2 Method
The method converts multi-head checkpoints by pooling key-value projections and briefly continuing pre-training, then uses grouped sharing to trade quality for inference efficiency. GQA spans the multi-head and multi-query endpoints while targeting favorable bandwidth and quality trade-offs.
- 2.1 Uptraining: Multi-query conversion first mean-pools key and value projection matrices into single projections, then continues pre-training on the same recipe.The conversion uses a small proportion α of the original training steps.
- 2.2 Grouped-query attention: GQA divides query heads into G groups, with each group sharing one key head and one value head.GQA-1 equals MQA, while GQA-H equals MHA.
- 2.2 Grouped-query attention: Intermediate group counts interpolate between MQA and MHA, offering higher quality than MQA with faster inference than MHA.
- 2.2 Grouped-query attention: GQA can reduce key-value cache loading and avoid replication waste from standard sharding of a single key and value head.
- 2.2 Grouped-query attention: GQA is not applied to encoder self-attention because encoder representations are computed in parallel and memory bandwidth is generally not the primary bottleneck.
3 Experiments
Experiments evaluate uptrained MQA and GQA models across summarization, translation, and question-answering, including inference-time and checkpoint-conversion ablations. Across tasks, GQA offers a strong quality–speed tradeoff, while conversion and uptraining choices materially affect performance.
- Experimental setup: The evaluation covers five summarization datasets, WMT translation, and TriviaQA question answering, using T5 Large, T5 XXL, and 5% uptrained MQA/GQA-8 XXL models.Timing uses per-sample TPUv4 measurements, while fine-tuning selects the checkpoint with highest dev performance.
- Main results: GQA achieves performance close to MHA-XXL with speed close to MQA, while uptrained MQA is higher-quality and faster than MHA-Large.These comparisons use average performance across all tasks and average inference time per sample.
- Checkpoint conversion: Mean pooling performs best for checkpoint conversion, followed by selecting one head and then random initialization.The ordering is attributed to how much information each method preserves from the pretrained model.
- Uptraining steps: GQA has reasonable performance immediately after conversion, whereas MQA requires uptraining to become useful.Both methods improve with 5% uptraining, but gains diminish at 10%.
- Number of groups: Increasing GQA groups from one initially causes modest inference slowdowns, with rising costs nearer MHA; the study selects eight groups as a middle ground.The number-of-groups experiment measures time per sample for GQA-XXL with input length 2048 and output length 512.
4 Related Work
Related work targets memory-bandwidth costs from loading keys and values through attention, while also addressing broader efficiency through alternative techniques. The paper’s uptraining procedure draws on prior checkpoint uptraining for sparse Mixture-of-Experts models.
- Attention efficiency: Prior work identifies multi-query attention as a way to reduce key-value memory-bandwidth overhead, especially for long inputs.The paper distinguishes this focus from methods that group attention heads without specifically targeting key-value heads.
- Alternative methods: Flash attention reduces memory and speeds training by avoiding materialized quadratic attention scores, while quantization reduces weight and activation size through lower precision.These methods address efficiency through computation or representation changes rather than the paper’s key-value-head grouping.
- Uptraining: The proposed uptraining procedure is inspired by work that converts standard T5 checkpoints into sparsely activated Mixture-of-Experts models.The cited prior work provides the methodological inspiration for checkpoint uptraining.
5 Conclusion
The paper proposes converting multi-head checkpoints to multi-query models with limited additional pre-training and introduces GQA to balance quality and inference speed. Its conclusion is that GQA approaches multi-head quality at multi-query-comparable speed.
- Conclusion: The paper converts multi-head models to multi-query models using a small fraction of original pre-training compute.This targets the memory-bandwidth overhead of loading keys and values during inference.
- Conclusion: Grouped-query attention interpolates between multi-query and multi-head attention, achieving quality close to multi-head at comparable speed to multi-query attention.GQA uses multiple key-value heads shared across groups of query heads.
Limitations
The evaluation is limited by uncertain summarization metrics, missing comparisons to models trained from scratch, and coverage restricted to encoder-decoder models.
- ROUGE is a flawed summarization measure, making the correctness of reported trade-offs difficult to establish.
- The XXL GQA model is not compared with a comparable model trained from scratch because computation was limited.
- Uptraining and GQA are evaluated only on encoder-decoder models, not decoder-only models.
- The authors expect GQA to have a stronger advantage over MQA in decoder-only models because they lack separate self-attention and cross-attention.
Acknowlegements
The acknowledgements thank Santiago Ontañón, Afroz Mohiuddin, William Cohen, and others at Google Research for insightful advice and discussion.
- Santiago Ontañón, Afroz Mohiuddin, and William Cohen are thanked by name.
- Others at Google Research are also thanked for insightful advice and discussion.
A Training Stability
The paper reports training instability for multi-query attention, especially on long-input tasks. Uptrained grouped-query attention appears stable, while uptrained multi-query models remain variable.
- Multi-query attention can cause training instability during fine-tuning, particularly on long-input tasks.
- T5-Large models trained with multi-query attention from scratch showed frequent pre-training loss spikes and immediate divergence during long-input fine-tuning.
- Uptrained multi-query models are more stable but still have high variance, so unstable-task results are averaged across three fine-tuning runs.
- Uptrained grouped-query attention models appear stable, but the root causes of multi-query instability were not investigated further.