GQA: Training Generalized Multi-Query Transformer Models
J. Ainslie, J. Lee-Thorp, M. de Jong, Y. Zemlyanskiy, F. Lebrón, S. Sanghai
- Date
- 22 May 2023
- What the date marks
- arXiv submission
- Kind
- Architectural idea
- Organisation
- Google Research
- Access
- Apache 2.0
- Other dates
- EMNLP 2023
- Learning phases
- 1. Transformer architecture; 6. Local models and deployment
A turning point: 6 later milestones descend directly from it, and its descendants reach 9 of the 10 threads. How levels are decided
Drafted with Claude from the paper on 15 October 2026; reviewed by Zeeshan Arif. Read from arxiv.org (arXiv version 3 (December 2023) of the EMNLP 2023 conference paper).
What it changed
It made text generation faster at almost no cost in quality, by letting groups of attention heads share one set of keys and values. This method, grouped-query attention (GQA), sits between standard attention and the faster but weaker multi-query attention. Many later open models use it, including Llama 2, Mistral 7B and Llama 3.
The problem before it
When a language model writes text, it produces one token at a time. At each step, it must load the keys and values of all earlier tokens from memory, for every attention head. This memory traffic, not the arithmetic, is often what makes generation slow.
Multi-query attention (MQA) had cut this traffic sharply: all the query heads share a single key head and a single value head. But the authors note three problems:
- Quality can drop, and training can become unstable.
- Training a separate fast model only for generation is expensive.
- Many existing models, such as T5 and LLaMA, use standard multi-head attention (MHA), so they cannot benefit.
The idea
The mechanism. The paper has two parts.
- Grouped-query attention. Split the query heads into
G groups. Each group shares one key head and one value head. With one group, this is multi-query attention. With as many groups as heads, it is standard multi-head attention. A number in between keeps most of the quality and most of the speed. The authors chose 8 groups.
- Uptraining. Instead of training a new model from scratch, convert an existing multi-head model:
- For each group, average (mean-pool) the key and value projections of the heads in that group into one.
- Continue pre-training for a small part of the original training, for example 5% of the steps.
Averaging the heads worked better than keeping only the first head or starting the new heads from random values.
The authors argue that GQA suits large models well. Larger models have more heads, so a single shared head cuts their capacity a lot. GQA keeps the same ratio as the model grows.
The notation. GQA-G means grouped-query attention with G groups. GQA-1 is MQA and GQA-H, with H heads, is MHA. Going from MHA to MQA makes the key-value cache, the stored keys and values, H times smaller.
What it showed
Reported in the paper. It was published as a peer-reviewed conference paper at EMNLP 2023. We read version 3 of the arXiv paper. All models are T5 encoder-decoder models, and GQA and MQA were applied only to the decoder.
- Close to full quality, close to full speed. An uptrained T5-XXL with GQA-8 scored an average of 47.1 over 7 tasks, against 47.2 for the original multi-head T5-XXL and 46.6 for the MQA version.
- Much faster. Time per sample was 0.28 for GQA-8, 0.24 for MQA and 1.51 for the original multi-head XXL model. (The paper labels this unit as seconds in Table 1 and as milliseconds in Figure 3.)
- Big but fast beats small. Even the MQA version of T5-XXL was better and faster than the multi-head T5-Large model.
- Cheap to convert. Uptraining on 5% of the original steps took about 600 TPUv3 chip-days. GQA worked reasonably well right after conversion, while MQA needed the extra training. Training on 10% instead of 5% gave only a small further gain.
- More stable. Multi-query models trained from scratch had frequent loss spikes and failed when fine-tuned on long inputs. Uptrained GQA models were stable.
Limits, and what came later
The limits come from the paper:
- Hard to measure. The speed matters most for long outputs, whose quality is hard to judge. The summaries were scored with ROUGE, which the authors call a flawed measure.
- No from-scratch comparison. For lack of compute, the authors did not train a GQA model from scratch, so they cannot say whether uptraining is as good.
- Encoder-decoder models only. They did not test decoder-only models, the most popular kind. They expect GQA to help those even more.
- One model family. All tests used T5 models.
GQA spread quickly. Llama 2 used it in its 70-billion-parameter model, and Mistral 7B, Llama 3, Llama 3.1 and Qwen2 also use it. DeepSeek-V2 compared its own latent attention against GQA and reported an improvement.
What to know first
- Multi-head attention in the Transformer: queries, keys and values.
- Multi-query attention, which GQA generalises.
- How a model generates text one token at a time, storing the keys and values of earlier tokens.
How to read it
The paper is short: 5 pages of main text, with one page of appendix.
- Read section 1 for the problem.
- Read section 2 with Figures 1 and 2. This is the whole method.
- Look at Table 1 and Figure 3 for the main results.
- Read the limitations section and appendix A on training stability.
Expect about 30 minutes for the steps above (an estimate).
Links
- The paper: arXiv 2305.13245
- Code: the paper does not release its own code. It names the Flaxformer library, which its models are built on.
- Video: none is linked here yet.