LayerNorm JVP and VJP adjointness

Spec.layerNormJvp is the forward-mode rule and Spec.layerNormBackward the reverse-mode rule shipped with Spec.layerNorm. This file proves that they are adjoint: pairing an input and parameter tangent with the JVP gives the same scalar as pairing the upstream gradient with the three backward outputs. LayerNormFDeriv identifies this JVP with the derivative of the actual specification for positive epsilon, then uses this pairing to prove that the backward rule is the adjoint derivative. The pairing itself is algebraic and allows any real epsilon.

The proof is entrywise. Both spec functions are unfolded to explicit matrix forms (rfl), every entry is rewritten as a real expression, and the per-row identity row_adjoint finishes with ring.

Entry lemmas

theorem Proofs.Autograd.LayerNorm.getScalar_mulSpec' {k : ℕ} (a b : TorchLean.Tensor ℝ [k]) (i : Fin k) : (a.mulSpec b).getScalar i = a.getScalar i * b.getScalar i

Entries of a pointwise vector product.

theorem Proofs.Autograd.LayerNorm.getScalar_reduceSum_zero {m n : ℕ} (x : TorchLean.Tensor ℝ [m, n]) (h : Spec.Shape.NonemptyAxis 0 (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (j : Fin n) : (TorchLean.Tensor.reduceSum 0 x h).getScalar j = ∑ i : Fin m, Spec.get2 x i j

Column sums along the sequence axis.

Explicit matrix forms

@[reducible, inline]

abbrev Proofs.Autograd.LayerNorm.bAS {m n : ℕ} (v : TorchLean.Tensor ℝ [m]) : TorchLean.Tensor ℝ [m, n]

Row broadcast of a length-m vector to an m × n matrix.

noncomputable def Proofs.Autograd.LayerNorm.lnInvStd {m n : ℕ} (x : TorchLean.Tensor ℝ [m, n]) (h1 : Spec.Shape.NonemptyAxis 1 (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (ε : ℝ) : TorchLean.Tensor ℝ [m]

1 / std, as computed by the spec.

noncomputable def Proofs.Autograd.LayerNorm.layerNormJvpMat {m n : ℕ} (x tangent : TorchLean.Tensor ℝ [m, n]) (gamma dgamma dbeta : TorchLean.Tensor ℝ [n]) (h1 : Spec.Shape.NonemptyAxis 1 (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (ε : ℝ) : TorchLean.Tensor ℝ [m, n]

Spec.layerNormJvp with its axis, broadcast, and cast evidence spelled out.

noncomputable def Proofs.Autograd.LayerNorm.layerNormBackwardMat {m n : ℕ} (x : TorchLean.Tensor ℝ [m, n]) (gamma : TorchLean.Tensor ℝ [n]) (grad : TorchLean.Tensor ℝ [m, n]) (h1 : Spec.Shape.NonemptyAxis 1 (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (h0 : Spec.Shape.NonemptyAxis 0 (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (ε : ℝ) : Spec.NormalizationGradients ℝ [m, n] [n]

Spec.layerNormBackward with its axis, broadcast, and cast evidence spelled out.

theorem Proofs.Autograd.LayerNorm.layerNormJvp_eq_mat {m n : ℕ} (hm : 0 < m) (hn : 0 < n) (x tangent : TorchLean.Tensor ℝ [m, n]) (gamma dgamma beta dbeta : TorchLean.Tensor ℝ [n]) (ε : ℝ) : Spec.layerNormJvp hm hn x tangent gamma dgamma beta dbeta ε = layerNormJvpMat x tangent gamma dgamma dbeta ⋯ ε

Spec.layerNormJvp is layerNormJvpMat by unfolding.

theorem Proofs.Autograd.LayerNorm.layerNormBackward_eq_mat {m n : ℕ} (hm : 0 < m) (hn : 0 < n) (x : TorchLean.Tensor ℝ [m, n]) (gamma : TorchLean.Tensor ℝ [n]) (grad : TorchLean.Tensor ℝ [m, n]) (ε : ℝ) : Spec.layerNormBackward hm hn x gamma grad ε = layerNormBackwardMat x gamma grad ⋯ ⋯ ε

Spec.layerNormBackward is layerNormBackwardMat by unfolding.

The per-row identity

theorem Proofs.Autograd.LayerNorm.row_adjoint {n : ℕ} (t g xh γ dγ dβ : Fin n → ℝ) (s : ℝ) : ∑ j : Fin n, (s * (t j - (∑ k : Fin n, t k) / ↑n - xh j * ((∑ k : Fin n, t k * xh k) / ↑n)) * γ j + xh j * dγ j + dβ j) * g j = ∑ j : Fin n, t j * (s * (g j * γ j - (∑ k : Fin n, g k * γ k) / ↑n - xh j * ((∑ k : Fin n, g k * γ k * xh k) / ↑n))) + ∑ j : Fin n, dγ j * (g j * xh j) + ∑ j : Fin n, dβ j * g j

Scalar adjointness of the LayerNorm rules on one row: s is the inverse standard deviation, xh the normalized row, γ the per-column scale.

Adjointness

theorem Proofs.Autograd.LayerNorm.get2_layerNormJvpMat {m n : ℕ} (x tangent : TorchLean.Tensor ℝ [m, n]) (gamma dgamma dbeta : TorchLean.Tensor ℝ [n]) (h1 : Spec.Shape.NonemptyAxis 1 (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (ε : ℝ) (i : Fin m) (j : Fin n) : Spec.get2 (layerNormJvpMat x tangent gamma dgamma dbeta h1 ε) i j = (lnInvStd x h1 ε).getScalar i * (Spec.get2 tangent i j - (∑ k : Fin n, Spec.get2 tangent i k) / ↑n - Spec.get2 (Norm.lnCentered x h1) i j * (lnInvStd x h1 ε).getScalar i * ((∑ k : Fin n, Spec.get2 tangent i k * (Spec.get2 (Norm.lnCentered x h1) i k * (lnInvStd x h1 ε).getScalar i)) / ↑n)) * gamma.getScalar j + Spec.get2 (Norm.lnCentered x h1) i j * (lnInvStd x h1 ε).getScalar i * dgamma.getScalar j + dbeta.getScalar j

Entries of the JVP output.

theorem Proofs.Autograd.LayerNorm.get2_inputGradient {m n : ℕ} (x : TorchLean.Tensor ℝ [m, n]) (gamma : TorchLean.Tensor ℝ [n]) (grad : TorchLean.Tensor ℝ [m, n]) (h1 : Spec.Shape.NonemptyAxis 1 (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (h0 : Spec.Shape.NonemptyAxis 0 (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (ε : ℝ) (i : Fin m) (j : Fin n) : Spec.get2 (layerNormBackwardMat x gamma grad h1 h0 ε).inputGradient i j = (lnInvStd x h1 ε).getScalar i * (Spec.get2 grad i j * gamma.getScalar j - (∑ k : Fin n, Spec.get2 grad i k * gamma.getScalar k) / ↑n - Spec.get2 (Norm.lnCentered x h1) i j / (Norm.lnStd x h1 ε).getScalar i * ((∑ k : Fin n, Spec.get2 grad i k * gamma.getScalar k * (Spec.get2 (Norm.lnCentered x h1) i k / (Norm.lnStd x h1 ε).getScalar i)) / ↑n))

Entries of the backward input gradient.

theorem Proofs.Autograd.LayerNorm.getScalar_scaleGradient {m n : ℕ} (x : TorchLean.Tensor ℝ [m, n]) (gamma : TorchLean.Tensor ℝ [n]) (grad : TorchLean.Tensor ℝ [m, n]) (h1 : Spec.Shape.NonemptyAxis 1 (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (h0 : Spec.Shape.NonemptyAxis 0 (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (ε : ℝ) (j : Fin n) : (layerNormBackwardMat x gamma grad h1 h0 ε).scaleGradient.getScalar j = ∑ i : Fin m, Spec.get2 grad i j * (Spec.get2 (Norm.lnCentered x h1) i j / (Norm.lnStd x h1 ε).getScalar i)

Entries of the backward scale gradient.

theorem Proofs.Autograd.LayerNorm.getScalar_biasGradient {m n : ℕ} (x : TorchLean.Tensor ℝ [m, n]) (gamma : TorchLean.Tensor ℝ [n]) (grad : TorchLean.Tensor ℝ [m, n]) (h1 : Spec.Shape.NonemptyAxis 1 (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (h0 : Spec.Shape.NonemptyAxis 0 (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (ε : ℝ) (j : Fin n) : (layerNormBackwardMat x gamma grad h1 h0 ε).biasGradient.getScalar j = ∑ i : Fin m, Spec.get2 grad i j

Entries of the backward bias gradient.

theorem Proofs.Autograd.LayerNorm.centered_mul_invStd {m n : ℕ} (x : TorchLean.Tensor ℝ [m, n]) (h1 : Spec.Shape.NonemptyAxis 1 (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (ε : ℝ) (i : Fin m) (j : Fin n) : Spec.get2 (Norm.lnCentered x h1) i j * (lnInvStd x h1 ε).getScalar i = Spec.get2 (Norm.lnCentered x h1) i j / (Norm.lnStd x h1 ε).getScalar i

1 / std times centered is centered / std.

theorem Proofs.Autograd.LayerNorm.layerNormJvp_layerNormBackward_adjoint {m n : ℕ} (hm : 0 < m) (hn : 0 < n) (x tangent gradOutput : TorchLean.Tensor ℝ [m, n]) (gamma dgamma beta dbeta : TorchLean.Tensor ℝ [n]) (ε : ℝ) : Spec.dot (Spec.layerNormJvp hm hn x tangent gamma dgamma beta dbeta ε) gradOutput = Spec.dot tangent (Spec.layerNormBackward hm hn x gamma gradOutput ε).inputGradient + Spec.dot dgamma (Spec.layerNormBackward hm hn x gamma gradOutput ε).scaleGradient + Spec.dot dbeta (Spec.layerNormBackward hm hn x gamma gradOutput ε).biasGradient

The LayerNorm reverse rule is adjoint to its forward differential.

Pairing the input tangent and both parameter tangents with Spec.layerNormJvp gives the same scalar as pairing the upstream gradient with the three outputs of Spec.layerNormBackward.