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

In plain words

When a Transformer writes text one token at a time, every step has to read back the keys and values it stored for all earlier tokens. With standard multi-head attention, each attention head keeps its own keys and values, so this cache is large and reading it from memory, not arithmetic, dominates the time. A known fix, multi-query attention, lets all heads share a single set of keys and values. That is much faster, but quality can drop and training can become unstable.

Grouped-query attention is the middle road. Query heads are split into a few groups, and each group shares one key head and one value head. Think of a library where, instead of every reader having a private copy of a book or everyone sharing one copy, each table gets its own copy: far fewer copies to fetch, but less crowding than a single shared one. With 8 groups in a T5-XXL model, the authors report quality almost equal to full multi-head attention at a speed close to multi-query attention.

The second idea is cheap conversion. Rather than training a new model from scratch, you take an existing multi-head checkpoint, average the key and value heads within each group into one, and continue pre-training for a small fraction (5%) of the original steps so the model adapts.

TL;DR

GQA partitions the H query heads into G groups; each group shares one key and one value projection, shrinking the KV cache by a factor H/G. GQA-1 is multi-query attention (MQA), GQA-H is multi-head attention (MHA). Existing MHA checkpoints are converted by mean-pooling the K/V projection matrices of the heads in each group, then "uptrained" for a fraction α (0.05 in the main results) of the original pre-training steps. On T5-XXL (encoder-decoder), uptrained GQA-8 reaches an average score of 47.1 vs 47.2 for MHA-XXL and 46.6 for MQA-XXL, at 0.28 s vs 1.51 s and 0.24 s inference time per sample (author-reported). Only encoder-decoder models were evaluated.

Problem

Autoregressive decoding is bottlenecked by memory bandwidth: each step loads the decoder weights plus all cached keys and values §1. MQA cuts the key/value load by using a single K/V head, but it can degrade quality and cause training instability, and many public models (T5, LLaMA) were trained with MHA, so getting an MQA model would mean training another model from scratch §1. Moving from MHA to MQA also cuts K/V capacity by a factor H, which grows with model size since larger models have more heads, making MQA an increasingly aggressive cut §2.2. Under standard model sharding, the single K/V head of MQA is replicated across partitions, wasting memory §2.2.

Method

Two independent contributions: a checkpoint conversion plus uptraining recipe, and the GQA attention variant §2.

Grouped-query attention §2.2 fig. 2. Query heads are divided into G groups; all query heads in a group attend using that group's single key head and single value head. GQA-G denotes G groups. GQA-1 equals MQA; GQA-H (G equal to the number of heads H) equals MHA. The KV cache, and thus the data loaded per decoding step, shrinks by H/G relative to MHA (the paper states the factor H for MQA §2.2; H/G for general G follows from the definition).

In the experiments GQA/MQA is applied to decoder self-attention and to cross-attention, but not to encoder self-attention, because encoder representations are computed in parallel and are not bandwidth-bound §2.2 §3.1.

Checkpoint conversion §2.1 fig. 1. To build a group's key (value) head, mean-pool the key (value) projection matrices of all original heads in that group. The paper compared this with taking the first head and with random initialization of the new K/V heads; mean pooling was best §3.3 fig. 4.

Uptraining §2.1 §3.1. Continue pre-training the converted checkpoint for a proportion α of its original pre-training steps, using the original pre-training recipe and dataset.

Algorithm: MHA -> GQA uptraining  [§2.1] [§2.2] [§3.1]
Input: MHA checkpoint with H heads per attention layer; number of groups G;
       uptraining proportion alpha; original pre-training recipe and data
1. For each decoder self-attention and cross-attention layer
   (encoder self-attention is left unchanged):
   a. Partition the H heads into G groups.
   b. For each group g: W_K[g] = mean of W_K[h] over heads h in group g;
                        W_V[g] = mean of W_V[h] over heads h in group g.
   c. Keep all H query projections and the output projection unchanged.
2. Continue pre-training with the original objective, data and optimizer
   settings for alpha * (original number of pre-training steps).
3. Fine-tune on downstream tasks as usual.
Output: GQA-G checkpoint

How heads are assigned to groups (contiguous blocks or otherwise) is not stated; step 1c is implied by the method (only K/V heads are pooled) but not spelled out. Attention per query head then works as in MHA, with query head h using the K/V head of its group:

