5 ResNet-34

Deeper networks used to stop training

Before 2015, the dominant assumption in image classification was “deeper is better, up to a point.” The “up to a point” part was not metaphorical. If you took a 20-layer CNN that trained to 70% val accuracy and added 12 more convolutional layers, you didn’t get a 32-layer model with 72% accuracy. You got a 32-layer model that, after the same training budget, fit the training set worse than the 20-layer one did. The 32-layer model had strictly more capacity. It still trained worse. Something about deep stacks of convolutions made gradient flow degrade in a way that the optimizer couldn’t navigate around.

He et al. 2015 (arXiv:1512.03385) had the right observation: if the 12 extra layers in the deeper network were initialized to compute the identity, the 32-layer network would have at least the 20-layer’s performance, because it could replicate it exactly. So the question was not “can a deep network represent what a shallow one represents?” It provably could. The question was: why is optimizing 12 free-form layers to recover identity so much harder than just leaving identity there?

5.1 Run it first

Before any of the math, train the thing. Four commands and about an hour and a half of GPU time:

lake exe cache get                            # Mathlib oleans, ~30 s
./download_imagenette.sh                      # 330 MB dl, 2.4 GB unpacked
lake build resnet34-verified-adam
./.lake/build/bin/resnet34-verified-adam data

On one RTX 4060 Ti (CUDA 12.9), from runs/2026-08-12-r34-imagenette-xla-cuda/. XLA’s startup banner is removed, the network description line is wrapped, and epochs 5 through 75 are elided:

[pjrt_ffi] XLA backend: PJRT 0.112, 1 device(s)
[pjrt_ffi] compiled verified_mlir/resnet34_adam_train_step.mlir
             (@resnet34_adam_train_step, 405 outputs, 1 replica) in 4952 ms
[pjrt_ffi] compiled verified_mlir/resnet34_fwd.mlir
             (@resnet34_fwd, 1 outputs, 1 replica) in 409 ms
[pjrt_ffi] compiled verified_mlir/resnet34_fwd_eval.mlir
             (@resnet34_fwd_eval, 1 outputs, 1 replica) in 195 ms
Real ResNet-34 on Imagenette 224² (7×7-s2 stem→3×3-s2 overlapping max pool→
  [3,4,6,3] blocks w/ batch-norm, He et al. option-B 1×1 projection shortcuts,
  no conv biases; 56→28→14→7→GAP→dense)
  via the VERIFIED renderer → XLA/PJRT → GPU
  train 9469, val 3925; bs 32, ResNet-34 adam
    (cosine+warmup 3ep, baseLR 0.001000), He init
  running-stats BN: 36 layers, 17024 stat floats → eval via @resnet34_fwd_eval
Epoch 1/80: loss=2.049761 lr=0.000333
  epoch 1: val_acc = 1083/3925 = 27.592357%  top5 = 2601/3925 = 66.267516%
Epoch 2/80: loss=1.530066 lr=0.000667
  epoch 2: val_acc = 1363/3925 = 34.726115%  top5 = 3024/3925 = 77.044586%
Epoch 3/80: loss=1.388131 lr=0.001000
  epoch 3: val_acc = 2250/3925 = 57.324841%  top5 = 3554/3925 = 90.547771%
Epoch 4/80: loss=1.300556 lr=0.001000
  epoch 4: val_acc = 2510/3925 = 63.949045%  top5 = 3642/3925 = 92.789809%
Epoch 76/80: loss=0.502019 lr=0.000007
  epoch 76: val_acc = 3522/3925 = 89.732484%  top5 = 3858/3925 = 98.292994%
Epoch 79/80: loss=0.502105 lr=0.000000
  epoch 79: val_acc = 3522/3925 = 89.732484%  top5 = 3857/3925 = 98.267516%
Epoch 80/80: loss=0.502114 lr=0.000000
  epoch 80: val_acc = 3521/3925 = 89.707006%  top5 = 3857/3925 = 98.267516%
done (trained ResNet-34 adam + cosine/warmup via packed threading).

Eighty epochs at about 65 seconds each, roughly an hour and a half in total, and 89.71% top-1 with 98.27% top-5 on Imagenette’s 3,925-image validation split. The best epoch reached \(89.86\% \). That is a thirty-four-layer network trained from scratch on ten thousand photographs, which is a different kind of problem from the handwritten digits three chapters back.

Thirty-four layers on, nothing about the claim has changed. That run did not execute a reimplementation of the network this chapter describes. It executed verified_mlir/resnet34_adam_train_step.mlir, which is pretty(provenGraph) off the renderer, and every one of the sixteen residual blocks in it computes the additive fan-in backward that Theorem 37 proves. The thirty-six BatchNorms compute the three-term backward of Chapter 4, and the convolutions compute the backward of Chapter 3. Every theorem in the rest of this chapter is about the graph that just trained.

Three compile lines rather than the two earlier chapters showed, and the third is the interesting one. @resnet34_fwd_eval is a separate forward, because BatchNorm behaves differently at training time than at evaluation time. Training normalizes each batch by its own statistics, while evaluation has to use the running averages accumulated over training, and the line running-stats BN: 36 layers, 17024 stat floats is the driver saying it is threading those averages out of the train step and into the eval forward. A net that skips this reports chance forever, since the eval forward normalizes by statistics nobody ever computed.

The final loss settles at \(0.502\) and stops. That is not the optimizer giving up. It is the label-smoothing floor derived in §5.5: with \(\varepsilon = 0.1\) over ten classes the smallest attainable cross-entropy is about \(0.485\), so a run that lands at \(0.502\) has essentially converged.

The residual block: ask for the difference, not the function

Their fix was a one-line architectural change. Instead of asking each block to learn a function \(y = f(x)\), ask it to learn a function \(y = f(x) + x\). The block’s job is no longer to become the new representation. Its job is to learn what to add to the current representation. At initialization, with randomly-small \(f\), each block is approximately the identity, and a stack of 16 such blocks is approximately a stack of 16 identities, which is to say, approximately the input. The network starts out near-identity at every depth and only gradually takes on structure as \(f\) trains. The optimization problem becomes “what residual should each block add” instead of “what function should each block be,” and the former turns out to be dramatically easier.

