Part VII · coda · the continuous cousin

Automatic differentiation

Dual numbers and forward mode, the tape and reverse mode, why a full gradient costs one backward pass, what it costs in memory, and where the answers quietly stop being derivatives.

needs 29, 28~1½ h slides + 2½ h lablab: uv run co lab 30

30 · overview

The route

1 · three kinds of derivative

Symbolic, numerical, automatic; dual numbers and forward mode.

2 · reverse mode

The Wengert tape, one backward sweep, and the cost measured.

3 · memory and higher order

Tape size, checkpointing, JVPs and VJPs, Hessian-vector products.

4 · real data, and lies

An MLP on terrain elevations; non-smooth points, branches, argmax.

getting around

Words with a dotted underline are glossary terms: click one for its definition. G lists every term in the unit, N shows the notes. Optional extra: interactive figures, linked from the slides they support by a small ▸ chip; skip them freely.

section 1

Three kinds of derivative

Only one of them is exact and cheap.

30 · kinds

Symbolic, numerical, automatic

method	what it does	failure mode
symbolic	rewrites the formula	expression swell; no loops or branches
numerical	$\frac{f(x+h) - f(x-h)}{2h}$	truncation error $O(h^2)$ against rounding error $O(\epsilon/h)$; one evaluation pair per input
automatic	chain rule on the operations the program executed	differentiates the program, not the function (section 4)

Numerical derivatives are the test oracle: every op in the lab is gradient-checked against central differences.

forward mode: dual numbers $$\begin{aligned}(a + b\varepsilon)(c + d\varepsilon) &= ac + (ad + bc)\varepsilon \\ g(a + b\varepsilon) &= g(a) + g'(a)\, b\,\varepsilon\end{aligned}$$

with $\varepsilon^2 = 0$.

Seed $x = x_0 + 1\varepsilon$, run the program, read $f'(x_0)$ off the $\varepsilon$ part. For $f: \mathbb{R}^n \to \mathbb{R}$, one pass per input direction: n passes for a gradient.

Lab step 1: twelve operators and five functions. The same code also accepts a Dual whose value is a Dual.

30 · kinds · example

Forward mode, on the example

$f(x, y) = xy + \sin x$ at $(x, y) = (1, 2)$. By hand: $\partial f/\partial x = y + \cos x = 2.5403$, $\ \partial f/\partial y = x = 1$.

pass 1: seed the first input

step	value	$\varepsilon$ part
$x = 1 + 1\varepsilon$	1	1
$y = 2 + 0\varepsilon$	2	0
$xy$	2	$1 \cdot 2 + 1 \cdot 0 = 2$
$\sin x$	0.8415	$\cos 1 \cdot 1 = 0.5403$
$f$	2.8415	2.5403

pass 2: seed the second input

step	value	$\varepsilon$ part
$x = 1 + 0\varepsilon$	1	0
$y = 2 + 1\varepsilon$	2	1
$xy$	2	$0 \cdot 2 + 1 \cdot 1 = 1$
$\sin x$	0.8415	0
$f$	2.8415	1

Two inputs, two passes. The value column is the ordinary evaluation, carried along for free.

explore interactively

section 2

Reverse mode

The whole gradient in one backward pass.

30 · reverse

The Wengert tape

Record every operation with its inputs and local derivatives (the tape). Then sweep backwards (the backward pass), with adjoints $\bar v = \partial f / \partial v$:

$$\bar u \mathrel{+}= \bar v \cdot \frac{\partial v}{\partial u}$$

for each edge $u \to v$, visiting nodes in reverse topological order.

why reverse is cheap

Each operation's backward step costs a small constant times its forward step. So a full gradient costs a constant multiple of one function evaluation, whatever the number of inputs: the cheap gradient principle (Griewank), with a constant of about 4 to 5.

Wengert (1964) for the recorded list; Linnainmaa (1970) for reverse accumulation; backpropagation (Rumelhart, Hinton & Williams 1986) is reverse mode on a network.

Lab then, gradient of $\sum_i \sin(x_i)\, x_{i+1}$:

n	forward (n passes)	reverse (1 sweep)	jax.grad, jit
10	0.000 s	0.0001 s	0.0002 s
100	0.014 s	0.0007 s	0.0004 s
1 000	1.55 s	0.0057 s	0.0021 s

Forward mode grows like $n^2$, reverse like $n$. At a million parameters, forward mode is a million passes and reverse mode is one.

two traps in the sweep

