Softmax Function and Attention Weights
A causal attention mask zeroes out attention weights for future positions so that each position attends only to itself and earlier positions.
Why a Language Model Must Not Look Ahead
A language model that generates text one token at a time has a strict rule: each position may be influenced only by the positions that came before it and by itself. A position must not be allowed to look ahead at tokens that have not been generated yet. The tool that enforces this rule inside the self-attention mechanism is the causal attention mask.
The mask does not compute attention from scratch. It is applied to a matrix of attention weights that has already been computed, and it changes that matrix so that the weight a position assigns to any future position becomes zero. This article follows a small matrix through that process, one stage at a time, so you can see exactly what each stage must produce.
Reading the Attention Matrix by Position
Think of the attention weights as a square table. Each row corresponds to a query position, and each column corresponds to a key position. The entry in row i and column j is the attention weight from query position i to key position j. Because a causal mask only cares about which positions come before which, the table is always square, with one row and one column for every position in the sequence.
Keeping only the entries where the column index j is less than or equal to the row index i means each position attends to itself and to earlier positions, and never to later ones. Those kept entries sit on and below the main diagonal. Everything above the diagonal is blocked. This is why the mask is called causal: the direction of the diagonal encodes the direction of time.
Building the Mask with torch.tril
The mask is built with torch.tril applied to a matrix of ones. The function keeps the lower triangle of the matrix, including the diagonal, and sets everything above the diagonal to zero. The result is a lower triangular matrix of ones and zeros.
tensor([[1., 0., 0.],
[1., 1., 0.],
[1., 1., 1.]])torch.ones(n, n)creates a 3 by 3 matrix filled with ones- torch.tril keeps the lower triangle, including the diagonal, and zeros everything above it
- Entries equal to 1 mark positions that may be attended to; entries equal to 0 mark blocked positions
What do you think happens?
Before running the code: what do the entries above the diagonal look like once torch.tril(torch.ones(3, 3)) is printed?
Reveal answer
They are all zero.
torch.tril keeps the lower triangle, including the diagonal, and zeros everything above the diagonal. The printed mask therefore has zeros in the upper-right portion, and the 1s form a lower triangle.
The Three Steps of Applying the Mask
Applying a causal mask is a three-step procedure. Each step has a specific job, and the final result depends on all three being done in order. First, the attention weights are computed with softmax. Second, the weights are multiplied by the lower triangular mask, which zeroes the future positions. Third, each row is renormalized by its sum so that every row is again a valid set of attention weights that sums to 1.
The renormalization step is written as a division by the row sum. For a weight w in row i and column j, with mask entry m, the renormalized weight is the masked weight divided by the sum of the masked weights in row i. In code, that row sum is computed with sum(dim=-1, keepdim=True).
Tracing a Three-by-Three Attention Matrix
The following example uses a small matrix of attention weights that has already been produced by softmax. Each row sums to 1. The trace below applies the mask and the renormalization, and prints the intermediate matrix after each stage so you can check each one against its required property.
tensor([[0.5000, 0.0000, 0.0000],
[0.1000, 0.6000, 0.0000],
[0.2000, 0.3000, 0.5000]])
tensor([[0.5000],
[0.7000],
[1.0000]])
tensor([[1.0000, 0.0000, 0.0000],
[0.1429, 0.8571, 0.0000],
[0.2000, 0.3000, 0.5000]])- weights * mask multiplies element by element, so the entries above the diagonal become 0
- The masked rows now sum to 0.5, 0.7, and 1.0, not to 1
- Dividing each row by its own sum restores a total of 1 per row
- Row 0 keeps only its first entry, which becomes 1 after renormalization
Checking the first row by hand
Row 0 of the softmax weights is [0.5, 0.3, 0.2]. Apply the causal mask and renormalize this row. What is the final row?
- 1Apply the mask
Row 0 of the mask is [1, 0, 0], so only the first position is visible. The masked row is [0.5, 0, 0].
- 2Compute the row sum
The sum of the masked row is 0.5 + 0 + 0 = 0.5.
- 3Renormalize
Divide each entry by 0.5. The first entry becomes 0.5 / 0.5 = 1, and the others stay at 0.
The final row is [1, 0, 0]. The single remaining weight is divided by itself and becomes 1, which matches the rule that position 0 attends only to itself.
What Renormalization Changes in a Row
After the multiplication with the mask, the elements above the diagonal are zero, which is the intended effect of the mask. That masked matrix is not yet the final attention matrix. Its rows no longer sum to 1, because the weights that were removed took part of the probability mass with them. Renormalization gives that mass back to the remaining positions in proportion to their weights, so the row becomes a valid set of attention weights again.
Row 1 shows the mechanism directly. Its future weight of 0.3 on key position 2 is removed by the mask, so the masked row is 0.1 and 0.6 with a sum of 0.7. Dividing each value by 0.7 gives approximately 0.143 and 0.857, which together sum to 1. The renormalized row is a different distribution from the original, but the relative size of the two visible weights is preserved.
Renormalization does not undo the mask. The zeros above the diagonal stay zero, because multiplying zero by any positive scale is still zero. Only the nonzero visible entries are rescaled.
Locating Where a Masked Result Goes Wrong
When a masked attention matrix looks wrong, the most useful question is which stage first produced a value that does not match the expected property. Each stage has a checkable property. The mask should be zero above the diagonal. The masked product should still be zero above the diagonal. The renormalized rows should sum to 1 while keeping those zeros. Checking these properties in order tells you which stage to inspect, instead of comparing final numbers and guessing.
| Stage | Expected property | Typical failure to look for |
|---|---|---|
| Mask from torch.tril | Entries above the diagonal are 0 | A mask built with triu or without tril, or a mask of the wrong size |
| Masked product | Entries above the diagonal are still 0 | A mask that was not applied, or a mask whose shape does not match the weights |
| Renormalized rows | Every row sums to 1 and keeps the zeros above the diagonal | Rows that sum to less than 1 because renormalization was skipped, or a row sum taken over the wrong dimension |
Diagnostic checks, run stage by stage
Using the three-by-three weights from the trace above, change the value in row 1, column 2 from 0.3 to 0.9 and rerun the masking steps. Before you run the code, predict which row sums will change. Then check whether the zeros above the diagonal are still in place after the multiplication, and whether every row still sums to 1 after renormalization.
Hints
- Column 2 is above the diagonal in row 1, so the mask removes that entry before renormalization
- Changing a value that the mask removes should not change the final renormalized row
- Check the masked product first, then the row sums, then the renormalized matrix
Mistakes Beginners Make with the Causal Mask
Common mistake Expecting the masked matrix to be the final attention matrix
What it looks likePrinting weights * mask and treating the output as finished attention weights.
Why it failsThe removed weights took part of the probability mass with them, so the rows of the masked product no longer sum to 1.
Do this insteadAlways run the renormalization step after multiplying by the mask.
Common mistake Keeping the diagonal out of the mask
What it looks likeAssuming that a position should not attend to itself and expecting zeros on the diagonal.
Why it failsA causal mask keeps the entries where column j is less than or equal to row i, so the diagonal is allowed. Each position attends to itself and earlier positions.
Do this insteadUse torch.tril, which keeps the lower triangle including the diagonal, and check that the 1s sit on and below the diagonal.
Common mistake Reading rows and columns the wrong way round
What it looks likeTreating column j as the query position and row i as the key position.
Why it failsRow i corresponds to a query position and column j corresponds to a key position. Swapping them flips the direction of the mask.
Do this insteadRead each row as the view from one query position, and each column as a key position it may or may not see.
Common mistake Checking only the final matrix instead of each stage
Key Takeaways
- A causal attention mask zeroes out attention weights for future positions, so each position attends only to itself and earlier positions.
- Applying the mask is three ordered steps: compute weights with softmax, multiply by a lower triangular mask built with torch.tril, and renormalize each row by its sum.
- Renormalization is needed because removing future weights removes probability mass; dividing by the row sum restores each row to a total of 1 while keeping the zeros.
- To find where a masked result diverges, check the properties stage by stage: zeros above the diagonal after the mask, zeros still present after the product, and row sums of 1 after renormalization.