LayerNorm graph

The explicit SSA graph used by the LayerNorm proofs: a seqLen × embedDim input is normalized across its last axis, the row statistics are broadcast back over each token, and the affine parameters gamma/beta are broadcast over the sequence axis. The graph is written as an explicit snoc chain so that every prefix has a name; the calculus proofs in NN.Proofs.Autograd.Tape.Ops.Norm.LayerNorm and the evaluation lemmas in NN.Proofs.Autograd.Tape.Ops.Norm.LayerNormEval refer to these prefixes directly.

The nodes are the shape-generic tape nodes: row mean, row and column broadcast, elementwise arithmetic, the clamped square root sqrt (max x 0), and inversion.

@[reducible, inline]

abbrev Proofs.Autograd.LayerNorm.MatShape (m n : ℕ) : Spec.Shape

Matrix shape m×n.

@[reducible, inline]

abbrev Proofs.Autograd.LayerNorm.VecShape (k : ℕ) : Spec.Shape

Rank-one tensor shape k.

@[reducible, inline]

abbrev Proofs.Autograd.LayerNorm.ΓLN (m n : ℕ) : List Spec.Shape

Input context shapes: [X, gamma, beta] for layer norm over the last axis.

@[reducible, inline]

abbrev Proofs.Autograd.LayerNorm.ssPrefix6 (m n : ℕ) : List Spec.Shape

First 6 intermediates in the LayerNorm computation (up to var_eps).

@[reducible, inline]

abbrev Proofs.Autograd.LayerNorm.ssPrefix7 (m n : ℕ) : List Spec.Shape

Prefix intermediates up to std (adds one more vector).

@[reducible, inline]

abbrev Proofs.Autograd.LayerNorm.ssLayerNorm (m n : ℕ) : List Spec.Shape

Full list of intermediates for the LayerNorm graph in this file.

def Proofs.Autograd.LayerNorm.idxX {m n : ℕ} {ss : List Spec.Shape} : Idx (ΓLN m n ++ ss) (MatShape m n)

Index of the input matrix X in the base LayerNorm context ΓLN m n ++ ss.

def Proofs.Autograd.LayerNorm.idxGamma {m n : ℕ} {ss : List Spec.Shape} : Idx (ΓLN m n ++ ss) (VecShape n)

Index of the scale vector gamma in the base LayerNorm context ΓLN m n ++ ss.

def Proofs.Autograd.LayerNorm.idxBeta {m n : ℕ} {ss : List Spec.Shape} : Idx (ΓLN m n ++ ss) (VecShape n)

Index of the shift vector beta in the base LayerNorm context ΓLN m n ++ ss.

noncomputable def Proofs.Autograd.LayerNorm.nodeMean {m n : ℕ} : Node (ΓLN m n) (VecShape m)

Mean over the last axis: mean : ℝ^{m×n} → ℝ^{m}.

noncomputable def Proofs.Autograd.LayerNorm.g1 {m n : ℕ} : Graph (ΓLN m n) [VecShape m]

Graph prefix producing [mean].

def Proofs.Autograd.LayerNorm.idxMean {m n : ℕ} : Idx (ΓLN m n ++ [VecShape m]) (VecShape m)

Index of mean in the extended context ΓLN ++ [mean].

noncomputable def Proofs.Autograd.LayerNorm.nodeMeanB {m n : ℕ} : Node (ΓLN m n ++ [VecShape m]) (MatShape m n)

Broadcast mean back to m×n (row-wise).

noncomputable def Proofs.Autograd.LayerNorm.g2 {m n : ℕ} : Graph (ΓLN m n) [VecShape m, MatShape m n]

Graph prefix producing [mean, mean_b].

def Proofs.Autograd.LayerNorm.idxMeanB {m n : ℕ} : Idx (ΓLN m n ++ [VecShape m, MatShape m n]) (MatShape m n)

Index of mean_b in ΓLN ++ [mean, mean_b].

noncomputable def Proofs.Autograd.LayerNorm.nodeCentered {m n : ℕ} : Node (ΓLN m n ++ [VecShape m, MatShape m n]) (MatShape m n)

Center: centered := X - mean_b.

noncomputable def Proofs.Autograd.LayerNorm.g3 {m n : ℕ} : Graph (ΓLN m n) [VecShape m, MatShape m n, MatShape m n]

Graph prefix producing [mean, mean_b, centered].

def Proofs.Autograd.LayerNorm.idxCentered {m n : ℕ} : Idx (ΓLN m n ++ [VecShape m, MatShape m n, MatShape m n]) (MatShape m n)

