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}^{n} \to \mathbb {R}^{m}\), \(\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 C 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}\) correctly declines to return a value at these points, but ML frameworks have to return something, so they all pick a convention. PyTorch and JAX both use \(0\) as ReLU’s subgradient at the kink and break maxPool ties by the lowest index. Those choices aren’t theorems. They’re conventions that make backprop computable everywhere.

We make those conventions formal. For a non-smooth \(f\), we don’t try to prove a closed-form \(\mathsf{HasVJP}\) from \(\operatorname {fderiv}\), 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. Wherever \(f\) is smooth, \(\operatorname {pdiv}\) agrees with \(\operatorname {fderiv}\), and that agreement is proved. At the non-smooth points \(\operatorname {pdiv}\) takes the standard convention by definition. The later chapters call this the pdiv-derived backward: ReLU, maxPool and the depthwise input VJP each get a \(\mathsf{HasVJP}\) whose backward is built from \(\operatorname {pdiv}\) and whose correct is rfl.

The codegen then substitutes the same subgradient and argmax convention every framework already uses. The benefit isn’t a different numerical result. It’s a paper trail that says we know exactly which choice we made at each non-smooth point, and the typechecker has verified the algebra around it.

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}^{n} \to \mathbb {R}^{m}\), that’s the \(m \times n\) matrix of partial derivatives \(J_{i,j} = \operatorname {pdiv}\, f\, x\, i\, j\). 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}^{m}\), a single vector. Apply that vector to the chain on the left:

Read this right-to-left. Start with one \(m\)-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 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 a cotangent \(dy \in \mathbb {R}^{n}\) to a vector 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:

@[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.