LayerNorm derivatives

LayerNorm subtracts each row's mean, divides by sqrt (variance + ε), and applies the per-column scale and bias. The input, scale, and bias all vary in these theorems; epsilon is a fixed positive real number. In particular, neither the scale nor its perturbation is replaced by an all-ones vector.

The graph calculus already proves that Spec.layerNorm is differentiable. Here we identify its derivative with the actual Spec.layerNormJvp formula. We differentiate a normalized row entry, apply the product rule for the scale, and add the bias perturbation. The existing JVP/VJP pairing then identifies Spec.layerNormBackward with the adjoint derivative, including both parameter cotangents. This also identifies the three tensors returned by the primitive with the backward result of the detailed proof graph.

Positive epsilon keeps every row's standard deviation nonzero, including constant rows and rows with a single feature. All statements here concern exact real arithmetic.

theorem Proofs.Autograd.LayerNorm.rowMeanE_specX {m n : ℕ} (q : CtxVec (ΓLN m n)) (i : Fin m) : Norm.rowMeanE (specX q) i = RowNorm.rowMean (valX q) i

The tensor and vector descriptions use the same row mean.

theorem Proofs.Autograd.LayerNorm.rowVarE_specX {m n : ℕ} (q : CtxVec (ΓLN m n)) (i : Fin m) : Norm.rowVarE (specX q) i = RowNorm.rowVar (valX q) i

The tensor and vector descriptions use the same population variance.

theorem Proofs.Autograd.LayerNorm.lnInvStd_specX {m n : ℕ} (hn : 0 < n) {ε : ℝ} (hε : 0 < ε) (q : CtxVec (ΓLN m n)) (i : Fin m) : (lnInvStd (specX q) ⋯ ε).getScalar i = RowNorm.invStd (valX q) ε i

With positive epsilon the standard-deviation clamp is inactive in every row.

theorem Proofs.Autograd.LayerNorm.specLayerNormVec_idxMN {m n : ℕ} (hm : 0 < m) (hn : 0 < n) {ε : ℝ} (hε : 0 < ε) (q : CtxVec (ΓLN m n)) (i : Fin m) (j : Fin n) : (specLayerNormVec hm hn ε q).ofLp (TapeNodes.Matmul.idxMN i j) = RowNorm.nrm (valX q) ε i j * (valGamma q).ofLp j + (valBeta q).ofLp j

One output entry is the normalized input times its own scale, plus its own bias.

theorem Proofs.Autograd.LayerNorm.fderiv_specLayerNormVec_idxMN {m n : ℕ} (hm : 0 < m) (hn : 0 < n) {ε : ℝ} (hε : 0 < ε) (q dq : CtxVec (ΓLN m n)) (i : Fin m) (j : Fin n) : ((fderiv ℝ (specLayerNormVec hm hn ε) q) dq).ofLp (TapeNodes.Matmul.idxMN i j) = RowNorm.nrmJvp (valX q) ε (valX dq) i j * (valGamma q).ofLp j + RowNorm.nrm (valX q) ε i j * (valGamma dq).ofLp j + (valBeta dq).ofLp j

Entrywise differentiation includes the input tangent and both affine parameter tangents.

The normalized row differential subtracts the tangent's row mean and its component along the normalized input. The product rule supplies the additional xhat * dgamma term.

theorem Proofs.Autograd.LayerNorm.get2_layerNormJvp_eq_row_differential {m n : ℕ} (hm : 0 < m) (hn : 0 < n) {ε : ℝ} (hε : 0 < ε) (q dq : CtxVec (ΓLN m n)) (i : Fin m) (j : Fin n) : Spec.get2 (Spec.layerNormJvp hm hn (specX q) (specX dq) (specGamma q) (specGamma dq) (specBeta q) (specBeta dq) ε) i j = RowNorm.nrmJvp (valX q) ε (valX dq) i j * (valGamma q).ofLp j + RowNorm.nrm (valX q) ε i j * (valGamma dq).ofLp j + (valBeta dq).ofLp j

The primitive JVP has exactly the row differential obtained by differentiating the forward formula, with the same epsilon, scale, and bias perturbations.

