Causal Attention in Large Language Models
A causal attention mask zeroes out attention weights for future positions so that each position attends only to itself and earlier positions.
Why Future Tokens Must Stay Hidden
A language model that generates text one token at a time must follow a strict rule: each position may be influenced only by the positions that came before it and by itself. A position must never look ahead at tokens that have not been generated yet. Inside the self-attention mechanism, the tool that enforces this rule is the causal attention mask.
The mask does not change the input and does not recompute the attention scores. It is applied to the attention weights themselves, so that the weight a position assigns to any future position becomes zero. The name causal comes from the shape of the mask: the kept entries form a staircase along the diagonal, reflecting the fact that the past causes the present but the present cannot reach into the future.
This article works on a matrix of attention weights that has already been computed. Before running any code, it helps to know what the result should look like. The goal is a matrix in which every value above the diagonal is zero and every row still sums to 1. The mask is the mechanism that zeroes those values, and the renormalization step is what restores the row sums.
Building the Lower Triangular Mask
The causal mask is built with torch.tril, which keeps the lower triangle of a matrix and sets the entries above the diagonal to zero. The trick is to apply it to a matrix of ones. The result is a lower triangular matrix of ones and zeros, with ones on and below the diagonal and zeros above it.
Read the diagram by rows. Row i is a query position, and column j is a key position. A cell is kept when column j is less than or equal to row i. So row 0 keeps only column 0, which is the position attending to itself. Row 2 keeps all three columns, because the last position is allowed to see itself and every earlier position. The diagonal runs from the top left to the bottom right, and everything above it is the future that must be hidden.
tensor([[1., 0., 0.],
[1., 1., 0.],
[1., 1., 1.]])torch.ones(3, 3)creates a 3 by 3 matrix of ones- torch.tril keeps the lower triangle, including the diagonal, and zeroes everything above it
- Each row i keeps columns 0 through i
The Three Steps of Applying a Causal 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. The first step computes the attention weights with softmax, so each row is a set of weights that sums to 1. The second step multiplies those weights by the lower triangular mask, which removes the future positions. The third step renormalizes each row so that it again sums to 1.
- Step 1, compute weights with softmax: each row of the raw attention weights is a probability distribution over the positions, so every row sums to 1 before any masking.
- Step 2, multiply by the lower triangular mask: every entry above the diagonal is multiplied by 0, so the weight from a position to any future position is removed.
- Step 3, renormalize each row: divide each remaining weight by the sum of the masked weights in its row, so the row again sums to 1 and remains a valid set of attention weights.
The second step alone is not enough. After multiplication, the elements above the diagonal are zero, which is the intended effect of the mask. But this is not the final attention matrix, because the rows no longer sum to 1. The weights that were removed took part of the probability mass with them. Renormalization gives back that mass to the visible positions in proportion to their existing weights.
In this expression, w with indices i and j is the attention weight from query position i to key position j, and m with the same indices is the mask entry, which is 1 on or below the diagonal and 0 above it. The denominator is the sum of the masked weights in row i. In code, that row sum is computed with sum(dim=-1, keepdim=True), and dividing by it is the renormalization step.
Worked Example: Masking a Three-Token Attention Matrix
Applying the causal mask to a 3 by 3 weight matrix
The softmax weights for three positions are given below, one row per query. Apply the causal mask and renormalize each row. Row 0 is [0.5, 0.3, 0.2], row 1 is [0.2, 0.5, 0.3], and row 2 is [0.1, 0.3, 0.6].
- 1Multiply by the mask
Keep only entries on or below the diagonal. Row 0 becomes [0.5, 0, 0]. Row 1 becomes [0.2, 0.5, 0]. Row 2 is unchanged at [0.1, 0.3, 0.6] because every position is visible to the last token.
- 2Check the row sums after masking
Row 0 sums to 0.5, row 1 sums to 0.7, and row 2 sums to 1.0. Rows 0 and 1 no longer sum to 1, because the removed future weights took probability mass with them.
- 3Renormalize row 0
Divide each entry by 0.5. The result is [1.0, 0, 0]. Only the first position is visible, so it is divided by itself and becomes 1.
- 4Renormalize row 1
Divide each entry by 0.7. The result is approximately [0.286, 0.714, 0]. The visible weights keep their relative proportions and now sum to 1.
- 5Confirm row 2
Row 2 already sums to 1.0, so dividing by 1.0 leaves it as [0.1, 0.3, 0.6].
Final causal weight matrix: row 0 is [1.0, 0, 0], row 1 is approximately [0.286, 0.714, 0], and row 2 is [0.1, 0.3, 0.6]. Every value above the diagonal is zero and every row sums to 1.
Checking Where a Result Diverges
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 property you can check. After the mask, the matrix should be zero above the diagonal. After the multiplication, it should still be zero above the diagonal. After renormalization, every row should sum to 1 while the zeros remain in place.
What do you think happens?
Before running torch.tril on a matrix of ones, what do the entries above the diagonal look like once the mask is printed?
Reveal answer
They are all zero.
torch.tril keeps the lower triangle, including the diagonal, and sets everything above the diagonal to zero. The printed mask has ones on and below the diagonal and zeros above it.
Mistakes When Applying a Causal Mask
Common mistake Skipping the renormalization step after masking
What it looks likeRow 1 of the masked matrix is [0.2, 0.5, 0] and is used as-is, so it sums to 0.7.
Why it failsThe removed future weights took probability mass with them, so the row no longer behaves as a set of attention weights.
Do this insteadDivide each row by its sum over the masked weights so the row sums to 1 again.
Common mistake Building the mask with the upper triangle instead of the lower
What it looks likeUsing torch.triu on a matrix of ones keeps the entries above the diagonal and zeroes the ones below it.
Why it failsThis lets each position attend to later positions, which breaks the rule that information from the future must be blocked.
Do this insteadUse torch.tril so that column j is kept only when j is less than or equal to row i.
Common mistake Checking only the final matrix
What it looks likeThe renormalized matrix is wrong, and the search starts with the softmax code.
Why it failsEach stage has its own checkable property, so the error may appear at the mask or at renormalization.
Do this insteadCheck that the mask is zero above the diagonal, that the masked product is still zero above the diagonal, and that the rows sum to 1 after renormalization.
Common mistake Treating the mask as a change to the scores before softmax, rather than to the weights after it
Practice: Build and Check a Causal Mask
Create a 4 by 4 lower triangular mask with torch.tril applied to a matrix of ones. Then take the 4 by 4 weight matrix that you expect to be causal and check three things: the entries above the diagonal are zero, the matrix is zero where the mask is zero, and each row sums to 1 after renormalization. Identify the first row and the last row where the check fails, if any.
Hints
- Row i should keep columns 0 through i.
- The renormalized row sum is the sum over the masked weights in that row.
- Row 0 should always become a single 1 followed by zeros.
Key Takeaways
- A causal attention mask zeroes out attention weights for future positions so that each position attends only to itself and earlier positions.
- The mask is built with torch.tril applied to a matrix of ones, producing a lower triangular matrix with zeros above the diagonal.
- Applying the mask takes three steps: compute the weights with softmax, multiply by the lower triangular mask, and renormalize each row so it sums to 1 again.
- Checking the property at each stage, zeros above the diagonal after masking and row sums of 1 after renormalization, is the fastest way to find where a result diverges.