Verified Deep Learning with Lean 4

1 MNIST: linear classifier

A VJP (vector-Jacobian product) is the backward function for a layer: takes an upstream gradient, returns the downstream one. Every backward function in this book is a VJP built by composing the foundation here. That foundation is the partial-derivative function \(\operatorname {pdiv}\) (defined via Mathlib’s \(\operatorname {fderiv}\)), its three structural rules of chain, sum, and product, and three VJP record types (\(\mathsf{HasVJP}\), \(\mathsf{HasVJPMat}\), \(\mathsf{HasVJP3}\), one per tensor rank) that bundle a backward function with its correctness claim.

This is the technically hardest chapter in the book, by design. Going from nothing to a complete trained model is the entire machinery, and we build all of it here: the \(\operatorname {pdiv}\) calculus, the VJP framework, a forward pass, a loss, a backward pass, an SGD step, all the way down to the GPU. We do it on the smallest possible network (one matrix multiply: the MNIST linear classifier) to demonstrate the pattern. Once that machinery exists, every later chapter is just adding a layer: one new primitive dropped into the same forward / loss / backward / optimize loop.

1.1 Run it first

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

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

On one RTX 4060 Ti (CUDA 12.9), from runs/2026-09-01-linear-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/linear_train_step.mlir
             (@linear_train_step, 2 outputs, 1 replica) in 604 ms
[pjrt_ffi] compiled verified_mlir/linear_fwd.mlir
             (@linear_fwd, 1 outputs, 1 replica) in 116 ms
MNIST-Linear via the VERIFIED renderer (pretty∘emit) → XLA/PJRT → GPU
  xla/pjrt verified_mlir/linear_train_step.mlir
  xla/pjrt verified_mlir/linear_fwd.mlir
  train 60000, test 10000; dense 784->10, bs 128, SGD
  epoch 1: test_acc = 8977/10000 = 89.770000%  [95% CI 89.16–90.35] (173ms)
  epoch 2: test_acc = 9081/10000 = 90.810000%  [95% CI 90.23–91.36] (175ms)
  epoch 3: test_acc = 9123/10000 = 91.230000%  [95% CI 90.66–91.77] (217ms)
  epoch 4: test_acc = 9157/10000 = 91.570000%  [95% CI 91.01–92.10] (242ms)
  epoch 5: test_acc = 9167/10000 = 91.670000%  [95% CI 91.11–92.20] (219ms)
  epoch 6: test_acc = 9176/10000 = 91.760000%  [95% CI 91.20–92.28] (239ms)
  epoch 7: test_acc = 9186/10000 = 91.860000%  [95% CI 91.31–92.38] (239ms)
  epoch 8: test_acc = 9194/10000 = 91.940000%  [95% CI 91.39–92.46] (239ms)
  epoch 9: test_acc = 9199/10000 = 91.990000%  [95% CI 91.44–92.51] (243ms)
  epoch 10: test_acc = 9198/10000 = 91.980000%  [95% CI 91.43–92.50] (246ms)
  epoch 11: test_acc = 9203/10000 = 92.030000%  [95% CI 91.48–92.54] (247ms)
  epoch 12: test_acc = 9210/10000 = 92.100000%  [95% CI 91.56–92.61] (205ms)
done (trained MNIST-Linear via the proof-rendered StableHLO).

Twelve epochs at about 224 ms each, and 92.10% on the full 10,000-image test set. That is the going rate for a linear classifier on MNIST, and LeCun’s original 1998 benchmark put a comparable model at around 92%.

The bracketed interval is worth reading once, because every accuracy in this book carries one. 92.10% is 9,210 images out of 10,000, and a different 10,000 handwritten digits would not have given exactly that number; the 95% Wilson interval says the rate underneath it lies plausibly anywhere between 91.56 and 92.61%. The six decimal places the trainer prints are formatting rather than precision — only the first three digits are real, and on a test set this size a gap of less than about half a point is not yet a result.

This chapter’s accuracy is also the one that does not move. Trained five times over, the linear model finishes at 92.10% every time. It is zero-initialized rather than randomly, so there is no seed for the initializer to spend, and the only variation left is the order in which the GPU accumulates its sums — worth a single image at epochs 3 and 9, and nothing at all by epoch 12. Every later chapter has a real spread across seeds and reports it. This one has none to report, which makes it the cleanest baseline in the book to start from.

1.2 How it works

The next twelve theorems all reduce to a single definition (\(\operatorname {pdiv}\)) plus chain rule, sum rule, and product rule. Before we drop you into the deep end, here’s the shape of what’s coming:

