Lesson 2. Multi-Head Attention

Mission link: Lesson 1's equation is the operation; this lesson is how the transformer actually uses it, running it several times in parallel on different learned projections of the same input, rather than once at full dimension.
Primary source: Paper: "Attention Is All You Need", Vaswani et al., 2017
Prerequisites: Lesson 1, Scaled dot-product attention

Warm-up

▢ Write the scaled dot-product attention equation from memory.

Check

Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V.

▢ Why does the / sqrt(d_k) scaling matter as d_k grows?

Check

The variance of the raw dot-product scores grows with d_k, which pushes softmax toward a near-one-hot output where its gradient is nearly flat. Dividing by sqrt(d_k) counteracts that growth, keeping scores in a range where softmax's gradient is still usable.

Know this

One attention operation gives one view

Scaled dot-product attention, computed once over the full input, produces exactly one weighted blend per query, one "view" of how each position relates to every other position. That's a real limitation: a sentence's structure carries multiple, different kinds of relationships at once (which word a pronoun refers to, which words are syntactically linked, which words are simply nearby), and forcing all of that into a single attention pattern per position means averaging relationships together that might be better kept separate.

Project into several smaller subspaces, attend in each separately

Multi-head attention fixes this by running attention several times in parallel, each time on its own, separately learned projection of the input. For h heads, each head i gets its own learned weight matrices W_Q^i, W_K^i, W_V^i, projecting the input down to a smaller dimension d_k = d_model / h:

Q_i = X W_Q^i,   K_i = X W_K^i,   V_i = X W_V^i
head_i = Attention(Q_i, K_i, V_i)

Each head applies lesson 1's exact equation, just inside its own lower-dimensional subspace. Because each head's projection matrices are learned independently, different heads are free to specialize, one might end up tracking nearby-word relationships, another something more syntactic, without anything forcing them to agree on a single shared attention pattern.

Concatenate the heads, then project once more

The heads' outputs are concatenated back together and passed through one final learned projection, W_O, back to the model's full dimension:

MultiHead(Q, K, V) = Concat(head_1, ..., head_h) W_O

The final projection isn't optional decoration: concatenation alone just places each head's output side by side in a bigger vector, with no way for information from one head to combine with another. W_O is what lets the model learn how to blend the different heads' separate views back into a single, unified representation for whatever comes next in the network.

flowchart TD X["input X"] --> P1["W_Q^1, W_K^1, W_V^1"] X --> P2["W_Q^2, W_K^2, W_V^2"] X --> Pn["W_Q^h, W_K^h, W_V^h"] P1 --> H1["head_1 = Attention(Q_1, K_1, V_1)"] P2 --> H2["head_2 = Attention(Q_2, K_2, V_2)"] Pn --> Hn["head_h = Attention(Q_h, K_h, V_h)"] H1 --> C["Concat(head_1, ..., head_h)"] H2 --> C Hn --> C C --> O["W_O projection"] O --> Out["MultiHead(Q, K, V)"]

The per-head dimension is what the scaling formula actually uses

Since each head operates at d_k = d_model / h, not the full d_model, lesson 1's / sqrt(d_k) scaling uses that smaller, per-head dimension. This matters because it's easy to assume the scaling should track the model's overall size; it doesn't, it tracks whatever dimension the dot products inside that specific attention computation are actually summed over, which is the per-head dimension once attention is split across heads.

Practice

▢ A model has d_model = 512 and uses 8 attention heads. What is d_k, the dimension each head operates in?

Check

d_k = 512 / 8 = 64.

▢ Why is a final linear projection (W_O) applied after concatenating the heads' outputs, rather than just using the concatenated vector directly?

Check

Concatenation alone places each head's output side by side with no interaction between them. W_O lets the model learn how to combine information across the separate heads into one unified representation, rather than leaving the heads' outputs as disconnected chunks of a larger vector.

▢ Why can different heads end up specializing in different kinds of relationships between positions, when all of them apply the exact same attention equation?

Check

Each head has its own, independently learned projection matrices (W_Q^i, W_K^i, W_V^i). Two heads applying the identical equation to differently projected versions of the same input can attend to entirely different aspects of the input, since the projections themselves shape what each head's queries and keys actually represent.

▢ A model splits attention across h heads, each operating at d_k = d_model / h. When computing the / sqrt(d_k) scaling inside one head's attention computation, should d_k be the per-head dimension or the full d_model? Why?

Hint

Ask what dimension the dot products being scaled are actually summed over.

Check

The per-head dimension. Lesson 1's scaling exists to counteract the variance growth of a dot product summed over d_k terms; inside one head, that dot product is summed over the head's own dimension, d_model / h, not the full model dimension. Using the wrong (full) dimension would under-scale the scores relative to what that head's own dot products actually need.

▢ Which claim is true of multi-head attention compared to computing attention once at the full model dimension?
- a) Multi-head attention runs attention once, then artificially splits the single output into pieces afterward
- b) Each head has its own learned projections and computes attention independently in a smaller subspace, before all heads are concatenated and projected back
- c) All heads share the same query, key, and value projection matrices, differing only in which positions they attend to
- d) The final projection W_O is optional and can be skipped without changing what the model can represent

Check

b) That's exactly the mechanism this lesson describes. (a) is false: each head computes its own full attention operation independently, on its own projected inputs, not a post-hoc split of one shared computation. (c) is false: each head's projection matrices are learned independently, which is what lets heads specialize. (d) is false: without W_O, concatenated heads have no way to combine information across each other.

Real-world reps

[ ] Implement multi-head attention from raw tensor operations: per-head linear projections, lesson 1's attention equation applied independently per head, concatenation, and the final W_O projection. Confirm the output shape matches d_model regardless of how many heads you choose.
[ ] For a small d_model and two different head counts (say 2 and 8), compute d_k for each and note how it shrinks as head count grows.
[ ] Tomorrow: read the "Multi-Head Attention" section of the primary source paper in full, and check that every matrix in its equations matches something you named in your own implementation.

Going further

Not landing? Reread the primary source at the top, since this lesson compresses it and compression is where understanding leaks. Check the glossary for any term that felt slippery.

If the lesson itself is unclear rather than the material, that is a defect: open an issue.