Automatic Differentiation

Automatic differentiation (AD), also called algorithmic differentiation, computes the derivatives of a function that is given as a computer program. It runs the program and, alongside every elementary operation (an addition, a product, an exponential, a shift of an image), applies that operation’s known derivative through the chain rule. The result is the numerical value of the derivative at the given input, exact up to floating-point rounding, at a cost of a small constant multiple of running the program. Baydin, Pearlmutter, Radul and Siskind (2018) describe it as a non-standard interpretation of the program, in which variables carry derivative values as well as ordinary values and operators are redefined to propagate them.

Modern deep-learning frameworks provide it in one form or another (Baydin et al., 2018). Backpropagation, the algorithm that gives the gradient of a network’s loss with respect to its weights, is the reverse mode of automatic differentiation applied to one kind of program. The same machinery gives Jacobians of image-formation models, Hessian-vector products for second-order methods, gradients through a renderer or through the solution of an optimization problem, and derivatives of derivatives. This article covers the two modes, what each costs, how PyTorch and JAX implement them, and where they are used beyond training networks; the backpropagation article covers the neural-network case, its layer-by-layer formulas and the history of the reverse mode in detail.

Definition

Baydin et al. (2018) separate four ways of obtaining derivatives in software: deriving them by hand and coding them; numerical differentiation by finite differences; symbolic differentiation, which manipulates expressions in a computer algebra system and returns a derivative expression; and automatic differentiation. AD returns numbers, not expressions, and uses the symbolic rules of differentiation only at the level of elementary operations, keeping the intermediate values rather than the expressions that produced them. In their words it is a family of techniques that compute derivatives “through accumulation of values during code execution”.

Two modes combine the local derivatives in opposite orders:

• Forward mode (also called tangent linear mode) carries, with every intermediate value, its derivative with respect to a chosen input direction. One run gives a Jacobian-vector product (JVP) $J\mathbf{v}$.
• Reverse mode (also called adjoint or cotangent linear mode) first runs the program and records what it did, then propagates derivatives of a chosen output combination backward through the record. One backward sweep gives a vector-Jacobian product (VJP) $\mathbf{u}^\top J$, equivalently $J^\top \mathbf{u}$.

Intuition

Take an image $\mathbf{x}$, blur it with a known kernel, compare it with an observed image and compute a per-pixel penalty; this is the objective of non-blind image deblurring. Forward mode asks: if I nudge the input pixels in direction $\mathbf{v}$, how does every output change? It pushes that nudge through the blur, the subtraction and the penalty in the same order as the computation, and needs nothing stored for later. Reverse mode asks the opposite question: for a weighted combination $\mathbf{u}$ of the outputs, how sensitive is it to every input pixel? Answering it starts at the outputs, so the computation must be run first and its intermediate values kept.

The Deep Learning book (Goodfellow, Bengio and Courville, 2016, §6.5.9) compares the two modes to multiplying a chain of Jacobians $ABCD$ from one end or the other: started from an end where the chain reduces to a vector, every step is a matrix-vector product, while the other order multiplies matrices. Forward mode multiplies from the input side and reverse mode from the output side, so the cheaper one depends on the shape of the overall Jacobian: forward mode suits few inputs and many outputs, reverse mode many inputs and few outputs. A training loss has millions of inputs and one output, which is why deep learning uses reverse mode.

Formal Definition

Evaluation trace and computational graph

Let $f : \mathbb{R}^n \to \mathbb{R}^m$ be computed by a program. Following the notation of Griewank and Walther (2008) that Baydin et al. adopt, the run is written as an evaluation trace, also called a Wengert list after Wengert (1964): inputs $v_{i-n} = x_i$ for $i = 1, \dots, n$, intermediate variables

$$v_i = \varphi_i\big(v_j\big)_{j \prec i}, \qquad i = 1, \dots, l,$$

where $\varphi_i$ is an elementary operation and $j \prec i$ means that $v_j$ is an argument of $\varphi_i$, and outputs $y_{m-i} = v_{l-i}$ for $i = m-1, \dots, 0$. Drawing an edge from each argument to each result gives the computational graph, a directed acyclic graph. Only the trace that actually ran matters: loops, branches and function calls leave behind a sequence of elementary operations on numbers, and AD differentiates that sequence, which is why it handles control flow that a closed-form expression cannot express (Baydin et al., 2018, §3).

Forward mode

Choose a direction $\dot{\mathbf{x}} = \mathbf{v} \in \mathbb{R}^n$ and associate with each variable its tangent $\dot v_i$, its derivative along $\mathbf{v}$. Each elementary operation updates the tangent by the chain rule as it computes the value:

$$\dot v_i = \sum_{j \prec i} \frac{\partial \varphi_i}{\partial v_j}\, \dot v_j .$$

At the end the output tangents are $\dot{\mathbf{y}} = J_f(\mathbf{x})\, \mathbf{v}$. Setting $\mathbf{v} = \mathbf{e}_i$, the $i$-th unit vector, gives the $i$-th column of the Jacobian, so the full $m \times n$ Jacobian takes $n$ forward runs.

Dual numbers. Forward mode is evaluation in the algebra of dual numbers $v + \dot v\,\epsilon$, where $\epsilon \neq 0$ and $\epsilon^2 = 0$, which Baydin et al. trace to Clifford (1873). Arithmetic on them reproduces the differentiation rules:

