Week 11: Automatic Differentiation

Why AD Exists#

Every optimiser in this course needs gradients. How do we get them?

• Numerical differentiation: $f'(x) \approx (f(x+h) - f(x))/h$. Simple but $O(d)$ evaluations per gradient and suffers from truncation vs roundoff trade-off. Unacceptable for $d = 10^7$ parameters.
• Symbolic differentiation: Apply calculus rules to the closed-form expression of $f$. Produces exact derivatives but suffers from expression swell — the derivative expression can be exponentially larger than the original function. Impractical for programs with loops and conditionals.
• Automatic (algorithmic) differentiation (AD): Compute derivatives by applying the chain rule to the sequence of elementary operations that the program executes. AD evaluates the derivative at a specific numeric input, producing exact (to machine precision) derivatives at a cost proportional to evaluating the function itself. This is what PyTorch, JAX, and TensorFlow do.

AD works because every computer program — no matter how complex — is a composition of elementary operations ($+, \times, \sin, \exp$, etc.) whose derivatives are known. The chain rule threads through this computational graph.

Forward Mode#

Forward-mode AD propagates derivatives from inputs to outputs. Associate each variable $v$ with a dual number $\dot{v}$ (its derivative with respect to a chosen input direction). For each elementary operation $w = \phi(u_1, \ldots, u_k)$, compute:

$\dot{w} = \sum_{j=1}^k \frac{\partial \phi}{\partial u_j} \dot{u}_j$

Initialise $\dot{x} = 1$ for the input of interest and $\dot{y} = 0$ for all others.

Worked example — Dual numbers for $\sin(x^2)$ at $x = 2$:

The elementary operations are $v_1 = x^2$, $v_2 = \sin(v_1)$. With $\dot{x} = 1$:

• $v_1 = 2^2 = 4$, $\dot{v}_1 = 2x \cdot \dot{x} = 2 \cdot 2 \cdot 1 = 4$
• $v_2 = \sin(4) \approx -0.7568$, $\dot{v}_2 = \cos(v_1) \cdot \dot{v}_1 = \cos(4) \cdot 4 \approx -2.6146$

Analytic check: $\frac{d}{dx}\sin(x^2) = 2x \cos(x^2)$, at $x=2$: $4 \cos(4) \approx -2.6146$. Matches.

Cost of forward mode. For $f: \mathbb{R}^d \to \mathbb{R}^p$, computing the full Jacobian requires $d$ forward passes (one per input dimension). Forward mode is efficient when $d \ll p$ (few inputs, many outputs).

Reverse Mode#

Reverse-mode AD (a.k.a. backpropagation) propagates derivatives from outputs back to inputs. For a scalar output $f(x)$, this computes the full gradient $\nabla f(x)$ in a single backward pass.

The algorithm works in two phases:

1. Forward pass: Evaluate the function, recording the computational graph (the "tape") — store every intermediate value and the operations that produced them.
2. Backward pass: Traverse the graph in reverse topological order. For each operation $w = \phi(u_1, \ldots, u_k)$, having already computed $\bar{w} = \partial f / \partial w$, compute the adjoint of each input:

$\bar{u}_j = \bar{w} \cdot \frac{\partial \phi}{\partial u_j}$

accumulating contributions from all uses of $u_j$. At the end, $\bar{x}_i = \partial f / \partial x_i$ for all $i$.

The backward pass computes a vector-Jacobian product (VJP): given $\bar{w}$, compute $\bar{u} = J_\phi(u)^T \bar{w}$, where $J_\phi$ is the Jacobian of $\phi$. The full gradient is the VJP of the scalar $1$ with respect to all parameters.

Worked Micro-Graph#

Consider $f(a, b, c) = (a \cdot b + c)^2$ with inputs $a = 1, b = 2, c = 0.5$.

Forward pass (building the graph):

• $v_1 = a \times b = 1 \times 2 = 2$
• $v_2 = v_1 + c = 2 + 0.5 = 2.5$
• $v_3 = v_2^2 = 2.5^2 = 6.25$ (output $f = v_3$)

Backward pass:

• $\bar{v}_3 = \partial f / \partial v_3 = 1$ (seed)
• Operation $v_3 = v_2^2$: $\partial v_3 / \partial v_2 = 2v_2 = 5$, so $\bar{v}_2 = 1 \cdot 5 = 5$
• Operation $v_2 = v_1 + c$:
  • $\partial v_2 / \partial v_1 = 1$, so $\bar{v}_1 = \bar{v}_2 \cdot 1 = 5$
  • $\partial v_2 / \partial c = 1$, so $\bar{c} = \bar{v}_2 \cdot 1 = 5$
