Automatic differentiation: forward mode and reverse mode

Before this lesson: lesson 04 (the chain rule, VJPs, computational graphs and the two sweeps). Linear algebra lesson 15 gives a short first view of JVPs and VJPs; this lesson goes much deeper. Time: about 80 minutes, plus the notebook.

Why this matters #

• Training. loss.backward() in PyTorch and jax.grad in JAX are reverse-mode AD. One call gives the gradient with respect to all parameters, for about 2–3 times the cost of the forward pass.
• Memory. Activation memory during training is the tape of reverse mode. Activation (gradient) checkpointing, used to train large language models, is the memory–time trade-off of section 3.
• Forward mode has its own uses: sensitivities of a simulator with few parameters, and Hessian-vector products (forward mode applied to a gradient, lesson 08).
• Interviews and debugging. "How does autograd work?", "why does training use reverse mode?", "why do I run out of memory?" all have exact answers in terms of the two modes.

1. Predict first #

2. The picture #

Three ways to get a derivative.

Symbolic differentiation (as in SymPy or Mathematica) works on formulas. Formulas can grow very large. Take the logistic map $l_{k+1} = 4\,l_k(1 - l_k)$, starting from $l_0 = x$. The program for $l_4$ has 12 operations (3 per step). Written as one formula, $l_4$ has 38 operations, because each step copies the previous formula twice; SymPy's derivative of that formula, before simplification, has 264 operations. This growth is called expression swell (Baydin et al. 2018, Table 1).

Numerical differentiation treats $f$ as a black box and uses $\big(f(x + \eta e_i) - f(x - \eta e_i)\big)/(2\eta)$. It needs two evaluations per input, and the result has truncation and rounding errors (lesson 07). It is the standard tool to check a gradient, not to train with.

Automatic differentiation does not build a formula. It applies the chain rule of lesson 04 to the primitive operations, while the program runs, at the current numbers. So loops, branches and function calls are no problem: AD differentiates the steps that were actually executed.

Forward mode: values and tangents move together. Choose an input direction $\dot x$. Next to every value $v_i$, carry its tangent $\dot v_i$: the rate at which $v_i$ changes when the input moves along $\dot x$. Each primitive updates the pair (value, tangent) with its local derivative. This is the forward sweep of lesson 04, run during the forward computation.

Dual numbers: the tangent as a second part of the number. Write a number with a tangent as $a + b\,\varepsilon$, where $\varepsilon$ is a symbol with $\varepsilon^2 = 0$ (but $\varepsilon \ne 0$). Multiply two such numbers:

The $\varepsilon$ part is exactly the product rule. If you evaluate a function at $x + 1\varepsilon$, the $\varepsilon$ part of the result is $f'(x)$. The figure runs small programs on dual numbers.

Reverse mode: first record, then walk back. Reverse mode needs the adjoints $\bar v_i = \partial y/\partial v_i$. The adjoint of a node depends on its children, which come later in the program. So reverse mode cannot run during the forward computation. It works in two phases:

1. Forward run. Run the program and record each primitive operation on a tape: which variables it used, and the values needed for its local derivatives.
2. Backward run. Seed $\bar y = 1$. Walk the tape from the end to the start. For each entry $v_i = \varphi(v_j, \dots)$, add $\bar v_i\,\partial v_i/\partial v_j$ to the adjoint of each input $v_j$.

3. Definitions and main results #

Let $F:\R^n\to\R^m$ be computed by a program of primitive operations, as in lesson 04: inputs $v_{1-n},\dots,v_0$, intermediate variables $v_1, v_2, \dots$ in topological order, and outputs $y = (y_1, \dots, y_m)$. Write $J = DF(x)\in\R^{m\times n}$.

• A Jacobian-vector product (JVP) is $Jv$ for a vector $v\in\R^n$ (an input direction, also called a tangent). With $v = e_j$ it is column $j$ of $J$.
• A vector-Jacobian product (VJP) is $u^\top J$ (or $J^\top u$) for a vector $u\in\R^m$ (an output weighting, also called a cotangent). With $u = e_i$ it is row $i$ of $J$.

