II · THE IDEA · ARTIFICIAL INTELLIGENCE
Multi-Query and Grouped-Query Attention
▶ Listen · narrated
A long conversation or a large document can fill memory not with the text itself but with the hidden states each token leaves behind for every attention head to consult later.
At a glance
- What it changes
- Number of separate key and value projections per layer
- Multi-query attention
- One key projection and one value projection shared by all query heads
- Grouped-query attention
- Query heads divided into groups; each group shares one key and one value projection
- Memory saving
- Proportional to the reduction in KV heads — up to 10–20× for multi-query in large models
Imagine a large meeting where everyone needs to refer back to everything said earlier. In the standard arrangement, each participant keeps their own complete transcript, which means 32 transcripts for 32 participants. Multi-query attention is like having one shared transcript that everyone consults. Grouped-query attention is a compromise: divide the 32 participants into, say, 8 groups, and each group shares one transcript. You now have 8 transcripts instead of 32, which takes up much less space, but each group still has its own record so they can focus on different aspects of the conversation. The trade-off is that people in the same group have to work from the same notes, which loses some nuance, but in practice the loss is small and the space saved is large.
In standard multi-head attention, each head has its own learned linear projections for queries, keys and values. During autoregressive generation, the keys and values for all previous tokens must be cached so each new token can attend to them without recomputing the projections. For a layer with h heads, hidden dimension d, and sequence length n, the KV cache stores 2 · h · n · (d/h) = 2 · n · d values per layer. Multi-query attention reduces this to 2 · n · (d/h) by using a single key projection and a single value projection shared across all query heads. Each query head still computes separate attention weights, but all heads attend over the same keys and retrieve the same values, which are then projected separately per head after the attention operation. Grouped-query attention uses g key-value heads where 1 < g < h, dividing the query heads into g groups. Each group shares one key-value head, so the cache stores 2 · g · n · (d/h) values per layer. The Ainslie et al. paper found that g = h/4 or g = h/8 recovered most of the quality of full multi-head attention while retaining most of the memory saving of multi-query attention. The uptraining procedure they describe takes a checkpoint with h key-value heads, averages them into g heads by mean-pooling within each group, then continues training for approximately 5% of the original pretraining tokens. The resulting model's KV cache is g/h times the size of the original, and generation throughput increases roughly in proportion because memory bandwidth during decoding is often the bottleneck.
Look closer
The cache grows with every token generated
During generation, each new token attends to all previous tokens in the sequence. To do that efficiently, the model stores the key and value vectors for every past token rather than recomputing them. In standard multi-head attention with, say, 32 heads, each token adds 32 key vectors and 32 value vectors to this cache. Over a long conversation or document, the cache can exceed the memory required to hold the model's own weights. Multi-query attention stores one key and one value per token instead of one per head, shrinking the cache by a factor equal to the head count.
Grouped-query sits between multi-head and multi-query
Multi-query attention uses a single key-value head for all queries, which saves the most memory but can degrade quality on some tasks. Grouped-query attention divides the query heads into groups — often four or eight — and gives each group its own key-value head. A model with 32 query heads and 4 KV heads stores a quarter as much cache as the standard arrangement, while keeping more representational capacity than a strict multi-query design. The GQA paper showed this compromise recovers most of the quality loss while retaining most of the memory saving.
You can convert an existing checkpoint without full retraining
The GQA paper describes a method called uptraining: take a model trained with standard multi-head attention, average its key and value projection weights into the smaller number of grouped heads, then continue training for a small fraction of the original compute budget. The result is a model with a much smaller KV cache that performs nearly as well as the original. This matters because training a large model from scratch is expensive, but adapting one that already exists is comparatively cheap.
The story
Standard transformer attention computes three projections for every head: a query, a key and a value. If a layer has 32 heads, it maintains 32 separate sets of these projections, and during generation each token leaves behind 32 key vectors and 32 value vectors that must stay in memory so future tokens can attend to them. This cache — the KV cache — grows linearly with sequence length and linearly with the number of heads, and in a long conversation or a large document it can become the dominant memory cost, dwarfing the model weights themselves.
Noam Shazeer's 2019 paper proposed multi-query attention: use many query heads as usual, but only one key projection and one value projection, shared across all of them. Each query head still computes its own attention scores, but every head is scoring and weighting the same set of values. The KV cache shrinks by a factor equal to the head count — if you had 32 heads, you now store one thirty-second the keys and values.
The cost is a modest drop in quality. Shazeer reported that translation models with multi-query attention performed slightly worse than their multi-head equivalents, but the gap was small enough that the memory and speed gains often justified it, especially for deployment where serving costs matter more than the last fraction of a percentage point on a benchmark.
Grouped-query attention, introduced by Ainslie and colleagues in 2023, offers a middle path. Instead of one shared key-value head, divide the query heads into groups and give each group its own key-value head. A model with 32 query heads might use 8 key-value heads, storing an eighth the cache of the standard design while preserving more representational diversity than pure multi-query attention. The paper demonstrated that this arrangement closes much of the quality gap while keeping most of the memory benefit.
The GQA paper also showed that you need not train from scratch. Their uptraining procedure takes a checkpoint trained with standard multi-head attention, averages the key and value weights to produce the smaller number of grouped heads, then continues training for a few per cent of the original token count. Models converted this way recovered nearly all their original performance. This matters because it makes the technique retroactive: a model that has already consumed millions of dollars of compute can be adapted rather than discarded.
The technique is now widespread. Many recent open-weight models use grouped-query attention by default, typically with four to eight key-value heads regardless of how many query heads they carry. The memory saving is most visible in long-context scenarios — a 100,000-token conversation with standard 32-head attention might require 40 gigabytes just for the KV cache, while the same conversation with 4 key-value heads needs 5 gigabytes. That difference determines whether the conversation fits in consumer hardware or requires a data centre.
Why it mattered then
Shazeer's 2019 paper was motivated by inference cost. Training a model is expensive but happens once; serving it to users happens millions of times, and every byte of memory and every millisecond of latency translates directly into hardware bills and user-facing lag. Multi-query attention addressed the KV cache because the cache was becoming the bottleneck: as models grew larger and context windows grew longer, the memory required to store past keys and values during generation was outpacing the memory required to hold the model itself. Shazeer's experiments on translation models showed that the quality cost was small and the speed and memory gains were large, making it a pragmatic trade for production systems where throughput mattered more than the last decimal place of BLEU score. The technique was adopted quietly in some commercial systems, but it did not become widespread in open research until grouped-query attention offered a more palatable compromise.
Why it matters now
Grouped-query attention is now a standard architectural choice in newly trained large models, particularly those designed for long contexts. The 2023 GQA paper made the technique practical by demonstrating that existing checkpoints could be converted cheaply and that a moderate number of key-value heads recovered most of the quality lost in pure multi-query designs. As context windows have grown from a few thousand tokens to hundreds of thousands, the KV cache has become the dominant memory cost during generation, and grouped-query attention directly addresses it without requiring algorithmic breakthroughs or exotic hardware. The technique also matters for on-device deployment: a model with a smaller KV cache can run longer conversations on a phone or laptop without swapping to disk or refusing the request. It is a rare example of a simple architectural change that improves both the economics of serving models at scale and the feasibility of running them locally.
The surprising detail
The uptraining procedure in the GQA paper works by averaging. If you have 32 key-value heads and want 8, you average heads 0–3 into the new head 0, heads 4–7 into the new head 1, and so on. This is a lossy compression: you are throwing away information that took weeks of training to learn. Yet after continuing training for just 5 per cent of the original token budget, the converted model performs nearly as well as the original. The result suggests that much of what those extra key-value heads learned was redundant, and that the model can relearn the non-redundant parts quickly once it adapts to the new architecture. It also means that the choice of how many key-value heads to use need not be made at the start of training, which decouples architectural decisions from the largest and least reversible part of the cost.
Remember this
Fewer key-value heads mean a smaller cache, and the cache is often the memory bottleneck during generation. The quality cost is real but small.
Test yourself
A model has 40 query heads per layer and uses grouped-query attention with 5 key-value heads. You want to halve the KV cache memory without retraining. Can you do it by reducing the number of query heads instead of the number of key-value heads, and why?
No. The KV cache stores keys and values, and the number of those is determined by the number of key-value heads, not the number of query heads. Halving the query heads would halve the computation during the attention operation itself, but it would not change the size of the cache at all. To halve the cache, you would need to reduce the key-value heads from 5 to 2 or 3, which would require retraining or uptraining because the existing queries are divided into groups that expect 5 key-value heads. The asymmetry is the point: you can have many queries sharing fewer keys and values, which is what makes the memory saving possible.
Go deeper
- Fast Transformer Decoding: One Write-Head is All You Need · arXiv · Noam Shazeer et al. · 2019-11-06
- GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints · arXiv · Joshua Ainslie et al. · 2023-05-22
Image: Original diagram, The Daily Triptych. Licence: Original work. Source.