8 ConvNeXt

ConvNets strike back

By 2022 the conventional read on image recognition had flipped. ResNet had been the default backbone for six years. Then in late 2020 ViT (Dosovitskiy et al., arXiv:2010.11929) showed that a pure transformer could match or beat ResNet on ImageNet at sufficient scale, and a year later Swin (Liu et al. 2021, arXiv:2103.14030) layered in hierarchical attention to win COCO/ADE20K on top of that. The research community’s working assumption became “transformers beat convnets for vision now.”

ConvNeXt (Liu et al. 2022, arXiv:2201.03545) was a deliberate stress test of that assumption. The authors started from a standard ResNet-50 and asked: which of the design differences between ResNet and Swin is actually doing the work? They proceeded by applying transformer-era choices to the ResNet one at a time and measuring each contribution in isolation. The full list:

1. Training recipe modernization (300 epochs, AdamW, RandAug, Mixup, CutMix, label smoothing, stochastic depth, layer scale)

2. Stage compute ratio change (from (3, 4, 6, 3) to (3, 3, 9, 3), matching Swin)

3. “Patchify” stem (replace the 7\(\times \)7 stride-2 conv with a 4\(\times \)4 stride-4 conv, matching ViT’s patch embedding)

4. Depthwise convolutions (split spatial and channel mixing, as in MobileNetV2)

5. Inverted bottleneck (1\(\times \)1 expand 4\(\times \), depthwise, 1\(\times \)1 project, MobileNetV2-style)

6. Larger kernel (depthwise 7\(\times \)7 instead of 3\(\times \)3, to approximate Swin’s window attention receptive field)

7. Fewer activations and norms (one per block instead of one per layer, again matching transformer blocks)

8. LayerNorm instead of BatchNorm

9. GELU instead of ReLU

The result: a model that’s pure convolution top to bottom, doesn’t contain a single attention operation, and matches or beats Swin-T on ImageNet at the same compute budget. The conclusion of the paper isn’t “convolutions are back.” It’s “the architectural gap between ResNet and Swin was mostly the training recipe, not the receptive field.” Most of what looked like transformer superiority was modernization that hadn’t been back-ported to the convnet baseline.

8.1 Run it first

Before the two new primitives, train the network that needs them. Four commands, and the heaviest per-epoch net in the book still finishes inside an hour and a half:

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

On one RTX 4060 Ti (CUDA 12.9), from runs/2026-08-12-convnext-imagenette-xla-cuda/. XLA’s startup banner is removed, the network description line is wrapped, and the middle of the run is elided:

[pjrt_ffi] XLA backend: PJRT 0.112, 1 device(s)
[pjrt_ffi] compiled verified_mlir/convnext_adam_train_step.mlir
             (@convnext_adam_train_step, 543 outputs, 1 replica)
             in 6176 ms
[pjrt_ffi] compiled verified_mlir/convnext_fwd.mlir
             (@convnext_fwd, 1 outputs, 1 replica) in 473 ms
ConvNeXt-T on Imagenette 224² (patchify /4 → stem channel-LN →
  [3,3,9,3] blocks @ [96,192,384,768] depthwise-7×7 + channel-LN +
  GELU + layerScale + 3 downsamples 56→7 → GAP → dense)
  via the VERIFIED renderer → XLA/PJRT → GPU
  train 9469, val 3925; bs 32, ConvNeXt-T adam
    (cosine+warmup 3ep, baseLR 0.001000), He init
[pjrt_ffi] RESIDENT: @convnext_adam_train_step holds 540 parameter
             tensors (318.4 MB) on 1 device; they stop crossing PCIe
             from here
Epoch 1/80: loss=2.322982 lr=0.000333
  epoch 1: val_acc = 1485/3925 = 37.834395%  top5 = 2973/3925 = 75.745223%
Epoch 2/80: loss=1.974630 lr=0.000667
  epoch 2: val_acc = 1977/3925 = 50.369427%  top5 = 3485/3925 = 88.789809%
Epoch 3/80: loss=1.691063 lr=0.001000
  epoch 3: val_acc = 2078/3925 = 52.942675%  top5 = 3571/3925 = 90.980892%
Epoch 4/80: loss=1.461296 lr=0.001000
  epoch 4: val_acc = 2295/3925 = 58.471338%  top5 = 3661/3925 = 93.273885%