$$(v + \dot v\epsilon) + (u + \dot u\epsilon) = (v + u) + (\dot v + \dot u)\epsilon, \qquad (v + \dot v\epsilon)(u + \dot u\epsilon) = vu + (v\dot u + \dot v u)\epsilon ,$$

and for a scalar function, $f(v + \dot v\epsilon) = f(v) + f'(v)\,\dot v\,\epsilon$. Because every elementary operation respects this rule, so does every composition of them, and the derivative is the coefficient of $\epsilon$ after evaluating the program on $x + 1\epsilon$. In code, a dual number is a pair (value, tangent) with overloaded operators.

Reverse mode

Choose an output weighting $\bar{\mathbf{y}} = \mathbf{u} \in \mathbb{R}^m$ and associate with each variable its adjoint $\bar v_j$, the derivative of $\mathbf{u}^\top \mathbf{y}$ with respect to $v_j$. After a forward run that records the trace, the adjoints are accumulated in reverse order:

$$\bar v_j = \sum_{i \,:\, j \prec i} \bar v_i\, \frac{\partial \varphi_i}{\partial v_j} ,$$

so that each variable collects contributions from every operation that used it. At the inputs, $\bar{\mathbf{x}} = J_f(\mathbf{x})^\top \mathbf{u}$. With $\mathbf{u} = \mathbf{e}_j$ one sweep gives the $j$-th row of the Jacobian, so the full Jacobian takes $m$ backward sweeps; for a scalar loss ($m = 1$) a single sweep gives the whole gradient. The partial derivatives $\partial \varphi_i / \partial v_j$ depend on the values from the forward run, which is why they are stored. The record of operations and the values they need is called the tape (ADOL-C, for example, records operations on tape data structures that are “played back” in the reverse pass, as Baydin et al. describe).

Jacobian-vector and vector-Jacobian products

The two products are the building blocks that torch.func and JAX expose:

$$\text{JVP: } (\mathbf{x}, \mathbf{v}) \mapsto J_f(\mathbf{x})\,\mathbf{v}, \qquad \text{VJP: } (\mathbf{x}, \mathbf{u}) \mapsto J_f(\mathbf{x})^\top \mathbf{u} .$$

Neither forms $J$. The JAX documentation calls them the pushforward of a tangent vector and the pullback of a cotangent vector, and points out that the linear map of a VJP is the transpose of the linear map of a JVP. One consequence is the identity $\langle \mathbf{u}, J\mathbf{v} \rangle = \langle J^\top \mathbf{u}, \mathbf{v} \rangle$ for every $\mathbf{u}$ and $\mathbf{v}$, which is a cheap test that a hand-written reverse rule is the true transpose of its forward rule (the example below runs it). For a linear operation such as a blur with zero padding, the VJP is the correlation with the kernel rotated by 180 degrees, as the convolution and backpropagation articles derive.

Cost

Writing $\operatorname{ops}(f)$ for the operation count of $f$, Baydin et al. state, citing Griewank and Walther (2008), that the full Jacobian costs $n\,c\,\operatorname{ops}(f)$ in forward mode and $m\,c\,\operatorname{ops}(f)$ in reverse mode, with a constant $c < 6$ that is typically 2 to 3. The JAX Autodiff Cookbook gives the same picture per product: a JVP or a VJP costs about three times the floating-point work of evaluating the function. The modes differ in memory. A JVP evaluates the derivatives as it goes, so it stores nothing for later and its memory does not depend on the depth of the computation (JAX Autodiff Cookbook); a VJP must keep the intermediate values until the backward sweep reaches them, so its memory grows with the length of the computation, in the worst case in proportion to the number of operations (Baydin et al., 2018, §3.2). Checkpointing trades this memory for recomputation; the backpropagation article describes it.

Central finite differences need two evaluations of $f$ per input direction, $2n$ for a Jacobian, and give only an approximation whose truncation error shrinks as the step $h$ shrinks while its rounding error grows (Baydin et al., 2018, §2.1). Neither forward nor reverse order need be optimal: other orders of eliminating the chain rule’s terms exist, and the Deep Learning book (§6.5.9) notes that finding the cheapest sequence of operations for a gradient is NP-complete.

Higher derivatives

Derivative code produced by AD is itself a program built from differentiable operations, so it can be differentiated again. For a scalar $f$ the full Hessian is an $n \times n$ matrix, too large to form when $n$ is the number of network weights; second-order methods and Krylov solvers need only the Hessian-vector product $H\mathbf{v}$ (Goodfellow et al., 2016, §6.5.10). AD gives it at a cost of a few function evaluations by composing the modes. Baydin et al. describe computing the directional derivative $\nabla f \cdot \mathbf{v}$ in forward mode and then applying reverse mode to it, crediting Pearlmutter (1994); the Deep Learning book gives the reverse-over-reverse form $H\mathbf{v} = \nabla_{\mathbf{x}}\big[(\nabla_{\mathbf{x}} f)^\top \mathbf{v}\big]$; and the JAX cookbook recommends the forward-over-reverse form, the JVP of the gradient, because forward mode has less overhead and the outer transformation differentiates the larger computation. Baydin et al. cite Christianson (2012) for the result that forward-on-reverse, reverse-on-forward and reverse-on-reverse can compute the second derivative with the same arithmetic, their taping overheads differing by implementation.

