LayerNorm

Pointwise analytic correctness for a LayerNorm graph.

This is spec-level over ℝ. It is the proof-tape counterpart of the runtime/spec LayerNorm in Spec.layerNorm: a seqLen × embedDim tensor is normalized across the last axis, the row-wise normalizer is broadcast back over each token, and affine parameters gamma/beta are broadcast over the sequence dimension. The runtime API and typed graph path both route through that spec definition; this file proves the corresponding reverse-mode graph rule.

Because the proof graph uses the differentiable scalar nodes sqrt (max x 0) and inv, the main theorem is pointwise (GraphFDerivCorrectAt). The _of_domain variants take the two domain assumptions (positive var + ε, nonzero std) at the execution point; the main statements take only 0 < ε, since the row variance is a mean of squares and therefore nonnegative (LayerNormEval.varEps_pos_of_eps_pos, LayerNormEval.std_ne_zero_of_eps_pos). Away from the clamp kink and zero denominator, backprop is the adjoint of the Fréchet derivative. The executable Spec.layerNorm additionally clamps the raw variance before adding epsilon as a numerical guard; over exact real variance this is the same contract on the positive branch used by the proof.

The last section connects the graph to the spec. outputCLM_evalVec_layerNormGraph shows that the output block of the graph evaluation is tensorToVec (Spec.layerNorm X gamma beta) for the packed inputs, so hasFDerivAt_specLayerNormVec and backpropVec_single_eq_adjoint_specLayerNorm state differentiability and the adjoint rule for the spec function itself. The adjointness of Spec.layerNormJvp and Spec.layerNormBackward is proved in NN.Proofs.Autograd.Tape.Ops.Norm.LayerNormAdjoint. The companion LayerNormFDeriv module identifies that JVP with the actual derivative and proves that the primitive backward rule returns the same cotangents as this graph.

PyTorch correspondence / citations

• Conceptually corresponds to torch.nn.LayerNorm (without batching/running stats): normalize along the last dimension, then apply affine parameters (gamma,beta). https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html

noncomputable def Proofs.Autograd.LayerNorm.layerNormGraphFderivCorrectAtOfDomain {m n : ℕ} (ε : ℝ) (xV : CtxVec (ΓLN m n)) (hVarEpsPos : ∀ (i : Fin (VecShape m).size), 0 < (CtxVec.get idxVarEps ((layerNormPrefix6 ε).evalVec xV)).ofLp i) (hStdNe0 : ∀ (i : Fin (VecShape m).size), (CtxVec.get idxStd ((layerNormPrefix7 ε).evalVec xV)).ofLp i ≠ 0) : GraphFDerivCorrectAt (layerNormGraph ε) xV

Pointwise proof that layerNormGraph satisfies GraphFDerivCorrectAt, from explicit domain assumptions.

The hypotheses hVarEpsPos and hStdNe0 ensure that sqrt and inv are differentiable at the execution point. layerNormGraphFderivCorrectAt discharges both from 0 < ε.

noncomputable def Proofs.Autograd.LayerNorm.layerNormGraphFderivCorrectAt {m n : ℕ} (ε : ℝ) (xV : CtxVec (ΓLN m n)) (hε : 0 < ε) : GraphFDerivCorrectAt (layerNormGraph ε) xV

Pointwise proof that layerNormGraph satisfies GraphFDerivCorrectAt whenever 0 < ε.

theorem Proofs.Autograd.LayerNorm.backprop_eq_adjoint_fderiv_layerNorm_at_of_domain {m n : ℕ} (ε : ℝ) (xV : CtxVec (ΓLN m n)) (seedV : CtxVec (ΓLN m n ++ ssLayerNorm m n)) (hVarEpsPos : ∀ (i : Fin (VecShape m).size), 0 < (CtxVec.get idxVarEps ((layerNormPrefix6 ε).evalVec xV)).ofLp i) (hStdNe0 : ∀ (i : Fin (VecShape m).size), (CtxVec.get idxStd ((layerNormPrefix7 ε).evalVec xV)).ofLp i ≠ 0) : (layerNormGraph ε).backpropVec xV seedV = (ContinuousLinearMap.adjoint (fderiv ℝ (layerNormGraph ε).evalVec xV)) seedV

Pointwise end-to-end result from explicit domain assumptions: backprop equals (fderiv eval)† for layerNormGraph.

The hypotheses hVarEpsPos and hStdNe0 are the domain assumptions needed for differentiability of sqrt (after clamp) and inv at the actual execution point.

