BatchNormChannelFirst

Pointwise analytic correctness for a channel-first BatchNorm-like graph.

This matches the existing spec/runtime operator Spec.batchNorm_channel_first used by Runtime.Autograd.Tape.batchnorm_channel_first: it normalizes each channel independently by computing mean/variance over the spatial dimensions (H,W) and then applying per-channel affine parameters (gamma,beta).

The proof is spec-level over ℝ. Because the graph uses sqrt (max x 0) and inv, the statement is pointwise (GraphFDerivCorrectAt) with explicit domain assumptions.

Note: this is not PyTorch BatchNorm2d over N×H×W with running statistics; it is closer to an InstanceNorm/GroupNorm-style normalization over spatial dimensions per channel.

PyTorch correspondence / citations

• Reference BatchNorm2d and InstanceNorm2d docs (for naming/background; this file is a simpler, stateless normalization graph). https://pytorch.org/docs/stable/generated/torch.nn.BatchNorm2d.html https://pytorch.org/docs/stable/generated/torch.nn.InstanceNorm2d.html

@[reducible, inline]

abbrev Proofs.Autograd.BatchNormChannelFirst.CHWShape (channels height width : ℕ) : Spec.Shape

Channel-first tensor shape C×H×W.

@[reducible, inline]

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

Matrix shape m×n.

@[reducible, inline]

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

Vector shape k.

@[reducible, inline]

abbrev Proofs.Autograd.BatchNormChannelFirst.ΓBN (channels height width : ℕ) : List Spec.Shape

Input context shapes: [x, gamma, beta] with x : C×H×W and gamma/beta : C.

@[reducible, inline]

abbrev Proofs.Autograd.BatchNormChannelFirst.hw (height width : ℕ) : ℕ

Flattened spatial size H*W.

@[reducible, inline]

abbrev Proofs.Autograd.BatchNormChannelFirst.ssPrefixVarEps (channels height width : ℕ) : List Spec.Shape

Prefix intermediates up to var_eps (after flattening spatial dimensions).

@[reducible, inline]

abbrev Proofs.Autograd.BatchNormChannelFirst.ssPrefixStd (channels height width : ℕ) : List Spec.Shape

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

@[reducible, inline]

abbrev Proofs.Autograd.BatchNormChannelFirst.ssBatchNorm (channels height width : ℕ) : List Spec.Shape

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

def Proofs.Autograd.BatchNormChannelFirst.idxX {channels height width : ℕ} {ss : List Spec.Shape} : Idx (ΓBN channels height width ++ ss) (CHWShape channels height width)

Index of the input x in the base BatchNorm context ΓBN channels height width ++ ss.

def Proofs.Autograd.BatchNormChannelFirst.idxGamma {channels height width : ℕ} {ss : List Spec.Shape} : Idx (ΓBN channels height width ++ ss) (VecShape channels)

Index of the scale vector gamma in the base BatchNorm context ΓBN ++ ss.

def Proofs.Autograd.BatchNormChannelFirst.idxBeta {channels height width : ℕ} {ss : List Spec.Shape} : Idx (ΓBN channels height width ++ ss) (VecShape channels)

Index of the shift vector beta in the base BatchNorm context ΓBN ++ ss.

def Proofs.Autograd.BatchNormChannelFirst.idxLast {Γ ss : List Spec.Shape} {τ : Spec.Shape} : Idx (Γ ++ ss ++ [τ]) τ

Index helper for the last element of an extended context Γ ++ ss ++ [τ].

Informal computation (per channel, flattening spatial dims):

Let x : C×H×W, flatten spatial dims to xMat : C×(H*W). Then for each channel c:

mean_c := (1/(H*W)) * ∑_{p} xMat[c,p] centered := xMat - mean_b var_c := (1/(H*W)) * ∑_{p} centered[c,p]^2 std_c := sqrt(var_c + ε) (implemented as sqrt_clamp) inv_std_c := 1/std_c normalized := centered ⊙ inv_std_b scaled := normalized ⊙ gamma_b yMat := scaled + beta_b yChw := reshape yMat back to C×H×W

This is the stateless, per-example normalization used by the runtime spec Spec.batchNorm_channel_first; it is closer to InstanceNorm/GroupNorm than to BatchNorm with running statistics.

