Causal Masking: How Decoder-Only Models Cheat-Proof Their Training

~12 min · causal-mask, decoder

Level 0Token

0 XP0/94 lessons0/10 achievements

0/120 XP to next level120 XP to go0% complete

In a decoder-only model trained on next-token prediction, position t must predict the token at position t+1. During training we want to compute losses for all positions in parallel — but the model must not be allowed to peek at future tokens during the forward pass. The fix is simple and load-bearing: a causal mask.

Concretely, we add a tensor of -∞ values to the upper triangle of the attention score matrix before softmax. After softmax, those entries become exactly 0 — the corresponding positions are completely ignored. Position 3 can attend to positions 0, 1, 2, 3 but not 4, 5, .... The model gets the parallelism of computing all positions at once while behaving, position by position, as if it had only seen the past.

Code

Causal mask construction·python

import torch

def causal_mask(seq_len, device='cpu'):
    # Returns (seq_len, seq_len) bool tensor
    # True at positions (i, j) where j > i (future)
    return torch.triu(torch.ones(seq_len, seq_len, device=device), diagonal=1).bool()

# Inside attention:
# scores = Q @ K.transpose(-2, -1) / d_k**0.5
# scores = scores.masked_fill(causal_mask(seq_len), float('-inf'))
# weights = torch.softmax(scores, dim=-1)
# After softmax, future positions have exactly 0 weight.

Inspect the resulting attention weights·python

seq_len = 6
mask = causal_mask(seq_len)
print(mask.int())
# tensor([[0, 1, 1, 1, 1, 1],
#         [0, 0, 1, 1, 1, 1],
#         [0, 0, 0, 1, 1, 1],
#         [0, 0, 0, 0, 1, 1],
#         [0, 0, 0, 0, 0, 1],
#         [0, 0, 0, 0, 0, 0]])
# After softmax with -inf substitution, position 0 attends only to itself,
# position 1 to {0, 1}, ... position 5 attends to all six.

External links

Exercise

Implement causal_mask. Plug it into the multi-head attention from earlier. Train a tiny model on a copy task. Then accidentally remove the mask and retrain. Observe how the model loss drops faster (it's cheating) and how generation quality collapses at inference time. Document the lesson.

Progress

Progress is local-only — sign in to sync across devices.

← PreviousSelf-Attention vs Cross-Attention Next →Attention Patterns: What Heads Actually Learn

Spotted a bug or have feedback on this page?Report an Issue

Comments 0

🔔 Reply notifications (sign in)

No comments yet — be the first.