Automatic Differentiation

Automatic differentiation is a family of techniques for evaluating the derivative of a numerically-specified function by decomposing it into elementary operations — addition, multiplication, exponentiation, trigonometric functions, and so on — each with a known derivative, and combining these via the chain rule. It stands apart from two other approaches. Numerical differentiation approximates derivatives via finite differences, e.g. $\frac{f(x+h)-f(x)}{h}$, which is cheap to implement but suffers from truncation error (bias from finite $h$) and cancellation error (loss of precision as $h \to 0$), and requires a separate function evaluation per input dimension. Symbolic differentiation manipulates expression trees to produce closed-form derivative expressions, as a computer algebra system would; it is exact but the resulting expressions can grow exponentially in size (expression swell) when applied to programs with shared subexpressions, loops, or branching. Autodiff sidesteps both problems: it evaluates derivatives numerically at a specific point, using the exact symbolic rule for each elementary step, and reuses shared intermediate values instead of re-expanding them. The formal object underlying AD is the computational graph (or Wengert list, or tape): a directed acyclic graph whose nodes are variables (inputs, intermediate values, outputs) and whose edges represent elementary operations. Every differentiable program, however it is written — with loops, conditionals, recursion — can be unrolled at a specific input into such a graph, because AD only needs to differentiate through the operations actually executed, not through all possible executions. This is the crucial insight: AD operates on the execution trace, not on the static source code, which is why it handles arbitrary control flow correctly as long as that control flow doesn't itself depend discontinuously on the differentiated variables in a way that breaks differentiability. Two dual computational strategies exploit this graph. Forward-mode AD computes, alongside every intermediate value $v_i$, its directional derivative $\dot v_i = \frac{\partial v_i}{\partial x_j}$ with respect to one designated input $x_j$, propagating from inputs to outputs in the same order as the original computation. Reverse-mode AD computes, for one designated output $y_k$, the adjoint $\bar v_i = \frac{\partial y_k}{\partial v_i}$ for every intermediate variable, propagating from outputs back to inputs — the opposite order from the original computation, which is precisely what "backpropagation" refers to. Both modes compute the exact same Jacobian matrix $J = \frac{\partial y}{\partial x}$; they differ only in whether they compute it column-by-column (forward) or row-by-row (reverse).

Forward mode. Extend every real number $v$ in the computation to a dual number pair $(v, \dot v)$, where $\dot v$ tracks the derivative of $v$ with respect to a chosen input direction. Elementary operations are overloaded to propagate both components simultaneously using the sum, product, and chain rules:

Seed the input whose derivative you want with $\dot x_j = 1$ (and every other input with $\dot x = 0$), run the program once, and every intermediate and output variable now carries its derivative with respect to $x_j$ for free, in the same pass as the value itself. To get the full Jacobian for $n$ inputs, repeat this $n$ times, once per seed direction (or propagate all $n$ seed directions simultaneously as vectors, at $n\times$ the per-variable memory).