Index of centered in ΓLN ++ [mean, mean_b, centered].

noncomputable def Proofs.Autograd.LayerNorm.nodeCenteredSq {m n : ℕ} : Node (ΓLN m n ++ [VecShape m, MatShape m n, MatShape m n]) (MatShape m n)

Square centered: centered_sq := centered ⊙ centered.

noncomputable def Proofs.Autograd.LayerNorm.g4 {m n : ℕ} : Graph (ΓLN m n) [VecShape m, MatShape m n, MatShape m n, MatShape m n]

Graph prefix producing [mean, mean_b, centered, centered_sq].

def Proofs.Autograd.LayerNorm.idxCenteredSq {m n : ℕ} : Idx (ΓLN m n ++ [VecShape m, MatShape m n, MatShape m n, MatShape m n]) (MatShape m n)

Index of centered_sq in the extended context.

noncomputable def Proofs.Autograd.LayerNorm.nodeVar {m n : ℕ} : Node (ΓLN m n ++ [VecShape m, MatShape m n, MatShape m n, MatShape m n]) (VecShape m)

Variance per row: var := mean(centered_sq) producing a length-m vector.

noncomputable def Proofs.Autograd.LayerNorm.g5 {m n : ℕ} : Graph (ΓLN m n) [VecShape m, MatShape m n, MatShape m n, MatShape m n, VecShape m]

Graph prefix producing [mean, mean_b, centered, centered_sq, var].

def Proofs.Autograd.LayerNorm.idxVar {m n : ℕ} : Idx (ΓLN m n ++ [VecShape m, MatShape m n, MatShape m n, MatShape m n, VecShape m]) (VecShape m)

Index of var in the extended context.

noncomputable def Proofs.Autograd.LayerNorm.nodeVarEps {m n : ℕ} (ε : ℝ) : Node (ΓLN m n ++ [VecShape m, MatShape m n, MatShape m n, MatShape m n, VecShape m]) (VecShape m)

Add epsilon: var_eps := var + ε.

noncomputable def Proofs.Autograd.LayerNorm.layerNormPrefix6 {m n : ℕ} (ε : ℝ) : Graph (ΓLN m n) (ssPrefix6 m n)

Graph prefix computing the first 6 intermediates (ssPrefix6).

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

Index of var_eps in ΓLN ++ ssPrefix6.

noncomputable def Proofs.Autograd.LayerNorm.nodeStd {m n : ℕ} : Node (ΓLN m n ++ ssPrefix6 m n) (VecShape m)

Standard deviation: std := sqrt_clamp(var_eps).

This is where the development becomes pointwise: differentiability depends on the (clamped) input.

noncomputable def Proofs.Autograd.LayerNorm.layerNormPrefix7 {m n : ℕ} (ε : ℝ) : Graph (ΓLN m n) (ssPrefix7 m n)

Graph prefix computing ssPrefix7 (adds std).

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

Index of std in ΓLN ++ ssPrefix7.

noncomputable def Proofs.Autograd.LayerNorm.nodeInvStd {m n : ℕ} : Node (ΓLN m n ++ ssPrefix7 m n) (VecShape m)

Inverse standard deviation: inv_std := 1/std.

noncomputable def Proofs.Autograd.LayerNorm.g8 {m n : ℕ} (ε : ℝ) : Graph (ΓLN m n) (ssPrefix7 m n ++ [VecShape m])

Graph prefix adding invStd.

def Proofs.Autograd.LayerNorm.idxInvStd {m n : ℕ} : Idx (ΓLN m n ++ (ssPrefix7 m n ++ [VecShape m])) (VecShape m)

Index of invStd in the extended context.

noncomputable def Proofs.Autograd.LayerNorm.nodeInvStdB {m n : ℕ} : Node (ΓLN m n ++ (ssPrefix7 m n ++ [VecShape m])) (MatShape m n)

Broadcast invStd back to m×n (row-wise).

noncomputable def Proofs.Autograd.LayerNorm.g9 {m n : ℕ} (ε : ℝ) : Graph (ΓLN m n) (ssPrefix7 m n ++ [VecShape m, MatShape m n])

Graph prefix adding inv_std_b.

def Proofs.Autograd.LayerNorm.idxCentered9 {m n : ℕ} : Idx (ΓLN m n ++ (ssPrefix7 m n ++ [VecShape m, MatShape m n])) (MatShape m n)

Index of centered in the stage-g9 context.

