9 Vision Transformer

Where Part 1 has been heading

Every previous chapter has introduced a new architectural primitive and proved its backward pass. Each one slotted into the VerifiedNetSpec type and earned its own hasVJP theorem, and that includes convolution, BatchNorm, residual, depthwise, SE, LayerNorm and GELU. This chapter does the same for attention, but it also does something larger: it composes every layer the framework has proved so far into one machine-checked backward pass over a complete state-of-the-art image classifier.

The destination is Theorem 138 at the end of the chapter. It says: the whole ViT-Tiny’s hand-written input-gradient chain — through the patch embedding, twelve transformer blocks with their own parameters, final LayerNorm and head, each slot the certified backward of its op at its own saved activation — equals the network’s Jacobian transpose at every input, assuming only \(\varepsilon {\gt} 0\) at the LayerNorms. The proof composes everything this chapter introduces, plus most of what the previous chapters introduced, into a tower of vjpMatComp applications with vjpComp at the joins. Read the proof’s dependency list and you see a tour of Part 1: dense, residual, layer norm, GELU, softmax, matrix multiply, plus the sixteen matrix-machinery lemmas this chapter sets up. None of them is new in its own right. What’s new is how high the composition rule can be pushed.

9.1 Run it first

Before any of the attention mathematics, train the network that needs it. Four commands, and the last network in Part 1 is also the quickest one in it:

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

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

[pjrt_ffi] XLA backend: PJRT 0.112, 1 device(s)
[pjrt_ffi] compiled verified_mlir/vit_adam_train_step.mlir
             (@vit_adam_train_step, 603 outputs, 1 replica)
             in 6104 ms
[pjrt_ffi] compiled verified_mlir/vit_fwd.mlir
             (@vit_fwd, 1 outputs, 1 replica) in 1644 ms
ViT-Tiny on Imagenette 224² (patch-16 → CLS+pos → 12 transformer
  blocks @ dim192/3heads/MLP768 → final LN → CLS-head 10)
  via the VERIFIED renderer → XLA/PJRT → GPU
  train 9469, val 3925; bs 32, ViT-Tiny adam
    (cosine+warmup 5ep, baseLR 0.000300), He init
[pjrt_ffi] RESIDENT: @vit_adam_train_step holds 600 parameter
             tensors (63.2 MB) on 1 device; they stop crossing PCIe
             from here
Epoch 1/80: loss=2.127541 lr=0.000060
  epoch 1: val_acc = 1381/3925 = 35.184713%  top5 = 3121/3925 = 79.515924%
Epoch 2/80: loss=1.869707 lr=0.000120
  epoch 2: val_acc = 1619/3925 = 41.248408%  top5 = 3350/3925 = 85.350318%
Epoch 3/80: loss=1.750333 lr=0.000180
  epoch 3: val_acc = 1591/3925 = 40.535032%  top5 = 3362/3925 = 85.656051%
Epoch 30/80: loss=1.004060 lr=0.000225
  epoch 30: val_acc = 2571/3925 = 65.503185%  top5 = 3698/3925 = 94.216561%
Epoch 79/80: loss=0.515797 lr=0.000000
  epoch 79: val_acc = 2698/3925 = 68.738854%  top5 = 3551/3925 = 90.471338%
Epoch 80/80: loss=0.515159 lr=0.000000
  epoch 80: val_acc = 2698/3925 = 68.738854%  top5 = 3549/3925 = 90.420382%
done (trained ViT-Tiny adam + cosine/warmup via packed threading).

Eighty epochs at about 18 seconds each, twenty-four minutes in total, and 68.74% top-1 with 90.42% top-5 on Imagenette’s 3,925-image validation split. The best epoch reached \(68.92\% \). That is the fastest network in this book’s Imagenette column by a wide margin and also the least accurate one in it (§9.4).

Two compile lines, and the second one is the whole eval story. Like ConvNeXt, and unlike the three chapters whose BatchNorm keeps running statistics, ViT emits no vit_fwd_eval. 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 top-5 column goes the wrong way, and that is the chapter’s result in miniature. Top-1 climbs the whole run, 35.2 to 68.7. Top-5 climbs to \(94.88\% \) at epoch 25 and then falls, finishing 4.5 points down at \(90.42\% \). On a ten-class problem top-5 is a weak question, so what that divergence measures is a network growing more confident about its top choice while its ranking of the remaining classes decays. A transformer with 5.5M parameters and 9,469 training images has enough capacity to memorize its way to a better top-1 and not enough signal to keep the tail honest, which is §9.4’s data-hunger argument showing up inside a single run rather than across the comparison table.

Attention is where you would expect this to break. What ran is pretty(provenGraph), the file verified_mlir/vit_adam_train_step.mlir, and every softmax in it backpropagates by Theorem 116, every scaled dot-product attention by Theorems 118–120, every LayerNorm by Theorem 126, every GELU by Theorem 97, and every residual skip by Theorem 104.

And the fold holds at every depth. vitForwardKVHasVJP chains the patch embedding, \(k\) transformer blocks with their own parameters, the final LayerNorm and the head through vjpComp for every \(k\), and vitTinyHasVJP_correct says the hand-written chain vitInputGradK is that fold’s Jacobian transpose at this network’s shape — twelve blocks, \(224/16\), any number of classes. Its only hypothesis is \(\varepsilon {\gt} 0\) at the LayerNorms. It is a single global HasVJP over the entire network for the same reason Chapter 8’s is — softmax, GELU and LayerNorm are all smooth — so ViT gets the strong sentence too: every operation in the graph that just trained carries a proved backward, and their composition into the whole network carries one too.

