PyTorch The Causal Mask
A causal mask is a triangular table of ones and zeros that marks which positions a reader may see. The mask turns scores for forbidden positions into negative infinity before the softmax, so those positions receive exactly zero weight.
tril draws the triangle
torch.tril returns the lower triangle of its input and replaces the rest with zeros. Row zero keeps a single one, and the last row keeps all four.
import torch
mask = torch.tril(torch.ones(4, 4))
print(mask.tolist())Negative infinity becomes zero weight
Two identical score rows receive the mask treatment. The first row keeps only its first score and takes weight 1.0 there. The second row spreads weight over its two visible scores.
import torch
from torch.nn import functional as F
scores = torch.tensor([[2.0, 1.0, 0.0, -1.0],
[2.0, 1.0, 0.0, -1.0]])
mask = torch.tril(torch.ones(4, 4))
masked = scores.masked_fill(mask[:2] == 0, float("-inf"))
print(masked.tolist())
print([round(v, 2) for v in F.softmax(masked, dim=-1)[0].tolist()])
print([round(v, 2) for v in F.softmax(masked, dim=-1)[1].tolist()])Example
import torch
from torch import nn
from torch.nn import functional as F
batch_size = 32
block_size = 64
n_embd = 128
n_head = 4
n_layer = 4
head_size = n_embd // n_head
torch.manual_seed(42)
x = torch.randn(1, 4, n_embd)
query = nn.Linear(n_embd, head_size, bias=False)
key = nn.Linear(n_embd, head_size, bias=False)
q = query(x)
k = key(x)
scores = q @ k.transpose(-2, -1) * head_size ** -0.5
mask = torch.tril(torch.ones(block_size, block_size))
scores = scores.masked_fill(mask[:4, :4] == 0, float("-inf"))
weights = F.softmax(scores, dim=-1)[0]
print("mask corner:")
for row in mask[:4, :4].tolist():
print(" ".join(f"{v:.0f}" for v in row))
print("masked weights:")
for row in weights.tolist():
print(" ".join(f"{v:.2f}" for v in row))
print("row sums:", [round(v, 2) for v in weights.sum(dim=-1).tolist()])
print("future weights all zero:", bool((weights.triu(diagonal=1) == 0).all()))Key points
- torch.tril keeps the lower triangle of a table and zeroes the rest
- The diagonal stays visible, so a position can always read itself
- Negative infinity is written float('-inf') in PyTorch code
- A softmax gives a score of negative infinity a weight of exactly zero
- masked_fill writes a value wherever a condition table is true
- The mask depends only on position, never on the text itself
In this example, the seeded window from the scaling lesson produces the same scaled scores as before. A triangular mask built at the full block size is cut down to the window size of four. masked_fill then writes negative infinity into every score above the diagonal. The softmax turns those entries into exact zeros, and the printed grid shows the staircase shape that results. The first position reads only itself, and the last position reads all four. The row sums still print as 1.0, and the final check confirms that every future weight is zero.
Try it Yourself
Edit the code, then press Run.
Running is simulated · Python Playground
Exercise
Why are masked scores set to negative infinity instead of zero before the softmax?