Automatic Differentiation

Training a neural network comes down to one question repeated millions of times: how should we nudge each weight to reduce the loss? The answer requires a derivative — the gradient — of a function that can have billions of inputs and layers of nonlinearities stacked on top of each other.

Symbolic differentiation (the kind taught in calculus class) produces exact formulas, but applied to code it generates expressions that explode in size. Finite differences approximate derivatives by evaluating f(x+ε) and f(x), but they are slow (one extra pass per parameter) and numerically fragile. Neither scales.

Automatic differentiation (autodiff) threads a third path. It does not manipulate formulas. It does not perturb inputs. Instead it applies the chain rule mechanically to each arithmetic operation the program performs, accumulating exact derivatives as a by-product of a normal execution. The result is exact (up to floating-point rounding) and costs only a small constant multiple of the original computation — regardless of how many parameters there are.

This idea, developed in its modern form through the 1960s–1980s, is the algorithm that makes deep learning tractable. Every major framework — PyTorch, TensorFlow, JAX — runs autodiff under the hood.

The demo below builds the expression (x · w + b)² — a tiny neural-network neuron squared — and computes its gradient with respect to every variable using reverse-mode autodiff.

<div class="controls">
  <label>x = <input id="xval" type="number" value="2" step="0.5"></label>
  <label>w = <input id="wval" type="number" value="3" step="0.5"></label>
  <label>b = <input id="bval" type="number" value="-1" step="0.5"></label>
</div>
<div class="btns">
  <button id="fwd" type="button">{{btn_fwd}}</button>
  <button id="bwd" type="button" disabled>{{btn_bwd}}</button>
  <button id="reset" type="button" class="ghost">{{btn_reset}}</button>
</div>
<div id="graph" class="graph"></div>
<p class="hint">{{hint}}</p>

* { box-sizing: border-box; }
body { font-family: system-ui, sans-serif; margin: 0; color: #222; }
.controls { display: flex; gap: 1rem; flex-wrap: wrap; margin-bottom: .6rem; }
.controls label { font-size: .9rem; }
.controls input { width: 64px; padding: .2rem .4rem; border: 1px solid #aaa; border-radius: 6px;
                  font-size: .9rem; margin-left: .3rem; }
.btns { display: flex; gap: .5rem; flex-wrap: wrap; margin-bottom: .8rem; }
button { font: 600 14px system-ui; padding: .4rem .85rem; border: 1px solid #1d3557;
         background: #1d3557; color: #fff; border-radius: 8px; cursor: pointer; }
button:disabled { opacity: .4; cursor: default; }
button.ghost { background: #fff; color: #1d3557; }
.hint { font-size: .85rem; color: #555; margin-top: .6rem; }
.graph { display: flex; flex-direction: column; gap: 6px; }
.row { display: flex; gap: 6px; flex-wrap: wrap; }
.node { background: #e8eef3; border: 1.5px solid #cdd9e3; border-radius: 10px;
        padding: .35rem .6rem; min-width: 90px; text-align: center; font-size: .82rem; line-height: 1.4; }
.node .op  { font-weight: 700; font-size: .78rem; color: #1d3557; }
.node .val { font-family: ui-monospace, monospace; font-size: .88rem; }
.node .grad { font-size: .78rem; color: #0a7d33; font-family: ui-monospace, monospace; }
.node.input { background: #dceefb; border-color: #8ab4cc; }
.node.output { background: #d4edda; border-color: #6abf80; }
.node.active { border-color: #e63946; }
.arrow { font-size: 1.1rem; color: #888; align-self: center; }

// Code not found

Click Run forward pass to evaluate the expression and build the computation graph. Then click Run backward pass to backpropagate gradients. Each node shows its value and its gradient (∂L/∂node). Notice that the gradient of each input is exact — no approximations, no extra function calls per parameter. Change the values of x, w, and b to see how the gradients shift.

Autodiff comes in two flavours that trade cost differently.

• Forward mode (tangent-linear): computes one directional derivative per pass. Cost: one extra pass per input. Efficient when there are few inputs and many outputs — e.g., sensitivity analysis.
• Reverse mode (adjoint / backpropagation): computes all partial derivatives in one backward sweep over a stored tape (the Wengert list of operations). Cost: one forward pass to record, one backward pass to accumulate. Efficient when there are many inputs and few outputs — exactly the shape of a loss function over millions of weights.
• The chain rule is the key. Each elementary operation (add, multiply, exp, …) contributes a known local derivative. Reverse mode multiplies these local Jacobians from output back to inputs, assembling the full gradient without ever forming the full Jacobian matrix.
• Overhead: reverse-mode autodiff costs roughly 3–5× the forward pass in time and proportional memory for the tape. This constant factor is independent of the number of parameters — a crucial property that symbolic differentiation cannot match.
• Limits: autodiff is decidable and polynomial for programs with fixed control flow. Differentiating through dynamic loops or recursion whose depth depends on data can require storing an unbounded tape; and differentiating through non-differentiable branches (e.g., if x > 0) yields subgradients, not classical derivatives.

Unlike problems in NP, autodiff itself is not a search problem — it is a linear-time traversal of the computation graph. Its status is solved-and-efficient, not hard. The challenge is engineering: building a system that traces arbitrary code and applies the chain rule correctly at every node. That engineering is what PyTorch's autograd and JAX's grad transform do, and why they matter so much to neural network training.

Automatic differentiation is the quiet engine under a remarkable range of modern computation:

• Deep learning: backpropagation is reverse-mode autodiff applied to a neural loss function. Without it, training networks with millions of weights would be computationally impossible. Every epoch of GPT, every forward-pass of a diffusion model, runs autodiff in reverse.
• Scientific computing: JAX and similar frameworks bring autodiff to numerical simulations. Differentiating a physics solver gives you gradients for inverse problems — fitting a climate model to data, optimizing a wing shape, reconstructing protein structures.
• Optimal control: gradient-based trajectory optimization (e.g., MuJoCo, Brax) differentiates through entire simulated rollouts to find control policies that minimize cost.
• Probabilistic programming: variational inference (ELBO maximization) and normalizing flows require gradients of stochastic computation graphs; autodiff handles this via the reparameterization trick.
• Implicit differentiation: by differentiating through fixed-point conditions rather than the iterative solver itself, implicit autodiff can compute gradients through optimization loops in $O(1)$ backward passes — the trick behind hyperparameter optimization at scale.

The pattern is always the same: wherever gradient-based optimization is needed for a function defined by a program, autodiff makes it exact and cheap.

Automatic differentiation is one of those rare ideas that is both simple and transformative. The chain rule is undergraduate calculus; applying it mechanically to every operation a program executes is a straightforward engineering task. Yet together they produce exact gradients through arbitrarily complex functions at constant overhead — a fact that unlocks gradient-based learning at any scale.

Every time a language model improves, every time a physics simulator finds an optimal trajectory, every time a variational model fits data, autodiff is silently threading the chain rule backward through a computation graph. It is not a hard problem in the complexity-theoretic sense — it is a solved one. That rarity is exactly what makes it so powerful.

For the problems that are hard — like the P vs NP question — we still hope that someday a clever algorithm will thread the needle. Autodiff is a reminder that sometimes the right idea really does make a hard-looking problem easy.