Spot the error: 'I apply the causal mask by setting masked positions to 0 in the attention scores before softmax.'

Spot the error·Medium·4.0 · 0·~2 min

Attempt it

Click any words you think contain an error. Click again to unmark.

Mark at least one word to submit.

The statement contains a common implementation bug. The author wants to mask out forbidden positions by setting their attention scores to zero before softmax, believing this 'zeros them out'. It does not. The confusion comes from treating the pre-softmax score and the post-softmax weight as the same thing. They are not. The weighted sum that produces the attention output uses the post-softmax weight, and softmax exponentiates the score before normalizing. Work through the arithmetic. Suppose the scores are [1, 2, 3, 4] and the author wants to mask position 3. They set position 3's score to zero, giving [1, 2, 3, 0]. Softmax exponentiates each: 2.72, 7.39, 20.09, 1.00. The sum is 31.20. Position 3's weight is 1.00 divided by 31.20, which is about 0.032: roughly 3 percent. The masked position still gets real attention. The reason is that the exponential function never returns zero for any finite input. The only value that drives e to the x to zero is negative infinity. So to truly zero out a position, you have to push its score to negative infinity (or a large finite negative like -1e9) before softmax. The correct pattern: build a mask tensor with zero at allowed positions and negative infinity at forbidden positions, add it to the scores, then run softmax. The masked entries become e to the minus billion, which underflows to zero, and the unmasked entries renormalize to sum to one. This bug is dangerous because it does not crash. The model trains, loss curves look fine, but attention leaks to masked positions and the causal property is partially violated, typically discovered only via behavioral testing or attention-pattern inspection.

The trap here is conflating the pre-softmax score with the post-softmax weight, score zero does not mean weight zero. The author sets the masked score to zero and assumes that produces zero attention weight. It does not: softmax exponentiates before normalizing, and e to the zero is one. So a masked position with score zero still earns roughly one over the sum of exponentials in weight, which for typical attention scores is a few percent. The masked position is still being attended to. The fix is to set the masked score to negative infinity (or in practice a large finite negative like -1e9 in fp32, or torch.finfo(dtype).min in fp16). Now e to the minus infinity is zero, the masked entry contributes nothing to either the numerator or the denominator of softmax, and the unmasked positions renormalize to a valid probability distribution. The pattern in code is additive: build a mask tensor with zero at allowed positions and negative infinity at forbidden positions, add it to the scores, then softmax. The insidious thing about this bug is that it doesn't crash. The model trains, the loss curves look fine, but attention leaks to forbidden positions. You typically catch it only by inspecting attention patterns or by summing weights over forbidden positions and seeing a nonzero result.

Quick reference

TL;DR

Setting masked scores to zero leaves them with about one-over-N weight after softmax; the correct mask value is negative infinity, applied before softmax.

Memory aid

Easy to grasp

Picture a vote where every candidate's score gets converted into a vote count by a rule that turns even a zero into one vote, because the conversion exponentiates first. A candidate you tried to disqualify by giving them a zero score still ends up on the ballot with one vote. To actually remove them, you have to push their score down to 'negative infinity bad', which is the only score the rule turns into zero votes. The intuition that 'zero score equals zero influence' is wrong here, and that misunderstanding is exactly the masking bug we are spotting.

Key concepts

Softmax function Causal mask Attention weights Common attention bugs

Detailed answer & concept explanation

~4 min read

Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.

Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example.

Quick scan · 2–4 min

Everything important, quickly.

5 min: domain confusion between score and weight; arithmetic for softmax of [1, 2, 3, 0]; correct additive masking with -inf; dtype-aware mask values; multiplicative masking equivalent bug; training inference mismatch signature; how to test for leakage.

Code example

python

import math
import torch

B, H, T, D = 1, 1, 4, 8
q = torch.randn(B, H, T, D)
k = torch.randn(B, H, T, D)
causal_mask = torch.tril(torch.ones(T, T))  # 1 at allowed, 0 at forbidden

# WRONG: set masked scores to 0
scores = q @ k.transpose(-2, -1) / math.sqrt(D)
bad = scores.masked_fill(causal_mask == 0, 0.0)
weights_bad = torch.softmax(bad, dim=-1)
# weights_bad has nonzero weight at future positions

# CORRECT: set masked scores to -inf
good = scores.masked_fill(causal_mask == 0, float('-inf'))
weights_good = torch.softmax(good, dim=-1)
assert weights_good[..., 0, 1:].sum().item() == 0.0  # row 0 attends only to col 0

Real-world examples

Real products, models, and research that use this idea.

PyTorch torch.nn.functional.scaled_dot_product_attention with is_causal=True handles masking correctly under all backends.
HuggingFace transformers uses torch.finfo(dtype).min as the mask value, which is roughly -3.4e38 for fp32.
A long tail of tutorial-blog attention implementations on the open web use 'scores times a binary mask' and silently leak attention, this remains one of the most common review findings in custom kernel code.
Karpathy's nanoGPT correctly uses tril plus masked_fill with float('-inf') as the canonical reference pattern.

Interviewer follow-ups

What an interviewer would ask next. Try answering before peeking at the approach.

QHow would you detect this bug in production?

QWhy prefer a large finite negative over literal float('-inf') in fused kernels?

QWhat is the equivalent bug if masking is implemented multiplicatively as scores times a 0/1 mask?

Don't say this

Red flags and common mistakes that signal junior thinking. Click to expand.

Most common mistake

Conflating pre-softmax score with post-softmax weight. A score of zero exponentiates to one, so the masked position still claims a real slice of the probability mass.

60-second cheat sheet

The night-before-the-interview bullets. Scan these on the way to the call.

References

Primary sources. Skim if you want the original framing.