Verified Deep Learning with Lean 4

How this book is organized

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.

  • Chapter 1: MNIST: linear classifier — Chain rule, sum rule, product rule for partial derivatives; every backward in the book is a VJP (vector-Jacobian product) composed from these. The book’s hardest chapter — it builds the whole pipeline (foundation calculus \(\to \) VJP framework \(\to \) forward/loss/backward/SGD \(\to \) GPU), used by the MNIST linear classifier, proven and trained. Everything after is just adding a layer.

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

  • Chapter 3: MNIST: 2D CNN — Convolution, 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.

  • Chapter 6: MobileNetV2 — Depthwise convolution. Same as regular conv, but diagonal in the channel dimension.

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

  • Chapter 8: ConvNeXt — LayerNorm, Batchnorm in disguise. GELU activation function. Data augmentation strategies.

  • Chapter 9: Vision Transformer — Self-attention. Softmax in the middle of the network, three-way fan-in at the input. The capstone.

Part 2: The bestiary (Chapter 10). A catalog of \(\sim \)37 additional architectures, from U-Net to Mamba to diffusion models, each decomposed into the framework’s primitives. Organized by the question you’re asking: “how do I do detection?”, “how do I do generation?”, etc.

Appendices. Data availability (A), toolchain setup (B), and the complete proof framework as a reference (C).

Theorem and definition budget per chapter

If you’re planning how much of the book to read, here’s the size of what each chapter proves about gradients. Each proves a handful of VJP contracts — machine-checked theorems that a layer’s backward function computes the right gradient (formally \(\texttt{backward}\, x\, dy\, i = \sum _j \operatorname {pdiv}f\, x\, i\, j \cdot dy_j\): the backward equals that layer’s \(\operatorname {pdiv}\)-Jacobian contracted against the incoming cotangent, which the Foundations chapter builds up from scratch). Across Chapters 1–9 there are 40 contracts — one per layer, operator, and whole-network architecture — plus 8 witnesses: concrete end-to-end instances (a specific net at specific weights), kept in their own column as sanity checks rather than citable contracts. Even Chapter 1’s toolkit is used straight away by a proven model, so no chapter is pure scaffolding.

Chapter

Contracts

Witnesses

New primitive / result

1 (Foundation)

1

mnistLinear: the whole classifier

2 (MNIST: 1D)

2

1

ReLU, MLP

3 (MNIST: 2D)

3

4

conv2d, maxPool2, whole-net MNIST CNN

4 (CIFAR: BN)

4

1

per-channel BN (flat \(+\) NCHW), CIFAR whole-nets \(\pm \)BN

5 (ResNet-34)

9

1

residual (\(+\)proj), GAP, stride-2 conv, block/stage, whole-net cnn

6 (MobileNetV2)

3

1

depthwise (stride 1\(+\)2), whole-net mobilenetv2

7 (EfficientNet)

5

swish, SE, whole-net \(\times \)2, full B0 (16 MBConv blocks)

8 (ConvNeXt)

7

GELU, LayerNorm, layer-scale, whole-net \(\times \)2, full ConvNeXt-T

9 (ViT)

6

MHSA, block, depth-2 (\(\times \)2 LN forms), every depth \(k\), vit_full

 

40

8

total 48

That’s the ladder at a glance; Figure C.1 in Appendix C draws it as three architecture spines, every node a Lean theorem, each tracing back to the Chapter 1 calculus. The per-chapter detail lives where the work does — each chapter names its contracts as it proves them — and how the claims are checked (the audits, the bridge theorems, the hardware oracles) is Appendix C’s subject. It reads most naturally once you’ve been through the Foundations chapter, so it’s safe to skip on a first pass.

Roadmap: skip to your target architecture

Every theorem and definition in the book carries a \uses{} annotation, so the dependency graph from “which target architecture do I care about?” back to which foundation pieces you need is explicit — we get it for free from doing the proofs in Lean. Chapter 1 itself is short enough that the right answer for every target is “just read all of it” (11 theorems + 1 definition, mostly chain rule and fan-in variants). The interesting subsetting decision is which later chapters to read, and whether you need Chapter 9’s matrix-machinery section.

