Evaluating the LayerNorm graph

Closed forms for every saved intermediate of layerNormGraph. Each val* definition is the value the corresponding node computes from the packed input [X, gamma, beta], and each get_* lemma identifies the block of Graph.evalVec at that stage with its closed form.

Two consequences drive the calculus: the per-row variance block is a mean of squares, so it is nonnegative, and therefore var + ε is positive and the clamped square root is nonzero as soon as 0 < ε. Those two facts are exactly the domain hypotheses of the LayerNorm graph theorem.

theorem Proofs.Autograd.LayerNorm.vecShape_size (k : ℕ) : (VecShape k).size = k

Flattened size of a length-k vector shape.

def Proofs.Autograd.LayerNorm.idxX0 {m n : ℕ} : Idx (ΓLN m n) (MatShape m n)

Index of X in the bare LayerNorm context.

def Proofs.Autograd.LayerNorm.idxGamma0 {m n : ℕ} : Idx (ΓLN m n) (VecShape n)

Index of gamma in the bare LayerNorm context.

def Proofs.Autograd.LayerNorm.idxBeta0 {m n : ℕ} : Idx (ΓLN m n) (VecShape n)

Index of beta in the bare LayerNorm context.

Closed forms of the intermediates

noncomputable def Proofs.Autograd.LayerNorm.valX {m n : ℕ} (xV : CtxVec (ΓLN m n)) : Vec (MatShape m n).size

The input matrix block.

noncomputable def Proofs.Autograd.LayerNorm.valGamma {m n : ℕ} (xV : CtxVec (ΓLN m n)) : Vec n

The scale vector block.

noncomputable def Proofs.Autograd.LayerNorm.valBeta {m n : ℕ} (xV : CtxVec (ΓLN m n)) : Vec n

The shift vector block.

noncomputable def Proofs.Autograd.LayerNorm.valMean {m n : ℕ} (xV : CtxVec (ΓLN m n)) : Vec (VecShape m).size

Row means of X.

noncomputable def Proofs.Autograd.LayerNorm.valMeanB {m n : ℕ} (xV : CtxVec (ΓLN m n)) : Vec (MatShape m n).size

Row means broadcast back to the matrix shape.

noncomputable def Proofs.Autograd.LayerNorm.valCentered {m n : ℕ} (xV : CtxVec (ΓLN m n)) : Vec (MatShape m n).size

Centered input X - mean_b.

noncomputable def Proofs.Autograd.LayerNorm.valCenteredSq {m n : ℕ} (xV : CtxVec (ΓLN m n)) : Vec (MatShape m n).size

Squared centered input.

noncomputable def Proofs.Autograd.LayerNorm.valVar {m n : ℕ} (xV : CtxVec (ΓLN m n)) : Vec (VecShape m).size

Row variances (means of the squared centered entries).

noncomputable def Proofs.Autograd.LayerNorm.valVarEps {m n : ℕ} (xV : CtxVec (ΓLN m n)) (ε : ℝ) : Vec (VecShape m).size

Row variances shifted by ε.

noncomputable def Proofs.Autograd.LayerNorm.valStd {m n : ℕ} (xV : CtxVec (ΓLN m n)) (ε : ℝ) : Vec (VecShape m).size

Clamped standard deviation sqrt (max (var + ε) 0).

noncomputable def Proofs.Autograd.LayerNorm.valInvStd {m n : ℕ} (xV : CtxVec (ΓLN m n)) (ε : ℝ) : Vec (VecShape m).size

Inverse standard deviation.

noncomputable def Proofs.Autograd.LayerNorm.valInvStdB {m n : ℕ} (xV : CtxVec (ΓLN m n)) (ε : ℝ) : Vec (MatShape m n).size

Inverse standard deviation broadcast back to the matrix shape.

noncomputable def Proofs.Autograd.LayerNorm.valNorm {m n : ℕ} (xV : CtxVec (ΓLN m n)) (ε : ℝ) : Vec (MatShape m n).size

Normalized input centered ⊙ inv_std_b.

noncomputable def Proofs.Autograd.LayerNorm.valGammaB {m n : ℕ} (xV : CtxVec (ΓLN m n)) : Vec (MatShape m n).size

