Rectangular attention with a Boolean mask

Query and key lengths are independent. The mask is fixed, and every row may have its own allowed keys, including no allowed keys at all. The statements below concern the actual attention JVP and backward specifications, with their existing scale and zero-row conventions.

The pairing theorem identifies the transpose relation between these two implementations. The analytic derivative of the row normalization comes from HardMaskedSoftmax; identifying the whole attention JVP with the derivative of the forward map additionally requires composing the matrix and row maps.

theorem Proofs.Autograd.HardMaskedAttention.unstack_hardMaskedSoftmaxSpec {nQ nK : ℕ} (scores : TorchLean.Tensor ℝ [nQ, nK]) (mask : TorchLean.Tensor Bool [nQ, nK]) (i : Fin nQ) : (Spec.hardMaskedSoftmaxSpec scores mask).unstack i = Spec.hardMaskedSoftmaxVecSpec (scores.unstack i) (mask.unstack i)

Each query row is normalized against its own row of the Boolean mask.

theorem Proofs.Autograd.HardMaskedAttention.unstack_softmaxBackwardFromWeightsSpec {nQ nK : ℕ} (weights gradient : TorchLean.Tensor ℝ [nQ, nK]) (i : Fin nQ) : (Spec.softmaxBackwardFromWeightsSpec weights gradient).unstack i = Spec.softmaxBackwardFromWeightsSpec (weights.unstack i) (gradient.unstack i)

The weighted-centering helper acts independently on each query row.

theorem Proofs.Autograd.HardMaskedAttention.dot_backward_comm {nQ nK : ℕ} (scores dx gradient : TorchLean.Tensor ℝ [nQ, nK]) (mask : TorchLean.Tensor Bool [nQ, nK]) : Spec.dot (Spec.softmaxBackwardFromWeightsSpec (Spec.hardMaskedSoftmaxSpec scores mask) dx) gradient = Spec.dot dx (Spec.softmaxBackwardFromWeightsSpec (Spec.hardMaskedSoftmaxSpec scores mask) gradient)

The concrete masked row helper is symmetric for rectangular score matrices.

There is no requirement that a row contain an allowed key. The all-false case is already part of the row derivative theorem, so summing its inner-product identity over query rows preserves it.

theorem Proofs.Autograd.HardMaskedAttention.jvp_backward_pairing {nQ nK d : ℕ} {hQ : nQ ≠ 0} {hK : nK ≠ 0} (ctx : Spec.AttentionContext ℝ nQ nK d hQ hK) (mask : TorchLean.Tensor Bool [nQ, nK]) (hmask : ctx.mask = some mask) (dQ : TorchLean.Tensor ℝ [nQ, d]) (dK dV : TorchLean.Tensor ℝ [nK, d]) (gradient : TorchLean.Tensor ℝ [nQ, d]) : Spec.dot (Spec.scaledDotProductAttentionJvp ctx dQ dK dV) gradient = Spec.dot dQ (Spec.scaledDotProductAttentionBackward ctx gradient).1 + (Spec.dot dK (Spec.scaledDotProductAttentionBackward ctx gradient).2.1 + Spec.dot dV (Spec.scaledDotProductAttentionBackward ctx gradient).2.2)

The spec SDPA JVP and backward formulas have the exact Frobenius transpose relation.

The three terms keep query, key, and value gradients separate. In particular, the key term uses the rectangular transposition dictated by Q Kᵀ; no equality of query and key lengths is assumed. This algebraic identity is stated separately from differentiability of the forward map.