Each node once: a graph with $2^{30}$ paths but 33 nodes must take 33 steps. No recursion: a 20 000-deep graph overflows Python's stack. Lab step 2 tests both.

30 · reverse · example

The tape and its adjoints, on the example

forward: record

node	operation	value	local derivatives
$v_0$	$x$	1
$v_1$	$y$	2
$v_2$	$v_0 \cdot v_1$	2	$\partial/\partial v_0 = v_1 = 2$, $\ \partial/\partial v_1 = v_0 = 1$
$v_3$	$\sin v_0$	0.8415	$\cos v_0 = 0.5403$
$v_4$	$v_2 + v_3$	2.8415	1 and 1

backward: one sweep

visit	adjoint updates
start	$\bar v_4 = 1$
$v_4$	$\bar v_2 \mathrel{+}= 1$, $\ \bar v_3 \mathrel{+}= 1$
$v_3$	$\bar v_0 \mathrel{+}= 1 \times 0.5403$
$v_2$	$\bar v_0 \mathrel{+}= 1 \times 2$, $\ \bar v_1 \mathrel{+}= 1 \times 1$
result	$\bar v_0 = 2.5403$, $\ \bar v_1 = 1$

Both partial derivatives from one backward sweep, where forward mode needed two passes. $x$ feeds two nodes, so its adjoint collects two contributions: that "+=" is the multivariable chain rule.

explore interactively

section 3

Memory and higher order

The tape is the price.

30 · memory

What the tape holds, and checkpointing

Reverse mode must keep every intermediate its backward steps need until the sweep reaches them: the tape memory, mostly activations. For the lab's MLP, the distinct arrays are:

the input batch: batch × in
per layer: W (m × n), matmul output and bias-add output (batch × n each)
per hidden layer: the tanh output (batch × n), which its vjp reuses
the loss's difference array (batch × out), and the scalar loss

MLP [2, 64, 64, 1], batch 64: predicted 233 480 bytes, measured 233 480 bytes. The count is on the next slide.

checkpointing

Keep only some intermediates and recompute the rest during the backward sweep. Keeping every $\sqrt{L}$-th layer of an $L$-layer chain costs $O(\sqrt L)$ memory and about one extra forward pass (Chen et al. 2016). Griewank's revolve gets $O(\log L)$ memory for $O(\log L)$ times the compute, which is optimal for chains.

why memory, not time, is the limit

Activation memory grows with batch × width × depth and must coexist on the device. Parameters are a fixed cost. That's why large-model training is designed around the tape.

30 · memory · example

Counting the tape, by hand

array	shape	numbers
input batch	64 × 2	128
layer 1: $W$, matmul out, bias-add out, tanh out	2×64, 64×64 (three)	$128 + 3 \times 4\,096 = 12\,416$
layer 2: $W$, matmul out, bias-add out, tanh out	64×64 (four)	$4 \times 4\,096 = 16\,384$
layer 3: $W$, matmul out, bias-add out	64×1 (three)	$3 \times 64 = 192$
loss: difference array, scalar	64×1, 1	65
total		29 185

At 8 bytes per float64: $29\,185 \times 8 = \mathbf{233\,480}$ bytes, exactly what the lab measures for the MLP.

what the count tells you

The 64×64 arrays dominate: 7 of them hold 28 672 of the 29 185 numbers. Double the batch and the six activations among them double; the weights don't.

The last layer's output is not followed by tanh (a regression output), and the biases are not saved at all, since adding a bias needs no stored value to differentiate.

30 · higher order

JVPs, VJPs, and Hessian-vector products

primitive	computes	by
JVP	$J v$ (one Jacobian column combination)	forward mode
VJP	$u^\top J$ (one row combination)	reverse mode
HVP	$H v = \frac{d}{d\varepsilon} \nabla f(x + \varepsilon v)$	forward over reverse

on the example

$H = \begin{pmatrix} -\sin x & 1 \\ 1 & 0 \end{pmatrix}$ at $(1, 2)$. With $v = (1, 0)$: $Hv = (-0.8415,\ 1)$, one tape run on dual inputs $(1 + \varepsilon,\ 2)$.

forward over reverse in the lab

Write the reverse-mode tape so its values can be dual numbers. Seed inputs $x_i + v_i\varepsilon$, sweep backwards, and the $\varepsilon$ parts of the gradient are $Hv$. One forward and one backward pass: no Hessian is ever formed.

Lab then, float64, 20 random points, largest absolute difference:

yours	against	difference
reverse-mode gradient	`jax.grad`	8.9e−16
forward-mode Jacobian	`jacfwd`, `jacrev`	8.9e−16
forward-over-reverse HVP	`jax.hessian` @ v	1.6e−15
an MLP batch's gradients	`jax.grad`	3.9e−16

HVPs are what Newton-CG and truncated-Newton methods use, and what "is this point a saddle?" needs, at the cost of about two gradients.

Agreement to machine precision is the expected result when both implementations apply the same chain rule to the same operations in double precision. A difference of 1e−8 would indicate a bug or a different operation (for example a fused kernel), not "numerical noise".

The dual-number trick generalises: nesting dual numbers gives higher derivatives (hyper-dual numbers), and forward-over-reverse is the standard composition in JAX (jax.jvp of jax.grad).

A VJP on the example. For $(x, y) \mapsto (xy,\ \sin x)$ and $u = (2, -1)$, one backward sweep seeded with $u$ gives $u^\top J = (2y - \cos x,\ 2x) = (3.4597,\ 2)$.

The example box. The reference hvp with inputs $(1, 2)$ and direction $(1, 0)$ returns $[-0.8414710, 1.0]$: the first column of $H$. The VJP: $J$ has rows $(y, x)$ and $(\cos x, 0)$, so $u^\top J = 2(2, 1) - (0.5403, 0) = (3.4597, 2)$.

section 4

Real data, and lies

What you get, and what you get wrong.

30 · real data

An MLP on terrain elevations

matplotlib's sample elevation model of the Jacksboro fault region: 344 × 403 heights, standard deviation 162.5 m. Learn elevation from grid coordinates: MLP 2 → 64 → 64 → 1, tanh, batch 64, 60 epochs of SGD. Train on 6 000 points, test on 6 000 others.

model	test RMSE	seconds
predict the mean (baseline)	161.3 m
linear in (x, y)	144.9 m
your MLP	90.4 m	1.8
same in JAX (jit)	90.4 m	2.6

It learns: 90 m against 145 m for a plane, on data it didn't train on (train RMSE 88.6 m).
Identical to JAX: same initial weights and batch order give the same test error, as they should when the gradients agree to 1e−16.
JAX was slower at this size. With 64-unit layers, per-step overhead plausibly outweighs compilation. That's a statement about this size, not about JAX.

The dataset choice: sklearn isn't installed and downloads need permission, so the lab uses real data already shipped with matplotlib.

30 · lies

Where autodiff quietly lies

expression	JAX's gradient	truth
`jnp.maximum(x, 0)` at 0	0.5	undefined (one-sided 0 and 1)
`jax.nn.relu` at 0	0.0	same function, other convention
`jnp.abs` at 0	1.0	undefined (one-sided −1, 1)
`sqrt(x²)` at 0	nan	undefined; 0 · ∞ in the chain rule
`where(x != 1, x, 1.0)` at 1	0.0	the identity: 1
`argmax([x, 0.5])` at 0.2	0.0	0 a.e., a jump at 0.5

the rule: differentiate the program

Autodiff differentiates the program that ran, branch by branch. Where the program branches on a value, or uses a function at a kink, the result is a convention, not a derivative.

Non-smooth points: a subgradient if you're lucky, a library convention otherwise.
Control flow: the derivative of the branch taken, blind to the other branch.
In-place mutation: saved values overwritten before backward.
argmax, sorting, rounding: derivatives exist and are zero almost everywhere, which is useless for learning.

30 · lab

Lab 30 — what you build

about 2½ hours · uv run co lab 30

Dual numbers and forward-mode gradients.
A scalar reverse-mode tape, and Hessian-vector products by forward over reverse.
An array tape with four ops, their vjps, and a gradient checker.
An MLP: loss, gradients, training, and its tape size predicted from the architecture.

uv run co test 30
uv run co then 30     # agreement with JAX · forward vs reverse cost · terrain MLP · where autodiff lies

30 · checkpoint

Done when

checkpoint

All four lab steps are ✓.
Your reverse-mode engine trains an MLP on a real dataset, and you can compute the tape's memory cost in bytes before you run it.
You can say, for a given program, where its autodiff gradient is a convention rather than a derivative.

Reading

Baydin et al., JMLR 18 (2018).
Griewank & Walther, Evaluating Derivatives (2008).
Chen, Xu, Zhang & Guestrin, "Training deep nets with sublinear memory cost", arXiv (2016).
Karpathy, micrograd (2020).

31 where discrete meets differentiable: Sinkhorn, Gumbel, argmin layers, and a learned branching rule.