2 MNIST: 1D MLP

One layer learned; now stack them

Chapter 1 built the calculus, proved the dense layer’s backward pass, and trained it. A single dense layer is a linear map, and a linear map can only cut pixel space with flat boundaries. This chapter adds the two things a hidden layer needs: a non-linearity between the dense layers, ReLU, and the chain rule applied more than once. Every theorem here is the same question as before — if I jiggle this input by \(\varepsilon \), how does the output jiggle? — asked about ReLU, and then about the whole stack.

2.1 Run it first

Before any of the math, train the thing. Four commands and about twelve seconds of GPU time:

lake exe cache get                         # Mathlib oleans, ~30 s
./scripts/datasets/download_mnist.sh       # ~11 MB
lake build mnist-mlp-verified
./.lake/build/bin/mnist-mlp-verified data

On one RTX 4060 Ti (CUDA 12.9), from runs/2026-09-01-mlp-verified-xla-cuda/. XLA’s startup banner is removed and two long lines are wrapped:

[pjrt_ffi] XLA backend: PJRT 0.114, 1 device(s)
[pjrt_ffi] command buffers: enabled (CUDA default)
[pjrt_ffi] compiled verified_mlir/mlp_train_step.mlir
             (@mlp_train_step, 7 outputs, 1 replica) in 2123 ms
[pjrt_ffi] compiled verified_mlir/mlp_fwd.mlir
             (@mlp_fwd, 1 outputs, 1 replica) in 157 ms
MNIST-MLP via the VERIFIED renderer (784→512→512→10) → XLA/PJRT → GPU
  xla/pjrt verified_mlir/mlp_train_step.mlir
  xla/pjrt verified_mlir/mlp_fwd.mlir
  train 60000, test 10000; bs 128, MNIST-MLP (6 params, 669706 floats),
    mean-loss SGD lr=0.100000, He init
  epoch 1: loss = 0.413496, test_acc = 9199/10000 = 91.990000%
             [95% CI 91.44–92.51] (870ms)
  epoch 2: loss = 0.191186, test_acc = 9526/10000 = 95.260000%
             [95% CI 94.83–95.66] (879ms)
  epoch 3: loss = 0.139057, test_acc = 9636/10000 = 96.360000%
             [95% CI 95.97–96.71] (888ms)
  epoch 4: loss = 0.109009, test_acc = 9690/10000 = 96.900000%
             [95% CI 96.54–97.22] (842ms)
  epoch 5: loss = 0.088619, test_acc = 9716/10000 = 97.160000%
             [95% CI 96.82–97.47] (847ms)
  epoch 6: loss = 0.073548, test_acc = 9732/10000 = 97.320000%
             [95% CI 96.98–97.62] (864ms)
  epoch 7: loss = 0.061923, test_acc = 9746/10000 = 97.460000%
             [95% CI 97.13–97.75] (802ms)
  epoch 8: loss = 0.052682, test_acc = 9762/10000 = 97.620000%
             [95% CI 97.30–97.90] (818ms)
  epoch 9: loss = 0.045157, test_acc = 9767/10000 = 97.670000%
             [95% CI 97.36–97.95] (841ms)
  epoch 10: loss = 0.038872, test_acc = 9771/10000 = 97.710000%
             [95% CI 97.40–97.99] (831ms)
  epoch 11: loss = 0.033574, test_acc = 9778/10000 = 97.780000%
             [95% CI 97.47–98.05] (847ms)
  epoch 12: loss = 0.029061, test_acc = 9781/10000 = 97.810000%
             [95% CI 97.50–98.08] (850ms)
done (trained MNIST-MLP via the proof-rendered StableHLO).

Twelve epochs at about 850 ms each, and 97.81% on the full 10,000-image test set. Adding two hidden layers to the linear model of Chapter 1 buys 5.7 points of accuracy over its 92.10% (runs/2026-09-01-linear-verified-xla-cuda/) at 850 ms an epoch against its 224.

