Multi-head attention
Many small attention patterns side by side, and how sharing their keys and values makes generation cheaper.
Read first:Self-attention
Step 1 of 6· Illustrative attention patterns and vectors; the head splits, sharing layouts and cache arithmetic are real
One attention pattern can only blend a word's context in one way. Multi-head attention runs several smaller attentions side by side over the same words, each with its own learned view, so one head can follow the previous word while another tracks what “it” refers to. Their results are joined and mixed into a single update. While generating, the model keeps every head's keys and values in memory, the KV cache, and that memory adds up quickly. Newer designs let groups of heads share one set of keys and values, which shrinks the cache several times over for a small loss in quality.
Why it matters for your product
During generation, every cached key and value is read again for each new token, and that memory traffic, more than arithmetic, often sets the speed. The cache also decides how many conversations fit on one GPU and how long each can be. Grouped-query attention cuts it by the number of heads sharing each key–value head, which is why models such as Llama 2 70B use it. When choosing a model to self-host, check its key–value head count, not just its parameter count.
For engineersShow the maths
MultiHead(X) = Concat(head₁, …, headₕ) · Wᴼ, headᵢ = Attention(X·Wᵢᑫ, X·Wᵢᴷ, X·Wᵢⱽ)
Run h attentions in parallel, each on its own narrow projections of the same input, then join their outputs and mix them with one more matrix.
Worked example: With a model width of 512 and h = 8, each head works in 512 ÷ 8 = 64 dimensions, and the eight 64-wide outputs join back into 512.
KV cache = 2 × layers × G × d_head × tokens × bytes per number
Two tensors (keys and values) per layer, one per key–value head G, for every cached token. G equals the number of query heads for multi-head attention and 1 for multi-query.
Worked example: 32 layers, heads of 128, 8,192 tokens and 16-bit numbers: 32 key–value heads need about 4.3 GB per sequence; 8 need about 1.1 GB, a 4× saving.
Heads do not always specialise cleanly: many learn overlapping or hard-to-name patterns, and some can be removed with little effect. Sharing keys and values trades some quality for memory, and multi-query attention gives up more than grouped-query attention does.
The words you will hear
Attention head
Primary sources
- 01Attention Is All You NeedVaswani et al. · 2017
- 02Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be PrunedVoita et al. · 2019
- 03Fast Transformer Decoding: One Write-Head is All You NeedShazeer · 2019
- 04GQA: Training Generalized Multi-Query Transformer Models from Multi-Head CheckpointsAinslie et al. · 2023
- 05Llama 2: Open Foundation and Fine-Tuned Chat ModelsTouvron et al. · 2023
Connected ideas
Generating text
KV cache
Why a model keeps the keys and values of earlier tokens instead of recomputing them, and what that memory costs.
ExploreInside the transformer
Self-attention
How every word looks at every other word to work out what it means here.
ExploreGenerating text
Prefill and decode
Why the first word of an answer takes a moment and the rest stream out at a steady pace.
ExploreSpeed and cost
FlashAttention
The same attention maths, done in tiles so the big score table never touches slow memory.
ExploreWant this working on your data?
We design and build the systems these ideas power: retrieval, agents, voice and the models behind them. Start with a free discovery call.