\mathrm{head}_h = \mathrm{softmax}\!\left(\frac{(X W_Q^{h})(X' W_K^{g(h)})^\top}{\sqrt{d_k}}\right) X' W_V^{g(h)}

This equation is our notation, not one from the paper (the paper has no numbered equations); g(h) is the group of query head h, X' is the key/value source (decoder states for self-attention, encoder outputs for cross-attention). The scaling factor follows standard attention and is not discussed in the paper.

Hyperparameters

Results

Author-reported, dev sets. T5-Large and T5-XXL with MHA vs T5-XXL uptrained with α=0.05 to MQA or GQA-8. T_infer is time per sample per TPUv4 chip §3.1 tab. 1.

Source: tab. 1. fig. 3 plots the same averages against inference time.

Ablations on the average of CNN/Daily Mail, MultiNews and TriviaQA §3.3. Values are the plotted data points from the figure sources.

GQA is usable right after conversion (α=0), whereas MQA needs uptraining; both improve with 5% uptraining, with smaller gains from 10% §3.3 fig. 5.

Groups vs speed (GQA-XXL, input 2048, output 512) fig. 6:

Paper vs code

No official code. The paper points to Flaxformer as its framework and to an independently developed GQA implementation by Markus Rabe in Flaxformer (flaxformer/components/attention/memory_efficient_attention.py) §3.1 §4, but neither is released as this paper's code.

Pitfalls

• Encoder-decoder only. All experiments are T5 encoder-decoder models; decoder-only models were not evaluated. The authors expect GQA's advantage over MQA to be larger there, but this is a conjecture [Limitations].
• Which layers. GQA/MQA replaces decoder self-attention and cross-attention; encoder self-attention stays MHA §3.1. Converting only self-attention would not reproduce the setup.
• Group assignment is unstated. The paper does not say which heads form a group. Contiguous blocks of H/G heads is the natural reading of fig. 2 but is not specified; with mean pooling, the choice changes the initial weights.
• Pool the projections, not the cache. Conversion averages the K and V projection matrices §2.1 fig. 1; query and output projections are untouched (implied, not stated).
• Timing setups differ. tab. 1 reports MQA-XXL at 0.24 s, but fig. 6 reports 0.489 s for 1 group: fig. 6 uses fixed input 2048 / output 512, while tab. 1 averages across tasks. The axis of fig. 3 is labelled "ms" although its values match the seconds in tab. 1.
• Non-monotone speed. In fig. 6, 4 groups (0.476 s) is slightly faster than 1 group (0.489 s); the paper does not comment. Per §2.2, MQA's single K/V head is replicated across model partitions, which may be related, but the paper does not tie the two together.
• MQA instability. T5-Large MQA models trained from scratch had frequent pre-training loss spikes and diverged immediately when fine-tuned on long-input tasks; uptrained MQA was more stable but high-variance (MQA scores on unstable tasks are averages of 3 runs). Uptrained GQA appeared stable app. A.
• Hidden cost of uptraining. α=0.05 for T5-XXL was ~600 TPUv3 chip-days §3.1, and requires access to the original pre-training data and recipe.
• No from-scratch baseline. GQA-XXL was never compared with a GQA model trained from scratch, so uptraining vs scratch is unknown [Limitations].
• Metric caveat. Summarization is scored with ROUGE-1, which the authors call flawed for judging these trade-offs [Limitations].
• Mixed ablation model sizes. The conversion ablation uses T5-Large fig. 4, the uptraining ablation T5-XXL fig. 5; the ablation averages (around 55 to 57) are over a 3-task subset and are not comparable with tab. 1 averages.
• Fine-tuning optimizer details (optimizer, warmup) beyond LR, batch size and dropout are not stated §3.1.

Lineage

• Builds on: multi-query attention, Shazeer, "Fast transformer decoding: One write-head is all you need" (arXiv:1911.02150); generalizes it to G K/V heads.
• Builds on: sparse upcycling (arXiv:2212.05055), which inspired the uptraining of existing T5 checkpoints §4.
• Uses: T5.1.1 checkpoints. MQA is used in PaLM §1.
• Related: speculative decoding is listed as another way to ease the memory-bandwidth bottleneck §4.
• Compared against: MHA-Large, MHA-XXL, uptrained MQA-XXL tab. 1.