4 CIFAR with BatchNorm

This chapter is the bridge between the MNIST chapters and the deep-network chapters that follow. Chapter 3 did MNIST with a two-convolution net and a 2\(\times \)512 dense head. Chapter 5 (ResNet-34) goes thirty-four layers deep. The network here sits exactly in between. It keeps MNIST’s same 2\(\times \)512 head and the same conv/ReLU/max-pool machinery, but stacks the convolutions four stages deep (eight in all) on a harder dataset (CIFAR-10). That is the first point in the book where depth is enough to make two things bite: BatchNorm starts to earn its keep, and the choice of optimizer starts to matter. The depth itself is what ResNet then scales to thirty-four layers. So the chapter has two jobs: prove the one new operator (BatchNorm), and use the deeper net to measure what actually governs training (§4.3).

Proving BN is also what makes this structurally the hardest chapter in the book. BN’s inverse-stddev term \(1/\sqrt{\sigma ^2 + \varepsilon }\) has a gradient that blows up as \(\sigma ^2 \to 0\). Proving the backward pass exists and is bounded requires ContinuousLinearMap-based real-analysis machinery from Mathlib that none of the previous chapters needed. If you’re new to formal math, skim the proofs and trust them. The takeaways are concrete:

• BN’s gradient has a closed-form 3-term formula (Theorem 34).

• The formula needs \(\varepsilon {\gt} 0\) to stay bounded (Theorem 32).

• BN’s payoff is speed and stability: a deeper network reaches the same accuracy in far fewer epochs with it than without, and under a strong optimizer the un-normalized net does not reliably finish training at all. The example in §4.3 measures both directly.

The proofs themselves use Mathlib’s HasFDerivAt.sqrt, (hasDerivAt_inv).comp_hasFDerivAt, and a centering CLM chained through the chain rule from Ch 1. They’re correct (the Lean kernel checks them) and they’re available in Proofs/BatchNorm.lean for the curious. You do not need to understand them line-by-line to use BN as a layer or to follow the rest of the book. Ch 5 (ResNet-34) is the easiest chapter in the book and follows immediately. This is a localized difficulty spike, not the new normal.

What BN actually is

BatchNorm (Ioffe & Szegedy, 2015) takes a batch of activations, \(x\), and does three things in sequence. Each is one line of code. Together they are the layer.

1. Center. Compute the batch mean \(\mu = \frac{1}{n} \sum _k x_k\) and subtract it from every sample: \(x - \mu \). Output has mean zero.

2. Normalize. Compute the batch variance \(\sigma ^2\), add a small \(\varepsilon \) for numerical safety, and divide: \(\hat{x} = (x - \mu )/\sqrt{\sigma ^2 + \varepsilon }\). Output has unit variance.

3. Affine. Scale and shift with learnable per-channel parameters \(\gamma \) and \(\beta \): \(\mathrm{bn}(x) = \gamma \hat{x} + \beta \).

The chapter’s first six theorems are the Jacobian of each step and the VJPs that fall out of them. Theorem 36 closes the loop by composing them. The forward is three steps, so the chain rule from Chapter 1 gives us three Jacobians to multiply, and the “BN three-term backward” is exactly that product written out.

4.1 Run it first

Before any of the math, train the thing. Four commands and about thirteen minutes of GPU time:

lake exe cache get                            # Mathlib oleans, ~30 s
./download_cifar.sh                           # ~170 MB
lake build cifar8w-bn-ablation
./.lake/build/bin/cifar8w-bn-ablation data

That one binary runs the same network three times, once per optimizer. Here is the second of the three, on one RTX 4060 Ti (CUDA 12.9), from runs/2026-09-01-cifar8w-bn-xla-cuda/cifar8w-bn.log. XLA’s startup banner is removed, the per-step loss lines are dropped, and epochs 6 through 36 are elided:

[pjrt_ffi] XLA backend: PJRT 0.114, 1 device(s)
[pjrt_ffi] compiled verified_mlir/cifar8w_bn_mom_train_step.mlir
             (@cifar8w_bn_mom_train_step, 117 outputs, 1 replica) in 1061 ms
[pjrt_ffi] compiled verified_mlir/cifar8w_bn_fwd.mlir
             (@cifar8w_bn_fwd, 1 outputs, 1 replica) in 237 ms