Target: MLP (Ch 2). Ch 1 + Ch 2. The MLP backward pass uses chain rule, additive fan-in, identity, and the Dense Jacobians (which are themselves built from this chapter’s finite-sum + const + reindex).

Target: CNN, with or without BN (Ch 34). Ch 1 + Ch 24. CNN adds Conv2d and MaxPool VJPs proved from this chapter’s foundation; BN adds the inverse-stddev smoothness chain.

Target: ResNet-34 (Ch 5). Ch 1–5. Adds one theorem on top of CNN+BN: the residual additive-fan-in VJP. This is a great stopping point. ResNet-34 was state of the art for years and still routinely gets 90% on Imagenette; its bigger brother ResNet-50 is the standard benchmark (🐐) for production image classification on consumer hardware. If your destination is a real image classifier that works, Ch 2–5 is a complete reading list; Chs 69 are strictly skippable.

Target: MobileNetV2 (Ch 6). ResNet subset \(+\) Ch 6. Adds depthwise conv VJPs (same recipe as Ch 3 with one fewer summation level).

Target: EfficientNet (Ch 7). MobileNet subset \(+\) Ch 7. Adds the SE block VJP (squeeze-and-excitation as elementwise product over global pool + 1\(\times \)1 convs).

Target: ConvNeXt (Ch 8). EfficientNet subset \(+\) Ch 8. Adds GELU and LayerNorm as new operator primitives; both proved from Ch 1 foundation rules.

Target: Attention / ViT (Ch 9). Everything above plus Ch 9. This is the only chapter that needs matrix-level machinery (matmul Jacobians and VJPs in both factor positions, transpose, scalar-scale, row-wise lifting, 3D extensions); they’re gathered in §9.1 at the start of Ch 9, just before the attention proofs that use them. Every earlier chapter routes around the matrix kit entirely.

If you’re unsure: read Ch 1 carefully (it’s short), pick an architecture target that fits your interest, follow the chapter chain to it. You can always extend the read list later as new architectures catch your eye.

For readers of the first book

If you’re coming to this book from Convolutional Neural Networks with Swift for TensorFlow, you already know most of the architectures and the forward-pass intuition: VGG, ResNet, MobileNet, and EfficientNet are all familiar, and you’ve built and trained them at least once.

Chapters 8 and 9 will be new material 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 in this book that looks like a transformer — LayerNorm, GELU, multi-head attention — postdates the first book and will be genuinely new for you. Otherwise, this book is the same architectures brought under the additional scrutiny of a formal backward pass.

Concretely, that means:

  • The architecture introductions in chapters 2-7 will feel redundant. The primary sequence is a trimmed version of the first-book progression, kept so the chapters stay consistent. Head straight for the *_has_vjp theorems and the \uses{} annotations that pin down the backward pass.

  • Chapter 1 is the new material for you. The tensor-calculus chapter is what the first book didn’t have — where the gradients get formalized. If you only read one chapter thoroughly, make it this one.

  • Training recipe upgraded. The first book’s codegen trained with SGD (momentum 0.9, learning rate 0.002) and no other tricks — one recipe held constant across architectures for consistency. This book adopts the modern recipe — Adam, cosine decay, linear warmup, weight decay, label smoothing, and basic data augmentation are how this book’s Imagenette results were achieved. That full recipe is assembled and ablated in Chapter 5 (ResNet) and ran against each architecture after as a common baseline. Later chapters add full ImageNet recipes, trained through a Lean 4 \(\to \) JAX bridge, with modern data augmentations.

  • The codegen is new. The first book’s Swift-for-TensorFlow pipeline is replaced with Lean 4 \(\to \) StableHLO MLIR \(\to \) IREE \(\to \) GPU. Mostly doesn’t affect the math; it changes where the gradients are computed (at codegen time in Lean, not 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 *_has_vjp theorem \(\to \) see what theorems it cites \(\to \) move on.” The dependency DAG is your friend.