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. 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 hardest 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). 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 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. Across Chapters 1–9 there are 40 contracts, one per layer, operator, and whole-network architecture. There are also 8 witnesses, which are concrete end-to-end instances 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.

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, 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 itself is short enough that the right answer for every target is “just read all of it” (11 theorems and 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, and reindex rules.

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. This is the only chapter that needs matrix-level machinery, meaning matmul Jacobians and VJPs in both factor positions, transpose, scalar-scale, row-wise lifting, and 3D extensions. They’re gathered in §9.2 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 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 *_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. 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 *_has_vjp theorem \(\to \) see what theorems it cites \(\to \) move on.” The dependency DAG is your friend.