Scale vector broadcast over rows.

noncomputable def Proofs.Autograd.LayerNorm.valScaled {m n : ℕ} (xV : CtxVec (ΓLN m n)) (ε : ℝ) : Vec (MatShape m n).size

Scaled normalized input.

noncomputable def Proofs.Autograd.LayerNorm.valBetaB {m n : ℕ} (xV : CtxVec (ΓLN m n)) : Vec (MatShape m n).size

Shift vector broadcast over rows.

noncomputable def Proofs.Autograd.LayerNorm.valY {m n : ℕ} (xV : CtxVec (ΓLN m n)) (ε : ℝ) : Vec (MatShape m n).size

LayerNorm output scaled + beta_b.

Stage evaluations

theorem Proofs.Autograd.LayerNorm.get_g1 {m n : ℕ} (xV : CtxVec (ΓLN m n)) : CtxVec.get idxMean (g1.evalVec xV) = valMean xV

Stage g1: the mean block.

theorem Proofs.Autograd.LayerNorm.get_g2 {m n : ℕ} (xV : CtxVec (ΓLN m n)) : CtxVec.get idxMeanB (g2.evalVec xV) = valMeanB xV

Stage g2: the broadcast mean block.

theorem Proofs.Autograd.LayerNorm.get_g3 {m n : ℕ} (xV : CtxVec (ΓLN m n)) : CtxVec.get idxCentered (g3.evalVec xV) = valCentered xV

Stage g3: the centered block.

theorem Proofs.Autograd.LayerNorm.get_g4 {m n : ℕ} (xV : CtxVec (ΓLN m n)) : CtxVec.get idxCenteredSq (g4.evalVec xV) = valCenteredSq xV

Stage g4: the squared centered block.

theorem Proofs.Autograd.LayerNorm.get_g5 {m n : ℕ} (xV : CtxVec (ΓLN m n)) : CtxVec.get idxVar (g5.evalVec xV) = valVar xV

Stage g5: the variance block.

theorem Proofs.Autograd.LayerNorm.get_prefix6 {m n : ℕ} (ε : ℝ) (xV : CtxVec (ΓLN m n)) : CtxVec.get idxVarEps ((layerNormPrefix6 ε).evalVec xV) = valVarEps xV ε

Stage layerNormPrefix6: the shifted variance block.

theorem Proofs.Autograd.LayerNorm.get_prefix7 {m n : ℕ} (ε : ℝ) (xV : CtxVec (ΓLN m n)) : CtxVec.get idxStd ((layerNormPrefix7 ε).evalVec xV) = valStd xV ε

Stage layerNormPrefix7: the standard deviation block.

theorem Proofs.Autograd.LayerNorm.get_g8 {m n : ℕ} (ε : ℝ) (xV : CtxVec (ΓLN m n)) : CtxVec.get idxInvStd ((g8 ε).evalVec xV) = valInvStd xV ε

Stage g8: the inverse standard deviation block.

theorem Proofs.Autograd.LayerNorm.get_g9 {m n : ℕ} (ε : ℝ) (xV : CtxVec (ΓLN m n)) : CtxVec.get idxInvStdB9 ((g9 ε).evalVec xV) = valInvStdB xV ε

Stage g9: the broadcast inverse standard deviation block.

theorem Proofs.Autograd.LayerNorm.get_centered9 {m n : ℕ} (ε : ℝ) (xV : CtxVec (ΓLN m n)) : CtxVec.get idxCentered9 ((g9 ε).evalVec xV) = valCentered xV

The centered block is still available at stage g9.

theorem Proofs.Autograd.LayerNorm.get_g10 {m n : ℕ} (ε : ℝ) (xV : CtxVec (ΓLN m n)) : CtxVec.get Idx.last ((g10 ε).evalVec xV) = valNorm xV ε

Stage g10: the normalized block.

theorem Proofs.Autograd.LayerNorm.get_g11 {m n : ℕ} (ε : ℝ) (xV : CtxVec (ΓLN m n)) : CtxVec.get idxGammaB11 ((g11 ε).evalVec xV) = valGammaB xV

Stage g11: the broadcast scale block.

