ZeroShotMind

Paper

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

Grouped-query attention interpolates between multi-head and multi-query attention: groups of query heads share one key/value head. It cuts the KV cache by the grouping factor while keeping nearly all of multi-head's quality, and it can be uptrained cheaply from an existing multi-head checkpoint.

Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, Sumit Sanghai — Google Research2023arXiv ↗Views:

attentionkv-cacheinference

The KV cache is what makes decoding expensive

Autoregressive generation is memory-bandwidth bound. Each new token requires reading the keys and values of every previous token from the KV cache, and that cache is large: in standard multi-head attention (MHA), every one of the model's attention heads stores its own key and value vector per token. With dozens of heads across dozens of layers, the cache for a long sequence runs into many gigabytes, and the time to stream it from memory — not the matmul — sets the pace of decoding. Shrinking the KV cache is therefore one of the highest-leverage things you can do for inference speed, and the key/value heads are the obvious target, because that is where the cache's size comes from.

The two endpoints: MHA and MQA

Multi-query attention (MQA), proposed earlier by Noam Shazeer, takes the aggressive route: keep all the query heads but collapse them onto a single shared key/value head. Every query head attends using the same keys and values. This shrinks the KV cache by the full head count — if a model has 64 query heads, MQA stores one KV head instead of 64, a 64× cache reduction — and decoding speeds up dramatically. The problem is quality. A single shared KV head is a real bottleneck: it measurably degrades the model on harder tasks, and it can make training less stable. MQA cut the cache, but it cut too far.

MHA and MQA are the two ends of a spectrum — one KV head per query head, versus one KV head for all of them. The paper's observation is that nothing forces you to pick an endpoint.

Grouped-query attention: the stable middle

Grouped-query attention (GQA) partitions the query heads into GG groups, and each group shares a single key/value head. With GG equal to the number of query heads you recover MHA; with G=1G = 1 you recover MQA; in between you get a tunable trade. A typical configuration uses 8 KV heads for 64 query heads — eight groups of eight query heads each — which cuts the KV cache by 8× relative to MHA while leaving quality almost untouched.

The reason this middle point works is that the quality cost of sharing KV heads is steeply nonlinear. Going from a KV head per query head down to a handful of KV heads sacrifices very little, because a modest number of distinct key/value subspaces is enough to capture most of what the full set provided; it is only the final collapse all the way to one shared head that hurts. GQA sits where the cache savings are nearly as large as MQA's but the quality loss has not yet kicked in — it captures most of the memory benefit of multi-query while staying close to the accuracy of multi-head.

Uptraining: you don't have to start over

The second half of the contribution is procedural and is what made GQA cheap to adopt. You do not need to pretrain a GQA model from scratch. Given an existing multi-head checkpoint, you can convert it: construct each group's shared key/value head by mean-pooling the original KV heads that fall into that group, then continue training — "uptraining" — for a small fraction of the original pretraining compute to let the model adapt to the shared heads. The paper shows this recovers essentially full quality after only a few percent of the original training budget. The same recipe converts an MHA checkpoint to MQA for comparison. This is what turned GQA from a clean idea into a default: labs sitting on expensive multi-head models could get the inference savings without paying to retrain, simply by mean-pooling and briefly continuing training.

Why it became standard

GQA is now the default attention scheme across the open frontier. LLaMA 2 introduced it on its 34B and 70B models; LLaMA 3 uses it across all sizes; Mistral, Mixtral, Qwen, and most other recent open models adopt it. The reason is that it lands on the right point of the trade-off and asks almost nothing in return: the KV cache shrinks by the grouping factor — directly translating into a larger batch you can hold in memory, longer contexts you can serve, and faster decoding because there are fewer bytes to stream per token — while downstream quality stays within noise of the multi-head baseline. It also composes cleanly with the rest of the modern stack: rotary position embeddings act on the query and key vectors and work unchanged under grouping, and FlashAttention's tiling handles grouped heads directly.

Where it sits in the lineage

GQA is best read as the resolution of a tension MQA opened. MQA proved the KV cache could be slashed by sharing key/value heads, but paid for it in quality; GQA found the dial between the extremes and showed the sweet spot is much closer to MHA's quality than MQA's aggressiveness suggested. It is part of a broader and continuing effort to make attention cheaper to serve rather than to compute — alongside KV-cache quantization, attention-based eviction, and the more radical latent-attention schemes that compress keys and values into a shared low-rank space. Among these, GQA is the one that is essentially free: a small architectural change, a cheap conversion recipe, and almost no quality cost.

Limitations

The grouping factor is a hyperparameter, and the right value depends on the model and the deployment — too few KV heads and quality starts to slip toward MQA's, too many and the cache savings shrink toward MHA's, so the 8-KV-head convention is a well-tested default rather than a derived optimum. GQA reduces the cache by a fixed factor but does not change its fundamental growth with sequence length and batch size, so it composes with rather than replaces quantization and eviction — and for the largest models chasing extreme cache compression, latent attention pushes further than grouping can. And the mean-pooling-then-uptrain recipe assumes you have a multi-head checkpoint to convert; a model designed for GQA from the start sidesteps the conversion entirely. These are refinements at the edges of a method that, for the central case, simply works — which is why nearly every model trained since adopts it.