Forward mode (tangent mode) #

Algorithm (forward mode). Input: $x$ and a tangent $\dot x$.

1. Set $v_{j-n} = x_j$ and $\dot v_{j-n} = \dot x_j$ for the inputs.
2. For each operation $v_i = \varphi_i(\text{parents})$ in program order, compute the value $v_i$ and the tangent $\dot v_i = \sum_{j\in\text{parents}(i)} \frac{\partial \varphi_i}{\partial v_j}\,\dot v_j$.
3. Return $y$ and $\dot y$.

Claim. $\dot y = J\dot x$, a JVP. Proof: this is the forward sweep of lesson 04, applied to each output $y_k$: $\dot y_k = \nabla y_k^\top\dot x$, which is row $k$ of $J$ times $\dot x$. $\blacksquare$

Step 2 needs only the parents of $v_i$, which were just computed. So a tangent can be freed as soon as no later operation needs it: forward mode needs no tape.

Dual numbers #

A dual number is $a + b\varepsilon$ with real $a, b$ and the rule $\varepsilon^2 = 0$. The arithmetic follows from this one rule:

(For the quotient, multiply the claimed result by $b + b'\varepsilon$ and check that you get $a + a'\varepsilon$.)

Theorem (dual numbers give derivatives). For every polynomial $p$, $p(a + b\varepsilon) = p(a) + p'(a)\,b\,\varepsilon$.

Proof.

1. For a power, the binomial theorem gives $(a + b\varepsilon)^k = a^k + k\,a^{k-1}b\,\varepsilon + \binom{k}{2}a^{k-2}b^2\varepsilon^2 + \cdots$.
2. Every term with $\varepsilon^2$ or a higher power is 0. So $(a + b\varepsilon)^k = a^k + k a^{k-1}\,b\,\varepsilon$.
3. A polynomial is a sum of terms $c_k x^k$. Add the results of step 2: $p(a + b\varepsilon) = \sum_k c_k a^k + \big(\sum_k k\,c_k a^{k-1}\big)b\,\varepsilon = p(a) + p'(a)\,b\,\varepsilon$. $\blacksquare$

For other primitives we define $\varphi(a + b\varepsilon) = \varphi(a) + \varphi'(a)\,b\,\varepsilon$, for example $\sin(a + b\varepsilon) = \sin a + (\cos a)\,b\,\varepsilon$. This is consistent with composition: $\varphi(\psi(a + b\varepsilon)) = \varphi(\psi(a)) + \varphi'(\psi(a))\,\psi'(a)\,b\,\varepsilon$, which is the chain rule. So running a whole program on $x + \dot x\varepsilon$ is exactly forward mode, and the $\varepsilon$ part of the output is $J\dot x$.

A dual number is not a small number. There is no step size, so there is no truncation error, unlike finite differences. In code, a dual number is a pair (value, tangent) with overloaded operators (__add__, __mul__, …). This technique is called operator overloading: you run your ordinary code, and the new number type does the calculus (section 5).

Reverse mode (adjoint mode) #

Algorithm (reverse mode). Input: $x$ and a cotangent $u\in\R^m$.

1. Forward run. Evaluate the program. For each operation $v_i = \varphi_i(\text{parents})$, push an entry on the tape: the index $i$, the parents, and the values needed to compute $\partial\varphi_i/\partial v_j$.
2. Seed. Set all adjoints to 0, then $\bar y = u$.
3. Backward run. Pop the tape entries in reverse order. For entry $i$ and each parent $j$: $\bar v_j \mathrel{+}= \bar v_i\,\partial\varphi_i/\partial v_j$.
4. Return the input adjoints $\bar x = J^\top u$.

Claim. The result is the VJP $J^\top u$. Proof: define the scalar $s = u^\top y$. Its gradient is $\nabla s = J^\top u$ (the chain rule for the linear map $y\mapsto u^\top y$). The backward sweep of lesson 04 for the scalar $s$ starts with $\bar y = \partial s/\partial y = u$ and gives $\partial s/\partial x$. The reverse order of the tape is a reverse topological order, so every adjoint is complete before it is used. $\blacksquare$

