Automatic Differentiation

Dive into Deep Learning · §1.5

From the chain rule to backpropagation
automatic differentiation for a complete network.

Record the forward pass; replay it in reverse

Motivation

Deriving a million-parameter network’s gradient by hand is impractical. Instead, the framework records each operation during the forward pass, then replays it in reverse, systematically applying the chain rule from the calculus section to obtain the gradient with respect to every input at once.

Training steps in this book generally include a forward pass and a backward pass over this graph.

The autograd loop1markflag an input2forwardrecords the graph3backwardchain rule, reversed4read .graduse it, then resetrepeated every optimization step

01

The mechanics

record forward · sweep backward

A function with a known answer: ∇y = 4x

Mechanics

Differentiate y = 2\,\mathbf{x}^\top\mathbf{x} w.r.t. the vector \mathbf{x}. The analytic answer, \nabla_\mathbf{x} y = 4\mathbf{x}, provides a check for the automatic-differentiation result.

x = torch.arange(4.0)
x
tensor([0., 1., 2., 3.])
Computational graph for y = 2 xᵀxdot× 2∂y/∂a = 2∂a/∂x = 2xxinputa = xᵀxscalary = 2aoutputforward →← backwardchain rule: ∂y/∂x = (∂y/∂a)(∂a/∂x) = 2 · 2x = 4x

Track x, and the graph builds itself

Mechanics

First tell the framework to track x (reserve a slot for its gradient), then run the forward pass; y is now the root of a recorded graph:

# Can also create x = torch.arange(4.0, requires_grad=True)
x.requires_grad_(True)
x.grad  # The gradient is None by default
y = 2 * torch.dot(x, x)
y
tensor(28., grad_fn=<MulBackward0>)

One backward call returns the whole gradient

Mechanics

One call sweeps the graph in reverse, and the result equals the promised 4\mathbf{x}, at every coordinate:

y.backward()
x.grad
tensor([ 0.,  4.,  8., 12.])
x.grad == 4 * x
tensor([True, True, True, True])

That reverse sweep is the calculus section’s chain rule, run from output to input.

02

Working with gradients

accumulation · non-scalar outputs · detaching · inference

Gradients accumulate: reset first

Gradients

PyTorch adds each new gradient into x.grad rather than replacing it (handy for summing losses). So zero it before a fresh computation:

x.grad.zero_()  # Reset the gradient
y = x.sum()
y.backward()
x.grad
tensor([1., 1., 1., 1.])

Forgetting .zero_() between iterations causes unintended gradient accumulation.

Vector outputs: differentiate their sum

Gradients

Gradients are defined for a scalar loss. For a vector y, the engine differentiates the sum of its components (a vector–Jacobian product), exactly what a per-example batch loss needs:

x.grad.zero_()
y = x * x
y.backward(gradient=torch.ones(len(y)))  # Equivalently: y.sum().backward()
x.grad
tensor([0., 2., 4., 6.])

detach freezes a value: ∂z/∂x = u, not 3x²

Gradients

Sometimes a value should count as a constant: gradients must not flow through it. detach (or stop_gradient) severs the graph above it, so z = u \cdot x differentiates to u, not to 3x^2:

x.grad.zero_()
y = x * x
u = y.detach()
z = u * x

z.sum().backward()
x.grad == u
tensor([True, True, True, True])
Detaching freezes a value to a constantsquaredetach× xxinputx²uconstantz = u·x∂z/∂x = u (gradient returns only through the direct edge)the path above u is cut, so z = x·x·x but ∂z/∂x = u, not 3x²

Inference skips the bookkeeping

Gradients

When we only need the value (prediction, evaluation, manual updates), we turn recording off and avoid its bookkeeping cost. This is the default mode for inference throughout the book:

with torch.no_grad():
    y = 2 * torch.dot(x, x)
y.requires_grad  # False: y is detached from the graph
False

03

Dynamic graphs

the graph is whatever actually ran

The graph records what actually ran

Dynamic graphs

In eager execution, autograd records the operations executed by if and while statements. This function’s loop count and branch both depend on its input (compiled transformations may impose additional constraints):

def f(a):
    b = a * 2
    while b.norm() < 1000:
        b = b * 2
    if b.sum() > 0:
        c = b
    else:
        c = 100 * b
    return c
The graph is built at runtimea = 0.30a×2×2×2b.sum()>0 → bca = −0.75a×2×2×2×2×2else → 100·bcdifferent input → different path → autograd records whatever ran

Branch or loop, the derivative is f(a)/a

Dynamic graphs

Each call realizes a concrete graph that backward can walk. Whichever branch ran, f scaled its input by some constant, so f(a) = k\,a and the gradient must equal f(a)/a. It does:

a = torch.randn(size=(), requires_grad=True)
d = f(a)
d.backward()
a.grad == d / a
tensor(True)

Reverse mode: the whole gradient for one extra pass

Beyond

A counting argument settles which way to sweep. With n inputs and m outputs, the full derivative matrix costs m reverse sweeps or n forward sweeps, each sweep priced at roughly one function evaluation.

A training loss has m = 1 and n in the millions: one reverse sweep delivers every parameter’s gradient, for the cost of about one extra forward pass. Forward mode wins the opposite regime (few inputs, many outputs) and Hessian–vector products.

Two ways to sweep the same graphforward mode: one input → all outputsxhyxhyreverse mode (backprop): one output → all inputsa scalar loss over millions of parameters → reverse wins

Differentiate the derivative: f″(2) = 12

Beyond

The gradient is itself a function on the graph, so we can differentiate it. For f(x) = x^3 at x = 2: f'(2) = 3x^2 = 12 and f''(2) = 6x = 12: the values coincide at this point, and autograd computes both:

x3 = torch.tensor(2.0, requires_grad=True)
dy = torch.autograd.grad(x3 ** 3, x3, create_graph=True)[0]  # 3x^2 = 12
d2y = torch.autograd.grad(dy, x3)[0]                          # 6x  = 12
dy, d2y
(tensor(12., grad_fn=<MulBackward0>), tensor(12.))

Recap

Wrap-up

  • Record forward, sweep backward: the chain rule, automated and verified against 4\mathbf{x}.
  • One reverse sweep = the whole gradient (m{=}1, n in millions).
  • detach / no-grad keep values out of the graph.
  • Mind per-framework gradient handling (PyTorch accumulates).
  • The graph is built at runtime; control flow needs no special handling.
  • Higher-order derivatives: differentiate the gradient again.

Backpropagation through real networks gets its full treatment in the backpropagation section; forward vs. reverse mode is derived in the matrix-calculus-and-automatic-differentiation section.