Lesson 2 - The Causal Mask

Welcome to The Causal Mask#

In the previous lesson you framed language modeling as an autoregressive task: at every position the model predicts the next token from the tokens that came before it. That objective has a strict rule hidden inside it. When your mini-GPT is asked to predict the character at position $t+1$, it is only allowed to look at positions $0$ through $t$. The self-attention head you built in Module 2 breaks that rule flagrantly: it lets every position attend to every position, future included. Position $t$ would get to see position $t+1$ — the very answer it is supposed to predict.

This lesson fixes that with one small, precise change: the causal mask. You set the attention scores for future positions to negative infinity before the softmax, so those positions get essentially zero weight and each query attends only to itself and the past. The result is a lower-triangular attention matrix, and it is the single mechanism that turns a generic attention block into a GPT-style decoder. You will implement it, watch the attention matrix go triangular, confirm position 0 attends only to itself, and then prove the masked backward pass is correct with a numerical gradient check.

By the end of this lesson, you will be able to:

• Explain why unmasked self-attention “cheats” at next-token prediction
• Build a causal mask with np.triu and apply it to the scores before softmax
• Read a lower-triangular attention matrix and verify each row still sums to 1
• Show that position 0 attends only to itself with weight exactly 1.0
• Re-apply the mask in the backward pass so masked positions receive zero gradient
• Gradient-check the causal attention backward pass to a tiny max relative error

You should be comfortable with the self-attention forward and backward passes from Module 2 and the softmax Jacobian. Everything here is one masking step layered on top of machinery you already trust. Let’s begin.

Why Unmasked Attention Cheats#

Recall the shape of scaled dot-product attention. For a single head, the scores are $S = Q K^\top / \sqrt{hs}$, a (T, T) matrix in which entry $S_{ij}$ is how much query position $i$ matches key position $j$. Softmax over each row turns those scores into weights, and each output is a weighted average of the values: position $i$ reads a blend of every position’s value, weighted by $A_{ij}$.

The trouble is the word every. Nothing in that computation stops $A_{ij}$ from being large when $j > i$. Position $i$ can pour most of its attention onto positions that come after it in the sequence. For a task like classification, where the whole sequence is available at once, that is exactly what you want — full bidirectional context. But our objective is next-token prediction. To train the model to predict the character at position $i+1$, we feed it positions $0 \ldots i$ and ask for position $i+1$. If attention at position $i$ is allowed to look ahead to position $i+1$, the model can simply copy the answer. It would score perfectly during training and produce nonsense at generation time, when the future genuinely does not exist yet.

So autoregressive modeling imposes a hard constraint: position $i$ may attend only to positions $j \le i$. We need a way to enforce that inside the attention head without changing anything else.

The Fix: Negative Infinity Before Softmax#

The elegant trick is to poison the forbidden scores before the softmax. Softmax exponentiates its inputs, and $e^{-\infty} = 0$. So if we set every future score $S_{ij}$ (those with $j > i$) to $-\infty$ before the softmax runs, the softmax assigns those positions a weight of exactly zero, and the remaining visible positions renormalize to sum to 1 on their own. In practice we use a large finite negative number like $-10^9$ rather than literal infinity, because np.exp(-1e9) underflows cleanly to 0.0 while -np.inf can produce nan if an entire row were ever masked.

The mask itself is a boolean (T, T) matrix that is True exactly where $j > i$ — the strict upper triangle, the future. NumPy’s np.triu with k=1 builds it in one line:

import numpy as np

T = 5
causal_mask = np.triu(np.ones((T, T), dtype=bool), k=1)  # True strictly ABOVE the diagonal
print(causal_mask)

[[False  True  True  True  True]
 [False False  True  True  True]
 [False False False  True  True]
 [False False False False  True]
 [False False False False False]]

Read row $i$: the True entries are the future positions $j > i$ that row $i$ must not see. Row 0 must not see positions 1, 2, 3, 4; row 4 sees everything, because for the last position the whole sequence is the past. Applying it to the scores is a single np.where:

scores = np.where(causal_mask, -1e9, scores)   # block the future before softmax

Here is the full causal forward pass. It is the same scaled dot-product attention from Module 2 with exactly two new lines — building the mask and applying it — inserted between the scores and the softmax.

import numpy as np

np.random.seed(42)
B, T, C, hs = 1, 5, 8, 8      # one sequence of length 5, a single head

def softmax(x, axis=-1):
    x = x - x.max(axis=axis, keepdims=True)   # stable: subtract row max
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)