The += in step 3 is the fan-out rule: a variable used $r$ times gets $r$ contributions. Libraries like micrograd and PyTorch implement this with a Value (or tensor) object that stores its data, its parents and a small backward function; backward() sorts the graph topologically and calls the backward functions in reverse order.

One product, two orders #

For a chain $F = f_L\circ\cdots\circ f_1$ with $J_k = Df_k(x_{k-1})\in\R^{n_k\times n_{k-1}}$, the Jacobian is $J = J_L\cdots J_1$. The two modes compute it in the two natural orders:

Every step is a matrix-vector product, never a matrix-matrix product. In fact the $J_k$ are never built: each primitive comes with a JVP rule ($\dot v \mapsto J_k\dot v$) and a VJP rule ($\bar v\mapsto J_k^\top\bar v$). For a dense layer both rules are one matrix-vector product, the same work as the layer itself.

Cost in time #

Write $\operatorname{ops}(F)$ for the number of operations to evaluate $F$.

Why one pass costs only a constant times $F$. Each primitive has a local derivative that costs a small constant times the primitive itself: for $v = ab$, the tangent $\dot a\,b + a\,\dot b$ costs two multiplications and one addition; for $v = \sin a$, the derivative needs one $\cos$. Summing over all operations gives a constant times $\operatorname{ops}(F)$. Griewank and Walther (2008) prove a bound with $c < 6$; in practice $c$ is about 2 to 3 (as quoted by Baydin et al. 2018).

Consequences.

• The cheap gradient principle. For a scalar function ($m = 1$), one reverse pass gives the whole gradient for at most $c\cdot\operatorname{ops}(F)$, independent of the number of inputs $n$. Forward mode needs $n$ passes. This is why every neural network is trained with reverse mode.
• Few inputs, many outputs ($n \ll m$): forward mode wins. Example: how a simulated field with $10^4$ values responds to 3 physical parameters needs 3 forward passes and no tape.
• Square maps ($n = m$): the pass counts are equal, and forward mode often wins on memory.

Cost in memory, and checkpointing #

In reverse mode, the backward run needs the stored values of the forward run. For a network with $L$ layers, the tape holds the input of every layer (the activations), so memory grows with depth and with batch size. Forward mode keeps only the current values and tangents.

Checkpointing trades memory for time. Number the layer inputs $a_0, \dots, a_{L-1}$ (layer $j$ needs $a_{j-1}$ in the backward run). Choose an interval $k$.

1. Forward run: keep only the checkpoints $a_0, a_k, a_{2k}, \dots$ (about $L/k$ of them).
2. Backward run, segment by segment from the end: recompute the activations of one segment from its checkpoint (at most $k - 1$ values), run the backward steps of that segment, then free them.

Count. The peak memory is about $L/k$ checkpoints plus $k$ values of one segment: $M(k) \approx L/k + k$. The extra work is one recomputation of every non-checkpoint activation: fewer than $L$ layer evaluations, so at most one extra forward pass. $M(k)$ is smallest where its derivative $-L/k^2 + 1$ is 0, at $k = \sqrt L$, with $M \approx 2\sqrt L$ instead of $L$ (Chen et al. 2016).

Full Jacobians and the order of the product #

To build the whole Jacobian, you can run $n$ JVPs (columns) or $m$ VJPs (rows). With dense Jacobians, you can also multiply the matrices $J_L\cdots J_1$ directly. Starting from the input side, step $k$ multiplies $J_k$ by the accumulated $n_{k-1}\times n_0$ block; starting from the output side, the accumulated $n_L\times n_k$ block by $J_k$. The numbers of multiplications are

The prefactors $n_0$ and $n_L$ usually decide: start on the side with the small dimension. When a narrow layer sits in the middle, a mixed order can beat both ends.

AD in practice #

The same two modes appear in every framework, under different names.

Note that backward() adds into .grad. This is the fan-out rule again, applied across calls: that is why a training loop calls optimizer.zero_grad() before each backward run.

