Chapter 16 Masks and Multi-Head Attention

Suppose the input is red key . and the output above red must predict key. Chapter 15 let that position read the representation of key itself. Training could reward a shortcut that will be unavailable when generating a new token. Putting the data into a different validation file does not close that shortcut if validation uses the same illegal visibility.

We need to specify which information exists at each prediction, then enforce that specification before mixing values. Once visibility is correct, a second limitation appears: one mixture gives all value coordinates the same set of position weights. Several separately projected mixtures can offer more flexibility. We will calculate both changes, keeping the position and row-vector conventions from Chapter 15.

16.1 Specify visibility before choosing a mask representation

Example 16.1 — Restrict the same three-by-three score matrix. Reuse Chapter 15’s X, Q, K, V, and scaled scores, with $s=1/\sqrt{2}$:

$$S=\begin{bmatrix}s&0&s\\s&s&2s\\2s&s&3s\end{bmatrix}.$$

The row still denotes a receiving query, and the column a contributing key. For next-token prediction, position 1 knows its current input but no future input; position 2 knows inputs 1–2; position 3 knows 1–3. This is causal visibility. The current input is allowed because the target is the next token. For [BOS,red,key,.], the aligned targets are [red,key,.,EOS]. The target is shifted by one; we do not remove the diagonal to compensate for an unshifted target.

A batch may also contain shorter records. If the third slot is padding, it does not represent a real source token and must not contribute to a real query. This is padding visibility. Table 16.1 states the relation without depending on an API’s choice of 0, 1, True, or False.

Table 16.1 — Allowed key positions for each query. The padding example has two real inputs and one trailing placeholder.

Query slot	Causal only	Padding keys only	Both restrictions	Query participates in loss?
1	1	1, 2	1	Yes, if it has a scored target
2	1, 2	1, 2	1, 2	Yes, if it has a scored target
3	1, 2, 3	1, 2	1, 2	No: this query is padding

The last column applies to the padded example. Key padding controls what a query may read. Query padding controls whether its output is used. The loss mask controls which target events are scored. They are related but not interchangeable. Merely removing padding from the loss does not prevent real queries from reading padded keys in a bidirectional calculation.

16.2 Remove forbidden contributors before normalization

A mask is the representation of that relation in the calculation. In this book’s functions, visible=True means allowed. An additive version adds zero to allowed scores and negative infinity to forbidden ones. Negative infinity is a mathematical way to make its exponential zero; it is not an extremely unfavorable but still possible contributor.

For the causal case, the masked score matrix and resulting probabilities are:

$$S_{\text{causal}}=\begin{bmatrix}s&-\infty&-\infty\\s&s&-\infty\\2s&s&3s\end{bmatrix},\qquad A_{\text{causal}}\approx\begin{bmatrix}1&0&0\\0.5&0.5&0\\0.283995&0.140029&0.575975\end{bmatrix}.$$

The first row has only one exponential left and therefore assigns it probability 1. The second row has two identical exponentials and divides each by their sum: one half each. The third row is unchanged from Chapter 15. Multiplying the unchanged V gives:

$$O_{\text{causal}}\approx\begin{bmatrix}1&0\\1&1\\1.575975&1.432009\end{bmatrix}.$$

The second output is exactly $0.5[1,0]+0.5[1,2]=[1,1]$. Its previous unmasked value [1.503490,1.503490] used the forbidden third position. Multiplying probabilities by a 0/1 mask after Softmax and stopping there would leave a row whose weights no longer sum to one. Mask first, then recompute the denominator over allowed keys.

Query 1 → next token

1. An operationKey 1: allowed; 2: forbidden; 3: forbidden
2. A list of numbersWeights [1, 0, 0] → output [1, 0]

The current input is known; its successor is the target.

Query 2 → next token

1. An operationKeys 1,2: allowed; key 3: forbidden
2. A list of numbersWeights [0.5, 0.5, 0] → output [1, 1]

The denominator includes the two allowed keys.

Query 3: real or padding?