The linear model has no seed to spend — it starts at zero — so it has no spread to report. This one does. Run at five seeds it averages \(97.73\% \) with a spread of \(0.06\) either side, and every one of those five sits inside the interval printed above — which is \(\pm 0.29\) points on its own. Worth noticing once, because it holds for the rest of the book: changing the seed moves these numbers less than a single run can pin them down. The interval is the wider of the two, and it is the one to check a claimed improvement against.

ReLU: the piecewise case

ReLU is the function \(\mathrm{relu}(x) = \max (x, 0)\) applied coordinatewise. Its Jacobian is a diagonal matrix: each output coordinate depends only on the matching input coordinate. The diagonal entry is \(1\) where \(x_i {\gt} 0\) and \(0\) where \(x_i {\lt} 0\). At \(x_i = 0\) the function is not differentiable in the classical sense, because the slope jumps from \(0\) to \(1\). For our purposes this matters less than you might fear. Theorem 26 states the Jacobian at smooth points. At the corner \(\operatorname {fderiv}\) returns the zero map and the codegen emits the framework’s \([x {\gt} 0]\) mask instead, so the two part company there and agree everywhere else, and every theorem that reaches the emitted mask says so with a smooth-point hypothesis. The corner is the only place in the chapter where “differentiable” becomes slightly subtle.

Composition

A three-layer MLP is \(\mathrm{dense}_2 \circ \mathrm{relu} \circ \mathrm{dense}_1 \circ \mathrm{relu} \circ \mathrm{dense}_0\). The VJP chain rule (Theorem 11) says the backward of a composition is the backwards run in reverse order, so the MLP’s backward is three dense backwards — each a multiplication by that layer’s \(W\) — with ReLU’s diagonal mask applied between them. That is the structural fact the rest of the book leans on: a network’s backward pass is its layers’ backward passes, composed.

There is one catch, and it is ReLU’s. Theorem 11 asks for both functions to be differentiable everywhere, and ReLU is not. So the whole-network result is stated at a point: wherever no pre-activation is exactly zero, the composed backward equals the \(\operatorname {pdiv}\)-contracted Jacobian (Theorem 29). The VJP record mlpHasVJP, with the \(\operatorname {pdiv}\)-derived backward (Definition 28), is the object later chapters name; the theorem is what makes naming it a claim.

2.2 The theorems

Figure 2-1 is the dependency graph of this section.

2.3 Example: MNIST MLP

The theorems above are the calculus. Here is a concrete architecture built from those pieces: a three-layer fully-connected classifier for 28\(\times \)28 MNIST digits.

Dataset overview

MNIST is the standard testbed for image-recognition learning. It’s a collection of 28\(\times \)28 grayscale images of handwritten digits 0–9: 60 000 training images and 10 000 test images, each with a label indicating which digit was drawn. The dataset has been around since 1998. At 784 pixels per image and 10 classes, MNIST is small enough that you can train a competitive model on a laptop CPU in minutes while still having a nontrivial learning problem.

Our goal in this chapter is to correctly classify a held-out test digit based on a model trained from the 60 000 training digits. We’re going to ignore the 2D spatial structure of the image entirely for now, so just flatten each 28\(\times \)28 image into a 784-dim vector and treat it as a plain supervised-learning classification problem. This is the multilayer perceptron (MLP). Chapter 3 revisits MNIST with convolutions that respect the spatial structure.

Architecture

Figure 2-2 is the network.

Three dense layers stacked with ReLU non-linearities between them: \(784 \to 512 \to 512 \to 10\). First layer ingests the flattened image. The two hidden layers let the network learn nonlinear features. The final layer maps to 10-dimensional logits, one per digit class.

The verified spec and program

The network above is a VerifiedNetSpec, the same object type Chapter 1 used for the linear model, now with two hidden layers and ReLU between them:

