Automatic Differentiation: How Gradients Are Actually Computed

Symbolic differentiation is what a computer algebra system does, and it is genuinely useful right up to the point where the function is 400 lines of Python with a for loop in it. Numerical differencing, from the first lesson, survives that but cannot pay the per-parameter cost.

Automatic differentiation is the third option, and its trick is to refuse to work with formulas at all.

Every program, however long, is ultimately a sequence of primitive operations: add, multiply, exponentiate, index, matrix-multiply. There are perhaps a hundred of them in a framework, and the derivative of each one is a line of code somebody wrote once.

Automatic differentiation records the sequence of primitives a particular execution performed, then applies the chain rule along that recorded sequence. It never needs an expression for $f$, only the trace of what $f$ did.

Key idea: Automatic differentiation is exact, like symbolic differentiation, and mechanical, like numerical differentiation. It gets both by changing the object it operates on: a straight-line list of primitive operations rather than a mathematical expression. Loops, branches and recursion disappear, because by the time you have run the program they have already resolved into a flat list.

Take $L = x_1 x_2 + \sin(x_1 x_2)$ evaluated at $x_1 = 2$, $x_2 = 3$. Running it produces this trace:

Three primitives, each with a derivative rule known in advance: a product rule, the derivative of sine, and a sum rule. The chain rule now has somewhere to run. The only remaining question is which direction to run it in, and that choice is the whole difference between the two modes of automatic differentiation.

Attach a second number to every value, holding its derivative with respect to one chosen input. These pairs are dual numbers, obeying $\varepsilon^2 = 0$:

$(a + b\varepsilon)(c + d\varepsilon) = ac + (ad + bc)\varepsilon$

The product rule appears in the $\varepsilon$ coefficient by itself. Forty lines of Python implement it:

import math

class Dual:
    def __init__(self, v, d=0.0):
        self.v, self.d = v, d          # value, derivative
    def __add__(self, o):
        return Dual(self.v + o.v, self.d + o.d)
    def __mul__(self, o):
        return Dual(self.v * o.v, self.v * o.d + self.d * o.v)

def sin(u):
    return Dual(math.sin(u.v), math.cos(u.v) * u.d)

x1 = Dual(2.0, 1.0)                    # seed: differentiate wrt x1
x2 = Dual(3.0, 0.0)
v1 = x1 * x2
L = v1 + sin(v1)
print(L.v, L.d)                        # 5.720584... 5.880510...

One sweep, forward, no tape. The catch is the seed: it differentiates with respect to $x_1$ only, so an $n$-input gradient takes $n$ sweeps.

Reverse mode evaluates the function first, recording each intermediate value on a tape. Then it walks the tape backwards, carrying an adjoint: how much the final output changes per unit change in that intermediate. Seed the output with 1 and every adjoint falls out by the chain rule. For our trace, $\bar{v}_2 = 1$, then $\bar{v}_1 = 1 + \cos(6) = 1.9602$, and finally $\bar{x}_1 = 1.9602 \times 3 = 5.8805$ and $\bar{x}_2 = 1.9602 \times 2 = 3.9203$. Both partial derivatives, from one backward sweep.

Forward mode costs one sweep per input. Reverse mode costs one sweep per output. Training has an extreme shape, many inputs and exactly one output, so reverse mode wins by whatever the parameter count happens to be.

JAX's own documentation puts it plainly: for a function $f: \mathbb{R}^n \to \mathbb{R}$, "we can do it in just one call. That's how grad is efficient for gradient-based optimization, even for objectives like neural network training loss functions on millions or billions of parameters."

The constant is small and, remarkably, it does not depend on the number of parameters. Baur and Strassen proved in 1983 that computing a full gradient costs at most about four times the cost of evaluating the function itself. JAX quotes a similar figure for its own transformations, "only about three times the cost of just evaluating the function".

The second bar never moves. That flat line is the reason a trillion-parameter model can be trained at all, and it is a theorem, not an engineering achievement.

The backward sweep needs the intermediate values the forward sweep produced, so they have to be kept. Memory therefore grows with the depth of the computation, which is exactly the cost forward mode does not pay. This is why activations, not weights, usually fill a training GPU.

The standard trade is gradient checkpointing: keep only some intermediates and recompute the rest on the way back. Chen and colleagues showed in 2016 that storing every $\sqrt{n}$-th layer gives $O(\sqrt{n})$ memory "with only the computational cost of an extra forward pass per mini-batch", and reported reducing a 1,000-layer residual network from 48 GB to 7 GB with a 30 percent runtime increase.

In practice: Checkpointing is one flag in every major framework. Reach for it when you are memory-bound rather than compute-bound, which for large models is most of the time.

Automatic differentiation returns the exact derivative of the program that ran. That is not always the derivative you wanted.

• Control flow is resolved before differentiation, so you get the derivative of the branch taken. A max(a, b) routes the gradient entirely to the winner and reports nothing about the switch itself.
• Non-differentiable points are still non-differentiable. Frameworks return a convention, ReLU at zero being the standard case, and that convention becomes part of your numerics.
• In-place mutation overwrites values the backward sweep still needed. PyTorch raises "a variable needed for gradient computation has been modified by an inplace operation" rather than returning something wrong, which is the kinder failure.
• Iterative solvers differentiate through every iteration you actually ran, not through the fixed point you were converging to. Implicit differentiation exists for precisely this case.

Gotcha: "The gradients are wrong" is nearly always a mismatch between the mathematical function you had in mind and the program you actually wrote.

They compose, which is where the second-order tools come from. A Hessian-vector product $Hv$ is the derivative of the gradient in the direction $v$: build the gradient with reverse mode, then push one forward-mode sweep through it. Forward-over-reverse gives you $Hv$ exactly, at roughly the cost of a gradient, without ever forming the $n \times n$ matrix that made curvature look unaffordable in the previous lesson.