Documentation

LeanMlir.Proofs.Architectures.PerChannelBN

Per-channel BatchNorm — the block-diagonal VJP #

Chapters 4–5 used a per-example global BatchNorm: one scalar (γ, β) over the whole oc·h·w activation (LayerNorm-shaped). Real ResNet wants per-channel BN: normalize each channel-slice independently with its own (γ_c, β_c), γ/β : Vec oc.

Because each channel is independent, the whole Jacobian is block-diagonal across the channel axis — the genuinely-new piece. We get it for free by generalizing the existing rowwiseHasVJPMat (Tensor.lean, multi-head attention) from a single per-row map to a per-row family g : Fin m → (Vec n → Vec p): viewing the activation as Mat oc (h·w) (row = channel), per-channel BN is exactly fun A => fun c => bnForward (h·w) ε (γ c) (β c) (A c). Its VJP runs each channel's bnHasVJP on that channel's cotangent slice; the cross-channel blocks vanish.

The file also holds inference BN (frozen statistics), batch BN on the [N,C,H,W] layout (bnBatchTensor4, chapter 7) and the sync-BN op at supplied statistics (bnSyncTensor4, its γ and input gradients) — the forms StableHLO's den reads. Sharding that layout across replicas (batchShard and the shard = global identities) is DataParallel.Sync's.

Everything closes under [propext, Classical.choice, Quot.sound].

noncomputable def Proofs.rowwisePerRowHasVJPMat {m n p : ℕ} (g : Fin m → Vec n → Vec p) (hg : (r : Fin m) → HasVJP (g r)) (hg_diff : ∀ (r : Fin m), Differentiable ℝ (g r)) :

HasVJPMat fun (A : Mat m n) (r : Fin m) => g r (A r)

Row-wise lifting of a per-row HasVJP family. Each row r gets its own map g r (with its own VJP); the matrix backward runs (g r).backward on row r's cotangent. The per-row peer of rowwiseHasVJPMat.

Equations

Proofs.rowwisePerRowHasVJPMat g hg hg_diff = { backward := fun (A : Proofs.Mat m n) (dY : Proofs.Mat m p) (r : Fin m) (c : Fin n) => (hg r).backward (A r) (dY r) c, correct := ⋯ }

Instances For

theorem Proofs.rowwisePerRow_flat_differentiable {m n p : ℕ} (g : Fin m → Vec n → Vec p) (h_g_diff : ∀ (r : Fin m), Differentiable ℝ (g r)) :

Differentiable ℝ fun (v : Vec (m * n)) => Mat.flatten ((fun (A : Mat m n) (r : Fin m) => g r (A r)) (Mat.unflatten v))

A per-row family flattens to a differentiable Vec → Vec map. The Differentiable witness vjpCompAt / the network composition needs to thread a per-channel BN through a block.

noncomputable def Proofs.bnPerChannelMat (oc m : ℕ) (ε : ℝ) (γ β : Vec oc) :

Mat oc m → Mat oc m

Per-channel BatchNorm (matrix view): BN each row (= channel-slice of m = h·w spatial cells) with its own (γ_c, β_c). The real-ResNet BN that bnForward's global scalar version approximates.

Equations

Proofs.bnPerChannelMat oc m ε γ β A c = Proofs.bnForward m ε (γ c) (β c) (A c)

Instances For

noncomputable def Proofs.bnPerChannelMatHasVJP (oc m : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) :

HasVJPMat (bnPerChannelMat oc m ε γ β)

Per-channel BN VJP (block-diagonal). Each channel runs its own bnHasVJP; the cross-channel Jacobian blocks vanish (pdivMat_rowIndep_perRow_at).

Equations

Proofs.bnPerChannelMatHasVJP oc m ε hε γ β = Proofs.rowwisePerRowHasVJPMat (fun (c : Fin oc) => Proofs.bnForward m ε (γ c) (β c)) (fun (c : Fin oc) => Proofs.bnHasVJP m ε (γ c) (β c) hε) ⋯

Instances For

noncomputable def Proofs.bnPerChannelFlat (oc m : ℕ) (ε : ℝ) (γ β : Vec oc) :

Vec (oc * m) → Vec (oc * m)

Per-channel BN as a flat-vector op Vec (oc·m) → Vec (oc·m) (row-major, channel c = the m-wide slab at flat positions finProdFinEquiv (c, ·)).

Equations