def mlpVerified : VerifiedNetSpec where
  name     := "MNIST-MLP"
  slug     := "mlp"
  inC      := 1
  imageH   := 28
  imageW   := 28
  data     := .mnist
  layers   := [.dense 784 512, .relu,
               .dense 512 512, .relu,
               .dense 512 10]
  blurb    := "MNIST-MLP via the VERIFIED renderer
                (784→512→512→10) → %LOWERER% → GPU"
  lossSlot := true

Five entries where Chapter 1 had one, and nothing else about the type changes. (layers and blurb are each one line in the source, wrapped here.) The slug binds the spec to verified_mlir/mlp_train_step.mlir and verified_mlir/mlp_fwd.mlir, and Proofs.StableHLO.mlpTrainStepText is what emits the train step. lossSlot is the field behind the loss column in §2.1’s transcript, which Chapter 1’s does not have: it tells the driver this render returns a trailing scalar to print. That scalar is also the one line of verified_mlir/mlp_train_step.mlir that is not pretty of a proven graph node, and the render says so itself, in a comment of its own: %loss below is REPORT-ONLY (logging), NOT pretty(AST node). Nothing reads it but the printer and no theorem depends on it. The six parameter outputs are the proved ones — MlpFold shows each denotes the certified gradient-descent step derived from the Mathlib \(\operatorname {fderiv}\) math.

What .relu contributes.

Nothing:

  | relu => #[]

ReLU carries no parameters, so its toSpecs arm is the empty array and the five-entry layers list above still yields six tensors — three weight matrices and three biases, the 6 params, 669706 floats the run reported. A layer with no state is still a layer: it holds a position in the list because the render needs it in order, not because the optimizer does.

The tie back to the proofs has the same shape as before, now at a point, because ReLU has a corner. In LeanMlir/Proofs/SpecVJP.lean:

noncomputable def mlpVerifiedHasVJPAt (W0 b0 W1 b1 W2 b2) (x)
    (h0 : ∀ k, dense W0 b0 x k ≠ 0)
    (h1 : ∀ k, dense W1 b1 (relu 512 (dense W0 b0 x)) k ≠ 0) :
    HasVJPAt (denoteMLP mlpVerified.layers W0 b0 W1 b1 W2 b2) x := ...

Again the backward — the five per-layer backwards composed, at an input where neither ReLU sits on its corner — is proved for the denotation of the spec’s own layers field, and that is the field the trainer reads to build its graph. The whole program is short enough to print:

import LeanMlir.Verified.NetsCore
import LeanMlir.Verified.Train

def mlpConfig : VerifiedConfig where
  epochs    := 12
  batchSize := 128

def main (argv : List String) : IO Unit :=
  mlpVerified.toNet.train mlpConfig (argv.head?.getD "data")

That is apps/mnist/MainMnistMlpVerified.lean. The only difference from Chapter 1’s driver is .train in place of .trainLinear, because the MLP carries six parameter tensors packed together rather than a separate W0 and b0.

The two compile lines in §2.1 are where that shows up: the train step compiles in 2,123 ms against the linear step’s 604 (its transcript in Chapter 1, from runs/2026-09-01-linear-verified-xla-cuda/), which is the cost of six parameter tensors and two ReLU subgradients rather than one dense layer.

Results

The run in §2.1 is the result. Figure 2-3 is its per-epoch training loss, on a log scale.

The loss falls \(14\times \) across the run, and test accuracy ends at 97.81%.

Both numbers come off the same binary this chapter’s theorems describe.

Return on width

The \(512{\to }512\) hidden size is a convention, not a measurement. The verified renderer is parametric in the layer dimensions and mlpHasVJP is polymorphic in \(d_0,d_1,d_2,d_3\), so one theorem covers every width and we can sweep the hidden size, training each point on its own proof-rendered StableHLO. The mnist-mlp-grid driver does exactly this: it renders \(784{\to }d{\to }d{\to }10\) from the faithful emitter and trains it end to end. Sweeping \(d\) over the powers of two from 8 to 4096 (each 12-epoch SGD, on the GPU) gives Figure 2-4, the accuracy-vs-width curve.