• Operation $v_1 = a \times b$:
  • $\partial v_1 / \partial a = b = 2$, so $\bar{a} = \bar{v}_1 \cdot 2 = 10$
  • $\partial v_1 / \partial b = a = 1$, so $\bar{b} = \bar{v}_1 \cdot 1 = 5$

Analytic check: $\partial f / \partial a = 2(ab + c) \cdot b = 2(2.5) \cdot 2 = 10$ ✓. $\partial f / \partial b = 2(ab + c) \cdot a = 5$ ✓. $\partial f / \partial c = 2(ab + c) = 5$ ✓.

Complexity#

Reverse mode for scalar-output functions (the case for most ML losses):

$\text{cost}(\nabla f) \leq c \cdot \text{cost}(f)$

where $c$ is a small constant (typically $c \approx 3$–$5$). The gradient of a scalar function of $d$ variables costs roughly the same as evaluating the function itself — independent of $d$. This is the cheap gradient principle, and it is why deep learning with millions of parameters is computationally feasible.

The trade-off:

• If $d$ (inputs) $\gg$ $p$ (outputs), use reverse mode. One backward pass computes all $\partial f_j / \partial x_i$ for a scalar $f$, vs $d$ forward passes.
• If $p \gg d$, use forward mode. Each forward pass computes one column of the Jacobian.
• Deep learning: $d$ (parameters) is enormous, $p = 1$ (scalar loss). Reverse mode wins by a factor of $d$.

Memory cost. The forward pass must store all intermediate activations for the backward pass (or recompute them). For deep nets with large activations, this is the primary memory bottleneck — not the parameters themselves. Gradient checkpointing trades memory for compute by selectively recomputing activations during the backward pass.

Browser lab#

Implement forward-mode AD (dual numbers) for $\sin(x^2)$ and compare with analytic derivative. Then implement manual reverse-mode AD on the micro-graph $f = (a \cdot b + c)^2$.

import numpy as np

# --- Forward-mode AD (dual numbers) ---
def forward_sin_x2(x):
    """Compute sin(x^2) and its derivative using dual numbers."""
    # v1 = x^2
    v1 = x**2
    dv1 = 2 * x * 1.0  # dx = 1 (seeded)

    # v2 = sin(v1)
    v2 = np.sin(v1)
    dv2 = np.cos(v1) * dv1

    return v2, dv2

x_vals = np.linspace(0.5, 3, 50)
fwd_vals = np.array([forward_sin_x2(x) for x in x_vals])

# Analytic: d/dx sin(x^2) = 2x * cos(x^2)
analytic_vals = 2 * x_vals * np.cos(x_vals**2)

print("Forward-mode AD: sin(x^2) and derivative")
for x in [1.0, 2.0, 3.0]:
    f, df = forward_sin_x2(x)
    an_df = 2 * x * np.cos(x**2)
    print(f"  x={x:.1f}: f={f:.6f}, AD df={df:.6f}, analytic={an_df:.6f}")

# --- Reverse-mode AD on f = (a*b + c)^2 ---
def reverse_micro(a, b, c):
    """Manual reverse-mode on f = (a*b + c)^2."""
    # Forward pass
    v1 = a * b          # v1 = a * b
    v2 = v1 + c         # v2 = v1 + c
    v3 = v2**2          # f = v3

    # Backward pass
    d_v3 = 1.0                           # seed
    d_v2 = d_v3 * 2 * v2                 # d(v3)/d(v2) = 2*v2
    d_v1 = d_v2 * 1.0                    # d(v2)/d(v1) = 1
    d_c  = d_v2 * 1.0                    # d(v2)/d(c) = 1
    d_a  = d_v1 * b                      # d(v1)/d(a) = b
    d_b  = d_v1 * a                      # d(v1)/d(b) = a

    return v3, d_a, d_b, d_c

a, b, c = 1.0, 2.0, 0.5
f, da, db, dc = reverse_micro(a, b, c)
print(f"\nReverse-mode on f = (a*b + c)^2 with (a={a}, b={b}, c={c})")
print(f"  f = {f:.4f}")
print(f"  df/da = {da:.4f}  (analytic: {2*(a*b + c)*b:.4f})")
print(f"  df/db = {db:.4f}  (analytic: {2*(a*b + c)*a:.4f})")
print(f"  df/dc = {dc:.4f}  (analytic: {2*(a*b + c):.4f})")