What Is Automatic Differentiation?
- 7 hours ago
- 35 min read
Every time a neural network gets a little better at its job, something below the surface just solved a calculus problem. Not an approximate one. An exact one, accurate to the limits of floating-point computing, computed automatically for a function with millions of moving parts. That capability is called automatic differentiation, and it is the quiet algorithm that makes modern machine learning — and a good chunk of modern science — actually work.
TL;DR
Automatic differentiation (AD, or "autodiff") computes exact derivatives of functions defined by computer programs by applying the chain rule to each elementary operation the program executes, not by approximating or symbolically expanding a formula.
AD comes in two primary modes: forward mode, which propagates tangent values alongside the computation, and reverse mode, which records a trace and propagates adjoint values backward through it.
Reverse-mode AD is the general mathematical technique behind backpropagation, the method used to train neural networks; the two terms are closely related but not perfectly interchangeable.
Frameworks such as JAX, PyTorch, and TensorFlow expose AD through operator overloading and taped or traced computational graphs, so developers rarely write derivative formulas by hand.
AD is not free of pitfalls — nondifferentiable operations, in-place mutation, control flow, and memory pressure from long reverse-mode tapes all require careful handling.
What Is Automatic Differentiation?
Automatic differentiation is a method computers use to calculate exact derivatives of functions written as programs. It breaks a computation into elementary operations, applies the chain rule to each step, and combines the results. Unlike numerical differentiation, it is not an approximation, and unlike symbolic differentiation, it does not build a giant formula — it differentiates the actual sequence of operations the program runs.
Table of Contents
What Automatic Differentiation Actually Means
Automatic differentiation, also called algorithmic differentiation or "autodiff," is a family of techniques for computing exact derivatives of numeric functions expressed as computer programs, by systematically applying the chain rule to the elementary operations the program actually executes. Baydin, Pearlmutter, Radul, and Siskind, in their widely cited Journal of Machine Learning Research survey, describe AD as a set of techniques <cite index="47-1">similar to but more general than backpropagation for efficiently and accurately evaluating derivatives of numeric functions expressed as computer programs</cite> [1]. The same survey notes AD has quietly served fields such as <cite index="47-2">computational fluid dynamics, atmospheric sciences, and engineering design optimization</cite> for decades [1], long before it became central to deep learning.
"Automatic" does not mean "approximate." AD does not perturb inputs and measure the resulting change, the way numerical differentiation does, nor does it expand a function into one giant symbolic formula, the way a computer algebra system does. Instead, AD treats a program as a sequence of simple operations — additions, multiplications, sine, exponentials, logarithms — each with a known, simple derivative rule, and walks through them in the order the program actually runs, tracking derivative information as it goes. The result is exact up to the floating-point precision of the arithmetic involved, not an estimate.
Why does this matter to a working programmer or data scientist? Because derivatives tell an optimization algorithm which direction reduces error — nearly all of training, curve fitting, and control tuning comes down to nudging parameters toward lower loss, and that direction is the gradient. Computing gradients by hand for a network with billions of parameters is not something any human should attempt or that scales as architectures change weekly. AD removes that burden: define the forward computation, and the derivative comes along for free, exactly the way a chain of known, local sensitivities multiplies together into one global sensitivity along any assembly line of simple steps.
The Calculus You Need to Follow Along
Before touching AD mechanics, it helps to fix notation. A derivative measures how much a function's output changes for a small change in its input; for $f: \mathbb{R} \to \mathbb{R}$, it is written $f'(x)$ or $\frac{df}{dx}$. When a function has multiple inputs, $f: \mathbb{R}^n \to \mathbb{R}$, the partial derivative $\frac{\partial f}{\partial x_i}$ measures sensitivity to one input while holding the others fixed. Stacking all partial derivatives into a vector gives the gradient, $\nabla f(x) \in \mathbb{R}^n$ — the direction of steepest increase of $f$ at $x$. For a vector-valued function $f: \mathbb{R}^n \to \mathbb{R}^m$, the natural generalization is the Jacobian, an $m \times n$ matrix $J$ whose entry $J_{ij} = \frac{\partial f_i}{\partial x_j}$ captures how the $i$-th output responds to the $j$-th input. This article uses the convention that $J$ has outputs along rows and inputs along columns, which matters later for Jacobian-vector and vector-Jacobian products.
Case | Function type | Derivative object | Shape |
Scalar to scalar | $f: \mathbb{R} \to \mathbb{R}$ | Derivative | Scalar |
Vector to scalar | $f: \mathbb{R}^n \to \mathbb{R}$ | Gradient $\nabla f$ | $n$-vector |
Scalar to vector | $f: \mathbb{R} \to \mathbb{R}^m$ | Derivative vector | $m$-vector |
Vector to vector | $f: \mathbb{R}^n \to \mathbb{R}^m$ | Jacobian $J$ | $m \times n$ matrix |
The chain rule is the single mathematical fact that makes AD possible. If $y = g(u)$ and $u = h(x)$, then
$$\frac{dy}{dx} = \frac{dy}{du} \cdot \frac{du}{dx}$$
For a composition of many functions, $y = f_k(f_{k-1}(\dots f_1(x)))$, the same idea extends: the total derivative is a product of local derivatives, one per step. AD is, at its core, nothing more than a disciplined, automated bookkeeping system for that product, applied to every elementary step a real program performs — a point Griewank and Walther develop rigorously in their SIAM reference text on algorithmic differentiation [2].
Four Ways to Get a Derivative, Compared
There are four broad approaches to obtaining derivatives from a program, and confusing them is the single most common source of misunderstanding about what AD is.
Hand-derived derivatives. A person works out the calculus manually and codes the formula. Exact and fast, but it does not scale to large or frequently changing programs and grows error-prone as functions become long or nested.
Numerical differentiation. The derivative is estimated with a finite difference such as $\frac{f(x+h) - f(x-h)}{2h}$. Easy to implement, but only an approximation, and accuracy depends heavily on step size $h$: too large introduces truncation error, too small introduces floating-point cancellation.
Symbolic differentiation. A computer algebra system manipulates the formula directly to produce a new formula. Exact, but the result can grow explosively with each differentiation step — expression swell — and it typically needs a closed-form expression rather than an arbitrary program with loops and branches.
Automatic differentiation. The derivative is computed by applying the chain rule to each elementary operation the program actually executes, accumulating results numerically as it runs. Exact to floating-point precision, immune to expression swell, and works directly on real code, though the exact scope of supported control flow depends on the implementation.
Criterion | Hand-derived | Numerical (finite differences) | Symbolic | Automatic differentiation |
Basic method | Manual calculus, coded directly | Perturb inputs, measure output change | Algebraic manipulation of a formula | Chain rule applied to elementary program operations |
Accuracy | Exact, if derived correctly | Approximate; sensitive to step size | Exact | Exact to floating-point precision |
Runtime for gradients | Fast, once written | One or more extra function evaluations per input dimension | Can be slow if the expression swells | Roughly a small constant multiple of one function evaluation (reverse mode) |
Memory use | Minimal | Minimal | Can grow rapidly with expression size | Forward mode: modest; reverse mode: proportional to trace length |
Implementation effort | High, and repeated for every function | Low | Moderate to high | Low once the AD system exists; near-zero per new function |
Scalability to large models | Poor | Poor (cost grows with input dimension) | Poor (expression swell) | Strong, especially reverse mode for many-parameter models |
Expression swell | Not applicable | Not applicable | Common | Avoided; intermediate values are numbers, not growing formulas |
Sensitivity to step size | None | High | None | None |
Suitability for large ML models | Impractical | Impractical for gradients with many parameters | Impractical | Standard approach |
Typical use cases | Small, fixed formulas; textbook derivations | Debugging and gradient checking | Symbolic math software (for example, computer algebra systems) | Neural network training, scientific computing, optimization |
A point AD newcomers frequently blur: AD is not numerical approximation, and it is not symbolic differentiation, even though both are exact. Symbolic differentiation manipulates and returns an explicit formula; AD evaluates the program at a specific point while carrying derivative information alongside the primal values, producing numbers, not a new formula — a distinction the Baydin et al. survey addresses at length because it is so often conflated in machine learning writing [1].
How Automatic Differentiation Works Under the Hood
Every program, however complex, can be decomposed into a sequence of elementary operations: additions, multiplications, divisions, and calls to elementary functions like $\sin$, $\exp$, or $\log$. AD systems exploit this decomposition through a few connected ideas.
Decomposition into elementary operations. A complicated expression like $z = \sin(x \cdot y) + e^{x}$ is broken into intermediate steps: multiply $x$ and $y$, take the sine of that result, exponentiate $x$, and add the two results. Each of these steps has a simple, known derivative rule.
The evaluation trace. As the program runs with specific input values, AD records the sequence of intermediate variables it computes — this record is the evaluation trace. Each entry corresponds to one elementary operation and stores which earlier variables fed into it.
The computational graph. The trace can be visualized as a directed acyclic graph, where nodes represent variables (inputs, intermediate results, and outputs) and edges represent the elementary operations connecting them. This is the same structure that PyTorch's documentation refers to as a dynamically built directed acyclic graph, or DAG, tracing operations from input tensors (the leaves) to output tensors (the roots) [6].
Local derivatives. For each elementary operation, AD applies a known, hard-coded derivative rule. The derivative of a multiplication, a sine, or an exponential is simple and fixed; the AD system does not need to "discover" these rules, it just applies a small library of them.
Repeated application of the chain rule. The overall derivative of the composed function is assembled by chaining these local derivatives together, exactly as the chain rule prescribes, following either the order the program executed (forward mode) or the reverse of that order (reverse mode).
Primal computation and derivative propagation. AD systems compute two things side by side: the primal values (the actual numeric results of the program, what you'd get running it normally) and the derivative-carrying values (tangents in forward mode, adjoints in reverse mode). This is the essential reason AD differentiates the executed computation rather than a static formula: primal and derivative information travel together through the exact sequence of operations the program performs on the exact inputs supplied, including any control flow that depends on those inputs.
This distinction — differentiating the trace of what actually ran, not a symbolic description of what could run — is why AD naturally handles loops, conditionals, and recursion in a way symbolic differentiation of a fixed formula generally cannot, though it also means an AD-computed derivative reflects only the particular execution path taken for the given inputs.
A Full Worked Example
To make all of this concrete, consider a function with two inputs and a mix of elementary operations:
$$f(x_1, x_2) = \sin(x_1 x_2) + x_1^2$$
Evaluate and differentiate this at $x_1 = 2$, $x_2 = 3$.
Step 1: Decompose into intermediate variables.
$$v_1 = x_1, \quad v_2 = x_2, \quad v_3 = v_1 \cdot v_2, \quad v_4 = \sin(v_3), \quad v_5 = v_1^2, \quad v_6 = v_4 + v_5 = f$$
Step 2: Forward (primal) evaluation.
Variable | Expression | Value |
$v_1$ | $x_1$ | $2$ |
$v_2$ | $x_2$ | $3$ |
$v_3$ | $v_1 v_2$ | $6$ |
$v_4$ | $\sin(v_3)$ | $\sin(6) \approx -0.279415$ |
$v_5$ | $v_1^2$ | $4$ |
$v_6$ | $v_4 + v_5$ | $\approx 3.720585$ |
Step 3: Local derivatives at each node.
Node | Local derivative rule | Value at this point |
$v_3 = v_1 v_2$ | $\frac{\partial v_3}{\partial v_1} = v_2$, $\frac{\partial v_3}{\partial v_2} = v_1$ | $3$, $2$ |
$v_4 = \sin(v_3)$ | $\frac{\partial v_4}{\partial v_3} = \cos(v_3)$ | $\cos(6) \approx 0.960170$ |
$v_5 = v_1^2$ | $\frac{\partial v_5}{\partial v_1} = 2 v_1$ | $4$ |
$v_6 = v_4 + v_5$ | $\frac{\partial v_6}{\partial v_4} = 1$, $\frac{\partial v_6}{\partial v_5} = 1$ | $1$, $1$ |
Step 4: Analytical check. By ordinary calculus, $\frac{\partial f}{\partial x_1} = x_2 \cos(x_1 x_2) + 2x_1$ and $\frac{\partial f}{\partial x_2} = x_1 \cos(x_1 x_2)$. Substituting $x_1=2$, $x_2=3$:
$$\frac{\partial f}{\partial x_1} = 3 \cdot \cos(6) + 4 \approx 3(0.960170) + 4 \approx 6.880511$$ $$\frac{\partial f}{\partial x_2} = 2 \cdot \cos(6) \approx 1.920341$$
These two numbers are the targets both AD modes must reproduce.
The Same Example in Forward Mode
Forward mode seeds one input direction and propagates a tangent alongside every primal value, using the notation $\dot v = \frac{\partial v}{\partial x_1}$ (seeding $\dot x_1 = 1$, $\dot x_2 = 0$ to get the derivative with respect to $x_1$):
Variable | Primal | Tangent computation | Tangent value |
$v_1$ | $2$ | $\dot v_1 = 1$ (seed) | $1$ |
$v_2$ | $3$ | $\dot v_2 = 0$ (seed) | $0$ |
$v_3$ | $6$ | $\dot v_3 = \dot v_1 v_2 + v_1 \dot v_2$ | $1(3) + 2(0) = 3$ |
$v_4$ | $-0.279415$ | $\dot v_4 = \cos(v_3)\dot v_3$ | $0.960170 \times 3 \approx 2.880511$ |
$v_5$ | $4$ | $\dot v_5 = 2v_1 \dot v_1$ | $4 \times 1 = 4$ |
$v_6$ | $3.720585$ | $\dot v_6 = \dot v_4 + \dot v_5$ | $2.880511 + 4 = 6.880511$ |
This matches the analytical result for $\frac{\partial f}{\partial x_1}$ exactly. A second forward sweep, seeding $\dot x_2 = 1$ and $\dot x_1 = 0$, would recover $\frac{\partial f}{\partial x_2} \approx 1.920341$.
The Same Example in Reverse Mode
Reverse mode runs the primal forward pass once (as in Step 2), storing intermediate values, then propagates adjoints $\bar v = \frac{\partial f}{\partial v}$ backward from the output:
Variable | Adjoint computation | Adjoint value |
$\bar v_6$ | Seed: $\frac{\partial f}{\partial v_6} = 1$ | $1$ |
$\bar v_4$ | $\bar v_6 \cdot \frac{\partial v_6}{\partial v_4} = \bar v_6 \cdot 1$ | $1$ |
$\bar v_5$ | $\bar v_6 \cdot \frac{\partial v_6}{\partial v_5} = \bar v_6 \cdot 1$ | $1$ |
$\bar v_3$ | $\bar v_4 \cdot \cos(v_3)$ | $1 \times 0.960170 = 0.960170$ |
$\bar v_1$ (partial, from $v_5$) | $\bar v_5 \cdot 2v_1$ | $1 \times 4 = 4$ |
$\bar v_1$ (partial, from $v_3$) | $\bar v_3 \cdot v_2$ | $0.960170 \times 3 = 2.880511$ |
$\bar v_1$ (total) | Sum of contributions | $4 + 2.880511 = 6.880511$ |
$\bar v_2$ | $\bar v_3 \cdot v_1$ | $0.960170 \times 2 = 1.920341$ |
A single reverse sweep produces both $\frac{\partial f}{\partial x_1} \approx 6.880511$ and $\frac{\partial f}{\partial x_2} \approx 1.920341$ at once, matching the forward-mode result and the analytical derivative. Notice that $v_1$ received contributions from two separate downstream paths ($v_3$ and $v_5$); reverse mode simply sums contributions arriving at a node from every path that used it, which is the operational meaning of the multivariate chain rule.
What is stored and reused: forward mode never needs to store the full trace — it can discard primal and tangent values for a node once its consumers have used them, so memory use tracks the width of the computation, not its length. Reverse mode must store every intermediate primal value computed on the forward pass (the "tape"), because the backward sweep needs them to evaluate local derivatives like $\cos(v_3)$ after the fact. This asymmetry is the origin of reverse mode's higher memory cost and is discussed further below.
Forward-Mode Automatic Differentiation
Forward-mode AD propagates a tangent value $\dot v$ alongside every primal value $v$ as the program executes, representing the derivative of $v$ with respect to a chosen input direction. The technique can be implemented cleanly using dual numbers, an extension of the reals of the form $v + \dot v \varepsilon$ where $\varepsilon^2 = 0$; arithmetic on dual numbers automatically produces correct tangent propagation, since $(v_1 + \dot v_1 \varepsilon)(v_2 + \dot v_2 \varepsilon) = v_1 v_2 + (\dot v_1 v_2 + v_1 \dot v_2)\varepsilon$ is exactly the product rule.
To compute a derivative with respect to a specific input, seed that input's tangent to $1$ and all others to $0$, then run the program once. The output tangent is the directional derivative along that seed direction — a Jacobian-vector product, written $Jv$. One forward sweep computes $Jv$ for one direction $v$.
Cost relative to input directions: because each sweep produces one column-direction's worth of derivative information, computing the full Jacobian of a function with $n$ inputs requires $n$ forward sweeps, one per standard basis direction. This makes forward mode most efficient when the number of independent inputs is small relative to the number of outputs — for instance, differentiating a function $\mathbb{R}^2 \to \mathbb{R}^{1000}$ takes only two forward sweeps to get the full Jacobian.
Advantages: low memory overhead (no need to store a full trace for later replay), simple implementation via dual numbers or operator overloading, and natural handling of Jacobian-vector products without ever materializing the full Jacobian matrix.
Disadvantages: cost scales linearly with the number of independent input directions, so it becomes expensive for functions with many inputs and few outputs — exactly the shape of a typical loss function in deep learning, which has millions of parameters (inputs) and one scalar loss (output).
A minimal Python-style illustration using a simple dual-number class (executable, using only the standard library):
import math
class Dual:
def __init__(self, val, tan=0.0):
self.val, self.tan = val, tan
def __add__(self, other):
other = other if isinstance(other, Dual) else Dual(other)
return Dual(self.val + other.val, self.tan + other.tan)
def __mul__(self, other):
other = other if isinstance(other, Dual) else Dual(other)
return Dual(self.val * other.val,
self.tan * other.val + self.val * other.tan)
def dual_sin(d):
return Dual(math.sin(d.val), math.cos(d.val) * d.tan)
# f(x1, x2) = sin(x1 * x2) + x1**2, derivative w.r.t. x1
x1 = Dual(2.0, 1.0) # seed: tangent = 1 for x1
x2 = Dual(3.0, 0.0) # seed: tangent = 0 for x2
result = dual_sin(x1 * x2) + x1 * x1
print(result.val, result.tan) # matches df/dx1 from the worked exampleThis snippet mirrors the forward-mode column of the worked example above: result.tan after running is the same $\approx 6.880511$ computed by hand. Real systems, such as JAX's jax.jvp transformation, implement the same underlying idea with far more operation coverage, hardware acceleration, and support for vector and tensor-valued programs [4].
Reverse-Mode Automatic Differentiation
Reverse-mode AD works in two passes. The forward pass runs the program normally, computing primal values and recording the operations performed — a record often called a tape. The reverse pass then walks the tape backward from the output, propagating adjoint values $\bar v = \frac{\partial f}{\partial v}$, the sensitivity of the final output to each intermediate variable.
Seeding the output adjoint to $1$ and propagating backward computes a vector-Jacobian product, written $v^{\mathsf T} J$. JAX's documentation frames this precisely: <cite index="59-1">the linear part of a VJP as the transpose (or adjoint conjugate) of the linear part of a JVP</cite> [5]. One reverse sweep, seeded on a scalar output, computes the gradient with respect to every input simultaneously — exactly why reverse mode dominates in machine learning, where loss functions are scalar and parameter counts are enormous.
Accumulating contributions from multiple paths: as demonstrated in the worked example, when an intermediate variable feeds into more than one downstream computation, its adjoint is the sum of the contributions flowing back from each path. This summation is the reverse-mode expression of the multivariable chain rule and requires no special handling — it falls out naturally from processing the tape in reverse topological order.
Memory requirements: because the reverse pass needs primal values computed during the forward pass to evaluate local derivatives, reverse mode must retain the entire computation trace (or enough of it to reconstruct needed values) until the backward sweep completes. This is the central memory-versus-computation trade-off in reverse-mode AD: for very deep or very long computations, the tape can become large, motivating techniques like gradient checkpointing, discussed in the advanced section below.
Advantages: a single backward sweep yields the full gradient of a scalar function with respect to arbitrarily many inputs, which is exactly the shape of a training loss. This efficiency is why reverse mode is central to training neural networks with billions of parameters against one scalar loss.
Disadvantages: the tape must be stored, so memory use grows with the depth and size of the computation; the technique is less naturally suited to functions with many outputs relative to inputs, where forward mode wins; and reverse mode is somewhat more complex to implement, since it requires either recording a full execution trace or performing a source-level program transformation.
A minimal, illustrative reverse-mode sketch in Python (pseudocode-level, for conceptual clarity rather than production use):
class Node:
def __init__(self, val, children=()):
self.val = val
self.children = children # list of (parent_node, local_derivative)
self.grad = 0.0
def add(a, b):
return Node(a.val + b.val, [(a, 1.0), (b, 1.0)])
def mul(a, b):
return Node(a.val * b.val, [(a, b.val), (b, a.val)])
def backward(node, seed=1.0):
node.grad += seed
for parent, local_deriv in node.children:
backward(parent, seed * local_deriv)
# f = x1 * x2, reverse-mode gradient
x1, x2 = Node(2.0), Node(3.0)
f = mul(x1, x2)
backward(f)
print(x1.grad, x2.grad) # 3.0, 2.0 — matches df/dx1 and df/dx2This is conceptually how PyTorch's torch.autograd engine operates at a high level: it constructs a dynamic directed acyclic graph during the forward pass and calls .backward() methods in reverse topological order to propagate gradients [7].
Forward Mode vs. Reverse Mode
Aspect | Forward mode | Reverse mode |
What it computes per sweep | Jacobian-vector product ($Jv$) | Vector-Jacobian product ($v^{\mathsf T}J$) |
Best input/output shape | Few inputs, many outputs | Many inputs, few outputs (especially one scalar) |
Sweeps to get full Jacobian | One per input dimension | One per output dimension |
Memory profile | Low; no full tape required | Higher; tape must be retained for the backward pass |
Typical application | Sensitivity analysis with few parameters, forward simulations, Jacobian-vector products in iterative solvers | Neural network training, any scalar-loss optimization |
Relationship to backpropagation | Not directly used for backprop | The generalized technique underlying backpropagation |
A practical decision guide: if a function has substantially more outputs than inputs, favor forward mode. If it has substantially more inputs than outputs — the near-universal case for training loss functions in machine learning, where there is one scalar loss and potentially billions of parameters — favor reverse mode. Neither mode is categorically faster; the right choice depends entirely on the shape of the function being differentiated, a caveat both the JAX documentation and the Baydin et al. survey are careful to state rather than asserting a universal winner [1] [4].
Mixed mode becomes valuable for structures like the Hessian, the matrix of second derivatives. JAX's autodiff cookbook demonstrates computing a Hessian-vector product efficiently by composing forward-mode over reverse-mode: differentiate once with reverse mode to get the gradient function, then apply forward mode to that gradient function to get its Jacobian-vector product in one additional sweep, avoiding ever materializing the full Hessian matrix <cite index="56-1">That's efficient, but we can do even better and save some memory by using forward-mode together with reverse-mode</cite> [4]. This "forward-over-reverse" composition is a standard pattern for Hessian-vector products in scientific and machine learning code.
Automatic Differentiation and Backpropagation
Backpropagation and automatic differentiation are related but not identical concepts, and precision here matters. Backpropagation, introduced to the machine learning community in its modern form by Rumelhart, Hinton, and Williams in their 1986 Nature paper, describes <cite index="108-1">a new learning procedure...for networks of neurone-like units</cite> that <cite index="108-1">repeatedly adjusts the weights of the connections in the network so as to minimize a measure of the difference between the actual output vector of the net and the desired output vector</cite> [9]. In practice, backpropagation is reverse-mode automatic differentiation applied specifically to the layered computational structure of a neural network, where a loss function depends on parameters through a chain of layer transformations.
The Baydin et al. survey is explicit that AD is the more general concept: it describes AD as a family of techniques <cite index="47-1">similar to but more general than backpropagation</cite>, precisely because reverse-mode AD applies to any composition of differentiable elementary operations, not only to the specific layer-by-layer structure of a conventional feedforward network [1]. Backpropagation is what you get when you specialize reverse-mode AD to a particular, common computational pattern; reverse-mode AD is the general mathematical machinery, applicable to arbitrary programs, not just networks organized into layers.
Weight sharing and reused values. In architectures where the same parameter is reused multiple times in a single forward pass — convolutional filters applied at many spatial locations, or recurrent weights reused across time steps — reverse mode's summation-of-contributions behavior (demonstrated in the worked example above) automatically and correctly accumulates the gradient contributions from every use of that shared parameter. This is not a special case requiring extra code; it falls directly out of how the chain rule handles a variable with multiple downstream consumers.
Beyond conventional neural networks. Reverse-mode AD generalizes past standard layered architectures to differentiate through control flow, iterative numerical solvers, physics simulators, and probabilistic models — anywhere a scalar objective depends on many inputs through a chain of differentiable operations. This generalization is precisely why frameworks built for AD, like JAX, are used well outside conventional neural network training, in areas like scientific computing and probabilistic programming, discussed further below.
How AD Systems Are Actually Built
There are several distinct engineering strategies for implementing AD, each with real trade-offs between flexibility, performance, and debuggability.
Operator overloading. The AD system defines a custom numeric type (like the Dual class shown earlier) that overloads standard arithmetic operators, so each operation both computes the primal result and propagates derivative information. This is how PyTorch's eager torch.autograd and JAX's tracing both work conceptually, integrating naturally with a host language's control flow.
Source transformation. Some AD tools instead analyze source code (or an intermediate representation) ahead of time and generate new derivative-computing code directly. This can produce highly optimized derivative code but needs more compiler-level tooling — LLVM Enzyme, which operates on compiled intermediate representations, works this way.
Tracing and tapes. In tracing-based systems, running a function once produces a tape: a concrete record of the operations performed, which reverse-mode AD replays backward. TensorFlow's tf.GradientTape makes this explicit in its name and API: it <cite index="70-1">"records" relevant operations executed inside the context of a tf.GradientTape onto a "tape"</cite> and <cite index="70-2">then uses that tape to compute the gradients of a "recorded" computation using reverse mode differentiation</cite> [8].
Static versus dynamic graphs. A static graph is built once and executed repeatedly, enabling aggressive optimization but making data-dependent shapes harder to express. A dynamic graph is rebuilt each run, trading some optimization for flexibility. PyTorch's documentation is explicit that its graphs are rebuilt from scratch: <cite index="63-1">DAGs are dynamic in PyTorch...the graph is recreated from scratch; after each .backward() call, autograd starts populating a new graph</cite> [6]. JAX traces functions into an intermediate representation that can still be compiled ahead of time, blending both approaches.
Primitives, custom rules, and gradient stopping. Every AD system rests on a library of primitive operations, each paired with a known derivative rule; complex programs are differentiated by decomposing them into these primitives. When code calls something the system cannot trace — an external library or a numerically unstable naive derivative — developers can supply a custom rule directly, as with JAX's jax.custom_jvp and jax.custom_vjp decorators, which <cite index="60-1">define custom differentiation rules for Python functions that are already JAX-transformable</cite> [12]. And sometimes a value should affect the forward result without receiving gradient flow; PyTorch's .detach() <cite index="67-1">creates a new tensor that shares the same underlying data storage as the original tensor but is explicitly detached from the computational graph history</cite> [7], with TensorFlow and JAX offering equivalent stop_gradient mechanisms.
The trade-off is consistent across strategies: compile-time structure (static graphs, source transformation) buys speed at some cost to flexibility and debuggability; runtime flexibility (dynamic graphs, eager overloading) buys ease of use at some cost to peak performance.
Automatic Differentiation in JAX, PyTorch, and TensorFlow
JAX
JAX exposes forward-mode AD through jax.jvp and reverse-mode AD through jax.vjp, with jax.grad built as a convenience wrapper around reverse mode for scalar-output functions. This is a reverse-mode example, computing a gradient:
import jax
import jax.numpy as jnp
def f(x):
return jnp.sin(x[0] * x[1]) + x[0]**2
x = jnp.array([2.0, 3.0])
grad_f = jax.grad(f)
print(grad_f(x)) # reverse-mode gradient, matches the worked exampleThis differentiates f, a function from a 2-vector to a scalar, using reverse mode — jax.grad is documented as returning the gradient of a scalar-valued function, consistent with JAX's autodiff cookbook description of grad as built on top of vjp internally [4]. For forward mode directly:
import jax
def f(x1, x2):
return jnp.sin(x1 * x2) + x1**2
primal_out, tangent_out = jax.jvp(f, (2.0, 3.0), (1.0, 0.0))
print(primal_out, tangent_out) # tangent_out matches df/dx1jax.jvp computes a Jacobian-vector product directly, seeding the tangent for x1 to 1.0 and x2 to 0.0, exactly mirroring the forward-mode worked example above.
PyTorch
PyTorch's default and primary mechanism is reverse mode, exposed through torch.autograd:
import torch
x1 = torch.tensor(2.0, requires_grad=True)
x2 = torch.tensor(3.0, requires_grad=True)
f = torch.sin(x1 * x2) + x1**2
f.backward()
print(x1.grad, x2.grad) # matches df/dx1 and df/dx2 from the worked exampleCalling .backward() triggers reverse-mode traversal of the dynamically built graph. PyTorch's own tutorial documentation frames this precisely: <cite index="63-2">To compute those gradients, PyTorch has a built-in differentiation engine called torch.autograd. It supports automatic computation of gradient for any computational graph</cite> [6]. PyTorch also provides a forward-mode AD API (via torch.func.jvp or the lower-level forward-mode context managers), which the documentation explicitly labels a beta API with less complete operator coverage than reverse mode [7].
TensorFlow
TensorFlow exposes reverse-mode AD through tf.GradientTape:
import tensorflow as tf
x1 = tf.Variable(2.0)
x2 = tf.Variable(3.0)
with tf.GradientTape() as tape:
f = tf.sin(x1 * x2) + x1**2
grad = tape.gradient(f, [x1, x2])
print(grad) # matches df/dx1 and df/dx2 from the worked exampleThe official guide states plainly that <cite index="69-1">TensorFlow provides the tf.GradientTape API for automatic differentiation; that is, computing the gradient of a computation with respect to some inputs, usually tf.Variables</cite>, and that TensorFlow uses the recorded tape <cite index="70-2">to compute the gradients of a "recorded" computation using reverse mode differentiation</cite> [8]. This example represents reverse mode; TensorFlow does not expose a general-purpose forward-mode API in the same central way that JAX and PyTorch do.
All three examples above compute the same gradient values derived analytically and by hand in the worked example section, which is a useful sanity check when learning any of these frameworks: hand-verify a small case, then trust the framework on larger ones.
Other notable AD systems worth knowing include Julia's ForwardDiff.jl and Zygote.jl, and LLVM Enzyme, which performs AD by transforming compiled LLVM intermediate representation rather than operating at the level of a host language's operators — a fundamentally different implementation strategy from the tracing and operator-overloading approaches above.
Higher-Order and Advanced Differentiation
Second derivatives and Hessians. The Hessian of $f: \mathbb{R}^n \to \mathbb{R}$ is the $n \times n$ matrix of second partial derivatives. Forming a full Hessian is expensive for large $n$, but many applications only need its action on a vector.
Hessian-vector products. A Hessian-vector product $Hv$ can be computed by differentiating the gradient function itself, without ever forming $H$. Pearlmutter's 1994 Neural Computation paper established this technique, showing a method that <cite index="88-1">directly calculates Hv, where v is an arbitrary vector</cite>, producing a result that is <cite index="88-2">exact and numerically stable...which takes about as much computation, and is about as local, as a gradient evaluation</cite> [3] — roughly the cost of one extra gradient evaluation, not a full Hessian.
Forward-over-reverse composition is the standard, memory-efficient modern pattern for Hessian-vector products, directly reflecting Pearlmutter's approach [3] [4]. Other nestings (reverse-over-forward, and so on) are also possible with different trade-offs.
Nested differentiation, implicit differentiation, and differentiable simulation. AD systems supporting arbitrary composition can differentiate a function that itself contains a differentiation operation — useful for meta-learning or differentiating through an inner optimization loop — but each nesting level multiplies compute and memory cost; it is not free simply because the syntax allows it. When a quantity is defined implicitly, as the solution to an equation or optimization problem, AD can still supply gradients via the implicit function theorem, differentiating the optimality conditions rather than every solver iteration — often far cheaper than naively unrolling the solver. This underlies a growing body of work differentiating entire simulators end to end for gradient-based design.
Checkpointing (rematerialization). For very deep or long computations, storing the full reverse-mode tape can exceed memory. Checkpointing stores only a subset of forward-pass values and recomputes the rest during the backward pass — a deliberate trade of extra compute for reduced memory, not a free win.
Differentiable programming. Taken to its conclusion, AD lets entire programs, not just formulas, be treated as differentiable objects optimized end to end — the framing Baydin et al. use to connect AD to programming-language research and compiler design [1]. None of these advanced techniques should be assumed cheap: higher-order and nested differentiation carry real, scaling costs.
Practical Limitations and Failure Modes
AD is powerful, but it is not magic, and it does not remove the need to understand what a computation is actually doing.
Floating-point error. AD computes exact derivatives of the floating-point computation actually performed, but that computation is still subject to ordinary rounding. Long operation chains can accumulate rounding error like any numerical computation.
Nondifferentiable operations and subgradients. Functions like abs, max, and ReLU-style activations are mathematically nondifferentiable at specific points. Frameworks handle this with a convention — returning a subgradient, one valid member of the generalized-derivative set at that point — rather than raising an error. This is a convention, not proof the chosen value is uniquely "correct."
Discontinuous control flow. If an if statement's outcome depends on inputs in a way that creates a genuine discontinuity, AD's returned derivative reflects only the branch actually executed — it cannot capture the discontinuity itself.
Detached tensors, in-place mutation, and missing rules. Unintentional detaching is a common source of silently zero gradients. Modifying a tensor's data in place after it feeds the forward pass can corrupt values a reverse-mode tape needs later; PyTorch's documentation warns that <cite index="65-1">autograd's aggressive buffer freeing and reuse makes it very efficient and there are very few occasions when in-place operations actually lower memory usage by any significant amount</cite> [7]. Calls into native code or custom operations without a registered derivative rule simply have no gradient unless one is supplied.
Stateful and stochastic operations. Random sampling needs special handling, such as the reparameterization trick, to be differentiable with respect to distribution parameters, since raw randomness is not itself a differentiable function of those parameters.
Memory pressure. Reverse mode's memory cost scales with computation length, and very deep networks or long simulation rollouts can exhaust memory without checkpointing.
Vanishing and exploding gradients. Multiplying many small or large local derivatives across a deep chain can shrink the accumulated gradient toward zero or let it grow unbounded. This reflects the mathematics of the function being differentiated, not an AD flaw.
Broadcasting and complex numbers. Shape mismatches "fixed" by broadcasting can silently accumulate gradients over the wrong dimensions. Differentiating complex-valued, non-holomorphic functions is genuinely subtler, with conventions still varying across frameworks [55].
Symptom | Likely cause | Diagnostic step | Possible remedy |
Gradient is exactly zero | Value was detached or stopped from the graph, or passed through a non-differentiable integer operation | Check for .detach(), stop_gradient, or integer dtypes along the path | Remove unnecessary detachment; recast to floating point where appropriate |
Gradient is None or missing | The tensor never had requires_grad set, or fell outside the traced graph | Confirm requires_grad=True or that the tensor was produced inside the tape/trace context | Explicitly mark inputs as differentiable before the forward pass |
Gradient is NaN | Numerical overflow, division by zero, or a log/sqrt evaluated at an invalid point during the backward pass | Check intermediate values for extreme magnitudes; test the forward pass with smaller inputs | Add numerical stabilization (clipping, epsilon terms) at the offending operation |
Gradient is much larger or smaller than expected | Vanishing or exploding gradients from a long chain of small or large local derivatives | Inspect per-layer gradient magnitudes across the chain | Use normalization, gradient clipping, or architectural changes suited to the underlying model |
Gradient check fails by a large margin | An incorrect custom derivative rule, or an unintended in-place mutation corrupting saved values | Re-run with a simpler, isolated test case around the suspected operation | Fix the custom JVP/VJP, or remove the offending in-place operation |
Training runs but never improves | Gradient exists but flows through the wrong variable due to a broadcasting or indexing bug | Verify shapes at each step; compute the gradient of a known, hand-checkable toy example | Correct shape alignment; add explicit shape assertions during development |
Gradient Checking
Gradient checking is a diagnostic technique, not a replacement for AD: it uses finite differences to sanity-check that an AD-computed (or hand-derived) gradient is plausibly correct, especially after implementing a custom derivative rule.
The standard approach uses the central difference formula, which is more accurate than a one-sided difference for a given step size:
$$\frac{\partial f}{\partial x_i} \approx \frac{f(x + h e_i) - f(x - h e_i)}{2h}$$
where $e_i$ is the standard basis vector in the $i$-th direction and $h$ is a small step size, commonly around $10^{-4}$ to $10^{-6}$ for double-precision arithmetic — small enough to keep truncation error low, but not so small that floating-point cancellation dominates the result.
Compare the finite-difference estimate to the AD-computed value using both absolute and relative tolerances, since a purely absolute tolerance is too strict for large-magnitude gradients and too loose for small ones. A common convergence check is:
$$\frac{|g_{\text{AD}} - g_{\text{numerical}}|}{\max(|g_{\text{AD}}|, |g_{\text{numerical}}|, \epsilon)} < \text{tolerance}$$
Why finite-difference checks can fail even when AD is correct: step size selection is delicate (too large introduces truncation bias, too small introduces cancellation error); functions with sharp curvature or near-nondifferentiable points can produce misleading finite-difference estimates; and stochastic or stateful operations do not have a stable finite-difference baseline to compare against in the first place. A failed gradient check is a signal to investigate, not automatic proof that the AD-computed gradient is wrong.
Compact example: checking the $\frac{\partial f}{\partial x_1}$ result from the worked example above, at $h = 10^{-5}$:
import math
def f(x1, x2):
return math.sin(x1 * x2) + x1**2
h = 1e-5
x1, x2 = 2.0, 3.0
numerical = (f(x1 + h, x2) - f(x1 - h, x2)) / (2 * h)
print(numerical) # approximately 6.880511, matching the AD resultThis should closely match the analytical and AD-computed value of approximately $6.880511$ derived earlier, confirming the AD implementation is behaving correctly on this small, hand-checkable case.
Because gradient checking is primarily a testing and debugging technique, it is best used sparingly during development on small, representative sub-computations — not as a routine part of a training loop, where its extra function evaluations would be prohibitively expensive at scale.
Real Applications Beyond Standard Neural Networks
Neural network training remains the most visible application: reverse-mode AD, applied as backpropagation, computes the gradients that optimizers like stochastic gradient descent or Adam use to update weights.
Scientific machine learning and physics-informed models. Physics-informed neural networks, introduced by Raissi, Perdikaris, and Karniadakis, embed differential equation constraints directly into a network's loss function, using AD to compute the exact spatial and temporal derivatives the governing equations require — the network's derivatives with respect to inputs like time and space <cite index="119-1">can be derived by applying the chain rule for differentiating compositions of functions using automatic differentiation</cite> [10]. This lets a single trained network satisfy a PDE's residual at sampled points without discretizing the domain the way classical solvers do.
Computational fluid dynamics and engineering design optimization are AD's original, pre-machine-learning home; the Baydin et al. survey lists these as long-standing application areas that predate AD's adoption in machine learning [1].
Probabilistic programming, differentiable rendering, and robotics. Many probabilistic programming systems use AD to compute gradients of variational objectives, often paired with reparameterization tricks to make sampling differentiable. Differentiable rendering lets gradients flow from a rendered image back to scene parameters like geometry and lighting, enabling optimization-based 3D reconstruction. Differentiable robotics simulators similarly let gradients flow from a control objective back through simulated physics to controller parameters, supporting gradient-based policy optimization alongside purely sample-based reinforcement learning.
Differential equations, parameter estimation, and meta-learning. AD underpins differentiable ODE and PDE solvers used to infer unknown physical constants from observed data, and it powers meta-learning and hyperparameter optimization techniques that differentiate through an inner training loop, relying on the nested-differentiation capabilities described above.
Common Misconceptions
"Automatic differentiation is the same as symbolic differentiation." They are both exact, but AD evaluates numerically at a specific point using the program's actual execution trace, while symbolic differentiation manipulates and returns an explicit formula. AD does not suffer from the expression-swell problem that can afflict symbolic approaches on long compositions.
"It estimates derivatives by perturbing inputs." That description applies to numerical differentiation (finite differences), not AD. AD applies exact differentiation rules to each elementary operation; it does not measure output changes from small input perturbations.
"Backpropagation and automatic differentiation always mean exactly the same thing." As discussed above, backpropagation is best understood as reverse-mode AD specialized to the layered structure of a conventional neural network; reverse-mode AD is the more general underlying technique [1].
"Reverse mode is always better." Reverse mode dominates in machine learning specifically because loss functions are typically scalar with vast numbers of parameters. For functions with many outputs and few inputs, forward mode is usually more efficient, and mode selection should follow the shape of the function, not a blanket rule.
"AD removes the need to understand calculus." AD automates the mechanical bookkeeping of the chain rule; it does not remove the need to understand what a gradient means, why a function might be ill-conditioned, why a particular operation is nondifferentiable, or how to diagnose a numerically unstable computation.
"Any program can be differentiated without restrictions." As the limitations section makes clear, control flow, in-place mutation, external native code, stochastic operations, and missing derivative rules can all break or complicate differentiation. The scope of what an AD system can actually handle depends on its specific implementation.
"An exact derivative is automatically numerically stable." AD produces the exact derivative of the floating-point computation it is given, but that does not guarantee the underlying computation is well-conditioned. A mathematically exact gradient of a poorly conditioned function can still behave badly under further arithmetic, such as vanishing or exploding through a long chain.
"Using AD makes optimization problems convex or easy." AD supplies gradients; it says nothing about the shape of the loss landscape those gradients are computed on. Nonconvex, poorly conditioned, or highly non-smooth objectives remain exactly as difficult to optimize as they would be with hand-derived gradients — AD only removes the burden of deriving and coding the gradient formula itself.
Choosing and Using an AD System
Several practical factors should guide the choice and use of an AD system for a given project:
Input/output shape and higher-order needs. Favor reverse-mode-oriented tools (most deep learning frameworks) for many-input, few-output problems. If Hessian-vector products or nested differentiation matter, confirm the framework composes forward and reverse transformations, as JAX does via its jvp/vjp model [4].
Control flow and compilation. Dynamic, data-dependent control flow is easiest in eager systems like PyTorch; static, compiled graphs trade some flexibility for optimization. If deployment targets specialized hardware or kernel fusion, ahead-of-time compilation paths like JAX's jit or TensorFlow's graph mode can matter.
Ecosystem, custom operations, and memory. Weigh the availability of pretrained models and deployment tooling, not just the AD engine. Confirm custom-JVP/VJP or custom-autograd-function support if hand-derived rules are needed for stability or performance, and check for built-in gradient checkpointing if models or rollouts are deep enough for tape memory to become the bottleneck.
Hardware, debugging, and interoperability. Verify GPU or accelerator support for the target hardware. Eager execution generally debugs more easily with standard tools than compiled graphs. If the project calls into existing Python or NumPy code, confirm which calls can be traced and which need custom rules or a rewrite.
Where This Is Heading
Automatic differentiation is no longer a niche numerical-computing technique; it has become foundational infrastructure connecting machine learning, scientific computing, and programming-language design. The broader idea the Baydin et al. survey terms "differentiable programming" treats gradient information as a first-class property of software itself, not something bolted on afterward for a specific model [1]. Compiler research continues to explore tighter integration of AD into general-purpose language toolchains, exemplified by projects like LLVM Enzyme that operate at the level of compiled intermediate representations rather than a single host language's operators. In scientific computing, differentiable simulators and solvers are expanding the reach of gradient-based optimization into domains — fluid dynamics, structural design, control — that historically relied on gradient-free or hand-derived-gradient methods. None of this suggests derivatives are becoming "solved" in some final sense; rather, the tooling for computing them correctly, efficiently, and at scale continues to mature alongside the models and simulations that depend on it.
FAQ
What is automatic differentiation in simple terms?
Automatic differentiation is a technique computers use to calculate exact derivatives of functions written as code. It breaks the function into small steps and applies the chain rule to each step in the order the program actually runs. The result is an exact derivative, not an estimate, computed from the program's own operations.
Is automatic differentiation the same as backpropagation?
Not exactly. Backpropagation is reverse-mode automatic differentiation applied specifically to the layered structure of a neural network. Reverse-mode AD is the more general technique, applicable to any program built from differentiable elementary operations, neural network or not.
How is automatic differentiation different from numerical differentiation?
Numerical differentiation estimates a derivative from a finite difference between nearby function evaluations, which introduces approximation error tied to the chosen step size. Automatic differentiation applies exact rules to each elementary operation a program performs, giving a result accurate to floating-point precision rather than an approximation.
How is automatic differentiation different from symbolic differentiation?
Symbolic differentiation manipulates a formula algebraically and returns a new formula, which can grow very large — expression swell. Automatic differentiation evaluates the program numerically at a specific point, carrying derivative information alongside the computation, avoiding expression swell and working directly on real code with loops and conditionals.
What are forward mode and reverse mode?
Forward mode propagates a tangent value in the same order the program executes, and is efficient for functions with few inputs relative to outputs. Reverse mode records the computation, then propagates adjoint values backward, and is efficient for many inputs relative to outputs — the typical shape of a machine learning loss function.
Why is reverse mode used in deep learning?
Training involves a scalar loss depending on potentially billions of parameters. Reverse mode computes the gradient with respect to all of them in a single backward sweep, whereas forward mode would need one sweep per parameter — computationally infeasible at that scale.
Is automatic differentiation exact?
Yes, in that it computes the exact derivative of the specific floating-point computation the program performs, using exact rules per elementary operation rather than approximating with finite differences. It remains subject to ordinary floating-point rounding, the same as any numerical computation.
Can automatic differentiation handle loops and conditionals?
It depends on the implementation. Tracing-based systems differentiate the specific execution path taken for given inputs, which naturally handles loops and conditionals the system can trace through; some require special constructs for control flow that would otherwise break tracing.
Can every function be automatically differentiated?
No. Genuine mathematical discontinuities, operations without a registered derivative rule, calls into opaque external code, and certain integer or stateful operations can all block or complicate differentiation. AD generally requires a program built substantially from operations it knows how to differentiate.
What is a computational graph?
A computational graph represents a computation as a directed graph, with nodes as variables (inputs, intermediate results, outputs) and edges as the elementary operations connecting them. AD systems use this structure, whether built explicitly or implicitly through tracing, to apply the chain rule systematically.
What is a Jacobian-vector product?
A Jacobian-vector product, written $Jv$, is a function's Jacobian matrix multiplied by a vector $v$, without ever forming the full Jacobian. It is exactly what one forward-mode AD sweep computes, seeded with direction $v$.
What is a vector-Jacobian product?
A vector-Jacobian product, written $v^{\mathsf T}J$, is a vector multiplied by a function's Jacobian from the left, again without forming the full matrix. It is exactly what one reverse-mode AD sweep computes, seeded with vector $v$ on the output side.
How are higher-order derivatives computed?
Higher-order derivatives, such as Hessians, are typically computed by composing AD with itself — applying forward mode to a function that is itself the result of reverse-mode differentiation, yielding a Hessian-vector product efficiently, as first shown by Pearlmutter in 1994. Forming a full, explicit Hessian is far more expensive and usually unnecessary.
Which programming frameworks support automatic differentiation?
Widely used frameworks include JAX, PyTorch, and TensorFlow, alongside Julia tools like ForwardDiff.jl and Zygote.jl, and compiler-level systems like LLVM Enzyme. Each implements AD with a different strategy, spanning operator overloading, tracing, and source-level transformation.
When should gradient checking be used?
Gradient checking, comparing an AD-computed gradient against a finite-difference estimate, is most useful when implementing or debugging a custom derivative rule, or verifying a new AD implementation on a small, hand-checkable case. It is a diagnostic technique, not something to run routinely inside a full training loop.
Key Takeaways
Automatic differentiation is exact bookkeeping of the chain rule applied to a program's actual elementary operations, not approximation and not symbolic formula manipulation.
Forward mode and reverse mode are two directions of propagating the same chain-rule information, and the right choice depends on whether a function has more inputs or more outputs.
Reverse mode's ability to compute a full gradient in one backward sweep is precisely why it underlies backpropagation and dominates neural network training at scale.
Backpropagation and automatic differentiation are related but distinct: backpropagation is reverse-mode AD specialized to layered neural network structures.
Memory, not just compute, is a first-class constraint in reverse-mode AD, which is why techniques like gradient checkpointing exist.
Common AD frameworks differ mainly in implementation strategy — dynamic tracing versus tapes versus source transformation — not in the underlying mathematics they compute.
Nondifferentiable operations, in-place mutation, stochastic operations, and missing derivative rules are the most common practical sources of broken or misleading gradients.
Gradient checking with finite differences remains a valuable diagnostic tool for validating custom derivative code, even though it is not a substitute for AD in production.
AD's reach extends well beyond deep learning into physics-informed modeling, differentiable simulation, robotics, and probabilistic programming, anywhere a scalar objective depends differentiably on many inputs.
Actionable Next Steps
Implement a basic finite-difference derivative for a simple one-variable function to build intuition for what numerical differentiation actually measures and where its error comes from.
Build the tiny dual-number forward-mode example from this guide by hand, then extend it to a function with three or four elementary operations of your own choosing.
Manually trace a small computational graph on paper for a function with at least two inputs and shared intermediate variables, then compute both forward-mode tangents and reverse-mode adjoints and confirm they agree.
Reproduce the worked example in this guide using JAX, PyTorch, and TensorFlow side by side, comparing each framework's API for the same reverse-mode gradient computation.
Run a gradient check against a small custom function using the central-difference formula given in this guide, and confirm the result matches your framework's automatic differentiation output within a reasonable tolerance.
Explore Jacobian-vector products and vector-Jacobian products directly, using jax.jvp and jax.vjp, to build intuition for what each mode actually returns before relying on higher-level wrappers like jax.grad.
Compute a Hessian-vector product using the forward-over-reverse composition shown in this guide, and compare its cost against explicitly forming a full Hessian for a small test function.
Write a custom derivative rule (a custom JVP or VJP, or a custom PyTorch autograd Function) for a simple operation, to understand how AD systems integrate hand-supplied derivatives into an otherwise automatic pipeline.
Glossary
Adjoint: In reverse-mode AD, the sensitivity of the final output to an intermediate variable, propagated backward.
Algorithmic differentiation: Another name for automatic differentiation.
Automatic differentiation (AD): Techniques for computing exact derivatives of programs by applying the chain rule to their elementary operations.
Backpropagation: Reverse-mode AD applied to a neural network's layered structure to compute gradients of a loss with respect to weights.
Chain rule: The derivative of a composition of functions is the product of the derivatives of each function in the composition.
Computational graph: A directed graph of a computation, with nodes as variables and edges as elementary operations.
Derivative: How much a function's output changes for a small change in its input.
Dual number: A number $v + \dot v \varepsilon$ with $\varepsilon^2 = 0$, used to implement forward-mode AD via ordinary arithmetic.
Evaluation trace: The recorded sequence of operations and intermediate values from running a program with specific inputs.
Expression swell: Rapid, unwieldy growth in a symbolic expression's size under repeated symbolic differentiation.
Forward mode: An AD strategy propagating tangents in execution order, efficient for few inputs relative to outputs.
Gradient: The vector of all partial derivatives of a scalar-valued function, pointing toward steepest increase.
Hessian: The matrix of second-order partial derivatives of a scalar-valued function.
Hessian-vector product: A Hessian matrix multiplied by a vector, computable via composed AD without forming the full Hessian.
Jacobian: The matrix of first-order partial derivatives of a vector-valued function.
Jacobian-vector product (JVP): A function's Jacobian times a vector, computed by one forward-mode AD sweep.
Local derivative: The derivative of a single elementary operation with respect to its immediate inputs.
Numerical differentiation: Estimating a derivative via finite differences between nearby function evaluations.
Primal value: The ordinary numeric result of a computation, as opposed to a derivative-carrying tangent or adjoint.
Reverse mode: An AD strategy recording a forward pass, then propagating adjoints backward; efficient for many inputs relative to outputs.
Subgradient: A generalized derivative used where a function is not classically differentiable, such as a kink in ReLU.
Symbolic differentiation: Computing a derivative by algebraic manipulation, producing a new formula rather than a number.
Tangent: In forward-mode AD, the derivative of a variable with respect to a seeded input direction.
Tape: A recorded trace of forward-pass operations used by reverse-mode AD for the backward pass.
Vector-Jacobian product (VJP): A vector times a function's Jacobian from the left, computed by one reverse-mode AD sweep.
Sources & References
Baydin, A.G., Pearlmutter, B.A., Radul, A.A., Siskind, J.M. (2018). Automatic Differentiation in Machine Learning: a Survey. Journal of Machine Learning Research, 18(153):1–43. Retrieved from https://jmlr.org/papers/v18/17-468.html
Griewank, A., Walther, A. (2008). Evaluating Derivatives: Principles and Techniques of Algorithmic Differentiation, 2nd ed. Society for Industrial and Applied Mathematics (SIAM). Retrieved from https://epubs.siam.org/doi/book/10.1137/1.9780898717761
Pearlmutter, B.A. (1994). Fast Exact Multiplication by the Hessian. Neural Computation, 6(1):147–160. Retrieved from https://doi.org/10.1162/neco.1994.6.1.147
The JAX Authors. The Autodiff Cookbook. JAX documentation. Retrieved from https://docs.jax.dev/en/latest/notebooks/autodiff_cookbook.html
The JAX Authors. Forward- and reverse-mode autodiff in JAX. JAX documentation. Retrieved from https://docs.jax.dev/en/latest/jacobian-vector-products.html
PyTorch. Automatic Differentiation with torch.autograd. PyTorch Tutorials. Retrieved from https://docs.pytorch.org/tutorials/beginner/basics/autogradqs_tutorial.html
PyTorch. Automatic differentiation package - torch.autograd. PyTorch documentation, version 2.11. Retrieved from https://docs.pytorch.org/docs/2.11/autograd.html
TensorFlow. Introduction to gradients and automatic differentiation. TensorFlow Core guide. Retrieved from https://www.tensorflow.org/guide/autodiff
Rumelhart, D.E., Hinton, G.E., Williams, R.J. (1986). Learning representations by back-propagating errors. Nature, 323:533–536. Retrieved from https://www.nature.com/articles/323533a0
Raissi, M., Perdikaris, P., Karniadakis, G.E. (2019). Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational Physics, 378:686–707. Retrieved from https://faculty.sites.iastate.edu/hliu/files/inline-files/PINN_RPK_2019_1.pdf
Google Search Central. Mark Up FAQs with Structured Data. Google for Developers documentation. Retrieved from https://developers.google.com/search/docs/appearance/structured-data/faqpage?hl=en
The JAX Authors. Custom derivative rules for JAX-transformable Python functions. JAX documentation. Retrieved from https://docs.jax.dev/en/latest/notebooks/Custom_derivative_rules_for_Python_code.html