Proofs.bnPerChannelFlat oc m ε γ β v = (Proofs.bnPerChannelMat oc m ε γ β (Proofs.Mat.unflatten v)).flatten

Instances For

noncomputable def Proofs.bnPerChannelFlatHasVJP (oc m : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) :

HasVJP (bnPerChannelFlat oc m ε γ β)

Per-channel BN flat VJP — the block-diagonal matrix VJP bridged to Vec.

Equations

Proofs.bnPerChannelFlatHasVJP oc m ε hε γ β = (Proofs.bnPerChannelMatHasVJP oc m ε hε γ β).toHasVJP

Instances For

theorem Proofs.bnPerChannelFlat_differentiable (oc m : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) :

Differentiable ℝ (bnPerChannelFlat oc m ε γ β)

Per-channel BN is differentiable everywhere (ε > 0). The composition witness.

theorem Proofs.bnPerChannelFlatHasVJP_correct (oc m : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) (v dy : Vec (oc * m)) (i : Fin (oc * m)) :

(bnPerChannelFlatHasVJP oc m ε hε γ β).backward v dy i = ∑ j : Fin (oc * m), pdiv (bnPerChannelFlat oc m ε γ β) v i j * dy j

Per-channel BN VJP correctness (ℝ-headline): the flat backward equals the pdiv-contracted (block-diagonal) Jacobian of per-channel BN.

noncomputable def Proofs.bnPerChannelGradInput (oc m : ℕ) (ε : ℝ) (γ : Vec oc) (x dy : Vec (oc * m)) :

Vec (oc * m)

Per-channel consolidated BN input-gradient — the renderable closed form: run the per-example three-term bnGradInput on each channel-slice (m = h·w spatial cells), reusing that channel's γ_c. This is exactly what a bnPerChannelBack SHlo op / renderLNBack-per-channel emits; the abstract bnPerChannelFlatHasVJP.backward is the spec it must match.

Equations

One or more equations did not get rendered due to their size.

Instances For

theorem Proofs.bnPerChannelGradInput_correct (oc m : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) (x dy : Vec (oc * m)) (i : Fin (oc * m)) :

bnPerChannelGradInput oc m ε γ x dy i = ∑ j : Fin (oc * m), pdiv (bnPerChannelFlat oc m ε γ β) x i j * dy j

Renderable backward is faithful (ℝ-headline): the per-channel consolidated gradient equals the pdiv-contracted Jacobian of per-channel BN, under 0 < ε. Each channel reduces to the per-example bn_input_grad_correct. The licence to render per-channel BN's backward as the three-term formula per channel.

The network carries its activations in the Tensor3 flat layout (oc*h)*w (flatConv etc.), but bnPerChannelFlat is defined on the Mat-split layout oc*(h*w) (row c = the h·w spatial cells of channel c). The two Vecs have the same size; they differ only in how finProdFinEquiv associates the product. So the bridge is a pure re-association reindex (a permutation of coordinates) — a reindexCLM whose VJP is the scatter pdiv_reindex gives, exactly like decimateFlat. Conjugating bnPerChannelFlat by this bridge yields per-channel BN acting on the network's Tensor3 activations, with its VJP for free via vjpComp.

noncomputable def Proofs.reassocFwdIdx (oc h w : ℕ) (mIdx : Fin (oc * (h * w))) :

Fin (oc * h * w)

Re-association index Fin (oc*(h*w)) → Fin (oc*h*w): a Mat-split flat index (c, s) with s ↔ (hi, wi) maps to the Tensor3 flat index ((c, hi), wi). A pure product re-association — no arithmetic.

Equations

One or more equations did not get rendered due to their size.

Instances For

noncomputable def Proofs.reassocBackIdx (oc h w : ℕ) (t : Fin (oc * h * w)) :

Fin (oc * (h * w))

Re-association index Fin (oc*h*w) → Fin (oc*(h*w)): the inverse direction, Tensor3 ((c, hi), wi) ↦ Mat-split (c, (hi, wi)).

Equations

One or more equations did not get rendered due to their size.

Instances For

theorem Proofs.reassocFwdIdx_reassocBackIdx (oc h w : ℕ) (t : Fin (oc * h * w)) :

reassocFwdIdx oc h w (reassocBackIdx oc h w t) = t

The two re-association indices are mutual inverses — the bridge is a genuine relabeling (so conjugating by it really is per-channel BN, just in Tensor3 coordinates). Pure finProdFinEquiv round-trip.