Attention in one equation

A transformer encoder block does two things to its input sequence \(X \in \mathbb {R}^{N \times d}\) (think of \(N\) image patches, each represented as a \(d\)-vector):

1. Multi-head self-attention (MHSA): each output row is a weighted average of all input rows, where the weights are computed from the input itself.

2. Per-token MLP: a small two-layer MLP applied independently to every row.

Both halves are wrapped in a LayerNorm + residual sandwich (exactly the same template as ConvNeXt’s blocks from Chapter 8, and the layout that transferred to convnets in 2022 actually came from transformers in 2017).

The interesting half is the first one. Scaled dot-product attention computes, for queries \(Q\), keys \(K\), and values \(V\) (all linear projections of \(X\)):

In words: every row of \(Q\) is compared to every row of \(K\) via inner product (the \(QK^\top \) matrix is \(N \times N\), one entry per query-key pair). Each row of that \(N \times N\) matrix is passed through softmax to get a probability distribution over the keys (which positions to attend to). Those probabilities then weight the rows of \(V\) to produce the output. The \(1 / \sqrt{d_h}\) scaling keeps the softmax from saturating at large \(d_h\). Multi-head attention just runs several copies of this in parallel on lower-dimensional projections, concatenating the outputs.

Each token, meaning each row of \(X\), ends up looking at every other token. That’s the qualitative break from CNNs: convolutions look only at a local neighborhood, attention looks everywhere. The receptive field is global from the first layer.

Why this chapter has thirty-seven theorems

The math of attention is matrix algebra. Chapter 1’s foundation rules (chain, sum, product, identity) were stated and proved for vector functions. Before we can prove backward passes through \(Q K^\top \) and the row-wise softmax, we need the matrix-level lifts of those rules: chain rule on matrices, additive fan-in on matrices, identity on matrices, plus matrix-specific moves like transpose, matmul-with-one-factor-fixed, and scalar scale. That’s the first section of theorems below (Matrix-level machinery, sixteen theorems). Every one of them is a mechanical lift of a vector-level theorem from Chapter 1, proved via the row-projection \(\mathrm{ContinuousLinearMap}\) and the existing chain rule. Pay the cost once, reuse forever.

The second section (Attention proofs, twenty-one theorems) applies the matrix machinery to the actual attention math. Row-wise softmax (a per-row softmax applied to the \(N \times N\) score matrix) has the well-known closed-form Jacobian \(p_j(\delta _{ij} - p_i)\), and we prove the row-wise lifting that says this holds across matrices. Scaled dot-product attention’s backward decomposes into three matmul-backward applications (one for each of \(Q\), \(K\), \(V\)) plus the row-wise softmax backward, and we prove each. Then multi-head, then the attention sublayer (with residual), then the MLP sublayer (with residual), then a single transformer block, then a \(k\)-deep transformer tower, then the ViT body, then the whole network at its real shape. Each step is one or two compositions of the previous step. The chain that started at vjpComp (Theorem 10) reaches all the way to a 5.5M-parameter image classifier here.

What’s actually proved

That’s the framework-side payoff of this chapter: every architectural piece used by any Part-1 network has a machine-checked backward. With Theorem 138 for the input gradient and Theorem 141 for the parameter gradients the trainer steps on, the sentence “the gradient computed by this trainer is the mathematically correct one” is no longer a hope or a folklore result. It is a theorem the Lean kernel verifies on every build.

The empirical-side payoff comes after the theorems, in the Example section: ViT-Tiny trained on Imagenette underperforms the ConvNets, fairly dramatically, because 9,469 training images is too small for a transformer’s lack of inductive bias. That’s informative on its own terms: the data-hungry-transformer story made concrete.

9.2 Matrix-level machinery

Attention is fundamentally matrix algebra: queries times keys transposed (\(QK^T\)), softmax over rows of the resulting matrix, matmul against values (\(\cdot V\)). Before we can prove backward passes through these operations, we need the matrix-level extensions of Ch 1’s foundation rules. Each theorem below is the matrix lift of a vector-level theorem from Ch 1, proved by composing the vector version with the row-projection \(\mathrm{ContinuousLinearMap}\). The Lean proofs all live in Proofs/Foundation/Tensor.lean alongside their vector counterparts. We collect them here, where the attention proofs need all of them at once.

9.3 Attention proofs

The ladder for the transformer: the tower’s VJP at every depth, the typed graph that denotes it, ViT-Tiny’s whole chain, and then the batched rungs — the rendered backward, the emitted step, the loss gradient.

9.4 Example: ViT-Tiny on Imagenette

The architecture

ViT-Tiny is DeiT’s (Touvron et al. 2021, arXiv:2012.12877) addition to the family — the original paper’s smallest was ViT-Base — and the shortest spec in Part 1: .patchEmbed chops the image into \(16 \times 16\) patches and projects each to a 192-dim token (196 of them, plus a prepended CLS token), .transformerEncoder runs twelve pre-norm blocks, and a final LayerNorm feeds the CLS token to a dense head. The inset shows one of the twelve transformer blocks, which is two residual sub-blocks, attention then MLP, each pre-normed.

