Automatic Differentiation

Why Does This Exist?

Nobody derives the gradient of a modern network by hand, and nobody needs to. You write a forward pass in Python and .backward() produces exact gradients for every parameter. That is not magic and it is not symbolic algebra. It is a specific technique with a specific cost model, and knowing the model tells you why some things are fast and others surprisingly are not.

Three concrete payoffs:

• Why training costs about two forward passes, regardless of parameter count. Reverse mode gets every gradient in one sweep. If it scaled per parameter, training a 7-billion-parameter model would be impossible rather than merely expensive.
• Why jacrev and jacfwd both exist, and how to pick. Choosing wrong can cost a thousandfold.
• Why a custom backward is sometimes necessary, and what you promise when you write one.

Autodiff is also the reason research moves quickly: you change an architecture and get correct gradients without re-deriving anything. The chain rule it mechanises and the graphs it records live in computational graphs, while the object it produces is in gradients.

The one idea you need first

Derivatives compose by the chain rule: if $y$ depends on $u$ and $u$ depends on $x$, then $dy/dx = (dy/du)(du/dx)$, where $dy/du$ is the outer slope and $du/dx$ is the inner slope. A gradient collects one such rate per parameter. The full rule with worked chains is in chain rule.

Think of It Like This

How It Actually Works

Autodiff is neither of the two things people assume.

It is not symbolic differentiation. Symbolic tools manipulate expressions and produce formulas, which suffer expression swell: the symbolic derivative of a deep composition is enormous and slow to evaluate.

It is not numerical differentiation. Finite differences approximate, need a step size, cost one evaluation per input, and lose precision to cancellation described in machine epsilon.

Autodiff is exact and cheap because it applies the chain rule to the sequence of operations that ran. Every primitive such as add, multiply, exp, and matmul has a known local derivative, and the framework composes them numerically as it goes. No formula is produced and no approximation is made: the result is exact to floating-point precision.

Two directions, one rule

Forward mode propagates derivatives alongside values, from inputs to outputs. Each variable carries a pair of its value and its derivative with respect to one chosen input. One sweep gives $\partial(\text{everything}) / \partial x_j$ for a single $j$, where $x_j$ is the seeded input.

The clean implementation is dual numbers: represent each quantity as $a + b\epsilon$ where $\epsilon^2 = 0$, with $a$ the value and $b$ the derivative. Then $(a + b\epsilon)(c + d\epsilon) = ac + (ad + bc)\epsilon$, and $ad + bc$ is exactly the product rule.

Reverse mode runs the function forward while recording a computational graph, where the graph is the tape of intermediates, then walks it backwards accumulating adjoints. One sweep gives $\partial(\text{one output}) / \partial(\text{everything})$.

The cost model that decides everything

For a function with $n$ inputs and $m$ outputs, in units of one function evaluation, where $n$ counts inputs and $m$ counts outputs:

• Forward mode costs about $n$ sweeps for the full Jacobian. Cheap when $n$ is small.
• Reverse mode costs about $m$ sweeps. Cheap when $m$ is small.

Training has $n$ in the billions and $m = 1$, because the loss is a scalar. Reverse mode gets the entire gradient in one sweep, while forward mode would need one per parameter. That asymmetry is the single most important fact about autodiff, and it is why backpropagation is just reverse mode applied to a network.

The trade is memory. Reverse mode must retain intermediate values for the backward pass, so activation memory scales with depth and batch size. Forward mode needs no tape at all, which is why it is used for Jacobian-vector products and for verifying gradients on small inputs.

Where the abstraction leaks

Autodiff is exact for the operations it sees, and three situations break that:

• Non-differentiable points such as ReLU at exactly zero, abs, max, and rounding. Frameworks pick a subgradient by convention and carry on silently.
• Control flow on data. A dynamic framework differentiates the branch that ran, which is correct for that input and means the gradient is only valid locally.
• Operations outside the framework such as a NumPy call, a C extension, a simulator, or a discrete sampling step. The chain silently stops and you get zero gradient rather than an error.

That last one is the expensive failure, and it is what custom backward implementations and reparameterisation tricks exist to repair.

Worked example

Take $f(x_1, x_2) = x_1 x_2 + \sin(x_1)$ at $(2, 3)$, where $x_1$ and $x_2$ are inputs and $f$ is the scalar output.

