Masks and Heads

Step 4 of 7. In: \(Q\), \(K\) and \(V\) from step 3, \(5 \times 4\) each. Out: the attention output \(a\), also \(5 \times 4\), which the block adds back to \(X\). ← Query, key, value · next step: The block →

Masks: causal, padding, and everything else

Sometimes a token must not look at another: the future, in a model that generates text; the <pad> filler, in a batch. The mask is added to the scores before the softmax, as \(-\infty\) (in practice, a large negative number), and since \(e^{-\infty} = 0\) the forbidden weights become exactly zero:

Turn on causal mask in the matrix panel of the previous page and the upper triangle goes dark. This one line is the entire difference between an encoder and a decoder, and the reason a decoder can be trained on all positions in parallel: with the mask, position \(i\)'s prediction cannot see its own answer, so one forward pass gives you \(n\) training examples instead of one. For the running sentence with its answer, "the cat climbed the roof", in tokens:

Five examples, a single pass through the network. Without the mask, position 1 would see ␣cat right next to it and learn to copy instead of predict.

The mask in the running example

On the previous page head 1 ran without a mask. With it, every score above the diagonal becomes \(-\infty\) before the softmax, and each row shares its weight among the tokens that are left. Row 2, ␣cat, had the scores \([-0.333, 0.471, -0.218, -0.577, -0.471]\); the mask keeps the first two, \(e^{-0.333} = 0.717\) and \(e^{0.471} = 1.602\), which add up to 2.319, so the weights become \(0.717/2.319 = 0.309\) and \(1.602/2.319 = 0.691\):

The first token can only look at itself. The last row does not move: it is the row that predicts ␣roof, and it was already allowed to see everything before it.

Multi-head, and what it turned into

One head produces one distribution per position — one relationship. But a token usually needs several at once: a verb wants to find its subject and look at the previous word; a pronoun wants the noun it refers to and the punctuation that closes the sentence. A single row of weights that sums to 1 cannot point strongly at all of that. Multi-head attention runs \(h\) heads in parallel, each with its own \(W_Q^i, W_K^i, W_V^i\) on \(d/h\)-dimensional slices, and concatenates the outputs:

At \(d = 4096\) and \(h = 32\), each head works with 128-dimensional vectors; the 32 concatenated outputs are 4096 wide again, and \(W^O\) mixes what each head found. The head options in the matrix panel are exactly this kind of specialization.

In code the heads are not separate layers. One \(d \times d\) projection computes the queries of every head at once, and its output is cut into \(h\) slices of width \(d/h\); the same for K and V. For a batch of \(B\) sequences:

The parameters are four \(d \times d\) matrices, \(W_Q\), \(W_K\), \(W_V\) and \(W^O\) (plus biases, when the model uses them), whatever \(h\) is.

The same thing as the lab writes it — split_heads, the attention function with its mask, the concatenation and \(W^O\) — line by line, on three tokens with \(d = 4\) and \(h = 2\). Head 1 is the bank example of the previous page (the row of "bank" ends at \([0.85,\ 0.15]\) again); turn the causal mask on, or set \(h = 1\), and watch which lines move:

Note what it does not cost: with \(d/h\) per head, \(h\) heads cost the same FLOPs as one head of full width. Heads are free parallelism, not extra capacity — which is why \(h\) grew and grew.

Then inference happened, and the design changed: the cost of attention is why.

Two heads in the running example

The toy has \(d = 4\) and \(h = 2\): head 1 owns columns 1–2 of \(Q\), \(K\) and \(V\) (the head of the previous page), head 2 owns columns 3–4. For the last token head 2 has the same query as head 1, \(q = [u_4, u_1] = [1.732, -0.577]\), but other keys, \(k = [u_1, u_4]\):

Head 1 found the verb (␣climb and ed, 0.94); head 2 found the subject, ␣cat, 0.88. Its values are columns 3–4 of \(u\), so its output for the last token is \(o^{(2)} = [-0.521, -0.520]\). The two outputs go side by side, and \(W^O\) — the identity in the toy, to keep the sums short — gives the attention output of the last token:

and, for the five tokens with the causal mask,

That matrix is what the attention sublayer hands to the block, and where the Transformers chapter picks the example up. With a single head of width 4 instead of two of width 2, the last row would split between ␣cat (0.39) and ␣climb (0.48): one row of weights that adds up to 1 cannot point strongly at both — the case for multi-head, in numbers.

← Query, key, value · next step: The block → · side page: The cost of attention