Verified Deep Learning with Lean 4

Introduction

On the next-to-last page of my first book, I wrote that I dreamed of a future where we could take our code and compile it for whatever backend we wanted, via MLIR. I meant it as a wish, something that would happen eventually but not soon enough to matter for that work.

It turned out “eventually” was about five years.

This is the book I wanted to write back then and couldn’t. The tools weren’t ready. Lean 4 was still in development, MLIR was a research project, and driving a GPU still meant going through Python. All three have shipped. The thing I wished for now sits in a GitHub repo: ten years of networks, ResNet through the Vision Transformer, trained on ImageNet on a consumer GPU, with every gradient checked by Lean.

What this book is about

This book teaches you how backpropagation actually works.

You already know the hand-wave version. Compute gradients, update weights, repeat. If you’ve trained a model you’ve called .backward() a thousand times. What you probably haven’t done is open the hood, look at what it computes layer by layer, and check that the math is right.

That’s what we’re going to do. We’ll take every layer in the modern stack: dense, ReLU, convolution, pooling, batch normalization, residual connections, depthwise convolution, squeeze-and-excitation, layer normalization, and self-attention. For each one we will:

  1. Derive the backward pass from the forward pass, using nothing but the chain rule and some index algebra.

  2. State the result formally in Lean 4, a programming language that doubles as a theorem prover.

  3. Show the MLIR that compiles to GPU bytecode, and verify it matches the math.

  4. Train the model on real data, end-to-end, on your hardware.

Three views of one computation: the math, the proof, and the code that runs. They agree, and the book is about why.

The thesis

Backpropagation is one operation (the vector-Jacobian product), composed using three rules:

  1. The chain rule. If you compose two functions, you compose their backward passes.

  2. Additive fan-in. If two paths feed into a sum, like a residual connection, their gradients add.

  3. Multiplicative fan-in. If two paths feed into a product, like squeeze-and-excitation or attention, the product rule gives you two backward contributions that add.

And five structural tricks for layers whose Jacobians are dense but exploitably structured:

  1. Diagonal. Elementwise activations like ReLU, GELU, and Swish have a diagonal Jacobian, so the backward pass is one multiply.

  2. Sparse Toeplitz. Convolutions have a block-sparse Jacobian with a sliding-window structure, and the backward pass is another convolution with a reversed kernel.

  3. Binary selection. Max pooling routes the gradient to whichever input was the max.

  4. Rank-1 correction to diagonal. Softmax, batch norm, and layer norm all have dense Jacobians with a clean closed form, because each one is a diagonal plus a rank-1 outer product.

  5. Outer product reductions. For dense layers and matmul, the weight gradient is the outer product of the input with the output gradient.

That’s it. Three rules, five tricks. Every architecture in this book is some combination of those eight things: the MLP, the CNN, ResNet, MobileNet, EfficientNet, ConvNeXt, and the Vision Transformer. So is everything in the bestiary of Chapters 10–13: U-Net, YOLO, Mamba, GAN, VAE, diffusion, CLIP, and GPT-style decoders. There is no ninth.

I didn’t set out to prove this. I set out to train a ResNet with verified gradients, and after enough architectures the pattern was impossible to ignore. The framework kept working. It never needed a new primitive.

Who this book is for

This book assumes you’ve trained a neural network. Not a fancy one. If you’ve fine-tuned a pretrained model, debugged a NaN loss, or watched a model overfit and wondered why, you’re the audience.

I’ll assume you know what a derivative is, what a matrix is, and roughly what a convolution does. I won’t assume you know Lean, MLIR, formal verification, or compilers. If you skip every Lean snippet you still get a complete book on how backpropagation works. The Lean is there for the reader who wants to check me.

This is not a beginner’s book. If you want a first introduction to neural networks there are good ones: my first book, fast.ai’s free courses, and plenty more on the web. Come back to this one when you’ve trained something and want to understand what you did.

It’s also not a research textbook. I’m not trying to write Goodfellow or Bishop. The proofs are here for credibility, not depth, and if you want the full formal story, Mathlib is open source and waiting. What this book adds is the bridge from the math to the GPU code, and the assurance that they agree.

If you’re me from ten years ago, interested in ML, willing to put in the work, and not sure where the math stops and the folklore begins, this is the book I wish I’d had.

Why now

Three things had to happen before this book was possible, and all three happened in the last few years.

Lean 4 reached maturity. Lean is a programming language that is also a theorem prover. You write a function, then a proof that it does what you claim, and the compiler checks the proof. Lean 4 came out in 2023 and is the version that made this practical for non-specialists, with a real package manager, real tooling, a growing math library called Mathlib, and an active community.

