The batch-BatchNorm seal kit — non-degeneracy machinery for the full-width nets #
planning/full_width_seals.md §3. The level-2/3 witnesses (Training/JacobianSeal.lean's
backward_nontrivial_of_fderiv_ne) have until now been exhibited on 2-channel per-example
proxies, whose BatchNorm 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 (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. - §4 the 3×3/s2 pool: it shifts with a uniform offset (
maxPool3s2_shift, no argmax argument), and it preserves nonnegativity. - §5 continuity odds and ends for the ray argument (
relu,bnIstd;relu6_continuoussits withrelu6itself). - §11 the SYMMETRIC strided depthwise (
EDiff_dwS2), MobileNetV4's; §3c's is the XLA one. - §12–§14 what MobileNetV4's swish needs, and nothing else does. ⭐⭐
EDiffcarries only the gap between the two examples, which is all a relu-in-the-window or a centre-tap conv reads. A stage that is smooth but NOT affine changes that gap by an amount depending on the values themselves, so the carrier has to know them:BUnifsays each example's slab is constant over the grid, one value per channel. Every op in these nets preserves that, andbnBatchLAthen puts the two values symmetrically aboutβ(bnBatchLA_pair) — so the swish's two outputs, and hence their gapswishGap, are functions of the gap alone.EDiff_of_BUnifhands the carrier back to §8 on the far side.
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
Adding a batch-uniform vector shifts every example's slab the same way — the shape a zeroed residual body contributes, and the reason it is transparent to the example difference.
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 one arithmetic fact: ((c,i),j) and (c,(i,j)) are the same offset, so the
mul_assoc cast that defines bnBatchLA sends the network cell to the bnBatchTensor4 cell.
finProdFinEquiv is row-major, so both sides are j + w·i + h·w·c + oc·h·w·n.
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
⭐⭐ The bridge: a cell of a bnBatchLA output is the scalar bnForward of that cell's
channel row, at that cell's position in the row. Everything else in this file is BN algebra on
one row.
⭐ The workhorse. A property of every cell of a bnBatchLA output, reduced to the scalar
bnForward on each channel's row — no index decomposition at the call site. Every clause of
the shape "this BN output is off the kink / inside a window" goes through here.
BN acts on coordinate differences by γ·istd — the exact identity that propagates the
carrier undamped through every BN, and the reason no BN-variance derivative is ever taken:
the difference is γ · (difference) · istd with istd evaluated at the same activation.
The two-sided BN margin |bn − β| ≤ |γ|·√n, with no mean/variance computation
(bnXhat_sq_le). bnForward_lb's symmetric form; what makes a large β keep a relu off its
kink at every input, so the structural net needs no eventually-argument for its relus.
A constant channel normalizes to β (variance 0, xhat = 0): a zeroed residual body is
the constant β₂, whatever γ₂ is. The batched peer of bnForward_const.
The batched margin: with γ, β channel-constant, every cell is within |g|√(N·h·w)
of b.
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.
Batch BN is injective within one example and channel (γ_c ≠ 0): it is the strictly
monotone affine map γ_c·istd_c·(· − μ_c) + β_c there. The stem pool's no-tie discharge —
equal pooled cells force equal pre-BN cells, and the witness's ramp is positionally injective.
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
The centre-tap conv value: b o + s · (input channel 0 at the same position), at every
output channel. ⚠ c₀ is passed in (rather than built from 0 < ic) so that the carrier's
channel index is the same term at every use site — ring needs those istds to be one atom.
The flat centre-tap conv, in cell coordinates.
The strided centre-tap conv: b o + s · (input channel 0 at the even position).
The batched centre-tap conv, in cell coordinates — the carrier's conv step at stride 1 (ResNet-50's stage-1 projection is the one site that needs it).
The batched strided centre-tap conv, in cell coordinates — the carrier's conv step.
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).
Odd decimation reads position (2i+1, 2j+1), channel for channel — the peer of
decimate_unflatten.
The strided XLA-SAME centre-tap conv: b o + s · (input channel 0 at the ODD position).
flatConvStride2Xla is decimateOddFlat ∘ flatConv, so this is flatConvStride2_ctK with
decimateOdd_unflatten in place of decimate_unflatten.
The batched strided XLA-SAME centre-tap conv, in cell coordinates — MobileNetV2's stem.
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
The flat centre-tap depthwise, in cell coordinates.
The batched centre-tap depthwise — the carrier's depthwise step at stride 1.
The strided XLA-SAME centre-tap depthwise, at the odd position.
A zeroed depthwise (zero kernel, zero bias) sends everything to the constant 0 — the
residual bodies of an inverted-residual net. ⛔ No strided peer is needed: a net's strided
blocks change channels, so they are on the carrier and never carry a zeroed kernel.
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.
⭐ The 3×3/s2 pool shifts with a uniform offset. If one slab's channel is another's plus
the constant δ, so are their pooled values — max of a uniformly shifted family
(Finset.apply_sup'_eq_sup'_comp at (· + δ)). No argmax or eventually-argument: this holds at
every point of the ray, which is what lets the carrier cross the only real kink in the net.
relu is continuous everywhere — it is max · 0; only its derivative has a kink.
A residual branch is continuous when its body is.
A projected residual is continuous when both branches are.
bnIstd is continuous in the activation (ε > 0).
maxPool3s2Flat is continuous (a sup' of coordinates).
batchMap of a continuous per-example op is continuous.
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 base slab: channel 0 carries the strictly decreasing ramp −(i·W + j) — positionally
injective, which is a stem pool's no-tie condition — and the other two channels are zero (a
centre-tap stem reads only channel 0). Both examples carry the same slab, so the carrier
vanishes at t = 0.
Instances For
Equations
- Proofs.BatchSeal.rayBase H W = Proofs.BatchSeal.bfrom fun (x : Fin 2) => Proofs.BatchSeal.rayRamp H W
Instances For
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 strided XLA-SAME peer of EDiff_dw. The carrier is spatially uniform, so decimation —
whichever phase it keeps — is transparent to it.
The 3×3/s2 pool keeps the carrier, at every t.
⭐ The pre-BN stem activation is positionally injective within each example and channel:
the centre tap decimates the ramp, and example 0's uniform +t shifts every position alike.
⭐ 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.
the two examples' swish outputs, as a function of half their gap.
Equations
- Proofs.BatchSeal.swishGap β u = Proofs.swishScalar (β + u) - Proofs.swishScalar (β - u)
Instances For
with ε = 1 a batch istd is at most 1, which keeps the ray's gap inside the window
swishGap_pos needs.