theorem Proofs.Autograd.BatchNormChannelFirst.hsz_chw_mat {channels height width : ℕ} : (CHWShape channels height width).size = (MatShape channels (hw height width)).size

noncomputable def Proofs.Autograd.BatchNormChannelFirst.nodeXMat {channels height width : ℕ} : Node (ΓBN channels height width) (MatShape channels (hw height width))

Reshape x : C×H×W into a matrix xMat : C×(H*W) (flatten spatial dimensions).

noncomputable def Proofs.Autograd.BatchNormChannelFirst.g1 {channels height width : ℕ} : Graph (ΓBN channels height width) [MatShape channels (hw height width)]

Graph prefix producing [xMat].

def Proofs.Autograd.BatchNormChannelFirst.idxXMat {channels height width : ℕ} : Idx (ΓBN channels height width ++ [MatShape channels (hw height width)]) (MatShape channels (hw height width))

Index of xMat in ΓBN ++ [xMat].

noncomputable def Proofs.Autograd.BatchNormChannelFirst.nodeMean {channels height width : ℕ} : Node (ΓBN channels height width ++ [MatShape channels (hw height width)]) (VecShape channels)

Per-channel mean over spatial dims: mean : C×(H*W) → C.

noncomputable def Proofs.Autograd.BatchNormChannelFirst.g2 {channels height width : ℕ} : Graph (ΓBN channels height width) [MatShape channels (hw height width), VecShape channels]

Graph prefix producing [xMat, mean].

def Proofs.Autograd.BatchNormChannelFirst.idxMean {channels height width : ℕ} : Idx (ΓBN channels height width ++ [MatShape channels (hw height width), VecShape channels]) (VecShape channels)

Index of mean in ΓBN ++ [xMat, mean].

noncomputable def Proofs.Autograd.BatchNormChannelFirst.nodeMeanB {channels height width : ℕ} : Node (ΓBN channels height width ++ [MatShape channels (hw height width), VecShape channels]) (MatShape channels (hw height width))

Broadcast mean back to C×(H*W) (row-wise).

noncomputable def Proofs.Autograd.BatchNormChannelFirst.g3 {channels height width : ℕ} : Graph (ΓBN channels height width) [MatShape channels (hw height width), VecShape channels, MatShape channels (hw height width)]

Graph prefix producing [xMat, mean, mean_b].

def Proofs.Autograd.BatchNormChannelFirst.idxMeanB {channels height width : ℕ} : Idx (ΓBN channels height width ++ [MatShape channels (hw height width), VecShape channels, MatShape channels (hw height width)]) (MatShape channels (hw height width))

Index of mean_b in the extended context.

def Proofs.Autograd.BatchNormChannelFirst.idxXMat3 {channels height width : ℕ} : Idx (ΓBN channels height width ++ [MatShape channels (hw height width), VecShape channels, MatShape channels (hw height width)]) (MatShape channels (hw height width))

Index of xMat in the extended context at stage g3.

noncomputable def Proofs.Autograd.BatchNormChannelFirst.nodeCentered {channels height width : ℕ} : Node (ΓBN channels height width ++ [MatShape channels (hw height width), VecShape channels, MatShape channels (hw height width)]) (MatShape channels (hw height width))

Center: centered := xMat - mean_b.

noncomputable def Proofs.Autograd.BatchNormChannelFirst.g4 {channels height width : ℕ} : Graph (ΓBN channels height width) [MatShape channels (hw height width), VecShape channels, MatShape channels (hw height width), MatShape channels (hw height width)]

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

def Proofs.Autograd.BatchNormChannelFirst.idxCentered {channels height width : ℕ} : Idx (ΓBN channels height width ++ [MatShape channels (hw height width), VecShape channels, MatShape channels (hw height width), MatShape channels (hw height width)]) (MatShape channels (hw height width))

Index of centered in the extended context at stage g4.

noncomputable def Proofs.Autograd.BatchNormChannelFirst.nodeCenteredSq {channels height width : ℕ} : Node (ΓBN channels height width ++ [MatShape channels (hw height width), VecShape channels, MatShape channels (hw height width), MatShape channels (hw height width)]) (MatShape channels (hw height width))

Square centered: centered_sq := centered ⊙ centered.

