After the Q, K, V projections produce d_model-wide vectors, a reshape + transpose splits each vector into (num_heads, d_head), a tensor view, not a copy.
Picture a long candy bar. You can cut it into 8 equal pieces or call it one bar with 8 labeled sections; the candy is the same either way. Multi-head attention does the second thing. The token vector is one long buffer in memory, and the model just relabels contiguous chunks of it as 'this part is head 1, this part is head 2', and so on. Nothing is copied or duplicated. Each head then does its own small attention computation on its slice, and the model concatenates the slices back together at the end.
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 those topics where the diagrams in textbooks suggest a much more complicated implementation than the actual code does. A first reading of the original transformer paper leaves people imagining num_heads parallel projection matrices, each producing a small per-head Q, K, V. The reality is one big projection followed by a free reshape. Knowing this is the difference between thinking the architecture is wasteful (it is not) and being able to debug shape errors quickly (which is most of the work in practice).
This deep dive walks the actual mechanics step by step, explains why the design is mathematically equivalent to per-head projections, surveys the GQA/MQA variations that change the K and V shapes, and closes with the failure modes that show up most often in production code.
The end to end shape transformation
Start with an input tensor x of shape (batch, seq, d_model). Track the shape at every step.
Step 1: project to Q, K, V
Q = x @ W_Q # (batch, seq, d_model)
K = x @ W_K # (batch, seq, d_model)
V = x @ W_V # (batch, seq, d_model)
Each W_* is a (d_model, d_model) learned matrix. Output is still d_model wide. No head structure yet.
Step 2: reshape to expose the head axis
Q = Q.view(batch, seq, num_heads, d_head)
This is a view, not a copy. PyTorch reinterprets the contiguous (d_model,) axis as two axes (num_heads, d_head) by adjusting strides. The data does not move. The constraint is d_model = num_heads * d_head, which is why you sometimes see configs like d_model=4096, num_heads=32, d_head=128.
Step 3: transpose to make per-head sequences contiguous
Q = Q.transpose(1, 2) # (batch, num_heads, seq, d_head)
Now each head's sequence is a contiguous (seq, d_head) block. This may trigger an actual reorder in memory, but it is a one-time cost and the result is the layout every attention kernel expects.
Step 4: per-head scaled dot-product
scores = Q @ K.transpose(-2, -1) / sqrt(d_head)
A = softmax(scores, dim=-1)
out = A @ V # (batch, num_heads, seq, d_head)
Each head operates independently. The batch + num_heads axes are just outer dimensions for the matmul; the actual attention happens over the (seq, d_head) per-head slice.
Step 5: concat and project out
out = out.transpose(1, 2).contiguous().view(batch, seq, d_model)
out = out @ W_O # (batch, seq, d_model)
The transpose and view is the inverse of step 2-3. W_O is (d_model, d_model) and learns how to combine the per-head outputs.
The whole pipeline is six matrix shapes and two reshapes. Memorize this and most multi-head debugging becomes trivial.
Situations where this technique stops working.
2–4 min · Everything important, quickly.
Real products, models, and research that use this idea.
- PyTorch's `nn.MultiheadAttention` uses exactly this view and transpose pattern internally.
- Llama 4 Maverick uses grouped-query attention: Q stays full multi-head while K and V share groups, but the reshape mechanism is identical.
What an interviewer would ask next. Try answering before peeking at the approach.
QIf multi-head attention is mathematically equivalent to one big head with the same total dimension, why use multiple heads at all?
It is NOT equivalent because of the per-head softmax. With one big head, softmax normalizes over the full d_model; with multiple heads, each head's softmax normalizes only over its d_head slice, allowing different heads to attend to different patterns without competing for the same probability budget. The output projection W_O then learns how to combine them.
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.
Believing each head has its own private d_model-wide projection. The projections produce one d_model vector that is then sliced into per-head pieces.
60 second bullets to scan on the way to the call.
Shape transformation: (d_model) to (num_heads, d_head) via reshape
Why the reshape is a view, not a copy
Primary sources. Browse if you want the original framing.
Same topic, related formats. Practice these next.