Why does setting masked positions to −∞ (not 0) work correctly after softmax?

What you need to know

Softmax for one row of scores s is:

weight_j = exp(s_j) / sum over all k of exp(s_k)

A mask must do two things: give the masked position a weight of exactly 0, and leave the other weights summing to 1.

Three ways to mask, worked by hand

One query, three keys; key 3 must be hidden. Raw scores are [1.0, -0.5, 3.0].

The "score = 0" row is worse than it looks: the hidden token gets more weight than key 2, because 0 is larger than -0.5. The "after softmax" row keeps the right ratios but the output vector is scaled down by a factor that depends on the hidden score, so it changes from example to example.

Practical details

• Additive or boolean masks. Frameworks either add a mask of 0 and -inf to the scores, or take a boolean mask and fill. Both give the same result.
• Low precision. Many implementations fill with torch.finfo(dtype).min (for fp16, -65504) instead of -inf. After softmax the result is effectively the same, and it avoids some inf - inf edge cases.
• Fully masked rows. If every key in a row is -inf, softmax computes 0 / 0 and returns NaN, which spreads through the whole batch. With the finite minimum, the row becomes uniform instead — wrong but harmless if that output is ignored.

softmax([-inf, -inf, -inf])      -> [nan, nan, nan]softmax([-65504, -65504, -65504]) -> [0.333, 0.333, 0.333]

A real-life example

A speech-to-text team fine-tunes an encoder-decoder model on call-centre audio. Transcripts in a batch are padded to the longest one, and padding keys are masked. Training runs for six hours, then the loss suddenly becomes NaN and never recovers.

The cause: a few clips were pure silence with an empty transcript. Their rows in the attention were fully masked, softmax produced NaN, and one NaN gradient corrupted the weights. The fixes were to filter empty transcripts in the data pipeline, use the dtype's minimum instead of -inf in their custom layer, and add an assertion that no row of the mask is entirely blocked. Knowing why -inf works — and when it does not — is what made the six-hour failure a ten-minute diagnosis.

Follow-up questions to expect

• "Why not a large negative like -1e9 everywhere?" — In fp16 the most negative finite number is about -65,504, so -1e9 overflows to -inf; use the dtype's own minimum.
• "Does masking change the gradient?" — Masked positions get zero weight, so no gradient flows to their keys or values through that query.
• "What about padding queries?" — Their outputs are garbage but harmless, as long as the loss ignores them.