Answers to section 1. (1) Forward mode needs $10^6$ passes, one per input direction. Reverse mode needs one pass, because the output is a scalar ($m = 1$); it costs a small constant (about 2–3) times one evaluation of the loss. (2) $\partial\ell/\partial W = \bar z\,a^\top$ has the shape of $W$, $512\times256$, and rank at most 1, because it is an outer product (lesson 04). It has rank exactly 1 when $\bar z\ne 0$ and $a\ne 0$.

4. Worked example #

(a) Both modes on $f(x_1, x_2) = \ln x_1 + x_1x_2 - \sin x_2$ at $(2, 5)$, written as the traces an AD tool produces (Baydin et al. 2018, Tables 2 and 3).

Forward mode, seed $\dot x = e_1 = (1, 0)$. Each line computes the value and the tangent:

Result: $\partial f/\partial x_1 = 5.5$. A second pass with $\dot x = e_2$ gives $\partial f/\partial x_2 = 1.716$.

Reverse mode. The forward run writes the left column of the table above (values only) to the tape. Then the backward run reads the tape from the bottom:

Result: both partial derivatives, $(5.5, 1.716)$, from one backward run. The stored values $v_0 = 5$, $v_{-1} = 2$ and $\cos v_0$ were needed: this is what the tape is for.

(b) A full Jacobian by JVPs and by VJPs. Let $F = f_2\circ f_1:\R^2\to\R^3$ with

1. Forward run: $u = f_1(1, 2) = (2, 1)$ and $F(x) = (3, 2, 1)$.
2. Local Jacobians: $J_1 = \begin{pmatrix} x_2 & x_1 \\ -1 & 1\end{pmatrix} = \begin{pmatrix} 2 & 1 \\ -1 & 1\end{pmatrix}$ and $J_2 = \begin{pmatrix} 1 & 1 \\ u_2 & u_1 \\ 0 & 2u_2\end{pmatrix} = \begin{pmatrix} 1 & 1 \\ 1 & 2 \\ 0 & 2\end{pmatrix}$.
3. Two JVPs (forward mode, one per input):
   • $e_1$: $J_1e_1 = (2, -1)$, then $J_2(2, -1) = (1,\ 2 - 2,\ -2) = (1, 0, -2)$: column 1.
   • $e_2$: $J_1e_2 = (1, 1)$, then $J_2(1, 1) = (2, 3, 2)$: column 2.
4. Three VJPs (reverse mode, one per output):
   • $e_1$: $J_2^\top e_1 = (1, 1)$, then $J_1^\top(1, 1) = (2 - 1,\ 1 + 1) = (1, 2)$: row 1.
   • $e_2$: $J_2^\top e_2 = (1, 2)$, then $J_1^\top(1, 2) = (2 - 2,\ 1 + 2) = (0, 3)$: row 2.
   • $e_3$: $J_2^\top e_3 = (0, 2)$, then $J_1^\top(0, 2) = (-2, 2)$: row 3.
5. Both give $J = \begin{pmatrix} 1 & 2 \\ 0 & 3 \\ -2 & 2\end{pmatrix}$. Here $n = 2 < m = 3$, so forward mode needs fewer passes.

Check one entry directly: $F_2(x) = x_1x_2(x_2 - x_1)$, so $\partial F_2/\partial x_1 = x_2(x_2 - x_1) - x_1x_2 = 2 - 2 = 0$. ✓

(c) Dual numbers for $f(x) = x^2\sin x$ at $x = 1.5$. Seed $x = 1.5 + 1\varepsilon$.

1. $v_1 = x\cdot x = (1.5 + \varepsilon)(1.5 + \varepsilon) = 2.25 + 3\varepsilon$.
2. $v_2 = \sin(1.5 + \varepsilon) = \sin 1.5 + (\cos 1.5)\varepsilon = 0.9975 + 0.0707\varepsilon$.
3. $v_3 = v_1v_2 = 2.25\cdot0.9975 + (3\cdot0.9975 + 2.25\cdot0.0707)\varepsilon = 2.2444 + 3.1516\varepsilon$.