theorem Proofs.Autograd.LayerNorm.get_norm11 {m n : ℕ} (ε : ℝ) (xV : CtxVec (ΓLN m n)) : CtxVec.get idxNorm11 ((g11 ε).evalVec xV) = valNorm xV ε

The normalized block is still available at stage g11.

theorem Proofs.Autograd.LayerNorm.get_g12 {m n : ℕ} (ε : ℝ) (xV : CtxVec (ΓLN m n)) : CtxVec.get Idx.last ((g12 ε).evalVec xV) = valScaled xV ε

Stage g12: the scaled block.

theorem Proofs.Autograd.LayerNorm.get_g13 {m n : ℕ} (ε : ℝ) (xV : CtxVec (ΓLN m n)) : CtxVec.get idxBetaB13 ((g13 ε).evalVec xV) = valBetaB xV

Stage g13: the broadcast shift block.

theorem Proofs.Autograd.LayerNorm.get_scaled13 {m n : ℕ} (ε : ℝ) (xV : CtxVec (ΓLN m n)) : CtxVec.get idxScaled13 ((g13 ε).evalVec xV) = valScaled xV ε

The scaled block is still available at stage g13.

theorem Proofs.Autograd.LayerNorm.get_idxY {m n : ℕ} (ε : ℝ) (xV : CtxVec (ΓLN m n)) : CtxVec.get idxY ((layerNormGraph ε).evalVec xV) = valY xV ε

The output block of the full LayerNorm graph.

Positivity of the denominators

theorem Proofs.Autograd.LayerNorm.rowSumCLM_apply {m n : ℕ} (x : Vec (TapeNodes.Matmul.matSize m n)) (i : Fin m) : (TapeNodes.MatrixLinear.rowSumCLM x).ofLp i = ∑ j : Fin (TapeNodes.Matmul.vecSize n), x.ofLp (finProdFinEquiv (i, j))

Coordinates of the row-sum map.

theorem Proofs.Autograd.LayerNorm.rowMeanCLM_apply {m n : ℕ} (x : Vec (TapeNodes.Matmul.matSize m n)) (i : Fin m) : (TapeNodes.MatrixLinear.rowMeanCLM x).ofLp i = 1 / ↑n * ∑ j : Fin (TapeNodes.Matmul.vecSize n), x.ofLp (finProdFinEquiv (i, j))

Coordinates of the row-mean map.

theorem Proofs.Autograd.LayerNorm.valVar_nonneg {m n : ℕ} (xV : CtxVec (ΓLN m n)) (i : Fin (VecShape m).size) : 0 ≤ (valVar xV).ofLp i

The variance block is a mean of squares, hence nonnegative.

theorem Proofs.Autograd.LayerNorm.valVarEps_pos {m n : ℕ} {ε : ℝ} (hε : 0 < ε) (xV : CtxVec (ΓLN m n)) (i : Fin (VecShape m).size) : 0 < (valVarEps xV ε).ofLp i

With positive ε, every shifted variance is positive.

theorem Proofs.Autograd.LayerNorm.valStd_pos {m n : ℕ} {ε : ℝ} (hε : 0 < ε) (xV : CtxVec (ΓLN m n)) (i : Fin (VecShape m).size) : 0 < (valStd xV ε).ofLp i

With positive ε, every clamped standard deviation is positive.

theorem Proofs.Autograd.LayerNorm.varEps_pos_of_eps_pos {m n : ℕ} {ε : ℝ} (hε : 0 < ε) (xV : CtxVec (ΓLN m n)) (i : Fin (VecShape m).size) : 0 < (CtxVec.get idxVarEps ((layerNormPrefix6 ε).evalVec xV)).ofLp i

The LayerNorm sqrt domain hypothesis follows from 0 < ε.

theorem Proofs.Autograd.LayerNorm.std_ne_zero_of_eps_pos {m n : ℕ} {ε : ℝ} (hε : 0 < ε) (xV : CtxVec (ΓLN m n)) (i : Fin (VecShape m).size) : (CtxVec.get idxStd ((layerNormPrefix7 ε).evalVec xV)).ofLp i ≠ 0

The LayerNorm inv domain hypothesis follows from 0 < ε.