# larger init (*0.5) so the attention weights differentiate visibly
X  = np.random.randn(B, T, C) * 0.5
Wq = np.random.randn(C, hs) * 0.5
Wk = np.random.randn(C, hs) * 0.5
Wv = np.random.randn(C, hs) * 0.5

causal_mask = np.triu(np.ones((T, T), dtype=bool), k=1)   # True above the diagonal

def causal_forward(X, Wq, Wk, Wv):
    Q = X @ Wq                                    # (B, T, hs)
    K = X @ Wk                                    # (B, T, hs)
    V = X @ Wv                                    # (B, T, hs)
    S = Q @ K.transpose(0, 2, 1) / np.sqrt(hs)    # (B, T, T) raw scores
    S = np.where(causal_mask, -1e9, S)            # THE MASK: block the future
    A = softmax(S, axis=-1)                       # (B, T, T) lower-triangular
    out = A @ V                                   # (B, T, hs)
    cache = (X, Wq, Wk, Wv, Q, K, V, A)
    return out, A, cache

out, A, cache = causal_forward(X, Wq, Wk, Wv)
print("out:", out.shape, " A:", A.shape, " dtype:", A.dtype)

out: (1, 5, 8)  A: (1, 5, 8)  dtype: float64

Nothing about the shapes changed — the head still maps (1, 5, 8) to (1, 5, 8). What changed is which positions each output is allowed to mix. We used the demo init scale * 0.5 (not the training scale * 0.02) precisely so the surviving weights come out uneven and worth inspecting, as the course conventions recommend for any lesson that looks at attention patterns.

The Attention Matrix Goes Triangular#

Now look at the attention matrix the masked head produces. This is the payoff — the (T, T) weights should be zero everywhere above the diagonal and a valid probability distribution over the visible past on and below it.

np.set_printoptions(precision=3, suppress=True)
print("attention matrix A[0] (rows = query i, cols = key j):")
print(A[0])
print()
print("row sums:", A[0].sum(axis=1))
print("max entry strictly above diagonal:", A[0][causal_mask].max())
print("lower-triangular?", np.allclose(A[0][causal_mask], 0.0))
print("position 0 attends only to itself? A[0,0,0] =", round(float(A[0, 0, 0]), 6))

attention matrix A[0] (rows = query i, cols = key j):
[[1.    0.    0.    0.    0.   ]
 [0.79  0.21  0.    0.    0.   ]
 [0.269 0.254 0.477 0.    0.   ]
 [0.291 0.217 0.204 0.289 0.   ]
 [0.271 0.144 0.279 0.104 0.202]]

row sums: [1. 1. 1. 1. 1.]
max entry strictly above diagonal: 0.0
lower-triangular? True
position 0 attends only to itself? A[0,0,0] = 1.0

Everything the mask promised is visible in those numbers. Every entry above the diagonal is exactly 0.0, so the matrix is lower-triangular. Every row still sums to 1.0, because softmax renormalized the surviving scores. And the most telling row is the first one: position 0 has no past to attend to — the only visible position is itself — so A[0, 0, 0] is exactly 1.0 and the rest of its row is zero. Position 0 always attends only to itself in a causal model, no matter what the scores are, because it is the one position with nothing behind it.

Compare this to the unmasked version of the very same head. Without the mask, row 0 comes out as [0.118, 0.321, 0.199, 0.200, 0.162] — it spreads its attention across all five positions, including the four that lie in its future. That row is exactly the cheating we set out to forbid: a next-token predictor at position 0 leaning heavily on positions 1 through 4, the answers it is meant to produce. The mask replaces that leaky row with a clean [1, 0, 0, 0, 0].

Backward Through the Mask#

If the forward pass sets future scores to a constant $-10^9$, what happens to their gradient in the backward pass? The masked entries were replaced by a constant before softmax, so they no longer depend on $Q$ or $K$ at all. A constant has zero derivative. That means the gradient flowing back into every masked score position must be exactly zero — those positions contributed nothing to the output, so they must receive nothing back. If we let a nonzero gradient leak into a masked score, we would be training $Q$ and $K$ to adjust a score that the forward pass throws away, which is both wrong and, in a gradient check, immediately detectable.

Concretely, the forward did S = np.where(mask, -1e9, S). The derivative of that np.where with respect to the original score is 1 where the mask is False (the score passed through) and 0 where the mask is True (the score was overwritten by a constant). So in the backward pass we re-apply the same mask, zeroing dS at every masked position, right after the softmax backward and before pushing the gradient into dQuery and dKey:

dS = np.where(causal_mask, 0.0, dS)   # masked cells get ZERO gradient