def Proofs.Autograd.LayerNorm.idxInvStdB9 {m n : ℕ} : Idx (ΓLN m n ++ (ssPrefix7 m n ++ [VecShape m, MatShape m n])) (MatShape m n)

Index of inv_std_b in the stage-g9 context.

noncomputable def Proofs.Autograd.LayerNorm.nodeNorm {m n : ℕ} : Node (ΓLN m n ++ (ssPrefix7 m n ++ [VecShape m, MatShape m n])) (MatShape m n)

Node computing normalized := centered ⊙ inv_std_b.

noncomputable def Proofs.Autograd.LayerNorm.g10 {m n : ℕ} (ε : ℝ) : Graph (ΓLN m n) (ssPrefix7 m n ++ [VecShape m, MatShape m n, MatShape m n])

Graph prefix producing normalized := centered ⊙ inv_std_b.

noncomputable def Proofs.Autograd.LayerNorm.nodeGammaB {m n : ℕ} : Node (ΓLN m n ++ (ssPrefix7 m n ++ [VecShape m, MatShape m n, MatShape m n])) (MatShape m n)

Broadcast gamma to m×n (column-wise).

noncomputable def Proofs.Autograd.LayerNorm.g11 {m n : ℕ} (ε : ℝ) : Graph (ΓLN m n) (ssPrefix7 m n ++ [VecShape m, MatShape m n, MatShape m n, MatShape m n])

Graph prefix adding gamma_b.

def Proofs.Autograd.LayerNorm.idxNorm11 {m n : ℕ} : Idx (ΓLN m n ++ (ssPrefix7 m n ++ [VecShape m, MatShape m n, MatShape m n, MatShape m n])) (MatShape m n)

Index of normalized in the context at stage g11.

def Proofs.Autograd.LayerNorm.idxGammaB11 {m n : ℕ} : Idx (ΓLN m n ++ (ssPrefix7 m n ++ [VecShape m, MatShape m n, MatShape m n, MatShape m n])) (MatShape m n)

Index of gamma_b at stage g11.

noncomputable def Proofs.Autograd.LayerNorm.nodeScaled {m n : ℕ} : Node (ΓLN m n ++ (ssPrefix7 m n ++ [VecShape m, MatShape m n, MatShape m n, MatShape m n])) (MatShape m n)

Scale: scaled := normalized ⊙ gamma_b.

noncomputable def Proofs.Autograd.LayerNorm.g12 {m n : ℕ} (ε : ℝ) : Graph (ΓLN m n) (ssPrefix7 m n ++ [VecShape m, MatShape m n, MatShape m n, MatShape m n, MatShape m n])

Graph prefix adding scaled.

noncomputable def Proofs.Autograd.LayerNorm.nodeBetaB {m n : ℕ} : Node (ΓLN m n ++ (ssPrefix7 m n ++ [VecShape m, MatShape m n, MatShape m n, MatShape m n, MatShape m n])) (MatShape m n)

Broadcast beta to m×n (column-wise).

noncomputable def Proofs.Autograd.LayerNorm.g13 {m n : ℕ} (ε : ℝ) : Graph (ΓLN m n) (ssPrefix7 m n ++ [VecShape m, MatShape m n, MatShape m n, MatShape m n, MatShape m n, MatShape m n])

Graph prefix adding beta_b.

def Proofs.Autograd.LayerNorm.idxScaled13 {m n : ℕ} : Idx (ΓLN m n ++ (ssPrefix7 m n ++ [VecShape m, MatShape m n, MatShape m n, MatShape m n, MatShape m n, MatShape m n])) (MatShape m n)

Index of scaled at stage g13.

def Proofs.Autograd.LayerNorm.idxBetaB13 {m n : ℕ} : Idx (ΓLN m n ++ (ssPrefix7 m n ++ [VecShape m, MatShape m n, MatShape m n, MatShape m n, MatShape m n, MatShape m n])) (MatShape m n)

Index of beta_b at stage g13.

noncomputable def Proofs.Autograd.LayerNorm.nodeY {m n : ℕ} : Node (ΓLN m n ++ (ssPrefix7 m n ++ [VecShape m, MatShape m n, MatShape m n, MatShape m n, MatShape m n, MatShape m n])) (MatShape m n)

Output: y := scaled + beta_b.

noncomputable def Proofs.Autograd.LayerNorm.layerNormGraph {m n : ℕ} (ε : ℝ) : Graph (ΓLN m n) (ssLayerNorm m n)

Full LayerNorm graph (as an explicit snoc chain).

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

Index of the final LayerNorm output in ΓLN ++ ssLayerNorm.