════════ cifar8w-BN (wide head) — Nesterov momentum (μ.9, lr 0.02) ════════
Deeper CIFAR-10 CNN + per-channel BatchNorm, MNIST-style wide head
  (8× conv→BN→relu → 128→512→512→10) via the VERIFIED renderer → XLA/PJRT → GPU
  xla/pjrt verified_mlir/cifar8w_bn_mom_train_step.mlir
  xla/pjrt verified_mlir/cifar8w_bn_fwd.mlir
  train 50000, test 10000; bs 128, CIFAR-CNN8-wide-BN mom
    (constant lr 0.020000), He init
Epoch 1/40: loss=1.611771 lr=0.020000
  epoch 1: test_acc = 4927/10000 = 49.270000%  top5 = 9074/10000 = 90.740000%
             [95% CI 48.29–50.25]
Epoch 2/40: loss=1.207551 lr=0.020000
  epoch 2: test_acc = 6135/10000 = 61.350000%  top5 = 9543/10000 = 95.430000%
             [95% CI 60.39–62.30]
Epoch 3/40: loss=1.027676 lr=0.020000
  epoch 3: test_acc = 6469/10000 = 64.690000%  top5 = 9660/10000 = 96.600000%
             [95% CI 63.75–65.62]
Epoch 4/40: loss=0.916538 lr=0.020000
  epoch 4: test_acc = 6653/10000 = 66.530000%  top5 = 9682/10000 = 96.820000%
             [95% CI 65.60–67.45]
Epoch 5/40: loss=0.849334 lr=0.020000
  epoch 5: test_acc = 6868/10000 = 68.680000%  top5 = 9731/10000 = 97.310000%
             [95% CI 67.76–69.58]
Epoch 37/40: loss=0.373874 lr=0.020000
  epoch 37: test_acc = 7550/10000 = 75.500000%  top5 = 9792/10000 = 97.920000%
             [95% CI 74.65–76.33]
Epoch 38/40: loss=0.365542 lr=0.020000
  epoch 38: test_acc = 7565/10000 = 75.650000%  top5 = 9812/10000 = 98.120000%
             [95% CI 74.80–76.48]
Epoch 39/40: loss=0.354800 lr=0.020000
  epoch 39: test_acc = 7585/10000 = 75.850000%  top5 = 9808/10000 = 98.080000%
             [95% CI 75.00–76.68]
Epoch 40/40: loss=0.349538 lr=0.020000
  epoch 40: test_acc = 7594/10000 = 75.940000%  top5 = 9817/10000 = 98.170000%
             [95% CI 75.09–76.77]
done (trained CIFAR-CNN8-wide-BN mom + constant lr via packed threading).

Forty epochs at about six seconds each, and 75.94% on the full 10,000-image CIFAR-10 test set, 98.17% inside the top five. CIFAR-10 is a much harder problem than MNIST, so this is the first chapter where the headline accuracy drops out of the nineties. BatchNorm’s backward is also the first in this book you can’t read off by inspection.

The learning rate is constant here — no warmup, no decay — which is not what a CIFAR recipe would normally do and is deliberate: §4.3 spends the rest of the chapter changing one thing at a time, and a schedule is one more thing changing. Chapter 5 adds cosine annealing back and shows what it buys. Run five times at this setting the net lands between \(75.8\) and \(76.7\% \), so the number above is an ordinary draw rather than a best one, and the tables later in the chapter report medians over five for the same reason.

§4.3 runs this same binary and its no-BN peer across all three optimizers, which is where the chapter’s measurements come from, and lake run cifar is those two binaries.

The centering term has an indirect path

When you jiggle a single input \(x_i\), you do not only jiggle \(x_i\). You also jiggle the batch mean \(\mu \), and \(\mu \) appears in every sample’s centered value \(x_k - \mu \). So jiggling \(x_i\) by \(\varepsilon \) shifts every centered value by \(-\varepsilon /n\) plus the direct \(\varepsilon \) on the \(i\)th sample.

The \(\delta _{ij}\) is the direct effect. The \(-1/n\) is the indirect effect through \(\mu \). Theorem 31 states this formally. Forgetting the \(-1/n\) is the most common hand-derivation mistake on BN, and is exactly what the formal proof prevents.

The inverse-stddev term

The normalize step divides by \(\sqrt{\sigma ^2 + \varepsilon }\), where \(\sigma ^2 = \frac{1}{n} \sum _k (x_k - \mu )^2\) is itself a function of every input. Jiggle \(x_i\) by \(\varepsilon \) and \(\sigma ^2\) changes, which means the divisor changes, which means every output \(\hat{x}_j\) changes, not just \(\hat{x}_i\).

Working through the chain rule (\(x \mapsto x^2 \mapsto \text{mean} \mapsto \sqrt{\cdot + \varepsilon } \mapsto 1/\cdot \)) gives

