The derivative of Boolean-masked softmax

The mask is fixed while the scores vary. A row with an allowed key is a normalized sum of exponentials over those keys; a row without one is the constant zero function. These are separate cases of the same derivative formula, because a blocked key has exactly zero weight.

The implementation subtracts a maximum before exponentiating. We first cancel that common factor over the reals, then differentiate the resulting quotient. In particular, ties for the maximum do not require a hypothesis: the maximum is an implementation detail of the stable evaluation, and the function being differentiated is smooth even where that maximum changes.

noncomputable def Proofs.Autograd.HardMaskedSoftmax.numerator {n : ℕ} (mask : TorchLean.Tensor Bool [n]) (x : Vec n) (i : Fin n) : ℝ

An allowed logit contributes its exponential; a blocked logit contributes zero.

noncomputable def Proofs.Autograd.HardMaskedSoftmax.denominator {n : ℕ} (mask : TorchLean.Tensor Bool [n]) (x : Vec n) : ℝ

The sum is positive exactly when the fixed mask has an allowed key.

noncomputable def Proofs.Autograd.HardMaskedSoftmax.weights {n : ℕ} (mask : TorchLean.Tensor Bool [n]) (x : Vec n) : Vec n

The real formula for the weights, including the zero row when every key is blocked.

theorem Proofs.Autograd.HardMaskedSoftmax.getScalar_hardMaskedSoftmaxVecSpec {n : ℕ} (scores : TorchLean.Tensor ℝ [n]) (mask : TorchLean.Tensor Bool [n]) (i : Fin n) : (Spec.hardMaskedSoftmaxVecSpec scores mask).getScalar i = numerator mask (getScalarE scores) i / denominator mask (getScalarE scores)

Removing the stable shift gives the ordinary exponential quotient on the allowed keys.

theorem Proofs.Autograd.HardMaskedSoftmax.spec_eq_weights {n : ℕ} (scores : TorchLean.Tensor ℝ [n]) (mask : TorchLean.Tensor Bool [n]) : getScalarE (Spec.hardMaskedSoftmaxVecSpec scores mask) = weights mask (getScalarE scores)

The implementation and the smooth real presentation have the same vector coordinates.

noncomputable def Proofs.Autograd.HardMaskedSoftmax.derivative {n : ℕ} (mask : TorchLean.Tensor Bool [n]) (x : Vec n) : Vec n →L[ℝ] Vec n

The row Jacobian is diag(w) - w wᵀ. This definition accepts the zero row as well.

@[simp]

theorem Proofs.Autograd.HardMaskedSoftmax.derivative_apply {n : ℕ} (mask : TorchLean.Tensor Bool [n]) (x dx : Vec n) (i : Fin n) : ((derivative mask x) dx).ofLp i = (weights mask x).ofLp i * (dx.ofLp i - ∑ j : Fin n, (weights mask x).ofLp j * dx.ofLp j)

Reading one coordinate exposes the weighted centering performed by the spec JVP and VJP.

theorem Proofs.Autograd.HardMaskedSoftmax.denominator_pos {n : ℕ} (mask : TorchLean.Tensor Bool [n]) (x : Vec n) (h : ∃ (i : Fin n), mask.getScalar i = true) : 0 < denominator mask x

A nonempty allowed set makes the exponential denominator strictly positive.

theorem Proofs.Autograd.HardMaskedSoftmax.hasFDerivAt_weights {n : ℕ} (mask : TorchLean.Tensor Bool [n]) (x : Vec n) : HasFDerivAt (weights mask) (derivative mask x) x

Hard-masked softmax is differentiable for every score vector and every fixed Boolean mask.

An all-false row is handled before applying the quotient rule. Consequently the proof neither divides by its zero denominator nor invents a derivative at a masked infinity.

theorem Proofs.Autograd.HardMaskedSoftmax.hasFDerivAt_spec {n : ℕ} (mask : TorchLean.Tensor Bool [n]) (x : Vec n) : HasFDerivAt (fun (y : Vec n) => getScalarE (Spec.hardMaskedSoftmaxVecSpec (ofFnE y) mask)) (derivative mask x) x

This derivative theorem is about the actual stable tensor implementation.

theorem Proofs.Autograd.HardMaskedSoftmax.backward_eq_derivative {n : ℕ} (scores gradient : TorchLean.Tensor ℝ [n]) (mask : TorchLean.Tensor Bool [n]) : getScalarE (Spec.softmaxBackwardFromWeightsSpec (Spec.hardMaskedSoftmaxVecSpec scores mask) gradient) = (derivative mask (getScalarE scores)) (getScalarE gradient)

The implementation's weighted-centering helper evaluates the derivative itself.

theorem Proofs.Autograd.HardMaskedSoftmax.inner_derivative {n : ℕ} (mask : TorchLean.Tensor Bool [n]) (x dx gradient : Vec n) : inner ℝ ((derivative mask x) dx) gradient = inner ℝ dx ((derivative mask x) gradient)

The row derivative is self-adjoint, including rows whose weights are all zero.

theorem Proofs.Autograd.HardMaskedSoftmax.backward_eq_adjoint_fderiv {n : ℕ} (scores gradient : TorchLean.Tensor ℝ [n]) (mask : TorchLean.Tensor Bool [n]) : getScalarE (Spec.softmaxBackwardFromWeightsSpec (Spec.hardMaskedSoftmaxVecSpec scores mask) gradient) = (ContinuousLinearMap.adjoint (fderiv ℝ (fun (x : Vec n) => getScalarE (Spec.hardMaskedSoftmaxVecSpec (ofFnE x) mask)) (getScalarE scores))) (getScalarE gradient)

The concrete backward helper is the adjoint of the derivative of the concrete forward.

This is stronger than an algebraic JVP/VJP pairing: hasFDerivAt_spec supplies the analytic derivative, and the equality below identifies the helper called by attention with its adjoint.