BugZoo: attention-mask semantics

Attention code has its own failure modes: mask polarity, head reshaping, Q/K/V layout, and KV-cache mismatches are easy to get wrong and hard to notice from accuracy tests alone. This file focuses on the mask part, because TorchLean already has precise theorems for it.

Here is the bug-shaped PyTorch pattern we want to rule out:

# Wrong for causal attention: the polarity is flipped, so future tokens are allowed.
scores = q @ k.transpose(-2, -1) / math.sqrt(d)
mask = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1)
weights = torch.softmax(scores.masked_fill(mask == False, -torch.inf), dim=-1)

The intended PyTorch version uses true negative infinity on blocked future entries:

future = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1)
weights = torch.softmax(scores.masked_fill(future, -torch.inf), dim=-1)
assert weights[i, j] == 0.0 for all j > i

Lean's ordinary ℝ does not contain a literal $-\infty$, but mathlib does provide extended reals EReal, where ⊥ is negative infinity and EReal.exp ⊥ = 0. We record that exact $-\infty$ fact first. TorchLean's ordinary tensor softmax then uses the computationally convenient equivalent: blocked logits get zero numerator before normalization. Both views lead to the exact theorem below.

References:

• Vaswani et al., “Attention Is All You Need”, NeurIPS 2017. https://arxiv.org/abs/1706.03762
• PyTorch scaled_dot_product_attention documentation, for the runtime-style mask interface: https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
• PyTorch issue #99282, where MultiheadAttention(is_causal=True) was reported ignored when need_weights=True: https://github.com/pytorch/pytorch/issues/99282
• PyTorch issue #160064, where fully masked attention heads were reported to produce NaNs when attention weights were requested: https://github.com/pytorch/pytorch/issues/160064

noncomputable def NN.Examples.BugZoo.AttentionMask.exactMaskedLogit (score : ℝ) (allowed : Bool) : EReal

Exact extended-real masked logit: allowed entries keep their real score, blocked entries are literal $-\infty$ (⊥ : EReal).

@[simp]

theorem NN.Examples.BugZoo.AttentionMask.exactMaskedLogit_blocked (score : ℝ) : exactMaskedLogit score false = ⊥

Blocking a logit really means assigning $-\infty$ in the extended-real presentation.

@[simp]

theorem NN.Examples.BugZoo.AttentionMask.exactMaskedLogit_blocked_exp_zero (score : ℝ) : (exactMaskedLogit score false).exp = 0

The key negative-infinity softmax fact: $\exp(-\infty)=0$.

noncomputable def NN.Examples.BugZoo.AttentionMask.exactCausalMaskedScore {n : ℕ} (scores : Spec.Tensor ℝ (Spec.Shape.dim n (Spec.Shape.dim n Spec.Shape.scalar))) (i j : Fin n) : EReal

Exact extended-real causal masking of one score-matrix coordinate.

theorem NN.Examples.BugZoo.AttentionMask.exactCausalMaskedScore_future_eq_bot {n : ℕ} (scores : Spec.Tensor ℝ (Spec.Shape.dim n (Spec.Shape.dim n Spec.Shape.scalar))) (i j : Fin n) (hij : ↑i < ↑j) : exactCausalMaskedScore scores i j = ⊥

For a strict-future position, exact causal masking assigns literal $-\infty$.

This is the formal version of the PyTorch operation scores.masked_fill(future, -torch.inf) at one matrix coordinate.

theorem NN.Examples.BugZoo.AttentionMask.exactCausalMaskedScore_future_exp_zero {n : ℕ} (scores : Spec.Tensor ℝ (Spec.Shape.dim n (Spec.Shape.dim n Spec.Shape.scalar))) (i j : Fin n) (hij : ↑i < ↑j) : (exactCausalMaskedScore scores i j).exp = 0

Therefore, the strict-future numerator is exactly zero.

This is why TorchLean's attention spec writes this zero numerator directly.

theorem NN.Examples.BugZoo.AttentionMask.trueInfinityMask_future_attention_weight_zero {n : ℕ} (scores : Spec.Tensor ℝ (Spec.Shape.dim n (Spec.Shape.dim n Spec.Shape.scalar))) (i j : Fin n) (hij : ↑i < ↑j) : Spec.get2 (Spec.hardMaskedSoftmaxSpec scores (Spec.causalMask n)) i j = 0

Causal attention with true $-\infty$ masking gets the exact zero-weight theorem.

This is the statement we want for formal output-causality arguments: every strict-future key has zero attention mass for the current query row. In TorchLean this is represented by hardMaskedSoftmaxSpec, not by a finite real sentinel treated as $-\infty$.