In practice the softmax backward already drives these entries nearly to zero on its own, because the masked weights A[i, j] are ~0 and the softmax gradient A * (...) is proportional to A. But “nearly” is not “exactly”: the -1e9 scores leave a residue on the order of 1e-9. Re-applying the mask makes those entries exactly zero, which is both mathematically correct and what keeps the gradient check crisp. Here is the full masked backward pass — it is the Module 2 attention backward with that single extra line inserted.

def loss_from_out(out):
    return np.sum(out * G)          # scalar loss; dL/dout is exactly G

G = np.random.randn(B, T, hs)       # fixed output directions, defines the loss
out, A, cache = causal_forward(X, Wq, Wk, Wv)
dout = G

def causal_backward(dout, cache):
    X, Wq, Wk, Wv, Q, K, V, A = cache
    dA = dout @ V.transpose(0, 2, 1)                          # (B, T, T)
    dValue = A.transpose(0, 2, 1) @ dout                     # (B, T, hs)
    dS = A * (dA - np.sum(dA * A, axis=-1, keepdims=True))   # softmax backward
    dS = np.where(causal_mask, 0.0, dS)                      # re-apply the mask
    dS_scaled = dS / np.sqrt(hs)                             # (B, T, T)
    dQuery = dS_scaled @ K                                   # (B, T, hs)
    dKey = dS_scaled.transpose(0, 2, 1) @ Q                  # (B, T, hs)
    dWq = np.sum(X.transpose(0, 2, 1) @ dQuery, axis=0)      # (C, hs)
    dWk = np.sum(X.transpose(0, 2, 1) @ dKey, axis=0)        # (C, hs)
    dWv = np.sum(X.transpose(0, 2, 1) @ dValue, axis=0)      # (C, hs)
    dX = dQuery @ Wq.T + dKey @ Wk.T + dValue @ Wv.T         # (B, T, C)
    return dWq, dWk, dWv, dX

dWq, dWk, dWv, dX = causal_backward(dout, cache)
print("dWq:", dWq.shape, " dWk:", dWk.shape, " dWv:", dWv.shape, " dX:", dX.shape)

dWq: (8, 8)  dWk: (8, 8)  dWv: (8, 8)  dX: (1, 5, 8)

Every gradient has the shape of the quantity it belongs to — dWq, dWk, dWv match their (C, hs) weights, and dX matches the (B, T, C) input. Notice we kept the course naming rule: the gradients of Q, K, V are dQuery, dKey, dValue, never a bare two-letter form. Shapes lining up proves the plumbing, not the numbers. For the numbers, we gradient-check.

Proving It: A Gradient Check Through the Mask#

The test is the same central finite difference you used for the unmasked head: nudge one parameter entry up by $\epsilon$, down by $\epsilon$, see how the scalar loss moves, and compare that estimate to the analytic gradient. If the masked backward pass is correct — including the zeroing of dS — the numerical and analytic gradients will agree to within floating-point noise for every entry of every parameter and for X.

def numerical_grad(param):
    eps = 1e-5
    grad = np.zeros_like(param)
    it = np.nditer(param, flags=["multi_index"])
    while not it.finished:
        idx = it.multi_index
        original = param[idx]
        param[idx] = original + eps
        Lp = loss_from_out(causal_forward(X, Wq, Wk, Wv)[0])
        param[idx] = original - eps
        Lm = loss_from_out(causal_forward(X, Wq, Wk, Wv)[0])
        param[idx] = original          # always restore
        grad[idx] = (Lp - Lm) / (2 * eps)
        it.iternext()
    return grad

def max_rel_err(analytic, numeric):
    denom = np.maximum(np.abs(analytic), np.abs(numeric)) + 1e-12
    return np.max(np.abs(analytic - numeric) / denom)

print("max relative error (analytic vs numerical), causal backward:")
print("  dWq:", max_rel_err(dWq, numerical_grad(Wq)))
print("  dWk:", max_rel_err(dWk, numerical_grad(Wk)))
print("  dWv:", max_rel_err(dWv, numerical_grad(Wv)))
print("  dX :", max_rel_err(dX,  numerical_grad(X)))

max relative error (analytic vs numerical), causal backward:
  dWq: 5.730509395118713e-09
  dWk: 1.579986735460768e-08
  dWv: 1.6717096446951037e-09
  dX : 1.1043587614076531e-09