Nesting differentiations brings a class of bugs Baydin et al. call perturbation confusion: when two differentiations run through the same code, the two infinitesimals $\epsilon$ they introduce must be kept apart, or a derivative from one leaks into the other. PyTorch’s forward-mode API attaches every tangent to a numbered “dual level” that the user enters and exits, but its version 2.14.1 source states that nested dual_level contexts are not supported, so only one forward-mode level is available there; for higher-order forward derivatives it points to torch.func.jvp.

Properties

• Exact for the program, not for the ideal function. AD differentiates the operations that ran. If a routine approximates a function, AD returns the derivative of the approximation, which need not approximate the true derivative well; Baydin et al. put the rule as “approximate the derivative, not differentiate the approximation”. A function with a known derivative can be given a custom rule for that reason: jax.custom_jvp and torch.autograd.Function exist to supply one.
• Points where the derivative does not exist. PyTorch’s autograd notes (version 2.14) list the rules it applies, in order: use the gradient if it exists; for a locally convex function use the subgradient of minimum norm; for a locally concave one the supergradient of minimum norm; otherwise define it by continuity, so that $\sqrt{x}$ at 0 gets $\infty$; and where the function is undefined, the value is arbitrary, usually NaN. For a ReLU at zero the minimum-norm subgradient is 0.
• Masking after the fact does not mask the derivative. The same notes show that dividing by a tensor containing a zero and then discarding the infinite entry still gives a NaN gradient, because the division was recorded; in our reading, its backward rule divides the discarded entry’s zero incoming gradient by zero. The fix is to mask before the operation.
• Numerical stability is not inherited. A program that is stable forward can produce an unstable derivative. The JAX documentation’s example is $\log(1 + e^x)$ at $x = 100$: composing the primitives’ rules effectively evaluates $\big(1/(1 + e^x)\big) \cdot e^x$, a product of values that floating point rounds to 0 and $\infty$, although the true derivative is close to 1. A custom JVP rule that evaluates the equivalent $1 - 1/(1 + e^x)$ fixes it. Baydin et al. add that a sum of many large terms whose total is near zero, as at a minimum, is numerically fraught in the derivative too.
• Control flow is fixed per run. The derivative is that of the branch taken. A comparison against a threshold contributes no derivative with respect to the threshold, and arg-max, rounding or sampling give zero or no useful gradient, the problem that differentiable rendering faces at triangle edges.

How Libraries Implement It

Baydin et al. (2018, §5) group implementations into elemental libraries, compiler-based methods, operator overloading and hybrids of these. In elemental libraries the user writes calls to AD-aware routines; in source-code transformation a tool rewrites the program into one that also computes derivatives (ADIFOR for Fortran, Tapenade for Fortran and C, Tangent for a subset of Python and NumPy); and in operator overloading new number types carry derivative information through ordinary code (ADOL-C, autograd for Python, PyTorch). Their table lists Google’s Ceres Solver, the C++ library used for bundle adjustment and other nonlinear least-squares problems, among the forward-mode tools. For deep learning they contrast two styles: define-and-run systems such as Theano, TensorFlow, Caffe and CNTK as they stood in 2018, in which the user first builds a static computational graph in a small domain-specific language, and define-by-run systems such as autograd, Chainer and PyTorch, in which the model is an ordinary program and its execution builds the graph anew each time. The Deep Learning book (§6.5.5) draws a related line between symbol-to-number differentiation, which returns gradient values (Torch, Caffe), and symbol-to-symbol differentiation, which adds nodes describing the derivatives to the graph (Theano, TensorFlow), so that the derivative graph can itself be differentiated for higher-order derivatives.

PyTorch

torch.autograd is a reverse-mode system whose graph of recorded operations is rebuilt on every forward pass, as the PyTorch article describes. Operations save the tensors their backward rule needs (x.pow(2) saves x), and hooks on saved tensors can move them elsewhere, for example to disk, to save memory. That article also covers the training-loop side and in-place checks; Distributed Training covers gradient accumulation.

Forward mode is a separate, beta API in torch.autograd.forward_ad: inside a dual_level() context, make_dual(x, v) attaches a tangent to a tensor and unpack_dual reads the tangent of a result, giving a JVP in one pass. The higher-level torch.func library, previously functorch, provides JAX-like composable function transforms: grad, vjp, jvp, jacrev, jacfwd, hessian and vmap. Its documentation names computing per-sample gradients, Jacobians and Hessians, and batching the inner loop of meta-learning as the motivating cases. In version 2.14.1’s source, jacrev computes the Jacobian by applying vmap to the VJP over the rows of an identity matrix, all at once by default, with a chunk_size option to do it in pieces or, with chunk_size=1, row by row in a loop; jacfwd applies vmap to the JVP over the identity basis, with no chunk_size option; and hessian is jacfwd(jacrev(f)), which the docstring calls forward-over-reverse and “a good default”. Forward-mode coverage is narrower than reverse-mode coverage: the jacfwd docstring warns that some operators may raise “forward-mode AD not implemented” and suggests jacrev instead.

JAX

