Hard masking with an all-true mask

Spec.scaledDotProductAttention normalizes scores through one of two code paths: plain axis softmax when ctx.mask = none, and Spec.hardMaskedSoftmaxSpec when a Boolean mask is supplied. This file proves that a mask allowing every position selects exactly the unmasked weights, so the two branches agree wherever both apply. The statement is over ℝ, where the max used by the stable softmax shift is the order maximum compared by the hard-mask row scan.

theorem NN.Proofs.Models.Attention.hardMaskedSoftmaxVecSpec_allTrue {n : ℕ} (scores : TorchLean.Tensor ℝ [n]) : Spec.hardMaskedSoftmaxVecSpec scores (TorchLean.Tensor.dim fun (x : Fin n) => TorchLean.Tensor.scalar true) = Activation.softmaxVecSpec scores

Hard-masked softmax with every entry allowed is the stable softmax kernel.

theorem NN.Proofs.Models.Attention.hardMaskedSoftmaxSpec_allTrueMask {nQ nK : ℕ} (scores : TorchLean.Tensor ℝ [nQ, nK]) : Spec.hardMaskedSoftmaxSpec scores (Spec.allTrueMask nQ nK) = Activation.softmaxSpec 1 scores

Hard-masked softmax with the all-true mask is the unmasked axis-1 softmax.

theorem NN.Proofs.Models.Attention.scaledDotProductAttention_allTrueMask {nQ nK dModel : ℕ} {h1 : nQ ≠ 0} {h2 : nK ≠ 0} (ctx : Spec.AttentionContext ℝ nQ nK dModel h1 h2) : Spec.scaledDotProductAttention { Q := ctx.Q, K := ctx.K, V := ctx.V, mask := some (Spec.allTrueMask nQ nK) } = Spec.scaledDotProductAttention { Q := ctx.Q, K := ctx.K, V := ctx.V, mask := none }

Supplying the all-true mask to scaled dot-product attention is the same as supplying no mask.

theorem NN.Proofs.Models.Attention.scaledDotProductAttentionBackward_allTrueMask {nQ nK dModel : ℕ} {h1 : nQ ≠ 0} {h2 : nK ≠ 0} (ctx : Spec.AttentionContext ℝ nQ nK dModel h1 h2) (dOut : TorchLean.Tensor ℝ [nQ, dModel]) : Spec.scaledDotProductAttentionBackward { Q := ctx.Q, K := ctx.K, V := ctx.V, mask := some (Spec.allTrueMask nQ nK) } dOut = Spec.scaledDotProductAttentionBackward { Q := ctx.Q, K := ctx.K, V := ctx.V, mask := none } dOut

The all-true mask leaves the attention backward pass unchanged.

theorem NN.Proofs.Models.Attention.scaledDotProductAttentionJvp_allTrueMask {nQ nK dModel : ℕ} {h1 : nQ ≠ 0} {h2 : nK ≠ 0} (ctx : Spec.AttentionContext ℝ nQ nK dModel h1 h2) (dQ : TorchLean.Tensor ℝ [nQ, dModel]) (dK dV : TorchLean.Tensor ℝ [nK, dModel]) : Spec.scaledDotProductAttentionJvp { Q := ctx.Q, K := ctx.K, V := ctx.V, mask := some (Spec.allTrueMask nQ nK) } dQ dK dV = Spec.scaledDotProductAttentionJvp { Q := ctx.Q, K := ctx.K, V := ctx.V, mask := none } dQ dK dV

The all-true mask leaves the attention forward-mode derivative unchanged.