The bars are the same interval the transcripts print, and about \(\pm 0.3\) points once accuracy is in the nineties.

The curve flattens hard after \(\sim \)64 neurons: widening \(8{\to }64\) buys \(+5.0\) points, but \(64{\to }4096\), at \(64\times \) the width and \(364\times \) the parameters, buys only \(+1.0\). A \(32{\to }32\) MLP already reaches 96.8% at 27K parameters, \(755\times \) smaller than the \(4096{\to }4096\) net (98.0%).

The bars are what make that plateau concrete rather than merely visible. Every point from \(d = 256\) rightward has an interval overlapping every other, so a \(16\times \) increase in width and \(74\times \) in parameters buys nothing this test set can resolve. The climb up to the knee is real — \(d = 64\) and \(d = 4096\) do not overlap — and the plateau past it is not. The canonical \(512{\to }512\) sits comfortably on that plateau, and anything past it is paying for the third decimal place.

Off the diagonal the story is capacity rather than failure. A \(784{\to }8{\to }4096{\to }10\) net averages 93.70% over five seeds, with a spread of \(0.56\) — about a point and a half above what the \(8{\to }8\) net manages on its own (92.0%). The 8-unit layer is a bottleneck, and width behind it cannot recover signal that the narrow layer already discarded.

That point is also the one place in this sweep where three runs would have been misleading, which is worth seeing once. The first three seeds land between 93.14 and 93.70%, a tight and plausible-looking range; the fourth comes in at 94.53%, most of a point above the top of it. Three observations of a quantity with this much spread will happily agree with each other and still be wrong about where the next one falls. The diagonal points plotted above are one run each, and are best read with half a point of give.

2.4 MLIR: Dense

What is already proven. mlpForward is the three-dense-layer network, and mlpHasVJPAt proves its reverse-mode derivative (the vector–Jacobian product) by chaining the per-layer chain rule vjpCompAt (Theorem 29). The parameter gradients are pinned by denseWeightGrad_correct and denseBiasGrad_correct, and the loss gradient by softmaxCE_grad. So what the gradient is, is not in question.

The gap and how we close it. The emitted backward is a value of a Lean datatype with a denotation \([\! [\cdot ]\! ]\) into the proofs’ own vector type, and a theorem equates that denotation with the proven derivative. The cotangent chain of the MLP’s committed step, verified_mlir/mlp_train_step.mlir, is the graph

rooted at the loss cotangent, with \([\! [\texttt{dotGeneral}\, W]\! ] = \texttt{Mat.mulVec}\, W\) and \([\! [\texttt{selectPos}\, p]\! ] = (v \mapsto \text{if } p{\gt}0 \text{ then } v \text{ else } 0)\). The per-example bridge

says that graph denotes the proven derivative, and Theorem 30 says the same of the batched graph the file is printed from. Here is the chain in the artifact, batch 128, \(784{\to }512{\to }512{\to }10\), followed by the two weight gradients that consume it (SSA names the artifact’s, default precision dropped, lines wrapped):

%v18 = stablehlo.subtract %v17, %onehot : tensor<128x10xf32>
%v19 = stablehlo.dot_general %v18, %W2, contracting_dims = [1] x [1]
      : (tensor<128x10xf32>, tensor<512x10xf32>) -> tensor<128x512xf32>
%v20 = stablehlo.constant dense<0.0> : tensor<128x512xf32>
%v21 = stablehlo.compare GT, %v7, %v20 : (tensor<128x512xf32>,
      tensor<128x512xf32>) -> tensor<128x512xi1>
%v22 = stablehlo.select %v21, %v19, %v20 : tensor<128x512xi1>,
      tensor<128x512xf32>