Epoch 79/80: loss=0.501143 lr=0.000000
  epoch 79: val_acc = 3339/3925 = 85.070064%  top5 = 3819/3925 = 97.299363%
Epoch 80/80: loss=0.501138 lr=0.000000
  epoch 80: val_acc = 3339/3925 = 85.070064%  top5 = 3819/3925 = 97.299363%
done (trained ConvNeXt-T adam + cosine/warmup via packed threading).

Eighty epochs at about 59 seconds each, one hour and nineteen minutes in total, and 85.07% top-1 with 97.30% top-5 on Imagenette’s 3,925-image validation split. The best epoch reached \(85.22\% \). That is the lowest top-1 in this book’s Imagenette column, from much the largest network in it, and §8.3 is about why that is the expected answer rather than a defect.

Two compile lines, where the previous three chapters each had three. ResNet-34, MobileNetV2 and EfficientNet-B0 all emit a third artifact, <slug>_fwd_eval, because BatchNorm behaves differently in training and in evaluation and the evaluation forward has to read frozen running statistics. ConvNeXt has no BatchNorm. LayerNorm computes its statistics from the activation in front of it every time, so there is nothing to freeze, and the training forward and the evaluation forward are the same function. The chapter’s central swap pays for itself in the build before it pays for itself in the mathematics.

That also shows in the first four epochs. Top-1 climbs 37.8, 50.4, 52.9, 58.5, smoothly. Chapter 7 had to explain a jump at epoch 4, where B0’s running statistics were averaging over weights that were still moving fast and the evaluation forward was briefly using statistics that no longer described the network. No such step exists here, because no such statistics exist here.

Three new operations this chapter, still three rules. That run did not execute a reimplementation of the network this chapter describes. It executed verified_mlir/convnext_adam_train_step.mlir, which is pretty(provenGraph) off the renderer. Every LayerNorm in it backpropagates by Theorem 46, every GELU by Theorem 47, every layer scale by its own diagonal adjoint, every depthwise convolution by Chapter 6’s, and every residual skip by Chapter 5’s additive fan-in.

And the fold holds at full depth. convNextForwardTCh_has_vjp chains the patchify stem, all eighteen ConvNeXt blocks, the three downsamples, the pool and the head through vjp_comp, with convNextForwardTCh_has_vjp_correct carrying it back to the forward the spec denotes. It assumes 22 positive LayerNorm epsilons and nothing else. Unlike Chapter 6, which folds a representative stem-plus-two-blocks because relu6’s kink makes a whole-network statement false, this one is a single global HasVJP over the entire network, and it can be because GELU is smooth everywhere. So for ConvNeXt the accurate sentence is the strong one: every operation in the graph that just trained carries a proved backward, and their composition into the whole network carries one too.

Two new primitives

ConvNeXt requires two activations / normalizations that the chapters so far haven’t needed:

LayerNorm (Ba, Kiros, Hinton 2016, arXiv:1607.06450). BatchNorm normalizes per-feature across a batch: it relies on the batch having enough samples for the mean and variance to be stable. That fails for small batches (memory-constrained training) and for variable-length sequences (where “the batch axis” isn’t a clean concept). LayerNorm solves both by flipping the axis: normalize per-sample across the feature dimension instead. Same three-term Jacobian cancellation as BN (Chapter 4), same closed-form backward. Just a different reduction axis. Theorem 46 states this formally.

GELU (Hendrycks & Gimpel 2016, arXiv:1606.08415). A smooth approximation of ReLU: \(\mathrm{GELU}(x) = x \cdot \Phi (x)\), where \(\Phi \) is the standard-normal CDF. ReLU has a kink at zero, which forces the codegen to substitute a subgradient convention (Chapter 2). GELU is differentiable everywhere, with a soft-gate behavior near zero that empirically helps transformers (and, it turns out, modernized CNNs). Diagonal activation Jacobian, same family as ReLU’s diagonal, just with a smooth scalar derivative instead of an indicator. Theorems 43–47 provide the scalar function, its derivative, the Jacobian, and the assembled VJP.

Both primitives are used by the upcoming Vision Transformer chapter (Chapter 9) without modification. ConvNeXt introduces them here so they’re available before ViT needs them.

The ConvNeXt block

Each ConvNeXt block stacks the modernization items into a single residual unit:

