Attention Dropout for Training LLMs

Attention scores above the main diagonal are masked by replacing them with negative infinity, using a triangle mask built with torch.triu and diagonal=1.

Programming

Why Future Positions Must Be Hidden

In causal attention, each token's attention weights are computed over a grid of scores. The grid has one row for each query position and one column for each key position. The rule of causal attention is that a position may only look at itself and at positions before it. Every score that falls above the main diagonal of this grid therefore has to be hidden from the final weights.

This article focuses on one efficient way to hide those scores, and on the reason the method works. The method has two moves. First, the scores above the diagonal are overwritten with negative infinity. Second, the softmax function is applied to every row. Keep this two-step order in mind, because the rest of the article depends on it.

What do you think happens?

Before reading on: if a score is replaced with negative infinity and then softmax runs over its row, what weight do you expect that position to receive?

Reveal answer

Answer:

Zero. Softmax treats a negative infinity entry as having zero probability, so the masked position receives no weight.

The exponential of negative infinity approaches 0, so the masked entry contributes essentially nothing to the normalized total.

Tracing the Scores Through the Pipeline

Before you learn each line, it helps to see the order in which data moves. Raw attention scores come first. The scores above the diagonal are then masked with negative infinity, which produces the masked attention scores. Finally, softmax is applied to each row, which produces the masked attention weights. The source presents this order as a data flow.

Which stage produces each set of values, and at which stage a masking mistake would first show up in the output?

The order matters because softmax needs to see the negative infinity values in order to hide them. If softmax were applied first and the zeros were inserted afterward, the rows would no longer be normalized correctly. The source describes masking before softmax as the more efficient route, and this ordering is the reason it works.

Building the Upper-Triangular Mask

The first line of the code builds the mask. It starts from a square grid of ones, with one row and one column for each position in the context. The function torch.triu keeps the upper triangle of that grid and sets the rest to zero. The argument diagonal=1 shifts the cut one step above the main diagonal, so the diagonal itself is excluded. The result is a grid with 1s strictly above the diagonal and 0s on and below it.

Which cells of the score matrix, indexed by row i and column j, get replaced with -inf, and how does diagonal=1 move the boundary of the masked region?

The diagram shows the consequence of diagonal=1. The masked region is the strict upper triangle, so the diagonal cells such as (0,0), (1,1) and (2,2) keep their scores. Each row keeps one more allowed position than the row above it, which matches the rule that a position may look at itself and earlier positions.

The diagonal argument is the place to look when a mask is wrong. Setting it to 0 would include the main diagonal in the selected region, so each position would be blocked from attending to itself. Negative values extend the selected region further below the diagonal. The mask is only correct when the boundary sits exactly one step above the diagonal, which is what diagonal=1 gives.

Applying the Mask with masked_fill

The second line uses that mask to fill in the scores. The mask is converted to a boolean with mask.bool(). Every position where the mask is true is replaced by -torch.inf. Positions where the mask is false keep their original attention score. The masking is therefore done by replacement, and the scores themselves are never recomputed.

python

import torch

scores = torch.zeros(3, 3)
mask = torch.triu(torch.ones(3, 3), diagonal=1)
masked_scores = scores.masked_fill(mask.bool(), -torch.inf)
print(masked_scores)

import torch

scores = torch.zeros(3, 3)
mask = torch.triu(torch.ones(3, 3), diagonal=1)
masked_scores = scores.masked_fill(mask.bool(), -torch.inf)
print(masked_scores)

Output (expected)

tensor([[0., -inf, -inf],
        [0., 0., -inf],
        [0., 0., 0.]])

Notes on the code

scores is a 3 by 3 grid; all zeros here so the effect of the mask is easy to see
torch.triu(torch.ones(3, 3), diagonal=1) produces 1s strictly above the diagonal
masked_fill replaces every position where the boolean mask is True with -inf

Code walkthrough

Walk through the masked scores line by line to see which positions were replaced.

1
Line 1 imports torch, the library that provides the tensor operations used here.
2
Line 2 creates a 3 by 3 grid of zeros, standing in for raw attention scores.
3
Line 3 builds the mask: ones strictly above the diagonal, zeros on and below it.
4
Line 4 calls masked_fill with the boolean mask, so the 1s in the mask select cells to overwrite with -inf.
5
Line 5 prints the result, where the upper-right cells are -inf and the rest keep their score of 0.

Why Softmax Gives a Masked Position Zero Weight