Reverse mode. Run the ordinary computation forward once, recording every elementary operation and its inputs on a tape (the computational graph). Then traverse the tape backward, maintaining for each variable an adjoint $\bar v = \partial y / \partial v$ initialized to $\bar y = 1$ at the single output of interest. At each node, propagate the adjoint to its parents (the graph inputs to that operation) using the local Jacobian-vector product (a "vector-Jacobian product", VJP) of that elementary operation, and accumulate contributions when a variable feeds into multiple downstream nodes (multivariate chain rule = sum over all paths). After one backward traversal, every input variable holds $\partial y/\partial x_i$ — the entire gradient, for the cost of roughly one extra forward-equivalent pass, regardless of the number of inputs. ```

# ---------- Shared representation ----------

# A computational graph node stores:

# value : the primal numeric value

# op : the elementary operation that produced it (or None for a leaf/input)

# parents : the input nodes to that operation

# local_fn : function(parents_values) -> value (forward rule)

# vjp_fn : function(grad_output) -> tuple(grad_wrt_each_parent) (reverse rule)

# jvp_fn : function(parents_tangents) -> tangent (forward rule)

# ---------- FORWARD MODE (dual numbers) ---------- function FORWARD_MODE_AD(f, x, seed_index):

# x: list of n input values; seed_index: which input's derivative we want n = length(x) v = copy(x) # primal values vdot = zeros(n) vdot[seed_index] = 1.0 # seed direction: d/dx_seed for each elementary step (op, inputs) in TRACE(f, v): # trace built by executing f primal_out = op.local_fn(values(inputs)) tangent_out = op.jvp_fn(tangents(inputs)) # chain rule, one direction record(primal_out, tangent_out) return primal_out, tangent_out # = f(x), df/dx_seed evaluated at x

# To get the full gradient of a SCALAR loss w.r.t. n inputs: call this n times -> O(n) passes.

# ---------- REVERSE MODE (tape + adjoints) ---------- function REVERSE_MODE_AD(f, x): tape = [] # ordered list of (node, op, parents)

# ---- Forward pass: build the graph, remember every op ---- for each elementary step (op, parents) executed while computing f(x): out_node = Node(value = op.local_fn(values(parents)), op = op, parents = parents) tape.append(out_node) y = tape[-1].value # scalar output (the loss)

# ---- Backward pass: propagate adjoints ---- adjoint = { node: 0 for node in tape } adjoint[tape[-1]] = 1.0 # dy/dy = 1 for node in reversed(tape): # topological order guarantees parents processed after children if adjoint[node] == 0: continue grads_wrt_parents = node.op.vjp_fn(adjoint[node], node.parents) for (parent, g) in zip(node.parents, grads_wrt_parents): adjoint[parent] += g # accumulate: multiple paths sum (multivariate chain rule) return y, { x_i: adjoint[x_i] for x_i in inputs(x) } # one pass gives ALL n partial derivatives ``` The key structural difference is visible in the pseudocode: forward mode's outer loop is over input dimensions (repeat the whole pass per seed), while reverse mode's outer loop is over nothing — one forward pass builds the tape, one backward pass yields every partial derivative simultaneously, because the adjoint accumulation rule naturally fans out to all parents at once.

Forward mode complexity. Let the primal computation cost $T$ (time) and $S$ (space, dominated by intermediate values), for a function $f:\mathbb R^n \to \mathbb R^m$. Propagating one tangent direction costs a constant factor more than the primal (each elementary op does one extra local multiply/add for the tangent), so one directional derivative costs $O(T)$ time and $O(S)$ space — no tape is needed, since tangents are consumed as the pass proceeds. To get the full Jacobian ($n$ columns), repeat $n$ times: total $O(nT)$ time. Forward mode is therefore efficient when $n$ is small (few inputs, e.g. sensitivity of an output to one or two physical parameters) or when $m \gg n$ (many outputs, few inputs — Jacobian is "wide"), since you'd need $O(nT)$ regardless of $m$.