One piece of Mathlib unblocked this book in particular: its derivative API, fderiv, which supplies the chain rule and the product rule. The central claim of this book is that the backward pass equals the true Jacobian, and without that API it would not have been provable. Avigad contributed the first-order calculus in 2019, and Gouëzel built the higher-order scaffolding around it over 2019–2025 (arXiv:2509.04922).

MLIR and StableHLO stabilized. MLIR is compiler infrastructure for machine learning, and StableHLO is a portable set of operations built on it that says what a network computes. Together they let you write a training step as math and compile it for CPU, CUDA, or ROCm without any backend-specific code. This is the layer that turns math on paper into code on the GPU.

XLA got a stable C API. XLA is the compiler behind JAX and TensorFlow, and it produces some of the fastest GPU kernels there are, but for years the only supported way in was Python. PJRT changed that. A backend now ships as an ordinary shared library with a C interface. Any language can open it, hand it a StableHLO module, and get back something that runs on the card. That is what lets a Lean binary train on a GPU with no interpreter anywhere in the pipeline. The book supports a second runtime too, IREE, which compiles every kernel itself and so runs on hardware no vendor library covers.

Put the three together and one person can go \(\mathbb {R} \twoheadrightarrow \text{RAM}\). You write the network in a language that can state what it’s supposed to compute. You compile that to GPU code through a standard IR. You train it on a card you can buy. At the end you have a proof that the gradients were right. That’s the pipeline this book is built around.

Why image recognition

Same reason as the first book. It’s the oldest, best-understood corner of deep learning, so the primitives can arrive one at a time in a sensible order.

We start with MNIST (digits, \(28 \times 28\) grayscale) and work up through CIFAR-10 (color, \(32 \times 32\)) to Imagenette and then full ImageNet-1k (real photos, \(224 \times 224\)). The architectures form a natural progression: MLP \(\to \) CNN \(\to \) ResNet \(\to \) MobileNet \(\to \) EfficientNet \(\to \) ConvNeXt \(\to \) ViT. Each one adds exactly one new primitive. By ViT you’ve seen every primitive in modern vision, and every primitive in modern language models too, because a transformer is the same architecture in both.

The bestiary of Chapters 10–13 shows the same primitives covering detection (YOLO), segmentation (U-Net), generation (diffusion models), vision–language models (CLIP, LLaVA), and sequence modeling (Mamba). Image recognition is the on-ramp. The whole modern stack is the destination.

Why Lean

I considered several options for the formal half of the book, including Rocq, Agda, Isabelle, and just doing the proofs on paper in LaTeX. I chose Lean for a practical reason. It reads like code.

If you’ve trained a model in PyTorch, you can read: 

def dense {m n : Nat} (W : Mat m n) (b : Vec n) (x : Vec m) : Vec n :=
  fun j => (∑ i : Fin m, x i * W i j) + b j

and understand it. It’s a function with typed arguments. {m n : Nat} are the dimensions, fun j => is a lambda, and ∑ i : Fin m is a for loop that adds up. You don’t need to learn a new paradigm. You need maybe twenty keywords and some notation, and then Lean reads like any other typed functional language.

The other reason is that Lean enforces honesty. If I claim the backward pass of batch norm is a specific three-term formula, Lean makes me prove it, and if the proof compiles the claim is correct. You don’t have to trust me. There is no hidden hand-waving. Every VJP in this book builds with zero sorrys: dense, convolution, batch norm, residual, attention, all of them. The compiler is the referee.

Let’s go

In 2018 I was on a small team that won the CIFAR-10 track of Stanford’s DAWNBench competition and trained ImageNet in 3 hours for $25 on commodity hardware. What I took from it was that you didn’t need powerful tools to change the world, only a firm understanding of what was going on under the hood.

The math doesn’t need a PhD. The received wisdom says backpropagation is complicated. It isn’t. It’s three rules, five tricks, and a lot of bookkeeping, and the bookkeeping is the part people skip, so the math gets taken on faith. This book hands the bookkeeping to the computer, not as a calculator but as a referee: every rule and every trick is a theorem Lean checks, and what’s left for you is bite-sized. Nobody has to tell you where the math stops and the folklore begins. The line is drawn, and you’re holding the result.

Three of my favorite citations of the first book are from high school students, and I don’t think it’s an accident that teenagers are publishing results. The barriers to real ML research used to be infrastructure: compute, data, institutional access. Those have mostly fallen. What’s left is the smaller barrier of understanding what you’re doing well enough to verify it, not just run it.

This book is for the person who wants to cross that barrier. It might be you. It might be the next kid who picks it up in a library somewhere and decides to see how far they can go.

Let’s find out.