Every error sits between $10^{-8}$ and $10^{-9}$, far below the $10^{-4}$ threshold that separates a correct gradient from a buggy one. That is the proof: the masked forward pass and the masked backward pass are consistent, right down to the gradient the numerical check sees flowing (or not flowing) through the future positions. The finite-difference loop naturally respects the mask, because when you perturb a Q or K entry that only feeds masked scores, the -1e9 swamps the perturbation and the loss does not move — so the numerical gradient there is zero too, exactly matching the analytic zero. Run the whole script a second time and the numbers are byte-for-byte identical, because np.random.seed(42) fixes every draw and there is no nondeterminism anywhere in the pipeline.

Practice Exercises#

Try these before checking the hints. They reuse causal_forward, causal_backward, numerical_grad, loss_from_out, causal_mask, and the tensors defined above.

Exercise 1: See the Cheating for Yourself#

Run the head without the mask (delete the np.where line, or set causal=False) and print row 0 of the resulting attention matrix. Confirm it spreads weight across future positions 1 through 4, then compare it to the masked row 0 of [1, 0, 0, 0, 0]. Explain in a comment why the unmasked row 0 is “cheating” for next-token prediction.

# Your code here: compute an unmasked A, print A_unmasked[0, 0], compare to A[0, 0]

Exercise 2: Confirm Every Row Sums to 1#

The mask removes probability mass from the future, and softmax is supposed to redistribute it over the visible past. Verify this directly: assert that every row of the masked A[0] sums to 1 to floating-point tolerance, and that the number of nonzero entries in row $i$ is exactly $i + 1$.

# Your code here: check A[0].sum(axis=1) and count nonzeros per row

Exercise 3: Check That Masked Scores Get Zero Gradient#

Inside causal_backward, after the line dS = np.where(causal_mask, 0.0, dS), the strict upper triangle of dS must be all zeros. Add a print (or return dS) and confirm dS[0][causal_mask] is exactly zero. Then remove the re-masking line and confirm the entries are no longer exactly zero, only tiny.

# Your code here: expose dS from causal_backward and inspect dS[0][causal_mask]

Summary#

You added the one mechanism that turns generic self-attention into autoregressive, GPT-style attention: the causal mask. Let’s review.

Key Concepts#

The Problem

• Plain self-attention lets position $i$ attend to every position, future included
• For next-token prediction that is cheating: position $i$ could read position $i+1$, the answer

The Mask (Forward)

• Build the future mask with np.triu(np.ones((T, T), bool), k=1) — True strictly above the diagonal
• Apply it to the scores before softmax: S = np.where(mask, -1e9, S)
• softmax(-1e9) ≈ 0, so future weights vanish and the visible past renormalizes to sum to 1
• The attention matrix becomes lower-triangular; position 0 attends only to itself with weight 1.0

The Mask (Backward)

• Masked scores were replaced by a constant, so their gradient must be exactly zero
• Re-apply the mask in backward: dS = np.where(mask, 0.0, dS) before pushing into dQuery, dKey
• Keep the course naming rule: gradients of Q, K, V are dQuery, dKey, dValue

Proving It

• Real max relative errors ran 5.7e-9 (dWq) to 1.1e-9 (dX), all far below the 1e-4 bar
• Identical on a second run — seed 42 makes the whole pipeline deterministic
• Mask before softmax, re-mask in backward: both bugs show up in a gradient check

Why This Matters#

The causal mask is a tiny piece of code with an outsized role: it is the entire reason a transformer can generate text one token at a time. Because position $i$ never sees the future during training, the model learns to predict each token from only what precedes it — which is exactly the situation it faces at generation time, when the future has not been written yet. Train without the mask and the model would score beautifully on your training set and then fall apart the moment you ask it to actually generate, because the future it leaned on no longer exists. The mask aligns training with inference.

There is a practical elegance here too. The same lower-triangular structure means a single forward pass computes the next-token prediction at every position at once — position 0’s prediction, position 1’s, all the way to position $T-1$ — each using only its own past. That is why transformers train so efficiently: one masked attention call yields $T$ supervised predictions in parallel, not one. When you assemble the full mini-GPT and watch its loss drop, the causal mask is what makes that parallel, leak-free training possible.

Continue Building Your Skills#

You now have a causal attention head: masked forward, masked backward, and gradient-checked to nine decimal places. What is still missing is the surrounding architecture. In the next lesson you will stack causal transformer blocks into the decoder-only design that GPT uses — token and position embeddings feeding a stack of blocks, each block wrapping this causal multi-head attention in the residual connections, layer norms, and feed-forward network you built in Module 5. You will see how the mask lives inside every attention sub-layer of every block, so the whole tower stays autoregressive from input to output. That decoder-only stack is the body of your mini-GPT; the language-modeling head and the next-token loss that sit on top of it come right after.