Worked examples — Masked attention for autoregression

Before we start, one word we lean on constantly:

The mask $M$ adds $-\infty$ to every cell above the diagonal (the future) and $0$ everywhere else. After softmax those $-\infty$ cells become exactly $0$. That is the whole machine. Now let us exhaust it.

---

The scenario matrix

Every situation the causal mask can meet falls into one of these cells. Each row is a "case class"; the last column names the example that nails it.

#	Case class	What is unusual about it	Worked in
A	First row ($i=1$)	Only one visible token — softmax over a single number	Example 1
B	Last row ($i=T$)	Full past visible, nothing masked	Example 2
C	Middle row ($1<i<T$)	Some past, some masked future	Example 3
D	Degenerate $T=1$	Whole sequence is one token	Example 4
E	Ties in scores	Two allowed positions have equal $S_{ij}$	Example 5
F	Limiting value	One masked cell is huge / one allowed cell dominates	Example 6
G	Real-world word problem	Generation loop, next-token prediction	Example 7
H	Exam twist / trap	Mask applied after softmax (the wrong way)	Example 8

The numbers below reuse one small running example so you can watch the mask act on the same data across cases. We fix $d_k=2$ (so $\sqrt{d_k}=\sqrt{2}\approx 1.414$) and this un-masked score matrix for $T=3$:

$S=\frac{Q{K}^{\top }}{\sqrt{2}}=\left[\begin{array}{ccc}2.83& 4.95& 2.83\\ 4.95& 4.24& 4.95\\ 4.24& 5.66& 4.24\end{array}\right]$

(This is exactly the matrix from Example 2 of the parent note, rounded to two decimals.)

Look at the picture: the red cells are the future — the mask will erase them. The green staircase is what each row is actually allowed to see. Notice the staircase grows by one cell per row. That growth is autoregression.

---

Worked examples

Example 1 — Cell A: the very first row

1. Write the raw row. Row 1 of $S$ is $[2.83,4.95,2.83]$. Why this step? We always start from the un-masked scores, then let the mask decide what survives.

2. Apply the mask. For $i=1$, only $j\le 1$ is allowed, so columns 2 and 3 get $-\infty$: $[2.83,-\infty ,-\infty ].$ Why this step? Token 1 is generated first; there is literally no past and no permitted future, so it can only attend to itself.

3. Softmax. $e^{-\infty }=0$, so ${\alpha }_1=\left[\frac{e^{2.83}}{e^{2.83}},0,0\right]=[1,0,0].$ Why this step? Softmax turns scores into weights summing to 1. With one survivor the answer is forced: it must be $1$.

Verify: a softmax over exactly one live entry is always $1$, regardless of that entry's size — because $e^x/e^x=1$. The output ${\text{out}}_1=1\cdot V_1$: the first token attends purely to itself. ✅

---

Example 2 — Cell B: the last row, nothing masked

1. Raw row. Row 3 is $[4.24,5.66,4.24]$. Why this step? Same starting point; the mask will simply add $0$ everywhere here.

2. Apply the mask. For $i=3$, all $j\le 3$ are allowed, so $M$ adds $[0,0,0]$ — the row is unchanged. Why this step? The last token's "future" is empty, so masking removes nothing. This is the one row where causal attention and full attention agree.

3. Softmax. Using $e^{4.24}\approx 69.4$, $e^{5.66}\approx 287.1$: ${\alpha }_3=\left[\frac{69.4}{425.9},\frac{287.1}{425.9},\frac{69.4}{425.9}\right]\approx [0.16,0.68,0.16].$ Why this step? Convert to a probability distribution over the three visible tokens.

Verify: $0.16+0.68+0.16=1.00$ ✅ — and the two equal scores ($4.24$) yield two equal weights ($0.16$), a sanity check that softmax preserves ties (we press on this in Example 5). Answer confirms: for the last row masked attention = standard attention.

---

Example 3 — Cell C: a middle row, half masked

1. Raw row. Row 2 is $[4.95,4.24,4.95]$. Why this step? Column 3 has a high raw score ($4.95$) — a perfect chance to see the mask do real work by deleting a tempting future cell.

2. Apply the mask. For $i=2$, $j=3$ is future $\to -\infty$: $[4.95,4.24,-\infty ].$ Why this step? Even though column 3 scored high, it is the future. The mask must zero it or the model would cheat by copying the answer.

3. Softmax over the two survivors. $e^{4.95}\approx 141.2$, $e^{4.24}\approx 69.4$, sum $\approx 210.6$: ${\alpha }_2\approx \left[\frac{141.2}{210.6},\frac{69.4}{210.6},0\right]\approx [0.67,0.33,0].$ Why this step? The two live scores compete; the higher one (column 1) wins the larger share.

Verify: $0.67+0.33+0=1.00$ ✅. The masked column contributes exactly $0$, so ${\text{out}}_2=0.67V_1+0.33V_2$ — no leakage from $V_3$. ✅

---

Example 4 — Cell D: the degenerate one-token sequence

1. Build the mask. $M$ is $1\times 1$ with $M_{11}=0$ (the diagonal is always allowed). Why this step? The diagonal is never masked — a token may always attend to itself (this is Mistake 3 in the parent note).

2. Softmax over one entry. $\alpha =\left[\frac{e^{3.0}}{e^{3.0}}\right]=[1]$. Why this step? Same forced-to-one logic as Example 1, now for the whole matrix.