The change to the architecture is one + sign. The change to the training dynamics is the difference between “18 layers is about as deep as we can go” and “152 layers trains stably, so let us also add another zero on the ImageNet leaderboard.”

Why the proof is a one-liner

From the framework’s perspective, what is a residual block? \(\mathrm{residual}(x) = f(x) + \mathrm{id}(x)\), a sum of two functions sharing the same input. We already proved in Chapter 1 that the VJP of an additive fan-in is “send the upstream gradient through both branches and add.” We already proved that the VJP of \(\mathrm{id}\) is the identity on the gradient (and that the VJP of \(f\) is whatever the inner block’s machinery gives us). Composition is just \(\mathrm{vjp\_ comp}\). So the residual VJP is the additive fan-in lemma, with \(\mathrm{id}\) plugged into one slot.

That is Theorem 37 below, and its proof is one line. This chapter has exactly one theorem, and that theorem is mechanical. The engineering revolution was not in the math. It was in the realization that this operation was the missing piece for training deep networks. The fact that the framework absorbs the new architecture without any new math is the value proposition: once the foundation rules are in place, the next decade of image recognition follows from composing them.

ResNet-34 is sixteen of these blocks

ResNet-34 (the smallest “deep” variant in the original paper) is built from 16 residual blocks arranged in four stages of depths 3, 4, 6, 3, at channel widths 64, 128, 256, 512. Each stage starts with a stride-2 downsample (halving spatial resolution) before its stack of stride-1 blocks. The stages are glued together by a stem (a 7\(\times \)7 conv at stride 2 plus a max-pool, taking 224\(\times \)224\(\times \)3 input to 56\(\times \)56\(\times \)64 features) and a head (global average pool plus a 512-to-10 dense layer). Total: 34 weight-bearing layers, 21.3M parameters.

Two pieces of this template are new and worth naming:

1. The 7\(\times \)7 stride-2 stem. Earlier chapters’ networks operated at full input resolution. ResNet aggressively downsamples at the input so the bulk of compute happens at lower spatial resolution. The stem is one big conv that gets us from 224\(\times \)224 to 56\(\times \)56 in a single step.

2. Global average pool instead of flatten-plus-dense. Chapters 3 and 4 ended in a .flatten followed by a \(\sim \)4096-to-512 dense layer that contained tens of millions of parameters. ResNet collapses the final 7\(\times \)7\(\times \)512 feature map to a 512-vector by spatial averaging, then runs a single 512-to-10 dense. No flatten layer. The architectural move alone drops the model from what would have been a \(\sim \)200M-parameter AlexNet-era spec down to 21.3M.

Every chapter after this one is a variation on this template. ConvNeXt swaps in different blocks. MobileNetV2 swaps in depthwise-separable blocks. EfficientNet adds compound scaling. ViT replaces the whole convolutional spine with patches and attention. The stem-stages-head template, and the residual fan-in pattern, persist.

5.2 The theorem

5.3 Example: ResNet-34 on Imagenette

Same template as the MNIST and CIFAR chapters, but at the scale the architecture was designed for. ResNet-34 on 224\(\times \)224 Imagenette: 10 classes, real photographs, the transition point from “classroom dataset” to “actual image recognition.”

The residual block itself is the whole trick: instead of \(y = f(x)\), compute \(y = f(x) + x\). That’s the additive-fan-in pattern of § 11 with \(\mathrm{id}\) on one branch, which is the composition § 37 proves. Stacking 34 layers worth of these blocks just composes the fan-in rule 34 times (16 residual blocks deep and 34 convs across all of them).

The architecture

The same vertical column as the previous chapters, now at production scale: a \(7 \times 7\) stride-2 stem and a max-pool drop the \(224 \times 224 \times 3\) input to \(56 \times 56 \times 64\), then four stages of residual blocks (depths \(3, 4, 6, 3\) at widths \(64, 128, 256, 512\)) do the work, and a global average pool replaces the flatten-plus-fat-dense head. The inset shows the one new piece: each residual block adds its input back, \(y = f(x) + x\).

The verified trainer, which is the pattern every later chapter reuses

This is the reference configuration. Chapters 6 through 9 each swap the layers list and change almost nothing else, so it is worth reading once here rather than four more times later.

import LeanMlir

-- 1. The network. A VerifiedNetSpec, not a NetSpec: `slug` names the
--    committed render, and `bnChannels` drives BN statistic threading.
def resnet34Verified : VerifiedNetSpec where
  name     := "ResNet-34"
  slug     := "resnet34"
  inC      := 3
  imageH   := 224
  imageW   := 224
  nClasses := 10
  data     := .imagenette
  layers   := [
    .convBnNB 3 64 7 2,          -- 7x7-s2 stem, no conv bias   224->112
    .maxPool 3 2,                -- He et al. 3x3-s2 OVERLAPPING 112->56
    .residualStage  64  64 3 1,  -- stage 1, 3 blocks            @56
    .residualStage  64 128 4 2,  -- stage 2, downsample + 3      56->28
    .residualStage 128 256 6 2,  -- stage 3, downsample + 5      28->14
    .residualStage 256 512 3 2,  -- stage 4, downsample + 2      14->7
    .globalAvgPool,              -- replaces flatten + fat dense
    .dense 512 10 ]
  bnChannels := #[64, 64,64, ... ]   -- all 36, in forward order

-- 2. The schedule. Everything the optimizer needs beyond the render.
def resnet34Config : VerifiedConfig where
  epochs    := 80
  batchSize := 32

-- 3. The entry point: AdamW, cosine with a 3-epoch warmup, baseLR 1e-3.
def main (argv : List String) : IO Unit :=
  resnet34Verified.toNet.trainAdamSched
    resnet34Config (argv.head?.getD "data") 0.001 0.9 0.999 3 "adam"

Four things in that listing carry weight, and they are the four to copy.

slug names the artifact. resnet34 resolves to verified_mlir/resnet34_adam_train_step.mlir, which is pretty(provenGraph): the text the printer emits from the graph the theorems are about. The trainer does not build a network from this spec at run time. It loads that file. The spec’s job is to say which file, and to be #guarded against it so the two cannot drift.