noncomputable def Proofs.Autograd.BatchNormChannelFirst.g5 {channels height width : ℕ} : Graph (ΓBN channels height width) [MatShape channels (hw height width), VecShape channels, MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width)]

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

def Proofs.Autograd.BatchNormChannelFirst.idxCenteredSq {channels height width : ℕ} : Idx (ΓBN channels height width ++ [MatShape channels (hw height width), VecShape channels, MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width)]) (MatShape channels (hw height width))

Index of centered_sq in the extended context at stage g5.

noncomputable def Proofs.Autograd.BatchNormChannelFirst.nodeVar {channels height width : ℕ} : Node (ΓBN channels height width ++ [MatShape channels (hw height width), VecShape channels, MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width)]) (VecShape channels)

Per-channel variance over spatial dims: var := mean(centered_sq) producing a length-channels vector.

noncomputable def Proofs.Autograd.BatchNormChannelFirst.g6 {channels height width : ℕ} : Graph (ΓBN channels height width) [MatShape channels (hw height width), VecShape channels, MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width), VecShape channels]

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

def Proofs.Autograd.BatchNormChannelFirst.idxVar {channels height width : ℕ} : Idx (ΓBN channels height width ++ [MatShape channels (hw height width), VecShape channels, MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width), VecShape channels]) (VecShape channels)

Index of var in the extended context at stage g6.

noncomputable def Proofs.Autograd.BatchNormChannelFirst.nodeVarEps {channels height width : ℕ} (ε : ℝ) : Node (ΓBN channels height width ++ [MatShape channels (hw height width), VecShape channels, MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width), VecShape channels]) (VecShape channels)

Add epsilon: var_eps := var + ε.

noncomputable def Proofs.Autograd.BatchNormChannelFirst.batchNormPrefixVarEps {channels height width : ℕ} (ε : ℝ) : Graph (ΓBN channels height width) (ssPrefixVarEps channels height width)

Graph prefix computing ssPrefixVarEps (up to var_eps).

def Proofs.Autograd.BatchNormChannelFirst.idxVarEps {channels height width : ℕ} : Idx (ΓBN channels height width ++ ssPrefixVarEps channels height width) (VecShape channels)

Index of var_eps in ΓBN ++ ssPrefixVarEps.

noncomputable def Proofs.Autograd.BatchNormChannelFirst.nodeStd {channels height width : ℕ} : Node (ΓBN channels height width ++ ssPrefixVarEps channels height width) (VecShape channels)

Standard deviation: std := sqrt_clamp(var_eps).

This is where the development becomes pointwise: differentiability depends on positivity of var_eps.

noncomputable def Proofs.Autograd.BatchNormChannelFirst.batchNormPrefixStd {channels height width : ℕ} (ε : ℝ) : Graph (ΓBN channels height width) (ssPrefixStd channels height width)

Graph prefix computing ssPrefixStd (adds std).

def Proofs.Autograd.BatchNormChannelFirst.idxStd {channels height width : ℕ} : Idx (ΓBN channels height width ++ ssPrefixStd channels height width) (VecShape channels)

Index of std in ΓBN ++ ssPrefixStd.

noncomputable def Proofs.Autograd.BatchNormChannelFirst.nodeInvStd {channels height width : ℕ} : Node (ΓBN channels height width ++ ssPrefixStd channels height width) (VecShape channels)

Inverse standard deviation: inv_std := 1/std.

noncomputable def Proofs.Autograd.BatchNormChannelFirst.g8 {channels height width : ℕ} (ε : ℝ) : Graph (ΓBN channels height width) (ssPrefixStd channels height width ++ [VecShape channels])

Graph prefix adding inv_std.

def Proofs.Autograd.BatchNormChannelFirst.idxInvStd {channels height width : ℕ} : Idx (ΓBN channels height width ++ (ssPrefixStd channels height width ++ [VecShape channels])) (VecShape channels)

Index of inv_std in the extended context after g8.

noncomputable def Proofs.Autograd.BatchNormChannelFirst.nodeInvStdB {channels height width : ℕ} : Node (ΓBN channels height width ++ (ssPrefixStd channels height width ++ [VecShape channels])) (MatShape channels (hw height width))

