A head owns its own Q, K, V projections of width d_head, runs its own softmax, and produces a d_head-dimensional context vector per token.
Think of a panel of judges scoring an audition. Each judge has a private notebook for what they care about, a private notebook for what each contestant offers, and a private notebook for the contestant's actual material. Every judge scores the contestants independently using only their own three notebooks, never peeking at another judge's notes. After scoring, each judge writes down a short personal verdict. The verdicts are then stapled together and handed to a single producer who reads the whole stack and writes the final result on the show. Each judge is one independent scorer in the layer; the producer at the end is the one who mixes the verdicts back into the show.
Concept explanation~2 min read
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
Concept explanation~2 min read
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
Multi-head attention is one of the most repeated terms in transformer literature and one of the most under-pictured. People know it has heads but cannot always say what one head physically owns or produces.
This deep dive defines a single attention head from the perspective of an implementation, walks through the dimensions and operations, contrasts standard MHA with MQA and GQA, and explains the operational consequences of the design.
Mental model: one head is a small, independent attention computation owning its own (d_model, d_head) projections; n_heads of them run in parallel; their outputs are concatenated and mixed exactly once by W_O.
What one head owns
The three projection matrices
A single attention head owns three matrices, one each for Q, K, and V projection:
- W_Q of shape
(d_model, d_head) - W_K of shape
(d_model, d_head) - W_V of shape
(d_model, d_head)
For a model with d_model = 4096 and n_heads = 32, you have d_head = d_model / n_heads = 128. So each per-head matrix is (4096, 128). With 32 heads, the per-layer Q parameter count is 32 x 4096 x 128 = 16M, exactly the same as a single full-width (4096, 4096) projection.
Why the split costs nothing
The key arithmetic: n_heads * d_head = d_model. The total parameter count of a multi-head Q projection equals the parameter count of a single d_model x d_model projection. Multi-head buys you specialization (different heads learn different routing patterns) at the same compute and memory cost as single-head.
Implementation packing
In real code, per-head matrices are almost always packed into a single tensor of shape (d_model, d_model) and reshaped to (d_model, n_heads, d_head) at compute time. Mathematically identical, just easier to fuse into a single matmul.
Situations where this technique stops working.
2–4 min · Everything important, quickly.
Real products, models, and research that use this idea.
- GPT-3 175B: d_model = 12288, n_heads = 96, d_head = 128. Standard multi-head, every head owns its own Q, K, V.
- Llama-3-70B: d_model = 8192, 64 Q heads, 8 KV heads via GQA, d_head = 128.
What an interviewer would ask next. Try answering before peeking at the approach.
QIf d_head is fixed at 128, what changes when you scale up d_model?
n_heads grows linearly with d_model so d_head stays at 128. This is the convention in Llama, Mistral, and Qwen because head dimension 128 plays well with tensor core shapes and existing FlashAttention kernel tiling.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Saying heads share Q, K, and V projections. They do not in standard multi-head attention. Each head has its own (d_model x d_head) projection matrices.
60 second bullets to scan on the way to the call.
Shape of W_Q, W_K, W_V for a single head given d_model and n_heads
Formula for d_head in terms of d_model and n_heads
Primary sources. Browse if you want the original framing.
Same topic, related formats. Practice these next.