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 has_vjp 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 84 at the end of the chapter. It says, roughly: the full ViT backbone is a \(\mathsf{HasVJPMat}\), and that backbone is the patch embedding, twelve transformer blocks and a final LayerNorm. The proof composes everything this chapter introduces, plus most of what the previous chapters introduced, into a single tower of vjpMat_comp applications. Read the proof’s dependency list and you see a tour of Part 1: dense, residual, layer norm, GELU, softmax, matrix multiply, plus the eight or ten 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
./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, and §9.4 is about why both halves of that sentence are the same fact.

Two compile lines, and the second one is the whole eval story. Like ConvNeXt and unlike the three BatchNorm chapters, 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. That run did not execute a reimplementation of the network this chapter describes. It executed verified_mlir/vit_adam_train_step.mlir, which is pretty(provenGraph) off the renderer. Every softmax in it backpropagates by Theorem 66, every scaled dot-product attention by §§68–70, every LayerNorm by Chapter 8’s, every GELU by that chapter’s, and every residual skip by Chapter 5’s additive fan-in.

And the fold holds at every depth. vitForward2_has_vjp chains the patch embedding, all twelve transformer blocks, the final LayerNorm and the head through vjp_comp, with vitForward2_has_vjp_correct carrying it back to the forward the spec denotes, and vitForwardKV_has_vjp generalizes it to depth \(k\). 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 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, \(\sim \)15 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, \(\sim \)15 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_i(\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 full ViT body. Each step is one or two compositions of the previous step. The chain that started at vjp_comp in Chapter 1 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. Once we land Theorem 84, 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/Tensor.lean alongside their vector counterparts. We collect them here because Ch 9 is the first and only chapter to use them.

9.3 Attention proofs

9.4 Example: ViT-Tiny on Imagenette

This is the capstone. Every layer proved in the preceding chapters composes into this one architecture, and that list runs dense, ReLU, softmax cross-entropy, convolution, maxPool, batch norm, residual, depthwise conv, squeeze-and-excitation, layer norm, GELU, and the attention mechanism itself. If the VerifiedNetSpec below compiles and the VJP of its body is proved (§ 84), then every layer in modern deep learning you’ve seen in this book has its backward pass machine-checked.

ViT-Tiny (Dosovitskiy et al. 2020) was the smallest Vision Transformer variant in the original paper. It’s also the least data-hungry, and the big ViT variants start beating ConvNets only at ImageNet-21K scale and above. On a 9469-image Imagenette training set, ViT-Tiny underperforms every ConvNet in Part 1. That’s part of the teaching point: transformers are not a universal improvement, they are a different scaling regime.

The architecture

ViT-Tiny is 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

§5.3 is the pattern, and every chapter since Chapter 6 has reused it. ViT drops the same one of its four annotated points that ConvNeXt does. 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 vitVerified : VerifiedNetSpec where
  name     := "ViT-Tiny"
  slug     := "vit"
  inC      := 3
  imageH   := 224
  imageW   := 224
  nClasses := 10
  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,
    .layerNorm 192,             -- final LayerNorm, per-channel
    .dense 192 10 ]             -- CLS-head 192 -> 10
  -- Stochastic depth, read only by the `*drop` variants. TWENTY-FOUR
  -- entries for TWELVE blocks: ViT drops a block's two residual
  -- branches independently, at one keep shared by the pair.
  dropKeeps := (Array.range 24).map (fun s =>
                 1.0 - 0.1 * (s / 2).toFloat / 11.0)

-- 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.
def main (argv : List String) : IO Unit :=
  vitVerified.toNet.trainAdamSched
    vitAdamConfig (argv.head?.getD "data")
    0.0003 0.9 0.999 5 "adam"

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 Chapter 3’s proved backward 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, not a stage. There is no repeat count, so twelve blocks are twelve entries, exactly as ConvNeXt’s eighteen were eighteen. Each carries 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. So the listing above is not a description of the network that got trained, it is the object the renderer consumed.

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:

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.

ViT-Tiny is the fastest per step by a factor of \(1.5\) over the next network and the lowest-accuracy by more than sixteen 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.31\% \) over a thousand classes, which is the regime the design was for.

It reproduces exactly. Chapter 3’s width sweep is where this book first had to take run-to-run spread seriously, so 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. That makes ViT the first network since the dense-only graphs of Chapter 2 to reproduce exactly, and for the same reason. 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 (§§ 68–70, sdpa_back_Q/K/V), all built on the softmax backward, which couples every position to every other (§ 66). 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 }\). sdpa_has_vjp_mat3 bundles the three, and transformerAttnSublayer_has_vjp_mat composes the whole sublayer (LayerNorm, the QKV projections, attention, the output projection, the residual) through vjp_comp_at.

The gap and how we close it. The softmax backward is non-elementwise and the three projections are coupled. The emitted graph denotes transformerAttnSublayer_has_vjp_mat’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 of Chapter 5 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 sdpa_back_Q/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 (softmax_has_vjp, § 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

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 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-flavored training recipe, so it is still patch16, embed 192, 12 transformer blocks and 3 heads, at 5,717,416 parameters. This is the phase-2 reference implementation, jax/MainVitImagenet.lean, quoted as it stands today:

-- 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
  learningRate      := 5e-4          -- proper DeiT batch-512 LR
  batchSize         := 512
  epochs            := 300           -- full DeiT-Ti schedule
  useAdam           := true          -- AdamW (decoupled weight decay)
  weightDecay       := 0.05
  wdExcludeNormBias := true          -- no_weight_decay: skip norm/bias/pos-embed/CLS
  cosineDecay       := true
  warmupEpochs      := 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)
  randomErasing     := true          -- Random-Erase p 0.25
  dropPath          := 0.1           -- stochastic depth (linear 0 -> 0.1)
  useEMA            := true          -- model EMA decay 0.99996 (eval on shadow)
  bf16              := true          -- bf16 matmul compute
  valEveryEpochs    := 5             -- ImageNet val is data-loading-bound
  repeatedAug       := 3             -- DeiT Repeated Augmentation (see below)

