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-08-12-cifar8w-6arm-xla-cuda/cifar8w-bn-p1.log. XLA’s startup banner is removed, the per-step loss lines are dropped, and epochs 8 through 34 are elided:

[pjrt_ffi] XLA backend: PJRT 0.112, 1 device(s)
[pjrt_ffi] compiled verified_mlir/cifar8w_bn_mom_train_step.mlir
             (@cifar8w_bn_mom_train_step, 117 outputs, 1 replica) in 1215 ms
[pjrt_ffi] compiled verified_mlir/cifar8w_bn_fwd.mlir
             (@cifar8w_bn_fwd, 1 outputs, 1 replica) in 261 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
    (cosine+warmup 3ep, baseLR 0.020000), He init
Epoch 1/40: loss=1.958118 lr=0.006667
  epoch 1: test_acc = 4209/10000 = 42.090000%  top5 = 8952/10000 = 89.520000%
Epoch 2/40: loss=1.418591 lr=0.013333
  epoch 2: test_acc = 5474/10000 = 54.740000%  top5 = 9449/10000 = 94.490000%
Epoch 3/40: loss=1.167188 lr=0.020000
  epoch 3: test_acc = 6022/10000 = 60.220000%  top5 = 9523/10000 = 95.230000%
Epoch 4/40: loss=0.992680 lr=0.019964
  epoch 4: test_acc = 6499/10000 = 64.990000%  top5 = 9694/10000 = 96.940000%
Epoch 5/40: loss=0.882675 lr=0.019856
  epoch 5: test_acc = 6876/10000 = 68.760000%  top5 = 9726/10000 = 97.260000%
Epoch 6/40: loss=0.811264 lr=0.019677
  epoch 6: test_acc = 6997/10000 = 69.970000%  top5 = 9771/10000 = 97.710000%
Epoch 7/40: loss=0.756665 lr=0.019429
  epoch 7: test_acc = 7208/10000 = 72.080000%  top5 = 9784/10000 = 97.840000%
Epoch 35/40: loss=0.258739 lr=0.000888
  epoch 35: test_acc = 7759/10000 = 77.590000%  top5 = 9836/10000 = 98.360000%
Epoch 36/40: loss=0.251148 lr=0.000571
  epoch 36: test_acc = 7738/10000 = 77.380000%  top5 = 9832/10000 = 98.320000%
Epoch 37/40: loss=0.247924 lr=0.000323
  epoch 37: test_acc = 7750/10000 = 77.500000%  top5 = 9840/10000 = 98.400000%
Epoch 38/40: loss=0.246187 lr=0.000144
  epoch 38: test_acc = 7753/10000 = 77.530000%  top5 = 9839/10000 = 98.390000%
Epoch 39/40: loss=0.241790 lr=0.000036
  epoch 39: test_acc = 7751/10000 = 77.510000%  top5 = 9841/10000 = 98.410000%
Epoch 40/40: loss=0.243040 lr=0.000000
  epoch 40: test_acc = 7748/10000 = 77.480000%  top5 = 9840/10000 = 98.400000%
done (trained CIFAR-CNN8-wide-BN mom + cosine/warmup via packed threading).

Forty epochs at about six seconds each, and 77.48% on the full 10,000-image CIFAR-10 test set, 98.40% 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 the first in this book you can’t read off by inspection. That run did not execute a reimplementation of the network this chapter describes. It executed verified_mlir/cifar8w_bn_mom_train_step.mlir, which is the text Proofs.StableHLO emits for this net, and Proofs/Cifar8Close.lean proves that each of its parameter updates denotes the certified softmax cross-entropy loss-descent step derived from the Mathlib \(\operatorname {fderiv}\) math. Every BatchNorm in it computes the three-term backward this chapter proves.

The two compile lines carry the whole in-process story. The train step and the eval forward are handed to XLA once, as text, and compiled in about 1.5 seconds together, after which all forty epochs run without leaving the process.

§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 whole ladder in bf16 and in fp8 leaves the optimizer ordering exactly where it was, with medians agreeing to within half a point — less than the spread between repeated runs of a single configuration. 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]

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]

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/Architectures/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. Per-epoch test accuracy, the median of five runs each, from the verified runs in runs/2026-08-12-cifar8w-6arm-xla-cuda/:

CIFAR-10, wide 8-conv net, plain SGD (lr 0.1, cosine\({+}\)warmup) 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.5\) to \(75.6\) with BN and \(71.6\) to \(73.4\) without (runs/2026-08-12-cifar8w-6arm-xla-cuda/, the SGD phase).

BN leads from the first epoch, 34% to the bare net’s 24%, and keeps the lead the whole way, finishing about two and a half points up (75.0% vs 72.6%). The un-normalized net traces the same arc roughly ten epochs behind. Most of that is speed, since the curves are the same shape shifted left, but at this depth on plain SGD it is also a real convergence margin, and it is larger than the run-to-run spread on either curve. The conditioning that the three-term backward (§34) buys is worth more the deeper the stack, exactly the trend that makes BN standard equipment by ResNet’s thirty-four layers. Under momentum and AdamW, next, the un-normalized net trains fast enough to close most of the accuracy gap, and pays for it in a way plain SGD never exposed.

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 and the same 40 epochs. Median final accuracy, with the observed range across runs:

\(n=5\) per cell, \(n=4\) for the AdamW column. Parenthesised figures are the range from lowest run to highest.

Momentum with BN is the best result on the board at 77.1%, about two points over plain SGD and two and a half over AdamW, and every one of those gaps is wider than the spread within a cell. AdamW lands between the two: its per-coordinate second-moment scaling is worth something over plain SGD, but less than momentum is, even with the wide head’s extra parameters to exploit. Reading down the columns recovers Lever 1, and BN is ahead in all three (\(+2.4\) under SGD, \(+0.7\) under AdamW, and under momentum the comparison does not finish).

The empty cell is the finding. Un-normalized plus momentum did not fail to reach a good accuracy. It reached one and then destroyed it. In all five runs the training loss fell normally to about \(0.27\) and then went to NaN between epochs 29 and 34, at a learning rate the cosine schedule had already decayed to roughly \(0.002\). Three of the five collapsed to 10% by epoch 40, which is chance on ten classes. The other two happened to be evaluated before the damage propagated and still reported about 76%. There is no honest single number for that cell, so the table does not print one. The same thing happened once in four AdamW runs without BN. Across all fourteen un-normalized runs the loss went to NaN six times, and across all fourteen normalized ones it never did.

That is the sharper version of Lever 1. On plain SGD, BN buys speed and a couple of points. Give the net a stronger optimizer and BN stops looking like an accelerator and starts looking like the thing holding the net together, which is what deep normalization is actually for. It is the same instability an eight-layer stack without normalization has always been prone to, and it is why the next chapter’s thirty-four layers do not attempt it.

The Lever-1 graph drew this for the SGD column, with the optimizer fixed and BN toggled. The same per-epoch picture for the other two columns, on the same axes, is what reading down the table looks like drawn out:

Same wide 8-conv net, same shared pipeline, same 40 epochs, and only the optimizer changes from the SGD panel above (runs/2026-08-12-cifar8w-6arm-xla-cuda/, the momentum and AdamW phases). Medians as before, over the runs that finished.

Under momentum BN takes the early lead, 42% to 36% at epoch 1, and the un-normalized net spends the next twenty-five epochs closing it, reaching \(75.8\) against BN’s \(77.0\) by epoch 28. Then it stops. The orange curve ends where it does because the training loss went to NaN in every one of the five runs between epochs 29 and 34, and there is nothing after that worth plotting. Under AdamW both nets run the full forty epochs, BN ahead by about seven tenths the whole way (\(74.4\) vs \(73.7\)), which is inside the run-to-run spread and is not a gap the text ranks. One of the four un-normalized AdamW runs also went to NaN.

Read top-to-bottom, SGD then momentum then AdamW, the accuracy gap between the curves does narrow panel by panel, exactly the old “BN helps most under a weak optimizer” reading. The momentum panel is where that reading breaks. The gap narrows right up to the point the un-normalized net destroys itself, which means the two curves converging was never evidence that BN had stopped mattering.

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 momentum 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 the un-normalized momentum run is what that distinction looks like when it bites.

One honest note on the comparison. The first time we ran this, the optimizers looked far more separated, with momentum and Adam beating SGD by some ten points. That gap was mostly an artifact: the momentum and Adam runs happened to go through a data pipeline with shuffling and flip augmentation that the plain-SGD baseline lacked. Holding the pipeline genuinely fixed, as the table above does, collapsed the difference to the few points that are really about the optimizer. It is exactly the kind of mistake a verified gradient does not catch. The math was correct in every run, and the experiment design was what needed fixing. An ablation measures the thing you varied only if everything else is held constant, and “everything else” includes the parts that live outside the network.

Lever 3: the arithmetic

Levers 1 and 2 changed the network and the update rule. This one changes neither. It holds both fixed and changes only the number format the convolutions are computed in, which is the sharpest version of the chapter’s question: if the optimizer ranking is a fact about optimization, it should survive being computed in a different arithmetic.

Three precisions, the same three optimizers, the same 40 epochs and the same pipeline. Median of five runs, with the observed range:

\(n=5\) per cell. Reading down each column recovers Lever 2, and it recovers it three times: SGD is last, AdamW is second, momentum is first, in every arithmetic. The ranking is invariant. Reading across the rows, the medians agree to within half a point — 72.6/72.5/72.2 and 74.3/73.9/74.3 — which is well inside the spread within a single cell. Sixteen-bit and eight-bit arithmetic land on the same curve as fp32.

The whole momentum row is Lever 1 again, and the accuracies are what mislead. In all three arithmetics, every one of the five momentum runs sent its training loss to NaN — fifteen runs out of fifteen. Only two of them had visibly collapsed to 10% by the time epoch 40 was evaluated; the other thirteen still reported something around 76%, which is precisely the trap Lever 2 warns about, and precisely why the loss and not the accuracy is the thing to read. Had we scored this row on test accuracy alone, we would have printed a best-in-table 76.0 for bf16 and called it a win.

That the failure is identical across fp32, bf16 and fp8 is the point. The un-normalized eight-layer stack with momentum is unstable as a network; the number format neither causes it nor cures it. Reduced precision is being blamed for nothing here, and credited with nothing either.

AdamW is a milder version of the same story: two of five fp32 runs and one of five in each low-precision column touch NaN without ever collapsing, matching the “once in four” Lever 2 saw. Plain SGD is clean in all fifteen runs. So the ordering the table reports rests on the two rows that finish healthy, and the third row is reported as diverged rather than scored.

What the two low-precision columns actually are, since they are not the same kind of object. The bf16 arm is genuine reduced precision in the graph: all twenty-three convolutions of the training step — the forward, the input gradient and the weight gradient — carry bf16 operands and a bf16-typed result, and only the parameter updates stay in fp32. The fp8 arm is emulated: weights and inputs are projected onto the E4M3 grid, but the graph itself is fp32 and accumulation is fp32. It is the right arithmetic and the wrong hardware path, which is enough for a statement about numerics and not enough for one about speed.

And there is no speed claim here. At this network’s size the convolutions are launch-bound — sixteen to thirty-two channels on grids from \(32^2\) down to \(4^2\) — so the tensor cores that make bf16 and fp8 worth having on ImageNet-scale models sit idle. Measured directly, bf16 runs at \(0.87\times \) fp32 across this conv stack: slightly slower. The payoff for low precision is a story about ResNet-34 and beyond (Ch 5); what CIFAR can show is that the mathematics survives the trip, and that is what this lever is for.

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). Past a 16-wide head the curve is essentially flat, and everything from \(d{=}16\) to \(d{=}4096\) lands inside a single point (\(71.1\) to \(72.2\% \)), and the \(256\times \) wider head buys nothing, while its \(17\)M-parameter classifier just overfits (train loss \(0.03\), test unmoved). The only real drop is at \(d{=}8\) (\(66.8\% \)), where 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.