bnChannels is not decoration. It lists every BatchNorm’s channel count in forward order, and it is what makes trainAdamSched thread running statistics out of the train step and into a separate evaluation forward. A net that declares BN layers and then trains through a driver without that threading evaluates against statistics nobody computed, and reports chance forever. That is a real failure mode, not a hypothetical, and §5.1’s third compile line is where you can see the machinery that avoids it.

trainAdamSched, not .train. The generic driver is where the optimizer schedule, the augmentation pipeline and the BN threading live. .train is the simpler peer from the MNIST chapters and it has none of the last one.

The recipe is arguments, not architecture. The learning rate, the momentum pair, the warmup length and the variant string are all passed in. Swapping AdamW for momentum is a different argument against the same proven gradient, which is exactly what Chapter 4 used to measure optimizers against each other.

What .residualStage expands to. It is a constructor, not a magic word, and the whole of it is twelve lines in LeanMlir/VerifiedSpec.lean. Every layer reports the parameter tensors it contributes as (dims, initKind) pairs, where 0 is He(fan-in), 1 is ones (\(\gamma \)) and 2 is zeros (\(\beta \) and biases):

/-- identity basic block @ `c`: two conv->BN->relu units, no projection. -/
private def idBlk (c : Nat) : Array (Array Nat × Nat) :=
  #[(#[c,c,3,3],0),(#[c],1),(#[c],2), (#[c,c,3,3],0),(#[c],1),(#[c],2)]

/-- downsampling basic block `cin->c`: two conv->BN->relu + the 1x1
    option-B projection shortcut (He et al. §3.3). -/