Check: $f'(x) = 2x\sin x + x^2\cos x = 3\cdot0.9975 + 2.25\cdot0.0707 = 3.1516$. ✓

5. In code #

A complete forward-mode engine (dual numbers) and a complete reverse-mode engine (a Value class with a tape), in plain Python. Both reproduce the numbers of section 4. The notebook builds them step by step.

import math

class Dual:
    """a + b·ε with ε² = 0. The ε part (b) carries the derivative."""
    def __init__(self, a, b=0.0):
        self.a, self.b = a, b
    def __add__(self, o):
        o = o if isinstance(o, Dual) else Dual(o)
        return Dual(self.a + o.a, self.b + o.b)
    __radd__ = __add__
    def __sub__(self, o):
        o = o if isinstance(o, Dual) else Dual(o)
        return Dual(self.a - o.a, self.b - o.b)
    def __mul__(self, o):
        o = o if isinstance(o, Dual) else Dual(o)
        return Dual(self.a * o.a, self.b * o.a + self.a * o.b)   # ε² term dropped
    __rmul__ = __mul__

def sin(d): return Dual(math.sin(d.a), math.cos(d.a) * d.b)
def log(d): return Dual(math.log(d.a), d.b / d.a)

def f(x1, x2):                       # ordinary code; works for Duals
    return log(x1) + x1 * x2 - sin(x2)

y = f(Dual(2.0, 1.0), Dual(5.0, 0.0))     # seed: tangent e1
print(round(y.a, 3), round(y.b, 3))       # 11.652 5.5
y = f(Dual(2.0, 0.0), Dual(5.0, 1.0))     # seed: tangent e2
print(round(y.b, 3))                      # 1.716

class Value:
    """A scalar that records how it was made (the tape) and can run the reverse sweep."""
    def __init__(self, data, parents=(), local=()):
        self.data, self.grad = data, 0.0
        self.parents, self.local = parents, local     # local[i] = d self / d parents[i]
    def __add__(self, o): return Value(self.data + o.data, (self, o), (1.0, 1.0))
    def __sub__(self, o): return Value(self.data - o.data, (self, o), (1.0, -1.0))
    def __mul__(self, o): return Value(self.data * o.data, (self, o), (o.data, self.data))
    def backward(self):
        order, seen = [], set()
        def visit(v):                                   # topological order by depth-first search
            if id(v) not in seen:
                seen.add(id(v))
                for p in v.parents: visit(p)
                order.append(v)
        visit(self)
        self.grad = 1.0                                 # seed: d y / d y = 1
        for v in reversed(order):                       # reverse topological order
            for p, d in zip(v.parents, v.local):
                p.grad += v.grad * d                    # add: fan-out sums

def vlog(v): return Value(math.log(v.data), (v,), (1.0 / v.data,))
def vsin(v): return Value(math.sin(v.data), (v,), (math.cos(v.data),))

x1, x2 = Value(2.0), Value(5.0)
y = vlog(x1) + x1 * x2 - vsin(x2)
y.backward()
print(round(y.data, 3), x1.grad, round(x2.grad, 3))    # 11.652 5.5 1.716

The Value objects form the tape: each one keeps its parents and its local derivatives, which use the values of the forward run (o.data, self.data, cos(v.data)).

The same computations in PyTorch, including the full Jacobian of example (b) by 2 JVPs and by 3 VJPs:

import torch
from torch.func import jvp, vjp

x = torch.tensor([2.0, 5.0], dtype=torch.float64, requires_grad=True)
f = lambda x: torch.log(x[0]) + x[0] * x[1] - torch.sin(x[1])
y = f(x)                     # forward run: PyTorch records the graph (the tape)
y.backward()                 # reverse sweep with seed 1
print(x.grad)                # tensor([5.5000, 1.7163], dtype=torch.float64)

v = torch.tensor([1.0, 2.0], dtype=torch.float64)
print(jvp(f, (x.detach(),), (v,))[1])      # forward mode, J v: tensor(8.9327, dtype=torch.float64)

