AI Engineering Foundations 05: A Million Derivatives, Two Passes — Chain Rule and Autodiff from Scratch

Forum topic · 小凯 · 2026-10-01

Summary

Lesson 5 of the 'AI Engineering Foundations' series, based on the 58k-star open-source course ai-engineering-from-scratch by Rohit Ghumare, covers the chain rule and automatic differentiation in a 90-minute build session. The course combines the chain rule (the math) with reverse-mode autodiff (the algorithm) to compute exact gradients for millions of parameters in one forward plus one backward pass—the engine behind PyTorch's backward(), TensorFlow's tape, and JAX's grad. Students hand-build a Value class implementing bookkeeping, local derivatives, topological sort, and backpropagation, then stack it into a 17-parameter MLP that solves XOR in pure Python (loss 4.15 to 0.29). Gradient checking with finite differences validates the implementation to within 3.66×10⁻¹⁰. The review notes three caveats: unreported memory costs of reverse mode (the trade that motivates gradient checkpointing), a pow() signature that fails for non-constant exponents, and uncredited derivation from Karpathy's micrograd. A bonus experiment sweeps the step size h, revealing the U-shaped trade-off between truncation and rounding error.

AI Engineering Foundations 05: A Million Derivatives, Two Passes

This is lesson 5 of the *AI Engineering Foundations* series, working through the open-source course ai-engineering-from-scratch (58k stars on GitHub, by Rohit Ghumare), one lesson at a time. Today's topic: the chain rule and automatic differentiation — a Build-type lesson, 90 minutes. The previous lesson established what a gradient is ("which way to turn the knob"); this one answers the more fundamental question: how do you compute slopes for millions of knobs at all?

Why not brute force

A million parameters means a million derivatives. Hand-derivation is impossible — no argument needed. Numerical differentiation? Nudge each parameter left and right and run two forward passes per parameter: two million passes for a million parameters. Also impossible. The course's solution combines two things: the chain rule provides the math, automatic differentiation provides the algorithm. Together they deliver exact gradients for *all* parameters at the cost of roughly one forward pass plus one backward pass. PyTorch's loss.backward(), TensorFlow's tape, and JAX's grad all run on this machinery — but this lesson builds it from scratch, no frameworks.

The chain rule

The chain rule handles nested functions. For y = sin(x²), there are two layers: inner g(x) = x², outer sin. The derivative is each layer's slope multiplied: cos(x²) × 2x. In plain terms: each link is only responsible for its own slope, and the full chain's slope is the product of all links. A neural network — matrix multiply, add bias, activation, repeat, loss — is just a dozen or so nested functions. Each layer is one link.

Drawing the chain as a graph

The lesson's best move: turn the chain into a computational graph. Each operation is a node; data flows left to right. With x₁ = 2, x₂ = 3, multiply, add one, apply ReLU, and you reach y = 7. Then walk back from the output, multiplying each node's local derivative: ReLU on a positive input has slope 1; addition has slope 1 for both inputs; multiplication has slope x₂ with respect to x₁ and x₁ with respect to x₂. At the end: dy/dx₁ = 3, dy/dx₂ = 2. This process has a name: backpropagation.

There are two ways to traverse the graph, and where you seed it is the dividing line:

Forward mode: seed at the input, push forward — one pass gives derivatives with respect to one input.
Reverse mode: seed at the output, pull back — one pass gives derivatives with respect to *all* inputs.

Neural networks are extreme: millions of inputs (weights), one output (loss). Reverse mode wins decisively — one backward() gets everything. Forward mode isn't useless: for few-output, many-input problems it's the cheaper choice, and exercise 4 has you implement it.

The engineering core: the Value class

The Value class does three things:

1. Bookkeeping ledger: every number is wrapped in an object carrying its value and its gradient. 2. Recording the entry: every operation records its inputs and local derivatives. 3. Settlement: during backward, first perform a topological sort — linearize the graph's dependencies so every node receives gradients only after everything that depends on it has collected — then walk the order in reverse.

One detail worth pausing on: gradients are accumulated by addition, not replaced. If a value feeds two paths, its gradient is the sum of both contributions. Getting this wrong is the most common failure point when hand-building autograd.

Stacking it into a network

Three layers of abstraction on top of Value: a Neuron is tanh(weighted sum + bias), a Layer is a row of neurons, and an MLP chains layers together. Tested against the classic XOR problem: four input pairs, output answers "are the two inputs different?" Using a 2-4-1 architecture — 17 parameters total, pure Python, learning rate 0.05, 100 steps — the loss dropped from 4.15 to 0.29, with all four predictions correct in sign. No frameworks, just this bookkeeping system.

Gradient checking

How do you trust the ledger? Use finite differences as a notary: nudge x right by h, then left by h, run both, approximate the derivative from the difference, and compare against autograd. Two fully independent paths arriving at the same number makes the ledger credible. The course uses h = 10⁻⁷; on (x³ + 2x + 1) passed through tanh, the analytic gradient was 0.15252426, the numerical gradient 0.15252426, difference 3.66×10⁻¹⁰ — four orders of magnitude better than the required threshold of 10⁻⁵.

Fact-check: clean numbers, three caveats

The lesson's numbers are unusually clean: XOR convergence, all chain-rule examples, a PyTorch comparison at 7.0, tanh'(2) = 0.0707, the derivative of x³ at 2 equaling 12.0 — all verified, zero discrepancy, much more presentable than the previous lesson's two broken demos. But three issues are worth flagging:

1. The slogan hides the memory bill. The opening claims autodiff gives all exact gradients "in time comparable to one forward pass." That's half the accounting — the order of magnitude is right, but the constant factor (2–3× compute) is glossed over. What's truly unreported is memory: reverse mode must retain every intermediate value on the forward path to pass gradients back. This is why activation values eat GPU memory when training large models, and why gradient checkpointing exists. Reverse mode trades memory for time — a transaction the lesson's blackboard never mentions.

2. The pow signature. The backward formula n·xⁿ⁻¹ only holds when the exponent is a constant. Passing a Value into the exponent slot throws a TypeError immediately. The failure mode is honest — it crashes rather than silently computing wrong answers — but the limitation isn't documented, and anyone trusting the interface signature will hit the wall.

3. Missing attribution. Step 5 ends with "this is micrograd" — but Karpathy's name and repository link never appear, and further reading omits them. The Value class is a direct replica of micrograd throughout; crediting it with a single word falls short.

A hands-on experiment worth running

Sweep the gradient-check step size h outward from 10⁻⁷. Measured results: at h = 10⁻², error ≈ 3×10⁻⁵; at h = 10⁻⁶, the valley, 3×10⁻¹¹; at h = 10⁻¹², back up to 2×10⁻⁵; at h = 10⁻¹⁵, collapse to 1.4×10⁻² — threshold blown. A U-shaped curve: too large a step, truncation error kills you; too small, floating-point rounding kills you. The course's 10⁻⁷ sits on the gentle right slope — good enough. The theoretical optimum lies near the cube root of machine epsilon, about 6×10⁻⁶.

Exercise 4 is even more worthwhile: implement forward mode with dual numbers via a Dual class and reconcile it against the reverse-mode engine — the same seed, growing from both ends toward the middle.