Two pieces: the inverted bottleneck (\(1 \times 1\) expand 4\(\times \), depthwise spatial conv at the wide point, \(1 \times 1\) project) is exactly the MobileNetV2 inverted residual from Chapter 6. ConvNeXt is paying the MobileNet idea forward. LayerScale (Touvron et al. 2021) is a learnable per-channel scalar \(\gamma \) initialized to \(10^{-6}\) that multiplies the block’s contribution before the residual add, keeping the block near-identity at init in the same spirit as ResNet’s residual. The spec’s .convNextStage primitive emits \(N\) of these blocks at fixed channel width. The .convNextDownsample primitive emits the LN + \(2 \times 2\) stride-2 conv that transitions between stages.

ConvNeXt-T is four stages, ResNet-style

The full ConvNeXt-T spec is the same stem–stages–head template as ResNet-34. A \(4 \times 4\) stride-4 “patchify” stem replaces the \(7 \times 7\) stride-2 stem. Four stages with block counts \((3, 3, 9, 3)\) at channel widths \((96, 192, 384, 768)\) replace the ResNet stages. A final global-average-pool plus 768-to-10 dense replaces the head. Total: 27.8M parameters, comparable to ResNet-50. Every architectural primitive in the spec is now codegen-backed and has a proved VJP. None of the math in this chapter is genuinely new beyond the LayerNorm-axis rotation and the GELU smoothness.

With a modernized convnet backbone, the dominant axis of variation between runs becomes the training recipe, not the architecture. The example section below makes that concrete. The chapter closes by putting the full recipe side-by-side with the 2018 ResNet-34 recipe of Chapter 5, knob by knob: the four-year jump in how convnets are trained, on one page.

8.2 The theorems

8.3 Example: ConvNeXt-T on Imagenette

ConvNeXt (Liu et al. 2022, arXiv:2201.03545) was the “ConvNets strike back” paper: a deliberate back-port of design choices from the Swin Transformer back into a pure CNN, asking which of those choices were doing the actual work. The answer: depthwise \(7 \times 7\) convs at lower channel counts, LayerNorm instead of BatchNorm, GELU instead of ReLU, and an inverted- bottleneck-style \(4\times \) expansion in every block. None of those ingredients are transformer-specific. The architecture below is pure convolution top to bottom. The modernization is in the recipe, not the receptive field.

ConvNeXt-T is the smallest variant in the paper, \(\sim \)28M params (roughly the parameter budget of ResNet-50). At ImageNet-1K scale the paper reports \(82.1\% \) top-1, beating Swin-T at the same compute budget. At our 9.5K-image Imagenette scale, the data regime is different but the architectural story still holds.

The architecture

ConvNeXt-Tiny replaces the convolutional stages with .convNextStage blocks at four widths (\(96, 192, 384, 768\)) and depths (\(3, 3, 9, 3\)), glued by .convNextDownsample transitions (LayerNorm \(+\) \(2 \times 2\) stride-2 conv). A \(4 \times 4\) stride-4 patchify stem and a global-average-pool head bracket the stack. The inset shows one ConvNeXt block: an inverted bottleneck with a single LayerNorm and GELU, scaled by a learnable \(\gamma \) before the residual add.

The verified trainer

§5.3 is the pattern, and every chapter since Chapter 6 has reused it. ConvNeXt drops one of its four annotated points. There is no bnChannels here, because there is no BatchNorm.

import LeanMlir

-- 1. The network. `slug` names the committed render. No `bnChannels`:
--    LayerNorm keeps no running statistics, so nothing has to be
--    threaded from training into evaluation.
def convnextVerified : VerifiedNetSpec where
  name     := "ConvNeXt-T"
  slug     := "convnext"
  inC      := 3
  imageH   := 224
  imageW   := 224
  nClasses := 10
  data     := .imagenette
  layers   := [
    .conv 3 96 4 4, .layerNorm 96,       -- patchify 4x4/s4 + stem LN
    .convNextBlockCh 96,  .convNextBlockCh 96,  .convNextBlockCh 96,
    .layerNorm 96,  .conv 96 192 2 2,    -- downsample  96->192, 56->28
    .convNextBlockCh 192, .convNextBlockCh 192, .convNextBlockCh 192,
    .layerNorm 192, .conv 192 384 2 2,   -- downsample 192->384, 28->14
    .convNextBlockCh 384, .convNextBlockCh 384, .convNextBlockCh 384,
    .convNextBlockCh 384, .convNextBlockCh 384, .convNextBlockCh 384,
    .convNextBlockCh 384, .convNextBlockCh 384, .convNextBlockCh 384,
    .layerNorm 384, .conv 384 768 2 2,   -- downsample 384->768, 14->7
    .convNextBlockCh 768, .convNextBlockCh 768, .convNextBlockCh 768,
    .globalAvgPool, .dense 768 10 ]      -- head
  -- Stochastic depth, read only by the `*drop` variants: one global
  -- block index, keep_i = 1 - 0.1*i/17 at all eighteen blocks.
  dropKeeps := (Array.range 18).map (fun i =>
                 1.0 - 0.1 * i.toFloat / 17.0)