theorem Proofs.reassocBackIdx_reassocFwdIdx (oc h w : ℕ) (mIdx : Fin (oc * (h * w))) :

reassocBackIdx oc h w (reassocFwdIdx oc h w mIdx) = mIdx

noncomputable def Proofs.reassocFwd (oc h w : ℕ) :

Vec (oc * h * w) → Vec (oc * (h * w))

Tensor3 → Mat-split reindex: read the ((c,hi),wi) cell at Mat position (c, (hi,wi)). A reindexCLM, hence continuous-linear / differentiable.

Equations

Proofs.reassocFwd oc h w y k = y (Proofs.reassocFwdIdx oc h w k)

Instances For

noncomputable def Proofs.reassocBack (oc h w : ℕ) :

Vec (oc * (h * w)) → Vec (oc * h * w)

Mat-split → Tensor3 reindex (the inverse relabeling).

Equations

Proofs.reassocBack oc h w y k = y (Proofs.reassocBackIdx oc h w k)

Instances For

Training BN computes μ/σ² from the activation it is normalizing; inference BN consumes frozen statistics (the driver's EMA'd running mean/var) and is therefore a plain affine map — pointwise in the activation, with no reduction at all. That is exactly why an eval forward built on it is class-batch-independent: an example's logits do not depend on which other examples share its batch.

The chain below mirrors bnForward → bnPerChannelMat → bnPerChannelFlat → bnPerChannelTensor3 one-for-one, so the eval op drops into the same layout bridge as the training op.

noncomputable def Proofs.bnEvalForward (m : ℕ) (ε γ β μ v : ℝ) (x : Vec m) :

Inference BN on one channel's m activations: yᵢ = γ · (xᵢ − μ) · (var + ε)^(−1/2) + β, with μ/var supplied rather than computed from x. The bnForward peer — note it takes x pointwise, where bnForward reduces over all of x to get its own μ/σ².

Equations

Proofs.bnEvalForward m ε γ β μ v x i = γ * ((x i - μ) * (1 / √(v + ε))) + β

Instances For

noncomputable def Proofs.bnPerChannelEvalMat (oc m : ℕ) (ε : ℝ) (γ β μ v : Vec oc) :

Mat oc m → Mat oc m

Per-channel inference BN (Mat layout) — row c gets channel c's frozen stats.

Equations

Proofs.bnPerChannelEvalMat oc m ε γ β μ v A c = Proofs.bnEvalForward m ε (γ c) (β c) (μ c) (v c) (A c)

Instances For

noncomputable def Proofs.bnPerChannelEvalFlat (oc m : ℕ) (ε : ℝ) (γ β μ v : Vec oc) :

Vec (oc * m) → Vec (oc * m)

Per-channel inference BN (flat layout) — the bnPerChannelFlat peer.

Equations

Proofs.bnPerChannelEvalFlat oc m ε γ β μ v x = (Proofs.bnPerChannelEvalMat oc m ε γ β μ v (Proofs.Mat.unflatten x)).flatten

Instances For

noncomputable def Proofs.bnPerChannelEvalTensor3 (oc h w : ℕ) (ε : ℝ) (γ β μ v : Vec oc) :

Vec (oc * h * w) → Vec (oc * h * w)

Per-channel inference BN (Tensor3 layout) — the bnPerChannelTensor3 peer, through the same reassoc bridge. This is what SHlo.bnPerChannelEvalF denotes.

Equations

Proofs.bnPerChannelEvalTensor3 oc h w ε γ β μ v = Proofs.reassocBack oc h w ∘ Proofs.bnPerChannelEvalFlat oc (h * w) ε γ β μ v ∘ Proofs.reassocFwd oc h w

Instances For

theorem Proofs.bnEvalForward_at_own_stats (n : ℕ) (hn : n ≠ 0) (ε γ β : ℝ) (x : Vec n) :

bnEvalForward n ε γ β (bnMean n x) (bnMeanSq n x - bnMean n x * bnMean n x) x = bnForward n ε γ β x

Frozen-stats BN at a channel's own statistics is the training BN.

Hand bnEvalForward the mean and the second moment of x itself and it reproduces bnForward exactly, the variance arriving through bnVar_eq_bnMeanSq_sub_sq.

This is the R = 1 anchor of synchronised BatchNorm, and the reason the sync render is a drop-in: a graph that normalises with handed-in statistics denotes the same function as one that computes them, whenever the handed-in ones are the right ones. At R = 1 every allReduceMeanF threads its operand, so the sync forward collapses to exactly this and the single-device artifacts need not move.

