Causal masking

The whole mask at once

Apply the rule to every pair and you get a grid. Each row is a query and each column is a key. Here it is for four tokens, with 1 for allowed and 0 for blocked.

           key 0  key 1  key 2  key 3
query 0      1      0      0      0
query 1      1      1      0      0
query 2      1      1      1      0
query 3      1      1      1      1

Row $i$ may look at columns 0 up to $i$, itself included. The allowed spots fill the bottom-left half of the grid, diagonal included. A grid with this shape is called lower-triangular.

The code below builds the mask for a six-token sentence. Then it applies the mask to random scores and draws the weights as a picture.

import torch
import matplotlib.pyplot as plt

tokens = ["It", " is", " light", " and", " pack", "able"]
T = len(tokens)
mask = torch.tril(torch.ones(T, T)).bool()          # True = allowed
print(mask.int())

torch.manual_seed(0)
scores = torch.randn(T, T)
weights = torch.softmax(scores.masked_fill(~mask, float("-inf")), dim=-1)
print("row 0:", weights[0].numpy().round(2), "  (the first token can only attend to itself)")
print("row 3:", weights[3].numpy().round(2))

plt.imshow(weights.numpy(), cmap="viridis")
plt.xticks(range(T), tokens, rotation=45)
plt.yticks(range(T), tokens)
plt.xlabel("key (attended to)")
plt.ylabel("query (attending)")

Here's what each piece does.

• T is the number of tokens, here 6. torch.ones(T, T) makes a T × T grid of 1s.
• torch.tril keeps the grid's lower triangle and sets the rest to 0. The name is short for "triangle, lower".
• .bool() turns 1 into True and 0 into False. True means allowed.
• ~mask flips every True and False. Now True marks the blocked spots.
• masked_fill writes minus infinity into those blocked spots.
• torch.softmax(..., dim=-1) turns each row into weights.

In the picture, each row is a query and each column is a key. Brighter squares get more weight. Everything above the diagonal is exactly zero. Each row still sums to 1 over the tokens it's allowed to see. (torch is the course's in-browser PyTorch subset.)

Count the allowed spots row by row. Row 0 may look at 1 token. Row 1 may look at 2. Each row adds one more.

Training on every position at once

Because each position only sees its past, one pass over the sentence gives many guesses at once. Each guess is honest, because it's made from that position's own prefix: the tokens up to and including it.

Here's the flashcard sentence laid out that way. The left column is what the model may see. The right column is the token it must guess.

Six tokens give five guesses. In general, the inputs are the sentence minus its last token. The targets are the same sentence minus its first token. So the targets are the inputs shifted by one place. Training grades every guess, then averages the grades into one number, the loss.

The model always gets the real text as its input, never its own earlier guesses. This "shift by one" setup is called teacher forcing. It's why training is so efficient. A training sequence of 2,048 tokens teaches the model about 2,048 predictions for the price of one pass.

At generation time, the future doesn't exist yet. The model writes one token, adds it to the text, then writes the next. The mask makes training work under those same conditions: every guess uses only what came before.

Proving it's causal

"Causal" means cause and effect only run one way. Earlier tokens can affect later outputs. Later tokens can never affect earlier ones. Put precisely, the output at position $i$ depends only on inputs 0 up to $i$.

You can test this. Make up six token vectors and run causal attention. Then change token 4 only and run it again. Outputs at earlier positions must not move.

import math
import torch

def causal_attention(x, W_q, W_k, W_v):
    q, k, v = x @ W_q, x @ W_k, x @ W_v
    T = x.shape[0]
    scores = q @ k.T / math.sqrt(q.shape[-1])
    scores = scores.masked_fill(~torch.tril(torch.ones(T, T)).bool(), float("-inf"))
    return torch.softmax(scores, dim=-1) @ v

torch.manual_seed(1)
x = torch.randn(6, 8)
W_q, W_k, W_v = torch.randn(8, 4), torch.randn(8, 4), torch.randn(8, 4)
before = causal_attention(x, W_q, W_k, W_v)

x2 = x.clone()
x2[4] = torch.randn(8)                     # change token 4 only
after = causal_attention(x2, W_q, W_k, W_v)
changed = [(after[i] - before[i]).abs().max().item() > 1e-6 for i in range(6)]
print("output changed at positions:", [i for i, c in enumerate(changed) if c])

The function is the scaled dot-product attention from the last lesson, with the mask added. The last line compares each output before and after the change. It counts a change only if some number moved by more than 1e-6, one millionth. That small margin ignores rounding noise.

Only positions 4 and 5 change. Positions 0 to 3 can't see token 4, so nothing about them moves.

Where causal masks are (and aren't) used

• Decoder models (GPT, Claude, Llama): causal attention everywhere. That's what lets them write text left to right, one token at a time.
• Encoder models (BERT, embedding models): no causal mask. Every token sees the whole input, which suits tasks about understanding a text. A later lesson covers these.
• Padding masks: sentences in a batch often have different lengths. Shorter ones get filler tokens, called padding, to even them out. A padding mask blocks those filler tokens, so real tokens never attend to them. It's a different mask with a different job.