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 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}\): the best linear approximation — the continuous linear map \(D\) with \(f(x+h) \approx f(x) + D\, h\) — defined over any normed space, not just \(\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\) — apply \(\operatorname {fderiv}\, \mathbb {R}\, f\, x\) to the \(i\)-th basis vector \(e_i\) and take component \(j\). So 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 — chain, sum, product — are theorems proved from Mathlib’s API, not axioms we ask you to trust. Earlier drafts of this book took the other route and axiomatized those eight rules as opaque facts; the current foundation flips it, and every later-chapter axiom has been pruned the same way: \(\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 — a \(\mathsf{Differentiable}\) hypothesis propagates through each chapter. The benefit is that \(\texttt{\# print axioms vit\_ full\_ has\_ vjp}\) bottoms out at Lean core (propext, Classical.choice, Quot.sound) and nothing else — no thumb on the scale. See Appendix C for the full elimination history.

What about non-smooth points?

A reader who has met ReLU will have spotted a gap. 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 — by definition, that’s exactly what backward computes. Wherever \(f\) is smooth, \(\operatorname {pdiv}\) agrees with \(\operatorname {fderiv}\) (proved); at the non-smooth points, \(\operatorname {pdiv}\) takes the standard convention by definition. We call this the canonical pdiv-derived witness pattern, and it shows up by name in the next few chapters for ReLU, maxPool, and the depthwise input VJP.

The codegen then substitutes the same subgradient/argmax convention every framework already uses. The benefit isn’t a different numerical result — it’s a paper trail that says “yes, 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 correct and totally impractical. For a 512-dimensional hidden layer, each \(J_k\) is a \(512 \times 512\) matrix; multiplying two of them is \(\sim 1.3 \times 10^8\) flops. ResNet-34 has 34 of these layers, 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:

\[ (\nabla \mathcal{L})^T \cdot J_L \cdot J_{L-1} \cdots J_1. \]

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 — never a matrix — and its dimension matches the layer’s input or output shape.

The single step

\[ v \mapsto v^T \cdot J_f \]

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; together with \(f\) itself, it implements one layer’s contribution to backpropagation.

This is what reverse-mode automatic differentiation does, mechanically, for any program. 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 \(\mathbb {R}^{m} \to \mathbb {R}^{n}\);
a correct field — a proof that \(\mathrm{backward}(v)_j = \sum _i v_i \cdot \operatorname {pdiv}\, f\, x\, i\, j\), i.e. that the function really does compute \(v^T J_f\).

In Lean that’s exactly three lines:

structure HasVJP (f : Vec m → Vec n) where
  backward : Vec m → Vec n → Vec m
  correct  : ∀ x dy i, backward x dy i = ∑ j, 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 \)). Unlike autograd, which records each forward op on a tape and replays its derivative at backward time, we ship the closed form: codegen reads off the backward field and emits a single fused StableHLO kernel — no tape, no replay.