That \(\mathrm{istd}^3\) is what makes the BN backward expensive: the gradient of one sample depends on every sample’s centered value, scaled by the cube of the inverse standard deviation. Theorem 33 states this.

Notice what happens to the formula as \(\sigma ^2 \to 0\): \(\mathrm{istd} \to 1/\sqrt{\varepsilon }\), bounded. Without \(\varepsilon \) the gradient diverges. Theorem 32 is the formal statement that \(\varepsilon {\gt} 0\) is sufficient to make this term differentiable. This is the one place in the book where the math actually requires real-analysis machinery beyond chain-sum-product. Everything else in the framework reduces to those three rules.

The three-term backward

Compose the three forward steps and apply the product rule on \(\hat{x} = (x - \mu ) \cdot \mathrm{istd}\). The cross-terms collapse (the centered sum \(\sum _k (x_k - \mu ) = 0\) is what saves us) and what falls out is a one-line backward:

Three terms inside the parentheses, one per forward step’s indirect effect:

• \(n\, dy\): the direct effect, every sample’s upstream gradient.

• \(-\sum _k dy_k\): the centering correction, subtracts the total upstream gradient because shifting the mean shifts every sample.

• \(-\hat{x} \cdot \sum _k (dy_k\, \hat{x}_k)\): the normalization correction, subtracts the projection of the upstream gradient onto \(\hat{x}\), because rescaling by the inverse stddev couples every sample’s gradient through the shared divisor.

Theorem 34 formalizes this. Every production BN implementation (PyTorch, JAX, TensorFlow, custom CUDA kernels) computes this exact expression. The formula has been known since 2015. The value of the formal proof is that the Lean kernel mechanically verifies we are computing the right thing at every training step.

The affine step is dense-with-broadcasting

The third step, \(\gamma \hat{x} + \beta \), is structurally a dense layer applied per-channel. Its Jacobian \(\partial (\gamma v + \beta )/\partial v_i = \gamma \delta _{ij}\) is exactly the dense-Jacobian computation from Chapter 2, just lifted to a tensor shape and broadcast across spatial dimensions. Theorem 30 and Theorem 35 are essentially corollaries of the dense theorems. The new content here is zero.

Putting it together

Theorem 36 is the composition: \(\mathrm{bn} = \mathrm{affine} \circ \mathrm{normalize}\), with VJPs chained via the same \(\mathrm{vjp\_ comp}\) rule from Chapter 1. The 3-term backward is the centerpiece, affine is a corollary, and composition is a one-line proof. The structural story of the chapter is: one new Mathlib-level analytic dependency (sqrt and recip differentiability), one formula, and the rest is the same chain rule we already had.

4.2 The theorems

4.3 Example: training dynamics on CIFAR

Chapter 3 did MNIST with two convolutions and a 2\(\times \)512 dense head. This chapter keeps that exact head and the same conv/ReLU/max-pool machinery, but stacks the convolutions four stages deep, eight \(3 \times 3\) convolutions in all, because CIFAR-10 is a genuinely harder problem. Same head, four times the body. That extra depth is the whole point of the chapter: it is where BatchNorm starts to pay, where the optimizer starts to matter, and the same depth ResNet-34 will scale to thirty-four layers. It also lets us ask a sharper question than “does it train?” We can ask what governs how it trains?

Two levers govern the answer, normalization and the optimizer; a third, the arithmetic the whole thing is computed in, turns out not to move it at all — which is its own kind of finding — and a knob, the width of the head, does not either. We move each against an identical, machine-checked gradient, holding everything else fixed: same architecture, same 40 epochs, and the same data pipeline (per-epoch shuffle, random horizontal flip, cosine learning-rate schedule with warmup). The findings, stated up front so the graphs are no surprise:

• Normalization buys stability, and speed on the way. BatchNorm (dropped between each conv and its ReLU) reaches a given accuracy in markedly fewer epochs. It also keeps the net alive: the un-normalized net’s training loss went to NaN in six of our fourteen runs, and the normalized net’s never did.

• The optimizer moves the ceiling most. Trading plain SGD for SGD-with-momentum is worth about two points of final accuracy, more than anything else here. AdamW sits between the two, and its per-coordinate adaptivity earns less than momentum does.