private def downBlk (cin c : Nat) : Array (Array Nat × Nat) :=
  #[(#[c,cin,3,3],0),(#[c],1),(#[c],2), (#[c,c,3,3],0),(#[c],1),(#[c],2),
    (#[c,cin,1,1],0),(#[c],1),(#[c],2)]

private def stageSpec (ic oc count stride : Nat)
    : Array (Array Nat × Nat) := Id.run do
  let mut a := if stride != 1 || ic != oc then downBlk ic oc else idBlk oc
  for _ in [0:count-1] do a := a ++ idBlk oc
  return a

def toSpecs : VLayer -> Array (Array Nat × Nat)
  | dense ic oc             => #[(#[ic,oc],0),(#[oc],2)]      -- Ch. 1
  | relu                    => #[]                            -- Ch. 2
  | conv ic oc k _          => #[(#[oc,ic,k,k],0),(#[oc],2)]  -- Ch. 3
  | maxPool _ _             => #[]                            -- Ch. 3
  | flatten                 => #[]                            -- Ch. 3
  | bnPerChannel oc         => #[(#[oc],1),(#[oc],2)]         -- Ch. 4
  | residualStage ic oc n s => stageSpec ic oc n s            -- here
  | globalAvgPool           => #[]                            -- here
  | bottleneckStage i o n s => bottleneckStageSpec i o n s    -- R50, below
  | ...                        -- 12 more, one per later chapter's block

Read it against the spec above. .residualStage 64 128 4 2 carries stride = 2, so the condition fires: the first block is a downBlk — two \(3\times 3\) convs plus the \(1\times 1\) projection that lets the skip cross a channel change — and the other three are idBlks. Nine tensors plus three sixes, twenty-seven in the stage. .residualStage 64 64 3 1 has stride = 1 and ic = oc, so the condition is false, no projection is emitted, and all three blocks are identity: eighteen tensors. That one-line conditional is the entire difference between a stage that changes shape and one that does not, and it is the same dispatch .bottleneckStage uses for ResNet-50.

It also lets you check the spec against itself. Counting BatchNorms out of those twelve lines — one for the stem, then \(3 \times 2\), \(3 + 3 \times 2\), \(3 + 5 \times 2\) and \(3 + 2 \times 2\) for the four stages — gives \(1 + 6 + 9 + 13 + 7 = 36\), which is exactly the length bnChannels declares and the thirty-six BatchNorms §5.3 says the graph computes. Nothing in the listing is asserted twice; the second number is derived from the first.

This is the first chapter that uses the full production recipe (Adam + cosine + warmup + weight decay + augmentation + label smoothing) rather than the s4tfBaseline. Also the first chapter where .globalAvgPool replaces the flatten-plus-fat-dense head. Global average pooling (GAP) collapses each \(C \times H \times W\) feature map to a length-\(C\) vector by averaging over the spatial dimensions,

so the classifier head sees one summary scalar per channel instead of the full \(C \cdot H \cdot W\) flattened activation. That’s the architectural move that drops ResNet-34 from what would have been a \(\sim \)200M parameter AlexNet-era spec down to 21.3M. We’ll see GAP again as the “Squeeze” step inside the squeeze-and-excitation block in Chapter 7, and at the head of every post-2015 vision architecture in the bestiary.

Results

§5.1 has the run: 80 epochs at about 65 seconds each on one RTX 4060 Ti, roughly an hour and a half, finishing at 89.71% top-1 and 98.27% top-5 on Imagenette’s 3,925-image validation split. That is the going rate for a scratch-trained ResNet-34 at this resolution and data size. Rather than reprint the log, here is what is worth noticing in it.

- 21.3M parameters, and 729 KB of MLIR (728 801 chars) for the entire AdamW training step. That is \(112\times \) the MNIST MLP’s render and only \(6\times \) the wide CIFAR BatchNorm net’s, against a network that is far more than six times the earlier one. Emission is per-op, not per-parameter, so \(3 \times 3\) convolutions at 64 and at 512 channels cost the same handful of lines. Depth grows the text linearly, and width does not grow it at all.

- Per-step time is about 220 ms at batch 32, 295 steps to the epoch, dominated by the four .residualStage groups where the \(3 \times 3\) convolutions at 64, 128, 256 and 512 channels are. The first stage alone at 64 channels does \(3 \times 3 \times 64 \times 64\) multiply-accumulates per output spatial location, about \(3.4 \times 10^8\) FLOPs per forward image, times the batch, times three blocks at that width, times two for the backward.

- Three artifacts, not one. The train step is accompanied by resnet34_fwd (98 383 chars) and resnet34_fwd_eval (76 720 chars). The last exists only because BatchNorm evaluates differently than it trains, and it is the concrete cost of the bnChannels line in the spec above.

- Final loss plateaus near 0.50, not zero. That’s the label-smoothing floor: with \(\epsilon = 0.1\) and 10 classes, the minimum possible cross-entropy against a smoothed target is \(-0.9 \log 0.9 - 0.1 \log (0.1/9) \approx 0.485\). The 0.50 value at epoch 80 means the model has essentially converged. More epochs would mostly just extend training time without budging the loss.

- BN layers: 36. Every residual block contains two normalized convolutions, and there are 16 residual blocks across the four stages (3+4+6+3). Plus the stem. That is 33 from the .residualStage expansions plus 3 stem and projection layers, for 36 total. Each one proves its own VJP via § 36 and composes with the rest of the network via § 10.

5.4 MLIR: Residual

What is already proven. A residual block is \(\mathrm{residual}(x) = f(x) + x\), an additive fan-in, where \(f\) is the conv \(\to \) BN \(\to \) relu \(\to \) conv \(\to \) BN inner block. As Theorem 37 states, its reverse-mode derivative is the additive-fan-in rule with the identity in one slot: send the upstream gradient through both branches and add. residual_has_vjp_at proves exactly this at the evaluation point (a one-line proof: the fan-in lemma with the identity’s VJP plugged into one slot), and the inner \(f\)’s VJP is the convolution and BatchNorm backward already proven in Chapters 3 and 4. The whole block, and a tower of sixteen of them, composes by the chain rule vjp_comp_at (§ 10).

The gap and how we close it. The residual fan-in is the structural heart of a deep ResNet. The emitted backward is given a denotation valued in the proofs’ own tensor type, shown equal to residual_has_vjp_at’s backward, and the structural move is visible in the emitted text. The forward saves \(\texttt{add} = f(x) + x\) and the block output \(\texttt{relu}(\texttt{add})\). The backward first pushes the incoming cotangent back through that relu (a compare/select at the post-add, giving %dadd), and then, and this is the fan-in, sends %dadd through both branches and adds (the inner block’s forward and \(f\)-backward elided, as they are the conv/BN nodes of Chapters 3–4):

// forward (saved): %out = f(%x)  [conv-BN-relu-conv-BN, Ch. 4-5]
//                  %add = %out + %x
%madd = stablehlo.compare GT, %add, %zc
          : (tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>)
            -> tensor<1x2x4x4xi1>
%dadd = stablehlo.select %madd, %dOut, %zc
          : tensor<1x2x4x4xi1>, tensor<1x2x4x4xf32>
// %dF = f.backward(%dadd)  [conv/BN/relu input-VJP chain, Ch. 4-5]
%dx = stablehlo.add %dF, %dadd : tensor<1x2x4x4xf32>
return %dx : tensor<1x2x4x4xf32>

Read it against the theorem: %dadd is the cotangent at the fan-in, and it appears twice in the last two lines, once as the input to \(f\)’s backward (%dF, the gradient through the residual function) and once added directly (%dx = %dF + %dadd, the gradient straight down the identity skip). That single stablehlo.add is “send the gradient through both branches and add”: the skip term is the same %dadd the block term started from, which is why the gradient survives all sixteen blocks instead of attenuating away. The inner \(f\) backward that computes %dF, the convolution and BatchNorm input-VJPs, is bridged op by op exactly as in Chapters 3 and 4. Only the fan-in is new here.

Caveats.

• The fan-in add is unconditional, because addition is linear, so that bridge holds at every input.

• The post-add ReLU is a smooth-point bridge (no pre-activation exactly zero), and the two BatchNorms inside \(f\) carry their own \(\epsilon {\gt} 0\) conditions from Chapter 4.

• Representative scale, a two-channel block at \(4\times 4\). The same %dadd fan-in repeats for all sixteen blocks.

5.5 What’s in the production recipe?

Six ingredients first appear here, none of which are layers. They all live in TrainConfig or training code. They don’t touch the network, but they’re what moves a scratch-trained model from 72.82% (plain SGD + momentum, no other tricks) to 90.29%.

Adam (Kingma & Ba, 2014, arXiv:1412.6980). Replaces vanilla SGD. Maintains per-parameter running estimates of the first and second moments of the gradient and uses them for an adaptive per-parameter learning rate. Much faster convergence than SGD on most tasks, much more forgiving of learning-rate choice.

Cosine learning-rate schedule (Loshchilov & Hutter, 2016, arXiv:1608.03983). After warmup, the learning rate decays from its peak to near-zero following \(\mathrm{lr}(t) = \mathrm{lr}_{\max } \cdot \tfrac {1}{2}(1 + \cos (\pi t / T))\). Smoother than step schedules, and consistently produces slightly better final losses in practice. Paired with warmup and progressive resizing (Howard et al., 2018), this was the winning recipe at Stanford’s DAWNBench in 2018. The lineage traces back to Leslie Smith’s cyclical learning rates (Smith 2015, arXiv:1506.01186). Popularized for ImageNet-scale training by He et al. 2018’s “Bag of Tricks” (arXiv:1812.01187).

Warmup (Goyal et al., 2017, arXiv:1706.02677). Linearly ramps the learning rate from zero to its peak over the first few epochs (3 for ResNet-34). Stabilizes early training when gradients are still reorganizing randomly-initialized weights. Becomes essential for larger models and larger batches, and by ViT-Tiny in Chapter 9 it’s a necessity.

Weight decay. Adds a small L2 shrink to all trainable weights each step: \(w \leftarrow w - \eta \lambda w\). Regularizer, discourages weights from growing without bound. Loshchilov & Hutter’s “decoupled” variant (AdamW, 2017, arXiv:1711.05101) is what actually ships in modern pipelines, and it’s what TrainConfig.weightDecay implements.

Augmentation. Our pipeline does random crops (\(256 \to 224\)) and random horizontal flips on Imagenette, hflip only on CIFAR, and nothing on MNIST. Lives in the data pipeline, not the network. We will explore data augmentation strategies in upcoming chapters.

Label smoothing (Szegedy et al., 2016, arXiv:1512.00567). Replaces the one-hot target with a smoothed \((\varepsilon / (K-1), \ldots , 1 - \varepsilon , \ldots )\). Stops the network from producing arbitrarily large logits on the correct class, which keeps training stable and improves calibration. Explains the 0.50 loss floor in the Results section above: \(-0.9 \log 0.9 - 0.1 \log (0.1/9) \approx 0.485\).

Every chapter after this one uses the same six ingredients. The NetSpec changes, and the recipe stays.

5.6 Ablation: what each ingredient contributes

So how much does each ingredient actually earn? We ran the full recipe plus six leave-one-out variants on ResNet-34 Imagenette, 80 epochs each, same init and batch order. Each ablation row keeps everything except the one named component, so the lift measures that component’s contribution given everything else is present. It is an honest answer to “is this piece pulling its weight?” rather than the order-dependent story you’d get from an additive ladder.

Three observations worth naming:

Augmentation is by far the largest single contribution (\(-\)7.58 points, more than 2\(\times \) any other knob). Imagenette has just 9469 training images for a 21M-parameter network, so without augmentation the model overfits hard, and training loss drops below \(0.001\) while val accuracy stalls. Augmentation is doing most of the work that the rest of the recipe gets credit for.

Cosine decay and Adam are roughly tied for largest non-augmentation contribution (\(-\)3.48 and \(-\)3.25 points). They operate on different axes, with cosine shaping the learning-rate trajectory and Adam the per-parameter update magnitude, but their late-training influence on val accuracy is comparable.

The contributions are essentially additive. Sum of all six leave-one-out deltas: \(-17.38\) points. Bare-recipe delta vs full: \(-17.47\) points. The two agree to within rounding, suggesting the modern recipe’s ingredients aren’t significantly interacting. Each one buys roughly the same lift independently of which others are present. Weight decay and label smoothing are the smallest contributors (\(-\)0.61 and \(-\)1.05) but show the overfit-then-catch-up signature mid-training (e.g., no-WD lands \(-\)6 points at epoch 10 before the gap narrows by epoch 80).

A note on the SGD configs: r34-no-adam replaces Adam with vanilla SGD at lr=0.01 (no momentum), keeping everything else from the full recipe. r34-bare uses SGD with momentum 0.9 at lr=0.01 and turns all other tricks off. The two SGD configs therefore differ on momentum as well as on the recipe ingredients, so the bare row should be read as a “what does the modern recipe buy us total” baseline rather than a strict additive component of the leave-one-out chain.

5.7 ImageNet recipe

The Imagenette runs above all trained on \(\sim \)9.5K images for 80 epochs. This section is the same 21.8M-parameter network on 1.28M images and 1000 classes, trained on the paper’s 90-epoch schedule. The architecture is unchanged. Only the dataset, the output head width, and the training schedule differ.

Ninety epochs at global batch 256 is what He et al. ran, and this section runs it twice: once on the verified codegen path and once on the JAX reference. Both paths share everything above the data loader — the architecture, the .train invocation, and every recipe knob — so the only thing that differs is how the graph reaches the GPU.

The recipe in full: SGD with momentum \(0.9\) at batch 256, cosine LR with a 5-epoch warmup from peak \(0.1\), label smoothing \(0.1\), Inception-style random-resized-crop (sampling \(8\)–\(100\% \) of the image area at a \(\frac{3}{4}\)–\(\frac{4}{3}\) aspect ratio, then resizing to \(224\times 224\)) plus a horizontal flip, and \(\ell _2\) weight decay \(10^{-4}\). Every one of those is either the original paper or the “bag of tricks” polish that became standard after it, and §4 is the table that says which is which.

The PJRT lowerer

The verified path is a second trusted lowerer, ffi/pjrt_ffi.c: it implements the same C surface as the IREE shim and hands the graph to XLA through the PJRT C API. Nothing above the shim changes. The backend is whichever .so the binary linked, and backend detection is a weak symbol, so a binary cannot disagree with the library it linked. There is no Python at run time. What the trainer executes is pretty(provenGraph): the same rendered artifact the proofs reason about, not a hand-written emitter.

Here is the trainer:

-- 1. Same backbone as the Imagenette net, 1000 classes instead of 10.
def resnet34ImagenetVerified : VerifiedNetSpec where
  name     := "ResNet-34 (ImageNet-1k)"
  slug     := "resnet34in"
  inC      := 3
  imageH   := 224
  imageW   := 224
  nClasses := 1000
  data     := .imagenet
  shimScript := "generated_resnet34_imagenet_shim.py"
  layers   := [
    .convBnNB 3 64 7 2,          -- 7x7-s2 stem -> BN -> relu    224->112
    .maxPool 3 2,                -- He et al. 3x3-s2, OVERLAPPING 112->56
    .residualStage  64  64 3 1,  -- stage1: 3 identity            @56
    .residualStage  64 128 4 2,  -- stage2: downsample + 3        56->28
    .residualStage 128 256 6 2,  -- stage3: downsample + 5        28->14
    .residualStage 256 512 3 2,  -- stage4: downsample + 2        14->7
    .globalAvgPool,
    .dense 512 1000 ]
  bnChannels := #[64, 64,64, 64,64, ... ]   -- all 36, in forward order

-- The head is the ONLY thing that may differ from the 10-class net.
-- Anything else moving means the spec drifted from its Imagenette twin.
#guard resnet34ImagenetVerified.toSpecs.size == resnet34Verified.toSpecs.size
#guard resnet34ImagenetVerified.toSpecs.pop.pop == resnet34Verified.toSpecs.pop.pop
#guard resnet34ImagenetVerified.toSpecs.back!  == (#[1000], 2)

-- 2. 90 epochs at batch 256 — the paper schedule.
def resnet34ImagenetConfig : VerifiedConfig where
  epochs    := 90
  batchSize := 256

-- 3. The entry point. Heavy-ball momentum, cosine with a 5-epoch warmup
--    at peak 0.1, weight decay 1e-4, label smoothing 0.1.
def main (argv : List String) : IO Unit :=
  resnet34ImagenetVerified.toNet.trainAdamSched
    resnet34ImagenetConfig (argv.head?.getD "data") 0.1 0.9 0.999 5 "mom256"

That is apps/imagenette/MainResnet34Imagenet.lean, built as lake build resnet34-imagenet-verified, and it is the program the numbers below came from. Three things in it are worth pulling out, because they are what makes the listing a claim rather than a description.

The slug is the artifact. resnet34in names verified_mlir/resnet34in_mom256_train_step.mlir, which is pretty(provenGraph) off the same renderer every earlier chapter used. Scaling to ImageNet needed nothing new in it. nClasses, the batch, the optimizer and the slug are all ordinary parameters, so the ImageNet artifacts are three #evals of machinery that was already proved. What had to be built was the data path, not the compiler.

The #guards pin it to the Imagenette net. The 1000-class spec must be the 10-class one with a different head, and the three guards fail the build if any other parameter shape moves. That is what licenses reading the Imagenette chapter’s theorems as being about this network’s body.

The shimScript is per-net and load-bearing. It names this net’s own generated augmentation shim. Every net used to stream ResNet-34’s, which meant an augmentation mismatch could hide inside a fair-looking comparison. Here it is the same shim the JAX reference consumes, which is what makes the two a matched pair rather than two separate stories.

It runs with:

CUDA_VISIBLE_DEVICES=0,2,3,4 PJRT_REPLICAS=4 LEAN_MLIR_REPLICAS=4 \
  SHIM_WORKERS=8 PJRT_FFI_RESIDENT=1 \
  LEAN_MLIR_VARIANT=momdp64 LEAN_MLIR_BATCH=64 \
  .lake/build/bin/resnet34-imagenet-verified data

Global batch 256 as \(4 \times 64\) replicas, fp32 throughout. The ROCm row is a [TODO]: the same render, the same C surface, a different backend .so.

The run above, epoch by epoch. One attempt, no thermal rest and no PCIe interruption across all 90.

The JAX lowerer

The JAX trainer lives in jax/MainResnetImagenet.lean, builds as lake build resnet34-imagenet, and wires .imagenet through tfds streaming. It runs the same recipe on the same cards, and its job here is to be the independent second implementation the verified number is checked against.

It runs in bf16: the spec casts both matmuls and convolutions to bfloat16 while keeping fp32 master weights. NVIDIA’s tensor cores run bf16 convolution about \(1.6\times \) faster than fp32 — the opposite of AMD’s MIOpen, where it is a wash, which is why conv precision sits behind its own bf16Conv flag. That is the whole reason the reference column below is roughly three times quicker than the verified one: a throughput choice, not a fidelity one.

Per-epoch throughput for the same recipe on other hardware. The A100 row is a wall-clock estimate. ^† MI300X ran at batch 1024 / lr 0.4 and before the current validation protocol, so only its throughput is quotable.

The same recipe, epoch by epoch, through the JAX lowerer. The kink near epoch 5 is the end of warmup; the steady late climb is the cosine anneal.

The two paths, side by side

Both columns are ResNet-34 on the 2018 recipe, 90 epochs, global batch 256, on the same four 4060 Ti cards, scored over the same \(50{,}000\) validation images under the same protocol. Everything above the rule differs between them; everything below it does not. A ROCm column (2\(\times \) 7900 XTX) is — pending hardware.

Global batch 256 as \(4 \times 64\), heavy-ball momentum with coupled \(\ell _2\), cosine with a 5-epoch warmup at peak \(0.1\), weight decay \(10^{-4}\), label smoothing \(0.1\) — the same recipe, knob for knob, on both lowerers. The paper schedule at 90 epochs, completed in \(37.1\) hours on one attempt, with no thermal rest and no PCIe interruption.

The two paths agree to \(0.02\) points. Same architecture, same recipe, same box, same \(50{,}000\) images under the same protocol — but one lowered through JAX and the other through pretty(provenGraph) into XLA with no Python at run time. It is the closest verified-versus-reference pair in this book.

It is also a harder test than it looks. The JAX trainer runs under @jit with a NamedSharding mesh, so its BatchNorm means carry global semantics and XLA reduces them across all four cards: the statistic group is the whole \(256\). This run is four replicas of \(64\) with no collective touching the batch statistics, so its group is \(64\) — a fourfold difference that moved the answer by two hundredths of a point.

The two lowerers agree as tightly as the two paths do. Run from the same initialization on the same box, IREE and PJRT scored 17,312 and 17,313 correct at epoch 5 — one image apart. The proof-carrying tier stops at Imagenette, so the claim is one architecture, two independent lowerings, agreeing, not “proven at ImageNet scale”.

A reasonable stopping point

For the fastest path through this book, the route is roughly:

• The intro chapters & Ch 1: framework setup (mandatory context).

• Ch 2–3: MLP and CNN worked examples, the simplest cases of the framework reaching real architectures.

• Skim Ch 4 (BatchNorm). It’s structurally the hardest chapter and the prose says so up front. The math is right (pdiv_bnNormalize’s three-term cancellation is formalized and matches the codegen), but it’s not the place to spend time if you’re not specifically interested in BN. Spend time running as many MNIST and CIFAR demos as possible.

• Read this chapter (Ch 5, ResNet-34). Residual blocks are where the chain rule shows it can carry composition through arbitrarily deep architectures, and the proof reduces to one additive-fan-in lemma you’ve already seen. Run the ResNet-50 RSB-A3 ImageNet training demo locally and/or in the cloud (\(\sim \)6–9 hrs, \(\sim \)$15 in compute).

• Stop here. Assume Chapters 6–9 (MobileNet, EfficientNet, ConvNeXt, ViT) and the Bestiary (Ch 10) are correct, because they all reuse the same framework with new operators slotting into the same proof tree. Appendix C walks through how that claim is auditable if you ever want to spot-check it.

Reading through Ch 5 is enough to know the framework works. Read the AlphaGo Zero paper (bestiary entry), which combines reinforcement learning with a large residual network of convolutions.

5.8 ResNet-50 on ImageNet: the 2018 recipe and RSB-A3

Why the phase-2 trainer? The ImageNet numbers in this book come from the phase-2 (Lean\(\to \)JAX) trainer: at this scale its job is to validate the framework’s logic end to end, and it is the reference the verified path is measured against. Phase-3 verified-IREE codegen led to phase 4 — PJRT, which is the verified path itself. Where a section has both, it reports both.

The basic block ResNet-34 is built from generalizes to the bottleneck block of ResNet-50, and the jax/ folder ships it as a full ImageNet demo (jax/MainResnet50Imagenet.lean) trained on the modern “ResNet Strikes Back” A3 recipe (Wightman et al. 2021): LAMB at effective batch 2048 (via gradient accumulation), BCE over Mixup/CutMix soft targets, RandAugment, train@160 / eval@224. On the same proven-backward bottleneck codegen it reached \(\mathbf{78.26\% }\) top-1 / \(\mathbf{93.79\% }\) top-5 at 100 epochs, clearing the paper’s \(78.1\% \) and the \(78.05\% \) that timm’s own resnet50.a3_in1k weights score under the same protocol.

The JAX lowerer

ResNet-50 on the same four 4060 Ti cards, both recipes, both scored over all \(50{,}000\) validation images. A3 accumulates \(4\times 512\) to reach effective batch \(2048\), so BatchNorm normalizes per \(512\)-image micro-batch (“Ghost-BN”); the 2018 recipe runs a true global batch of \(256\).

The first row is the 2018 recipe on the same network. It is what makes the recipe comparison symmetric: both paths get both recipes, so a difference between two rows is attributable to the one thing that changed between them. Here that difference is \(\mathbf{+1.31}\) points for A3 (\(76.95 \to 78.26\)), with architecture, box, protocol and denominator all held fixed — and A3 reaches it in less wall clock, because it trains at \(160^2\).

ResNet-50 / ImageNet-1k validation accuracy per epoch on the 2018 recipe, four 4060 Ti cards, on the same axes as the A3 curve in §5.8 below. It shows the three-part shape ResNet-34 does: a fast climb, a long shallow middle, and a final lift as the cosine anneals.

Against A3 the two are closer through the middle than the \(+1.31\) endpoint gap suggests. From roughly epoch 25 they trade the lead repeatedly, and A3 does not hold it outright until epoch 86; at the matched epoch 90 A3 leads by \(0.77\) (\(77.72\) against \(76.95\)), and its remaining ten epochs of anneal add the rest. A3 also starts far lower — \(0.32\% \) at epoch 1 against \(5.62\% \) here — because BCE over mixup targets gives almost no signal until the logits separate. Compare the middles with care: the schedules are different lengths, so at a shared epoch index the two cosines are at different points in their decay. The endpoints are what the \(+1.31\) is measured from.

What the re-run corrected, and what it was worth. Beyond the protocol change above, two optimizer corrections read out of timm’s own source went in: LAMB’s gradient clip, which timm enables by default (max_grad_norm \(= 1.0\)) and we did not, and its guard excluding the no_weight_decay group from layer adaptation, which we had been applying to every parameter. Those two, the padding fix and the new evaluation are together what moved this section’s ResNet-50 onto timm’s own footing — a fidelity correction, not a recipe change, which is why the \(78.26\% \) is quotable against timm’s \(78.05\% \) rather than an improvement on it.

The PJRT lowerer

The same recipe on the verified codegen path, which is not a hand-written emitter but pretty(provenGraph), the rendered MLIR the proofs reason about, lowered through ffi/pjrt_ffi.c to XLA over the PJRT C API with no Python at run time. The artifact is resnet50in160_lambaccdp8x64wxclipbce_train_step: LAMB \(\times \) BCE-with-logits \(\times \) gradient accumulation \(k=8\), four replicas at per-replica batch 64 for an effective batch of 2048, train@160 / eval@224.

The first row is the same network on the 2018 recipe, and it is the row that makes the comparison below a controlled one: it is bit-for-bit the ResNet-34 configuration of §5.7 with the bottleneck backbone substituted, so the only difference between it and the A3 row is the recipe. It runs with:

CUDA_VISIBLE_DEVICES=0,1,2,3 PJRT_REPLICAS=4 LEAN_MLIR_REPLICAS=4 \
  SHIM_WORKERS=8 PJRT_FFI_RESIDENT=1 \
  LEAN_MLIR_VARIANT=momdp64bf16 LEAN_MLIR_BATCH=64 LEAN_MLIR_EPOCHS=90 \
  LEAN_MLIR_RECIPE=2018 LEAN_MLIR_BASE_LR_U=100000 \
  .lake/build/bin/resnet50-imagenet-verified data

LEAN_MLIR_RECIPE selects the augmentation. shimScript is a field on the network, not on the recipe, so the \(224^2\) spec has one shim per recipe and picking the wrong one trains the right optimizer on the wrong data. scripts/jobs/ encodes that and the schedule knobs as refusing prechecks.

How to quote these. The A3 row lands \(0.35\) points behind the JAX reference above (\(78.26\% \) / \(93.79\% \)) and \(0.14\) behind timm’s own \(78.05\% \), measured over the same \(50{,}000\) images under the same protocol. The recipe is not yet A3 to the letter — the ledger below is what is in and what is still out — so it is the A3 recipe on the verified path, with these deltas, rather than “RSB-A3 reproduced”. The 2018 row needs no such hedge: it is the recipe in full, and it lands \(0.12\) points ahead of its own JAX reference.

The verified path’s own curve, on the same axes as the 2018 panel above so the two can be read against each other. The endpoint is \(0.35\) points below the reference, but the shape getting there is not a uniformly worse version of it: the verified path is ahead for 56 of the 100 epochs, by as much as \(14.3\) points at epoch 31, and the reference does not take a durable lead until epoch 86 — the same late-anneal window in which A3 finally passes the 2018 recipe. Two runs can differ by a third of a point at the end and by fourteen in the middle, so a mid-training comparison between these panels says almost nothing about where they land.

What that gap is not evidence of: the two paths differ in more than the codegen. The reference normalises BatchNorm over \(512\) and this run over \(64\), an eightfold difference in the statistic group and much the largest known asymmetry between them, and this run is fp32 where the reference is bf16. Attributing the \(0.35\) to the verified codegen requires closing those first.

In, and in the \(77.91\% \) above. What this run carries that a bare A3 transcription would not:

• wdExcludeNormBias — BN \(\gamma /\beta \) and biases are not decayed; 54 decayed / 107 excluded of 161. Largest single contributor.

• LAMB’s gradient clip (timm’s max_grad_norm \(= 1.0\) default), applied after the all-reduce and after the accumulation, at threshold \(kC\) (Proofs.clipFactor_accum).

• The no_weight_decay group is not layer-adapted.

• BN running-statistic momentum compensated for accumulation. Evaluation-only.

• Evaluation scores all \(50{,}000\) images. Every top-1 outside this chapter is still over \(49{,}920\).

• RandAugment Posterize, plus timm’s validation protocol and resampler (appendix A).

• Verification: the optimizer is tied to the reference at the update as well as the gradient — six variants to \(\sim \! 10^{-7}\).

Still out, in rough order of expected effect:

• Ghost-BN group 64 against the reference’s 512; no collective touches the batch statistics. The largest known difference between the two paths, and untested: the \(-0.35\) points between them is its upper bound only if nothing else contributes.

• fp32, not bf16. The 2018 row above is bf16, but the LAMB + BCE shape at \(160^2\) still has no bf16 twin to run. (A2 and A1 at \(224^2\) now have theirs, so this is A3’s gap alone rather than the shared one it was.) Worth a measured \(1.54\times \) on this network, so most of the \(2.1\times \) wall clock against the JAX row: a throughput lever, not a fidelity one.

• \(0.08\% \) of each epoch dropped (\(8 \nmid 5004\)); structural to \(k\).

• Mixup \(\lambda \) drawn host-side from NumPy, so agreement is distributional, never per-step.

The two paths, side by side

Both recipes have now been run on both lowerers, which is what makes this a grid rather than a pair of anecdotes.

The two lowerers agree within \(0.35\) points of top-1 on both recipes, and both agree to within \(0.05\) on top-5. That is the claim the grid licenses and the reason the second lowerer exists: with one backend the proofs decorate a code generator and nothing checks the generator, while a second independent lowering makes the verified path the artifact and JAX the oracle it is checked against.

And the recipe delta replicates. A3 beats 2018 by \(+1.31\) points on the JAX lowerer and by \(+0.84\) on the verified one. Agreement at a single configuration can be a coincidence of that configuration; a difference that reproduces across two independent implementations is harder to arrange by accident, and it is what a one-row comparison cannot show.

What this is not evidence of. Each cell is one run at one seed, and \(0.35\) points is inside what seed variance alone produces on ImageNet — the agreement is consistent with the lowerers matching, not proof that they do. The A3 column trains at \(160\) and evaluates at \(224\) while the 2018 column does both at \(224\), so the recipe delta confounds resolution with recipe, by A3’s own design. And the PJRT column is not internally uniform: its 2018 entry is bf16 on the four-3060 box, its A3 entry fp32 on the four-4060 Ti box. Neither difference touches the within-recipe comparison, which is the one the table is for, but both would matter to anyone reading down the PJRT column.

What A3 changes from the 2018 recipe

The 2018-recipe row above is the original paper’s SGD-with-momentum plus the “bag of tricks” polish (cosine schedule, warmup, label smoothing, random-resized-crop). RSB-A3 is the 2021 “ResNet Strikes Back” recipe (Wightman et al.). Both rows are ResNet-50, which is the point: every knob that A3 sets differently is one lever of the three-year jump in how these networks are trained, and holding the architecture fixed is what stops the comparison confounding the recipe with the backbone. The spec, the render and the entry point are identical between them, and only the arguments differ. Here is the full diff:

What carries over unchanged: cosine LR decay, the 5-epoch warmup, the base augmentation in the row above, bf16 matmul and convolution, and running-BN evaluation. The through-line of the 2018\(\to \)2021 shift is regularize harder, train bigger: heavier augmentation (RandAugment + Mixup + CutMix), a loss that accepts soft targets (BCE), an order-of-magnitude more weight decay applied selectively, and a large-batch optimizer (LAMB) to make the longer, noisier training converge, with FixRes buying back the compute the extra epochs would cost. Every one of these is a one-line TrainConfig change on the same verified backbone.

5.9 Side quest: the A2 and A1 recipes

A3 is the cheap tier. The other two are recipes — a2-accum and a1 — and both now render on the verified path:

Eight artifacts: each tier at one and four replicas, in both precisions, all at \(4\times 128\) for the reason the batch-norm paragraph gives. A1 is not “the same render as A2”: the decay is baked, so its \(0.01\) against A2’s \(0.02\) is a re-render, kept apart so one line among \(18{,}000\) cannot be lost.

The two regularisers these renders used to be missing are now in them. The ledger §5.8 keeps for A3, before anything is run:

Both were structural, and A3 met neither — its recipe switches both off. The shadow and the accumulator used to be the same fourth region of \([\theta |m|v|\cdot ]\), and accumulation is not optional at \(224^2\): EMA was unreachable, not unrendered.

Stochastic depth’s site sits on the residual branch, and no structural check can tell that from a site on the block output: at an all-ones mask \(1 \odot (b+x)\) and \(1 \odot b + x\) agree bit-for-bit, so a misplaced render matches on names, operation count and arity and passes every endpoint gate. Moving all sixteen separates them by eight orders of magnitude at a real mask, and by exactly zero at an all-ones one. The reference had it wrong the other way round — a scalar bernoulli for the whole batch where timm’s is per sample, same expectation, which is why it survived.

Which leaves the batch-norm group, larger than it looks. Nothing here all-reduces activation statistics, so each replica normalises over its own \(64\); the reference’s jnp.mean reduces over a mesh-sharded axis, so XLA inserts one and its group is the full \(512\) — its per-device tensor is \(128\), and only the group is the statistic. These renders take the first half of that gap: \(k=4\) at \(128\) per device is the reference’s own factorisation, and it is the faster arrangement besides. The last \(4\times \) needs synchronised batch-norm — a new operator, a changed backward, a variance that does not average, and \(53\) collectives a step.

And what the phase-2 peers cost, A3 the control:

And the verified column is measured, not scaled from theirs. A3’s graph runs \(191.4\) ms/step on four cards and the complete A2 at \(4\times 128\) bf16 runs \(376.9\), while A3’s finished run fixes what the trainer adds at \(38.8\) ms per \(256\) images — additive, since the device figures already contain the all-reduce and only the host’s blob patch sits outside. That is \(19.0\) verified minutes an epoch, and run backwards it puts A3 itself at \(23\) hours rather than \(32.1\).