-- 3. Weight init. 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. vitInit
--    selects trunc_normal(0.02) for every transformer Linear and leaves
--    the patch-embed conv on PyTorch's Conv2d default. The run reported
--    below used THIS recipe, not the bare config above.
def vitTinyImagenetConfigDeitInit : TrainConfig :=
  { vitTinyImagenetConfig with vitInit := true }

-- 4. Two named recipes, chosen by a positional CLI arg. This is the
--    phase-2 reference driver; the verified path's entry point is
--    trainAdamSched, in the phase-4 subsection at the end of this chapter.
def vitTinyImagenetRecipes : List Recipe := [
  { name := "default",   cfg := vitTinyImagenetConfig,         ... },
  { name := "deit-init", cfg := vitTinyImagenetConfigDeitInit, ... }
]

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.

That driver wires .imagenet through tfds streaming. §4 is the verified path’s side of it.

Gradient clipping is load-bearing. 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. The run:

The 4060 Ti row is a per-epoch rate carried out to the full schedule, not a completed run, hence no accuracy. The 3060 row is the completed run: 72.31% top-1 / 91.12% top-5 over all 50,000 validation images, 300 epochs in 34.2 hours, one attempt, no restarts, no thermal pauses. DeiT-Ti without distillation publishes 72.2% / 91.1%, so this is \(+0.11\) and \(+0.02\) — on the paper, not past it in any meaningful sense. Per-epoch validation:

ViT-Tiny / ImageNet-1k validation accuracy per epoch (bf16, 4\(\times \) RTX 3060, CUDA), sampled every 5 epochs across the 300-epoch run and scored over all 50,000 validation images.

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); schedule (300 epochs, 5-epoch warmup, cosine to zero); optimizer (AdamW, weight decay \(0.05\), no decay on norms, biases, CLS or positional embedding); loss (cross-entropy with label smoothing \(0.1\)); the augmentation stack (Mixup \(0.8\), CutMix \(1.0\), rand-m9-mstd0.5-inc1 with the SolarizeAdd\(+\)Invert op-set and per-op apply probability \(0.5\), Random-Erase \(p{=}0.25\), repeated augmentation \(3\times \)); stochastic depth \(0.1\); and model EMA at \(0.99996\). Weight initialization now matches as well — every transformer Linear gets \(\mathrm{trunc\_ normal}(\sigma {=}0.02)\) and the patch-embed convolution 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. Four things, in rough order of how much they could matter:

1. Batch \(512\), not DeiT’s \(1024\). The learning rate is scaled consistently for it (\(5\times 10^{-4}\) at \(512\) is DeiT’s own \(\mathrm{lr} = 5\times 10^{-4}\cdot B/512\)), so this is a smaller-batch run of the same recipe rather than a mis-tuned one — but it is not the paper’s batch, and repeated augmentation interacts with batch size by construction.

2. Gradient clipping at global norm \(1.0\) is ours. It is load-bearing: without it the model collapses to chance the moment warmup ramps past \({\sim }1.6\times 10^{-4}\). Whether DeiT itself clips is not verified in this repo. The config comment calls it “the DeiT default” and that claim has never been checked against DeiT’s main.py; until someone reads the reference implementation, treat it as an addition of ours.

3. RandomResizedCrop and Random-Erasing are a from-scratch tf.data reimplementation, not timm’s code, so they agree by specification rather than bit-for-bit. Random-Erase fills with zero where timm fills with noise.

4. Geometric RandAugment interpolates bilinearly, where timm resolves bicubic from the model’s data config; appendix A carries the measurement and the decision to leave it.

None of these is a bug, and the loss curve is textbook DeiT. But “matches the paper’s number” and “matches the paper” are different claims, and only the first is established here.

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.

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 §9.1 ran at Imagenette scale, exists and builds. It 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"
  nClasses   := 1000
  data       := .imagenet
  shimScript := "generated_vit_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 128 PER DEVICE, so four
-- replicas give the reference's global 512 exactly.
def vitImagenetConfig : VerifiedConfig where
  epochs    := 300
  batchSize := 128

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, and it is worth saying which. 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 the same combination is missing here as there. Weight-decay exclusion, grad clipping and stochastic depth are spelled wx, clip and drop, and vitin_adamdp128x4wxclipdrop carries all three at once on top of AdamW and the four-way all-reduce. EMA is a variant too (vitin_emadp128x4), but no committed artifact combines it with the other three. Mixup, CutMix and repeated augmentation 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:

The columns are the two arms of adamdp128x4wxclipdrop at \(128\) per device, so four cards give DeiT’s own global \(512\) and \(2{,}502\) steps per epoch — the reference’s own figure. fp32 is the default; the run will be bf16, so read the right-hand column, and it will also carry EMA, which vit-default-emabf16-4gpu names and which measures \(151\) ms — one under the arm above, so the shadow is free.

[TODO: run vit-imagenet-verified.]

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,316 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.

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. All three run \(128\) per device, so four cards is DeiT’s own global \(512\) and the recipe’s \(5\times 10^{-4}\) is the rate that batch was set for — which took four #eval lines rather than a renderer feature, because the batch was never the obstacle.

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: