One Head Isn't Enough

Last experiment we built one attention head: softmax(QKᵀ/√d_k)·V, and watched “bank” reach across a sentence and become river-bank. That single head felt like the whole story. It isn’t. This experiment attacks a two-part hypothesis I walked in believing:

“One head is enough, and a word may look anywhere in the sentence it likes.”

Both halves are wrong, and they fail for opposite reasons — one is a capacity problem, the other a cheating problem. We fix the first with multi-head attention, the second with a causal mask, and along the way meet cross-attention and one of the nastiest bugs in the whole architecture. Verdict up front: FAR FROM IT.

A single head produces exactly one attention distribution — one set of weights, one weighted average of the values. But look at what “bank” actually needs in “the river bank was steep”: it has to attend to river to fix its meaning, and to was to satisfy subject–verb grammar, at the same time. One average can’t peak in two places. Force it and you get a blurry compromise that’s good at neither.

↑ Drag the heads dial. At 1 head, one tangled fan of arcs tries to serve all three relations at once — mush. Add heads and each relation peels off into its own clean specialist. Click a head chip to isolate it.

The fix is almost embarrassingly simple: don’t make one big head, make h small ones. Split the model width d_model into h slices of d_model/h, give each its own W_Q, W_K, W_V, run attention in parallel, then concatenate the outputs and mix them with one more matrix W_O:

$\text{MultiHead}(Q,K,V) = \text{Concat}(\text{head}_1, \dots, \text{head}_h)\,W_O$

$\text{head}_i = \text{Attention}(QW_Q^i,\; KW_K^i,\; VW_V^i)$

Note the accounting: we didn’t add width, we sliced it. Eight heads over a 512-dim model are eight 64-dim attentions, not eight 512-dim ones. Same parameter budget — spent on specialists instead of one generalist.

So far every word has queried its own sentence — self-attention. But think about translation: to produce the French word “rivière”, the decoder needs to look back at the English “river” it came from. The query lives in one sequence, the keys and values in another. That’s cross-attention — and here’s the punchline: it’s the exact same operation. The only thing that changes is where K and V come from.

↑ Press ▶ generate and watch the decoder write French one token at a time. In cross mode each new word reaches up into the English source it came from; toggle to self and it attends only to the French already written. One formula, two wirings.

Self-attention lets a sequence enrich itself; cross-attention lets one sequence read another. Encoders use self-attention to understand the input; a translation decoder uses both — self-attention over what it’s written so far, then cross-attention into the source. Same softmax(QKᵀ/√d_k)·V, different plumbing.

Now the second half of the hypothesis — that a word may look anywhere. For a model generating text one token at a time, that’s fatal. If word 3 is allowed to attend to word 5 during training, it can just copy the answer it’s supposed to predict. At inference word 5 doesn’t exist yet, so the model face-plants. A generator must only ever look backward.

We enforce it by adding a mask M to the scores before softmax, where every future position (column j > row i) is set to −∞:

$\text{Attention} = \text{softmax}\!\left(\frac{QK^{\top}}{\sqrt{d_k}} + M\right)V, \qquad M_{ij} = \begin{cases} 0 & j \le i \\ -\infty & j > i \end{cases}$

Why −∞ and not just 0? This is the trap. Softmax exponentiates: e^0 = 1, so a “blocked” score of 0 still hands the future token a full unit of weight. Only e^−∞ = 0 makes it truly vanish. Zeroing the score looks masked and isn’t.

↑ Toggle the three modes. “no mask” leaks ~40% of attention into the future. ”−∞” drives future leakage to a true 0%. “×0” looks masked but still leaks — because softmax(0)=1.

Here’s the mock question that caught me. “Your language model trains beautifully — train and validation loss both drop and stay low — but generated text is garbage. What happened?” My first instinct was overfitting. Wrong, and the giveaway is that validation loss looks fine too.

So you can’t debug this with a metric — you debug it structurally. Print the attention matrix (or its gradients) and look at the strict upper triangle: it must be exactly zero. Any non-zero weight there — the classic cause being someone wrote scores.masked_fill(mask, 0) instead of −inf — is the leak. It’s a bug you find by staring at the triangle, not at the curve.

• ①Both train and val loss look healthy → rules out ordinary overfitting (which spreads the two apart). Healthy-everything + broken generation is the fingerprint of train/inference mismatch.
• ②The diagnostic is a heatmap, not a number: strict-upper-triangle attention (or gradient) must be 0. Non-zero = leakage.

Three placements of one operation, distinguished only by where the query and key/value sequences live and whether the future is masked:

And each of those runs multi-head underneath — a stack of parallel specialists, concatenated and mixed by W_O. Capacity from the heads; honesty from the mask.

Time to place everything on the original blueprint. This is Figure 1 of Attention Is All You Need (Vaswani et al., 2017) — the whole machine. The three attention blocks we’ve now built by hand are lit in coral; hover any block to see its job, and click a lit one to jump to the experiment that derives it.

↑ Three attention blocks, one operation. The only dashed block left in our path is positional encoding — the subject of the next experiment.

One thing we’ve quietly assumed the whole time: that attention knows the order of the words. It doesn’t. A dot product doesn’t care whether “river” came before or after “bank” — shuffle the sentence and self-attention returns the same answer, just permuted. It’s permutation-invariant. So how does a transformer tell “dog bites man” from “man bites dog”? Next experiment: positional encoding — how order gets injected back in (EXP-015).