1. An operationReal: keys 1,2,3 allowed
2. An operationPadding: query unused; key 3 excluded

A real empty row is an error. An unused padded row returns zero by explicit convention.

Figure 16.1 — The inclusive diagonal preserves known input; future edges are forbidden. Original teaching diagram, CC BY-SA 4.0. Values and rules are declared in the chapter and shared data.

Figure 16.1 makes the diagonal and forbidden future edges explicit. When slot 3 is padding, we additionally exclude it as a key. Its query output is explicitly zeroed and omitted from loss; outputs 1 and 2 remain [1,0] and [1,1]. This is a declared convention for unused rows, not a probability distribution of all zeros.

What if a real query has no visible key? There is no mixture to normalize: every exponential is zero, giving a zero denominator. With maximum subtraction, an all-negative-infinity row can also create undefined intermediate arithmetic. Our implementation raises ValueError("a valid query has no visible key"). Padded queries take a separate safe path and return zero. A batch containing only ignored targets is also rejected when loss is computed. Quietly returning uniform weights for a real empty row would invent contributors that the specification forbids.

16.3 A model that appears excellent by reading the answer

Example 16.2 — A deliberately invalid future-copy oracle. Use two candidate labels in fixed order [red,blue]. At the first input position, the known token is BOS and its supplied value is [0.5,0.5]. At the second position, the token is either red, with supplied value [1,0], or blue, with supplied value [0,1]. The first position must predict this second token. We deliberately supply its score row as [0,ln(99)]. No parameter learning occurs in this example.

Allowing both positions gives exponentials [1,99], weights [0.01,0.99]. For the red record, the output is $0.01[0.5,0.5]+0.99[1,0]=[0.995,0.005]$. For the blue record it is [0.005,0.995]. These special V rows are candidate distributions, so their mixture can be scored directly; ordinary attention values are not vocabulary probabilities. This construction isolates the leak without hiding it in a language-model head.

Table 16.2 — Low loss without learning, when the answer is visible. The loss is negative natural log probability, in nats per first-position target.

Visibility	Correct-label probability, red record	Correct-label probability, blue record	Mean loss
Future visible	0.995	0.995	0.005013
Causal, BOS only	0.5	0.5	0.693147

A separately stored held-out record would look just as easy to this oracle: its future input still contains its target. The low loss proves only that we supplied the answer through a forbidden path. It establishes neither a useful predictor nor generalization.

An effective implementation test changes future inputs while holding a prefix fixed. With fixed parameters and no training-mode randomness, all logits computed for that prefix must remain unchanged under causal attention. Our tests also intentionally remove or reverse the triangle and require the test fixture’s prefix output to change. The bad version failing the intended invariant is evidence that the test detects this leak; shape assertions alone would accept it.

16.4 Why one mixture may be too restrictive

A single head uses one row of position weights for every coordinate it receives. Perhaps different combinations of input coordinates need different comparisons and mixtures. We can give each combination its own Q/K/V projections and compute them in parallel. Each such calculation is a head; their combination is multi-head attention. A head is defined by its parameters and calculation, not by a promised human-interpretable specialty.

Example 16.3 — Two small heads and a projection that combines them. Keep the three [2] rows of X from Chapter 15. This time use two heads of width 1, instead of one head of width 2. For head 1, each of its three projection matrices is the column [1,0]; for head 2 each is [0,1]. Written as matrices:

$$W_{Q,1}=W_{K,1}=W_{V,1}=\begin{bmatrix}1\\0\end{bmatrix},\qquad W_{Q,2}=W_{K,2}=W_{V,2}=\begin{bmatrix}0\\1\end{bmatrix}.$$

Thus head 1’s Q, K, and V are each the column [1,0,1]; head 2’s are each [0,1,1]. A projection has shape [2,1], a projected array [3,1], and each score matrix [3,3]. The divisor is $\sqrt{1}=1$. Apply the same causal triangle to both heads. They share visibility rules, but calculate different scores:

