Masked attention for autoregression: Exercises

This is your self-testing gym for masked attention. Work each problem before opening the solution. Difficulty climbs in five levels:

L1 Recognition — can you spot the right object?
L2 Application — can you run the machine?
L3 Analysis — can you explain why it behaves that way?
L4 Synthesis — can you combine masking with other ideas?
L5 Mastery — can you handle the edge cases and design decisions?

Every symbol used here is defined the moment it appears. If you ever feel a term arrived unearned, that is a bug — but let us re-anchor the objects we lean on most:

Definition The three objects we keep using

Attention score matrix $S$ : a grid of numbers. Row $i$ , column $j$ — written $S_{ij}$ — is "how strongly the token at position $i$ wants to listen to the token at position $j$ ." Bigger number = louder listening.
Scaling by $d_{k}$ : before we do anything with $S$ we divide it by $d_{k}$ , where $d_{k}$ is the length of each query/key vector. Why? Each score $S_{ij}$ is a dot product of two $d_{k}$ -long vectors — a sum of $d_{k}$ product terms. If the entries are roughly unit-sized, that sum grows with $d_{k}$ (its standard deviation scales like $d_{k}$ ). Large scores push softmax into a near one-hot spike where gradients vanish. Dividing by $d_{k}$ cancels that growth and keeps the scores in a healthy range — this is exactly Attention Score Scaling. We call the divided matrix the scaled scores.
Causal mask $M$ : a second grid, same shape, whose only job is to silence the future. Where the future lives ( $j > i$ , meaning column is to the right of the diagonal) it holds $- \infty$ ; everywhere else it holds $0$ . We add $M$ to the scaled scores before softmax so that the future's scores drop to $- \infty$ and softmax turns them into exactly $0$ .

Before the problems, fix the geometry in your head with the figure below. Refer back to it whenever a solution says "row $i$ " or "the diagonal" — it is the map every exercise navigates.

Figure — Masked attention for autoregression

$softmax (x)_{j} = \frac{e ^{x_{j}}}{\sum _{k} e ^{x_{k}}}$

We use softmax (and not, say, "just normalize by dividing") because we want a differentiable function that always yields positive weights summing to one — that is what lets us treat attention weights as probabilities and train them with gradients.

Level 1 — Recognition

Recall Solution

Answer: $A$ . A valid causal mask keeps the diagonal and everything below it at $0$ (past + self allowed), and sets everything above the diagonal ( $j > i$ ) to $- \infty$ (future silenced). That is a lower-triangular-of-zeros pattern — exactly the teal+orange cells of Figure s01.

$A$ : diagonal and below are $0$ , strictly-upper is $- \infty$ . ✅
$B$ : this is the transpose — it silences the past instead of the future. Wrong.
$C$ : it masks the diagonal too ( $C_{11} = - \infty$ ), which would forbid a token from attending to itself. That is the "diagonal-mask" error — a token must be allowed to attend to itself (we prove why in L5.2, where masking the diagonal makes the very first row empty and produces NaN). Wrong.

Recall Solution

Position $i$ may attend to ${1, 2, \dots, i}$ . For $i = 3$ : ${1, 2, 3}$ . It may not see ${4, 5}$ (the future). Count of allowed keys $= i = 3$ .

Level 2 — Application

Recall Solution

Position $i = 2$ may see ${1, 2}$ ; positions $3, 4$ are future $\Rightarrow$ set to $- \infty$ . Masked row: $[2, 4, - \infty, - \infty]$ . Softmax over the two survivors: $e^{2} = 7.389$ , $e^{4} = 54.598$ , sum $= 61.987$ . $[\frac{7.389}{61.987}, \frac{54.598}{61.987}, 0, 0] = [0.12, 0.88, 0, 0]$ Weights sum to $1$ , and the future columns are exactly $0$ . ✅

Recall Solution