• The arithmetic does not move the ranking. Recomputing the normalized ladder in bf16 leaves the optimizer ordering exactly where it was, with medians moving by \(+0.31\), \(-0.02\) and \(-0.33\) — every one of them smaller than the spread between repeated runs of a single configuration. Drop the normalization and bf16 starts to show: the medians move by up to \(1.4\) points and one seed collapses. The eight-layer stack is fragile, and the same runs show it, but it is fragile in fp32 too. What breaks the network is depth without normalization, not the number format.

• The head barely matters. This 2\(\times \)512 head carries about 25\(\times \) the parameters of a narrow 64-wide one, yet trains to within a point of it. The convolutional body, not the head, is doing the work. That is also why we can borrow MNIST’s head wholesale: it was never the bottleneck.

Every run below shares one proof-rendered backward pass, and only the BN layers or the rendered optimizer tail change. The numbers are from that verified training step running on a GPU.

On repeatability. Convolutional networks do not reproduce run to run on CUDA, because XLA selects convolution algorithms per process, so two runs of the identical command differ. Every accuracy in this section is therefore the median of five independent runs of the same command at the same seed (four for the AdamW rows), and each table states the observed range alongside it. That spread is about a point, which is wide enough to swallow small differences, so the text below only ranks gaps that are clearly outside it.

Architecture

The same vertical column as the MNIST CNN in Chapter 3, just deeper: eight \(3 \times 3\) convolutions in four stages (a max-pool after each) lift the \(32 \times 32 \times 3\) input to a \(2 \times 2 \times 32\) feature grid, which flattens to 128 and runs through the same 2\(\times \)512 dense head as that net (only the first layer changes width, \(6272\) vs. \(128\), because the deeper stack pools the grid down further). Eight convolutions is four times the MNIST net’s two, and that growth in depth, not width, is the lineage ResNet-34 (Ch 5) carries to thirty-four layers.

The ConvBN boxes are the with-BN variant. The no-BN net is identical with each ConvBN replaced by a plain Conv2D + ReLU. Both specs are written out next.

The two specs, differing by one keyword per layer

Eight \(3 \times 3\) convolutions in four stages (channel widths \(16, 16, 32, 32\), a max-pool after each stage), then the 2\(\times \)512 dense head lifted straight from the MNIST CNN. Without BN, each convolution is followed by a plain ReLU. With BN, a per-channel BatchNorm sits between the convolution and the ReLU. That one keyword per conv layer is the entire difference.

Without BN:

def cifar8wVerified : VerifiedNetSpec where
  name     := "CIFAR-CNN8-wide"
  slug     := "cifar8w"
  inC      := 3
  imageH   := 32
  imageW   := 32
  nClasses := 10
  data     := .cifar
  layers   := [.conv 3 16 3 1, .relu,
               .conv 16 16 3 1, .relu, .maxPool 2 2,
               .conv 16 16 3 1, .relu,
               .conv 16 16 3 1, .relu, .maxPool 2 2,
               .conv 16 32 3 1, .relu,
               .conv 32 32 3 1, .relu, .maxPool 2 2,
               .conv 32 32 3 1, .relu,
               .conv 32 32 3 1, .relu, .maxPool 2 2,
               .flatten,
               .dense 128 512, .relu, .dense 512 512, .relu, .dense 512 10]
  blurb    := "Deeper CIFAR-10 CNN, MNIST-style wide head (8 convs,
                [16,16,32,32], 4 pools 32→2 → 128→512→512→10)
                via the VERIFIED renderer → %LOWERER% → GPU"

With BN:

def cifar8wBnVerified : VerifiedNetSpec where
  name     := "CIFAR-CNN8-wide-BN"
  slug     := "cifar8w_bn"
  inC      := 3
  imageH   := 32
  imageW   := 32
  nClasses := 10
  data     := .cifar
  layers   := [.conv 3 16 3 1, .bnPerChannel 16, .relu,
               .conv 16 16 3 1, .bnPerChannel 16, .relu, .maxPool 2 2,
               .conv 16 16 3 1, .bnPerChannel 16, .relu,
               .conv 16 16 3 1, .bnPerChannel 16, .relu, .maxPool 2 2,
               .conv 16 32 3 1, .bnPerChannel 32, .relu,
               .conv 32 32 3 1, .bnPerChannel 32, .relu, .maxPool 2 2,
               .conv 32 32 3 1, .bnPerChannel 32, .relu,
               .conv 32 32 3 1, .bnPerChannel 32, .relu, .maxPool 2 2,
               .flatten,
               .dense 128 512, .relu, .dense 512 512, .relu, .dense 512 10]
  blurb    := "Deeper CIFAR-10 CNN + per-channel BatchNorm, MNIST-style
                wide head (8× conv→BN→relu → 128→512→512→10)
                via the VERIFIED renderer → %LOWERER% → GPU"

Each layers list and each blurb is one line in the source, wrapped here; the second blurb is the banner §4.1’s transcript opens with. The diff is eight .bnPerChannel entries, one inserted between each convolution and its ReLU. The dense head, the max-pools, and the training config are all identical.

What .bnPerChannel contributes.

One arm, and the only new one this chapter needs — .conv, .maxPool, .flatten, .dense and .relu are all Chapter 3’s, unchanged:

  | bnPerChannel oc => #[(#[oc],1),(#[oc],2)]   -- gamma (ones), beta (zeros)

This is where initKind stops being bookkeeping. BatchNorm’s \(\gamma \) initialises to ones and \(\beta \) to zeros, so a freshly built net’s BN is the identity and the surrounding convolutions see exactly what they would have seen without it. Zero-init \(\gamma \) instead and the whole trunk outputs \(\beta \) regardless of its input. Note also what is not here: the running mean and variance. Those are statistics, not parameters — no gradient reaches them — which is why they are threaded separately through bnChannels rather than living in this list.

The spec-to-proof tie. These are VerifiedNetSpecs, which is what makes the listing above load bearing rather than descriptive. The slug names the committed render, so cifar8w_bn is verified_mlir/cifar8w_bn_{mom,sgd,adam}_train_step.mlir, the files the run in §4.1 handed to XLA. The layers list folds to a parameter layout through toSpecs, and a #guard in the same file pins that layout against the one the renderer emits, so a spec that drifts from its render fails the build rather than training something else. The gradient itself is Proofs.cifarCnn8_has_vjp_at, proved once and parametrically in the head width, which is why the \(512\)-wide head here and the \(64\)-wide one behind cifar8-bn-verified need no separate proof between them. cifarCnn8_has_vjp_at_correct (Proofs/Nets/Small/CifarCNN.lean) is the statement that its backward is the derivative, and Proofs/Cifar8Close.lean carries that through to the emitted training step. Both nets run the same XLA/PJRT path on the GPU.

Lever 1: normalization

Fix the optimizer at plain SGD and toggle BN. Both nets train for 40 epochs on the shared pipeline at a constant learning rate — no warmup, no decay — because a schedule is a fourth thing changing and this section is about the third. Per-epoch test accuracy, the median of five runs each, from runs/2026-09-01-cifar8w-6arm-constlr/:

CIFAR-10, wide 8-conv net, plain SGD at a constant lr 0.1 on the shared pipeline, 40 epochs, BN vs no-BN, through the verified renderer. Median of \(n{=}5\) runs per curve. The final points span 74.1 to 75.2 with BN and 68.0 to 70.7 without.

BN leads from the first epoch — 41% against the bare netś 13% — and never gives the lead back, finishing 5.8 points up (74.5% vs 68.8%). The un-normalized net traces the same arc roughly ten epochs behind, and the gap never closes: it is still losing at epoch 40 by more than the spread of either set of runs.

That first epoch is worth a second look. Thirteen per cent is barely above chance on ten classes, and the un-normalized net spends four epochs getting to where the normalized one starts. The conditioning the three-term backward (§34) buys is what that is, and it is worth more the deeper the stack — exactly the trend that makes BN standard equipment by ResNet’s thirty-four layers.

Lever 2: the optimizer

Now hold the architecture fixed and change only the update rule. Three optimizers, each at its own tuned learning rate, all on the identical pipeline, the same 40 epochs and the same constant rate. Median final accuracy, with the observed range across five runs:

\(n=5\) per cell. Parenthesised figures are the range from lowest run to highest.

Momentum with BN is the best result on the board at 76.3%, about two points over plain SGD and two over AdamW, and that gap is wider than the spread within any cell. AdamW lands level with SGD rather than between the two: its per-coordinate second-moment scaling is worth something on the un-normalized net and nothing once BN is doing the conditioning.

Reading down the columns recovers Lever 1, and this time it recovers it everywhere. BN is ahead in all three — \(+5.8\) under SGD, \(+4.2\) under momentum, \(+1.5\) under AdamW. The benefit shrinks as the optimizer gets stronger, which is the familiar reading: a better update rule can partly substitute for better conditioning, and never entirely.

Where the failure actually is. Not in the accuracies, which all look reasonable, but in the losses behind them. Three of the five un-normalized AdamW runs sent their training loss to NaN and still reported 72.8% at epoch 40 — a number that would have gone in the table unremarked if we had only scored accuracy. None of the SGD or momentum runs did it, with or without BN. The un-normalized net is not unstable in general; it is unstable under the optimizer with the most internal state to corrupt, and only the loss says so.