$$D_1=\begin{bmatrix}1&0&1\\0&0&0\\1&0&1\end{bmatrix},\qquad D_2=\begin{bmatrix}0&0&0\\0&1&1\\0&1&1\end{bmatrix}.$$

For head 1’s second query, the allowed scores are [0,0], so its output is $(1+0)/2=0.5$. For head 2’s second query, allowed scores [0,1] give weights $[1,e]/(1+e)$, and output $e/(1+e)\approx0.731059$. At position 3, both outputs happen to equal $2e/(2e+1)\approx0.844638$, though their weights attach to different positions. Head 1 weights are [e,1,e]/(2e+1); head 2 weights are [1,e,e]/(2e+1). Equal outputs do not imply equal reads.

Place head 1’s output first and head 2’s second. This concatenation joins coordinates; it does not average heads. Then apply a stored output matrix:

$$J\approx\begin{bmatrix}1&0\\0.5&0.731059\\0.844638&0.844638\end{bmatrix},\qquad W_O=\begin{bmatrix}1&1\\1&-1\end{bmatrix},\qquad Y=JW_O\approx\begin{bmatrix}1&1\\1.231059&-0.231059\\1.689275&0\end{bmatrix}.$$

J has shape [3,2]; W_O [2,2]; Y [3,2]. The second output’s first coordinate adds the two head results, while its second subtracts them. This output projection lets the next layer use combinations of head coordinates. Its entries, like the input projections, are ordinarily learned. We chose them explicitly here.

Head 1, width 1

1. Stored model numbersW_Q = W_K = W_V = [1, 0]ᵀ
2. A list of numbersAllowed scores [0, 0] → [0.5, 0.5]
3. An operation0.5×1 + 0.5×0 = 0.5

All three positions are processed; this lane follows position 2.

Head 2, width 1

1. Stored model numbersW_Q = W_K = W_V = [0, 1]ᵀ
2. A list of numbersAllowed scores [0, 1] → [0.268941, 0.731059]
3. An operation0.268941×0 + 0.731059×1 = 0.731059

Same allowed keys; different comparison coordinates.

Join, then project

1. A list of numbersJ₂ = [0.5, 0.731059]
2. Stored model numbersW_O rows: [1, 1]; [1, -1]
3. A list of numbersY₂ = [1.231059, -0.231059]

Concatenation order must agree with the input rows of W_O.

Figure 16.2 — Two mixtures become a sum and a difference. Original teaching diagram, CC BY-SA 4.0. Values and rules are declared in the chapter and shared data.

Figure 16.2 follows position 2 through both mixtures. Changing the concatenation order without changing W_O changes the meaning of its input rows. To preserve the function when swapping the heads, swap the corresponding rows of W_O too.

16.5 Preserve the axes in code

For a conventional implementation with B examples, T positions, width C, H heads, and per-head width d, choose $C=Hd$. For example, C=16 and H=2 give d=8. Project to [B,T,C], reshape to [B,T,H,d], then exchange the T and H axes to obtain [B,H,T,d]. Reshape groups coordinates; transpose relocates an axis. They do not perform the same job.

Q times K with its last two axes swapped gives scores [B,H,T,T]. A visibility array [B,1,T,T] can broadcast the same relation across heads while preserving different score values. After weights multiply V, the head outputs are [B,H,T,d]. Transpose back to [B,T,H,d], join the last two axes, then multiply W_O [C,C]. The result is [B,T,C]. Using .T on a four-dimensional NumPy array reverses all axes, which is not the required key transpose.

At fixed C, increasing H does not automatically multiply the total projection width: each d gets smaller. But the explicit score array still stores $BHT^2$ numbers. With B=1, H=2, T=3, it contains 18 scores. If stored as 8-byte float64 numbers, the scores alone occupy 144 bytes, excluding weights, gradients and other arrays. Doubling T gives 72 scores and 576 bytes. This is an array-size calculation, not a runtime or full-training-memory benchmark.