Forward pass, recording each intermediate $v$, where $v$ names a stored value:

• $v_1 = x_1 = 2$
• $v_2 = x_2 = 3$
• $v_3 = v_1 v_2 = 6$
• $v_4 = \sin(v_1) = \sin(2) \approx 0.9093$
• $v_5 = v_3 + v_4 \approx 6.9093$

Reverse pass, starting from $\partial f / \partial v_5 = 1$ and walking back, where $\bar v$ means $\partial f / \partial v$:

• Addition sends its incoming derivative to both inputs unchanged: $\bar v_3 = 1$ and $\bar v_4 = 1$.
• $v_4 = \sin(v_1)$ has local derivative $\cos(v_1)$, contributing $\bar v_4 \cos(2) = -0.4161$ to $\bar v_1$, where $\cos(2) \approx -0.4161$.
• $v_3 = v_1 v_2$ has local derivatives $v_2$ and $v_1$, contributing $\bar v_3 \cdot 3 = 3$ to $\bar v_1$ and $\bar v_3 \cdot 2 = 2$ to $\bar v_2$.

Note $v_1$ receives contributions from two paths, and they add:

Verify analytically: $\partial f/\partial x_1 = x_2 + \cos(x_1) = 3 + \cos 2$ and $\partial f/\partial x_2 = x_1 = 2$.

The cost comparison is the point. Reverse mode produced both partials in one backward sweep. Forward mode would need two sweeps, one seeded on $x_1$ and one on $x_2$. With two inputs that is a factor of $2$, and with 7 billion it is the difference between possible and not.

Code

1import math2
3
4class Dual:5    """Forward mode: value plus derivative, propagated by the arithmetic itself."""6
7    def __init__(self, val: float, der: float = 0.0):8        self.val, self.der = val, der9
10    def __mul__(self, o: "Dual") -> "Dual":11        return Dual(self.val * o.val, self.val * o.der + self.der * o.val)12
13    def __add__(self, o: "Dual") -> "Dual":14        return Dual(self.val + o.val, self.der + o.der)15
16
17def sin_dual(d: Dual) -> Dual:18    return Dual(math.sin(d.val), math.cos(d.val) * d.der)19
20
21# Seed the derivative on x1 only: one forward sweep gives df/dx1.22x1, x2 = Dual(2.0, 1.0), Dual(3.0, 0.0)23out = x1 * x2 + sin_dual(x1)24print(round(out.val, 4), round(out.der, 4))      # -> 6.9093 2.583925
26# Seeding x2 instead needs a SECOND sweep -- that is the cost of forward mode.27x1, x2 = Dual(2.0, 0.0), Dual(3.0, 1.0)28print(round((x1 * x2 + sin_dual(x1)).der, 4))    # -> 2.029
30print(round(3 + math.cos(2), 4), 2.0)            # -> 2.5839 2.0  analytic check

$2.5839$ and $2.0$ match the hand calculation, and the analytic check confirms them. The two separate seedings are the forward-mode cost made explicit: one sweep per input, where reverse mode would have taken one sweep total.

Watch Out For

The Quick Version

• Autodiff is neither symbolic nor numerical. It applies the chain rule to the operations that actually ran, exactly.
• Every primitive has a known local derivative, and the framework composes them numerically.
• Forward mode carries derivatives with values, costs one sweep per input, and needs constant memory. Dual numbers are the clean implementation.
• Reverse mode records a graph forward, walks it back, costs one sweep per output, and needs the tape, so memory grows with depth.
• On $x_1 x_2 + \sin(x_1)$ at $(2, 3)$, forward values $6$ and $0.9093$ sum to $6.9093$, and one reverse sweep gives $2.5839$ and $2.0$.
• Cost is $n$ sweeps forward and $m$ sweeps reverse. Training has $n$ in the billions and $m = 1$.
• Training therefore costs about two forward passes regardless of parameter count.
• Backpropagation is reverse-mode autodiff applied to a neural network.
• The chain breaks silently at NumPy conversions, .detach(), external code, and discrete sampling. Assert nonzero gradients after the first step.
• Non-differentiable points get a convention, not an error. Verify custom backwards with gradcheck in float64.

What to Read Next

• The Chain Rule
• Computational Graphs
• Gradients
• Backpropagation