The reason we can make this comparison and trust it is that the optimizer is one swappable rendered tail. The forward pass, the softmax–cross-entropy loss, the backward pass, and every parameter gradient are the same proof-rendered graph in each column. Only the final per-parameter update op changes:

• SGD: \(\theta \leftarrow \theta - \mathrm{lr}\cdot g\).

• Momentum (Nesterov): \(v \leftarrow \mu v + g\), then \(\theta \leftarrow \theta - \mathrm{lr}\, (\mu v + g)\).

• AdamW: the bias-corrected first/second-moment step, rendered op-for-op as Proofs.adamWParam.

Each tail is emitted onto the same certified gradient (emitSgd, emitMomentum, ViTRender.emitAdamV), so the ablation is honest in the strong sense: identical, machine-checked gradients with different arithmetic stacked on top. And for plain SGD, that the binary32 step actually decreases the loss is itself a proved theorem.

It is worth being precise about what that does and does not cover, because the AdamW column just exercised the gap. The proofs say the emitted graph computes the gradient this chapter derives, and they say it in exact real arithmetic. They say nothing about whether forty epochs of binary32 updates stay in range. A verified gradient is not a guarantee of numerical stability, and three NaN losses under a correct AdamW step are what that distinction looks like when it bites.

One caution about what this table measures. An ablation measures the thing you varied only if everything else is held constant, and “everything else” includes the data pipeline, the learning-rate schedule and the op family the graph is rendered through. All three live outside the network, none of them is visible to the proofs, and each of them is capable of moving these numbers by more than the lever does. The gradients being machine-checked does not make the experiment design correct.

Lever 3: the arithmetic

The obvious third lever is the number format. Weights and activations do not have to be float32: bfloat16 keeps float32’s exponent range and throws away sixteen bits of mantissa, so it holds roughly three significant decimal digits instead of seven while overflowing at the same place. Modern GPUs multiply it about twice as fast and store it in half the memory, which is why the ImageNet-scale runs from Chapter 5 onward are trained that way. (The Imagenette runs those chapters lead with stay in float32; bf16 is what the full-ImageNet jobs need, and the tables there say so per row.) The question this chapter would like to answer is whether that costs any accuracy.

It can, on the net that matters, and getting there took one change worth a paragraph. The renderer has two convolution op families — the per-example one these MNIST and CIFAR chapters use, and the batched one the ImageNet chapters use — and the bf16 twins were built for the batched family only. The normalized net was rendered per-example, so until it moved across there was no bf16 arm to run on it. Moving it needed no new arithmetic: BatchNorm itself stays float32, which is not a shortcut but the recipe every ImageNet network in this book uses — low precision through the convolutions and the dense layers, full precision through the normalization.

So: the eight-convolution net with BatchNorm, one renderer, one constant rate, forty epochs, five seeds, and precision the only thing that differs between a pair of rows. All twenty-three convolutions are bf16 in the low-precision arms — forward, input gradient and weight gradient alike. Six arms, five seeds each, from runs/2026-09-01-bnprec-seeds/:

Medians move by \(+0.31\), \(-0.02\) and \(-0.33\). Each of those is smaller than the spread within the corresponding fp32 row — \(1.89\), \(0.75\) and \(1.18\) — so on this net, at this scale, the number format is not measurable against the noise of rerunning the same configuration. Thirty arm-runs, no NaN, nothing collapsed to chance. That is the result Chapter 5 was going to assume, and it no longer has to.

The un-normalized net answers differently, and the difference is the interesting part. Run the same comparison there and the medians move by \(-1.40\), \(-0.56\) and \(-0.09\), and one bf16 momentum seed collapses to chance where no fp32 seed does. Normalization is doing more than accelerating training: it is absorbing the perturbation. A net that was already fragile in fp32 — the three NaN losses under a correct AdamW step from the previous section — is the one where sixteen fewer mantissa bits show up at all.

The same pattern holds for a lever we did not mean to pull. Changing the op family at fixed fp32 — per-example to batched, the same arithmetic throughout — moves the un-normalized net by \(+2.0\) on SGD and \(+1.7\) on momentum, which is larger than anything precision does to it. On the normalized net the identical change is worth \(-0.11\), \(+0.37\) and \(+0.19\). The rendering path is a lever, it is invisible to the proofs, and normalization damps it too.

