Attention Masking › Worked Example: Tracing a Masked Forward Pass

To make the mechanics concrete, let's trace through a complete masked attention computation step by step with small numerical values you can verify by hand.

We will use a toy sequence of three tokens: "I", "love", "AI". We want to apply causal masking so that during training, each position predicts the next token using only past context.

Step 1: Set up the scores. Suppose we have the following raw attention scores (these would normally come from the dot product of queries and keys, but we will use fixed values for clarity):

$$S = \begin{bmatrix} 1.0 & 0.5 & 0.2 \\ 0.8 & 1.2 & 0.3 \\ 0.6 & 0.9 & 1.5 \end{bmatrix}$$

where $S_{ij}$ is the score from query position $i$ to key position $j$.

Step 2: Construct the causal mask. For a 3-token sequence, the causal mask is:

$$M = \begin{bmatrix} 0 & -\infty & -\infty \\ 0 & 0 & -\infty \\ 0 & 0 & 0 \end{bmatrix}$$

In practice we use $-10^9$ instead of $-\infty$.

Step 3: Add the mask to the scores. The masked scores are $S + M$:

$$S + M = \begin{bmatrix} 1.0 & -\infty & -\infty \\ 0.8 & 1.2 & -\infty \\ 0.6 & 0.9 & 1.5 \end{bmatrix}$$

The upper triangle has been driven to $-\infty$, making those positions invisible.

Step 4: Apply softmax row by row. Recall that softmax divides each exponential by the sum of exponentials in the row. For row 0 (position "I"):

$$\begin{aligned} \text{softmax}([1.0, -\infty, -\infty]) &= \left[\frac{e^{1.0}}{e^{1.0} + e^{-\infty} + e^{-\infty}}, \frac{e^{-\infty}}{\ldots}, \frac{e^{-\infty}}{\ldots}\right] \\ &\approx [1.00, 0.00, 0.00] \end{aligned}$$

Position "I" attends only to itself: weight 1.0. For row 1 (position "love"):

$$\begin{aligned} \text{softmax}([0.8, 1.2, -\infty]) &= \left[\frac{e^{0.8}}{e^{0.8} + e^{1.2}}, \frac{e^{1.2}}{e^{0.8} + e^{1.2}}, 0\right] \\ &= \left[\frac{2.225}{2.225 + 3.320}, \frac{3.320}{2.225 + 3.320}, 0\right] \\ &\approx [0.40, 0.60, 0.00] \end{aligned}$$

Position "love" splits its attention between "I" (40%) and itself (60%). For row 2 (position "AI"):

$$\text{softmax}([0.6, 0.9, 1.5]) \approx [0.17, 0.23, 0.60]$$

Position "AI" can attend to all three positions and distributes attention based on score magnitude, attending most strongly to itself (1.5 is the highest score).

Step 5: Interpret the result. The final attention weight matrix is:

$$A = \begin{bmatrix} 1.00 & 0.00 & 0.00 \\ 0.40 & 0.60 & 0.00 \\ 0.17 & 0.23 & 0.60 \end{bmatrix}$$

Each row is a valid probability distribution summing to 1.0. The upper triangle is exactly zero. The model is now constrained to make predictions using only the causal context it should have. When predicting the word that follows "love," the model sees only "I" and "love"; when predicting the word after "AI," it sees all three. This is precisely the progressive left-to-right context structure that makes autoregressive language modeling work.