Foundations

Mathlib’s fderiv

Backprop has exactly one job: compute the gradient of the loss with respect to every parameter. That gradient is the del operator \(\nabla \mathcal{L}\), the vector of partials \(\partial \mathcal{L}/\partial \theta \). So before we write a single layer, we need a rigorous, machine-checkable notion of “the derivative of this function at this point.”

Mathlib supplies it as the Fréchet derivative, \(\operatorname {fderiv}\). It is the best linear approximation, the continuous linear map \(D\) with \(f(x+h) \approx f(x) + D\, h\), and it is defined over any normed space rather than only \(\mathbb {R}^n\). For a layer \(f : \mathbb {R}^{m} \to \mathbb {R}^{n}\), \(\operatorname {fderiv}\, \mathbb {R}\, f\, x\) is the linear map whose matrix is the Jacobian.

We read individual partials off it. A single \((i,j)\) Jacobian entry is \(\operatorname {pdiv}\, f\, x\, i\, j = \big(\operatorname {fderiv}\, \mathbb {R}\, f\, x\, (e_i)\big)_j\), so you apply \(\operatorname {fderiv}\, \mathbb {R}\, f\, x\) to the \(i\)-th basis vector \(e_i\) and take component \(j\). The Jacobian and the gradient \(\nabla \mathcal{L}\) you already picture are assembled entry-by-entry from \(\operatorname {pdiv}\), and \(\operatorname {pdiv}\) is \(\operatorname {fderiv}\) underneath. Every layer in this book is built on this one accessor.