An fp8 (E4M3) arm exists in the repository and is left out of this table on purpose: its render rounds the forward convolutions only and leaves the backward in float32, so it is a lowering probe rather than a trainable arm, and reporting it beside bf16 would imply a comparison that was never run.

Does the head width matter?

Almost not at all, which is exactly why we could borrow MNIST’s head untouched. The 2\(\times \)512 head carries about \(334{,}000\) of the net’s \(374{,}000\) parameters. Swap it for a narrow 64-wide head (\(13{,}000\) params, whole net down to \(53{,}000\)) and the best cell on the board barely moves: BN with momentum gives \(77.1\% \) wide against \(77.1\% \) narrow, and the optimizer ordering is unchanged. Seven times the parameters buys no accuracy, and costs about \(1.6\times \) the wall-clock per epoch (\({\approx }6.3\) vs \({\approx }3.9\) seconds on an RTX 4060 Ti). The eight convolutions over the \(32 \times 32\) maps dominate the compute, while the head is cheap matmul however big it is. That is the lesson the bridge makes concrete: at this scale the depth of the convolutional body is the lever, not the width of the classifier on top. It is why the head can stay MNIST’s, and why the next chapter spends its budget on thirty-four layers of more convolution rather than a bigger head.

The two-point comparison above (64 vs 512) is worth drawing out in full, because the parametric renderer makes it cheap: cifar8-bn-grid holds the eight-convolution backbone fixed and renders the AdamW train step at any head width \(d\) (the same \(D_1\) that was hard-wired to 64 is now a parameter of the verified emitter), so we can sweep \(d\) from 8 to 4096 and train each point on its own proof-rendered StableHLO. Split across the two gfx1100s, the whole curve is one short run:

Dense-head width sweep for the 8-conv cifar8 net with per-channel BatchNorm, AdamW, 25 epochs, conv backbone held at \([16,16,32,32]\) (runs/cifar8bn_grid_results.tsv). Each point is one run, so the whiskers are the 95% Wilson interval for that accuracy on the \(9{,}984\) evaluated test images — about \(\pm 0.9\) here — and not a spread over repeats. Read that way the curve makes a sharper claim than “essentially flat”: every width from \(d{=}16\) to \(d{=}4096\) spans only \(71.1\) to \(72.2\% \), which is narrower than a single point’s own interval, so the nine plateau intervals all overlap and none of the differences between them is resolvable from this experiment. The \(256\times \) wider head buys nothing measurable, while its \(17\)M-parameter classifier just overfits (train loss \(0.03\), test unmoved). The one difference that does survive is \(d{=}8\) (\(66.8\% \), interval \(65.9\)–\(67.7\)), which clears the lowest plateau bound of \(70.2\) by more than two points: there the \(128{\to }8\) first layer throttles the 128-dim feature map the convolutions produce. The canonical \(d{=}64\) (circled) sits squarely on the plateau: the head width genuinely does not matter here, which is the whole reason the net could borrow MNIST’s classifier untouched. (Absolute accuracy is a couple of points below the 40-epoch board above, because this is a 25-epoch single-optimizer sweep, but the shape is the point.)

Why the levers work

Both levers do the same underlying thing by attacking different sources of noise. They make each step’s gradient a more reliable guide to the next. Each layer is tuned for the distribution of its inputs, but those inputs are the outputs of every layer below, which shift every step, so each layer chases a moving target, and the step that helps one layer can wreck the next. BN removes the moving part by pinning every layer’s input to mean-zero, unit-variance before the learnable \(\gamma ,\beta \) get a say. Ioffe & Szegedy framed this as reducing “internal covariate shift,” while Santurkar et al. (2018) argued the sharper effect is a smoother loss landscape. Momentum attacks a different noise: by averaging successive gradients it cancels the mini-batch jitter and accumulates the consistent direction. Both make the per-step direction more trustworthy, and reliable progress compounds across epochs, which is what the curves and the table measure.

The two levers interact, and the momentum column is where that shows. Momentum makes each step longer and more consistent, which is why it wins on accuracy, and a longer step through eight unnormalized layers is also exactly what runs the activations out of binary32 range. Normalization is what makes the aggressive optimizer survivable. That is the standard account of why the two arrived together, and why every image architecture since 2015 bakes a normalization layer in by default. At larger depth the effect is stronger still, and it is what opens up the learning rates that diverge without it. Between them the two levers decide not only how many passes the net needs, but whether it finishes them at all.

4.4 MLIR: BatchNorm

