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

Autoregressive decoding is often limited less by FLOPs than by memory bandwidth: every new token reloads model weights and the growing key/value (KV) cache. Multi-query attention (MQA) attacks that cost by keeping many query heads but only one key head and one value head. The catch is quality — and the practical annoyance that you may already own a strong multi-head (MHA) checkpoint you do not want to retrain from scratch.

This paper’s move is twofold. First, a cheap uptraining recipe converts an MHA checkpoint into MQA (or GQA) with about 5% of original pre-training compute. Second, grouped-query attention (GQA) interpolates the two extremes: query heads are partitioned into $G$ groups, and each group shares a single key head and value head. GQA-1 is MQA; GQA-H (groups equal to heads) is MHA.

What this paper explains

Decoder inference repeatedly reads keys and values for every past token. With multi-head attention you store and load $H$ key heads and $H$ value heads. MQA collapses that to one of each, cutting KV traffic by roughly a factor of $H$ — which is why Shazeer’s multi-query work (2019) and systems like PaLM leaned on it.

Ainslie et al. ask two engineering questions that matter if you ship models:

1. Can we retrofit existing MHA checkpoints instead of training a separate fast model?
2. Is there a better point on the tradeoff curve than “all heads” vs “one KV head”?

GQA is the architectural answer to (2). Mean-pooling KV projections and continuing pre-training for a small fraction $\alpha$ of original steps is the recipe for (1).

Prior limits

• MHA preserves capacity per head but makes the KV cache large; long-context decode pays for it on every step.
• MQA shrinks the cache aggressively, but the paper notes quality degradation and training instability risks, and many public models (T5, early LLaMA) were released as MHA.
• Training a second full model only for faster inference is expensive. Checkpoint surgery without adaptation also fails: the paper finds mean-pooling KV heads beats keeping a single head or random re-init, but still needs follow-on pre-training.

The mechanism

Grouped-query attention. Split $H$ query heads into $G$ groups. Inside a group, every query head attends using the same key head and value head for that group. Attention math stays familiar; only the KV multiplicity changes.

• $G = 1$ → MQA (one shared KV for all queries)
• $1 < G < H$ → GQA (the useful middle)
• $G = H$ → MHA (each query keeps its own KV)

Going from MHA to MQA reduces KV heads by $H$. Larger models often scale $H$ up, so that cut becomes brutal. GQA keeps a proportional number of KV heads so memory-bandwidth savings stay large while representation capacity does not collapse to a single KV pathway.

Uptraining. To convert an MHA checkpoint:

1. Mean-pool the key (and value) projection matrices within each intended group into one projection per group.
2. Continue the original pre-training recipe for $\alpha$ of the original step budget (main experiments use $\alpha = 0.05$).
3. Apply MQA/GQA to decoder self-attention and cross-attention (encoder self-attention stays multi-head in their T5 setup).

Algorithm walkthrough

1. Choose group count $G$ (paper highlights GQA-8 for large T5).
2. Partition the $H$ query heads into $G$ contiguous groups of size $H/G$ (assume $H$ divisible by $G$).
3. For each group, replace the group’s key projections with their mean (same for values).
4. Run a short continued pre-training phase with the new attention layout.
5. At decode time, load only $G$ key heads and $G$ value heads instead of $H$.

The opening figure places GQA between an MHA tower (private KV per query) and an MQA tower (one KV for everyone). The slider above lets you feel the same continuum by changing $G$.

What to notice when reading

• Where the paper draws the line between architecture (GQA) and adaptation (uptraining) — they solve different halves of the retrofit problem.
• How aggressively MQA collapses capacity as $H$ grows, and why a middle $G$ restores most quality in their tables.
• That encoder self-attention stays multi-head in their T5 recipe; the savings target decoder paths that dominate autoregressive cost.

Results and evidence

From the paper’s T5-XXL comparisons (their Table 1 / timing setup — treat as the authors’ reported numbers, not timeless absolutes):

Average aggregates their summarization (CNN/DailyMail, arXiv, PubMed, MediaSum, MultiNews; ROUGE-1), WMT En→De (BLEU), and TriviaQA (F1) dev results. Ablations show GQA needing less uptraining than MQA to recover after conversion, with diminishing returns past roughly 10% uptraining proportion.

Limitations

Stated by the authors, compressed:

• Rouge (and similar automatic metrics) are imperfect proxies for long-form generation quality — exactly where KV bandwidth matters most.
• They did not compare XXL GQA uptraining against an equally large GQA model trained from scratch.
• Experiments are on encoder–decoder T5; decoder-only models may change the MQA vs GQA gap (the authors expect GQA’s relative advantage over MQA to be stronger there because there is no separate cross-attention).
• The paper targets KV memory-bandwidth, not the full modern serving stack (paging, continuous batching, speculation).

How to read the paper

1. Abstract + §1 — why KV bandwidth dominates decode.
2. §2.1–2.2 — uptraining conversion and the GQA definition (Figures 1–2).
3. §3.1–3.2 — T5 setup, Table 1, Figure 3 tradeoff plot.
4. §3.3 ablations — uptraining proportion, group count, init choices.
5. Limitations + conclusion — what they did not claim.

Knowledge check

Keep reading

1. Original paper (arXiv:2305.13245) — abstract, Figures 1–3, Table 1.
2. Shazeer, Fast Transformer Decoding (MQA) — arXiv:1911.02150.
3. Related Fanout Daily: FlashAttention, PagedAttention, KV cache management, Orca continuous batching.

Sources

• Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, EMNLP 2023 — arXiv:2305.13245
• Shazeer, Fast Transformer Decoding: One Write-Head is All You Need (MQA), 2019 — arXiv:1911.02150