Because \(\operatorname {pdiv}\) is defined from \(\operatorname {fderiv}\) rather than assumed, the structural rules of chain, sum, and product are theorems proved from Mathlib’s API rather than axioms we ask you to trust, and the same holds all the way up: \(\texttt{conv2d}\), multi-head SDPA, patch embedding, the BN inverse-stddev broadcast, GELU, the softmax Jacobian, and every transformer-level composition chain are defined concretely and proved from the foundation. The cost is mild, because a \(\mathsf{Differentiable}\) hypothesis propagates through each chapter. The benefit is that \(\texttt{\# print axioms vitFullHasVJP}\) bottoms out at Lean core (propext, Classical.choice, Quot.sound) and nothing else. No thumb on the scale. Appendix B shows how that claim is independently audited.

What about non-smooth points?

ReLU is \(\max (0, x)\), which has no derivative at \(x=0\). Absolute value flunks the same test for the same reason. maxPool over a window has no derivative when two entries tie. Mathlib’s \(\operatorname {fderiv}\) is a total function, so at such a point it returns a value anyway: its default, the zero map. ML frameworks also have to return something, and they pick a different convention. PyTorch and JAX both use \(0\) as ReLU’s slope at the corner, which keeps the gradient flowing on every other coordinate, and they break maxPool ties by the lowest index. Those choices aren’t theorems. They’re conventions that make backprop computable everywhere.

Both conventions appear in this book, and it matters which is which. \(\operatorname {pdiv}\) is \(\operatorname {fderiv}\) by definition, everywhere, so for a non-smooth \(f\) we don’t try to prove a closed-form \(\mathsf{HasVJP}\) from it, which would fail at the singular set. Instead we define the \(\mathsf{HasVJP}\) directly, with backward function \(\mathrm{backward}(x, dy)_i := \sum _j \operatorname {pdiv}\, f\, x\, i\, j \cdot dy_j\). The correct field then holds by rfl, because by definition that is exactly what backward computes. The later chapters call this the pdiv-derived backward: ReLU, maxPool and the MLP composition each get a \(\mathsf{HasVJP}\) built this way. Where \(f\) is differentiable it is the true Jacobian. Where it is not, every entry is \(0\), so at a point with even one coordinate exactly zero the pdiv-derived ReLU backward is the zero vector, in every coordinate, and at a point with even one tied pooling window the pdiv-derived maxPool backward is zero everywhere.

The codegen emits the framework convention instead: the mask \([x {\gt} 0]\) and the lowest-index tie-break. Off the singular set the two are the same number, and that is a theorem (relu_codegen_matches_canonical, maxPool2_codegen_matches_canonical). On it they differ, and that difference is the one place where the proven backward and the emitted text part company. The benefit of writing it this way is a paper trail: every theorem that reaches an emitted ReLU or max-pool carries a smooth-point hypothesis naming that set, so the typechecker has verified the algebra everywhere the two agree and has marked exactly where it has not.

These conventions are invoked by ReLU, by maxPool and by ReLU6, which is to say by Chapters 1–6, and a whole-network theorem in those chapters names an input and assumes every activation stayed off its corner there. From EfficientNet on, every activation in Part 1 is smooth — swish, the sigmoid of a squeeze-and-excitation gate, GELU, softmax — and there is no max-pool, so there is no smooth-point hypothesis to carry: the statements quantify over every input, as Chapter 7’s fold is the first to do. \(\operatorname {pdiv}\) stays on as the notation the Jacobian lemmas are written in, and from there on it is the true Jacobian everywhere; \(\operatorname {fderiv}\)’s default at a corner is never consulted again.

Why VJPs, not Jacobians?

If you’ve taken multivariable calculus, the natural object to track through a network is the Jacobian. For \(f : \mathbb {R}^{m} \to \mathbb {R}^{n}\), that’s the \(n \times m\) matrix of partial derivatives, one row per output and one column per input, and in the accessor above it is \(J_{j,i} = \operatorname {pdiv}\, f\, x\, i\, j\): \(i\) is always the input index and \(j\) the output, the order the Lean definition fixes. The chain rule for a composition \(h = g \circ f\) is matrix multiplication: \(J_h = J_g \cdot J_f\). So if your network is \(L\) layers stacked, the Jacobian of the output with respect to the input is the matrix product \(J_L \cdot J_{L-1} \cdots J_1\).

This is both correct and totally impractical. For a 512-dimensional hidden layer, each \(J_k\) is a \(512 \times 512\) matrix, and multiplying two of them is \(\sim 1.3 \times 10^8\) flops. ResNet-34 has 34 of these layers and ViT has more. You’d be materializing intermediate matrices that are bigger than the network itself, and most of whose entries you never need.

The fix is to notice what you actually want. Training minimizes a scalar loss \(\mathcal{L}\), so the only output gradient you ever need is \(\nabla \mathcal{L} \in \mathbb {R}^{n}\), a single vector. Apply that vector to the chain on the left:

Read this right-to-left. Start with one \(n\)-dimensional vector. Multiply on the right by \(J_L\), get a vector. Multiply by \(J_{L-1}\), get a vector. Continue. At every step you carry a vector and never a matrix, and its dimension matches the layer’s input or output shape.

The single step

is the vector–Jacobian product, or VJP. It takes the gradient flowing in from above and produces the gradient flowing out below. The vector flowing in is called the cotangent, written \(dy\) from here on: the loss’s gradient in the layer’s output, one entry per output. The VJP returns the cotangent one layer down, the loss’s gradient in the layer’s input, one entry per input. The cost is \(O(mn)\) per layer instead of \(O(d^3)\) for a matrix–matrix product, and you never form \(J_f\) at all. You just need a function that, given \(v\), returns \(v^T J_f\). That function is the backward function, and together with \(f\) itself it implements one layer’s contribution to backpropagation.

This is what reverse-mode automatic differentiation does, mechanically, for any program. It’s the same right-to-left vector-pass story: forward to compute activations, backward to compute gradients, both linear in the number of layers.

What \(\mathsf{HasVJP}\) packages. A \(\mathsf{HasVJP}\) record for \(f\) holds two things:

• a backward function, taking the input \(x\) and the cotangent \(dy \in \mathbb {R}^{n}\) to the cotangent below, in \(\mathbb {R}^{m}\)

• a correct field, which is a proof that \(\mathrm{backward}(x, dy)_i = \sum _j \operatorname {pdiv}\, f\, x\, i\, j \cdot dy_j\), meaning that the function really does compute \(dy^T J_f\)

In Lean that’s the whole definition. Two spellings recur in every Lean block of the book: Vec n is a function from the indices below \(n\) to the reals, so a vector is read by applying it to an index, and Fin n is that index type, so \(\sum \) over Fin n is a plain sum of \(n\) terms.

@[ext] structure HasVJP {m n : Nat} (f : Vec m → Vec n) where
  backward : Vec m → Vec n → Vec m
  correct : ∀ (x : Vec m) (dy : Vec n) (i : Fin m),
    backward x dy i = ∑ j : Fin n, pdiv f x i j * dy j

A structure with two fields: a function, and a proof obligation about that function. Constructing a \(\mathsf{HasVJP}\) for some \(f\) means writing both. You don’t get the record without the proof.

The correctness field is the whole point. The backward function is the kernel that ships into the codegen, and the proof is the guarantee that what it computes matches what its name claims. Every later chapter proves a closed-form \(\mathsf{HasVJP}\) for a layer family (dense, conv, BN, attention, \(\ldots \)). Autograd records each forward op on a tape and replays its derivative at backward time. We ship the closed form instead: codegen reads off the backward field and emits a single fused StableHLO kernel, with no tape and no replay.

Inference and training

A deployed network is its forward pass. The weights \(\theta \) are fixed, an image goes in, a prediction \(f(x;\theta )\) comes out, and that is inference: all a trained model ever does. Nothing in this chapter so far is about it. The derivative, the VJP and the \(\mathsf{HasVJP}\) record exist for the other job, which is choosing \(\theta \).

Training adds one step to the forward pass. Pick a loss \(\mathcal{L}\) that scores the prediction against the label. After the forward pass, run the backward pass of the previous section: the cotangent \(\nabla \mathcal{L}\) enters at the output and each layer’s VJP carries it one layer down, until every parameter has its gradient. Then move each parameter a small step against its gradient. Forward, loss, backward, update. In this book those four are one function, \(\texttt{train\_ step} : (\theta , \text{batch}) \mapsto \theta '\), and training is that function applied batch after batch:

θ := init
for epoch in 0..E:
  for batch in batches(data):
    θ := train_step(θ, batch)   // forward, loss, backward, update
  eval(θ)                       // forward only

Five lines of host control flow around two compiled functions. The loop counts epochs and slices batches; it never computes a gradient or touches a tensor. eval is the forward pass, inference again, run over the held-out set to see how \(\theta \) is doing. train_step is the same forward pass with the loss, the backward pass and the update appended. The two are compiled separately, and one is the other’s opening: the linear classifier’s forward render is the first three of its training step’s twenty-seven operations, and Chapter 1 prints both.

Every architecture in Part 1 trains under this loop. Two things change as the chapters go. The data the loop draws its batches from grows, from MNIST to CIFAR to ImageNet. And in every chapter the inside of train_step changes: the forward pass gains a layer, the backward pass gains that layer’s VJP, and the proof that the VJP is correct is the chapter’s theorem. Chapter 1 builds the function for a single matrix multiply, proves it, compiles it to one GPU kernel and runs the loop; §1.6 is the loop as Lean and §1.7 is the compiled function.

What the theorems cover, and what they leave to the run

Five things every later chapter leans on are settled here, so that no chapter has to stop and settle them.

Parameters are inputs. The \(\mathsf{HasVJP}\) record differentiates in \(x\), and training needs the derivative in \(\theta \). There is no second record. A layer is a function of both, \(\mathrm{dense}(W, b)\, x\), and holding \(x\) fixed while \(W\) varies gives another function of a vector, the flattened weights, to which the same definition applies. Theorem 18 is the first instance, the dense layer’s Jacobian in its weights, and Theorem 20 reads the familiar outer product off it: \(dW = x \otimes dy\). Every weight gradient in the book is that move, the record applied along the parameter axis.

The theorems are over \(\mathbb {R}\); the card computes in floats. \(\operatorname {fderiv}\) lives in a normed space, and every statement in Part 1 is about real numbers: exact arithmetic, no rounding. The GPU runs bf16 convolutions under an fp32 classifier head. Nothing in Lean speaks about that gap, and nothing here pretends to. What carries a theorem across it is a second kind of evidence: the StableHLO the GPU runs is printed from the same datatype the proofs reason about, so the code cannot drift from the math, and what the hardware then does with it is cross-checked against independent oracles. Appendix B names the three kinds of certainty, proven, by construction and cross-checked, and §B.5 is the rounding budget. When a verified run matches its reference, or misses it, a wrong gradient is not among the explanations; a rounding story might be.

One example at a time, then the batch. Every record above is per example: one \(x\) in, one \(dy\) back. The loss is a mean over the batch, so the batch gradient is the sum of the per-example VJPs scaled by \(1/N\), and the compiled step of §1.7 does exactly that. The batch is an axis the per-example theorem is lifted over, not a new theorem. Two statements do the lifting for every certificate in the book, the matrix form of the VJP record and the lemma that applies a per-example backward to each row of a batch; they are proved in §9.2.1, and because they are bookkeeping rather than content, the chapter maps of Figures 1-1 and B-1 leave them out. The one exception is BatchNorm, which normalizes each example by statistics of the whole batch, so the batch enters the function itself and the per-example theorem no longer describes it. Chapter 4 proves it twice, per example and at batch statistics, which is why the budget table counts those as two functions, and trains it twice, so the two forms meet on one net. From Chapter 5 on every BatchNorm net is the batch form, which is why each of those five nets carries a data-parallel theorem of its own: the statistic has to cross the replicas before it is the batch’s.

The gradient is where the proofs stop. The update above is plain SGD, and it is the only update the book proves anything about: one step decreases the loss, Chapter 1, with the binary32 gradient allowed in place of the exact one. Everything after the gradient — momentum, Adam, the cosine schedule, warmup, weight decay, label smoothing, an average of the weights — is recipe, chosen because it trains well, and §5.5.1 assembles and ablates it. A recipe is a specification, and the book treats it as one, but convergence is a measurement, never a theorem.

Two trainers from one spec. Every run in the chapters, MNIST through Imagenette, is the verified path: Lean emitting StableHLO from the proven backward, handed to the GPU through PJRT with no interpreter in the loop. §5-2 is that trainer’s shape, which every chapter reuses. The ImageNet runs, from Chapter 5 on, are trained twice. The second trainer is the reference, Lean emitting a JAX program from the same spec and the same recipe, differentiated by JAX’s own autograd and run from Python. The reference is the oracle the verified path has to match, and each ImageNet section’s table puts the two side by side, first at Table 5-3, with the rows that differ above a rule and the rows held equal by construction between the rules, so that a gap, when there is one, has one place to come from.