JAX is organized around transformations of functions: jax.grad, jax.jvp, jax.vjp, jax.jacfwd, jax.jacrev, jax.vmap and jax.jit, which take a Python function and return a new one and compose freely (grad(grad(f)), jit(vmap(grad(f)))). To transform a function, JAX calls it with tracers, stand-ins for arrays that record the sequence of primitive operations, and represents that sequence in its own intermediate language, the jaxpr. The tracing documentation explains a practical difference between the transformations: jax.jit traces with abstract values (shape and dtype only), so a Python if on an array’s value fails unless that argument is marked static, whereas jax.grad and jax.jvp trace with concrete values and can resolve such a branch on the actual number.

The JAX cookbook builds jacfwd and jacrev by pushing forward or pulling back a whole identity basis at once with vmap, and notes that the authors’ earlier Autograd library had to pull back one vector at a time. Custom rules come as jax.custom_jvp and jax.custom_vjp; for a function with only a custom JVP rule, JAX derives reverse mode by transposing the linear computation on tangents. The same documentation uses custom_vjp for implicit differentiation: a fixed point $x^\star(a)$ computed by iterating $x \leftarrow f(a, x)$ in a while_loop is not differentiated through its iterations, which JAX cannot do in reverse mode because the compiled loop would need unbounded memory, but through the implicit function theorem at the solution, solving a linear equation for the VJP instead.

Examples

The example below writes both modes from scratch in NumPy, uses them on image functions with different numbers of inputs and outputs, and checks them against finite differences and against torch.func. The image is a synthetic 16 × 16 picture of a disk and a square, $s$, blurred by a known asymmetric 3 × 3 kernel $k$ (zero padding, correlation: output pixel $(i, j)$ is $\sum_{u,v} k_{uv}\, x_{i+u-1,\, j+v-1}$), with Gaussian noise of standard deviation 0.02 added to give the observation $t$. The derivatives are taken at $\mathbf{x} = t$, where a deblurring method might start. The per-pixel penalty is the Charbonnier function $\rho(r) = \sqrt{r^2 + \varepsilon^2}$ with $\varepsilon = 0.01$, applied to the residual $r = k \star \mathbf{x} - t$. Four functions are differentiated:

Function	Inputs $n$	Outputs $m$	What it is
F_map	256 pixels	256	the per-pixel penalty map
F_block	256	16	the map averaged over 4 × 4 blocks
F_mean	256	1	the mean penalty, a scalar loss
G_render	2 ($\sigma$, gain)	256	the image $s$ blurred by a normalized 3 × 3 Gaussian of width $\sigma$ and scaled by the gain, at $\sigma = 0.8$, gain 1.2

Dual holds a value and a tangent and implements forward mode by operator overloading. Var appends itself to a global tape and stores, for each argument, a function that maps the result’s adjoint to that argument’s adjoint; backward walks the tape in reverse. Setting __array_ufunc__ = None makes NumPy hand mixed expressions such as t - Dual(...) to these classes. The image functions are written once and run unchanged on plain arrays (for finite differences), on Dual values and on Var values. Jacobians are stored as $m \times n$ arrays, and each error is the largest absolute difference divided by the largest absolute entry of the reference. The first block holds the AD code and the functions:

import warnings
import numpy as np
import torch
from torch.func import jacfwd, jacrev, jvp, vjp, grad, hessian

warnings.filterwarnings("ignore", category=FutureWarning)
torch.set_num_threads(1)
H = W = 16
EPS = 0.01                                     # Charbonnier penalty parameter

# ---------------- forward mode: dual numbers ----------------
class Dual:
    """A value and its tangent: x + x_dot * eps with eps**2 = 0."""
    __array_ufunc__ = None                     # make NumPy hand mixed operations to us
    def __init__(self, val, dot):
        self.val, self.dot = np.asarray(val, float), np.asarray(dot, float)
    def __add__(a, b):
        b = b if isinstance(b, Dual) else Dual(b, 0.0)
        return Dual(a.val + b.val, a.dot + b.dot)
    __radd__ = __add__
    def __neg__(a): return Dual(-a.val, -a.dot)
    def __sub__(a, b): return a + (-b)
    def __rsub__(a, b): return (-a) + b
    def __mul__(a, b):
        b = b if isinstance(b, Dual) else Dual(b, 0.0)
        return Dual(a.val * b.val, a.dot * b.val + a.val * b.dot)
    __rmul__ = __mul__
    def __truediv__(a, b):
        b = b if isinstance(b, Dual) else Dual(b, 0.0)
        q = a.val / b.val
        return Dual(q, (a.dot - q * b.dot) / b.val)
    def __rtruediv__(a, b): return Dual(b, 0.0) / a
    def sqrt(a):
        r = np.sqrt(a.val)
        return Dual(r, a.dot / (2 * r))
    def exp(a):
        e = np.exp(a.val)
        return Dual(e, e * a.dot)
    def shift(a, du, dv): return Dual(np_shift(a.val, du, dv), np_shift(a.dot, du, dv))
    def blockmean(a, b): return Dual(np_blockmean(a.val, b), np_blockmean(a.dot, b))
    def sum(a): return Dual(a.val.sum(), a.dot.sum())

# ---------------- reverse mode: a tape ----------------
TAPE = []                                      # every Var in the order it was computed