F = lambda x: torch.stack([x[0] * x[1] + x[1] - x[0], x[0] * x[1] * (x[1] - x[0]), (x[1] - x[0]) ** 2])
x0 = torch.tensor([1.0, 2.0], dtype=torch.float64)
cols = [jvp(F, (x0,), (e,))[1] for e in torch.eye(2, dtype=torch.float64)]
_, pull = vjp(F, x0)
rows = [pull(e)[0] for e in torch.eye(3, dtype=torch.float64)]
print(torch.stack(cols, 1).tolist())   # [[1.0, 2.0], [0.0, 3.0], [-2.0, 2.0]]
print(torch.stack(rows, 0).tolist())   # [[1.0, 2.0], [0.0, 3.0], [-2.0, 2.0]]

In JAX the same calls are jax.grad(f)(x), jax.jvp(f, (x,), (v,)) and jax.vjp(f, x); with jax.config.update("jax_enable_x64", True) they print the same numbers.

6. Where it shows up #

• Training neural networks. Every optimizer step is one reverse-mode pass: the cheap gradient principle makes billions of parameters affordable. Lesson 06 writes the backward pass of an MLP by hand; deep learning lesson 03 covers backpropagation in networks.
• Activation memory. The tape is the main memory cost of training large models. Activation checkpointing (torch.utils.checkpoint, jax.checkpoint), recomputation and reversible layers reduce it; see the ML systems track for memory budgets of real models.
• Second derivatives. A Hessian-vector product is forward mode applied to the reverse-mode gradient ("forward-over-reverse") and costs a few gradients (lesson 08). Gradients through solvers and fixed points use implicit differentiation instead of a long tape (lesson 09).
• Probabilistic programming. Hamiltonian Monte Carlo (Stan, NumPyro, PyMC) needs the gradient of a log density at every step; it gets it from reverse-mode AD.
• ML for science. Differentiable simulators (fluid, climate, molecular dynamics) use reverse mode through the solver, which is the adjoint method of computational science; with many time steps, memory forces checkpointing. Forward mode gives the sensitivities of a simulation to a few physical parameters. See the ML for science track.

7. Common mistakes #

• "AD is symbolic differentiation" or "AD is finite differences". It is neither: it applies the chain rule to the executed operations at the current numbers. It builds no formula and uses no step size.
• Using forward mode for a loss with many parameters. It needs one pass per parameter. Use reverse mode for $\R^n\to\R$ and forward mode when the number of inputs is small.
• Forgetting that .grad accumulates. backward() adds into .grad. Without optimizer.zero_grad(), gradients of several steps are summed.
• Keeping the tape alive by accident. total_loss += loss keeps every graph in memory; use total_loss += loss.item() (a plain number) when you only want to log the value.
• Calling backward() on a vector output without a seed. For non-scalar $y$, PyTorch needs y.backward(u): it computes the VJP $J^\top u$, not the full Jacobian.
• Trusting AD at a kink or across a branch. AD differentiates the branch that ran. The code x if x != 0 else 0.0 equals $x$ everywhere, but AD gives derivative 0 at $x = 0$. For ReLU at 0, frameworks return 0 by convention.

8. Check yourself #

1. Compute $(2 + \varepsilon)^3$ with dual numbers, and say which derivative you get.
   Answer

   $(2 + \varepsilon)^3 = 8 + 3\cdot 4\,\varepsilon = 8 + 12\varepsilon$ (all terms with $\varepsilon^2$ vanish). The $\varepsilon$ part is $\frac{d}{dx}x^3 = 3x^2 = 12$ at $x = 2$.
2. A program computes $F:\R^{50}\to\R^{50}$. You need the full Jacobian. How many passes in each mode, and which do you choose?
   Answer

   50 JVPs or 50 VJPs: the same number of passes. Choose forward mode: it needs no tape, so its memory does not grow with the length of the program.