The practice implementation in code/part-iv/mechanisms.py exposes every stage. Its masked_softmax uses True for allowed positions. Do not transfer that convention blindly: in the versioned PyTorch comparison, functional.scaled_dot_product_attention uses True for allowed entries, while nn.MultiheadAttention Boolean attention and key-padding masks use True for forbidden or ignored entries. The precise adapters and CPU tests live with the practice code; mathematical visibility does not depend on those names. Functional API, module API.

16.6 What the restrictions and heads do not supply

A causal mask removes future positions from the allowed computation. It cannot repair a target copied into an earlier input field, a contaminated data split, or a preprocessing step that used future labels. Multiple heads allow distinct projections and mixtures; they do not guarantee that one head handles color, another syntax, or that all heads are useful.

We now have controlled information exchange and an output with the original model width. We still need a clear account of position, a way to preserve the incoming representation, and computation within each position beyond the attention mixtures. Those are the missing pieces of Chapter 17’s complete block.

16.7 Exercises

1. A four-slot record contains [BOS,red,key,PAD], with targets [red,key,EOS,ignored]. List allowed keys and scored targets under causal plus padding restrictions. Is the diagonal allowed at the third slot?
2. Someone applies the causal mask after Softmax to Chapter 15’s first row, obtaining [0.401112,0,0]. Compute its mixture with V and explain why it differs from the specified causal result.
3. In Example 16.2, replace 99 with 9 in the exponential ratio. Compute the leaked correct-label probability and loss. What remains unchanged when the future is forbidden?
4. Swap Example 16.3’s head order but leave W_O fixed. Compute position 2’s new result. How should W_O change to recover the original result?
5. An implementation uses a strict lower triangle, excluding the diagonal. Explain what happens to the first real query and why returning a uniform row silently changes the task.
6. A validation report shows equal shapes and decreasing loss but no future-perturbation test. State one mistake those observations cannot rule out, and one limitation that remains even if the perturbation test passes.

16.8 Reference answers

1. Allowed keys are [1], [1,2], [1,2,3] for the three real queries. The fourth query is unused and returns a zero mixture by convention. Targets red, key, and EOS are scored; the padded target is not. The third diagonal is allowed: key is known when predicting EOS. A placeholder is not a vocabulary event to learn.
2. The mixture is approximately [0.401112,0], since V row 1 is [1,0]. The weight sum is only 0.401112. Proper causal normalization gives [1,0,0] and output [1,0]. Dividing the remaining positive weights by their sum can recover this simple case, but masking scores before Softmax directly expresses the intended rule and handles numerical scale more reliably.
3. Weights become [0.1,0.9]; correct probability is $0.1\times0.5+0.9\times1=0.95$, and loss is $-\ln(0.95)\approx0.051293$. Causal visibility still leaves only BOS, yielding probability 0.5 and loss 0.693147. Changing the strength of the forbidden shortcut does not make it legitimate.
4. The concatenated row becomes [0.7310585786300049,0.5]. Its output is [1.2310585786300048,0.2310585786300049]; the second coordinate changes sign. Swap the two rows of W_O to compensate, producing rows [1,-1] and [1,1]. Then the same head value again multiplies the same output coefficients.
5. The first query has no allowed key. Its probability normalizer is zero, so our valid-row check raises an error. Uniform weights would let it read forbidden positions. With shifted next-token targets, the inclusive diagonal is the intended relation; a different modeling convention needs its own specified inputs and visibility.
6. Reversed or absent causal masking can yield the same shapes and low loss through future access. Perturbing future inputs tests the intended prefix invariant. It does not prove that labels were absent from earlier fields, that data splits are independent, or that the model generalizes to new tasks.

16.9 References

• Vaswani et al. (2017), Attention Is All You Need. Sections 3.2.2–3.2.3 specify projected heads, concatenation/output projection, and masking before Softmax. Our tiny heads and leakage oracle are independently constructed examples.
• PyTorch, scaled_dot_product_attention, 2.7 and MultiheadAttention, 2.7. Their parameter descriptions support the explicit comparison of Boolean-mask meanings; neither convention defines mathematical causality.