theorem Proofs.bnEvalForward_differentiable (m : ℕ) (ε γ β μ v : ℝ) :

Differentiable ℝ (bnEvalForward m ε γ β μ v)

Inference BN is differentiable everywhere — it is affine in x, so unlike the training BN this needs no 0 < ε hypothesis: ε only enters the constant scale factor.

noncomputable def Proofs.bnPerChannelGradGamma (oc m : ℕ) (ε : ℝ) (v dy : Vec (oc * m)) :

Vec oc

The rendered per-channel γ gradient: dγ_c = Σ_{s} dy_(c,s) · x̂_(c,s) (the reduce over batch/spatial of dy·x̂ that the bnGammaSgd op emits). x̂ is recomputed from the saved BN input v (the conv output). Lives here (not PerChannelBNGrad) so the bnGammaSgd SHlo op's den can reference it.

Equations

Proofs.bnPerChannelGradGamma oc m ε v dy c = ∑ s : Fin m, dy (finProdFinEquiv (c, s)) * Proofs.bnXhat m ε (Proofs.Mat.unflatten v c) s

Instances For

noncomputable def Proofs.bnPerChannelGradBeta (oc m : ℕ) (dy : Vec (oc * m)) :

Vec oc

The rendered per-channel β gradient: dβ_c = Σ_{s} dy_(c,s).

Equations

Proofs.bnPerChannelGradBeta oc m dy c = ∑ s : Fin m, dy (finProdFinEquiv (c, s))

Instances For

theorem Proofs.reassocFwd_differentiable (oc h w : ℕ) :

Differentiable ℝ (reassocFwd oc h w)

theorem Proofs.reassocBack_differentiable (oc h w : ℕ) :

Differentiable ℝ (reassocBack oc h w)

noncomputable def Proofs.reassocFwdHasVJP (oc h w : ℕ) :

HasVJP (reassocFwd oc h w)

VJP of the forward reindex — the scatter pdiv_reindex gives. Mirrors decimateFlatHasVJP.

Equations

Proofs.reassocFwdHasVJP oc h w = Proofs.reindexVJP (Proofs.reassocFwdIdx oc h w)

Instances For

noncomputable def Proofs.reassocBackHasVJP (oc h w : ℕ) :

HasVJP (reassocBack oc h w)

VJP of the inverse reindex reassocBack, by reindexVJP.

Equations

Proofs.reassocBackHasVJP oc h w = Proofs.reindexVJP (Proofs.reassocBackIdx oc h w)

Instances For

theorem Proofs.reassocBackHasVJP_backward_eq (oc h w : ℕ) (v : Vec (oc * (h * w))) (dy : Vec (oc * h * w)) :

(reassocBackHasVJP oc h w).backward v dy = reassocFwd oc h w dy

The bridge is a permutation, so each reindex's VJP backward is just the inverse reindex (the single matching delta survives the scatter). These two collapse the vjpComp backwards into a clean closed form for bnPerChannelTensor3.

theorem Proofs.reassocFwdHasVJP_backward_eq (oc h w : ℕ) (v : Vec (oc * h * w)) (dy : Vec (oc * (h * w))) :

(reassocFwdHasVJP oc h w).backward v dy = reassocBack oc h w dy

noncomputable def Proofs.bnPerChannelTensor3 (oc h w : ℕ) (ε : ℝ) (γ β : Vec oc) :

Vec (oc * h * w) → Vec (oc * h * w)

Per-channel BatchNorm on the Tensor3 (oc*h)*w activation layout. Conjugate the Mat-split bnPerChannelFlat by the layout bridge: relabel to Mat-split, normalize each channel over its h·w spatial cells, relabel back. This is what SHlo.bnPerChannelF denotes.

Equations

Proofs.bnPerChannelTensor3 oc h w ε γ β = Proofs.reassocBack oc h w ∘ Proofs.bnPerChannelFlat oc (h * w) ε γ β ∘ Proofs.reassocFwd oc h w

Instances For

theorem Proofs.bnPerChannelTensor3_differentiable (oc h w : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) :

Differentiable ℝ (bnPerChannelTensor3 oc h w ε γ β)

noncomputable def Proofs.bnPerChannelTensor3HasVJP (oc h w : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) :

HasVJP (bnPerChannelTensor3 oc h w ε γ β)

Per-channel BN (Tensor3 layout) VJP — block-diagonal across channels, lifted through the layout bridge by vjpComp (twice).

Equations

One or more equations did not get rendered due to their size.

Instances For

theorem Proofs.bnPerChannelTensor3HasVJP_correct (oc h w : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) (x dy : Vec (oc * h * w)) (i : Fin (oc * h * w)) :

(bnPerChannelTensor3HasVJP oc h w ε hε γ β).backward x dy i = ∑ j : Fin (oc * h * w), pdiv (bnPerChannelTensor3 oc h w ε γ β) x i j * dy j

Per-channel BN (Tensor3 layout) VJP correctness (ℝ-headline): the backward equals the pdiv-contracted (block-diagonal) Jacobian of per-channel BN on the network's activation layout. The licence to wire per-channel BN into ResNet-34.

theorem Proofs.bnPerChannelTensor3HasVJP_backward_eq (oc h w : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) (x dy : Vec (oc * h * w)) :

(bnPerChannelTensor3HasVJP oc h w ε hε γ β).backward x dy = reassocBack oc h w ((bnPerChannelFlatHasVJP oc (h * w) ε hε γ β).backward (reassocFwd oc h w x) (reassocFwd oc h w dy))

The composed vjpComp backward collapses (the bridge reindexes are permutations): per-channel BN's Tensor3 backward is the Mat-split block-diagonal backward, conjugated by the layout bridge.

noncomputable def Proofs.bnPerChannelTensor3GradInput (oc h w : ℕ) (ε : ℝ) (γ : Vec oc) (x dy : Vec (oc * h * w)) :

Vec (oc * h * w)

Renderable per-channel BN backward on the Tensor3 (oc*h)*w layout — relabel to Mat-split, run the per-channel consolidated three-term bnPerChannelGradInput, relabel back. This is exactly what the bnPerChannelBack SHlo op emits (per-channel renderLNBack, reducing over the spatial axis); its faithfulness spec is below.

Equations

Proofs.bnPerChannelTensor3GradInput oc h w ε γ x dy = Proofs.reassocBack oc h w (Proofs.bnPerChannelGradInput oc (h * w) ε γ (Proofs.reassocFwd oc h w x) (Proofs.reassocFwd oc h w dy))

Instances For

theorem Proofs.bnPerChannelTensor3GradInput_correct (oc h w : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) (x dy : Vec (oc * h * w)) (i : Fin (oc * h * w)) :

bnPerChannelTensor3GradInput oc h w ε γ x dy i = ∑ j : Fin (oc * h * w), pdiv (bnPerChannelTensor3 oc h w ε γ β) x i j * dy j

Renderable Tensor3 backward is faithful (ℝ-headline): equals the pdiv-contracted (block-diagonal) Jacobian of per-channel BN on the network's activation layout, under 0 < ε. The licence to render per-channel BN's backward in ResNet-34.

noncomputable def Proofs.bnchwFwdIdx (N oc h w : ℕ) (mIdx : Fin (oc * (N * (h * w)))) :

Fin (N * (oc * (h * w)))

Transpose-reindex Fin (oc*(N*(h*w))) → Fin (N*(oc*(h*w))): a per-channel Mat index (c, (n, s)) maps to the network [N,C,H,W] flat index (n, (c, s)) — swap the batch and channel axes (s ↔ (hi,wi) carried along). A permutation.

Equations

One or more equations did not get rendered due to their size.

Instances For

noncomputable def Proofs.bnchwBackIdx (N oc h w : ℕ) (t : Fin (N * (oc * (h * w)))) :

Fin (oc * (N * (h * w)))

Transpose-reindex Fin (N*(oc*(h*w))) → Fin (oc*(N*(h*w))): the inverse, network (n, (c, s)) ↦ per-channel Mat (c, (n, s)).

Equations

One or more equations did not get rendered due to their size.

Instances For

theorem Proofs.bnchwFwdIdx_bnchwBackIdx (N oc h w : ℕ) (t : Fin (N * (oc * (h * w)))) :

bnchwFwdIdx N oc h w (bnchwBackIdx N oc h w t) = t

theorem Proofs.bnchwBackIdx_bnchwFwdIdx (N oc h w : ℕ) (mIdx : Fin (oc * (N * (h * w)))) :

bnchwBackIdx N oc h w (bnchwFwdIdx N oc h w mIdx) = mIdx

noncomputable def Proofs.bnchwFwd (N oc h w : ℕ) :

Vec (N * (oc * (h * w))) → Vec (oc * (N * (h * w)))

[N,C,H,W] → [C,N·H·W] reindex (gather the network cell at the Mat position).

Equations

Proofs.bnchwFwd N oc h w y k = y (Proofs.bnchwFwdIdx N oc h w k)

Instances For

theorem Proofs.bnchwFwd_apply (N oc h w : ℕ) (y : Vec (N * (oc * (h * w)))) (k : Fin (oc * (N * (h * w)))) :

bnchwFwd N oc h w y k = y (bnchwFwdIdx N oc h w k)

noncomputable def Proofs.bnchwBack (N oc h w : ℕ) :

Vec (oc * (N * (h * w))) → Vec (N * (oc * (h * w)))

[C,N·H·W] → [N,C,H,W] reindex (the inverse relabeling).

Equations

Proofs.bnchwBack N oc h w y k = y (Proofs.bnchwBackIdx N oc h w k)

Instances For

theorem Proofs.bnchwFwd_differentiable (N oc h w : ℕ) :

Differentiable ℝ (bnchwFwd N oc h w)

theorem Proofs.bnchwBack_differentiable (N oc h w : ℕ) :

Differentiable ℝ (bnchwBack N oc h w)

noncomputable def Proofs.bnchwFwdHasVJP (N oc h w : ℕ) :

HasVJP (bnchwFwd N oc h w)

VJP of the [N,C,H,W] → [C,N·H·W] reindex bnchwFwd, by reindexVJP.

Equations

Proofs.bnchwFwdHasVJP N oc h w = Proofs.reindexVJP (Proofs.bnchwFwdIdx N oc h w)

Instances For

noncomputable def Proofs.bnchwBackHasVJP (N oc h w : ℕ) :

HasVJP (bnchwBack N oc h w)

VJP of the [C,N·H·W] → [N,C,H,W] reindex bnchwBack, by reindexVJP.

Equations

Proofs.bnchwBackHasVJP N oc h w = Proofs.reindexVJP (Proofs.bnchwBackIdx N oc h w)

Instances For

theorem Proofs.bnchwBackHasVJP_backward_eq (N oc h w : ℕ) (v : Vec (oc * (N * (h * w)))) (dy : Vec (N * (oc * (h * w)))) :

(bnchwBackHasVJP N oc h w).backward v dy = bnchwFwd N oc h w dy

theorem Proofs.bnchwFwdHasVJP_backward_eq (N oc h w : ℕ) (v : Vec (N * (oc * (h * w)))) (dy : Vec (oc * (N * (h * w)))) :

(bnchwFwdHasVJP N oc h w).backward v dy = bnchwBack N oc h w dy

noncomputable def Proofs.bnBatchTensor4 (N oc h w : ℕ) (ε : ℝ) (γ β : Vec oc) :

Vec (N * (oc * (h * w))) → Vec (N * (oc * (h * w)))

Batch-norm per channel on the network's [N,C,H,W] layout. Conjugate the Mat-split bnPerChannelFlat (with m = N·h·w, the whole batch's cells per channel) by the transpose bridge: relabel [N,C,H,W] → [C, N·H·W], normalize each channel over ALL its batch+spatial cells, relabel back. The EfficientNet normalization.

Equations

Proofs.bnBatchTensor4 N oc h w ε γ β = Proofs.bnchwBack N oc h w ∘ Proofs.bnPerChannelFlat oc (N * (h * w)) ε γ β ∘ Proofs.bnchwFwd N oc h w

Instances For

noncomputable def Proofs.bnSyncTensor4 (N oc h w : ℕ) (ε : ℝ) (γ β μ m2 : Vec oc) :

Vec (N * (oc * (h * w))) → Vec (N * (oc * (h * w)))

Synchronised batch-norm on the [N,C,H,W] layout — statistics handed in.

bnBatchTensor4's peer, conjugated by the same [N,C,H,W] → [C, N·H·W] bridge, but the per-channel normalisation reads μ and the second moment m2 from its arguments instead of reducing x for them. Under data parallelism those arguments are the ALL-REDUCED global statistics — which is how one replica normalises over a batch it cannot see.