class Var:
    """A value plus, for each input, a function mapping the output adjoint to that input's adjoint."""
    __array_ufunc__ = None
    def __init__(self, val, parents=()):
        self.val, self.parents = np.asarray(val, float), parents
        TAPE.append(self)
    def __add__(a, b):
        if not isinstance(b, Var):
            return Var(a.val + b, [(a, lambda g: g)])
        return Var(a.val + b.val, [(a, lambda g: unbroadcast(g, a.val)),
                                   (b, lambda g: unbroadcast(g, b.val))])
    __radd__ = __add__
    def __neg__(a): return Var(-a.val, [(a, lambda g: -g)])
    def __sub__(a, b): return a + (-b)
    def __rsub__(a, b): return (-a) + b
    def __mul__(a, b):
        if not isinstance(b, Var):
            return Var(a.val * b, [(a, lambda g: unbroadcast(g * b, a.val))])
        return Var(a.val * b.val, [(a, lambda g: unbroadcast(g * b.val, a.val)),
                                   (b, lambda g: unbroadcast(g * a.val, b.val))])
    __rmul__ = __mul__
    def __truediv__(a, b):
        if not isinstance(b, Var):
            return a * (1.0 / b)
        q = a.val / b.val
        return Var(q, [(a, lambda g: unbroadcast(g / b.val, a.val)),
                       (b, lambda g: unbroadcast(-g * q / b.val, b.val))])
    def __rtruediv__(a, b):
        q = b / a.val
        return Var(q, [(a, lambda g: unbroadcast(-g * q / a.val, a.val))])
    def sqrt(a):
        r = np.sqrt(a.val)
        return Var(r, [(a, lambda g: g / (2 * r))])
    def exp(a):
        e = np.exp(a.val)
        return Var(e, [(a, lambda g: g * e)])
    def shift(a, du, dv):                      # adjoint of a shift is the opposite shift
        return Var(np_shift(a.val, du, dv), [(a, lambda g: np_shift(g, -du, -dv))])
    def blockmean(a, b):
        return Var(np_blockmean(a.val, b),
                   [(a, lambda g: np.repeat(np.repeat(g, b, 0), b, 1) / b**2)])
    def sum(a):
        return Var(a.val.sum(), [(a, lambda g: np.full(a.val.shape, g))])

def unbroadcast(g, like):                      # a scalar used against an image collects the sum
    return np.sum(g) if np.ndim(like) == 0 and np.ndim(g) > 0 else g

def backward(out, seed):
    """One reverse sweep over the tape: returns the adjoint of every Var reached."""
    adj = {id(out): np.asarray(seed, float)}
    for node in reversed(TAPE):
        if id(node) in adj:
            for parent, vjp_rule in node.parents:
                g = vjp_rule(adj[id(node)])
                adj[id(parent)] = adj[id(parent)] + g if id(parent) in adj else g
    return adj

# ---------------- the image functions, written once ----------------
def np_shift(img, du, dv):                     # out[i, j] = img[i + du, j + dv], zero outside
    out = np.zeros_like(img)
    h, w = img.shape
    out[max(0, -du):h - max(0, du), max(0, -dv):w - max(0, dv)] = \
        img[max(0, du):h - max(0, -du), max(0, dv):w - max(0, -dv)]
    return out