theorem Proofs.Autograd.LayerNorm.backprop_eq_adjoint_fderiv_layerNorm_at {m n : ℕ} (ε : ℝ) (xV : CtxVec (ΓLN m n)) (seedV : CtxVec (ΓLN m n ++ ssLayerNorm m n)) (hε : 0 < ε) : (layerNormGraph ε).backpropVec xV seedV = (ContinuousLinearMap.adjoint (fderiv ℝ (layerNormGraph ε).evalVec xV)) seedV

Pointwise end-to-end result: backprop equals (fderiv eval)† for layerNormGraph whenever 0 < ε.

structure Proofs.Autograd.LayerNorm.Inputs (Γ : List Spec.Shape) (m n : ℕ) : Type

LayerNorm inputs inside an arbitrary tape context.

This is the model-level interface we use once LayerNorm is no longer the root graph. For example, in a post-norm Transformer block, x is the residual stream produced by an earlier SSA node, while gamma and beta are carried parameters in the surrounding context.

• x : Idx Γ (MatShape m n)

  Sequence/residual matrix normalized across its last axis.

• gamma : Idx Γ (VecShape n)

  Affine scale vector.

• beta : Idx Γ (VecShape n)

  Affine shift vector.

noncomputable def Proofs.Autograd.LayerNorm.packInputsCLM {Γ : List Spec.Shape} {m n : ℕ} (inputs : Inputs Γ m n) : CtxVec Γ →L[ℝ] CtxVec (ΓLN m n)

Linear map that packs arbitrary-context LayerNorm inputs into the canonical context [X, gamma, beta].

noncomputable def Proofs.Autograd.LayerNorm.outputCLM {m n : ℕ} : CtxVec (ΓLN m n ++ ssLayerNorm m n) →L[ℝ] Vec (MatShape m n).size

Project the final LayerNorm output from the full canonical graph context.

noncomputable def Proofs.Autograd.LayerNorm.wholeNode {Γ : List Spec.Shape} {m n : ℕ} (inputs : Inputs Γ m n) (ε : ℝ) : Node Γ (MatShape m n)

LayerNorm as one reusable pointwise node over arbitrary context indices.

Internally this node runs the already-proved detailed LayerNorm graph. Its JVP is defined as the Fréchet derivative of that composed map at the current point, and its VJP is the adjoint of that derivative. This is exactly the block-level abstraction needed for large model proofs: the detailed LayerNorm proof remains in this file, while Transformer/GPT/ViT proofs can treat LayerNorm as a single pointwise node with explicit domain assumptions.

noncomputable def Proofs.Autograd.LayerNorm.wholeNodeFDerivCorrectAt {Γ : List Spec.Shape} {m n : ℕ} (inputs : Inputs Γ m n) (ε : ℝ) (xV : CtxVec Γ) (hε : 0 < ε) : NodeFDerivCorrectAt (wholeNode inputs ε) xV

Pointwise derivative certificate for wholeNode.

Only 0 < ε is needed: the row variance computed inside the packed LayerNorm graph is a mean of squares, so var + ε is positive and the clamped standard deviation is nonzero at every point.

The graph computes Spec.layerNorm

The packed context [X, gamma, beta] determines three spec tensors. Evaluating the LayerNorm graph on it and reading the output block gives exactly tensorToVec (Spec.layerNorm X gamma beta). Consequently the spec function is differentiable wherever the graph is, with the graph's backprop as the adjoint of its derivative.

noncomputable def Proofs.Autograd.LayerNorm.specX {m n : ℕ} (xV : CtxVec (ΓLN m n)) : TorchLean.Tensor ℝ [m, n]

The input matrix of a packed LayerNorm context, as a spec tensor.

noncomputable def Proofs.Autograd.LayerNorm.specGamma {m n : ℕ} (xV : CtxVec (ΓLN m n)) : TorchLean.Tensor ℝ [n]

The scale vector of a packed LayerNorm context, as a spec tensor.

noncomputable def Proofs.Autograd.LayerNorm.specBeta {m n : ℕ} (xV : CtxVec (ΓLN m n)) : TorchLean.Tensor ℝ [n]

The shift vector of a packed LayerNorm context, as a spec tensor.

noncomputable def Proofs.Autograd.LayerNorm.specLayerNormVec {m n : ℕ} (hm : 0 < m) (hn : 0 < n) (ε : ℝ) (xV : CtxVec (ΓLN m n)) : Vec (MatShape m n).size

Spec.layerNorm as a map on the packed context vector [X, gamma, beta].