theorem Proofs.Autograd.LayerNorm.fderiv_specLayerNormVec_eq_layerNormJvp {m n : ℕ} (hm : 0 < m) (hn : 0 < n) {ε : ℝ} (hε : 0 < ε) (q dq : CtxVec (ΓLN m n)) : (fderiv ℝ (specLayerNormVec hm hn ε) q) dq = tensorToVec (Spec.layerNormJvp hm hn (specX q) (specX dq) (specGamma q) (specGamma dq) (specBeta q) (specBeta dq) ε)

The Fréchet derivative of the actual LayerNorm specification is its supplied JVP.

Both arguments use the graph's [input, scale, bias] packing, so the identity can be used directly in a graph proof without changing the parameter layout.

theorem Proofs.Autograd.LayerNorm.fderiv_layerNorm_eq_layerNormJvp {m n : ℕ} (hm : 0 < m) (hn : 0 < n) {ε : ℝ} (hε : 0 < ε) (x dx : TorchLean.Tensor ℝ [m, n]) (gamma dgamma beta dbeta : TorchLean.Tensor ℝ [n]) : (fderiv ℝ (specLayerNormVec hm hn ε) (packLN x gamma beta)) (packLN dx dgamma dbeta) = tensorToVec (Spec.layerNormJvp hm hn x dx gamma dgamma beta dbeta ε)

Tensor form of the derivative bridge, with arbitrary input, scale, and bias tangents.

theorem Proofs.Autograd.LayerNorm.exists_packLN {m n : ℕ} (q : CtxVec (ΓLN m n)) : ∃ (x : TorchLean.Tensor ℝ [m, n]) (gamma : TorchLean.Tensor ℝ [n]) (beta : TorchLean.Tensor ℝ [n]), packLN x gamma beta = q

Every vector in the canonical LayerNorm context packs one input matrix and two parameter vectors. This also applies to arbitrary test directions in the adjoint proof.

theorem Proofs.Autograd.LayerNorm.inner_packLN {m n : ℕ} (x dx : TorchLean.Tensor ℝ [m, n]) (gamma dgamma beta dbeta : TorchLean.Tensor ℝ [n]) : inner ℝ (packLN x gamma beta) (packLN dx dgamma dbeta) = Spec.dot x dx + Spec.dot gamma dgamma + Spec.dot beta dbeta

The packed inner product is the sum of the input, scale, and bias tensor dot products.

theorem Proofs.Autograd.LayerNorm.adjoint_fderiv_layerNorm_eq_layerNormBackward {m n : ℕ} (hm : 0 < m) (hn : 0 < n) {ε : ℝ} (hε : 0 < ε) (x gradOutput : TorchLean.Tensor ℝ [m, n]) (gamma beta : TorchLean.Tensor ℝ [n]) : (ContinuousLinearMap.adjoint (fderiv ℝ (specLayerNormVec hm hn ε) (packLN x gamma beta))) (tensorToVec gradOutput) = have gradients := Spec.layerNormBackward hm hn x gamma gradOutput ε; packLN gradients.inputGradient gradients.scaleGradient gradients.biasGradient

The primitive backward rule is the adjoint of the actual LayerNorm derivative.

The three output tensors occupy the original input, scale, and bias slots. In particular, the scale and bias cotangents sum over rows because those parameters are shared by every row.

theorem Proofs.Autograd.LayerNorm.backpropVec_single_eq_layerNormBackward {m n : ℕ} (hm : 0 < m) (hn : 0 < n) {ε : ℝ} (hε : 0 < ε) (x gradOutput : TorchLean.Tensor ℝ [m, n]) (gamma beta : TorchLean.Tensor ℝ [n]) : (layerNormGraph ε).backpropVec (packLN x gamma beta) (CtxVec.single idxY (tensorToVec gradOutput)) = have gradients := Spec.layerNormBackward hm hn x gamma gradOutput ε; packLN gradients.inputGradient gradients.scaleGradient gradients.biasGradient

Backpropagating through the detailed LayerNorm graph returns exactly the same three cotangents as the primitive backward rule.