\begin{tikzpicture} [ node distance=1.2cm, every node/.style={font=\small, align=center}, group/.style={rectangle, draw, fill=blue!5, rounded corners, inner sep=4pt, minimum width=5cm, minimum height=1.5cm}, root/.style={rectangle, draw, fill=red!10, rounded corners, inner sep=4pt, minimum width=4cm, minimum height=0.8cm}, payoff/.style={rectangle, draw, fill=green!10, rounded corners, inner sep=4pt, minimum width=6.5cm, minimum height=1cm}, arrow/.style={->, >=stealth, thick, gray} ] \node[root] (pdiv) at (0, 5) {\(\pdiv\) \\ \footnotesize defined via Mathlib's \(\fderiv\)}; \node[group] (foundation) at (0, 3) {\textbf{Foundation rules} \\ chain $\cdot$ sum $\cdot$ product \\ identity $\cdot$ const $\cdot$ reindex \\ finite-sum}; \node[group] (vjp) at (0, 0.5) {\textbf{VJP transpose} \\ \texttt{vjp\_comp} $\cdot$ biPath \\ elemwiseProduct $\cdot$ identity}; \node[payoff] (payoff) at (0, -2) {Used by the MNIST linear classifier here \\ (proven + trained); then reused by every \\ architecture in Chapters 2--9.}; \draw[arrow] (pdiv) -- (foundation); \draw[arrow] (foundation) -- (vjp); \draw[arrow] (vjp) -- (payoff); \end{tikzpicture}

Read top-to-bottom, this is the order things get proved. Read bottom-to-top, this is what every theorem in Chapters 3–9 unfolds to. (Ch 9 attention adds matrix-level machinery on top. See §9.2.) The full clickable dependency graph is in the blueprint web view.

How the proofs are written.

The proofs in this book follow Lamport’s structured style (How to Write a 21st Century Proof). Hypotheses are named up front in an assume:/prove: header, with the matching Lean hypothesis name in brackets. The proof is a numbered sequence of steps. Each step carries its own proof naming exactly the facts it follows from, and the final step is always q.e.d., which explains why the steps prove the goal. In equational chains, \(\mathord {@}\) stands for the expression established by the previous step. Each step mirrors one tactic block of the corresponding Lean proof, so the informal argument can be audited against the formal one step by step. One-line lemmas (pdiv_id, pdiv_const, pdiv_reindex, the identity VJP) stay unstructured: structuring a one-simp proof would be ceremony.

1.3 The theorems

Definition 1 Partial derivative
✓
#

Lean declarations

The partial derivative function. For \(f : \mathbb {R}^{m} \to \mathbb {R}^{n}\), \(\operatorname {pdiv}\, f\, x\, i\, j\) is the \((i, j)\)-entry of the Jacobian at \(x\), defined as \(\operatorname {fderiv}_{\mathbb {R}}\, f\, x\, (\mathbf{e}_i)\, j\) — the \(j\)-th coordinate of Mathlib’s Fréchet derivative applied to the \(i\)-th standard basis vector.

Theorem 2 Chain rule
✓
#

Uses

Lean declarations

assume:

  1. \(f : \mathbb {R}^{m} \to \mathbb {R}^{n}\) is differentiable at \(x\) [hf]

  2. \(g : \mathbb {R}^{n} \to \mathbb {R}^{p}\) is differentiable at \(f(x)\) [hg]

prove: for all \(i, k\):

\[ \operatorname {pdiv}(g \circ f)\, x\, i\, k = \sum _j \operatorname {pdiv}f\, x\, i\, j \cdot \operatorname {pdiv}g\, (f\, x)\, j\, k. \]
Proof ▶

Sketch: reduce to Mathlib’s chain rule for \(\operatorname {fderiv}\), then decompose over the standard basis to turn a composite of linear maps into the sum over the middle index.

  1. \(\operatorname {pdiv}(g \circ f)\, x\, i\, k = \operatorname {fderiv}_{\mathbb {R}}\, (g \circ f)\, x\, (\mathbf{e}_i)\, k\).
    proof: Definition 1.

  2. \(\operatorname {fderiv}_{\mathbb {R}}\, (g \circ f)\, x = \operatorname {fderiv}_{\mathbb {R}}\, g\, (f\, x) \circ \operatorname {fderiv}_{\mathbb {R}}\, f\, x\).
    proof: Mathlib’s fderiv_comp, applicable by assumptions 1 and 2.

  3. Let \(v := \operatorname {fderiv}_{\mathbb {R}}\, f\, x\, (\mathbf{e}_i) \in \mathbb {R}^{n}\). Then \(v = \sum _j v_j \cdot \mathbf{e}_j\).
    proof: Pointwise: at coordinate \(j'\) the sum collapses to \(v_{j'}\), because \((\mathbf{e}_j)_{j'}\) is the Kronecker delta.

  4. \(\operatorname {fderiv}_{\mathbb {R}}\, g\, (f\, x)\, v\, k = \sum _j v_j \cdot \operatorname {fderiv}_{\mathbb {R}}\, g\, (f\, x)\, (\mathbf{e}_j)\, k\).
    proof: By 3 and linearity of \(\operatorname {fderiv}_{\mathbb {R}}\, g\, (f\, x)\) (map_sum, map_smul).

  5. q.e.d.
    proof: By Definition 1, \(v_j = \operatorname {pdiv}f\, x\, i\, j\) and \(\operatorname {fderiv}_{\mathbb {R}}\, g\, (f\, x)\, (\mathbf{e}_j)\, k = \operatorname {pdiv}g\, (f\, x)\, j\, k\); chaining 1, 2, and 4 gives the goal.

Theorem 3 Sum rule
✓
#

Uses

Lean declarations

assume:

  1. \(f, g : \mathbb {R}^{m} \to \mathbb {R}^{n}\) are both differentiable at \(x\) [hf, hg]

prove: for all \(i, j\):

\[ \operatorname {pdiv}(f + g)\, x\, i\, j = \operatorname {pdiv}f\, x\, i\, j + \operatorname {pdiv}g\, x\, i\, j. \]
Proof ▶

  1. \(\operatorname {pdiv}(f + g)\, x\, i\, j = \operatorname {fderiv}_{\mathbb {R}}\, (f + g)\, x\, (\mathbf{e}_i)\, j\).
    proof: Definition 1; the pointwise sum \(\lambda y\, k.\; f\, y\, k + g\, y\, k\) is definitionally the function \(f + g\).

  2. \(\operatorname {fderiv}_{\mathbb {R}}\, (f + g)\, x = \operatorname {fderiv}_{\mathbb {R}}\, f\, x + \operatorname {fderiv}_{\mathbb {R}}\, g\, x\).
    proof: Mathlib’s fderiv_add, applicable by the assumption.

  3. q.e.d.
    proof: Evaluate 2 at \((\mathbf{e}_i, j)\); by Definition 1 the two terms are \(\operatorname {pdiv}f\, x\, i\, j\) and \(\operatorname {pdiv}g\, x\, i\, j\).

Theorem 4 Product rule
✓
#

Uses

Lean declarations

assume:

  1. \(f, g : \mathbb {R}^{m} \to \mathbb {R}^{n}\) are both differentiable at \(x\) [hf, hg]

prove: for all \(i, j\):

\[ \operatorname {pdiv}(f \odot g)\, x\, i\, j = \operatorname {pdiv}f\, x\, i\, j \cdot g\, x\, j + f\, x\, j \cdot \operatorname {pdiv}g\, x\, i\, j, \]

where \((f \odot g)\, y\, k = f\, y\, k \cdot g\, y\, k\) is the elementwise product.

Proof ▶

  1. \(\operatorname {pdiv}(f \odot g)\, x\, i\, j = \operatorname {fderiv}_{\mathbb {R}}\, (f \cdot g)\, x\, (\mathbf{e}_i)\, j\).
    proof: Definition 1; the elementwise product is definitionally the algebra product \(f \cdot g\) in \(\mathbb {R}^{n}\) (Pi.normedAlgebra), so Mathlib’s calculus of algebra-valued maps applies.

  2. \(\operatorname {fderiv}_{\mathbb {R}}\, (f \cdot g)\, x = f\, x \cdot \operatorname {fderiv}_{\mathbb {R}}\, g\, x + g\, x \cdot \operatorname {fderiv}_{\mathbb {R}}\, f\, x\) (scalar action taken pointwise).
    proof: Mathlib’s fderiv_mul, applicable by the assumption.

  3. q.e.d.
    proof: Evaluate 2 at \((\mathbf{e}_i, j)\): the pointwise scalar action gives \(f\, x\, j \cdot \operatorname {pdiv}g\, x\, i\, j + g\, x\, j \cdot \operatorname {pdiv}f\, x\, i\, j\) by Definition 1; commute the second summand (ring) to match the goal.

Theorem 5 Identity Jacobian
✓
#

Uses

Lean declarations

\(\operatorname {pdiv}(\mathrm{id})\, x\, i\, j = \delta _{ij}\).

Proof ▶

Mechanical; see Proofs.pdiv_id.

Theorem 6 Constant has zero Jacobian
✓
#

Uses

Lean declarations

Proof ▶

Mechanical; see Proofs.pdiv_const.

Theorem 7 Gather / reindex Jacobian
✓
#

Uses

Lean declarations

Covers permutations, reshapes, slicing. Generalizes pdiv_id.

Proof ▶

Mechanical; see Proofs.pdiv_reindex.

Theorem 8 Finite-sum rule
✓
#

Uses

Lean declarations

assume:

  1. \(S\) is a finite index set with a function \(f_s : \mathbb {R}^{m} \to \mathbb {R}^{n}\) for each \(s \in S\)

  2. every \(f_s\) is differentiable at \(x\) [hdiff]

prove: for all \(i, j\):

\[ \operatorname {pdiv}\Bigl(\sum _{s \in S} f_s\Bigr)\, x\, i\, j = \sum _{s \in S} \operatorname {pdiv}f_s\, x\, i\, j. \]
Proof ▶

Sketch: induction on \(S\); the two-summand sum rule does each inductive step.

  1. case \(S = \emptyset \): both sides are \(0\).
    proof: The empty sum is the constant zero function, whose Jacobian vanishes by Theorem 6; the right side is an empty sum.

  2. case \(S = \{ a\} \cup T\) with \(a \notin T\), assuming the claim holds for \(T\) (induction hypothesis): \(\operatorname {pdiv}\bigl(\sum _{s \in S} f_s\bigr)\, x\, i\, j = \operatorname {pdiv}f_a\, x\, i\, j + \sum _{s \in T} \operatorname {pdiv}f_s\, x\, i\, j\).
    proof: Split the sum as \(f_a + \sum _{s \in T} f_s\) (Finset.sum_insert, using \(a \notin T\)). The tail \(\sum _{s \in T} f_s\) is differentiable at \(x\) as a finite sum of functions differentiable by assumption 2 (DifferentiableAt.fun_sum), so the sum rule (Theorem 3) splits the Jacobian; the induction hypothesis rewrites the tail.

  3. q.e.d.
    proof: By induction on the finite set \(S\) (Finset.induction_on): 1 is the base case, and re-merging the sum in 2 gives the inductive step.

Definition 9 VJP record
✓
#

Uses

Lean declarations

For \(f : \mathbb {R}^{m} \to \mathbb {R}^{n}\), \(\mathsf{HasVJP}\, f\) bundles a backward function \(B : \mathbb {R}^{m} \to \mathbb {R}^{n} \to \mathbb {R}^{m}\) with its correctness claim: for all \(x\), \(dy\), \(i\),

\[ B(x, dy)_i = \sum _j \operatorname {pdiv}f\, x\, i\, j \cdot dy_j. \]

Exhibiting a \(\mathsf{HasVJP}\, f\) is exactly the statement “this backward function computes the vector–Jacobian product of \(f\).”

Theorem 10 VJP chain rule
✓
#

Uses

Lean declarations

assume:

  1. \(B_f\) is a correct backward function for \(f\) (\(\mathsf{HasVJP}\, f\)) [hf]

  2. \(B_g\) is a correct backward function for \(g\) (\(\mathsf{HasVJP}\, g\)) [hg]

  3. \(f\) is differentiable everywhere [hf_diff]

  4. \(g\) is differentiable everywhere [hg_diff]

prove: \(\mathsf{HasVJP}\, (g \circ f)\).

Proof ▶

Sketch: the composite backward is “run \(B_g\), feed the result to \(B_f\)”; correctness is the two correct fields glued by the chain rule for \(\operatorname {pdiv}\).

  1. Define \(B(x, dy) := B_f\bigl(x,\; B_g(f\, x,\; dy)\bigr)\).

  2. suffices: for all \(x\), \(dy\), \(i\): \(B(x, dy)_i = \sum _k \operatorname {pdiv}(g \circ f)\, x\, i\, k \cdot dy_k\).
    proof: Definition 9, with \(B\) as the candidate backward function.

  3. \(B(x, dy)_i = \sum _j \operatorname {pdiv}f\, x\, i\, j \cdot B_g(f\, x, dy)_j\).
    proof: Correctness of \(B_f\) (assumption 1).

  4. \(\mathord {@}= \sum _j \operatorname {pdiv}f\, x\, i\, j \sum _k \operatorname {pdiv}g\, (f\, x)\, j\, k \cdot dy_k\).
    proof: Correctness of \(B_g\) (assumption 2).

  5. \(\mathord {@}= \sum _k \Bigl(\sum _j \operatorname {pdiv}f\, x\, i\, j \cdot \operatorname {pdiv}g\, (f\, x)\, j\, k\Bigr)\, dy_k\).
    proof: Distribute and swap the two sums (Finset.mul_sum, Finset.sum_comm); both are finite, so no convergence question arises.

  6. q.e.d.
    proof: By the chain rule (Theorem 2), applicable at every \(x\) by assumptions 3 and 4, the inner sum in 5 is \(\operatorname {pdiv}(g \circ f)\, x\, i\, k\); with 2 this is the goal.

Theorem 11 Additive fan-in VJP
✓
#

Uses

Lean declarations

Used for residual connections. assume:

  1. \(B_f\) is a correct backward function for \(f\) (\(\mathsf{HasVJP}\, f\)) [hf]

  2. \(B_g\) is a correct backward function for \(g\) (\(\mathsf{HasVJP}\, g\)) [hg]

  3. \(f\) is differentiable everywhere [hf_diff]

  4. \(g\) is differentiable everywhere [hg_diff]

prove: \(\mathsf{HasVJP}\, (f + g)\).

Proof ▶

  1. Define \(B(x, dy)_i := B_f(x, dy)_i + B_g(x, dy)_i\).

  2. suffices: for all \(x\), \(dy\), \(i\): \(B(x, dy)_i = \sum _j \operatorname {pdiv}(f + g)\, x\, i\, j \cdot dy_j\).
    proof: Definition 9, with \(B\) as the candidate backward function.

  3. \(B(x, dy)_i = \sum _j \bigl(\operatorname {pdiv}f\, x\, i\, j + \operatorname {pdiv}g\, x\, i\, j\bigr) \cdot dy_j\).
    proof: Correctness of \(B_f\) and \(B_g\) (assumptions 1, 2); merge the two finite sums (Finset.sum_add_distrib) and factor out \(dy_j\) (ring).

  4. q.e.d.
    proof: The sum rule (Theorem 3), applicable at every \(x\) by assumptions 3 and 4, rewrites the bracket in 3 to \(\operatorname {pdiv}(f + g)\, x\, i\, j\); with 2 this is the goal.

Theorem 12 Multiplicative fan-in VJP
✓
#

Uses

Lean declarations

Used for Squeeze-and-Excitation. assume:

  1. \(B_f\) is a correct backward function for \(f\) (\(\mathsf{HasVJP}\, f\)) [hf]

  2. \(B_g\) is a correct backward function for \(g\) (\(\mathsf{HasVJP}\, g\)) [hg]

  3. \(f\) is differentiable everywhere [hf_diff]

  4. \(g\) is differentiable everywhere [hg_diff]

prove: \(\mathsf{HasVJP}\, (f \odot g)\), where \((f \odot g)\, x\, i = f\, x\, i \cdot g\, x\, i\).

Proof ▶

  1. Define \(B(x, dy)_i := B_f\bigl(x,\; g(x) \odot dy\bigr)_i + B_g\bigl(x,\; f(x) \odot dy\bigr)_i\), where \(\bigl(g(x) \odot dy\bigr)_j = g\, x\, j \cdot dy_j\).

  2. suffices: for all \(x\), \(dy\), \(i\): \(B(x, dy)_i = \sum _j \operatorname {pdiv}(f \odot g)\, x\, i\, j \cdot dy_j\).
    proof: Definition 9, with \(B\) as the candidate backward function.

  3. \(B(x, dy)_i = \sum _j \bigl(\operatorname {pdiv}f\, x\, i\, j \cdot g\, x\, j + f\, x\, j \cdot \operatorname {pdiv}g\, x\, i\, j\bigr) \cdot dy_j\).
    proof: Correctness of \(B_f\) at upstream gradient \(g(x) \odot dy\) and of \(B_g\) at \(f(x) \odot dy\) (assumptions 1, 2); merge the two finite sums and factor out \(dy_j\) (Finset.sum_add_distrib, ring).

  4. q.e.d.
    proof: The product rule (Theorem 4), applicable at every \(x\) by assumptions 3 and 4, identifies the bracket in 3 with \(\operatorname {pdiv}(f \odot g)\, x\, i\, j\); with 2 this is the goal.

Theorem 13 Identity VJP
✓
#

Uses

Lean declarations

Proof ▶

Mechanical; see Proofs.identity_has_vjp.

1.4 Example: MNIST linear classifier

The smallest network that learns: a single dense layer mapping 784-dim images to 10-dim logits, trained with softmax cross-entropy and plain SGD. About 7,850 parameters. \(\sim \)92% test accuracy in seconds.

\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.2cm}, io/.style = {col, draw=blue!55!black, fill=blue!8}, 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: the whole "network" is one dense layer. \node[io] (input) {Input \;\; $28\times28$ flatten $\to 784$}; \node[head, below=0.24cm of input] (d1) {\textbf{Dense} $784\to10$ \;(identity)}; \node[logits, below=0.24cm of d1] (out) {Logits \;\; 10 classes, softmax-CE}; \foreach \a/\b in {input/d1, d1/out} \draw[arr] (\a) -- (\b); \end{tikzpicture}

The forward pass and loss, in math.

For an input image \(x \in \mathbb {R}^{784}\) and true class index \(t \in \{ 0, \dots , 9\} \):

\begin{align*} z & = W x + b & & \text{(dense layer; $W \in \mathbb {R}^{10 \times 784}$, $b \in \mathbb {R}^{10}$)} \\ \hat{y}_i & = \mathrm{softmax}(z)_i = \frac{e^{z_i}}{\sum _j e^{z_j}} & & \text{(class probabilities)} \\ \ell (W, b;\, x, t) & = -\log \hat{y}_t & & \text{(cross-entropy on the true class)} \end{align*}

The trainable parameters are \(W\) and \(b\), which is 7,850 floats in total.

The same thing, as a VerifiedNetSpec.

def linearVerified : VerifiedNetSpec where
  name     := "MNIST-Linear"
  slug     := "linear"
  inC      := 1
  imageH   := 28
  imageW   := 28
  nClasses := 10
  data     := .mnist
  layers   := [.dense 784 10]
  blurb    := "MNIST-Linear via the VERIFIED renderer
                (pretty∘emit) → %LOWERER% → GPU"

The layer emits raw logits \(z = W x + b\). Softmax and cross-entropy are added by the training driver, because this is a classification task. The slug names the committed render, so this spec is bound to verified_mlir/linear_train_step.mlir and linear_fwd.mlir, which are the two files the run at the top of the chapter compiled. The blurb is the banner the run prints before its first epoch, %LOWERER% filled in with whichever lowerer was selected; it is one line in the source and wrapped here.

What .dense 784 10 contributes.

Layers are constructors of a VLayer inductive, and one function, toSpecs, says what each contributes to the parameter list as (dims, initKind) pairs — 0 is He(fan-in), 1 is ones, 2 is zeros: 

def toSpecs : VLayer -> Array (Array Nat × Nat)
  | dense ic oc => #[(#[ic,oc],0),(#[oc],2)]
  | ...

Two tensors: a \(784 \times 10\) weight matrix and a length-10 bias. That is this network’s entire parameter list, and d0, the input width the render needs, is derived from the same list rather than declared separately. Every chapter after this one adds an arm to that match. None of them changes this one, so this is the last time you will see it.

The init codes are the shared driver’s defaults, and this chapter is the one place they do not apply. trainLinear sets both tensors to zero outright rather than reading them off toSpecs, which is why the run at the top of the chapter lands on the same accuracy whatever the seed. From Chapter 2 onward the He(fan-in) in code 0 is what actually runs.

That binding is why the type exists. The spec lives in LeanMlir/VerifiedNets.lean rather than beside the trainer, because the theorem and the program have to name the same object. In LeanMlir/Proofs/Foundation/SpecVJP.lean: 

noncomputable def linearVerified_has_vjp (W : Mat 784 10) (b : Vec 10) :
    HasVJP (denoteLinear linearVerified.layers W b) :=
  dense_has_vjp W b

Read the type. The VJP is proved for the denotation of linearVerified.layers, and that is the same field the trainer reads to build its graph. Copy the layer list into a second definition and the proof would be about a different network than the one that runs.

The whole program.

The trainer is 54 lines including its docstring, so here it is entire, with the comments stripped: 

import LeanMlir.VerifiedNets

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

def main (argv : List String) : IO Unit :=
  linearVerified.trainLinear linearConfig (argv.head?.getD "data")

That is apps/mnist/MainMnistLinearVerified.lean, all of it. One spec, one config, one call. The GPU backend does not appear in it at all, because the shim is opened with dlopen at run time.

Notice also what linearConfig does not carry: a learning rate. The real rate is baked into verified_mlir/linear_train_step.mlir at render time, so the recipe that trains is the recipe the proofs describe, and the two cannot drift apart.

The gradient, in math.

Backpropagation through this network produces a single outer product:

\[ \nabla _W \ell = (\hat{y} - e_t) \otimes x, \qquad \nabla _b \ell = \hat{y} - e_t \]

where \(e_t \in \mathbb {R}^{10}\) is the one-hot vector for the true class. Ch 1’s theorems (chain rule + identity Jacobian/VJP) plus the Dense Jacobian and softmax-CE gradient (formalized in Ch 2, both themselves proved from this chapter’s foundation rules) are what guarantee these formulas are correct. The verified renderer emits exactly these as fused StableHLO, and that rendering is verified_mlir/linear_train_step.mlir, the file the run at the top of the chapter compiled.

Results.

The run at the top of this chapter is this network. Accuracy climbs from 89.77% after one epoch to 92.10% after twelve, and the curve is essentially flat from epoch 9 on, which is what a model with no hidden layer looks like once it has learned everything it can. Adding hidden layers and ReLU in Chapter 2 takes the same dataset and the same recipe to 97.83%. That gap of about 5.7 points is the value of non-linearity on this task.

Where this goes.

Every architecture in Part 2 extends this template the same way, one new primitive at a time:

  • Ch 2 stacks dense layers and inserts ReLU between them. One new operator VJP.

  • Ch 3 swaps the dense forward for conv2d. One new operator VJP.

  • Ch 4 adds BatchNorm. One new operator VJP.

  • Ch 5 adds the residual skip. Additive fan-in, Theorem 11, and no new operator at all.

  • Ch 6 factors standard conv into depthwise plus pointwise. One new operator VJP, recombined through the chain rule.

  • Ch 7 adds the SE channel-attention block, which is an elementwise product over global pooling and dense layers. One new VJP composed from pieces we already have.

  • Ch 8 replaces ReLU with GELU. One new operator.

  • Ch 9 adds attention, which needs new matrix-level machinery and softmax-VJP chains.

The structural rules from this chapter never change. Later chapters add operator-specific theorems and compose them through the chain rule we just proved. The driver that runs all of them is the same Lean either way, and it is shorter than you would expect, which is the subject of §1.6.

1.5 MLIR: Linear

The StableHLO in this section is the artifact the proofs are about, and it is what XLA compiles and runs. Nothing between the theorem and the GPU rewrites it.

The forward half is short enough to print whole. This is verified_mlir/linear_fwd.mlir exactly as committed, and it is one of the two files the run at the top of the chapter compiled: 

module @m {
  func.func @linear_fwd(%x: tensor<128x784xf32>,
                        %W0: tensor<784x10xf32>,
                        %b0: tensor<10xf32>) -> tensor<128x10xf32> {
    %v0 = stablehlo.dot_general %x, %W0,
            contracting_dims = [1] x [0],
            precision = [DEFAULT, DEFAULT]
            : (tensor<128x784xf32>, tensor<784x10xf32>)
              -> tensor<128x10xf32>
    %v1 = stablehlo.broadcast_in_dim %b0, dims = [1]
            : (tensor<10xf32>) -> tensor<128x10xf32>
    %v2 = stablehlo.add %v0, %v1 : tensor<128x10xf32>
    return %v2 : tensor<128x10xf32>
  }
}

Three operations for \(z = Wx + b\): a contraction, a broadcast of the bias across the batch, and an add. The shapes are the real ones, batch 128 against 784 inputs and 10 classes, because this file is the render that trains and not an illustration of one.

The backward pass is one operation. With the loss cotangent \(dy = \mathrm{softmax}(\mathrm{logits}) - \text{onehot}\), the input gradient is \(dx = dy\, W^{\! \top }\), a single matrix multiply. Here is what the code generator emits for it, shown at a small representative size so the shapes stay readable: 

func.func @linear_back(%dy: tensor<2x3xf32>, %W0: tensor<4x3xf32>)
    -> tensor<2x4xf32> {
  %bk0 = stablehlo.dot_general %dy, %W0, contracting_dims = [1] x [1]
           : (tensor<2x3xf32>, tensor<4x3xf32>) -> tensor<2x4xf32>
  return %bk0 : tensor<2x4xf32>
}

That dot_general contracts the output axis of \(dy\) against the output axis of \(W\), which is multiplication by \(W^{\! \top }\), and by a machine-checked theorem that is the linear layer’s exact reverse-mode derivative. Everything else in this book scales that one move up. Each architecture chapter closes with an “MLIR: operator” section that takes the chapter’s new operator and shows that its emitted backward is likewise the rendering of a proven derivative. Those operators are convolution, BatchNorm, the residual fan-in, depthwise convolution, squeeze-excitation, layer scale, and attention.

Each of those listings is shown at a small representative size, meaning a few channels, tokens, or units of width. The per-operator proof beneath is dimension-parameterized and the printer emits the identical graph at production scale, so each section states its toy shape and moves on. Appendix C covers the shared machinery underneath them: the denoted intermediate representation, the bridge theorems that tie it to the proofs, the printer that turns it into the text above, and the execution oracle that checks the result on the GPU.

1.6 What’s inside .train?

The section above treated linearVerified.trainLinear as a black box. You named the network, you named the epochs and the batch size, and training happened. That is deliberately the user-facing interface, but there is no magic underneath it. The driver is ordinary Lean in LeanMlir/VerifiedTrain.lean, and the interesting thing about it is how little it does.

The training loop, the way a math book would write it.

Algorithm: Mini-batch SGD, with the step compiled into the graph
Input:  net (a VerifiedNetSpec), cfg (epochs, batch size),
        dataset D = (X_train, y_train, X_test, y_test)
Output: trained parameters (W, b)

// 1. Open sessions over the two committed renders
ts  <- session(verified_mlir/<slug>_train_step.mlir)
fwd <- session(verified_mlir/<slug>_fwd.mlir)

// 2. Load data; zero-initialise the parameters
(X, y), (Xt, yt) <- load(D)
(W, b) <- 0

// 3. Epoch loop
for epoch = 1 to cfg.epochs:

    // 4. Batch loop: ONE call per batch
    for each mini-batch (x, t) of size B:
        (W, b) <- ts(x, W, b, onehot(t))

    // 5. Evaluate on the whole test set
    eval(fwd, W, b, Xt, yt)

The same loop, in Lean.

Here is VerifiedNet.trainLinear, the driver that produced the run at the top of this chapter. Three things are left out: the IO.getStdout flushes, two environment-variable overrides (LEAN_MLIR_PERTURB_R, which displaces the initialisation for a residency check, and LEAN_MLIR_MAX_EPOCHS, which caps the epoch count), and everything after the loop, which dumps the final parameters so the two lowerers can be diffed tensor by tensor. Two long lines are wrapped; every line that computes the answer is here: 

def VerifiedNet.trainLinear (net : VerifiedNet) (cfg : VerifiedConfig)
    (dataDir : String) : IO Unit := do
  let bs := cfg.batchSize
  let d0 := net.d0
  let d1 := net.nClasses
  net.printBlurb

  -- 1. One session per committed render, and the entry point in each
  let tsSess  <- mkSession s!"{net.mlirDir}/{net.slug}_train_step.mlir"
  let fwdSess <- mkSession s!"{net.mlirDir}/{net.slug}_fwd.mlir"
  let tsFn    := s!"m.{net.slug}_train_step"
  let fwdFn   := s!"m.{net.slug}_fwd"

  -- 2. Data, then zero-initialised parameters
  let (trainImg, trainLbl, nTrain, evalImg, evalLbl, nEval, _, _) <-
    loadData net dataDir
  let evalName := match net.data with | .imagenette => "val" | _ => "test"
  IO.println s!"  train {nTrain}, {evalName} {nEval};
                  dense {d0}->{d1}, bs {bs}, SGD"
  let mut W0 <- F32.const (d0 * d1).toUSize 0.0
  let mut b0 <- F32.const d1.toUSize 0.0
  let nb     := nTrain / bs
  let nbt    := (nEval + bs - 1) / bs   -- ceil: last eval batch is zero-padded
  let nResident : USize := 2            -- W0 and b0 stay on the device
  let pBytes := (d0 * d1 + d1) * 4
  let shapes := net.shapesBA            -- packed [W0|b0] layout for the forward
  let xShape := net.xShape bs
  let mut packed := W0 ++ b0

  -- 3. Epoch loop
  for ep in [0:cfg.epochs] do
    let tEp0 <- IO.monoMsNow

    -- 4. Batch loop: forward, loss, backward, SGD -- one call
    for bi in [0:nb] do
      let xb := F32.sliceImages trainImg (bi * bs) bs d0
      let yb := F32.sliceLabels trainLbl (bi * bs) bs
      let out <- LowererSession.linearTrainStepV tsSess tsFn
                   xb W0 b0 yb bs.toUSize d0.toUSize d1.toUSize nResident
      packed := out
      W0 := out.extract 0 (d0 * d1 * 4)
      b0 := out.extract (d0 * d1 * 4) pBytes

    -- 5. Evaluate on the full test set, every epoch
    packed <- LowererSession.readParams tsSess packed pBytes.toUSize
    W0 := packed.extract 0 (d0 * d1 * 4)
    b0 := packed.extract (d0 * d1 * 4) pBytes
    let params := packed
    let mut correct := 0
    for bi in [0:nbt] do
      let xb := F32.sliceImagesPad evalImg (bi * bs) bs d0 nEval
      let logits <- LowererSession.forwardF32 fwdSess fwdFn params shapes
                      xb xShape bs.toUSize d1.toUSize nResident (ep+1).toUSize
      for j in [0:min bs (nEval - bi * bs)] do  -- real rows only, not the pad
        let pred := (F32.argmaxN logits (j * d1).toUSize d1.toUSize).toNat
        if pred == F32.readLabel evalLbl (bi * bs + j) then
          correct := correct + 1
    let acc := correct.toFloat / nEval.toFloat * 100.0
    let epMs := (<- IO.monoMsNow) - tEp0
    IO.println s!"  epoch {ep+1}: {evalName}_acc = {correct}/{nEval} = {acc}%
                    [95% CI {wilson95 correct nEval}] ({epMs}ms)"

Walking through the numbered sections:

1. Sessions. mkSession takes the path to a committed .mlir, hands it to the shim, and gets back something callable. XLA compiles the StableHLO in process, which is the 604 ms the run reported.

2. Data and initialisation. loadData mmaps the on-disk binaries into F32Array buffers. trainImg holds flattened pixels and trainLbl holds integer class labels. Weights start at zero rather than He-initialised, which is fine for a model with no hidden layer and no symmetry to break.

3. Epoch loop. Twelve passes over the training set. Notice what is not here: no learning-rate computation, because the rate is baked into the render.

4. Batch loop, the core of training. Slice a batch of images and labels, then call linearTrainStepV. That one call does everything: the forward pass, the cross-entropy loss, the backward pass built from the VJPs proved earlier in this chapter, and the SGD update. It returns the new parameters packed as [W0|b0].

This one line is where the whole pipeline pays off. There is no Python. There is no per-step graph construction. There is no autograd interpreter walking a tape. The training step was compiled once at startup, in 604 ms, and every batch after that is a single dispatched call. Chapter 1’s run does 468 of those per epoch and finishes an epoch in about 224 ms including a full pass over the 10,000 test images.

5. Validation. Every epoch, not every tenth, because a forward pass over the test set is cheap next to the training epoch. Two details are worth naming now because every later chapter inherits them. The eval loop reads the parameters back with readParams rather than using the W0/b0 it already has, because under device residency those two are stale by design and packed is the authoritative copy. And nbt rounds up: the last eval batch is zero-padded to bs and the inner loop scores only the real rows, so a test set whose size is not a multiple of the batch size is fully scored rather than silently truncated. Ten thousand images and a batch of 128 would otherwise drop the last 16. The forward-only render is a separate artifact rather than a branch inside the training graph, which matters more in Chapter 4 where BatchNorm has to use running statistics at inference and per-batch estimates during training.

That is the whole function. The reason it is short is that everything heavyweight lives in one compiled graph on the GPU. Lean’s job here is not to compute gradients. It is to specify what the training step is, prove that specification correct, and invoke it. Three concerns, cleanly separated, and only the middle one is unusual.

Every chapter after this one hands a different VerifiedNetSpec to the same driver. Only the spec changes. That is the framework pitch in concrete form: once the driver is written and the layers are proved one at a time, every architecture in the book runs through it with no further infrastructure code. It all runs on your hardware, and Appendix B walks you through the toolchain setup.

1.7 MLIR: Training Step

Section 1.6 made a claim about the inner call. One dispatched StableHLO graph does the forward pass, the loss, the backward pass, and the optimizer update, with no tape and no per-step graph construction. The “MLIR: Linear” section above rendered a single operator’s backward in isolation. Here is the whole thing the linear classifier compiles to. This is verified_mlir/linear_train_step.mlir as committed, with the default precision attributes dropped and the long type signatures wrapped for the page. Nothing else is removed, including one band that repeats: 

func.func @linear_train_step(
    %x: tensor<128x784xf32>,  %W0: tensor<784x10xf32>,
    %b0: tensor<10xf32>,      %onehot: tensor<128x10xf32>)
    -> (tensor<784x10xf32>, tensor<10xf32>) {
  // -- forward + softmax-CE cotangent (rendered from lossCotGraph) --
  %v0 = stablehlo.dot_general %x, %W0, contracting_dims = [1] x [0]
          : (tensor<128x784xf32>, tensor<784x10xf32>) -> tensor<128x10xf32>
  %v1 = stablehlo.broadcast_in_dim %b0, dims = [1]
          : (tensor<10xf32>) -> tensor<128x10xf32>
  %v2 = stablehlo.add %v0, %v1 : tensor<128x10xf32>
  %v3 = stablehlo.exponential %v2 : tensor<128x10xf32>
  %v4 = stablehlo.constant dense<0.0> : tensor<f32>
  %v5 = stablehlo.reduce(%v3 init: %v4)
          applies stablehlo.add across dimensions = [1]
          : (tensor<128x10xf32>, tensor<f32>) -> tensor<128xf32>
  %v6 = stablehlo.broadcast_in_dim %v5, dims = [0]
          : (tensor<128xf32>) -> tensor<128x10xf32>
  %v7 = stablehlo.divide %v3, %v6 : tensor<128x10xf32>   // softmax
  %v8 = stablehlo.subtract %v7, %onehot : tensor<128x10xf32>   // dy
  // -- W0:  dW0 = x^T . dy,  then theta' = theta - lr*dW0 --
  %v9 = stablehlo.dot_general %x, %v8, contracting_dims = [0] x [0]
          : (tensor<128x784xf32>, tensor<128x10xf32>) -> tensor<784x10xf32>
  %v10 = stablehlo.constant dense<0.00078125> : tensor<784x10xf32>
  %v11 = stablehlo.multiply %v9, %v10 : tensor<784x10xf32>
  %v12 = stablehlo.subtract %W0, %v11 : tensor<784x10xf32>
  // -- v13-v21: v0-v8 emitted a SECOND time, verbatim --
  %v13 = stablehlo.dot_general %x, %W0, contracting_dims = [1] x [0]
          : (tensor<128x784xf32>, tensor<784x10xf32>) -> tensor<128x10xf32>
  %v14 = stablehlo.broadcast_in_dim %b0, dims = [1]
          : (tensor<10xf32>) -> tensor<128x10xf32>
  %v15 = stablehlo.add %v13, %v14 : tensor<128x10xf32>
  %v16 = stablehlo.exponential %v15 : tensor<128x10xf32>
  %v17 = stablehlo.constant dense<0.0> : tensor<f32>
  %v18 = stablehlo.reduce(%v16 init: %v17)
          applies stablehlo.add across dimensions = [1]
          : (tensor<128x10xf32>, tensor<f32>) -> tensor<128xf32>
  %v19 = stablehlo.broadcast_in_dim %v18, dims = [0]
          : (tensor<128xf32>) -> tensor<128x10xf32>
  %v20 = stablehlo.divide %v16, %v19 : tensor<128x10xf32>   // softmax again
  %v21 = stablehlo.subtract %v20, %onehot : tensor<128x10xf32>   // dy again
  // -- b0:  db0 = sum_batch dy,  then the same three-op update --
  %v22 = stablehlo.constant dense<0.0> : tensor<f32>
  %v23 = stablehlo.reduce(%v21 init: %v22)
          applies stablehlo.add across dimensions = [0]
          : (tensor<128x10xf32>, tensor<f32>) -> tensor<10xf32>
  %v24 = stablehlo.constant dense<0.00078125> : tensor<10xf32>
  %v25 = stablehlo.multiply %v23, %v24 : tensor<10xf32>
  %v26 = stablehlo.subtract %b0, %v25 : tensor<10xf32>
  return %v12, %v26 : tensor<784x10xf32>, tensor<10xf32>
}

One function, one straight line of dataflow, no control flow: the signature takes the parameters and a batch of data and returns the updated parameters, \((\theta ,\, \text{data}) \mapsto \theta '\). Read it in four bands. The first (%v0–%v8) fuses the forward pass and the loss cotangent, so it holds dot_general for the logits, the exp/reduce/divide softmax, and the subtract %onehot that yields \(dy = \partial \mathrm{CE}/\partial \text{logits}\). That whole band is the rendering of a single verified graph, lossCotGraph, certified equal to the softmax–cross-entropy gradient by lossCotGraph_isCEgrad. The second (%v9–%v12) is the weight half: dot_general %x, %v8 is exactly the move from the “MLIR: Linear” section, \(x^{\! \top } dy\), now serving as the weight gradient, and the three operations after it scale by the learning rate and subtract.

The third band (%v13–%v21) is the first band over again, operation for operation. The printer emits one cotangent chain per parameter rather than computing \(dy\) once and using it twice, so the forward pass and the softmax appear a second time before the fourth band (%v22–%v26) sums \(dy\) over the batch for the bias gradient and applies the same three-operation update. Whether the copy survives to the GPU is the compiler’s business — common-subexpression elimination is a standard pass — but the certified artifact contains both, and what this section shows is the artifact. Nothing is gained by pretending otherwise, and the duplication is a fair picture of what naive rendering costs: 27 operations where 19 would do.

There is no fifth band, no copy-back, and no separate update kernel. Forward, loss, backward, and step are one compiled function, which is exactly why § 1.6’s loop can treat a training step as a single dispatch.

Why SGD here. linearVerified trains with plain \(\theta ' = \theta - \alpha \nabla \) at a constant learning rate, which is three operations per tensor — the listing’s last band is the entire optimizer. Adam is available here and the deeper chapters use it, but it costs roughly 2.3 times the graph: Chapter 4’s CIFAR network, rendered both ways, goes from 1,001 StableHLO operations under SGD to 2,293 under AdamW, because the first- and second-moment buffers enter and leave the signature and an rsqrt appears. Size is not the deciding reason, though. Both optimizers are proved faithful, each emitted update certified to be its own update law applied to the same proven gradient. Only SGD also gets a descent theorem, that the step decreases the loss — and Adam cannot have one, because it is not monotone. That is a fact about Adam, the AMSGrad counterexample, and not a gap in the formalization. SGD is the optimizer here where both halves close, which is why it is the one to read first. Appendix C carries both certificates, and how far up the depth the descent half actually reaches.

The learning rate is folded in. The weight-gradient dot_general contracts the batch axis, so \(dW_0\) is a batch sum rather than a mean. The constant 0.00078125 is \(\alpha /N = 0.1/128\), the per-sample rate that turns that sum back into the mean update. Everything else in the listing is, line for line, the rendering of a theorem proved earlier in this chapter.

That one function is the only heavy thing in the loop that drives it. Stripped of its logging and its session setup, § 1.6’s driver is five lines: 

(W, b) := 0                            // zero-initialised, no Adam buffers
for epoch in 0..E:
  for batch in batches(data):
    (W, b) := train_step(W, b, batch)  // <- the function above
  eval(W, b)                           // the forward-only render

Five lines of host control flow around one compiled dispatch. The loop counts epochs and slices batches. It never computes a gradient or touches a tensor. The lone train_step call is the function above, and it is the only place arithmetic happens.

Where this stands. The scheme above covers more than the linear classifier. The shallow chapters’ train steps (MLP, CNN, CIFAR\(\pm \)BN) are closed both ways, and all five Part I architectures are proved at full architecture: ResNet-34, MobileNetV2, EfficientNet-B0, ConvNeXt-T, and ViT. That means machine-checked forward graphs throughout, whole-network backward at full depth, the non-triviality seal at full depth for all five (which is that the proven backward is nonzero at a witness), and three classical axioms. All five committed training steps were then re-rendered at the production trainers’ exact signatures and tied: every parameter update is proved to consume the cotangent its own backward pass delivers, threaded through the real forward, and each agrees on GPU to float rounding. There is one open gap, and it is narrower than it used to be. The theorems are over \(\mathbb {R}\); the GPU runs a finite precision. What sits between them is not a float32-specific argument but a budget over an arbitrary rounding model (§ C.4), which float32, bf16-mixed and fp8-E4M3 all instantiate. Part II’s bestiary is deliberately lighter, decomposed into the same proved primitives but not carried to whole-step closes. The book’s official claim is exactly this, no more and no less: every architecture in Part I trained end-to-end under one machine-checked scheme.