How this book is organized

A short Foundations chapter comes first: 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 it.

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. The chapters are also the book’s runnable spine: lake run mnist, lake run cifar, lake run imagenette and lake run imagenet train them in this order, one tier per scale, and “Getting started” is where those commands live.

• Chapter 1: MNIST: linear classifier. Chain rule, sum rule, and product rule for partial derivatives. Every backward in the book is a VJP (vector-Jacobian product) composed from these. This is the book’s densest chapter, because it builds the whole pipeline (foundation calculus \(\to \) VJP framework \(\to \) forward/loss/backward/SGD \(\to \) GPU) and then uses it on 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, 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.

• 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, which is 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). Forty-five 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 asks for beyond Part 1’s toolbox: none for VGG or AlphaZero, one for Mamba, never more than three. The entries are grouped by what they are for — vision classifiers, detection, segmentation, language, generation, physics, signal processing — and eleven of the groups carry a demo trained on a dataset from that field; a closing section on game theory trains three games — blackjack, Pong, tic-tac-toe — against their exact solutions, with the AlphaGo family as its catalogue.

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 is the size of what each chapter proves about gradients. Each one proves a handful of VJP contracts, which are machine-checked theorems that a layer’s backward function computes the right gradient. Formally that means \(\texttt{backward}\, x\, dy\, i = \sum _j \operatorname {pdiv}f\, x\, i\, j \cdot dy_j\), so the backward equals that layer’s \(\operatorname {pdiv}\)-Jacobian contracted against the incoming cotangent, which the Foundations chapter builds up from scratch. A contract is that theorem with the weights of \(f\) left as variables, so it holds at every weight setting; there is one per function, where a function is a layer, an operator, a block, or a whole network at a stated depth (a network at two depths, or per example and at batch statistics, is two functions). A witness is the same theorem at one fixed set of weights: a hand-built tiny instance, a live-shaped one at \(224 \times 224\), or the weights a training run produced. It shows the contract’s hypotheses can be met, and is kept in its own column rather than counted as a contract. A step certificate is a theorem about the program rather than the math alone: that the train step the book emits denotes the certified gradient step (a fold for the small nets, a step tie at full width), that a hand-written backward chain is the VJP it was written to be, that a data-parallel render’s step is one device’s step on the global batch, that a gradient node is the loss’s derivative in its parameter, or that one SGD step decreases the loss. A theorem that only restates a counted one at fixed dimensions or over the committed spec, and a smooth-point form that a later global theorem superseded, is not counted again. Across Chapters 1–9 there are 44 contracts, 9 witnesses and 37 step certificates. Even Chapter 1’s toolkit is used straight away by a proven model, so no chapter is pure scaffolding.

That’s the ladder at a glance. Every chapter’s theorems section opens with its dependency graph, every box a Lean theorem and a link, and Figure 1.1 in Chapter 1 follows the five chapter networks’ certificates back to the Chapter 1 calculus, chapter by chapter. The per-chapter detail lives where the work does, because each chapter names its contracts as it proves them. Appendix C covers how those claims get checked, through the audits, the bridge theorems, and the hardware oracles. 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 the foundation pieces you need is explicit. We get it for free from doing the proofs in Lean. Chapter 1 is the one chapter every target reads in full: the foundation rules, the dense layer and the loss. 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, all of them Chapter 1’s.

Target: CNN, with or without BN (Ch 3–4). Ch 1 + Ch 2–4. CNN adds Conv2d and MaxPool VJPs proved from this chapter’s foundation, and 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, and its bigger brother ResNet-50 is the standard benchmark (the GOAT) for production image classification. If your destination is a real image classifier that works, Ch 2–5 is a complete reading list and Chs 6–9 are strictly skippable.

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

Target: EfficientNet (Ch 7). MobileNet subset \(+\) Ch 7. Adds the SE block VJP, which is squeeze-and-excitation as an elementwise product over global pool and 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, which proves the matrix-level machinery — matmul Jacobians and VJPs in both factor positions, transpose, scalar-scale, row-wise lifting, and 3D extensions — in §9.2, just before the attention proofs that use it.

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 postdates the first book and will be genuinely new for you, including LayerNorm, GELU, and multi-head attention. 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 *HasVJP theorems and the \uses{} annotations that pin down the backward pass.

• Chapter 1 is the new material for you. Its first half — the calculus every later chapter composes — is what the first book didn’t have. It’s 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 paired Lean 4 \(\to \) JAX and Lean 4 \(\to \) PJRT trainers, 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 \) XLA \(\to \) GPU. That mostly doesn’t affect the math. It changes where the gradients are 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.