Verified Deep Learning with Lean 4

How this book is organized

Getting started comes next: four install steps, then one command per tier. A short Foundations chapter follows it: what a derivative is in Mathlib, why the book tracks vector–Jacobian products rather than Jacobians, and what training adds to inference. Everything after assumes both.

Part 1: The framework (Chapters 1–9). One new primitive per chapter, proved correct as we go. The first half pins down the composition rules (chain, additive fan-in, product) on dense layers, convolution, and normalization. The second half climbs the ladder to ViT. Each chapter opens by running its network and closes with a recap, and between them are the theorems, the example on real data, and the MLIR. From Chapter 5 on, each chapter also carries an ImageNet recipe, trained once through the JAX reference and once through the verified trainer, and the two scored against each other.

  • Chapter 1: MNIST: linear classifier. Chain rule, sum rule, and product rule for partial derivatives, the VJP framework, and the whole forward / loss / backward / SGD / GPU pipeline, used straight away on the linear classifier. The densest chapter in the book. Everything after is adding a layer.

  • Chapter 2: MNIST: 1D MLP. ReLU and hidden layers on top of Chapter 1’s dense layer, which gives us the first multi-layer backward pass.

  • Chapter 3: MNIST: 2D CNN. Convolution and max pooling. The backward pass of conv is another conv with a reversed kernel.

  • Chapter 4: CIFAR with BatchNorm. The first dense Jacobian with a clean closed form. Three terms, one cancellation, and you never have to derive it again.

  • Chapter 5: ResNet-34. Residual connections. Gradients from two paths add at the input. The ImageNet recipe is assembled and ablated here, and ResNet-50 trains on it.

  • Chapter 6: MobileNetV2. Depthwise convolution. Same as regular conv, but diagonal in the channel dimension. MobileNetV4 is the side quest.

  • Chapter 7: EfficientNet. Squeeze-and-excitation. The product rule: main path times gate, gradients from both paths add.

  • Chapter 8: ConvNeXt. LayerNorm, which is BatchNorm in disguise. GELU, layer scale, and the \(7 \times 7\) depthwise kernel. ConvNeXt-S and B are the side quest.

  • Chapter 9: Vision Transformer. Self-attention. Softmax in the middle of the network, three-way fan-in at the input. The capstone. ViT-S and B are the side quest.

Part 2: The bestiary (Chapters 10–13). Forty-six further architectures, from U-Net to Mamba to Stable Diffusion, each written as a NetSpec in about fifteen lines of Lean and led by how many new primitives it needs beyond Part 1: none for VGG or AlphaZero, one for Mamba, never more than three. Chapter 10 is the kit, the primitives the bestiary adds. The entries are then grouped by what they are for, from vision classifiers through detection, segmentation, language, generation, physics and signal processing, and twelve of the groups carry a demo trained on data from that field: recognition in Chapter 11, everything else in Chapter 12. A closing chapter on game theory trains Pong, Leduc poker and tic-tac-toe against their exact solutions, with the AlphaGo family and Deep CFR as its catalogue.

Appendices. The rest of the toolchain and the six GPU platforms (A), how every claim is checked, with the tally of what each chapter proves (B), the ImageNet tier (C), and data availability (D).

Roadmap: skip to your target architecture

Every theorem and definition in the book carries a \uses{} annotation, so the path from “which architecture do I care about?” back to the pieces you need is explicit. It comes free from doing the proofs in Lean, and Table 1 reads it off. Chapter 1 is the one chapter every target reads in full: the foundation rules, the dense layer and the loss.

ResNet-34 is a great stopping point. It was state of the art for years and still routinely gets 90% on Imagenette, and its bigger brother ResNet-50 is the standard benchmark for production image classification. If your destination is a real image classifier that works, Chapters 1–5 are the whole reading list and Chapters 6–9 are optional.

Table 1 What each target architecture needs, read off the \uses graph. Figure B-1 in Appendix B draws the same graph at the scale of chapters.

Target

Read

What it adds

MLP

Ch 1, 2

ReLU on Chapter 1’s dense layer; the first multi-layer backward pass

CNN, with or without BN

Ch 1–4

the Conv2d and MaxPool VJPs; BN’s inverse-stddev smoothness chain

ResNet-34

Ch 1–5

residual fan-in, global average pooling, the stride-2 convolution and the stem pool; one batch-lifting lemma borrowed from §9.2.1 and taken as read

MobileNetV2

Ch 1–4, 6

depthwise conv, the Chapter 3 recipe with one fewer summation level; nothing from Chapter 5

EfficientNet

Ch 1–4, 6, 7, and Chapter 5’s residual VJP

squeeze-and-excitation as an elementwise product over a global pool and \(1\times 1\) convs

ConvNeXt

Ch 1–6, 8

GELU and LayerNorm, both from Chapter 1’s rules; the \(7 \times 7\) depthwise kernel is Chapter 6’s; nothing from Chapter 7

Attention / ViT

Ch 1, 4’s inverse-stddev chain, 8’s GELU and LayerNorm, 9

the matrix-level machinery of §9.2.1 and the attention proofs; no convolution chapter on its chain

If you’re unsure: read Chapter 1 carefully, pick a target that fits your interest, and follow its row. You can always extend the list later as new architectures catch your eye.

For readers of the first book

If you’re coming from Convolutional Neural Networks with Swift for TensorFlow, you know the architectures of Chapters 2–7 and have trained them. Chapters 8 and 9 will be new even on the forward-pass side. I shipped the final draft of the first book in July 2020. In late August the grapevine started whispering that something attention-shaped was coming, and the ViT paper landed on arXiv that October. Everything here that looks like a transformer postdates the first book: LayerNorm, GELU, multi-head attention. Otherwise this is the same architectures, with the backward pass now formally checked.

What changes for you:

  • Chapter 1 is the new material. Its first half, the calculus every later chapter composes, is what the first book didn’t have. If you read one chapter thoroughly, make it this one.

  • The architecture introductions in Chapters 2–7 will feel redundant. Head for the *HasVJP theorems and the \uses{} annotations that pin down the backward pass.

  • The training recipe is modern. Adam, cosine decay, warmup, weight decay, label smoothing and augmentation, assembled and ablated in Chapter 5 and run as the common Imagenette baseline; the ImageNet recipes keep each paper’s own optimizer.

  • The codegen is new. Lean 4 \(\to \) StableHLO \(\to \) XLA \(\to \) GPU replaces Swift for TensorFlow. The math is unchanged; the gradients are now computed at codegen time in Lean rather than at runtime by a framework.

Practical reading style: work through Chapter 1 slowly, then skim the rest at the pace of “forward pass I know \(\to \) new *HasVJP theorem \(\to \) see what theorems it cites \(\to \) move on.” The dependency DAG is your friend.