The verified trainer

The pattern is §5.3’s, minus bnChannels, as ConvNeXt’s was: there is no BatchNorm.

import LeanMlir.Verified.NetsCore
import LeanMlir.Verified.Train
-- 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 vitVerified : VerifiedNetSpec where
  name     := "ViT-Tiny"
  slug     := "vit"
  inC      := 3
  imageH   := 224
  imageW   := 224
  data     := .imagenette
  layers   := [
    .conv 3 192 16 16,          -- patch embed 16x16/s16, 224 -> 14x14 = 196
    .param #[192] 2,            -- CLS token
    .param #[197, 192] 2,       -- positional embedding, 196 + 1
    .transformerBlock 192 768,  -- 12 pre-norm blocks @ dim 192, MLP 768
    .transformerBlock 192 768,  .transformerBlock 192 768,
    .transformerBlock 192 768,  .transformerBlock 192 768,
    .transformerBlock 192 768,  .transformerBlock 192 768,
    .transformerBlock 192 768,  .transformerBlock 192 768,
    .transformerBlock 192 768,  .transformerBlock 192 768,
    .transformerBlock 192 768,
    .layerNorm 192,             -- final LayerNorm, per-channel
    .dense 192 10 ]             -- CLS-head 192 -> 10
  -- Stochastic depth, read only by the `*drop` variants. TWENTY-FOUR
  -- sites for TWELVE blocks: ViT drops a block's two residual
  -- branches independently, at one keep shared by the pair,
  -- keep = 1 - 0.1 * (site/2) / 11.
  dropSites := (Array.range 24).map (· / 2)
  dropDenom := 11
  dropRate  := 0.1
-- 2. The schedule. Note it is NOT the other four nets' schedule.
def vitAdamConfig : VerifiedConfig where
  epochs    := 80
  batchSize := 32
-- 3. The entry point: AdamW, cosine with a 5-epoch warmup, baseLR 3e-4.
--    Batch, rate, EMA decay and drop rate are env knobs; the variant's
--    render must already exist.
def runViTAdam (argv : List String) : IO Unit := do
  let variant := (← IO.getEnv "LEAN_MLIR_VARIANT").getD "adam"
  let tsPath := s!"verified_mlir/vit_{variant}_train_step.mlir"
  if !(← System.FilePath.pathExists tsPath) then
    throw (IO.userError s!"{tsPath} missing — it is written by \
LeanMlir/Proofs/Codegen/ViTRender.lean; run `lake build LeanMlir.Proofs.Codegen.ViTRender` first")
  let bs := ((← IO.getEnv "LEAN_MLIR_BATCH").bind (·.toNat?)).getD vitAdamConfig.batchSize
  let emaDecay := match (← IO.getEnv "LEAN_MLIR_EMA_DECAY_U").bind (·.toNat?) with
    | some u => u.toFloat * 1e-6
    | none   => 0.99996
  let baseLR := match (← IO.getEnv "LEAN_MLIR_BASE_LR_U").bind (·.toNat?) with
    | some u => u.toFloat * 1e-6
    | none   => 0.0003
  let dropNet := match (← IO.getEnv "LEAN_MLIR_DROP_RATE_U").bind (·.toNat?) with
    | some u => { vitVerified.toNet with
                  dropKeeps := vitVerified.dropKeepsAt (u.toFloat * 1e-6) }
    | none   => vitVerified.toNet
  dropNet.trainAdamSched { vitAdamConfig with batchSize := bs }
    (argv.head?.getD "data") baseLR 0.9 0.999 5 variant 0.0 1.0 emaDecay
def main (argv : List String) : IO Unit := runViTAdam argv

The body is a patch embedding, twelve blocks and a head. The patch embed is spelled .conv 3 192 16 16, a \(16 \times 16\) kernel at stride 16, which chops the \(224 \times 224 \times 3\) image into \(14 \times 14 = 196\) non-overlapping patches and projects each to 192 dimensions. That it is a convolution rather than a distinct primitive is the point: a patch embedding is a strided convolution whose kernel equals its stride, so it reuses the conv2d backward (Theorem 34) and needs nothing new.

The CLS token and the positional embedding are .param entries. They are learned weights that no layer computes from an input, and the spec’s job is the parameter layout, so it names them rather than folding them into a block. The \(197\) is \(196\) patches plus the prepended CLS row.

.transformerBlock 192 768 is one block per line, twelve of them, each the pre-norm sandwich drawn in the inset above: LayerNorm, MHSA, residual, LayerNorm, MLP, residual.

