The batch-BatchNorm seal kit — non-degeneracy machinery for the full-width nets #
The level-2/3 witnesses (Training/JacobianSeal.lean's
backward_nontrivial_of_fderiv_ne) on 2-channel per-example proxies see a BatchNorm that is
bnForward over one activation. The nets the ImageNet artifacts
actually run normalize with StableHLO.bnBatchLA — bnBatchTensor4 at the network's left-assoc
index, i.e. bnPerChannelFlat oc (N·h·w): each channel over all its batch-and-spatial cells.
This file is the machinery for sealing those, shared by the four kinked full-width nets
(ResNet-34/50, MobileNetV2/V4).
Why the carrier has to change, and what that costs #
A proxy seal carries a channel difference: channel 0 = channel 1 + δ at every position. Under
per-channel batch BN that carrier dies — a channel-uniform offset is exactly what the channel's own
mean subtracts — and at N = 1 the structural net is constant in its input outright. So the
carrier here is an example difference (EDiff in the per-net files): example 0's slab is
example 1's slab plus δ, per channel. Batch BN keeps it and scales it by γ_c · istd_c
(bnBatchLA_exdiff), because the two examples share one mean and one istd. That forces N = 2
and makes the witness exercise the one op that couples examples.
The cost is that every fact now has to be read at a cell (n, c, i, j) of a left-assoc flat index.
§1 pays that once: bcell is the per-example Tensor3 view, bnRowLA is the row BN normalizes,
and bnBatchLA_bcell is the bridge. Everything after it is the usual BN algebra on one row.
What is here #
- §0
bcell/bfrom— the[N,C,H,W]cell view, andbatchMap's action on it. - §1 the index bridge:
laIdx,bnRowLA,bnBatchLA_bcell, andbnBatchLA_pointwise(the workhorse: a channel-independent property of every cell of abnBatchLAoutput reduces to the scalarbnForwardon one row, with no index decomposition at the call site). - §2 the BN consequences the seals need: constant channel ↦
β, the|bn − β| ≤ |γ|√nmargin (bnForward_abs_sub_le, inBatchNormwith the other op-level BN facts; hence positivity, hence relu off its kink at every input; hence, two-sided, the relu6 windowbnBatchLA_window/bnBatchLA_smooth6), the example-difference identity, and within-example injectivity (the stem pool's no-tie). - §3 the centre-tap kernel
ctKand its conv value: the one weight shape that carries a signal through a channel-changing conv while staying transparent to a uniform offset. Only the centre tap is nonzero, which is what makes it padding-proof — a conv of a constant is not constant near a zero-padded border, but a centre tap is always in range. - §3b the XLA-
SAMEpeers, which keep the odd spatial positions where the symmetric ops keep the even ones (decimateOdd_unflatten,flatConvStride2Xla_ctK,bcell_convS2Xla_ctK). - §3c the centre-tap depthwise kernel
ctDW. A depthwise cannot broadcast, so wherectKcollapses the carrier to one value at every output channel,ctDWscales it channel by channel. - The op-level facts the ray argument leans on live with their ops, not here: the 3×3/s2 pool
shifts with a uniform offset and keeps nonnegativity (
maxPool3s2_shift,maxPool3s2_nonneg,MaxPool3s2), GAP shifts likewise (globalAvgPool_shift,CNN), andrelu,residual, the pool,batchMapandbnIstdare continuous (MLP,Residual,MaxPool3s2,Batched,BatchNorm). §5 keepsbnRowLA_continuous, which is about this file's own reindex. - §11 the SYMMETRIC strided depthwise (
eDiff_dwS2), MobileNetV4's; §3c's is the XLA one. - §12
BUnif, a grid-constant slab, andeDiff_of_bUnif, the carrier it starts.
Assemble a batched activation from per-example slabs — the inverse of bcell, used to write
the witness input down.
Equations
- Proofs.BatchSeal.bfrom f = Proofs.Mat.flatten fun (n : Fin N) => (f n).flatten
Instances For
The network's left-assoc [N,C,H,W] flat index.
Equations
- Proofs.BatchSeal.laIdx N oc h w n c i j = finProdFinEquiv (n, finProdFinEquiv (finProdFinEquiv (c, i), j))
Instances For
The row batch BN normalizes: channel c's N·h·w cells, over the whole batch.
Equations
- Proofs.BatchSeal.bnRowLA N oc h w v c = Proofs.Mat.unflatten (Proofs.bnchwFwd N oc h w (v ∘ Fin.cast ⋯)) c
Instances For
Batch BN keeps equal cells equal within one example and channel: the per-channel affine
map reads nothing of a cell but its value. What lets a tie between two stem cells that read
identical input patches survive every BN parameter (StemConvTwin in ResNet34ParamGrad).
A constant channel normalizes to β (variance 0, xhat = 0): a zeroed residual body is
the constant β₂, whatever γ₂ is. The batched peer of bnForward_const.
Positivity from the margin — |g|√(N·h·w) < b makes the whole BN output strictly
positive, at every input. The relu clause of every conv-bn-relu stage in the witness.
The relu6 window from the margin — the two-sided twin of bnBatchLA_pos. With β = b
strictly inside (|g|√(N·h·w), 6 − |g|√(N·h·w)) the whole BN output sits strictly inside
(0, 6), at every input, so relu6_id_window collapses the stage that follows it.
b = 3 centres the window and makes both hypotheses the single check |g|√(N·h·w) < 3.
Both relu6 clauses at once, in the ≠ 0 ∧ ≠ 6 shape every MobileNet smoothness bundle
is stated in. Since the bound is input-independent, this discharges a relu6 clause without
reading the activation — which is why a relu6 net's whole clause bundle is weight-only.
The carrier step. Two examples of one channel share the channel's mean and istd, so
batch BN keeps their difference and multiplies it by γ_c · istd_c. This is what a
channel-difference carrier cannot do under per-channel batch BN, and it is why the witness is
at N = 2.
The centre-tap broadcast kernel. Every output channel reads input channel 0 through the
kernel's centre tap, scaled by s; every other tap is zero.
Only the centre tap, and that is the point: conv2d pads with zeros, so a conv of a constant
is not constant near the border — but the centre tap kh = (kH−1)/2 reads position hi
itself, which is in range at every output cell. So this kernel is transparent to a uniform
offset at every position, which a multi-tap kernel is not.
Equations
Instances For
A zeroed conv (zero kernel, zero bias) sends everything to the constant 0 — the residual
bodies. flatConv_eq_zero lifted across the batch.
A zeroed STRIDED conv sends everything to the constant 0 (decimation of a constant).
The centre-tap depthwise kernel: every channel reads itself through the kernel's centre
tap, scaled by s. At s = 1, b = 0 it is the identity. The depthwise peer of ctK.
Equations
Instances For
Cellwise ⇒ flatwise. A property of every cell (n, c, i, j) holds at every flat index —
the bridge from the bcell view back to the ∀ k shape every clause is stated in.
A channel-constant BN parameter. Every structural witness's γ and β are one of these,
which is what lets bnBatchLA_pointwise's property be channel-independent.
Equations
- Proofs.BatchSeal.kv c x x✝ = x
Instances For
The zero kernel — every residual body.
Equations
- Proofs.BatchSeal.zk oc ic kH kW x✝³ x✝² x✝¹ x✝ = 0
Instances For
The zero depthwise kernel — an inverted-residual net's zeroed bodies.
Equations
- Proofs.BatchSeal.dzk c kH kW x✝² x✝¹ x✝ = 0
Instances For
1 · √n < 160 whenever n < 25600 — the margin bnBatchLA_pos consumes, at γ = 1,
β = 160. Every BN width of a 224×224 ResNet witness clears it (the widest is the stem's
2·112² = 25088).
|1/64|·√n < 3 whenever n < 36864 = (3·64)² — the relu6 margin bnBatchLA_smooth6 consumes,
at γ = 1/64, β = 3. It clears both of that lemma's hypotheses at once, β = 3 being the
centre of (0, 6). Every relu6 BN width of a 224×224 MobileNetV2 witness fits: the widest is
2·112² = 25088, shared by the stem, b1's depthwise and b2's expand.
The perturbation: all of example 0's channel 0. Uniform over the spatial grid, so it
survives a max-pool for every t (maxPool3s2_shift) with no argmax argument.
Equations
Instances For
The witness input, as a ray through the base.
Equations
- Proofs.BatchSeal.rayX H W t = Proofs.BatchSeal.rayBase✝ H W + t • Proofs.BatchSeal.rayV H W
Instances For
The ray is continuous in its parameter.
The carrier: example 0's slab is example 1's plus the per-channel constant δ. The
replacement for a channel difference, which per-channel batch BN annihilates.
δ is a FUNCTION of the channel, not one scalar, and that is what makes it cheap: BN
multiplies channel c's offset by γ_c · istd_c with no need to prove the channels share an
istd, a centre-tap conv collapses the whole function to fun _ => s · δ 0, and only δ 0 is
ever read (by the head, and at each channel-changing conv).
Equations
- Proofs.BatchSeal.EDiff δ v = ∀ (ci : Fin c) (i : Fin h) (j : Fin w), Proofs.BatchSeal.bcell v 0 ci i j = Proofs.BatchSeal.bcell v 1 ci i j + δ ci
Instances For
Batch BN scales the carrier by γ_c · istd_c — the two examples share the channel's
mean and istd, so centring keeps their difference (bnBatchLA_exdiff).
A centre-tap conv copies channel 0's offset to every output channel, and is transparent to the zero padding because only the centre tap is nonzero.
The strided peer of eDiff_conv.
The XLA-SAME strided peer of eDiff_conv — MobileNetV2's stem.
A centre-tap depthwise scales the carrier channel by channel. Unlike eDiff_conv, which
collapses δ to the single value s * δ c₀ at every output channel, a depthwise reads only its
own channel, so the whole function δ survives, scaled.
The 3×3/s2 pool keeps the carrier, at every t.
The stem pool has no tie at the witness: BN is injective within a channel
(bnBatchLA_cell_inj) and the pre-BN activation is positionally injective. Stated in the
∀ example, MaxPool3s2Smooth (slab) shape the nets' *PoolSmoothAt unfolds to.
GAP and the dense head deliver the carrier to one class: GAP of a uniformly shifted
channel is shifted by the same constant, and a Wd that reads channel c₀ into class j
turns the per-channel carrier into δ c₀. Stated on the batchMaps a *HeadB unfolds to, so
every net's head instantiates it.
A centre-tap 1×1 (or kH×kW) conv, then its BatchNorm: the carrier is channel c₀'s offset
times γ0 · istd, at every output channel.
The strided peer of eDiff_convBn.
A centre-tap depthwise conv, then its BatchNorm: channel by channel, δ ci · γ0 · istd.
The strided peer of eDiff_dwBn.
The XLA-SAME strided peer of eDiff_dwBn.
Stage facts first needed by ResNet-34's seal, shared by every conv-net seal #
The zero-kernel collapse, the relu-free strided stage, the centre-tap projection witness
(sealProj). Stage continuity needs no lemma here: the op files tag their continuity facts
@[fun_prop], and each seal's Rr_continuous is one fun_prop. The namespace is ResNet-34's, kept so that every
citation keeps its name.
A zeroed final conv makes a body the constant β₂ — projB at a zero kernel is
bnBatchLA of the constant 0, which is β₂ (variance 0). Used for both block kinds.
The structural downsample's collapsed form: its centre-tap projection.
Equations
- Proofs.BatchSeal.sealProj N h w ic oc = Proofs.StableHLO.projStridedB N (Proofs.BatchSeal.ctK oc ic 1 1 1) (Proofs.BatchSeal.kv oc 0) 1 (Proofs.BatchSeal.kv oc 1) (Proofs.BatchSeal.kv oc 160)
Instances For
A strided conv-bn-relu stage whose BN is everywhere positive has no relu left.
sealProj, unfolded — bn of the centre-tap strided conv.