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 desire, via MLIR. I meant it as a wishlist item, something I expected to 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 are production-quality now, and the thing I wished for is sitting in a GitHub repo, training ResNets on a single GPU, with machine-checked proofs that every gradient is correct.

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, but what you probably haven’t done is open the hood, look at what .backward() is actually computing layer by layer, and verify that the computation is mathematically correct.

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 the same computation: the math, the proof, the code that runs. They agree, and the entire book is about understanding 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 decomposes into some combination of those eight things, and that includes the MLP, the CNN, ResNet, MobileNet, EfficientNet, and the Vision Transformer. So does every architecture in the bestiary at the end, which covers U-Net, YOLO, Mamba, GAN, VAE, diffusion models, 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 by the time I’d done it for enough architectures the pattern became impossible to ignore. The framework kept working without needing new primitives.

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’m going to assume you know what a derivative is, what a matrix is, and roughly what a convolutional layer does, and I’m not going to assume you know Lean, MLIR, formal verification, or compiler theory. If you skip every Lean snippet you still get a complete book on the mechanics of backpropagation. The Lean is there for the reader who wants to check me.

This is not a beginner’s book. If you’re looking for your first introduction to neural networks there are several good ones, including my first book, fast.ai’s free courses, and plenty of good resources all over the web. Come back to this one when you’ve trained something and want to understand what you did.

This is also not a research textbook, and I’m not trying to write Goodfellow or Bishop. The proofs here are for credibility, not for theoretical depth, and if you want the full formal story then Mathlib is open-source and you can extend it. What this book adds is the bridge from the math to the actual 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 in Lean, then you write a proof that the function does what you claim, and if that proof compiles then the claim is correct. You don’t have to trust the author, because the compiler checked it. 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, the Fréchet-derivative API called fderiv, which gives you the chain rule, the product rule, and composition of continuous linear maps. The central claim of this book is that the backward pass equals the true Jacobian, and without that API it would not be feasible to prove. Avigad contributed the first-order calculus in 2019, and Gouëzel stabilized the higher-order \(C^n\) scaffolding around it over 2019–2025 (arXiv:2509.04922).

MLIR and StableHLO stabilized. MLIR is a compiler infrastructure for machine learning, and StableHLO is a portable operation set built on MLIR that describes what neural networks compute. Together they let you write a training step as a sequence of mathematical operations and compile it for CPU, CUDA, or ROCm without writing 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 available, but for years the only supported way to reach it was from Python. PJRT changed that. A backend now ships as an ordinary shared library with a C interface, so a program written in any language can open it with dlopen, hand it a StableHLO module, and get back something that runs on the card. That is what lets a Lean binary train a network on a GPU with no interpreter anywhere in the pipeline. This book supports a second runtime as well, IREE, which compiles every kernel itself and so runs on hardware that 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, most well-understood corner of deep learning, which means we can introduce primitives one at a time in a logical order.

We start with MNIST (digits, \(28 \times 28\) grayscale) and work up through CIFAR-10 (color, \(32 \times 32\)) to Imagenette (real photos, \(224 \times 224\)). The architectures form a natural progression: MLP \(\to \) CNN \(\to \) ResNet \(\to \) MobileNet \(\to \) EfficientNet \(\to \) ViT. Each one adds exactly one new structural primitive to the framework. By the time you reach ViT you’ve seen every primitive there is for modern vision, and you’ve also seen every primitive for modern language models, because transformers are the same architecture in both domains.

The bestiary at the end demonstrates that the same primitives cover detection (YOLO), segmentation (U-Net), generation (diffusion models), self-supervised learning (MAE, CLIP), sequence modeling (Mamba), and multi-modal models (LLaVA). 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 => finSum m (fun i => 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 finSum is a for loop. 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 normalization is a specific three-term formula, the Lean compiler requires me to prove it, and if that proof compiles then the claim is correct. You don’t have to trust me, because the type checker verified it. There is no hidden hand-waving. Every VJP in this book builds with zero sorrys, and that means 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 Stanford’s DAWNBench competition, training ImageNet in 3 hours for $25 on commodity hardware. The lesson from that experience was simple. Most of the apparent difficulty in modern deep learning isn’t fundamental. It’s accumulated folklore, missing tools, and a culture that takes the math on faith.

This book applies that same lesson to the math itself. The folklore says backpropagation is complicated. It isn’t. It’s three rules, five tricks, and a lot of bookkeeping. The tools to verify this claim finally exist, 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 by accident that teenagers are publishing results. The barriers to doing real ML research used to be infrastructure. Compute, data, institutional access. Those barriers have mostly fallen, and what remains is the smaller barrier of understanding what you’re doing at the level where you can 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.