The arm, and why sixteen tensors is the whole block.

  | param dims kind    => #[(dims, kind)]              -- Ch. 9: CLS, pos-embed
  | transformerBlock d m =>                            -- Ch. 9
    #[(#[d],1),(#[d],2),                               -- LN1
      (#[d,d],0),(#[d],2), (#[d,d],0),(#[d],2),        -- Wq, Wk
      (#[d,d],0),(#[d],2), (#[d,d],0),(#[d],2),        -- Wv, Wo
      (#[d],1),(#[d],2),                               -- LN2
      (#[d,m],0),(#[m],2), (#[m,d],0),(#[d],2)]        -- MLP d->m->d

.param is the escape hatch: a bare learned tensor with no layer around it, which is how the CLS token and the positional embedding get into the list at all. Everything else here is a dense layer wearing a different name — the four attention projections are [d,d] matrices with biases, indistinguishable in the parameter list from Chapter 1’s .dense. What makes them attention is the graph, not the shapes. Sixteen tensors a block, twelve blocks, plus the patch embed, the two .params, the final LayerNorm and the head: 200.

200 parameter tensors, 5,526,346 scalars. vitVerified.toSpecs is kernel-#guarded against ViTLayout.specs, an independently audited hand-list, and the ImageNet spec of §9.6 is pinned to this one at everything but its head. At 1000 classes the count is 5,717,416.

Twenty-four keeps for twelve blocks, and the pairing is the content. ViT drops each block’s attention branch and MLP branch independently, but at the same keep probability, so sites \(2i\) and \(2i+1\) share \(\mathrm{keep}_i = 1 - 0.1i/11\). The driver needs one entry per site because it draws one Bernoulli stream per mask input. Deriving twenty-four evenly spaced keeps from the site ordinal instead would silently unpair them, and no structural check would notice.

Note the learning rate is 0.0003 (a third of what the ConvNets used) and warmup is 5 epochs (longer than the ConvNets’ 3). ViTs need gentler optimization schedules, because the attention softmax is prone to collapse if gradients push logits too hard early, and warmup keeps the first few epochs slow enough to avoid that failure mode.

Results

§9.1 has the run: 80 epochs at about 18 seconds each, twenty-four minutes, \(\mathbf{68.74\% }\) top-1 and \(\mathbf{90.42\% }\) top-5. That completes the Part-1 column, and every row of it is now a measurement on the verified XLA path at the same dataset, the same eighty epochs, and one RTX 4060 Ti:

Columns as in Chapter 7’s table; ResNet-34’s row is the five-seed mean of §5.6 (runs/2026-09-12-r34-ablation-fp32-seeds/) with one seed’s timing.

ViT-Tiny is the fastest per step by a factor of \(1.5\) over the next network and the lowest-accuracy by more than thirteen points. That gap is the data-hunger effect: 9,469 training images is tiny for a transformer. The ViT paper reached ImageNet-competitive accuracy only when trained on ImageNet-21K (14M images) or JFT-300M (300M images). At Imagenette scale the ConvNets’ inductive bias, locality and translation equivariance, actively helps, and the ViT has to learn those properties from data it doesn’t have. §9.6 is the control: the same architecture on 1.28M images reaches \(72.42\% \) over a thousand classes, which is the regime the design was for.

It reproduces exactly. This number was measured three times before being printed, and all three passes came back bit-identical on three different cards: the same \(68.738854\% \) and \(90.420382\% \) at epoch 80, and the same 160 epoch lines the whole way down. The reason is the graph’s shape. XLA selects convolution algorithms per process, which is what moves the convnet chapters’ numbers by a point or two between runs. This graph holds 435 dot_generals and two convolutions, so there is almost nothing for that selection to vary.

9.5 MLIR: Attention

What is already proven. Scaled dot-product attention is \(\mathrm{sdpa}(Q,K,V) = \mathrm{softmax}(QK^{\! \top }/\sqrt{d})\, V\). Its reverse-mode derivative produces three gradients (\(dQ\), \(dK\), \(dV\)), each proven separately (Theorems 118–120, sdpaBackQ/K/V), all built on the softmax backward, which couples every position to every other (Theorem 116). Because \(Q\), \(K\), and \(V\) are three linear projections of the same input \(a\), the input gradient is a three-way fan-in, \(da = dQ\, W_q^{\! \top } + dK\, W_k^{\! \top } + dV\, W_v^{\! \top }\). sdpaHasVJPMat3 bundles the three, and transformerAttnSublayerHasVJPMat composes the whole sublayer (LayerNorm, the QKV projections, attention, the output projection, the residual) through vjpMatComp.

The gap and how we close it. The softmax backward is non-elementwise and the three projections are coupled. The emitted graph denotes transformerAttnSublayerHasVJPMat’s backward. Here is the three-way fan-in at the input (two tokens, model dimension four, single head, with the SDPA backward that produces \(dQ,dK,dV\) elided to its comment):

// dQ,dK,dV = SDPA.back(dOut . Wo^T)  (softmax-weighted attention grads)
%daQ = stablehlo.dot_general %dQ, %Wq, contracting_dims = [1] x [1]
         : (tensor<2x4xf32>, tensor<4x4xf32>) -> tensor<2x4xf32>
%daK = stablehlo.dot_general %dK, %Wk, contracting_dims = [1] x [1]
         : (tensor<2x4xf32>, tensor<4x4xf32>) -> tensor<2x4xf32>
%daV = stablehlo.dot_general %dV, %Wv, contracting_dims = [1] x [1]
         : (tensor<2x4xf32>, tensor<4x4xf32>) -> tensor<2x4xf32>
%daQK = stablehlo.add %daQ, %daK : tensor<2x4xf32>
%da = stablehlo.add %daQK, %daV : tensor<2x4xf32>
// %dxa = LN.back(%da);  %dx = %dOut + %dxa  (residual skip)

Read it against the fan-in: each dot_general carries one of the attention gradients back through its projection (\(dQ\) through \(W_q\), and so on, contracting the model axis), and the two adds sum the three. This is the residual fan-in (Theorem 58) with three branches instead of one, because attention reads the input thrice, so the gradient returns thrice and adds. The softmax-weighted \(dQ,dK,dV\) that feed the dot_generals come from the proven sdpaBackQ/K/V, and the surrounding LayerNorm and output projection chain through their bridges.

One honesty note: the hard part of attention is in that first comment. This listing is the projection fan-in (three linear adjoints and two adds), while the softmax backward that couples every token to every other (softmaxHasVJP, § 9.3) is what // SDPA.back stands in for. We show the fan-in because it is the new structural move. The coupled core is carried by the proofs of § 9.3.

With attention pinned, the verified-codegen thread covers every operator in the modern vision stack, and that means dense, convolution, BatchNorm, residual, depthwise, squeeze-excite, layer scale and attention. Each is emitted as the rendering of a machine-checked derivative, from the same three pieces: denoted IR, bridge theorem, printer.

Caveats.

• The attention bridge is unconditional. Softmax and the linear projections are smooth everywhere, with no smooth-point clause.

• Representative scale (two tokens, \(d = 4\), one head).

9.6 ImageNet recipe

The Imagenette ViT-Tiny above trained on \(\sim \)9.5K images. The full 1000-class run is the same architecture with a 1000-class head, the .imagenet dataset, and the DeiT-Ti training recipe, so it is still patch16, embed 192, 12 transformer blocks and 3 heads, at 5,717,416 parameters. As in every ImageNet section, the recipe runs on both paths — the reference first, then the verified path, then the two side by side — and this is the chapter where the pair is cleanest, because a transformer has no BatchNorm to normalize differently.

The reference (JAX)

The reference trainer is jax/MainVitImagenet.lean, quoted as it stood for the run below (today it carries DeiT’s rows as the marker names them, for the rerun):

-- 1. Same ViT-Tiny backbone, 1000-class head.
def vitTinyImagenet : NetSpec where
  name   := "ViT-Tiny (ImageNet, bf16)"
  imageH := 224
  imageW := 224
  layers := [
    .patchEmbed 3 192 16 196,            -- (224/16)^2 = 196 patches
    .transformerEncoder 192 3 768 12,    -- 12 blocks, 3 heads, MLP 768
    .dense 192 1000 .identity            -- 1000-class head
  ]

-- 2. Full DeiT-Ti recipe: AdamW, cosine, the aug suite, stochastic depth, EMA.
def vitTinyImagenetConfig : TrainConfig where
  vitInit           := true          -- timm/DeiT trunc_normal(0.02) init, not Xavier
  learningRate      := 0.0005        -- proper DeiT batch-512 LR
  batchSize         := 512
  epochs            := 300           -- full DeiT-Ti schedule
  optimizer         := .adam         -- AdamW (decoupled weight decay)
  weightDecay       := 0.05
  wdExcludeNormBias := true          -- no_weight_decay: skip norm/bias/pos-embed/CLS
  cosineDecay       := true
  warmupEpochs      := 5
  minLR             := 0.00001       -- DeiT --min-lr 1e-5
  labelSmoothing    := 0.1
  gradClipNorm      := 1.0           -- the unlock for the 5e-4 LR (see below)
  useMixup          := true          -- Mixup alpha 0.8
  useCutmix         := true          -- CutMix alpha 1.0 (alternates with Mixup)
  useRandAugment    := true          -- rand-m9-mstd0.5-inc1 (color + geometric)
  randAugmentGeometric := true
  augBicubic        := true          -- PIL-bicubic geometry, as timm
  randomErasing     := true          -- Random-Erase p 0.25
  randomErasingProb := 0.25
  erasingPixel      := true          -- timm mode='pixel', N(0,1) fill
  dropPath          := 0.1           -- stochastic depth (linear 0 -> 0.1)
  useEMA            := true          -- model EMA (eval on shadow)
  emaDecay          := 0.99996       -- DeiT default
  bf16              := true          -- bf16 matmul compute
  lnEps             := 0.000001      -- DeiT's LayerNorm eps, 1e-6
  f32StemHead       := true          -- patch embed + head in fp32
  valEveryEpochs    := 5             -- ImageNet val is data-loading-bound
  repeatedAug       := 3             -- DeiT Repeated Augmentation (see below)

-- 3. Weight init. vitInit above selects trunc_normal(0.02) for every
--    transformer Linear and gives the patch-embed conv PyTorch's Conv2d
--    default for its weight and a zero bias; Xavier-uniform is the
--    emitter's generic path, and
--    at ViT-Ti's d=192 it is 3.6x wider than timm's fixed 0.02. The run
--    reported below is the default; xavier is the A/B arm.
def vitTinyImagenetConfigXavier : TrainConfig :=
  { vitTinyImagenetConfig with vitInit := false }

-- 4. Two named recipes, chosen by a positional CLI arg. This is the
--    reference driver; the verified path's entry point is trainAdamSched,
--    in the next section.
def vitTinyImagenetRecipes : List Recipe := [
  { name := "default", cfg := vitTinyImagenetConfig,       ... },
  { name := "xavier",  cfg := vitTinyImagenetConfigXavier, ... }
]

def main (args : List String) : IO Unit :=
  runRecipeMain "vit-tiny-imagenet" vitTinyImagenet .imagenet
    vitTinyImagenetRecipes args

The two recipes differ in exactly one field, and the trainers they generate differ in exactly one function: init_params, 59 initialization sites, with the schedule, augmentation, EMA and evaluation identical byte-for-byte.

Gradient clipping is not optional. At the DeiT learning rate (peak \(5\times 10^{-4}\) for batch 512) the model collapses to chance the moment warmup ramps past \({\sim }1.6\times 10^{-4}\), train loss pinned at \(\ln (1000)\approx 6.9\) and never recovering. gradClipNorm := 1.0 is the whole fix; the same LR then trains through warmup to convergence.

bf16. ViT is matmul-bound — patch embed, QKV, scores, output, MLP and head are all dense matmuls, no convolution — so bf16 is the only precision flag that applies and there is no bf16Conv to set. It buys \(\mathbf{1.47\times }\): \(205.1 \to 139.2\) ms/step, \(624 \to 919\) img/s, measured on one RTX 3060 at batch 128 (jax/scripts/jax_vit_bench.py tiny 128, synthetic). The multiple is below the convnets’ because ViT-Ti’s matmuls are narrow (\(d{=}192\)) and proportionally more time goes to the fp32 LayerNorm/softmax that bf16 leaves alone.

Global batch 512 as \(4 \times 128\) replicas, the EMA weights scored over all \(50{,}000\) validation images; one attempt, no restarts. DeiT-Ti without distillation publishes \(72.2\% \) / \(91.1\% \), and publishes it for the live weights, not a shadow: this run’s live weights score \(71.86\% \) / \(90.88\% \), \(0.34\) and \(0.23\) under the paper, and the EMA is worth about \(+0.2\) on top of them (McNemar \(p = 0.002\) over the \(50{,}000\) images, so the shadow’s gain is real, if small). At \(n = 50{,}000\) the Wilson 95% interval is \(\pm 0.39\) points either way: on the paper’s number, not under it in any meaningful sense.

The run, sampled every 5 epochs. The slow start is Mixup and CutMix blending the labels, compounded by the EMA shadow lagging the live weights; after that it climbs to the end of the schedule without an overfitting turn.

Where this still differs from the paper. The number matches. The recipe does not yet match exactly, and the honest list is short but not empty.

What matches. Architecture (patch16, \(d{=}192\), 12 blocks, 3 heads, 5,717,416 parameters, LayerNorm at \(\varepsilon = 10^{-6}\)); schedule (300 epochs, 5-epoch warmup, cosine to DeiT’s floor of \(10^{-5}\)); optimizer (AdamW, weight decay \(0.05\), no decay on norms, biases, CLS or positional embedding); the augmentation stack (Mixup \(0.8\) and CutMix \(1.0\) with label smoothing \(0.1\) folded into the mixed target, rand-m9-mstd0.5-inc1 with the SolarizeAdd\(+\)Invert op-set, per-op apply probability \(0.5\) and PIL-matched bicubic geometry, Random-Erase \(p{=}0.25\) in timm’s pixel mode, repeated augmentation \(3\times \)); stochastic depth \(0.1\); model EMA at \(0.99996\); and the patch embedding and classifier head in fp32 under the bf16 blocks. Weight initialization matches as well — every transformer Linear, the CLS token and the position embedding get \(\mathrm{trunc\_ normal}(\sigma {=}0.02)\), and the patch-embed convolution’s weight PyTorch’s Conv2d default — which is what vitInit in the listing above selects. It is worth saying why that one mattered: Xavier-uniform scales as \(1/\sqrt{d}\) against timm’s fixed \(0.02\), so the generic path was \(1.8\times \) too wide at ViT-B, \(2.6\times \) at ViT-S and \(3.6\times \) at ViT-Ti, while the patch embed came out \({\sim }6\times \) too narrow (it divided by the output fan \(d\cdot p^2\) rather than the input fan \(c\cdot p^2\)). The smallest model was the worst case.

What does not match is the marker after the side-by-side table; the verified path inherits every item, since it runs the same \(4\times 128\), clips at the same \(1.0\) inside the render, and streams the same tf.data pipeline. “Matches the paper’s number” and “matches the paper” are different claims, and only the first is established here, on either path.

The further \(+2.3\) points to DeiT-Ti’s \(74.5\% \) headline are a different thing entirely, the distillation variant (a second “distillation token” trained against a strong CNN teacher), which we do not implement.

The verified path (PJRT)

The verified trainer is apps/imagenette/MainViTImagenet.lean, built as lake build vit-imagenet-verified:

-- Identical backbone, 1000 classes instead of 10.
def vitImagenetVerified : VerifiedNetSpec where
  name       := "ViT-Tiny (ImageNet-1k)"
  slug       := "vitin"
  data       := .imagenet
  shimScript := "generated_vit_tiny_imagenet_shim.py"
  layers     := [ ... the Imagenette stack, 1000-class head ... ]

-- 300 epochs, because that is the schedule the reference number
-- above was measured on. Batch is 128 PER DEVICE, so four
-- replicas give the reference's global 512 exactly.
def vitImagenetConfig : VerifiedConfig where
  epochs    := 300
  batchSize := 128
  vitInit   := true  -- timm/DeiT init, as the reference's default recipe above

The head is the only parameter shape that moves, from \(192 \times 10\) to \(192 \times 1000\), which takes the count from 5,526,346 to 5,717,416. The same three kernel #guards that hold ConvNeXt’s two specs together hold these: equal toSpecs.size, equal toSpecs.pop.pop, and a back! of (#[1000], 2).

The collectives are gated one way here. tests/TestViTDpCheck.lean pins the data-parallel step against its single-device peer by handing both replicas the same rows, where the all-reduce mean is an identity. ViT has no BatchNorm, so that identity is exact rather than approximate. But a gate built on duplicated rows is structurally blind to a shard-offset bug, and the check that closes that hole works by giving the replicas genuinely different data. That one now covers the Imagenette vit spec, and it does not yet cover vitin. Two gates that sound alike are not one property, and this chapter has one and a half of them rather than ConvNeXt’s two.

The regularizer knobs are render variants, and here they compose. Weight-decay exclusion, grad clipping and stochastic depth are spelled wx, clip and drop, EMA is a variant too, and vitin_emadp128x4wxclipdropeps0000001bf16 is the render that carries all four at once on top of AdamW and the four-way all-reduce, at LayerNorm \(\varepsilon = 10^{-6}\) (the last field spells the \(\varepsilon \)). It is the artifact the run used: four regions, \([\theta \mid m \mid v \mid \mathrm{ema}]\), 831 operands and 200 all-reduces, with an fp32 evaluation forward at the same \(\varepsilon \) beside it. Theorem 140 reaches its bf16 gradient nodes at its twenty-four drop-path sites, with only the EMA tail outside it. Mixup, CutMix and repeated augmentation never enter the graph at all, because they are data-side and ride the shim. The job, which today launches the rerun’s render named in the marker below:

lake run vit-default-emabf16-4gpu

Global batch 512 as \(4 \times 128\) replicas, the EMA weights scored over the same \(50{,}000\) images; one attempt, no restarts, every epoch scored. The live weights score \(72.15\% \) / \(91.11\% \), which is DeiT-Ti’s own \(72.2\) / \(91.1\) to within a twentieth of a point.

The two paths, side by side

Both columns are ViT-Tiny on the DeiT-Ti recipe, 300 epochs, on the same four RTX 3060 cards, scored over the same \(50{,}000\) validation images. Rows above the first rule differ; rows between the rules are held equal by construction. [TODO: DeiT in five places, on both paths: batch \(1024\) (four replicas of \(256\), one shot), no gradient clip, the scheduler stepped once an epoch from \(10^{-6}\) over all \(300\) epochs, Mixup or CutMix drawn per batch, torchvision’s crop; plus the exact erf GELU, RandAugment’s fill at the image mean, the \(3\times \) repeats placed by an index-level sampler, the EMA decay constant from step \(0\), an fp32 evaluation forward, and the live weights as the headline where the EMA shadow is. In code since 2026-10-11 and not yet run; one rerun of both halves takes all of it — vit-default-jax-4gpu / vit-default-emabf16-4gpu, \({\sim }48 + 70\) h. Still out, disclosed: the patch-embedding bias at PyTorch’s default.]

The two paths agree within \(0.34\) points on top-1 and \(0.21\) on top-5, inside the 95% Wilson interval (\(\pm 0.39\) at \(72\% \) on \(50{,}000\) images), at one seed each. Above the rule sit the lowering, what it costs, and the two things the runs did not share: where the Mixup and CutMix draws are made — once over the global batch on the reference, once per replica’s \(128\) rows on the verified path, whose four producers each mix their own shard, which is also how timm mixes under DDP — and the precision of the evaluation forward, which a re-score of the reference in fp32 puts at \(+0.01\). Everything the first pair did not share — the label smoothing, the precision of the patch embedding and head, the CLS and position initialization, the cosine’s phase — is now held equal, and the batch is the reference’s own \(512\) with no BatchNorm to normalize over a different group. Against that first pair (\(72.31\) / \(72.35\)) the reference moved \(-0.24\) and the verified path \(+0.07\), each rerun also carrying the bicubic augmentation and DeiT’s \(\varepsilon \) and floor, so neither delta is a re-run of its predecessor.

Interval or no, these are not the same model. Both runs score the same \(50{,}000\) images, so the question can be put per image: the reference is right on \(2{,}257\) images the verified model misses and the verified model on \(2{,}438\) the reference misses, and McNemar’s exact test puts that at \(p = 0.009\) (\(0.043\) on the live weights). The verified checkpoint is ahead on this validation set. With one seed per arm that is a statement about two trained models, not about the two paths, and the per-shard Mixup draw is the first place to look for it; what the pair was rerun to settle is narrower, and settled: matched row for row, the verified path does not trail its reference.

The curves run together, and the verified path leads from early on. The two are tied through epoch 35 — \(13.89\) against \(13.99\) at epoch 5, \(48.27\) against \(48.24\) at 25 — and from epoch 40 the verified path sits between a tenth and half a point ahead, to the end, never reversing:

One recipe, two paths to XLA, sampled every 5 epochs and scored over all \(50{,}000\) validation images.

The remaining \(64.9\) against \(45.9\) hours is the data feed, not the renderer. With the shim removed the verified step is \(153\) ms — a \(38.5\)-hour schedule, under the reference’s own — and the fed step is \(299\); the graph has not changed since that was measured on the first pair’s render, which this one matches but for \(\varepsilon \). The gap is the shim: its producers run the same bicubic augmentation that cost the reference \(1.34\times \) over its pre-bicubic run, out of process, and then stream \(602\) KB per image of fp32 through a subprocess pipe, \(308\) MB per step at global \(512\), which is host work the JAX trainer does in-process. Two knobs close most of the pipe’s share and both are throughput-only: producer count, and tf.data op determinism, which serializes the augmentation map and costs \(5.3\times \) on the producers; it moves the order of random draws and nothing else, and the reference has never paid it either. The same accounting holds for every verified ImageNet run in this book, and the convnets’ gaps are read against it. [TODO: re-measure the gap on the shim as it stands, which mixes in place and writes from its own thread; the uint8 wire does not apply to this recipe, and PIL rotate is not built. At DeiT’s \(1024\) the 2026-10-11 smoke of vit-default-emabf16-4gpu read a \(330\) ms floor under a \(638\) ms fed step with four producers; the rerun measures it end to end.]

A note on the training loss, which the first pair could not compare: its curves sat \({\sim }0.47\) apart at epoch 50 while top-1 agreed to a tenth of a point, and that was the smoothing, because label-smoothed cross-entropy sums \(\log p\) over all \(1{,}000\) classes and is hypersensitive to the small-probability tail, which top-1 cannot see. With the smoothing matched the two curves run within \(0.08\) of each other for the whole schedule, the verified path’s lower on average from the third 25-epoch window on and by \(0.03\) at the end — so the model that generalizes better here also fits its training stream better, a stream that is not quite the reference’s, since each producer mixes its own \(128\) rows. The validation accuracies remain the comparison.

9.7 Side quest: scaling to ViT-S and B

The axis. ViT scales on one axis where ConvNeXt scaled on two. Tiny, Small and Base are the same twelve blocks at a wider model dimension, with the MLP at \(4D\) throughout. Depth never moves.

Counts #guard against DeiT-S’s \(22.05\) M and DeiT-B’s \(86.57\) M. These are the graph’s times, all-reduce in and no trainer around them, at four replicas because S and B have no single-device render at all, on the first renders at \(128\) per device. The jobs below run DeiT’s \(1024\) instead — S as four replicas of \(256\) in one shot, B as two accumulated micro-batches of \(4 \times 128\) — at the \(10^{-3}\) that batch sets, which took one renderer feature, the accumulating optimizer tail, beyond the #eval lines.

A full schedule, four cards, \(300\) epochs at the jobs’ own shape — DeiT’s \(1024\), \(1{,}251\) optimizer steps an epoch — with each path’s own trainer timed on real ImageNet (the 2026-10-11 smokes, on the 4060 Ti box):

Steps are \(1{,}024\) rows; the JAX side accumulates two of \(512\), and ViT-B’s verified render two micro-steps of \(4 \times 128\) at \(859\) ms each. Totals include the per-epoch evaluation. These are fed steps, not graph times, and both verified steps are compute-bound, \(9\) and \(5\) ms of them waiting on the shim.

What is queued. Both pairs exist. The verified jobs are vits-default-emabf16-4gpu on vitsin_emadp256x4wxdropeps0000001erfbf16, DeiT’s \(1024\) as four replicas of \(256\) in one shot, and vitb-default-emabf16-4gpu on vitbin_emaaccdp2x128x4wxdropeps0000001erfbf16, the same \(1024\) as two accumulated micro-batches of \(4 \times 128\), exact under LayerNorm — Tiny’s recipe as its marker has it: vitInit, LayerNorm \(\varepsilon = 10^{-6}\), the \(10^{-5}\) learning-rate floor, EMA at \(0.99996\) constant from step \(0\), no clip, the exact erf GELU, bf16 under an fp32 patch embedding and head — and their references are vits-accum-jax-4gpu and vitb-accum-jax-4gpu, both accumulating \(2\times 512\) to the same \(1024\) (the JAX ViT-S step does not fit the arena in one shot where the verified render does). ViT-S’s render peaks at \(10.75\) GiB of the \(11.68\) GiB default arena; ViT-B’s at \(13.83\) GiB, so its job raises the allocator fraction to \(0.97\) (\(15.11\) GiB).

[TODO: as Tiny’s; neither has run —
vits-accum-jax-4gpu / vits-default-emabf16-4gpu,
vitb-accum-jax-4gpu / vitb-default-emabf16-4gpu.]

What Part 1 has established

Every layer primitive shipped in Part 1’s trainers (MLP, CNN, ResNet, MobileNet, EfficientNet, ConvNeXt, ViT) has a machine-checked backward pass. The same VerifiedNetSpec / VerifiedConfig / trainAdamSched pipeline trains all of them with at most config-level changes. The MLIR each emits is between 7 KB (MLP) and 1,317 KB (EfficientNet-B0), and the lowerer compiles all of them the same way.

The bestiary in Part 2 then shows that every other architecture you’ve seen in modern deep learning composes from the same small set of primitives plus a handful of architecture-specific bundled layers. That list runs UNet, YOLO, DETR, Mask R-CNN, DCGAN, CycleGAN, Pix2Pix, DDPM, Stable Diffusion, VAE, AlphaGo, AlphaZero, MuZero, Mamba, BERT, GPT, Whisper, CLIP, LLaVA, SAM, SegFormer, DeepLab v3+, Nyström-former, QANet and Evoformer. Three rules (chain, additive fan-in, multiplicative fan-in). Five Jacobian tricks (diagonal, sparse Toeplitz, binary selection, rank-1 correction, outer product).

That’s the whole framework. Part 2’s bestiary takes the same set of primitives and tours how they compose into every well-known modern architecture, a survey of the field expressed in the proven kit, with a small library of bundled idioms for the patterns that recur.