It is stated at μ and the second moment m2, and forms the variance as m2 − μ². The emitted sync forward exchanges μ and σ² (Chan's parallel variance, bnVar_shard_chan), and SHlo.bnSyncF's den supplies m2 := σ² + μ², so m2 − μ² is that σ² in ℝ.

Equations

One or more equations did not get rendered due to their size.

Instances For

theorem Proofs.bnPerChannelEvalFlat_apply (oc m : ℕ) (ε : ℝ) (γ β μ v : Vec oc) (z : Vec (oc * m)) (idx : Fin (oc * m)) :

bnPerChannelEvalFlat oc m ε γ β μ v z idx = γ (finProdFinEquiv.symm idx).1 * ((z idx - μ (finProdFinEquiv.symm idx).1) * (1 / √(v (finProdFinEquiv.symm idx).1 + ε))) + β (finProdFinEquiv.symm idx).1

Frozen-stats per-channel BN is POINTWISE, with the channel read off the index: the only thing the Mat round-trip does is decide which channel's γ β μ v a cell gets.

theorem Proofs.bnSyncTensor4_at_own_stats (N oc h w : ℕ) (hm : N * (h * w) ≠ 0) (ε : ℝ) (γ β : Vec oc) (x : Vec (N * (oc * (h * w)))) :

bnSyncTensor4 N oc h w ε γ β (fun (c : Fin oc) => bnMean (N * (h * w)) (Mat.unflatten (bnchwFwd N oc h w x) c)) (fun (c : Fin oc) => bnMeanSq (N * (h * w)) (Mat.unflatten (bnchwFwd N oc h w x) c)) x = bnBatchTensor4 N oc h w ε γ β x

R = 1: the sync forward at the batch's own statistics is bnBatchTensor4.

At R = 1 every allReduceMeanF threads its operand, so the sync graph hands in exactly the statistics the batch would have computed, and this says that graph denotes bnBatchTensor4, the per-batch forward.

theorem Proofs.bnBatchTensor4_differentiable (N oc h w : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) :

Differentiable ℝ (bnBatchTensor4 N oc h w ε γ β)

noncomputable def Proofs.bnBatchTensor4HasVJP (N oc h w : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) :

HasVJP (bnBatchTensor4 N oc h w ε γ β)

Batch-norm (network layout) VJP — block-diagonal across channels (now coupling the whole batch within each channel), lifted through the transpose bridge.

Equations

One or more equations did not get rendered due to their size.

Instances For

theorem Proofs.bnBatchTensor4HasVJP_correct (N oc h w : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) (x dy : Vec (N * (oc * (h * w)))) (i : Fin (N * (oc * (h * w)))) :

(bnBatchTensor4HasVJP N oc h w ε hε γ β).backward x dy i = ∑ j : Fin (N * (oc * (h * w))), pdiv (bnBatchTensor4 N oc h w ε γ β) x i j * dy j

theorem Proofs.bnBatchTensor4HasVJP_backward_eq (N oc h w : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) (x dy : Vec (N * (oc * (h * w)))) :

(bnBatchTensor4HasVJP N oc h w ε hε γ β).backward x dy = bnchwBack N oc h w ((bnPerChannelFlatHasVJP oc (N * (h * w)) ε hε γ β).backward (bnchwFwd N oc h w x) (bnchwFwd N oc h w dy))

noncomputable def Proofs.bnBatchTensor4GradInput (N oc h w : ℕ) (ε : ℝ) (γ : Vec oc) (x dy : Vec (N * (oc * (h * w)))) :

Vec (N * (oc * (h * w)))

Renderable batch-norm backward on the [N,C,H,W] layout — relabel to the per-channel Mat, run the consolidated three-term bnPerChannelGradInput over the whole batch (m = N·h·w), relabel back. Exactly what the batched batch-norm backward StableHLO fragment emits (reduce over [0,2,3] per channel).

Equations

Proofs.bnBatchTensor4GradInput N oc h w ε γ x dy = Proofs.bnchwBack N oc h w (Proofs.bnPerChannelGradInput oc (N * (h * w)) ε γ (Proofs.bnchwFwd N oc h w x) (Proofs.bnchwFwd N oc h w dy))

Instances For

noncomputable def Proofs.bnSyncPerChannelGradInput (oc m : ℕ) (ε : ℝ) (γ μ m2 mdy mdyx : Vec oc) (x dy : Vec (oc * m)) :

Vec (oc * m)

Per-channel SYNC backward (flat layout) — bnPerChannelGradInput's peer, with each channel's μ, E[x²] and two reduction means supplied rather than reduced out of x/dy. Under data parallelism those are the all-reduced global ones.

Equations

One or more equations did not get rendered due to their size.

Instances For

noncomputable def Proofs.bnSyncPerChannelGradGamma (oc m : ℕ) (ε : ℝ) (μ m2 : Vec oc) (v dy : Vec (oc * m)) :

Vec oc

The per-channel γ gradient with x̂ at handed-in statistics — bnPerChannelGradGamma's peer. Under sync-BN the forward normalised with the all-reduced global statistics, so ∂L/∂γ_c = Σ dy·x̂ must use the SAME x̂; bnPerChannelGradGamma rebuilds it from the shard's own statistics (bnXhat), which is a different function once R > 1. β's gradient reads no statistic and needs no peer.

Equations

Proofs.bnSyncPerChannelGradGamma oc m ε μ m2 v dy c = ∑ s : Fin m, Proofs.Mat.unflatten dy c s * Proofs.bnSyncXhat m ε (μ c) (m2 c) (Proofs.Mat.unflatten v c) s

Instances For

theorem Proofs.bnSyncPerChannelGradGamma_at_own_stats (oc m : ℕ) (hm : m ≠ 0) (ε : ℝ) (v dy : Vec (oc * m)) :

bnSyncPerChannelGradGamma oc m ε (fun (c : Fin oc) => bnMean m (Mat.unflatten v c)) (fun (c : Fin oc) => bnMeanSq m (Mat.unflatten v c)) v dy = bnPerChannelGradGamma oc m ε v dy

R = 1: the sync γ gradient at its own statistics IS bnPerChannelGradGamma. The γ-gradient anchor beside bnSyncTensor4_at_own_stats: a single-device sync render's γ node denotes what today's bnGammaGradB denotes.

noncomputable def Proofs.bnSyncTensor4GradInput (N oc h w : ℕ) (ε : ℝ) (γ μ m2 mdy mdyx : Vec oc) (x dy : Vec (N * (oc * (h * w)))) :

Vec (N * (oc * (h * w)))

The sync batch-norm input-VJP on [N,C,H,W] — bnBatchTensor4GradInput's peer, through the same bnchwFwd/bnchwBack bridge. What a replica emits for its shard of the backward, given the four all-reduced per-channel statistic vectors.

Equations

One or more equations did not get rendered due to their size.

Instances For

theorem Proofs.bnSyncTensor4GradInput_at_own_stats (N oc h w : ℕ) (hm : N * (h * w) ≠ 0) (ε : ℝ) (γ : Vec oc) (x dy : Vec (N * (oc * (h * w)))) :

bnSyncTensor4GradInput N oc h w ε γ (fun (c : Fin oc) => bnMean (N * (h * w)) (Mat.unflatten (bnchwFwd N oc h w x) c)) (fun (c : Fin oc) => bnMeanSq (N * (h * w)) (Mat.unflatten (bnchwFwd N oc h w x) c)) (fun (c : Fin oc) => bnMean (N * (h * w)) fun (k : Fin (N * (h * w))) => γ c * Mat.unflatten (bnchwFwd N oc h w dy) c k) (fun (c : Fin oc) => bnMean (N * (h * w)) fun (k : Fin (N * (h * w))) => bnXhat (N * (h * w)) ε (Mat.unflatten (bnchwFwd N oc h w x) c) k * (γ c * Mat.unflatten (bnchwFwd N oc h w dy) c k)) x dy = bnBatchTensor4GradInput N oc h w ε γ x dy

R = 1: the sync backward at its own statistics is bnBatchTensor4GradInput.

The [N,C,H,W] lift of bnSyncGradInput_at_own_stats, and the backward half of the drop-in claim: a single-device sync render computes the same gradient as the per-batch bnBatchTensor4GradInput.

theorem Proofs.bnBatchTensor4GradInput_correct (N oc h w : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) (x dy : Vec (N * (oc * (h * w)))) (i : Fin (N * (oc * (h * w)))) :

bnBatchTensor4GradInput N oc h w ε γ x dy i = ∑ j : Fin (N * (oc * (h * w))), pdiv (bnBatchTensor4 N oc h w ε γ β) x i j * dy j

Renderable batch-norm backward is faithful (ℝ-headline): equals the pdiv-contracted (block-diagonal-across-channels, batch-coupled) Jacobian of batch-norm on the network's [N,C,H,W] layout, under 0 < ε. The licence to render EfficientNet's batch-norm backward as the per-channel three-term formula over the batch.