Step 1 — scores $Q K^{⊤}$ . ( $Q K^{⊤}$ answers "similarity of each query to each key" — a dot product is large when two vectors point the same way, which is why it measures alignment.) $Q K^{⊤} = [1001] [1120] = [1120]$ Step 2 — scale by $2$ (recall from the top of the page: this divides by $d_{k}$ to stop the dot products from growing with $d_{k}$ ): $[0.707 0.707 1.414 0]$ . Step 3 — mask. Row 1 ( $i = 1$ ) may see only ${1}$ : $[0.707, - \infty]$ . Row 2 ( $i = 2$ ) sees ${1, 2}$ : $[0.707, 0]$ . Step 4 — softmax. Row 1: only one survivor $\Rightarrow [1, 0]$ . Row 2: $e^{0.707} = 2.028$ , $e^{0} = 1$ , sum $= 3.028$ $\Rightarrow [0.6698, 0.3302]$ . Step 5 — weights $\times V$ . Position 1: $1 \cdot [10, 0] + 0 \cdot [0, 10] = [10, 0]$ . Position 2: $0.6698 \cdot [10, 0] + 0.3302 \cdot [0, 10] = [6.698, 3.302]$ . Output $= [10 6.70 0 3.30]$ . Notice: token 1's output is pure $V_{1}$ — it had no choice but itself. That is the causal constraint made visible.

Level 3 — Analysis

Recall Solution

No mask: every position attends to all $T$ positions $\Rightarrow T \times T = T^{2}$ edges. Causal mask: position $i$ attends to $i$ positions, so $\sum_{i = 1}^{T} i = \frac{T ( T + 1 )}{2} .$ Ratio (causal / full) $= \frac{T ( T + 1 ) /2}{T ^{2}} = \frac{T + 1}{2 T}$ . For $T = 1000$ : $\frac{1001}{2000} = 0.5005$ . So masking keeps just over half the edges — the lower triangle plus the diagonal. Insight: the cost is still $Θ (T^{2})$ ; masking halves the count but not the asymptotic order. This is why long-context attention is expensive with or without a mask.

Recall Solution

The masked cell's weight after softmax is $\frac{e ^{- 1 0^{9}}}{row sum} .$ $e^{- 1 0^{9}}$ underflows to $0.0$ in floating point, so numerically it is identical to $- \infty$ for any realistic real score $s$ (typically $∣ s ∣ < 100$ ). The difference only ever appears in pathological arithmetic where a genuine score also reaches $\sim - 1 0^{9}$ ; that never happens for scaled dot products. So $- 1 0^{9}$ is a safe, common implementation choice — and it avoids the $- \infty - \infty = NaN$ hazard that can occur if an entire row is masked.

Recall Solution

Softmax computes $\frac{e ^{- \infty}}{e ^{- \infty} + \dots} = \frac{0}{0} =$ NaN (not-a-number). A single NaN propagates through the whole network and poisons the loss. Cure: never mask a query's own diagonal (a token always keeps itself — see L5.2), or drop fully-padded query rows before softmax, or use the $- 1 0^{9}$ trick from L3.2 which yields a uniform-but-finite row instead of NaN. This is the degenerate case masking must survive.

Level 4 — Synthesis

Recall Solution

At step $t$ we have only ever created keys/values for positions $1.. t$ ; positions $t + 1.. T$ do not exist yet, so there is physically nothing future to attend to. The cache holds exactly $K_{1 : t}, V_{1 : t}$ . Softmaxing $q_{t}$ over these $t$ keys is identical to row $t$ of the masked training matrix (the survivors are ${1, \dots, t}$ ). So the mask is "baked into" the geometry of incremental decoding — we get causality for free. This is why KV Caching and causal masking are consistent: training with the lower-triangular mask makes each row match what inference will later see.

Recall Solution

