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 48).

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

• BN’s payoff is speed and accuracy: a deeper network reaches a given accuracy in far fewer epochs with it than without, and finishes three to eight points higher depending on the optimizer. Whether five runs of one command land close together is decided by something else, the group the statistics are taken over, and the example in §4.3 measures all three.

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/Architectures/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 \).

One thing to fix before the math, because the chapter’s measurements depend on it. Ioffe & Szegedy take \(\mu \) and \(\sigma ^2\) over the batch: every image in the mini-batch, at every spatial position, for one channel, \(N \times H \times W\) cells per statistic. That is the layer this chapter trains, with running statistics kept for evaluation, as ResNet-34 and every BatchNorm net after it do. The chapter also trains a second form of the same layer, per image: for each example and each channel, over that image’s \(32 \times 32 = 1{,}024\) spatial cells, with the batch axis never reduced. Nothing crosses from one example to another in that form, so it has no running statistics, computes the same thing in training and at evaluation, and keeps every op in the emitted graph per-example, which is what lets the step theorems below be stated one example at a time. The three steps and their Jacobians are the same for any group of values that is normalized together; \(n\) below is the group size, and the batched form, bnBatchTensor4, is Theorem 51. Where this chapter says “BatchNorm” without qualification it means the batch form; the per-image form is named each time it appears.

The chapter’s first six theorems are the Jacobian of each step and the VJPs that fall out of them. Theorem 50 closes the loop by composing them. The forward is three steps, so the chain rule (Theorem 2) 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 seven minutes of GPU time:

lake exe cache get                            # Mathlib oleans, ~30 s
./scripts/datasets/download_cifar.sh          # ~170 MB
lake build cifar8wb-bbn-ablation
./.lake/build/bin/cifar8wb-bbn-ablation data

That one binary runs the same network six times: three optimizers, each in fp32 and again in bf16. Here is the first of the six, plain SGD in fp32, on one RTX 4060 Ti (CUDA 12.9). The log is runs/2026-10-24-cifar8wb-bbn-xla-cuda/cifar8wb-bbn.log, with XLA’s startup banner removed, the per-step loss lines dropped, and epochs 6 through 36 elided:

[pjrt_ffi] XLA backend: PJRT 0.114, 1 device(s)
[pjrt_ffi] command buffers: enabled (CUDA default)
[pjrt_ffi] compiled verified_mlir/cifar8wb_bbn_sgd_train_step.mlir
             (@cifar8wb_bbn_sgd_train_step, 133 outputs, 1 replica) in 2620 ms
[pjrt_ffi] compiled verified_mlir/cifar8wb_bbn_fwd.mlir
             (@cifar8wb_bbn_fwd, 1 outputs, 1 replica) in 781 ms
[pjrt_ffi] compiled verified_mlir/cifar8wb_bbn_fwd_eval.mlir
             (@cifar8wb_bbn_fwd_eval, 1 outputs, 1 replica) in 198 ms
════════ cifar8wb-bbn (batch-statistics BN) f32 — SGD (lr 0.1) ════════
Wide-head CIFAR-10 CNN + per-channel BatchNorm at BATCH statistics
  (8 convs, BATCHED op family, d1=512, running-stat eval)
  via the VERIFIED renderer → XLA/PJRT → GPU
  xla/pjrt verified_mlir/cifar8wb_bbn_sgd_train_step.mlir
  xla/pjrt verified_mlir/cifar8wb_bbn_fwd.mlir
  xla/pjrt verified_mlir/cifar8wb_bbn_fwd_eval.mlir
  train 50000, test 10000; bs 128, CIFAR-CNN8-wide-BN-batchstat sgd
    (constant lr 0.100000), He init
Epoch 1/40: loss=1.497842 lr=0.100000
  epoch 1: test_acc = 4703/10000 = 47.030000%  top5 = 9113/10000 = 91.130000%
             [95% CI 46.05–48.01]
Epoch 2/40: loss=1.106227 lr=0.100000
  epoch 2: test_acc = 5861/10000 = 58.610000%  top5 = 9601/10000 = 96.010000%
             [95% CI 57.64–59.57]
Epoch 3/40: loss=0.957487 lr=0.100000
  epoch 3: test_acc = 6366/10000 = 63.660000%  top5 = 9637/10000 = 96.370000%
             [95% CI 62.71–64.60]
Epoch 4/40: loss=0.861241 lr=0.100000
  epoch 4: test_acc = 6334/10000 = 63.340000%  top5 = 9737/10000 = 97.370000%
             [95% CI 62.39–64.28]
Epoch 5/40: loss=0.794502 lr=0.100000
  epoch 5: test_acc = 5424/10000 = 54.240000%  top5 = 9053/10000 = 90.530000%
             [95% CI 53.26–55.21]
Epoch 37/40: loss=0.320061 lr=0.100000
  epoch 37: test_acc = 7603/10000 = 76.030000%  top5 = 9823/10000 = 98.230000%
             [95% CI 75.18–76.86]
Epoch 38/40: loss=0.310950 lr=0.100000
  epoch 38: test_acc = 7820/10000 = 78.200000%  top5 = 9860/10000 = 98.600000%
             [95% CI 77.38–79.00]
Epoch 39/40: loss=0.305377 lr=0.100000
  epoch 39: test_acc = 7784/10000 = 77.840000%  top5 = 9862/10000 = 98.620000%
             [95% CI 77.02–78.64]
Epoch 40/40: loss=0.298878 lr=0.100000
  epoch 40: test_acc = 7741/10000 = 77.410000%  top5 = 9847/10000 = 98.470000%
             [95% CI 76.58–78.22]
done (trained CIFAR-CNN8-wide-BN-batchstat sgd + constant lr via packed threading).

Forty epochs in about a minute, and 77.41% on the full 10,000-image CIFAR-10 test set, 98.47% 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 test accuracy is not monotone: it can fall several points from one epoch to the next while the training loss keeps falling. That is the evaluation rather than the training: the test forward runs on running statistics that lag the weights at this learning rate. §4.3 comes back to it, because it is where the chapter’s spreads come from.

The learning rate is constant here — no warmup, no decay. Run five times at this setting the net lands between \(74.9\) and \(77.9\% \), 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, its per-image sibling and the no-BN peer across all three optimizers, which is where the chapter’s measurements come from, and lake run cifar is those three 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 45 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 47 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 46 is the formal statement that \(\varepsilon {\gt} 0\) is sufficient to make this term differentiable. This is the first place in the book where the math requires real-analysis machinery beyond chain-sum-product (ConvNeXt’s GELU and the transformer’s softmax are the others); 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 48 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 of Theorem 17, just lifted to a tensor shape and broadcast across spatial dimensions. Theorem 44 and Theorem 49 are essentially corollaries of the dense theorems. The new content here is zero.

Putting it together

Theorem 50 is the composition: \(\mathrm{bn} = \mathrm{affine} \circ \mathrm{normalize}\), with VJPs chained via the same \(\mathrm{vjpComp}\) rule (Theorem 11). 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

Figure 4-1 is the dependency graph of this section.

The network, at two depths: the two-convolution CIFAR net’s fold and loss gradient, BatchNorm’s two updates, then the eight-convolution net with and without BatchNorm — its fused step, the packed steps this chapter’s runs train (the same backward chain handing each gradient to a separate optimizer), the same steps with BatchNorm over the batch, their loss gradients, and one descent step at the last convolution.

4.3 Example: the eight-conv CNN on CIFAR-10

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, the same data pipeline (per-epoch shuffle, random horizontal flip) and one constant learning rate. The findings, stated up front so the graphs are no surprise:

