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.ssVarEps (m n : ℕ) : List Spec.Shape

LayerNorm intermediates through the epsilon-shifted variance.

@[reducible, inline]

abbrev Proofs.Autograd.LayerNorm.ssStd (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.graphMean {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.graphMeanBroadcast {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.graphCentered {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.graphCenteredSq {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.graphVar {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.graphVarEps {m n : ℕ} (ε : ℝ) : Graph (ΓLN m n) (ssVarEps m n)

Graph prefix through the epsilon-shifted variance (ssVarEps).

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

Index of var_eps in ΓLN ++ ssVarEps.

noncomputable def Proofs.Autograd.LayerNorm.nodeStd {m n : ℕ} : Node (ΓLN m n ++ ssVarEps 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.graphStd {m n : ℕ} (ε : ℝ) : Graph (ΓLN m n) (ssStd m n)

Graph prefix computing ssStd (adds std).

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

Index of std in ΓLN ++ ssStd.

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

Inverse standard deviation: inv_std := 1/std.

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

Graph prefix adding invStd.

def Proofs.Autograd.LayerNorm.idxInvStd {m n : ℕ} : Idx (ΓLN m n ++ (ssStd 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 ++ (ssStd m n ++ [VecShape m])) (MatShape m n)

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

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

Graph prefix adding inv_std_b.

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

Index of centered in the stage-graphInvStdBroadcast context.

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

Index of inv_std_b in the stage-graphInvStdBroadcast context.

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

Node computing normalized := centered ⊙ inv_std_b.

noncomputable def Proofs.Autograd.LayerNorm.graphNormalized {m n : ℕ} (ε : ℝ) : Graph (ΓLN m n) (ssStd 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 ++ (ssStd 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.graphGammaBroadcast {m n : ℕ} (ε : ℝ) : Graph (ΓLN m n) (ssStd m n ++ [VecShape m, MatShape m n, MatShape m n, MatShape m n])

Graph prefix adding gamma_b.

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

Index of normalized in the context at stage graphGammaBroadcast.

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

Index of gamma_b at stage graphGammaBroadcast.

noncomputable def Proofs.Autograd.LayerNorm.nodeScaled {m n : ℕ} : Node (ΓLN m n ++ (ssStd 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.graphScaled {m n : ℕ} (ε : ℝ) : Graph (ΓLN m n) (ssStd 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 ++ (ssStd 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.graphBetaBroadcast {m n : ℕ} (ε : ℝ) : Graph (ΓLN m n) (ssStd 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.idxScaledForOutput {m n : ℕ} : Idx (ΓLN m n ++ (ssStd 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 graphBetaBroadcast.

def Proofs.Autograd.LayerNorm.idxBetaBroadcast {m n : ℕ} : Idx (ΓLN m n ++ (ssStd 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 graphBetaBroadcast.

noncomputable def Proofs.Autograd.LayerNorm.nodeY {m n : ℕ} : Node (ΓLN m n ++ (ssStd 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.