The softmax function converts its inputs into a probability distribution. When a row contains negative infinity values, softmax treats those entries as having zero probability. The source explains this with the expression $e^{-\infty}$, which approaches 0. So a position whose input is negative infinity contributes essentially nothing to the total, and the remaining positions share the probability between them.

What happens to a -inf input as it moves through exp and the normalization step, and why does its output become exactly 0?

This explains why replacing values with negative infinity is a clean way to hide them. You do not need to delete entries or change the shape of the matrix. The masked positions simply carry no weight after softmax, and the surviving positions in each row are still normalized together.

Checking That Each Row Sums to One

After softmax, every row of the attention weights should sum to 1. You can check this on the printed output by adding the values in one row, or by asking the code to do it. You should get 1 within the rounding shown. Because the masking happens before softmax and softmax normalizes each row, the rows already sum to 1, so no further normalization is necessary.

How do the unmasked weights in one row change before and after softmax, and why does the row total equal 1 with no added rescaling?

Worked example

Checking the row sums of an all-zero score matrix

A 3 by 3 score matrix is all zeros. Apply the causal mask with diagonal=1, then softmax each row, and check that every row sums to 1.

1
Mask the scores
Cells (0,1), (0,2) and (1,2) become -inf. Row 0 is [0, -inf, -inf], row 1 is [0, 0, -inf], row 2 is [0, 0, 0].
2
Apply softmax to row 0
Only one entry is finite, so its exponential is compared with the exponentials of the -inf entries, which are 0. Row 0 becomes [1, 0, 0].
3
Apply softmax to row 1
Two finite entries of 0 each have exponential 1, so the total is 2 and each gets weight 0.5. Row 1 becomes [0.5, 0.5, 0].
4
Apply softmax to row 2
All three entries are finite and equal, so each gets weight 1/3. Row 2 becomes [1/3, 1/3, 1/3].
5
Sum each row
Row 0 sums to 1, row 1 sums to 1, and row 2 sums to 1, with no extra rescaling step.

Result

Every row of the masked attention weights sums to 1, and every masked cell has weight 0.

Finding Where a Masking Bug First Appears

When an output looks wrong, the first step is to decide which stage produced the wrong value. Compare the masked scores against the weights, using the pipeline above as a map. If the upper-right cells of the masked scores are not -inf, the bug is in the mask or in the fill step. If the masked scores are correct but a masked cell still has nonzero weight, the problem is in the softmax step or the order of operations.

The row-sum check also helps locate bugs. If the row sums are not close to 1, the problem is in the softmax step or in the dimension it runs over. The code uses dim=1, which normalizes across each row. A different dimension would normalize in the wrong direction and break the per-row distribution that causal attention depends on.

Common mistake
Setting diagonal=0 in torch.triu and expecting the usual causal mask
What it looks like
torch.triu(torch.ones(3, 3), diagonal=0) includes the main diagonal, so position (0,0) is also masked.
Why it fails
The diagonal is where a position looks at itself, and the rule of causal attention allows that.
Do this instead
Use diagonal=1 so the cut sits one step above the main diagonal.
Common mistake
Applying softmax first and inserting zeros afterward
Why it fails
Softmax needs to see the negative infinity values in order to hide them. Zeros inserted after normalization leave the rows no longer normalized correctly.
Do this instead
Mask with negative infinity before softmax, then apply softmax to each row.
Common mistake
Softmax over the wrong dimension
What it looks like
Calling softmax with dim=0 normalizes down each column instead of across each row.
Why it fails
Causal attention needs each row to be its own distribution, so column normalization breaks it.
Do this instead
Use dim=1 so each row is normalized.
Common mistake
Adding an extra division to force row sums to 1 after softmax is unnecessary

Where Dropout Fits Later

Masking and dropout both affect attention weights, but they act at different points. Masking is applied before softmax and removes future positions by replacing their scores with negative infinity. Dropout is applied to the attention weights during training to reduce overfitting when training LLMs. This article stops at the masked weights, and the dropout tweak is covered in the next section.

Practicing the Mask Check

MEDIUM

Build a 4 by 4 mask with torch.triu(torch.ones(4, 4), diagonal=1). List which of the 16 cells are 1s, then apply masked_fill to a 4 by 4 zero score matrix with -inf. Finally, apply softmax with dim=1 and confirm that every row sums to 1 and that every cell above the diagonal is 0.

Hints

The 1s should sit strictly above the main diagonal, so the diagonal cells are not in the mask.
Row 0 should have three masked cells, row 1 should have two, row 2 should have one, and row 3 should have none.
Each row's weights should add up to 1, and every masked cell should be exactly 0.

Your attempt (not saved)