• Normalization buys accuracy, and speed on the way. BatchNorm (dropped between each conv and its ReLU) reaches a given accuracy in markedly fewer epochs and finishes three to eight points higher. What it does not buy is repeatability: five runs of the batch-statistics net spread over 2.2 to 3.0 points, about what the un-normalized net’s 2.7 to 3.3 are. The per-image form lands within 1.1, and the reason is where its statistics come from, not that it normalizes.

• The optimizer moves the ceiling less than the normalization does. Momentum finishes first on both normalized nets, but its lead over plain SGD is half a point at batch statistics, inside the spread of either. AdamW lands a point under SGD once BN is in, 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.21\), \(+0.57\) and \(+1.02\) — 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 runs at seeds 1–5 (LEAN_MLIR_SEED, which moves the initialisation, the shuffle and the flips), and each table states the observed range alongside it. That spread is seed-to-seed variance plus the nondeterminism: two to three points at batch statistics and without normalization, about a point per image, 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. Figure 4-2 is the network.

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
  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
  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"

At batch statistics, the spec the transcript in §4.1 trained is the second one with a single field changed:

def cifar8wbBbnVerified : VerifiedNetSpec :=
  { cifar8wBnVerified with
    name      := "CIFAR-CNN8-wide-BN-batchstat"
    slug      := "cifar8wb_bbn"
    runningBN := true
    blurb     := "Wide-head CIFAR-10 CNN + per-channel BatchNorm at BATCH statistics
                  (8 convs, BATCHED op family, d1=512, running-stat eval)
                  via the VERIFIED renderer → %LOWERER% → GPU" }

Each layers list and each blurb is one line in the source, wrapped here; the third blurb is the banner §4.1’s transcript opens with. The diff between the first two specs is eight .bnPerChannel entries, one inserted between each convolution and its ReLU. The diff between the second and the third is runningBN: the group each of those eight normalizes over, and whether the net keeps running statistics. 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: running mean and variance. They are state, not parameters — no gradient reaches them — so the renderer threads them through separate bnChannels slots rather than this list, sixteen of them for the batch-statistics spec, derived from runningBN and the eight widths. The per-image spec leaves runningBN off and has no slots, because its statistics never leave the example. The ImageNet nets of Chapters 5–7 are the batch-statistics form throughout.

The spec-to-proof tie. These are VerifiedNetSpecs, which is what makes the listing above binding rather than descriptive. The slug names the committed render, so cifar8wb_bbn names the six cifar8wb_bbn_*_train_step.mlir files in verified_mlir/, the ones the run in §4.1 handed to XLA, and cifar8w_bn names the per-image set. 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.cifarCnnBn8HasVJPAt, 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. cifarCnnBn8HasVJPAt_correct (Proofs/Nets/Small/CifarCNN.lean) is the statement that its backward is the derivative, and Proofs/Nets/Small/Cifar8BnStepTie.lean carries that through to the emitted per-image training step. The batch-statistics step is tied node by node in Theorem 58, its backward chain is the certified whole-net backward in Theorem 59, and each node is the loss’s derivative in its parameter in Theorem 60. All three nets run the same XLA/PJRT path on the GPU.

Lever 1: normalization

Fix the optimizer at plain SGD and change the normalization: none, per image, batch statistics. All three 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. Figure 4-3 is the per-epoch test accuracy, the median of five runs each.

The final points span 74.9 to 77.9 at batch statistics, 74.1 to 75.2 per image, and 68.0 to 70.7 without.

Both normalized nets lead from the first epoch — 44% and 41% against the bare net’s 13% — and never give the lead back, the batch form finishing 8.0 points up (76.8% vs 68.8%) and the per-image form 5.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 any set of runs. The batch curve is also the jagged one: its test reads swing where the per-image curve’s do not, because its test forward runs on running statistics that lag the weights, while the per-image net evaluates with the statistics it trained on.

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 ones start. The conditioning the three-term backward (Theorem 48) 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. Table 4-1 is the median final accuracy, with the observed range across five runs.

Momentum at batch statistics is the best result on the board at 77.4%, half a point over plain SGD and 1.7 over AdamW. The first gap is inside the spread of either cell and the second is about at its edge, so the order is momentum, SGD, AdamW on both normalized nets, and only the AdamW gap is a result. AdamW lands under 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. Batch statistics are ahead in all three — \(+8.0\) under SGD, \(+5.2\) under momentum, \(+2.9\) under AdamW — and the per-image form by \(+5.8\), \(+4.2\) and \(+1.5\). 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 spread is. Read the parenthesised ranges. Without BN, five runs of the same command land 2.7 to 3.3 points apart under every optimizer; at batch statistics, 3.0, 2.2 and 2.3; per image, 1.1, 1.0 and 0.2. So normalization is not what makes a single run’s number mean something. What does is where the evaluation’s statistics come from. The per-image net scores each test image with the statistics it trained on. The batch-statistics net scores on running means and variances, an average over roughly the last hundred steps (bnMomentum \(0.99\)), and at a constant learning rate the weights move faster than that average follows, so the final test read carries the position of the last few hundred steps as well as the seed. The dips in Figure 4-3 are the same effect at the epoch scale. The ImageNet chapters run this layer under a cosine schedule, which is what settles it. The un-normalized AdamW runs also printed loss=NaN in late epochs in three of the five seeds while their accuracy held, and that turned out to be the diagnostic, not the training: the progress line’s loss is \(\sum \mathrm{onehot}\cdot \log (\mathrm{softmax})\) with no log-sum-exp guard, so once an off-class probability underflows to zero it prints \(0\cdot (-\infty )\). The gradient is \(\mathrm{softmax}-\mathrm{onehot}\) and never takes the log; the next epoch’s loss came back finite and the accuracy never moved.

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 (Theorem 63).

Each tail is emitted onto the same certified gradient (sgdParamF, momParamF with momVNextF, adamWParamF), 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.

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, and nothing at all about the progress line: the scalar loss it prints is a report-only op outside the proven graph, which is exactly where the NaN above came from. A verified gradient is not a guarantee of numerical stability, and the low-precision seed that collapses to chance in the next section is 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.

The renderer has two convolution op families — the per-example one the MNIST chapters use, and the batched one the ImageNet chapters use — and the bf16 twins exist in the batched family only. The batch-statistics net lives there already, since a statistic over the batch needs the batch axis in the graph, so its bf16 arms come with the family. 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. Every convolution op in the graph is bf16 in the low-precision arms — twenty-three of them: eight forward, eight weight gradients, and the seven input gradients (the first layer’s input is the image and needs none). Table 4-2 is the six arms, five seeds each.

Medians move by \(-0.21\), \(+0.57\) and \(+1.02\). Each of those is smaller than the spread within the corresponding fp32 row — \(2.96\), \(2.18\) and \(2.25\) — 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. The two low bf16 AdamW seeds at \(72\) and \(73\) make that the widest row on the board, and the bf16 momentum row, five seeds within \(0.6\) of each other, is the narrowest.

The un-normalized net answers differently, and the difference is the interesting part. Run the same comparison there (the fp32b and bf16 arms of runs/2026-09-01-cifar8w-6arm-constlr/) 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 the less repeatable one in fp32 — the three-point ranges of 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 per-image 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, under AdamW on the sweep in Figure 4-4, the two land a point and a half apart with overlapping intervals: \(76.40\% \) narrow, \(77.85\% \) wide. Seven times the parameters buys no accuracy the test set can resolve. The eight convolutions over the \(32 \times 32\) maps dominate the compute, and the head is cheap matmul at these widths. 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.

That comparison is worth drawing out in full, because the parametric renderer makes it cheap. The grid driver, cifar8wb-bbn-grid, holds the eight-convolution backbone fixed and renders the AdamW train step, the forward and the evaluation forward 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. Dealt across three of the box’s four cards, the whole curve, Figure 4-4, is a four-minute run.

The recipe is the grid driver’s own; the levers above ran at a constant rate. A whisker is about \(\pm 0.8\) here, and is 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 \(76.4\) to \(77.9\% \), about one point’s own interval, so all nine of those intervals overlap the best point’s and none of the differences between them is resolvable from this experiment. The \(128\times \) wider head buys nothing measurable. One difference does survive: \(d{=}8\) (\(76.0\% \), interval \(75.1\)–\(76.8\)) stops short of the best point’s interval, which starts at \(77.0\); there the \(128{\to }8\) first layer throttles the 128-dim feature map the convolutions produce. The canonical \(d{=}64\) (circled) sits on the plateau: the head width genuinely does not matter here, which is the whole reason the net could borrow MNIST’s classifier untouched. (The \(d{=}512\) point is the board’s own net, and lands two points over its 40-epoch AdamW cell in Table 4-1; that is the cosine schedule against the constant rate, not the width, and 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 finishes first, and a longer step through eight unnormalized layers is also what let one low-precision momentum run collapse to chance. 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 Theorem 48: with \(\hat{x} = (x-\mu )\, \mathrm{istd}\),

bnHasVJP proves it, composing bnNormalizeHasVJP (the rank-1 wringer) with bnAffineHasVJP (the \(\gamma \, dy\) half). The one subtlety, isolated in Theorem 46, is that the inverse-stddev term carries an \(\mathrm{istd}^3\) that needs \(\varepsilon {\gt} 0\) to stay differentiable. That is the first 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 in a group to every other. The emitted graph is given a denotation in the proofs’ own vector type, and bn_back_bridge proves that denotation equal to bnHasVJP’s backward: the reduce/broadcast/elementwise graph is, by machine check, the three-term formula. Here it is in the chapter’s trained artifact, verified_mlir/cifar8w_bn_sgd_train_step.mlir, at the eighth convolution’s BatchNorm (32 channels at \(4\times 4\), batch 128), with the recompute of the forward statistics elided down to the rsqrt (SSA names the artifact’s, lines wrapped):

%v289 = stablehlo.reshape %v288 : (tensor<128x512xf32>)
      -> tensor<128x32x4x4xf32>
%v290 = stablehlo.reshape %v219 : (tensor<128x512xf32>)
      -> tensor<128x32x4x4xf32>
%v292 = stablehlo.constant dense<16.0> : tensor<128x32x4x4xf32>
// ...
%v303 = stablehlo.rsqrt %v302 : tensor<128x32x4x4xf32>
%v304 = stablehlo.multiply %v297, %v303 : tensor<128x32x4x4xf32>
%v305 = stablehlo.broadcast_in_dim %g8, dims = [1] : (tensor<32xf32>)
      -> tensor<128x32x4x4xf32>
%v306 = stablehlo.multiply %v305, %v289 : tensor<128x32x4x4xf32>
%v307 = stablehlo.reduce(%v306 init: %v291)
      applies stablehlo.add across dimensions = [2, 3]
      : (tensor<128x32x4x4xf32>, tensor<f32>) -> tensor<128x32xf32>
%v308 = stablehlo.broadcast_in_dim %v307, dims = [0, 1]
      : (tensor<128x32xf32>) -> tensor<128x32x4x4xf32>
%v309 = stablehlo.multiply %v304, %v306 : tensor<128x32x4x4xf32>
%v310 = stablehlo.reduce(%v309 init: %v291)
      applies stablehlo.add across dimensions = [2, 3]
      : (tensor<128x32x4x4xf32>, tensor<f32>) -> tensor<128x32xf32>
%v311 = stablehlo.broadcast_in_dim %v310, dims = [0, 1]
      : (tensor<128x32xf32>) -> tensor<128x32x4x4xf32>
%v312 = stablehlo.multiply %v306, %v292 : tensor<128x32x4x4xf32>
%v313 = stablehlo.subtract %v312, %v308 : tensor<128x32x4x4xf32>
%v314 = stablehlo.multiply %v304, %v311 : tensor<128x32x4x4xf32>
%v315 = stablehlo.subtract %v313, %v314 : tensor<128x32x4x4xf32>
%v316 = stablehlo.divide %v303, %v292 : tensor<128x32x4x4xf32>
%v317 = stablehlo.multiply %v316, %v315 : tensor<128x32x4x4xf32>

Read it against the formula. %v289 is the incoming cotangent \(dy\) and %v290 the layer’s input \(x\); %v303 is \(\mathrm{istd}\) and %v304 is \(\hat{x}\), recomputed from \(x\) rather than saved. %v306 is the affine backward \(\gamma \, dy\), with \(\gamma \) (%g8) broadcast along the channel axis. Each reduce across dimensions = [2, 3] followed by a broadcast_in_dim is one of the cross-coordinate sums: \(\sum _k \gamma \, dy_k\) as %v308 and \(\sum _k \hat{x}_k\, \gamma \, dy_k\) as %v311. %v312 is \(n\, \gamma \, dy\) with \(n = 16\) (%v292), %v313 subtracts the centering correction, %v315 subtracts the rank-1 normalization correction \(\hat{x}\sum \hat{x}\, \gamma \, dy\), and %v317 scales the bracket by \(\mathrm{istd}/n\) (%v316). The graph folds \(\gamma \) into %v306 up front, which is why that leading scale is \(\mathrm{istd}/n\) and not \(\mathrm{istd}\, \gamma /n\): the same formula with \(\gamma \) pulled inside the parenthesis.

The reduction axes say which layer this is. [2, 3] are the two spatial axes, so each statistic is taken per image and per channel over the 16 cells of a \(4\times 4\) map, which is the per-image normalization the section “What BN actually is” above describes; a batch-statistics BatchNorm reduces over [0, 2, 3], and that is what the ResNet and MobileNet artifacts of the next chapters emit. Both are instances of the one theorem, over groups of different width. The leaf theorem here is bnPerChannelBack_faithful: the per-channel three-term graph denotes the input gradient of bnPerChannelTensor3, the block-diagonal Jacobian contracted against \(dy\), under \(\varepsilon {\gt} 0\); its batch-statistics twin is bnBatchTensor4GradInput_eq_backward. Theorems 56 and 57 state that every parameter update in this file is the certified step at the cotangent the file’s own backward delivers, the eight BatchNorms included.

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 holds at every input, with no smooth-point exclusion, unlike ReLU and max-pool.

• Shape. The listing is one of the eight BatchNorms of the trained net; the theorem is stated over any group width \(n\).

4.5 Recap

BatchNorm is the hardest proof in the book and the cleanest statement. Its backward is a three-term formula, the direct term, a centering correction and a rank-one normalization correction, and the inverse-standard-deviation inside it is what needs Mathlib’s analysis: given \(\varepsilon {\gt} 0\) the gradient is bounded and the bridge is unconditional, smooth everywhere, with none of the smooth-point exclusions ReLU and max-pool carry. Theorem 50 is that formula; the per-channel and batch-statistics forms are its instantiations; and the eight-convolution step is tied with BatchNorm and without (Theorems 55 and 56), the batch-statistics step node by node and as the loss’s gradient (Theorems 58 and 60), with descent at the last convolution. The same emitted graph is LayerNorm’s backward, since LayerNorm is BatchNorm along a different axis, which is why Chapters 8 and 9 add no normalization proof of their own.

The experiment half says what the proofs do not: at eight convolutions, normalization is what governs training, the optimizer moves the result by less than a seed does, the group the statistics are taken over decides whether one run’s number can be trusted, and the number format is not measurable against seed noise. Eight is also about where a plain stack stops helping. The next chapter adds the skip connection, whose VJP is additive fan-in, the second of the three composition rules, and composes it with convolution and BatchNorm to go thirty-four layers deep, at ImageNet scale, with no new analysis at all.