PyTorch Attention Masks
An attention mask is a table of numbers added to the scores before softmax runs. A masked score sinks to minus infinity, so softmax gives that position a weight of exactly zero.
Hide one padding position
The last position here is padding, a filler token with no meaning of its own. Adding minus infinity to its score removes it from the mix.
import torch
scores = torch.tensor([2.0, 1.0, 0.0, 1.0])
mask = torch.tensor([0.0, 0.0, 0.0, float("-inf")])
print("Masked scores:", scores + mask)Build a causal mask
triu is the function that keeps the upper triangle of a table. diagonal=1 starts that triangle one step right, so every position keeps its own score and loses the future ones.
causal = torch.triu(torch.full((4, 4), float("-inf")), diagonal=1)
print(causal)Masked weights hit exactly zero
The first row may only see itself once the causal mask lands. Its weight row puts everything on position 1 and exactly zero elsewhere.
row = torch.tensor([2.0, 1.0, 0.0, 1.0])
weights = torch.softmax(row + causal[0], dim=0)
print("First row weights:", weights)Example
import torch
scores = torch.tensor([[2.0, 1.0, 0.0, 1.0],
[1.0, 2.0, 1.0, 0.0],
[0.0, 1.0, 2.0, 1.0],
[1.0, 0.0, 1.0, 2.0]])
padding_mask = torch.zeros(4, 4)
padding_mask[:, 3] = float("-inf")
causal_mask = torch.triu(torch.full((4, 4), float("-inf")), diagonal=1)
padding_weights = torch.softmax(scores + padding_mask, dim=1)
causal_weights = torch.softmax(scores + causal_mask, dim=1)
print("Padding-masked weights:")
print(padding_weights)
print("Causal-masked weights:")
print(causal_weights)Key points
- Padding: filler tokens added so that short sequences match the batch length
- Padding mask: a mask that hides the filler positions from every query
- Causal mask: a mask that hides later positions, so a step sees only its past
- Masks are added to the scores before softmax, never after it
- A minus infinity score becomes a weight of exactly 0 after softmax
- In short, a mask edits the scores so that forbidden positions win no weight
In this example, one scores table is masked in the two standard ways. The padding mask fills the last column with minus infinity, because token 4 is padding here. Every weight in that column prints as 0.0000 after softmax. The causal mask fills everything above the diagonal with minus infinity. Each row then keeps weight only on itself and the positions before it. The first causal row is 1.0000 followed by zeros, since position 1 has no past.
Note
The library module in the next lesson accepts its mask in this same additive form
Try it Yourself
Edit the code, then press Run.
Running is simulated · Python Playground
Exercise
When is a mask applied in scaled dot-product attention?