%v23 = stablehlo.dot_general %v22, %W1, contracting_dims = [1] x [1]
      : (tensor<128x512xf32>, tensor<512x512xf32>) -> tensor<128x512xf32>
%v24 = stablehlo.constant dense<0.0> : tensor<128x512xf32>
%v25 = stablehlo.compare GT, %v2, %v24 : (tensor<128x512xf32>,
      tensor<128x512xf32>) -> tensor<128x512xi1>
%v26 = stablehlo.select %v25, %v23, %v24 : tensor<128x512xi1>,
      tensor<128x512xf32>
// ...
%v36 = stablehlo.dot_general %v4, %v22, contracting_dims = [0] x [0]
      : (tensor<128x512xf32>, tensor<128x512xf32>) -> tensor<512x512xf32>
// ...
%v45 = stablehlo.dot_general %x, %v26, contracting_dims = [0] x [0]
      : (tensor<128x784xf32>, tensor<128x512xf32>) -> tensor<784x512xf32>

Read it against the graph. %v18 is the loss cotangent, softmax minus one-hot. Each dot_general with contracting_dims = [1] x [1] (%v19, %v23) is a dotGeneral \(W\) node: it contracts the output axis of the cotangent against the output axis of the weight, which is multiplication by \(W^{\! \top }\). Each compare GT + select pair (%v21/%v22, %v25/%v26) is a selectPos \(p\) node, the ReLU subgradient, and the tensor it compares is the saved pre-activation from the forward half of the same file: %v7 is the second layer’s \(Wx+b\) and %v2 the first’s. The two dot_generals with contracting_dims = [0] x [0] contract the batch axis instead: %v36 is \(a_1^{\! \top }\, dy_2\) and %v45 is \(x^{\! \top }\, dy_1\), the weight gradients of Theorem 20 summed over the batch. Theorem 30 states that each of the six returned updates is \(\theta \) minus \(\alpha /N\) times the certified gradient at the cotangent these very nodes carry: the slice lemmas it rests on say that example \(k\)’s slice of %v22 and %v26 is the one-example chain above, so the bridge and the batched tie are one claim at two widths.

Caveats.

• ReLU is a smooth-point bridge. Where a pre-activation is exactly zero ReLU has no derivative: the pdiv-derived backward is the zero vector there, while the emitted compare GT 0 still passes \(dy\) on the positive coordinates, so the bridge is stated only off that set. For generic weights the set is null.

• Trusted surface. The printer’s per-token text and the lowerer’s translation are tested rather than proved. Rounding \(\approx \mathbb {R}\) is not on this list: it is a theorem layer over an arbitrary rounding model (§ B.5), conditional on two measured constants rather than assumed outright. Appendix B.3.1 has the full accounting.

2.5 Recap

One new operation, and one new habit. The operation is ReLU, whose derivative does not exist at zero; its \(\mathsf{HasVJP}\) is built directly from \(\operatorname {pdiv}\), which is the mask \([x {\gt} 0]\) off the corner and \(\operatorname {fderiv}\)’s zero map on it, and its correct field holds by definition. The habit is applying the chain rule more than once: the whole MLP’s backward is the dense backward and the ReLU backward strung together in reverse, and Theorem 30 says the six updates the trainer emits are exactly that, summed over the batch at every batch size, with Theorem 32 the descent on the first layer at one example. Two witnesses, one hand-built and one at the trained weights, show the gradient is not the zero function.

The price of the corner is a hypothesis. The bridge from the proof to the emitted compare GT 0 holds where no pre-activation is exactly zero, a null set for generic weights, and the whole-network statement names the input it holds at. That hypothesis rides along until the activations turn smooth from Chapter 7 on, where there is none to carry. The limit of this chapter’s network is what it throws away: the image was flattened to a vector, so two pixels that were neighbours on the page are unrelated coordinates to a dense layer. The next chapter keeps the structure, and pays for it with a second convention, one whose excluded set is not null.