Broadcast inv_std back to C×(H*W) (row-wise), producing inv_std_b.

noncomputable def Proofs.Autograd.BatchNormChannelFirst.g9 {channels height width : ℕ} (ε : ℝ) : Graph (ΓBN channels height width) (ssPrefixStd channels height width ++ [VecShape channels, MatShape channels (hw height width)])

Graph prefix adding inv_std_b.

def Proofs.Autograd.BatchNormChannelFirst.idxCentered9 {channels height width : ℕ} : Idx (ΓBN channels height width ++ (ssPrefixStd channels height width ++ [VecShape channels, MatShape channels (hw height width)])) (MatShape channels (hw height width))

Index of centered in the context at stage g9.

This is obtained by weakening the earlier idxCentered along the extended intermediate list.

def Proofs.Autograd.BatchNormChannelFirst.idxInvStdB9 {channels height width : ℕ} : Idx (ΓBN channels height width ++ (ssPrefixStd channels height width ++ [VecShape channels, MatShape channels (hw height width)])) (MatShape channels (hw height width))

Index of inv_std_b in the context at stage g9.

noncomputable def Proofs.Autograd.BatchNormChannelFirst.nodeNorm {channels height width : ℕ} : Node (ΓBN channels height width ++ (ssPrefixStd channels height width ++ [VecShape channels, MatShape channels (hw height width)])) (MatShape channels (hw height width))

Normalize: normalized := centered ⊙ inv_std_b.

noncomputable def Proofs.Autograd.BatchNormChannelFirst.g10 {channels height width : ℕ} (ε : ℝ) : Graph (ΓBN channels height width) (ssPrefixStd channels height width ++ [VecShape channels, MatShape channels (hw height width), MatShape channels (hw height width)])

Graph prefix adding normalized.

def Proofs.Autograd.BatchNormChannelFirst.idxNorm10 {channels height width : ℕ} : Idx (ΓBN channels height width ++ (ssPrefixStd channels height width ++ [VecShape channels, MatShape channels (hw height width), MatShape channels (hw height width)])) (MatShape channels (hw height width))

Index of normalized in the extended context at stage g10.

noncomputable def Proofs.Autograd.BatchNormChannelFirst.nodeGammaB {channels height width : ℕ} : Node (ΓBN channels height width ++ (ssPrefixStd channels height width ++ [VecShape channels, MatShape channels (hw height width), MatShape channels (hw height width)])) (MatShape channels (hw height width))

Broadcast gamma : C to C×(H*W) (row-wise), producing gamma_b.

noncomputable def Proofs.Autograd.BatchNormChannelFirst.g11 {channels height width : ℕ} (ε : ℝ) : Graph (ΓBN channels height width) (ssPrefixStd channels height width ++ [VecShape channels, MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width)])

Graph prefix adding gamma_b.

def Proofs.Autograd.BatchNormChannelFirst.idxGammaB11 {channels height width : ℕ} : Idx (ΓBN channels height width ++ (ssPrefixStd channels height width ++ [VecShape channels, MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width)])) (MatShape channels (hw height width))

Index of gamma_b in the extended context at stage g11.

noncomputable def Proofs.Autograd.BatchNormChannelFirst.nodeScaled {channels height width : ℕ} : Node (ΓBN channels height width ++ (ssPrefixStd channels height width ++ [VecShape channels, MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width)])) (MatShape channels (hw height width))

Scale: scaled := normalized ⊙ gamma_b.

noncomputable def Proofs.Autograd.BatchNormChannelFirst.g12 {channels height width : ℕ} (ε : ℝ) : Graph (ΓBN channels height width) (ssPrefixStd channels height width ++ [VecShape channels, MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width)])

Graph prefix adding scaled.

def Proofs.Autograd.BatchNormChannelFirst.idxScaled12 {channels height width : ℕ} : Idx (ΓBN channels height width ++ (ssPrefixStd channels height width ++ [VecShape channels, MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width)])) (MatShape channels (hw height width))

Index of scaled in the extended context at stage g12.

noncomputable def Proofs.Autograd.BatchNormChannelFirst.nodeBetaB {channels height width : ℕ} : Node (ΓBN channels height width ++ (ssPrefixStd channels height width ++ [VecShape channels, MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width)])) (MatShape channels (hw height width))

Broadcast beta : C to C×(H*W) (row-wise), producing beta_b.

noncomputable def Proofs.Autograd.BatchNormChannelFirst.g13 {channels height width : ℕ} (ε : ℝ) : Graph (ΓBN channels height width) (ssPrefixStd channels height width ++ [VecShape channels, MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width)])

Graph prefix adding beta_b.

def Proofs.Autograd.BatchNormChannelFirst.idxBetaB13 {channels height width : ℕ} : Idx (ΓBN channels height width ++ (ssPrefixStd channels height width ++ [VecShape channels, MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width)])) (MatShape channels (hw height width))

Index of beta_b in the extended context at stage g13.

noncomputable def Proofs.Autograd.BatchNormChannelFirst.nodeYMat {channels height width : ℕ} : Node (ΓBN channels height width ++ (ssPrefixStd channels height width ++ [VecShape channels, MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width)])) (MatShape channels (hw height width))

Add bias: yMat := scaled + beta_b.

noncomputable def Proofs.Autograd.BatchNormChannelFirst.g14 {channels height width : ℕ} (ε : ℝ) : Graph (ΓBN channels height width) (ssPrefixStd channels height width ++ [VecShape channels, MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width)])

Graph prefix adding yMat.

def Proofs.Autograd.BatchNormChannelFirst.idxYMat {channels height width : ℕ} : Idx (ΓBN channels height width ++ (ssPrefixStd channels height width ++ [VecShape channels, MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width)])) (MatShape channels (hw height width))

Index of yMat in the extended context after g14.

noncomputable def Proofs.Autograd.BatchNormChannelFirst.nodeYChw {channels height width : ℕ} : Node (ΓBN channels height width ++ (ssPrefixStd channels height width ++ [VecShape channels, MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width), MatShape channels (hw height width)])) (CHWShape channels height width)

Reshape the matrix output yMat : C×(H*W) back into yChw : C×H×W.

noncomputable def Proofs.Autograd.BatchNormChannelFirst.batchNormGraph {channels height width : ℕ} (ε : ℝ) : Graph (ΓBN channels height width) (ssBatchNorm channels height width)

Full BatchNormChannelFirst graph (explicit snoc chain).

noncomputable def Proofs.Autograd.BatchNormChannelFirst.batchNormGraphFderivCorrectAt {channels height width : ℕ} (ε : ℝ) (xV : CtxVec (ΓBN channels height width)) (hVarEpsPos : ∀ (i : Fin (VecShape channels).size), 0 < (CtxVec.get idxVarEps ((batchNormPrefixVarEps ε).evalVec xV)).ofLp i) (hStdNe0 : ∀ (i : Fin (VecShape channels).size), (CtxVec.get idxStd ((batchNormPrefixStd ε).evalVec xV)).ofLp i ≠ 0) : GraphFDerivCorrectAt (batchNormGraph ε) xV

Pointwise proof that batchNormGraph satisfies GraphFDerivCorrectAt.

As with LayerNorm, the hypotheses are explicit domain assumptions needed for differentiability of sqrt (after clamp) and inv at the actual execution point.

theorem Proofs.Autograd.BatchNormChannelFirst.backprop_eq_adjoint_fderiv_batchNorm_channel_first_at {channels height width : ℕ} (ε : ℝ) (xV : CtxVec (ΓBN channels height width)) (seedV : CtxVec (ΓBN channels height width ++ ssBatchNorm channels height width)) (hVarEpsPos : ∀ (i : Fin (VecShape channels).size), 0 < (CtxVec.get idxVarEps ((batchNormPrefixVarEps ε).evalVec xV)).ofLp i) (hStdNe0 : ∀ (i : Fin (VecShape channels).size), (CtxVec.get idxStd ((batchNormPrefixStd ε).evalVec xV)).ofLp i ≠ 0) : (batchNormGraph ε).backpropVec xV seedV = (ContinuousLinearMap.adjoint (fderiv ℝ (batchNormGraph ε).evalVec xV)) seedV

Pointwise end-to-end result: backprop equals (fderiv eval)† for batchNormGraph.

This is the BatchNormChannelFirst analogue of the global DAG theorem, specialized to the explicit graph construction and with explicit domain assumptions for sqrt/inv.