What is already proven. BatchNorm factors as \(\mathrm{bnForward} = \mathrm{bnAffine} \circ \mathrm{bnNormalize}\), and its reverse-mode derivative is the three-term formula of § 34: with \(\hat{x} = (x-\mu )\, \mathrm{istd}\),

bn_has_vjp proves it, composing bnNormalize_has_vjp (the rank-1 wringer) with bnAffine_has_vjp (the \(\gamma \, dy\) half). The one subtlety, isolated in § 32, is that the inverse-stddev term carries an \(\mathrm{istd}^3\) that needs \(\varepsilon {\gt} 0\) to stay differentiable. That is the single place in the book where the math reaches past chain-sum-product into real analysis.

The gap and how we close it. The three-term backward is not elementwise, because the two \(\sum _k\) reductions couple every coordinate to every other. The emitted backward graph is given a denotation in the proofs’ own vector type, and bn_back_bridge proves that denotation equal to bn_has_vjp’s backward. The emitted reduce/broadcast/elementwise graph is, by machine check, the three-term formula. Here is what the printer emits, with the forward-statistic recompute (mean, variance, rsqrt, normalize) elided to its one comment line, at \(n=4\):

// forward stats recomputed: %mu, %istd, %xhat = (x-mu)*istd
%dxhat = stablehlo.multiply %gb, %dy : tensor<2x4xf32>   // dxhat = g*dy
%sdx_r = stablehlo.reduce(%dxhat init: %sc)
           applies stablehlo.add across dimensions = [1]
           : (tensor<2x4xf32>, tensor<f32>) -> tensor<2xf32>
%sdx = stablehlo.broadcast_in_dim %sdx_r, dims = [0]
           : (tensor<2xf32>) -> tensor<2x4xf32>          // sum dxhat
%xd = stablehlo.multiply %xhat, %dxhat : tensor<2x4xf32>
%sxdx_r = stablehlo.reduce(%xd init: %sc)
           applies stablehlo.add across dimensions = [1]
           : (tensor<2x4xf32>, tensor<f32>) -> tensor<2xf32>
%sxdx = stablehlo.broadcast_in_dim %sxdx_r, dims = [0]
           : (tensor<2xf32>) -> tensor<2x4xf32>          // sum xhat*dxhat
%t1 = stablehlo.multiply %dxhat, %nf : tensor<2x4xf32>   // N*dxhat
%i1 = stablehlo.subtract %t1, %sdx : tensor<2x4xf32>     //   - sum dxhat
%xs = stablehlo.multiply %xhat, %sxdx : tensor<2x4xf32>
%i2 = stablehlo.subtract %i1, %xs : tensor<2x4xf32>      //   - xhat*(sum)
%s = stablehlo.divide %istd, %nf : tensor<2x4xf32>       // istd/N
%dx = stablehlo.multiply %s, %i2 : tensor<2x4xf32>
return %dx : tensor<2x4xf32>

Read it against the formula. %dxhat is the affine backward \(\gamma \, dy\), and each reduce along dimensions = [1] followed by a broadcast_in_dim is one of the cross-coordinate sums (\(\sum _k dy_k\) as %sdx, \(\sum _k\hat{x}_k\, dy_k\) as %sxdx). %i1 assembles \(n\, dy - \sum dy\) (the direct term minus the centering correction), %i2 subtracts the rank-1 normalization correction \(\hat{x}\sum \hat{x}\, dy\), and %dx scales the whole bracket by \(\mathrm{istd}/n\). The graph folds \(\gamma \) into %dxhat up front, which is why that leading scale is \(\mathrm{istd}/n\) and not \(\mathrm{istd}\, \gamma /n\). It is the same formula with \(\gamma \) pulled inside the parenthesis. The bridge theorem is precisely the claim that this text computes bn_has_vjp’s backward.

Because LayerNorm is BatchNorm along a different axis, the very same emitted graph denotes the LayerNorm backward (layernorm_back_bridge is literally bn_back_bridge). That is the normalization sitting inside the residual, depthwise, and ConvNeXt blocks built on this foundation.

Caveats.

• Needs \(\varepsilon {\gt} 0\). The inverse-stddev term carries an \(\mathrm{istd}^3\). Given \(\varepsilon {\gt} 0\) the bridge is unconditional (smooth everywhere, with no smooth-point exclusion, unlike ReLU and max-pool).

• Representative scale (\(n = 4\)).

The next chapter (§ 5) adds residual connections and the same mechanical approach: prove that the VJP of a skip connection is additive fan-in, compose with BN and conv, and the rest of the ResNet family falls out without introducing any new math.