noncomputable def Proofs.Autograd.LayerNorm.packLN {m n : ℕ} (X : TorchLean.Tensor ℝ [m, n]) (gamma beta : TorchLean.Tensor ℝ [n]) : CtxVec (ΓLN m n)

Pack three spec tensors into a LayerNorm context vector.

theorem Proofs.Autograd.LayerNorm.valX_idxMN {m n : ℕ} (xV : CtxVec (ΓLN m n)) (i : Fin m) (j : Fin n) : (valX xV).ofLp (TapeNodes.Matmul.idxMN i j) = Spec.get2 (specX xV) i j

Entries of the input block.

theorem Proofs.Autograd.LayerNorm.valGamma_apply {m n : ℕ} (xV : CtxVec (ΓLN m n)) (j : Fin n) : (valGamma xV).ofLp j = (specGamma xV).getScalar j

Entries of the scale block.

theorem Proofs.Autograd.LayerNorm.valBeta_apply {m n : ℕ} (xV : CtxVec (ΓLN m n)) (j : Fin n) : (valBeta xV).ofLp j = (specBeta xV).getScalar j

Entries of the shift block.

theorem Proofs.Autograd.LayerNorm.valMean_apply {m n : ℕ} (xV : CtxVec (ΓLN m n)) (i : Fin m) : (valMean xV).ofLp (Fin.cast ⋯ i) = Norm.rowMeanE (specX xV) i

The mean block holds the spec row means.

theorem Proofs.Autograd.LayerNorm.valMeanB_idxMN {m n : ℕ} (xV : CtxVec (ΓLN m n)) (i : Fin m) (j : Fin n) : (valMeanB xV).ofLp (TapeNodes.Matmul.idxMN i j) = Norm.rowMeanE (specX xV) i

The broadcast mean block.

theorem Proofs.Autograd.LayerNorm.valCentered_idxMN {m n : ℕ} (xV : CtxVec (ΓLN m n)) (i : Fin m) (j : Fin n) : (valCentered xV).ofLp (TapeNodes.Matmul.idxMN i j) = Spec.get2 (specX xV) i j - Norm.rowMeanE (specX xV) i

The centered block.

theorem Proofs.Autograd.LayerNorm.valVar_apply {m n : ℕ} (xV : CtxVec (ΓLN m n)) (i : Fin m) : (valVar xV).ofLp (Fin.cast ⋯ i) = Norm.rowVarE (specX xV) i

The variance block holds the spec row variances.

theorem Proofs.Autograd.LayerNorm.valInvStdB_idxMN {m n : ℕ} (xV : CtxVec (ΓLN m n)) (ε : ℝ) (i : Fin m) (j : Fin n) : (valInvStdB xV ε).ofLp (TapeNodes.Matmul.idxMN i j) = (√(max (Norm.rowVarE (specX xV) i + ε) 0))⁻¹

The broadcast inverse standard deviation block.

theorem Proofs.Autograd.LayerNorm.valY_idxMN {m n : ℕ} (xV : CtxVec (ΓLN m n)) (ε : ℝ) (i : Fin m) (j : Fin n) : (valY xV ε).ofLp (TapeNodes.Matmul.idxMN i j) = (Spec.get2 (specX xV) i j - Norm.rowMeanE (specX xV) i) * (√(max (Norm.rowVarE (specX xV) i + ε) 0))⁻¹ * (specGamma xV).getScalar j + (specBeta xV).getScalar j

The output block, entrywise.

theorem Proofs.Autograd.LayerNorm.specLayerNormVec_eq_valY {m n : ℕ} (hm : 0 < m) (hn : 0 < n) (ε : ℝ) (xV : CtxVec (ΓLN m n)) : specLayerNormVec hm hn ε xV = valY xV ε

The spec LayerNorm of the packed inputs is the closed form of the graph output.

theorem Proofs.Autograd.LayerNorm.outputCLM_evalVec_layerNormGraph {m n : ℕ} (hm : 0 < m) (hn : 0 < n) (ε : ℝ) (xV : CtxVec (ΓLN m n)) : outputCLM ((layerNormGraph ε).evalVec xV) = specLayerNormVec hm hn ε xV

Forward bridge: the output block of the LayerNorm graph is Spec.layerNorm of the packed inputs.

theorem Proofs.Autograd.LayerNorm.specLayerNormVec_eq {m n : ℕ} (hm : 0 < m) (hn : 0 < n) (ε : ℝ) : specLayerNormVec hm hn ε = fun (xV : CtxVec (ΓLN m n)) => outputCLM ((layerNormGraph ε).evalVec xV)