-- 2. The schedule. ResNet-34's, MobileNetV2's and B0's, unchanged.
def convnextAdamConfig : 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 :=
  convnextVerified.toNet.trainAdamSched
    convnextAdamConfig (argv.head?.getD "data")
    0.001 0.9 0.999 3 "adam"

The structure follows the paper’s macro-design prescription: 4 stages with block counts \((3, 3, 9, 3)\) and channels doubling on every downsample \((96 \to 192 \to 384 \to 768)\), with the heavy 9-block stage at the \(14 \times 14\) / 384-channel resolution. The stem is a \(4 \times 4\) stride-4 conv (mirrored from ViT’s patch-embed), not the \(7 \times 7\) stride-2 stem of a ResNet, and that is one of the seven modernization-recipe items called out in the paper. The same base recipe as the EnetB0 chapter (Adam @ 0.001, cosine, warmup-3, WD 1e-4, label smoothing 0.1) keeps the cross-architecture comparison clean.

.convNextBlockCh c is one block, not a stage. There is no repeat count, so the eighteen blocks are eighteen entries and the stage boundaries are visible as the places where the width changes. Ch names the norm: the LayerNorm inside the block reduces over the c channels at each spatial position, which is ConvNeXt’s real normalization rather than the scalar reduction the bare name suggests.