def np_blockmean(img, b):
    return img.reshape(img.shape[0] // b, b, img.shape[1] // b, b).mean(axis=(1, 3))

def shift(z, du, dv): return np_shift(z, du, dv) if isinstance(z, np.ndarray) else z.shift(du, dv)
def sqrt(z): return np.sqrt(z) if isinstance(z, np.ndarray) else z.sqrt()
def exp(z): return np.exp(z) if isinstance(z, (np.ndarray, float)) else z.exp()

def blur(img, ker):                            # zero-padded 3x3 correlation
    return sum(ker[u][v] * shift(img, u - 1, v - 1) for u in range(3) for v in range(3))

yy, xx = np.mgrid[0:H, 0:W]
s = np.full((H, W), 0.2)                       # sharp synthetic image: disk and square
s[(yy - 5) ** 2 + (xx - 5) ** 2 <= 16] = 0.9
s[9:14, 8:14] = 0.6
k = np.array([[0.00, 0.05, 0.10], [0.05, 0.30, 0.20], [0.00, 0.10, 0.20]])   # known blur
noise = np.random.default_rng(0).normal(0, 0.02, (H, W))
t = blur(s, k) + noise                         # observed blurry, noisy image
x0 = t.copy()                                  # point of differentiation: the observation
d2 = [[(u - 1) ** 2 + (v - 1) ** 2 for v in range(3)] for u in range(3)]

def loss_map(x):                               # Charbonnier penalty of the residual
    r = blur(x, k) - t
    return sqrt(r * r + EPS ** 2)
def F_map(x): return loss_map(x)               # 256 -> 256
def F_block(x): return loss_map(x).blockmean(4) if not isinstance(x, np.ndarray) \
    else np_blockmean(loss_map(x), 4)          # 256 -> 16
def F_mean(x): return loss_map(x).sum() * (1.0 / (H * W))   # 256 -> 1
def G_render(theta):                           # (sigma, gain) -> 256 pixels
    sigma, gain = theta
    g = [[exp(-d2[u][v] / (2 * sigma * sigma)) for v in range(3)] for u in range(3)]
    Z = sum(g[u][v] for u in range(3) for v in range(3))
    return gain * blur(s, [[g[u][v] / Z for v in range(3)] for u in range(3)])

The second block builds full Jacobians in both modes, counting forward passes and backward sweeps, and compares them with central differences ($h = 10^{-6}$), with the same functions written in PyTorch and differentiated by torch.func.jacfwd and jacrev, with single random JVPs and VJPs from torch.func.jvp and vjp, and through the adjoint test. It ends with a Hessian-vector product of the mean loss, computed in torch.func as the JVP of the gradient and compared with a central difference of our own reverse-mode gradient and with the full Hessian from torch.func.hessian.

# ---------------- full Jacobians, counting passes ----------------
def jac_forward(f, p):
    cols, passes = [], 0
    for i in range(p.size):                    # one forward pass per input direction
        e = np.zeros(p.size); e[i] = 1.0
        z = Dual(p.reshape(H, W), e.reshape(H, W)) if p.size == H * W else \
            [Dual(p[j], e[j]) for j in range(p.size)]
        cols.append(np.ravel(f(z).dot)); passes += 1
    return np.stack(cols, axis=1), passes

def jac_reverse(f, p):
    TAPE.clear()
    inputs = [Var(p.reshape(H, W))] if p.size == H * W else [Var(p[j]) for j in range(p.size)]
    out = f(inputs[0] if p.size == H * W else inputs)
    nodes, floats = len(TAPE), sum(v.val.size for v in TAPE)
    rows, sweeps = [], 0
    for j in range(out.val.size):              # one reverse sweep per output component
        seed = np.zeros(out.val.size); seed[j] = 1.0
        adj = backward(out, seed.reshape(out.val.shape))
        rows.append(np.concatenate([np.ravel(adj.get(id(v), np.zeros(v.val.shape))) for v in inputs]))
        sweeps += 1
    return np.stack(rows), sweeps, nodes, floats

def jac_fd(f, p, h=1e-6):                      # central differences: 2n evaluations
    def fp(q): return np.ravel(f(q.reshape(H, W) if q.size == H * W else list(q)))
    return np.stack([(fp(p + h * e) - fp(p - h * e)) / (2 * h) for e in np.eye(p.size)], axis=1)

# ---------------- the same functions in PyTorch ----------------
def t_shift(img, du, dv):
    return torch.nn.functional.pad(img, (1, 1, 1, 1))[1 + du:1 + du + H, 1 + dv:1 + dv + W]
def t_blur(img, ker):
    return sum(ker[u][v] * t_shift(img, u - 1, v - 1) for u in range(3) for v in range(3))
T, S = torch.tensor(t), torch.tensor(s)
def T_loss_map(x):
    r = t_blur(x.reshape(H, W), k.tolist()) - T
    return torch.sqrt(r * r + EPS ** 2)
T_funcs = {
    "F_map": lambda x: T_loss_map(x).reshape(-1),
    "F_block": lambda x: T_loss_map(x).reshape(4, 4, 4, 4).mean(dim=(1, 3)).reshape(-1),
    "F_mean": lambda x: T_loss_map(x).mean().reshape(1),
    "G_render": lambda th: (th[1] * t_blur(S, (lambda g: [[g[u][v] / g.sum() for v in range(3)]
                            for u in range(3)])(torch.exp(-torch.tensor(d2, dtype=torch.float64)
                            / (2 * th[0] ** 2))))).reshape(-1),
}

def rel(A, B): return np.abs(A - B).max() / np.abs(B).max()

print(f"torch {torch.__version__}, numpy {np.__version__}")
print(f"kernel sum {k.sum():.3f}, symmetric under 180-degree rotation: {np.allclose(k, k[::-1, ::-1])}, "
      f"noise std {noise.std():.4f}")
cases = [("F_map", F_map, x0.ravel()), ("F_block", F_block, x0.ravel()),
         ("F_mean", F_mean, x0.ravel()), ("G_render", G_render, np.array([0.8, 1.2]))]
print(f"{'function':9s}{'n':>4s}{'m':>5s}{'fwd passes':>11s}{'rev sweeps':>11s}{'tape nodes':>11s}"
      f"{'tape floats':>12s}{'fwd-rev':>9s}{'vs FD':>9s}{'vs torch':>9s}")
rng = np.random.default_rng(1)
for name, f, p in cases:
    Jf, npass = jac_forward(f, p)
    Jr, nsweep, nodes, floats = jac_reverse(f, p)
    Jfd = jac_fd(f, p)
    tp = torch.tensor(p)
    Jtf = jacfwd(T_funcs[name])(tp).numpy()
    Jtr = jacrev(T_funcs[name])(tp).numpy()
    err_fd = max(rel(Jf, Jfd), rel(Jr, Jfd))
    err_t = max(rel(Jf, Jtf), rel(Jr, Jtr))
    print(f"{name:9s}{p.size:4d}{Jf.shape[0]:5d}{npass:11d}{nsweep:11d}{nodes:11d}{floats:12d}"
          f"{rel(Jf, Jr):9.1e}{err_fd:9.1e}{err_t:9.1e}")
    # one JVP and one VJP with random directions, against torch.func.jvp / vjp
    v, u = rng.normal(size=p.size), rng.normal(size=Jf.shape[0])
    if p.size == H * W:
        jv = np.ravel(f(Dual(p.reshape(H, W), v.reshape(H, W))).dot)
    else:
        jv = np.ravel(f([Dual(p[j], v[j]) for j in range(p.size)]).dot)
    TAPE.clear()
    inputs = [Var(p.reshape(H, W))] if p.size == H * W else [Var(p[j]) for j in range(p.size)]
    out = f(inputs[0] if p.size == H * W else inputs)
    adj = backward(out, u.reshape(out.val.shape))
    uj = np.concatenate([np.ravel(adj.get(id(q), np.zeros(q.val.shape))) for q in inputs])
    _, jv_t = jvp(T_funcs[name], (tp,), (torch.tensor(v),))
    _, vjp_fn = vjp(T_funcs[name], tp)
    uj_t = vjp_fn(torch.tensor(u))[0]
    print(f"   JVP vs torch.func.jvp {rel(jv, jv_t.numpy()):.1e}, VJP vs torch.func.vjp "
          f"{rel(uj, uj_t.numpy()):.1e}, adjoint test |<u,Jv>-<J'u,v>|/|<u,Jv>| "
          f"{abs(u @ jv - uj @ v) / abs(u @ jv):.1e}")

# Hessian-vector product of the mean loss: forward-over-reverse in torch.func,
# against a central difference of our reverse-mode gradient
def our_grad(xv):
    TAPE.clear()
    xin = Var(xv.reshape(H, W))
    return np.ravel(backward(F_mean(xin), 1.0)[id(xin)])
v = rng.normal(size=H * W)
hv_fd = (our_grad(x0.ravel() + 1e-6 * v) - our_grad(x0.ravel() - 1e-6 * v)) / 2e-6
fm = lambda z: T_loss_map(z).mean()
hv_t = jvp(grad(fm), (torch.tensor(x0.ravel()),), (torch.tensor(v),))[1].numpy()
hv_h = (hessian(fm)(torch.tensor(x0.ravel())) @ torch.tensor(v)).numpy()
print(f"HVP of F_mean: torch forward-over-reverse vs FD of our gradient {rel(hv_t, hv_fd):.1e}, "
      f"vs torch.func.hessian @ v {rel(hv_t, hv_h):.1e}")

Output (PyTorch 2.14.1, NumPy 2.5.4, float64; the script ran in about 2 seconds on our machine, a figure that varies with the machine, and printed identical output on two runs):

torch 2.14.1, numpy 2.5.4
kernel sum 1.000, symmetric under 180-degree rotation: False, noise std 0.0202
function    n    m fwd passes rev sweeps tape nodes tape floats  fwd-rev    vs FD vs torch
F_map     256  256        256        256         32        8192  1.9e-16  2.4e-10  0.0e+00
   JVP vs torch.func.jvp 0.0e+00, VJP vs torch.func.vjp 0.0e+00, adjoint test |<u,Jv>-<J'u,v>|/|<u,Jv>| 1.4e-16
F_block   256   16        256         16         33        8208  2.2e-16  2.2e-10  2.2e-16
   JVP vs torch.func.jvp 1.3e-16, VJP vs torch.func.vjp 0.0e+00, adjoint test |<u,Jv>-<J'u,v>|/|<u,Jv>| 0.0e+00
F_mean    256    1        256          1         34        8194  2.2e-16  1.3e-09  1.1e-16
   JVP vs torch.func.jvp 1.9e-16, VJP vs torch.func.vjp 0.0e+00, adjoint test |<u,Jv>-<J'u,v>|/|<u,Jv>| 0.0e+00
G_render    2  256          2        256         75        4920  5.6e-16  2.5e-10  8.0e-16
   JVP vs torch.func.jvp 1.9e-15, VJP vs torch.func.vjp 9.9e-16, adjoint test |<u,Jv>-<J'u,v>|/|<u,Jv>| 2.2e-16
HVP of F_mean: torch forward-over-reverse vs FD of our gradient 4.1e-09, vs torch.func.hessian @ v 3.7e-16

The pass counts are the costs of the two modes made concrete. Forward mode needed one pass per input, 256 for every image function and 2 for the renderer; reverse mode needed one sweep per output after a single recorded run, from 256 for the penalty map down to a single sweep for the scalar loss. For F_mean one reverse sweep after one recorded run replaces 256 forward passes, which is the case training relies on; for G_render, two scene parameters and 256 pixels, 2 forward passes replace 256 reverse sweeps, the shape of a renderer’s sensitivity to a few scene parameters; for the square F_map the two tie. Whichever mode builds it, the Jacobian is the same: the two agree to within a few units of float64 rounding.

The finite-difference column is larger, $2 \times 10^{-10}$ to $10^{-9}$. Since the two AD Jacobians agree with each other and with PyTorch to about $10^{-16}$, the gap belongs to the finite differences, whose truncation and rounding errors at $h = 10^{-6}$ are of this order; the scalar loss shows the largest relative gap, in our reading because its gradient entries are small (the mean divides by 256) while the rounding error of each difference does not shrink with them. The comparison with torch.func is zero or at rounding level for every function: for F_map, in our reading, the PyTorch version performs the same operations in the same order, so even the rounding matches. The adjoint test passes at rounding level, which matters here because the kernel is not symmetric under a 180-degree rotation: in our reasoning, a VJP that used the forward shift instead of the opposite one would not be the transpose and would fail it.

The tape columns show the memory side. Recording F_map kept 32 arrays of 256 values, 8,192 floats, 32 times the size of the input image, and every one of them stays alive until the backward sweep has passed it; the forward-mode passes kept nothing beyond the current value and tangent. G_render recorded 75 nodes, many of them the scalars that build the Gaussian kernel from $\sigma$. The last line shows second-order differentiation by composition: one forward-mode JVP of the reverse-mode gradient gives $H\mathbf{v}$ without forming the $256 \times 256$ Hessian, agreeing with the explicit Hessian to rounding and with a finite difference of our own gradient to $4 \times 10^{-9}$.

Common Misconceptions

• “Automatic differentiation is symbolic differentiation.” Symbolic differentiation returns an expression, and applied naively to nested products it can return expressions exponentially larger than the original, the expression swell that Baydin et al. illustrate with iterations of the logistic map. AD keeps numerical values of shared subexpressions instead, and works on code with loops and branches that has no closed form.
• “Automatic differentiation is numerical differentiation.” It involves no step size and no truncation error. The example’s finite differences were off by up to $10^{-9}$; the two modes agreed to $10^{-16}$.
• “Reverse mode is always the right choice.” It is for scalar losses. For functions with few inputs and many outputs, such as the sensitivity of a rendered image to a handful of scene parameters, forward mode needs fewer passes and no tape, as G_render shows. The JAX cookbook adds that for near-square Jacobians jacfwd probably has an edge over jacrev.
• “A full Jacobian costs one backward pass.” It costs $m$ VJPs or $n$ JVPs. Libraries batch them with vmap into one call, which uses the hardware better but does not remove the work.
• “Backpropagation and automatic differentiation are different things.” Backpropagation is reverse-mode AD specialized to a network’s loss; general AD also differentiates programs that are not networks, in either mode and to any order (Baydin et al., 2018, §4.2).
• “If autograd returns a gradient, it is the gradient I want.” It is the gradient of the operations that ran. Zero gradients from a hard rasterizer, NaN from a masked division and derivatives of an approximating routine are all correct derivatives of the program and wrong for the intended function.

Where It Is Used

• Training networks. Every gradient step in training a neural network uses reverse-mode AD of the loss, as described in the backpropagation article, and the optimizers of gradient descent consume the result.
• Differentiable rendering and inverse graphics. Fitting scene parameters to images needs derivatives of pixels with respect to geometry, materials and cameras. Baydin et al. list OpenDR (Loper and Black, 2014), which fitted a 3D body model to images and range data, among the uses of AD in computer vision; the differentiable rendering article explains where plain AD of a rasterizer fails and how renderers repair it, and neural radiance fields and 3D Gaussian splatting are fitted to photographs by reverse-mode AD through a renderer that is differentiable by design.
• Classical vision optimization. Baydin et al. cite Srajer et al.’s benchmark of AD tools on bundle adjustment, Gaussian mixture model fitting and hand tracking; uses of AD for denoising, segmentation and stereo (Pock et al., 2007); and GPU-accelerated 2D/3D medical image registration with reverse-mode AD (Grabner et al., 2008), whose authors, Baydin et al. report, found it six times faster than central differences. Least-squares solvers such as Gauss–Newton and Levenberg–Marquardt need the Jacobian of the residuals; in our reading, residuals with few parameters each, as in bundle adjustment, are the case where forward mode fits.
• Optimization layers and implicit differentiation. OptNet (Amos and Kolter, 2017) makes the solution of a quadratic program a layer of a network. Instead of differentiating the solver’s iterations, it differentiates the optimality (KKT) conditions at the solution and obtains the backward pass by solving one linear system with the factorization the interior-point solver already computed, which the authors describe as adding virtually no cost on top of the solve; they trained such a layer to solve 4 × 4 mini-Sudoku puzzles from examples alone. JAX’s custom-VJP fixed-point example applies the same implicit-function idea to iterative solvers.
• Second derivatives in meta-learning. MAML (Finn, Abbeel and Levine, 2017) learns an initialization through a gradient step, so its meta-gradient differentiates through a gradient and needs Hessian-vector products from an extra backward pass; the authors’ first-order approximation drops them. The few-shot recognition article reports the image-classification results of both.
• Hyperparameters and probabilistic models. Baydin et al. survey further uses: Hessian-vector products for stochastic second-order methods, gradients of a training objective with respect to hyperparameters through nested AD, Hamiltonian Monte Carlo and variational inference in probabilistic programming systems such as Stan and Pyro, and the reparameterization trick for gradients through random sampling.

History

Baydin et al. (2018, §3.3) trace the ideas behind AD to the 1950s, and credit Wengert (1964) with essentially discovering forward mode as a general method for evaluating partial derivatives. They identify the essence of the reverse mode, in continuous time, with the Pontryagin maximum principle of control theory; its first published descriptions and its reinventions in machine learning as backpropagation are traced in the backpropagation article. Speelpenning (1980) gave the first implementation that accepted a program in a general-purpose language and produced the reverse-mode derivative automatically. Interest revived in the 1980s, mostly through Griewank’s work. Baydin et al. describe the AD and machine-learning communities as largely unaware of each other until shortly before their survey, and general-purpose AD as entering machine learning, with libraries such as autograd, Chainer and PyTorch, under the names “dynamic computational graphs” and “differentiable programming”.