The spec LayerNorm map is the output projection of the graph evaluation.

theorem Proofs.Autograd.LayerNorm.hasFDerivAt_specLayerNormVec {m n : ℕ} (hm : 0 < m) (hn : 0 < n) {ε : ℝ} (hε : 0 < ε) (xV : CtxVec (ΓLN m n)) : HasFDerivAt (specLayerNormVec hm hn ε) (outputCLM ∘SL fderiv ℝ (layerNormGraph ε).evalVec xV) xV

Spec.layerNorm is differentiable in [X, gamma, beta] for 0 < ε, with derivative the output projection of the graph derivative.

theorem Proofs.Autograd.LayerNorm.differentiableAt_specLayerNormVec {m n : ℕ} (hm : 0 < m) (hn : 0 < n) {ε : ℝ} (hε : 0 < ε) (xV : CtxVec (ΓLN m n)) : DifferentiableAt ℝ (specLayerNormVec hm hn ε) xV

Spec.layerNorm is differentiable in [X, gamma, beta] for 0 < ε.

theorem Proofs.Autograd.LayerNorm.adjoint_outputCLM {m n : ℕ} (δ : Vec (MatShape m n).size) : (ContinuousLinearMap.adjoint outputCLM) δ = CtxVec.single idxY δ

The adjoint of the output projection injects a cotangent into the output block.

theorem Proofs.Autograd.LayerNorm.backpropVec_single_eq_adjoint_specLayerNorm {m n : ℕ} (hm : 0 < m) (hn : 0 < n) {ε : ℝ} (hε : 0 < ε) (xV : CtxVec (ΓLN m n)) (δ : Vec (MatShape m n).size) : (layerNormGraph ε).backpropVec xV (CtxVec.single idxY δ) = (ContinuousLinearMap.adjoint (fderiv ℝ (specLayerNormVec hm hn ε) xV)) δ

Reverse-mode bridge: backprop of the LayerNorm graph seeded on the output block is the adjoint derivative of Spec.layerNorm in [X, gamma, beta].

Packed tensors

theorem Proofs.Autograd.LayerNorm.valX_packLN {m n : ℕ} (X : TorchLean.Tensor ℝ [m, n]) (gamma beta : TorchLean.Tensor ℝ [n]) : valX (packLN X gamma beta) = tensorToVec X

The input block of a packed context.

theorem Proofs.Autograd.LayerNorm.get_idxGamma0_packLN {m n : ℕ} (X : TorchLean.Tensor ℝ [m, n]) (gamma beta : TorchLean.Tensor ℝ [n]) : CtxVec.get idxGamma0 (packLN X gamma beta) = tensorToVec gamma

The scale block of a packed context.

theorem Proofs.Autograd.LayerNorm.get_idxBeta0_packLN {m n : ℕ} (X : TorchLean.Tensor ℝ [m, n]) (gamma beta : TorchLean.Tensor ℝ [n]) : CtxVec.get idxBeta0 (packLN X gamma beta) = tensorToVec beta

The shift block of a packed context.

theorem Proofs.Autograd.LayerNorm.specX_packLN {m n : ℕ} (X : TorchLean.Tensor ℝ [m, n]) (gamma beta : TorchLean.Tensor ℝ [n]) : specX (packLN X gamma beta) = X

Packing and unpacking round-trip on the input matrix.

theorem Proofs.Autograd.LayerNorm.specGamma_packLN {m n : ℕ} (X : TorchLean.Tensor ℝ [m, n]) (gamma beta : TorchLean.Tensor ℝ [n]) : specGamma (packLN X gamma beta) = gamma

Packing and unpacking round-trip on the scale vector.

theorem Proofs.Autograd.LayerNorm.specBeta_packLN {m n : ℕ} (X : TorchLean.Tensor ℝ [m, n]) (gamma beta : TorchLean.Tensor ℝ [n]) : specBeta (packLN X gamma beta) = beta

Packing and unpacking round-trip on the shift vector.

theorem Proofs.Autograd.LayerNorm.outputCLM_evalVec_layerNormGraph_packLN {m n : ℕ} (hm : 0 < m) (hn : 0 < n) (ε : ℝ) (X : TorchLean.Tensor ℝ [m, n]) (gamma beta : TorchLean.Tensor ℝ [n]) : outputCLM ((layerNormGraph ε).evalVec (packLN X gamma beta)) = tensorToVec (Spec.layerNorm X gamma beta hm hn ε)

Forward bridge in tensor form: evaluating the LayerNorm graph on packed spec tensors and reading the output block gives Spec.layerNorm.