2 MNIST: 1D MLP
The pdiv we built last chapter, now applied
In Chapter 1 we defined a function \(\operatorname {pdiv}\) that captures the partial derivative of any sufficiently smooth function. We proved that three structural rules, chain and sum and product, suffice to compose new partials out of old ones. That was machinery without a target. This chapter picks the target: we’re going to compute the partial derivative of every component of a small image classifier, end-to-end, and the goal is for that classifier to come out the other side able to recognize handwritten digits.
To say the same thing more concretely: a neural network is a long chain of functions, and “training” means using the partials of those functions to nudge their parameters in a direction that reduces a loss. Two different motions, and this book keeps two words for them: you nudge a parameter to move the loss, and you jiggle an input to find out what would move it. Every theorem in this chapter is the answer to the same question, asked about a different building block: if I jiggle this input by \(\varepsilon \), how does the output jiggle?
2.1 Run it first
Before any of the math, train the thing. Four commands and about ten seconds of GPU time:
lake exe cache get # Mathlib oleans, ~30 s ./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-08-11-mlp-verified-xla-cuda/. XLA’s startup banner is removed and two long lines are wrapped:
[pjrt_ffi] XLA backend: PJRT 0.112, 1 device(s)
[pjrt_ffi] compiled verified_mlir/mlp_train_step.mlir
(@mlp_train_step, 7 outputs, 1 replica) in 2103 ms
[pjrt_ffi] compiled verified_mlir/mlp_fwd.mlir
(@mlp_fwd, 1 outputs, 1 replica) in 118 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.413492, test_acc = 9199/10000 = 91.990000% (849ms)
epoch 2: loss = 0.191197, test_acc = 9525/10000 = 95.250000% (901ms)
epoch 3: loss = 0.139092, test_acc = 9635/10000 = 96.350000% (902ms)
epoch 4: loss = 0.109095, test_acc = 9690/10000 = 96.900000% (938ms)
epoch 5: loss = 0.088633, test_acc = 9714/10000 = 97.140000% (903ms)
epoch 6: loss = 0.073492, test_acc = 9733/10000 = 97.330000% (904ms)
epoch 7: loss = 0.061921, test_acc = 9749/10000 = 97.490000% (931ms)
epoch 8: loss = 0.052623, test_acc = 9765/10000 = 97.650000% (943ms)
epoch 9: loss = 0.045073, test_acc = 9764/10000 = 97.640000% (903ms)
epoch 10: loss = 0.038765, test_acc = 9775/10000 = 97.750000% (919ms)
epoch 11: loss = 0.033494, test_acc = 9774/10000 = 97.740000% (946ms)
epoch 12: loss = 0.028984, test_acc = 9783/10000 = 97.830000% (942ms)
done (trained MNIST-MLP via the proof-rendered StableHLO).
Twelve epochs at about 900 ms each, and 97.83% on the full 10,000-image test set. Adding two hidden layers to Chapter 1’s linear model buys 5.7 points of accuracy at about three times the wall clock per epoch.
The same claim, now with two hidden layers under it. That run did not execute a reimplementation of the network this chapter describes. It executed verified_mlir/mlp_train_step.mlir, which is the text Proofs.StableHLO.mlpTrainStepFaithfulV emits. Every line of that file is pretty of a proven graph node, with one marked exception, and MlpFaithfulPoC proves that each of the six parameter outputs denotes the certified gradient-descent step derived from the Mathlib \(\operatorname {fderiv}\) math. The exception is the trailing %loss scalar, which the module labels REPORT-ONLY in a comment of its own. It exists so the trainer can print the loss you see above, it is read by nothing else, and no theorem depends on it. The two hidden layers, their ReLUs, the softmax cross-entropy cotangent, the six parameter gradients, and the SGD update are all proof-backed operations.
The two compile lines are worth a second look. XLA took the StableHLO and compiled it in process in about two and a half seconds, with no separate compiler invocation and nothing left on disk. The larger of the two numbers belongs to the training step, which is the graph carrying all six parameter gradients.
The Jacobian as multidimensional “how does it jiggle?”
For a scalar function \(f : \mathbb {R} \to \mathbb {R}\) the answer is a single number: the derivative. For our networks, no function is scalar-in scalar-out. The smallest layer maps an input vector to an output vector, so \(n\) numbers in and \(m\) numbers out. “Jiggle the input” now means: pick a direction in \(\mathbb {R}^n\) and push a small distance along it. The textbook word for that is a perturbation, and it is the word Appendix C switches to once the size of the jiggle has to be bounded rather than imagined. Same motion, and the informal name is kept here only because nothing yet depends on how big it is. “The output jiggles” now means: a vector in \(\mathbb {R}^m\). The full picture of how every output coordinate responds to every input coordinate is an \(m \times n\) table of numbers, one slope per (output, input) pair. That table is the Jacobian. There is nothing more to it. It is just the multidimensional generalization of “the derivative is the slope.”
The previous chapter’s \(\operatorname {pdiv}\) is the one entry of this table at a chosen \((i, j)\). The Jacobian is \(\operatorname {pdiv}\) applied across both indices at once, organized so it can be multiplied into other matrices later.
The dense layer’s Jacobian is the weight matrix
Our smallest building block is a dense (or “fully connected”) layer: \(y = Wx + b\), where \(W\) is an \(m \times n\) weight matrix, \(x\) is an \(n\)-vector, and \(b\) is an \(m\)-vector. Pick any output coordinate \(y_j\) and write it out:
Now jiggle \(x_i\) by \(\varepsilon \). Of the \(n\) terms in the sum, only the \(k = i\) term notices, and it changes by \(W_{ji}\, \varepsilon \). So \(\partial y_j / \partial x_i = W_{ji}\). The Jacobian of \(y = Wx + b\) with respect to \(x\) is the weight matrix \(W\) itself. No new structure and no surprise, because the Jacobian of a linear map is the linear map, written as a matrix.
We will do the same exercise for the partial with respect to \(W\), and we will get an analogous answer: the dependence is local, the entries of the Jacobian are just \(x_{i'}\, \delta _{jj'}\). These two, the input-Jacobian and the weight-Jacobian, are the only objects we need for a dense layer. Theorems 14 and 15 below formalize them. We have already done the substantive work.
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 16 states the Jacobian at smooth points, and the codegen substitutes the standard subgradient convention at the kink and we move on. The kink is the only place in the chapter where “differentiable” becomes slightly subtle.
Softmax cross-entropy: where vectors collapse to a scalar
The last building block is the loss: softmax cross-entropy between the model’s output \(z \in \mathbb {R}^{10}\) and the true class label \(y \in \{ 0, \ldots , 9\} \). The full computation is
The loss is a scalar, so its Jacobian with respect to \(z\) is a vector, not a matrix. Working through the algebra (chain rule on \(-\log \circ \mathrm{softmax}\) at the label index) produces a remarkably clean answer:
The gradient of the loss is the difference between the model’s predicted distribution and the truth. Theorem 17 makes this formal.
From Jacobians to VJPs
Training does not multiply Jacobians together directly. The loss is a scalar, and the quantity we actually want is “how does each parameter affect that scalar?” That quantity is the Jacobian of the loss with respect to the parameters transposed and applied to the upstream gradient, which is a vector-Jacobian product, or VJP. Concretely: the upstream gradient is a vector \(dy\), and we want to compute the corresponding \(dx\) and \(dW\). For a dense layer the answers fall out by transposing the picture we already have:
The remaining theorems in this chapter (18 through 22) are the formalizations of these identities, plus the proof that VJPs of composed layers are themselves VJPs obtained by composing the building blocks in reverse order. That last claim is what makes a 3-layer MLP’s backward pass exactly three transposed matrix multiplies, and it is the structural fact that the rest of the book leans on.
2.2 The theorems
For the dense layer \(\mathrm{dense}(W, b)\, x = \lambda j.\; \bigl(\sum _i x_i\, W_{ij}\bigr) + b_j\):
No hypotheses: dense is affine, so the differentiability obligations are discharged inside the proof rather than assumed.
Sketch: factor the layer as (finite sum of bilinear summands) + constant, distribute \(\operatorname {pdiv}\), apply the product rule per summand, collapse the Kronecker \(\delta \). Every foundation rule from Chapter 1 except the chain rule fires exactly once.
\(\operatorname {pdiv}\bigl(\mathrm{dense}(W, b)\bigr)\, x\, i\, j = \operatorname {pdiv}\bigl(\lambda y\, j'.\, \textstyle \sum _{i'} y_{i'} W_{i'j'}\bigr)\, x\, i\, j\).
proof: Split the layer as \(\bigl(\sum _{i'} \cdots \bigr) + (\text{const } b)\) and apply the sum rule (Theorem 3) and the constant rule (Theorem 6). Their differentiability hypotheses hold because each summand \(y \mapsto y_{i'} \cdot W_{i'j'}\) is (reindex) \(\times \) (constant) — differentiable since reindexing is a continuous linear map — and the finite sum inherits differentiability (DifferentiableAt.fun_sum).\(\mathord {@}= \sum _{i'} \operatorname {pdiv}\bigl(\lambda y\, j'.\, y_{i'} \cdot W_{i'j'}\bigr)\, x\, i\, j\).
proof: Finite-sum rule (Theorem 8), with the same per-summand differentiability.For each \(i'\): \(\operatorname {pdiv}\bigl(\lambda y\, j'.\, y_{i'} \cdot W_{i'j'}\bigr)\, x\, i\, j = \delta _{i i'} \cdot W_{i'j}\).
proof: Product rule (Theorem 4) on (reindex) \(\times \) (constant); the reindex Jacobian (Theorem 7) with \(\sigma = \lambda \_ .\, i'\) contributes the \(\delta \), the constant factor’s Jacobian vanishes (Theorem 6); case split on \(i = i'\).q.e.d.
proof: Substitute 3 into 2: \(\sum _{i'} \delta _{i i'} W_{i'j} = W_{ij}\) (Finset.sum_ite_eq); with 1 this is the goal.
The symmetric counterpart of Theorem 14, differentiating in \(W\) instead of \(x\). Since \(\operatorname {pdiv}\) works on vectors, view the layer as a function of the flattened weights: for \(v \in \mathbb {R}^{m \cdot n}\), let \(F(v) := \mathrm{dense}(\mathrm{unflatten}\, v,\, b)\, x\), and write \(\varphi \) for the index bijection \((i, j) \leftrightarrow \varphi (i, j)\) (finProdFinEquiv). prove: for all \(i, j', j\):
Sketch: same skeleton as Theorem 14 — split, distribute, product rule, collapse — with the reindex step now going through the flatten bijection.
\(F(v)_{j_o} = \bigl(\sum _{i'} x_{i'} \cdot v_{\varphi (i', j_o)}\bigr) + b_{j_o}\).
proof: Unfold \(\mathrm{dense}\) and \(\mathrm{unflatten}\).\(\operatorname {pdiv}F\, (\mathrm{flatten}\, W)\, \varphi (i, j')\, j = \operatorname {pdiv}\bigl(\lambda w\, j_o.\, \textstyle \sum _{i'} x_{i'} \cdot w_{\varphi (i', j_o)}\bigr)\, (\mathrm{flatten}\, W)\, \varphi (i, j')\, j\).
proof: Drop the constant bias: sum rule (Theorem 3) and constant rule (Theorem 6); each summand is (constant) \(\times \) (reindex), differentiable, and the finite sum inherits differentiability (DifferentiableAt.fun_sum).\(\mathord {@}= \sum _{i'} \operatorname {pdiv}\bigl(\lambda w\, j_o.\, x_{i'} \cdot w_{\varphi (i', j_o)}\bigr)\, (\mathrm{flatten}\, W)\, \varphi (i, j')\, j\).
proof: Finite-sum rule (Theorem 8).For each \(i'\): the summand equals \(\text{if } i = i' \wedge j' = j \text{ then } x_i \text{ else } 0\).
proof: Product rule (Theorem 4) on (constant \(x_{i'}\)) \(\times \) (reindex \(\sigma = \lambda j_o.\, \varphi (i', j_o)\)); the constant factor’s Jacobian vanishes (Theorem 6), and the reindex Jacobian (Theorem 7) is \(1\) exactly when \(\varphi (i, j') = \varphi (i', j)\), which by injectivity of \(\varphi \) is \(i = i' \wedge j' = j\).q.e.d.
proof: Sum 4 over \(i'\): if \(j' = j\) the Kronecker condition picks the single term \(x_i\) (Finset.sum_ite_eq); if \(j' \neq j\) every term is \(0\). Both match \(\delta _{jj'}\, x_i\); with 2 and 3 this is the goal.
assume:
\(x\) is a smooth point of \(\mathrm{ReLU}\): every coordinate \(x_k \neq 0\) [h_smooth]
prove: for all \(i, j\):
where \([P]\) is the Iverson bracket (\(1\) if \(P\) holds, else \(0\)).
Sketch: near a smooth point, ReLU is a fixed linear map (each coordinate is committed to its branch of the \(\max \)); compute that map’s derivative and transport it.
Define \(\Lambda _x := \Pi _k\, (\text{if } x_k {\gt} 0 \text{ then } \mathrm{proj}_k \text{ else } 0)\), the diagonal indicator CLM (reluLinearPart).
Let \(r := \min _k |x_k|\). Then \(r {\gt} 0\), and \(\mathrm{ReLU}\) agrees with \(\Lambda _x\) on the ball \(B(x, r)\).
proof: \(r {\gt} 0\) by assumption 1. For \(y \in B(x, r)\) and every \(k\): \(|y_k - x_k| \le \lVert y - x \rVert {\lt} r \le |x_k|\), so \(y_k\) has the sign of \(x_k\); both functions then return \(y_k\) where \(x_k {\gt} 0\) and \(0\) where \(x_k {\lt} 0\).\(\mathrm{ReLU}\) has Fréchet derivative \(\Lambda _x\) at \(x\).
proof: A continuous linear map is its own derivative; by 2 the two functions agree on a neighborhood of \(x\), and HasFDerivAt.congr_of_eventuallyEq transports the derivative across that agreement.q.e.d.
proof: By Definition 1 and 3, \(\operatorname {pdiv}(\mathrm{ReLU})\, x\, i\, j = \Lambda _x(\mathbf{e}_i)_j = [x_j {\gt} 0] \cdot (\mathbf{e}_i)_j\); case split on \(i = j\) (when they coincide, \([x_j {\gt} 0] = [x_i {\gt} 0]\)) gives the goal.
Write \(p := \mathrm{softmax}(z)\) and \(\mathrm{CE}(z, \ell ) := -\log p_\ell \) (viewed as \(\mathbb {R}^{1}\)-valued so \(\operatorname {pdiv}\) applies; we read off the only output coordinate). prove: for all \(j\):
Sketch: chain rule on \(-\log \circ (z \mapsto p_\ell )\), with the softmax Jacobian (proved in Ch 9) supplying the inner derivative; the \(1/p_\ell \) from \(\log \) cancels the \(p_\ell \) the Jacobian produces.
\(p_\ell {\gt} 0\), in particular \(p_\ell \neq 0\).
proof: \(p_\ell \) is a positive exponential over a positive finite sum of exponentials.\(\operatorname {pdiv}\bigl(\mathrm{CE}(\cdot , \ell )\bigr)\, z\, j\, 0 = \operatorname {fderiv}_{\mathbb {R}}\, \bigl(z' \mapsto \mathrm{CE}(z', \ell )\bigr)\, z\, (\mathbf{e}_j)\).
proof: Definition 1; the \(\mathbb {R}^{1}\) wrapper just evaluates the single output coordinate (fderiv_apply), legitimate because the wrapper is differentiable — \(\mathrm{softmax}\) is differentiable and \(\log \) is differentiable away from \(0\), which 1 grants.\(\operatorname {fderiv}_{\mathbb {R}}\, \bigl(z' \mapsto \mathrm{CE}(z', \ell )\bigr)\, z = -\bigl(p_\ell ^{-1} \cdot \operatorname {fderiv}_{\mathbb {R}}\, (z' \mapsto \mathrm{softmax}(z')_\ell )\, z\bigr)\).
proof: \(\mathrm{CE}(\cdot , \ell ) = -\log \circ (z' \mapsto \mathrm{softmax}(z')_\ell )\); HasFDerivAt.log with 1 differentiates the \(\log \), then negate.\(\operatorname {fderiv}_{\mathbb {R}}\, (z' \mapsto \mathrm{softmax}(z')_\ell )\, z\, (\mathbf{e}_j) = \operatorname {pdiv}(\mathrm{softmax})\, z\, j\, \ell = p_\ell \, (\delta _{j\ell } - p_j)\).
proof: Definition 1, then the softmax Jacobian (Theorem 65).q.e.d.
proof: Chain 2–4: \(-p_\ell ^{-1} \cdot p_\ell \, (\delta _{j\ell } - p_j) = p_j - \delta _{j\ell }\), cancelling by 1; and \(\mathrm{onehot}(\ell )_j = \delta _{j\ell }\) by definition.
\(\mathsf{HasVJP}\, \bigl(\mathrm{dense}(W, b)\bigr)\): the input-gradient backward of a dense layer is multiplication by \(W\).
Define \(B(x, dy) := W\, dy\), i.e. \(B(x, dy)_i = \sum _j W_{ij}\, dy_j\).
suffices: for all \(x\), \(dy\), \(i\): \(B(x, dy)_i = \sum _j \operatorname {pdiv}\bigl(\mathrm{dense}(W, b)\bigr)\, x\, i\, j \cdot dy_j\).
proof: Definition 9, with \(B\) as the candidate backward function.q.e.d.
proof: The Dense Jacobian (Theorem 14) gives \(\operatorname {pdiv}\bigl(\mathrm{dense}(W, b)\bigr)\, x\, i\, j = W_{ij}\); substituting into 2 leaves exactly \(\sum _j W_{ij}\, dy_j = B(x, dy)_i\).
\(dW = x \otimes dy\), with \(F\) and \(\varphi \) as in Theorem 15, prove: for all \(i, j\):
Each summand: \(\operatorname {pdiv}F\, (\mathrm{flatten}\, W)\, \varphi (i, j)\, k \cdot dy_k = (\text{if } k = j \text{ then } x_i \text{ else } 0) \cdot dy_k\).
proof: Dense Jacobian wrt weight (Theorem 15).q.e.d.
proof: The sum in 1 collapses at \(k = j\) (Finset.sum_eq_single) to \(x_i \cdot dy_j\), which is \((x \otimes dy)_{ij}\) by definition of the outer product.
\(db = dy\). prove: for all \(i\):
\(\operatorname {pdiv}\bigl(b' \mapsto \mathrm{dense}(W, b')\, x\bigr)\, b\, i\, j = \delta _{ij}\).
proof: As a function of \(b'\), the layer is \((\text{constant in } b') + b'\): sum rule (Theorem 3), constant rule (Theorem 6), and identity Jacobian (Theorem 5). (This is the Lean lemma pdiv_dense_b.)q.e.d.
proof: Substitute 1: \(\sum _j \delta _{ij}\, dy_j = dy_i\) (Finset.sum_eq_single).
noncomputable def over the canonical pdiv-derived witness; HasVJP.correct holds by rfl since \(\operatorname {pdiv}\) is a def over \(\operatorname {fderiv}\). At non-smooth points the canonical backward is \(\operatorname {fderiv}\)’s junk default of \(0\); the codegen substitutes the standard subgradient convention.
noncomputable def over the canonical pdiv-derived witness; same shape as relu_has_vjp. Codegen routes the ReLU subgradient at the kink.
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
![\begin{tikzpicture} [
>={Stealth[length=1.8mm]},
every node/.style={font=\sffamily\scriptsize},
col/.style = {align=center, rounded corners=2pt, inner sep=2pt, minimum height=0.62cm, minimum width=4.0cm},
io/.style = {col, draw=blue!55!black, fill=blue!8},
dense/.style = {col, draw=orange!65!black, fill=orange!12},
head/.style = {col, draw=red!60!black, fill=red!8},
logits/.style = {col, draw=green!50!black, fill=green!14, very thick},
arr/.style = {->, thick, gray!60, shorten >=1pt, shorten <=1pt},
]
% Vertical layer column: input at top -> logits at bottom.
\node[io] (input) {Input \;\; $28\times28$ flatten $\to 784$};
\node[dense, below=0.24cm of input] (d1) {\textbf{Dense} $784\to512$, ReLU};
\node[dense, below=0.24cm of d1] (d2) {\textbf{Dense} $512\to512$, ReLU};
\node[head, below=0.24cm of d2] (d3) {\textbf{Dense} $512\to10$ \;(identity)};
\node[logits, below=0.24cm of d3] (out) {Logits \;\; 10 classes, softmax-CE};
\foreach \a/\b in {input/d1, d1/d2, d2/d3, d3/out}
\draw[arr] (\a) -- (\b);
\end{tikzpicture}](images/img-0003.svg)
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
nClasses := 10
data := .mnist
layers := [.dense 784 512, .relu,
.dense 512 512, .relu,
.dense 512 10]
Five entries where Chapter 1 had one, and nothing else about the type changes. The slug binds the spec to verified_mlir/mlp_train_step.mlir and verified_mlir/mlp_fwd.mlir, and Proofs.MlpRender.mlpTrainStepFaithfulV is what emits the train step.
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. In LeanMlir/Proofs/Foundation/SpecVJP.lean:
noncomputable def mlpVerified_has_vjp (W0 b0 W1 b1 W2 b2) :
HasVJP (denoteMLP mlpVerified.layers W0 b0 W1 b1 W2 b2) := ...
Again the VJP 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.VerifiedNets def mlpConfig : VerifiedConfig where epochs := 12 batchSize := 128 def main (argv : List String) : IO Unit := mlpVerified.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.
XLA compiles that train step in about 5.4 seconds and the forward in about 40 milliseconds. The train step takes roughly seven times as long to compile as Chapter 1’s, 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. Its per-epoch training loss, on a log scale:
![\begin{tikzpicture}
\begin{axis}[
width=0.92\linewidth, height=6.5cm,
xlabel={Epoch}, ylabel={Training loss},
ymode=log,
xmin=0, xmax=13, ymin=0.02, ymax=0.5,
xtick={0,2,4,6,8,10,12},
grid=major, grid style={gray!18},
tick label style={font=\small},
label style={font=\small},
every axis plot/.append style={line width=1pt, mark size=1pt},
]
\addplot[blue, mark=*, mark options={fill=blue}] coordinates {
(1,0.413492) (2,0.191197) (3,0.139092) (4,0.109095) (5,0.088633) (6,0.073492) (7,0.061921) (8,0.052623) (9,0.045073) (10,0.038765) (11,0.033494) (12,0.028984)
};
\end{axis}
\end{tikzpicture}](images/img-0004.svg)
MNIST MLP (\(784{\to }512{\to }512{\to }10\)), SGD 0.1, 12 epochs, log-scale mean training loss per epoch (runs/2026-08-11-mlp-verified-xla-cuda/). The loss falls \(14\times \) across the run, and test accuracy ends at 97.83%.
Both numbers come off the same binary this chapter’s theorems describe. An earlier edition reported 98.57% here, measured with an unverified ablation runner that shared the architecture but not the proof-rendered graph, and that runner was kept because it was the only thing that printed a loss. The verified train step now returns one, so the second program is gone.
2.4 Return on width
The \(512{\to }512\) hidden size is a convention, not a measurement. Because the verified renderer is parametric in the layer dimensions, so the same mlp_has_vjp theorem covers every width, being polymorphic in \(d_0,d_1,d_2,d_3\) carry through, we can sweep the hidden width and train 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 the accuracy-vs-width curve below.
![\begin{tikzpicture}
\begin{axis}[
width=0.92\linewidth, height=6.5cm,
xlabel={Hidden width $d$ (both hidden layers)},
ylabel={MNIST test accuracy (\%)},
xmode=log, log basis x=2,
xmin=6.5, xmax=5000, ymin=90.5, ymax=98.7,
xtick={8,16,32,64,128,256,512,1024,2048,4096},
xticklabels={8,16,32,64,128,256,512,1024,2048,4096},
ytick={91,92,93,94,95,96,97,98},
grid=major, grid style={gray!18},
tick label style={font=\small},
label style={font=\small},
every axis plot/.append style={line width=1pt, mark size=1.5pt},
]
\addplot[blue, mark=*, mark options={fill=blue}] coordinates {
(8,91.99) (16,94.01) (32,96.76) (64,97.02) (128,97.25) (256,97.67) (512,97.83) (1024,97.89) (2048,98.00) (4096,98.01)
};
\addplot[only marks, mark=o, mark size=4pt, red, line width=1pt, forget plot] coordinates {(64,97.02)};
\node[anchor=west, font=\footnotesize, red!70!black] at (axis cs:74,95.5)
{knee $\approx 64$: 97.0\% at 55K params};
\end{axis}
\end{tikzpicture}](images/img-0005.svg)
Return on width for the verified MNIST MLP (\(784{\to }d{\to }d{\to }10\)): test accuracy versus hidden width \(d\) (log scale), 12-epoch SGD on the proof-rendered StableHLO (runs/mlp_grid_results_xla.tsv). 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 canonical \(512{\to }512\) sits comfortably on the 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 reaches 93.1%, and at three initialization seeds it lands between 93.1% and 93.7%. That is only about a point 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.
2.5 MLIR: Dense
What is already proven. mlpForward is the three-dense-layer network, and mlp_has_vjp_at proves its reverse-mode derivative (the vector–Jacobian product) by chaining the per-layer chain rule vjp_comp_at. The parameter gradients are pinned by dense_weight_grad_correct and dense_bias_grad_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, Back, with a denotation \([\! [\cdot ]\! ]\) into the proofs’ own vector type. For the MLP the graph is
rooted at the incoming 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)\), and the bridge theorem
says the denotation of that graph is the proven derivative. The forward, the loss cotangent, and the parameter gradients are covered the same way. The printer walks the graph and emits one stablehlo op per node:
Here is exactly what the printer emits for that graph, with nothing hand-written, one stablehlo op per IR node (default precision attributes elided for the page):
func.func @mlp_back(%dy: tensor<2x2xf32>, %W0: tensor<4x3xf32>,
%W1: tensor<3x3xf32>, %W2: tensor<3x2xf32>,
%p0: tensor<2x3xf32>, %p1: tensor<2x3xf32>) -> tensor<2x4xf32> {
%bk0 = stablehlo.dot_general %dy, %W2, contracting_dims = [1] x [1]
: (tensor<2x2xf32>, tensor<3x2xf32>) -> tensor<2x3xf32>
%bk1 = stablehlo.constant dense<0.0> : tensor<2x3xf32>
%bk2 = stablehlo.compare GT, %p1, %bk1
: (tensor<2x3xf32>, tensor<2x3xf32>) -> tensor<2x3xi1>
%bk3 = stablehlo.select %bk2, %bk0, %bk1
: tensor<2x3xi1>, tensor<2x3xf32>
%bk4 = stablehlo.dot_general %bk3, %W1, contracting_dims = [1] x [1]
: (tensor<2x3xf32>, tensor<3x3xf32>) -> tensor<2x3xf32>
%bk5 = stablehlo.constant dense<0.0> : tensor<2x3xf32>
%bk6 = stablehlo.compare GT, %p0, %bk5
: (tensor<2x3xf32>, tensor<2x3xf32>) -> tensor<2x3xi1>
%bk7 = stablehlo.select %bk6, %bk4, %bk5
: tensor<2x3xi1>, tensor<2x3xf32>
%bk8 = stablehlo.dot_general %bk7, %W0, contracting_dims = [1] x [1]
: (tensor<2x3xf32>, tensor<4x3xf32>) -> tensor<2x4xf32>
return %bk8 : tensor<2x4xf32>
}
Read it against the graph: each dot_general (%bk0, %bk4, %bk8) is a dotGeneral \(W\) node, whose denotation is \(\texttt{Mat.mulVec}\, W\). Each compare GT + select pair (%bk2/%bk3, %bk6/%bk7) is a selectPos \(p\) node, which is the ReLU subgradient. The bridge theorem is precisely the claim that this text computes mlp_has_vjp_at’s backward.
Caveats.
ReLU is a smooth-point bridge. Where a pre-activation is exactly zero ReLU has no derivative. The emitted compare GT 0 sends that case to \(0\), and the bridge is permitted to fail on that measure-zero set.
Representative scale. Shown small. The trained net is \(784{\to }512{\to }512{\to }10\).
Trusted surface. The printer’s faithfulness and the lowerer’s translation are tested rather than proved. Rounding \(\approx \mathbb {R}\) is no longer on this list: it is a theorem layer over an arbitrary rounding model (§ C.4), conditional on two measured constants rather than assumed outright. Appendix C.2.1 has the full accounting.