Masked attention for autoregression: Visual walkthrough

We will build up the machinery of self-attention just far enough to see where a leak happens, then plug it — one step, one picture at a time.

Step 1 — What "attending" even means: a grid of glances

WHAT. A sequence is just a row of tokens. Give each token a position number $1, 2, 3, \dots$ — this is the same numbering that Positional Encoding injects into the vectors. When token number $i$ "reads" the sentence, it decides how much attention to pay to every other token $j$ . Collect all those glances into a square grid.

WHY a grid? Because "every token looks at every token" is exactly a table with one row per looker and one column per looked-at. Call it $A$ . The entry in row $i$ , column $j$ is one number: how strongly $i$ glances at $j$ .

PICTURE. Look at the grid. The looker sits on the row, the target sits on the column. The red cell is "token 3 glancing at token 4" — a glance into the future.

Figure — Masked attention for autoregression

Step 2 — Where the score comes from: query · key

WHAT. Each token produces two little vectors: a query $Q_{i}$ ("what am I looking for?") and a key $K_{j}$ ("what do I offer?"). The glance score is their dot product.

WHY the dot product and not something else? The dot product $Q_{i} \cdot K_{j}$ is large when the two arrows point the same way and small when they point apart. It is the cheapest possible "how aligned are these two?" meter — exactly the question "does what token $i$ wants match what token $j$ offers?". That is why attention uses it rather than, say, a distance.

PICTURE. Two arrows: a query and a key. When they line up (small angle) the shaded overlap — the dot product — is big and red; when they are perpendicular it collapses to zero.

Step 3 — Taming the size: divide by $d_{k}$

WHAT. Divide every score by $d_{k}$ , where $d_{k}$ is the length of the query/key vectors (how many numbers each holds).

WHY divide, and why that number? A dot product of two $d_{k}$ -long random vectors grows roughly like $d_{k}$ . If we let scores balloon, the next step (softmax) turns tiny score gaps into all-or-nothing spikes and the gradients die. Dividing by $d_{k}$ cancels that growth exactly — this is the whole point of Attention Score Scaling.

PICTURE. Left: tall spiky bars (unscaled). Right: the same bars gently flattened after dividing — the red bar is still the tallest, but nothing is off the chart.

Step 4 — Spotting the leak

WHAT. Look again at the full scaled grid $\tilde{s}$ . Split it along the diagonal. Everything above the diagonal ( $j > i$ ) is a token peeking at something that comes later in the sentence.

WHY this is fatal. At generation time we produce tokens left-to-right; when we are predicting token $i$ , tokens $i + 1, i + 2, \dots$ do not exist yet. If training lets row $i$ read those future columns, the model learns the shortcut "copy the answer sitting in column $i + 1$ ", then falls apart at inference. This is the train–test mismatch the parent note warns about, and it is exactly why Teacher Forcing alone is not enough.

PICTURE. The grid, diagonal drawn in bold. The red upper-triangle is labelled "THE FUTURE — must not be read". The lower part is labelled "past + present — allowed".

Step 5 — The mask: add $- \infty$ to the future

WHAT. Build a second grid $M$ , the same size, filled with $0$ everywhere we allow and $- \infty$ everywhere we forbid. Add it to the scores before softmax.

WHY $- \infty$ and not "multiply by zero"? Because the very next step exponentiates: $e^{- \infty} = 0$ , cleanly and automatically. If instead we zeroed weights after softmax, the surviving weights would no longer add up to $1$ and we'd have to renormalise by hand (the parent's Mistake 1). Adding $- \infty$ before softmax makes the renormalisation happen for free.

WHY only the strict upper triangle? The diagonal ( $j = i$ ) stays $0$ — a token is allowed to look at itself (Mistake 3 in the parent). We only kill $j > i$ .

PICTURE. Two stacked grids: the score grid, a big red "+", and the mask grid ( $0$ in the lower triangle, $- \infty$ glowing red in the upper triangle).

Step 6 — Softmax turns $- \infty$ into a clean zero

WHAT. Apply softmax across each row. Softmax takes a row of numbers and returns positive weights that sum to $1$ : $α_{ij} = \frac{e ^{\tilde{s}_{ij} + M_{ij}}}{\sum _{k = 1}^{T} e ^{\tilde{s}_{ik} + M_{ik}}}$

WHY softmax? We need the glances to become a probability distribution — a set of positive numbers adding to $1$ — so the output is a proper weighted average. Exponentiating makes everything positive; dividing by the row sum makes it total $1$ .

WHAT IT LOOKS LIKE. The masked cells had $- \infty$ , so $e^{- \infty} = 0$ : they contribute nothing to the sum and receive weight $0$ . The allowed cells share all the probability among themselves.

PICTURE. A single row before and after softmax. The two future cells (red) collapse to a flat $0$ ; the past+self cells rise into bars that together fill the bar of height $1$ .

Step 7 — The output only mixes the past

WHAT. Each token's output is the weighted blend of the value vectors $V_{j}$ , using those weights: $out_{i} = \sum_{j = 1}^{T} α_{ij} V_{j} = \sum_{j = 1}^{i} α_{ij} V_{j}$

WHY this proves causality. Every $α_{ij}$ with $j > i$ is $0$ , so those $V_{j}$ drop out. The sum genuinely stops at $i$ . Token $i$ 's new representation is built only from tokens $1 \dots i$ — the leak is sealed.

WHAT IT LOOKS LIKE. The final grid is lower-triangular: solid weights below the diagonal, blank white above it — a perfect staircase. This same staircase is what lets KV Caching reuse past keys/values, and it is the defining feature of the Decoder Architecture (versus the full grid used inside Encoder-Decoder Models).

The one-picture summary

Here is the whole journey on one canvas: raw scores $\to$ add mask $\to$ softmax $\to$ lower-triangular weights $\to$ output from the past only.

Recall Feynman retelling — say it back in plain words

Every token asks a question (its query) and every token wears a label (its key). We measure how well each question matches each label with a dot product, giving a grid of glance scores. We shrink the scores by $d_{k}$ so the next step doesn't explode. Now the danger: a token could glance forward at words that haven't been generated yet, so it would learn to copy instead of predict. To stop this we add a mask that dumps $- \infty$ onto every forward-looking cell. Softmax then exponentiates, and $e^{- \infty} = 0$ , so those cells politely become zero and the remaining weights re-sum to one for free. The final blend of value vectors therefore uses only the current token and everything before it. The picture is a clean staircase — filled below the diagonal, empty above it — and that staircase is masked attention.

Related reveals:

Why add

- \infty

instead of multiplying by 0 after softmax?

Adding

- \infty

before softmax makes those cells exponentiate to 0 and keeps the surviving weights normalised to sum to 1 automatically.

Is the diagonal masked?

No — a token may attend to itself; only the strict upper triangle (future,

j > i

) is masked.

Why divide by

d_{k}

Dot products grow like

d_{k}

; dividing cancels that so softmax stays smooth and gradients survive.

Parent: Masked attention for autoregression · Hinglish: 4.1.11 Masked attention for autoregression (Hinglish)