Reverse mode complexity. The forward pass costs $O(T)$ and must be stored in full — every intermediate value on the tape is kept because the backward pass needs it to evaluate local derivative rules — so space is $O(S)$ where $S$ scales with the number of operations executed (not just the primal's own working memory), a real and often binding cost for very deep networks. The backward pass processes each tape entry exactly once, and each entry's `vjp_fn` costs a constant factor more than its forward counterpart (Baur–Strassen-type argument: differentiating one elementary op costs at most a small constant times evaluating it), so the backward pass is also $O(T)$. One backward pass yields the entire row of the Jacobian for a chosen output — i.e., the gradient of a scalar loss with respect to all $n$ inputs simultaneously — at total cost $O(T)$, independent of $n$. This is why training a network with $10^8$ parameters is feasible: reverse mode gives all $10^8$ partial derivatives for about $2$–$3\times$ the cost of one forward evaluation, whereas forward mode or finite differences would need $10^8$ separate passes.

Why reverse mode wins for deep learning. A training step computes one scalar loss $L$ (so $m=1$) as a function of millions of parameters (so $n \gg 1$). Forward mode costs $O(nT)$ — intractable. Reverse mode costs $O(T)$ — one backward pass. This asymmetry, $m=1 \ll n$, is exactly the regime reverse mode is built for, and it is precisely the mathematical content of Backpropagation: applying reverse-mode AD to the graph of a layered network's loss.

Correctness argument. The correctness of reverse mode follows from the multivariate chain rule applied inductively over a topological order of the DAG. Define the adjoint of node $v$ as $\bar v = \partial y/\partial v$ for the single scalar output $y$. If $v$ feeds into successors $w_1,\dots,w_k$, the chain rule gives

Processing nodes in reverse topological order guarantees that when we reach $v$, every successor $w_i$ has already had its adjoint $\bar w_i$ fully accumulated (all of its successors were processed earlier, since reverse order visits sinks before sources). Thus the accumulation loop in the pseudocode — `adjoint[parent] += g` executed for every node in reverse order — computes exactly this sum, term by term, and by induction on the topological order, every node's adjoint equals its true partial derivative once all of its successors have been visited. The base case is the output itself, $\bar y = 1$. This is a finite, exact sequence of arithmetic operations — no truncation, no approximation — so reverse-mode AD is exact up to the floating-point precision of the arithmetic used, unlike numerical differentiation which has an unavoidable $O(h)$ or $O(h^2)$ truncation bias.

Honest tradeoffs. Reverse mode trades time for memory: the tape must retain every intermediate activation for the backward pass, which is the dominant memory cost of training large networks (addressed in practice by graph-level tricks like gradient checkpointing/activation recomputation, trading recomputation time for reduced memory). Forward mode needs no tape and has flat memory, but its cost scales linearly in input dimension, making it a poor fit for gradient descent over huge parameter vectors, though it remains the tool of choice for Jacobian-vector products, Hessian-vector products (forward-over-reverse), and sensitivity analysis with few parameters.

Let $f(x,y) = \sin(xy) + (xy)^2$, evaluated at $x=2, y=3$. Let $a = xy$, $b=\sin a$, $c = a^2$, $z = b + c$. We want $\partial z/\partial x$ and $\partial z/\partial y$. Analytically: $z = \sin(xy) + (xy)^2$, so $\partial z/\partial x = y\cos(xy) + 2xy^2$ and $\partial z/\partial y = x\cos(xy) + 2x^2y$. At $(x,y)=(2,3)$: $xy = 6$, $\cos 6 \approx 0.9602$, so $\partial z/\partial x = 3(0.9602) + 2(2)(9) = 2.8806 + 36 = 38.8806$, and $\partial z/\partial y = 2(0.9602) + 2(4)(3) = 1.9204 + 24 = 25.9204$.

Reverse mode trace (one forward pass, one backward pass gives both derivatives):

Both partial derivatives match the analytic values exactly, and step 9 shows the accumulation rule in action: $a$ feeds two successors ($b$ and $c$), so its adjoint is the sum of contributions from both — the multivariate chain rule enforced automatically by the `+=` accumulation, exactly as the correctness argument requires.

Forward mode, for comparison, needs two separate passes (one per input) to get the same two numbers. Seeding $\dot x = 1,\dot y=0$: $\dot a = \dot x y + x\dot y = 3$; $\dot b = \cos(a)\dot a = 0.9602\cdot 3=2.8806$; $\dot c = 2a\dot a = 12\cdot 3=36$; $\dot z = \dot b+\dot c = 38.8806$ — matches $\partial z/\partial x$. Re-run seeding $\dot x=0,\dot y=1$ to get $\dot z = 25.9204 = \partial z/\partial y$. Two full passes for two derivatives versus one backward pass for both — the gap only widens as the number of inputs grows, which is the entire complexity argument made concrete.

Automatic differentiation's reverse mode was independently discovered several times before becoming central to deep learning: in optimal control theory, in numerical computing, and finally as the engine of neural network training under the name backpropagation. Seppo Linnainmaa gave the first complete reverse-mode AD algorithm for arbitrary differentiable programs in his 1970 thesis, describing it as an error-propagation method for numerical rounding analysis, without yet connecting it to machine learning. The idea was reinvented and popularized for training multilayer neural networks over the following decade, culminating in the 1986 paper by Rumelhart, Hinton, and Williams that demonstrated backpropagation could train hidden layers to learn useful internal representations, catalyzing the connectionist revival.

Every major deep learning framework — PyTorch's `autograd`, TensorFlow's `GradientTape`, JAX's `grad` — is a reverse-mode AD engine operating over a computational graph, whether built dynamically at runtime (define-by-run) or compiled ahead of time (define-and-run). Training any neural network, from a small multilayer perceptron to a trillion-parameter transformer, relies on reverse-mode AD to compute the gradient of the loss with respect to every weight in a single backward pass, which is the concrete mechanism behind Backpropagation. Beyond neural networks, forward-mode AD is standard in engineering simulation and sensitivity analysis (computing how a small number of design parameters affect many outputs, e.g. in computational fluid dynamics or circuit design), and mixed forward-over-reverse AD computes Hessian-vector products used in second-order optimization and in curvature-aware training diagnostics. Physics-informed machine learning, differentiable rendering, and probabilistic programming languages all depend on AD to differentiate through complex, control-flow-heavy programs exactly. Once a framework can compute exact gradients cheaply, the next practical question becomes how to initialize the parameters those gradients will update, which leads directly into Weight Initialization Strategies.

Ask the tutor: How does gradient checkpointing trade recomputation time for reduced memory in reverse-mode AD, and how do you choose which activations to checkpoint?