Verify: the weight is $1$ for any $s_{11}$; the value of $s_{11}=3.0$ never enters the answer. Output is ${\text{out}}_1=V_1$. A one-token sequence is a fixed point of masked attention — nothing to leak, nothing to weigh. ✅

---

Example 5 — Cell E: exact ties among allowed positions

1. Confirm the mask leaves both ties alive. Columns 1 and 2 are $j\le i=2$, so both survive; only column 3 is $-\infty$. Why this step? We want to isolate the effect of ties, not of masking, so we make sure the mask does not break the tie for us.

2. Softmax. $e^{2.0}\approx 7.389$ for each; sum $=2\times 7.389=14.778$: $\alpha =\left[\frac{7.389}{14.778},\frac{7.389}{14.778},0\right]=[0.5,0.5,0].$ Why this step? Softmax is symmetric — equal inputs give equal outputs. The masked cell still contributes $0$.

Verify: $0.5+0.5+0=1.00$ ✅ and the two weights are bit-for-bit equal. Ties split evenly; the mask does not disturb this. ✅

---

Example 6 — Cell F: limiting behaviour (a dominating cell)

1. Compute exactly. $e^{10}\approx 22026.5$, $e^0=1$, sum $\approx 22028.5$: $\alpha \approx \left[\frac{22026.5}{22028.5},\frac{1}{22028.5},\frac{1}{22028.5}\right]\approx [0.99991,0.0000454,0.0000454].$ Why this step? We test the softmax at an extreme so we can see its limiting shape.

2. Read the limit. As the gap $\to \infty$, the weights approach $[1,0,0]$ — softmax becomes a hard "argmax". Why this step? This tells us how masking interacts with confidence: a very sharp distribution is fine, and a masked $-\infty$ is just the extreme case of a very small score forced to weight $0$.

Verify: $0.99991+0.0000454+0.0000454\approx 1.00$ ✅. And notice a masked cell ($-\infty$) is exactly this limit taken all the way: $e^{-\infty }=0$, contributing precisely $0$ — the mask is "the argmax limit for the future." ✅

---

Example 7 — Cell G: a real generation loop

1. List the query's allowed keys. Row $i=6$ allows $j\le 6$, i.e. columns 1–6. But column 6 is the token we are predicting — it is not yet generated. Why this step? During generation the future simply doesn't exist, so there is nothing to mask beyond what already isn't there. Compare with training, where all 6 tokens sit in the buffer and the mask must actively hide column 6 to prevent copying.

2. Suppose the (already scaled) visible scores over columns 1–5 are $[0.5,1.2,0.3,2.0,0.4]$. Softmax them: $e^{0.5}=1.6487,e^{1.2}=3.3201,e^{0.3}=1.3499,e^{2.0}=7.3891,e^{0.4}=1.4918$; sum $\approx 15.2996$. $\alpha \approx [0.108,0.217,0.088,0.483,0.098].$ Why this step? This is the actual distribution the model uses to weight $V_1\dots V_5$ before producing "Paris". "France" (column 4) dominates — exactly the context that should trigger "Paris".

3. Predict. The blended value ${\sum }_j{\alpha }_{6j}{V}_j$ feeds the output layer, which emits "Paris". No cell for token 6 appears — the answer was inferred, never copied. Why this step? This is the whole payoff of masking: training with the mask made the model behave identically to this leak-free inference loop (teacher forcing supplied the ground-truth prefix during training).

Verify: $0.108+0.217+0.088+0.483+0.098=0.994\approx 1.00$ (rounding) ✅ and the largest weight sits on "France", the semantically correct cue. ✅

---

Example 8 — Cell H: the exam trap (mask after softmax)

1. Softmax first. $e^2=7.389,e^3=20.086,e^4=54.598$; sum $=82.073$: $\text{softmax}=[0.0900,0.2447,0.6652].$ Why this step? This is the (wrong) order the student used.

2. Zero the future with the binary mask. Multiply by $[1,1,0]$: $[0.0900,0.2447,0],\text{sum}=\mathrm{0.3348.}$ Why this step? The masked weight is removed, but nobody re-scaled the rest — the distribution now leaks probability mass, summing to $0.33$, not $1$.

3. Contrast with the correct order. Applying $-\infty$ before softmax gives $\text{softmax}([2,3,-\infty ])=[0.269,0.731,0]$, which does sum to 1. Why this step? Masking before softmax renormalizes automatically — the denominator only sums the survivors. That is why the parent note insists on $-\infty$-before-softmax, not $0$-or-$1$-after (Mistake 1).

Verify: wrong order sums to $0.0900+0.2447+0=0.3348\ne 1$ ✅ (trap confirmed); correct order sums to $0.269+0.731=1.000$ ✅.

---

Recall

Recall Which rows of the mask actually delete a cell?

Every row except the last ::: row $i$ masks columns $i+1,\dots ,T$; the last row ($i=T$) has no future, so nothing is deleted (Example 2).

Recall A softmax over a single surviving score always equals…

$1$ ::: because $e^x/e^x=1$ regardless of $x$ (Examples 1 and 4).

Recall Why is the diagonal never masked?

A token must be able to attend to itself for self-reference ::: masking hides only future positions $j>i$, never the current one $j=i$ (Example 4).

Recall What breaks if you mask after softmax with a 0/1 mask?

The weights no longer sum to 1 ::: e.g. $[0.09,0.24,0]$ sums to $0.33$; you'd have to renormalize by hand (Example 8).

Related: Decoder Architecture, Encoder-Decoder Models, Positional Encoding.