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.

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

Index of X in the bare LayerNorm context.

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

Index of gamma in the bare LayerNorm context.

def Proofs.Autograd.LayerNorm.idxInputBeta {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_graphMean {m n : ℕ} (xV : CtxVec (ΓLN m n)) : CtxVec.get idxMean (graphMean.evalVec xV) = valMean xV

Stage graphMean: the mean block.

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

Stage graphMeanBroadcast: the broadcast mean block.

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

Stage graphCentered: the centered block.

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

Stage graphCenteredSq: the squared centered block.

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

Stage graphVar: the variance block.

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

Stage graphVarEps: the shifted variance block.

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

Stage graphStd: the standard deviation block.

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

Stage graphInvStd: the inverse standard deviation block.

theorem Proofs.Autograd.LayerNorm.get_graphInvStdBroadcast {m n : ℕ} (ε : ℝ) (xV : CtxVec (ΓLN m n)) : CtxVec.get idxInvStdBroadcast ((graphInvStdBroadcast ε).evalVec xV) = valInvStdB xV ε

Stage graphInvStdBroadcast: the broadcast inverse standard deviation block.

theorem Proofs.Autograd.LayerNorm.get_centeredForNorm {m n : ℕ} (ε : ℝ) (xV : CtxVec (ΓLN m n)) : CtxVec.get idxCenteredForNorm ((graphInvStdBroadcast ε).evalVec xV) = valCentered xV

The centered block is still available at stage graphInvStdBroadcast.

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

Stage graphNormalized: the normalized block.

theorem Proofs.Autograd.LayerNorm.get_graphGammaBroadcast {m n : ℕ} (ε : ℝ) (xV : CtxVec (ΓLN m n)) : CtxVec.get idxGammaBroadcast ((graphGammaBroadcast ε).evalVec xV) = valGammaB xV

Stage graphGammaBroadcast: the broadcast scale block.

theorem Proofs.Autograd.LayerNorm.get_normForScale {m n : ℕ} (ε : ℝ) (xV : CtxVec (ΓLN m n)) : CtxVec.get idxNormForScale ((graphGammaBroadcast ε).evalVec xV) = valNorm xV ε

The normalized block is still available at stage graphGammaBroadcast.

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

Stage graphScaled: the scaled block.

theorem Proofs.Autograd.LayerNorm.get_graphBetaBroadcast {m n : ℕ} (ε : ℝ) (xV : CtxVec (ΓLN m n)) : CtxVec.get idxBetaBroadcast ((graphBetaBroadcast ε).evalVec xV) = valBetaB xV

Stage graphBetaBroadcast: the broadcast shift block.

theorem Proofs.Autograd.LayerNorm.get_scaledForOutput {m n : ℕ} (ε : ℝ) (xV : CtxVec (ΓLN m n)) : CtxVec.get idxScaledForOutput ((graphBetaBroadcast ε).evalVec xV) = valScaled xV ε

The scaled block is still available at stage graphBetaBroadcast.

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.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 ((graphVarEps ε).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 ((graphStd ε).evalVec xV)).ofLp i ≠ 0

The LayerNorm inv domain hypothesis follows from 0 < ε.