The arm, and the one tensor with no partner.

  | layerNorm d      => #[(#[d],1),(#[d],2)]           -- Ch. 8: gamma, beta
  | convNextBlockCh c =>                               -- Ch. 8
    #[(#[c,1,7,7],0),(#[c],2),(#[c],1),(#[c],2),       -- dw 7x7 + LN(gamma,beta)
      (#[4*c,c,1,1],0),(#[4*c],2),                     -- expand 1x1 (4x)
      (#[c,4*c,1,1],0),(#[c],2),                       -- project 1x1
      (#[c],1)]                                        -- layer scale

Eight of those nine tensors pair off into weight-and-bias or \(\gamma \)-and-\(\beta \). The ninth, the trailing (#[c],1), has no partner: layer scale is a single learned per-channel multiplier with nothing added after it, initialised to ones so a fresh block is the identity on its residual branch. It is the smallest parameter in the chapter and the one whose backward is a pure diagonal.

The stem and the downsamples are spelled out, not packaged. A .convNextDownsample 96 192 is two entries here, a .layerNorm 96 and a .conv 96 192 2 2, and the patchify stem is likewise a .conv followed by a .layerNorm. The spec’s job is the parameter layout, not the block vocabulary, so it names weights rather than concepts.

182 parameter tensors, 27,827,818 scalars. convnextVerified.toSpecs is kernel-#guarded against ConvNeXtLayout.specs, an independently audited hand-list, and the ImageNet spec of §8.5 is pinned to this one at everything but its head. At 1000 classes the count is 28,589,128, which is timm’s figure for convnext_tiny exactly. So the listing above is not a description of the network that got trained, it is the object the renderer consumed.

Twenty-three LayerNorms are the entire hypothesis budget. One in the stem, one per block, one per downsample, and one on the head between the pool and the classifier. Theorem 46 wants \(\varepsilon {\gt} 0\) at each, and those 23 positivities are the only assumptions anywhere in the chain.

Results

§8.1 has the run: 80 epochs at about 59 seconds each, one hour and nineteen minutes, \(\mathbf{85.07\% }\) top-1 and \(\mathbf{97.30\% }\) top-5, with the loss settling on the same \({\sim }0.50\) label-smoothing floor as the three chapters before it. The four-chapter comparison — same dataset, same recipe, same eighty epochs, and every row a measurement on the verified XLA path, on one RTX 4060 Ti:

Step time is the epoch wall-clock divided by 295 batches, so it includes each epoch’s validation pass. MLIR sizes are the committed <slug>_adam_train_step.mlir in 1000-byte units.

• The largest network in the column finishes last, and by a clear margin: \(85.07\% \) against EfficientNet-B0’s \(89.96\% \), from \(6.9\times \) the parameters. Two things contribute. ConvNeXt-T carries 27.83M parameters against 9,469 training images, which is not the data scale that capacity wants. And the recipe here is the shared one, Adam at \(10^{-3}\) with cosine and label smoothing, which is deliberately not the recipe ConvNeXt was designed around.

• That is the chapter’s thesis arriving as a result rather than a claim. ConvNeXt’s paper argues the modernization lives in the training recipe, so a ConvNeXt run on somebody else’s recipe should underperform, and it does. §8.5 is the control: the same architecture, given the full DeiT pack and 300 epochs, reaches \(81.10\% \) on ImageNet-1k and beats every other network in this book. The architecture was never the variable.

• It is no longer the slowest per step. At 196 ms ConvNeXt-T sits under ResNet-34’s 220 ms, which reverses the ordering the phase-3 numbers showed, where ConvNeXt was slowest by \(1.45\times \). The depthwise \(7 \times 7\) and the channel LayerNorms still touch every spatial position twice per block, but that costs less than ResNet-34’s dense \(3 \times 3\) stacks once both graphs go through the same compiler.

• The MLIR is mid-pack at 985 KB, above ResNet-34’s 729 and below MobileNetV2’s 1,047 and EfficientNet-B0’s 1,316. Block count is not what separates them, since ConvNeXt’s eighteen is more than either. What separates them is what one block costs to emit: ConvNeXt spends three convolutions and a single LayerNorm per block, while an inverted residual spends three convolutions and three BatchNorms, and an MBConv adds the squeeze-excitation pair on top of that. One normalization per block instead of one per layer was item seven on the modernization list, and it is visible in the file size.

8.4 MLIR: Layer Scale

Layer scale is a per-channel learnable diagonal, \(\mathrm{layerScale}(\gamma , x) = \gamma \odot x\), and it is the simplest operator in this series. Its Jacobian is diagonal, so it is its own adjoint: layerScale_has_vjp proves \(\mathrm{back}(x, dy) = \gamma \odot dy\), the backward multiplying by the same \(\gamma \) the forward did. There is nothing to bridge past that identity, so the emitted graph is two multiplys against one tensor:

// ConvNeXt block tail:  out = x + layerScale(gamma, project(...))
%ls  = stablehlo.multiply %gls, %pr   : tensor<1x2x4x4xf32> // gamma (*) pr
%out = stablehlo.add %x, %ls          : tensor<1x2x4x4xf32> // + identity skip
// backward, cotangent %dOut:  d(pr) = gamma (*) d(out)
%dpr = stablehlo.multiply %gls, %dOut : tensor<1x2x4x4xf32>

The forward multiply %gls, %pr and the backward multiply %gls, %dOut are the same op against the same \(\gamma \) tensor: a diagonal map is its own transpose, so its VJP is itself. The identity skip is the residual fan-in of Chapter 5. The surrounding GELU, LayerNorm, and convolution backwards reuse their bridges. Uniquely in this book, every one is unconditional (the only hypothesis anywhere is LayerNorm’s \(\epsilon {\gt} 0\)). The full block is \(\mathrm{residual}(\mathrm{layerScale} \circ \mathrm{project} \circ \mathrm{gelu} \circ \mathrm{expand} \circ \mathrm{LN} \circ \mathrm{depthwise})\), composed by convNextBlock_has_vjp.

8.5 ImageNet recipe

[TODO: run with timm tweaks.] The numbers in this section were measured before the trainers moved onto timm’s validation protocol and its antialiased resampler (appendix A). That change moves the training distribution as well as the evaluation, so it needs a re-run rather than a re-score, and none of it is reflected below.

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.

ConvNeXt’s architecture reached this book through the Imagenette example earlier in this chapter. The phase-2 port wires the two ConvNeXt-specific layers (.convNextStage, .convNextDownsample) into the JAX codegen. Because that path differentiates with jax.value_and_grad, only the forward had to be written, and the backward comes for free. The spec mirrors jax/MainConvNeXtImagenet.lean:

-- Faithful patchify stem, stage/downsample backbone, 1000 classes.
def convNeXtTinyImagenet : NetSpec where
  name   := "ConvNeXt-T (ImageNet, bf16)"
  imageH := 224
  imageW := 224
  layers := [
    .convNextStem 3 96 4,              -- patchify: 4x4 s4 conv -> channel-LN (no BN/ReLU)
    .convNextStage 96 3 .ln .gelu,     -- stage 1: 3 blocks @ 96  (56x56)
    .convNextDownsample 96 192,        -- 56 -> 28
    .convNextStage 192 3 .ln .gelu,    -- stage 2: 3 blocks @ 192
    .convNextDownsample 192 384,       -- 28 -> 14
    .convNextStage 384 9 .ln .gelu,    -- stage 3: 9 blocks @ 384
    .convNextDownsample 384 768,       -- 14 -> 7
    .convNextStage 768 3 .ln .gelu,    -- stage 4: 3 blocks @ 768
    .globalAvgPool,
    .dense 768 1000 .identity          -- 1000-class head
  ]

-- ConvNeXt needs AdamW, not SGD; full DeiT-style recipe.
def convNeXtTinyImagenetConfig : TrainConfig where
  learningRate      := 2.5e-4       -- 4e-3 @ batch 4096, linearly scaled to 256
  batchSize         := 256
  epochs            := 80           -- validation tier of an 80->300 ladder
  useAdam           := true         -- AdamW (decoupled weight decay)
  weightDecay       := 0.05
  wdExcludeNormBias := true         -- timm no_weight_decay: skip norm g/b, bias, LayerScale
  cosineDecay       := true
  warmupEpochs      := 20           -- ConvNeXt paper warmup (was 5)
  gradClipNorm      := 1.0          -- cheap insurance (unlocked the ViT run)
  augment           := true
  useRandAugment    := true         -- geometric RandAugment N=2, M=9, mstd0.5, inc1
  randAugmentGeometric := true
  useMixup          := true         -- Mixup alpha 0.8 ...
  useCutmix         := true         -- ... + CutMix alpha 1.0 (alternates per step) ...
  randomErasing     := true         -- ... + Random Erasing p 0.25 (full DeiT pack)
  labelSmoothing    := 0.1
  bf16              := true
  bf16Conv          := true         -- dw-7x7 + 1x1s; LN/GELU stay fp32
  useEMA            := true         -- weight averaging, decay 0.9999
  dropPath          := 0.1          -- stochastic depth (ConvNeXt-T paper value)

-- 300-epoch "full" tier: same config, longer schedule.
def convNeXtTinyImagenetConfigFull : TrainConfig :=
  { convNeXtTinyImagenetConfig with epochs := 300 }

Unlike MobileNet and EfficientNet, ConvNeXt trains with AdamW and a small decoupled weight decay, not SGD. The modernized-convnet recipe is the architecture’s whole thesis. bf16Conv pays off especially here: the \(7\times 7\) depthwise is \(\sim 2.3\times \) faster in bfloat16 (a large kernel has enough arithmetic for the tensor cores to bite, unlike the \(3\times 3\) depthwise, which is a wash), and the inverted-bottleneck \(1\times 1\)s are matmuls that love it. The channel LayerNorm and GELU stay in fp32.

Every element of the paper recipe is now wired and on:

• the faithful patchify stem (.convNextStem: \(4\times 4\) stride-4 conv \(\to \) channel-LayerNorm, no BN/ReLU)

• stochastic depth (drop-path 0.1) and EMA (decay 0.9999), together \({{\lt}}5\% \) per epoch

• decoupled weight decay that excludes the norm/bias/LayerScale 1-D params (wdExcludeNormBias, timm’s no_weight_decay)

• the full DeiT-style augmentation pack: Mixup \(\alpha 0.8\), CutMix \(\alpha 1.0\), Random Erasing \(p0.25\), and geometric RandAugment (\(N{=}2\), \(M{=}9\), mstd0.5-inc1)

The peak LR is \(2.5\mathrm{e}{-}4\) (the \(4\mathrm{e}{-}3\)@4096 official value linearly scaled to batch 256) over a 20-epoch warmup. No architectural deviation from the paper remains, and the schedule is no longer short of it either. The resampler and validation protocol still differ, which is what the marker at the top of this section is about. The run below is the full 300-epoch tier.

Compute budget. ConvNeXt-T is the heaviest of the three at \(28.6\)M parameters, the canonical count, verified by an init/forward/backward/save/reload round-trip. The run below used four of the six RTX 4060 Ti (CUDA, bf16, batch 256 = 4\(\times \)64): after the input-pipeline pass (device-side prefetch double-buffering, so the host loop no longer blocks on each step’s host-to-device copy) steady-state throughput is \(\sim \)177 ms per step, about \(14.9\) minutes per epoch. That is the most expensive per-epoch of any net in this book, driven by the \(7\times 7\) depthwise and the channel LayerNorms. Being compute-bound, ConvNeXt is one of only two nets that still scales to six GPUs (\(\sim \)136 ms/step, \(\sim \)1.3\(\times \)) rather than saturating the host loop the way the lighter nets do. We ran it on four anyway, trading \(\sim \)23% throughput for thermal headroom and staying off the two AER-prone cards (idx 1, 5) so the run could go unattended. A thermal duty cycle held the box stable start to finish with zero AER kills: nine 30-minute cooldowns, each resuming bit-for-bit from a full-state checkpoint (params \(+\) optimizer \(+\) EMA \(+\) step, so the cosine schedule and Adam moments are untouched).

(\(^{\dagger }\)CUDA, bf16. \(\sim \)75 hr pure training and \(\sim \)80 hr actual wall-clock across nine thermal rests plus one benign watchdog restart, a corrected CPU-cache event mis-read as a PCIe fault, with no real hardware fault the entire run.)

The full 300-epoch tier (convnext-tiny-imagenet full, convNeXtTinyImagenetConfigFull) reached \(\mathbf{81.10\% }\) top-1 and \(\mathbf{95.37\% }\) top-5 on the held-out validation pass. That is the strongest result in this book’s sweep by a wide margin, ahead of EfficientNet-B0 (\(72.3\% \)) and ResNet-34 (\(72.1\% \)), and it lands within \(\sim \)1% of ConvNeXt-T’s \(82.1\% \) paper headline from a from-scratch Lean\(\to \)JAX-codegen reimplementation. The residual \(\sim \)1% is EMA-eval / test-crop / augmentation-detail territory rather than architecture — and test-crop is one of the things the pending re-run moves, so read that as the candidate list, not a settled split.

The validation curve tells the augmentation story directly: a deliberately slow start, since Mixup and CutMix blend labels, so the net under-trains for the first \(\sim \)20 epochs, compounded by the EMA shadow lagging the live weights. Then it climbs for the rest of the schedule without an overfitting turn, which is what the heavy-regularization regime predicts and what makes the long schedule worth buying.

Learning-rate robustness. The full pack was also a test of whether ConvNeXt’s AdamW recipe is as learning-rate-robust as EfficientNet’s RMSProp (Chapter 7) proved fragile. It is: at the \(2.5\mathrm{e}{-}4\) peak with global-norm gradient clipping at \(1.0\), top-1 climbed across the whole 300-epoch schedule with no erosion and no divergence, the opposite of B0’s RMSProp blow-up. AdamW’s bounded, decoupled step plus the grad-clip makes the peak LR a non-event, which is what lets the heavy augmentation run to schedule instead of destabilizing it early.

ConvNeXt-T / ImageNet-1k validation accuracy per epoch over the full 300-epoch schedule (bf16, 4\(\times \) 4060 Ti). Top-1 climbs to \(81.10\% \) and top-5 to \(95.37\% \) with no late-schedule plateau or overfitting dip. The nine 30-minute thermal rests are invisible, each resuming bit-for-bit.

Phase 4: the verified trainer

Everything above this point is the phase-2 Lean\(\to \)JAX trainer. The phase-4 peer, which trains the same network through the proof-rendered StableHLO that §8.1 ran at Imagenette scale, exists and builds. It is apps/imagenette/MainConvNeXtImagenet.lean, built as lake build convnext-imagenet-verified:

-- Identical backbone, 1000 classes instead of 10.
def convnextImagenetVerified : VerifiedNetSpec where
  name       := "ConvNeXt-T (ImageNet-1k)"
  slug       := "convnextin"
  nClasses   := 1000
  data       := .imagenet
  shimScript := "generated_convnext_tiny_imagenet_shim.py"
  layers     := [ ... the Imagenette stack, 1000-class head ... ]

-- 300 epochs, because that is the schedule the phase-2 number
-- above was measured on. Batch is 32 PER DEVICE.
def convnextImagenetConfig : VerifiedConfig where
  epochs    := 300
  batchSize := 32

The head is the only parameter shape that moves, from \(768 \times 10\) to \(768 \times 1000\), which takes the count from 27,827,818 to 28,589,128. Three kernel #guards hold the two specs together at everything else: equal toSpecs.size, equal toSpecs.pop.pop, and a back! of (#[1000], 2).

The collectives are gated by two gates, and they prove different things. tests/TestConvNeXtDpCheck.lean pins the data-parallel forward bit-exactly against its single-device peer. That gate hands both replicas the same rows, so it is structurally blind to a shard-offset bug, and tests/TestConvNeXtShardCheck.lean closes exactly that hole by giving the replicas genuinely different data. ConvNeXt is where the shard check was written before being generalized to the other nets, so this is the one chapter where both gates are native rather than inherited.

Four of the paper’s six extra knobs are rendered, and one combination is missing. Weight-decay exclusion, grad clipping and stochastic depth are render variants, spelled wx, clip and drop, and convnextin_adamdpwxclipdrop carries all three at once on top of AdamW and the four-way all-reduce. EMA is a render variant too (convnextin_emadp), but no committed artifact combines it with the other three, so there is no single graph carrying the whole regularizer stack. Mixup and CutMix never enter the graph at all, because they are data-side and ride the shim.

What it has not done is run. No phase-4 ImageNet result exists for this network, which is why every accuracy above comes from phase 2. What is measured is the throughput, and the schedule follows from it:

[TODO: run convnext-imagenet-verified.]

What the ConvNeXt recipe changes from the 2018 recipe

The ResNet-34 trainer of Chapter 5 is the 2018-era recipe: the original paper’s SGD-with-momentum plus the “bag of tricks” polish (cosine schedule, warmup, label smoothing, random-resized-crop). The config above is the 2022 ConvNeXt recipe, DeiT’s transformer-training playbook applied to a convnet. Where Chapter 5’s RSB-A3 side quest showed the 2018\(\to \)2021 jump on a fixed ResNet, this table shows the 2018\(\to \)2022 jump with the block modernized in step. Since “the recipe is the thesis” is this architecture’s whole argument, the diff below is the chapter in table form. The trainAdamSched entry point is identical in both columns, and every row is a one-line change to the VerifiedNetSpec, the VerifiedConfig, or the arguments beside them:

What carries over unchanged: cosine LR decay, the warmup-then-anneal structure, random-resized-crop + horizontal flip as the base augmentation, label smoothing \(0.1\), batch 256, train and eval both at \(224\) (no FixRes), and bf16 matmul and convolution.

The augmentation row is the full DeiT pack: Mixup (Zhang et al. 2017, arXiv:1710.09412), CutMix (Yun et al. 2019, arXiv:1905.04899), Random Erasing (Zhong et al. 2017, arXiv:1708.04896), and RandAugment (Cubuk et al. 2019, arXiv:1909.13719). All four are implemented as data-only kernels that mutate the input tensor before the graph sees it, tier 1 in our codegen scope, no MLIR plumbing required. Mixup and CutMix route through a soft-label train-step variant since they produce fractional labels. Random Erasing and RandAugment leave labels alone.

The through-line is the same regularize harder, train bigger as RSB-A3, arrived at from the transformer side: keep cross-entropy and soften the targets instead of switching losses, AdamW instead of LAMB (the batch stays moderate), stochastic depth and EMA from the ViT toolkit, and a \(3\times \) longer schedule to let the regularization pay.

8.6 Side quest: scaling to ConvNeXt-S and B

The axis. Scaling ConvNeXt-T is two moves. S deepens stage 3 from 9 blocks to 27 and touches nothing else; B does that and widens every stage, \(96/192/384/768 \to 128/256/512/1024\). On the JAX side both are parameter edits. On the verified side they are not the same job at all.

Every number in this section was measured in one session on 4\(\times \) RTX 4060 Ti (16 GB, CUDA 12.9). Counts are #guarded against \(50.22\) M and \(88.59\) M.

A full schedule, four cards, \(300\) epochs, totals including \(37.5\) s per epoch of evaluation and checkpointing. Trainer steps, all measured on real ImageNet over \(40\) steps: