Grouped-query attention keeps separate query heads while letting groups of them share key and value heads. Its main practical attraction is a smaller KV cache and less cached data to read during decoding.
Read the diagram in head counts
The card compares eight query heads with two different arrangements. In ordinary multi-head attention, each query head has its own key and value head: eight of each. In the illustrated GQA arrangement, each pair of query heads shares one key and value head: eight query heads, four key heads, and four value heads.
Holding sequence length, head dimension, and data type fixed, that halves the key and value storage. Each query still computes its own attention weights. Sharing keys and values does not mean that both queries must attend to the same tokens with the same weights.
MHA, GQA, and MQA are points on one axis
- Multi-head attention (MHA): as many KV heads as query heads.
- Grouped-query attention (GQA): several query heads share each KV head.
- Multi-query attention (MQA): all query heads share one KV head.
For a second example, reducing 32 KV heads to eight gives a fourfold reduction in the KV-head-dependent cache storage. It does not give a fourfold reduction in model weights, total GPU memory, or end-to-end latency. Those include other costs.
What is the trade-off?
Fewer KV heads limit the distinct key and value representations available to queries. How this affects quality depends on the model, training procedure, and task. The GQA paper demonstrated configurations approaching multi-head quality while achieving decoding speed closer to multi-query attention. That result is evidence for a useful design choice, not a guarantee for every head count.
GQA normally belongs to the model’s architecture and training. Treating an arbitrary trained multi-head model as GQA by discarding heads is not an equivalent change. The paper also studies a procedure for converting and further training existing checkpoints.
Source: Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. Use the KV cache memory example to see where head count enters the budget.