3. Why can reverse mode not compute $\bar v_j$ at the moment $v_j$ is computed in the forward run?
   Answer

   $\bar v_j = \sum_{\text{children}} \bar v_i\,\partial v_i/\partial v_j$ needs the adjoints of the children, which are computed later in the program. So reverse mode must finish the forward run first and keep (on the tape) the values that the local derivatives need.
4. A chain has widths $(n_0, n_1, n_2, n_3) = (3, 100, 100, 10^4)$. Count the multiplications to form the full Jacobian from the input side and from the output side.
   Answer

   Input side: $3\,(100\cdot100 + 10^4\cdot100) = 3.03\cdot10^6$. Output side: $10^4\,(100\cdot100 + 100\cdot3) = 1.03\cdot10^8$. The input side (forward mode) is about 34 times cheaper.
5. A network has $L = 64$ layers. With checkpointing, which interval $k$ minimizes peak memory, and about how many activations are stored at the peak?
   Answer

   $k = \sqrt{64} = 8$. The peak is about $L/k + k = 8 + 8 = 16$ activations (exactly 15 with the count in the figure: 7 earlier checkpoints plus a segment of 8), instead of 64. The extra work is 56 recomputed layers, less than one forward pass.
6. In the Value class, what goes wrong if p.grad += v.grad * d is replaced by p.grad = v.grad * d?
   Answer

   At a fan-out the second contribution overwrites the first. For the example, $x_1$ would get only $0.5$ (the path through $\ln x_1$, processed last) instead of $5.5$, and $x_2$ only $2$ instead of $1.716$.

9. Interview angle #

• "How does autograd work?" — The forward run records a graph (tape) of primitive operations and the values they need. backward() seeds $\bar y = 1$ and visits the nodes in reverse topological order, applying each operation's VJP rule and adding into the adjoints of its inputs.
• "Why do we train with reverse mode?" — The loss is a scalar: one VJP gives the full gradient for about 2–3 forward passes. Forward mode needs one pass per parameter. The price is activation memory.
• "When would you use forward mode?" — Few inputs and many outputs; tight memory; JVPs inside Hessian-vector products (forward-over-reverse).
• "How do you reduce activation memory?" — Checkpointing (store every $k$-th activation, recompute; $k\approx\sqrt L$), smaller batches or micro-batching, mixed precision, reversible layers.
• "Implement a tiny autograd." — A Value class with data, grad, parents and local derivatives; a topological sort; a reverse loop with +=. Then test it against finite differences.

10. Practice #

The notebook for this lesson runs in Colab. It has five exercises, each with a check cell and a folded solution:

1. A Dual class with +, *, sin, exp, log; the derivative of a polynomial program with a loop.
2. A Value class (tape, topological sort, backward()) with +, *, tanh, exp, log; the gradient of $\ln x_1 + x_1x_2 - \sin x_2$.
3. jacobian_fwd and jacobian_rev: full Jacobians from the two engines, compared on a map $\R^2\to\R^3$.
4. chain_flops(dims): the dense costs from the input side and the output side; find the cheaper end.
5. A checkpointing simulation: count stored activations and recomputations for every interval $k$.

11. Watch and read #

• Baydin, Pearlmutter, Radul and Siskind, Automatic differentiation in machine learning: a survey, JMLR 2018: §2 (AD vs symbolic and numerical, expression swell), §3 (forward and reverse mode, costs, Tables 2–3).
• Blondel and Roulet, The Elements of Differentiable Programming, chapter 8: forward and reverse mode on chains and graphs, complexity, implementation, checkpointing (§8.5).
• Bright, Edelman and Johnson, Matrix Calculus (for Machine Learning and Beyond), chapter 8: dual numbers, AD on computational graphs, forward vs reverse mode.
• JAX documentation, The Autodiff Cookbook: grad, jvp, vjp, jacfwd, jacrev and their costs.
• PyTorch documentation, Autograd mechanics: how the graph is recorded, saved tensors, and no_grad.
• Chen, Xu, Zhang and Guestrin, Training Deep Nets with Sublinear Memory Cost, 2016: checkpointing with $O(\sqrt L)$ memory.
• Karpathy, micrograd: a complete reverse-mode engine in about 100 lines.