Teacher Forcing hands the model the whole ground-truth sequence at once so all $T$ predictions can be computed in one parallel pass (fast). But parallelism means every position could peek at every other — including the answer sitting one slot to the right. The causal mask restores the promise "position $i$ only sees $1.. i$ ", so the parallel computation produces exactly the same conditioning $P (y_{i} ∣ y_{< i})$ that the slow, one-at-a-time inference of L4.1 uses. Mask = the bridge that lets us train fast but behave sequentially.

Recall Solution

(a) Decoder self-attention — masked. The decoder is generating autoregressively, so it must not see its own future tokens. (b) Cross-attention — not masked. The encoder has already read the entire source sentence (it is the input, fully available), so a decoder token is allowed to attend to all encoder positions. Blocking future source positions would throw away legitimately available information. Rule of thumb: mask only where a "future" would be cheating — that is only inside the decoder's self-attention.

Level 5 — Mastery

Recall Solution

Because each $V_{j}$ is $10$ on a distinct axis, the weighted sum $out_{3} = α_{1} V_{1} + α_{2} V_{2} + α_{3} V_{3}$ separates coordinate-by-coordinate: $out_{3} = [10 α_{1}, 10 α_{2}, 10 α_{3}] .$ Matching to $[3, 5, 2]$ : $10 α_{1} = 3, 10 α_{2} = 5, 10 α_{3} = 2 \Rightarrow α = [0.3, 0.5, 0.2]$ . Sanity checks: all $\geq 0$ ✅, sum $= 1.0$ ✅ (a valid softmax distribution), no future column (there is none for $i = 3 = T$ ) ✅.

Recall Solution

For $i = 1$ the allowed set becomes ${1, \dots, 0} = \emptyset$ — empty (there is no position $0$ in our 1-based numbering). The score row is entirely $- \infty$ , and (L3.3) softmax returns $0/0 =$ NaN. The first token would have nothing to attend to, poisoning training instantly. This is precisely why the correct mask keeps the diagonal ( $j \leq i$ ): every position, including the first, always has at least itself to attend to.

Recall Solution

The output is $out_{i} = \sum_{j = 1}^{T} α_{ij} V_{j}$ with $α_{ij} = \frac{e x p ( s _{ij} + M _{ij} )}{\sum _{k = 1}^{T} e x p ( s _{ik} + M _{ik} )} .$ For $j > i$ , $M_{ij} = - \infty \Rightarrow exp (s_{ij} - \infty) = 0 \Rightarrow α_{ij} = 0$ . Therefore $out_{i} = \sum_{j = 1}^{i} α_{ij} V_{j} + \sum_{j = i + 1}^{T} 0 \cdot V_{j} = \sum_{j = 1}^{i} α_{ij} V_{j} .$ No $V_{j}$ with $j > i$ appears. Its coefficient is exactly $0$ , independent of the value $s_{ij}$ . Causality is guaranteed for every row, every quadrant of the matrix — the future has zero influence by construction. $■$

Recall Quick self-quiz (cover the answers)

What triangle of the mask holds $- \infty$ ? ::: The strictly-upper triangle ( $j > i$ ), the future. Why mask before softmax, not after? ::: So weights renormalize over survivors and still sum to $1$ . Why keep the diagonal unmasked? ::: A token must attend to at least itself; masking it gives an empty row and NaN for $i = 1$ . Why divide scores by $d_{k}$ ? ::: A dot product of $d_{k}$ -long vectors grows like $d_{k}$ ; dividing keeps softmax out of the vanishing-gradient spike. Does cross-attention in an encoder-decoder get a causal mask? ::: No — the full source is available; only decoder self-attention is masked. Fraction of edges kept for large $T$ ? ::: About one half, $\frac{T + 1}{2 T} \to \frac{1}{2}$ .

Related pages: Self-Attention Mechanism · Decoder Architecture · Attention Score Scaling · Positional Encoding · KV Caching · Teacher Forcing · Encoder-Decoder Models.