A large language model has billions of parameters, and every training step needs the derivative of the loss with respect to every one of them. If you estimated those derivatives with finite differences, as in §0.2, you would need one extra forward pass per parameter, billions of passes per step. Training would never finish.

Yet loss.backward() computes all of them at a cost of roughly two forward passes. That is not an approximation. It is exact, and it works for any program built from differentiable pieces. The method is reverse-mode automatic differentiation, and the neural-network special case is called backpropagation. It is arguably the single most important algorithm in modern AI, and it is short enough that you will build a working version of it in this section, then a second version that trains a real network.

A computation is a graph

Any computation built from simple operations can be drawn as a directed acyclic graph (a DAG). Each node is a value, and each edge says "this value was used to compute that one". Take

with $a = 1$, $b = -0.5$, $c = 2$. The forward pass evaluates nodes in an order where every node comes after its inputs, giving $e = -0.5$, $d = 0.5$, $t = \tanh(0.5) = 0.4621$, and $\mathcal{L} = 0.9242$. Notice that $a$ is used twice, once in $e$ and once in $d$. That detail will matter.

Now we want $\partial\mathcal{L}/\partial a$, $\partial\mathcal{L}/\partial b$, and $\partial\mathcal{L}/\partial c$. Step through the widget below, first forward, then backward, and watch each node receive its gradient.

Reverse mode, one local rule at a time

The backward pass walks the graph in the opposite order, from the loss toward the inputs. For every node $v$ it computes the adjoint $\bar v = \partial\mathcal{L}/\partial v$, how much the loss would change per unit change in $v$. It starts with $\bar{\mathcal{L}} = 1$, since the loss changes one-for-one with itself. Then each node passes gradient to its inputs using only its local derivative, by the chain rule of §0.2.

• $\mathcal{L} = t \cdot c$, so $\bar t = c \cdot \bar{\mathcal{L}} = 2$ and $\bar c = t \cdot \bar{\mathcal{L}} = 0.4621$.
• $t = \tanh(d)$, so $\bar d = (1 - t^2)\, \bar t = 0.7864 \times 2 = 1.5729$.
• $d = e + a$. Addition passes gradient through unchanged, so $\bar e = \bar d = 1.5729$, and $a$ receives $1.5729$ from this edge.
• $e = a \cdot b$, so $\bar b = a\, \bar e = 1.5729$, and $a$ receives $b\, \bar e = -0.7864$ from this edge.

So $a$ gets gradient from two places, and its total is the sum, $\bar a = 1.5729 - 0.7864 = 0.7864$. This is the multivariate chain rule, which says a variable's influence on the loss is the sum of its influences along every path. In code it means adjoints accumulate with +=, never with =. Getting this wrong is the most common bug in hand-written autograd engines. (PyTorch sums the contributions of every path inside its engine during a backward pass, and separately it adds each backward pass's result into .grad rather than overwriting it, which we come back to below.)

The order matters, so we sort

A node may only pass gradient to its inputs once its own adjoint is complete, meaning every node that uses it has already sent its contribution. In our graph, $a$ must wait for both $e$ and $d$. The order that guarantees this is a topological order of the graph, every node after all its inputs, run in reverse.

A depth-first search produces one. Visit a node, recursively visit all its inputs first, then append the node to a list. Each node lands in the list after everything it depends on. Reversing the list gives an order in which every node comes before its inputs, which is exactly what the backward pass needs.

def topo_order(root, inputs_of):
    raise NotImplementedError

A scalar autograd engine in about fifty lines

We now have everything. Each value is an object that stores its number (data), its adjoint (grad), the values it was computed from, and a small closure, _backward, that applies its local rule. Every operation builds a new value and attaches that closure. This is the design of Karpathy's micrograd (linked below), and in miniature it is the design of PyTorch.

The backward pass sorts the graph and calls every closure in reverse order.

These match the hand computation, including the two-path gradient of $a$. A tiny engine is already enough to train a neural network. Here a network with one hidden layer of 4 tanh units learns XOR, the function no single linear unit can represent, using nothing but Value objects.

The outputs head toward $(-1, 1, 1, -1)$. Note the line that zeroes the gradients before each backward pass. Because adjoints accumulate, skipping it would add this step's gradient to the last one's, which is exactly the optimizer.zero_grad() you call in PyTorch.

import math

class Value:
    def __init__(self, data, inputs=()):
        self.data, self.grad = float(data), 0.0
        self._inputs = inputs
        self._backward = lambda: None

Forward mode, and why machine learning runs in reverse

There is a second way to apply the chain rule automatically. Instead of sending adjoints backward from the output, forward mode sends derivatives forward from one input. Attach to every value $v$ a tangent $\dot v = \partial v / \partial x$ for one chosen input $x$, set $\dot x = 1$, and update tangents alongside values. For $v = a \cdot b$, the product rule gives $\dot v = \dot a\, b + a\, \dot b$. For $v = \tanh(a)$, it is $\dot v = (1 - v^2)\, \dot a$. When the computation finishes, the output's tangent is the derivative.

A neat way to implement this is with dual numbers, pairs $a + \dot a\,\epsilon$ with the rule $\epsilon^2 = 0$. Multiplying two of them gives $(a + \dot a \epsilon)(b + \dot b\epsilon) = ab + (\dot a b + a \dot b)\epsilon$, the product rule for free. Overloading arithmetic on such pairs gives exact derivatives of any program.

The two modes compute the same derivatives at very different costs. For a function $f: \mathbb{R}^n \to \mathbb{R}^m$, one forward-mode pass computes one column of the $m \times n$ Jacobian (a Jacobian-vector product), and one reverse-mode pass computes one row (a vector-Jacobian product). Each pass costs a small constant multiple of evaluating $f$ (Griewank & Walther, 2008, show that for the usual operation sets, the function and its full gradient together cost at most about 3 to 4 times the function alone in reverse mode).

Training is the extreme case of the second row. There are $n = 10^9$ inputs (the parameters) and $m = 1$ output (the loss). Reverse mode gets the entire gradient in one backward pass. Forward mode would need a billion passes. That asymmetry is the whole reason deep learning frameworks are built around reverse mode.

The price of reverse mode is memory. The backward pass needs values from the forward pass (multiplication's rule needs both factors, tanh's needs its output), so the engine must keep them until the backward pass reaches them. For a large model those stored activations are often the biggest consumer of GPU memory, more than the weights. You already have one number. The attention scores of GPT-2 small, kept for the backward pass, take about 4.8 GB at a batch of 8 sequences of 1,024 tokens (§0.1), ten times the fp32 weights. Trading compute for memory by recomputing some of them during the backward pass is called activation checkpointing, and it is a standard tool in where the memory goes (§10.1).

import math

class Dual:
    def __init__(self, val, eps=0.0):
        self.val, self.eps = float(val), float(eps)

def derivative(f, x):
    raise NotImplementedError

From scalars to tensors

A scalar engine creates one Python object per number, which is far too slow for real networks. Real frameworks apply the same idea to whole tensors. Each node holds an array, each operation's backward rule is a tensor VJP (the ones we derived in §0.2), and a matrix multiply is one node instead of millions.

One new issue appears, and it is worth understanding properly. Tensor operations broadcast (§0.1). When a bias $\mathbf{b} \in \mathbb{R}^{d}$ is added to activations of shape $(N, d)$, the bias is used $N$ times, once per row. By the accumulation rule, its gradient is the sum of the $N$ row gradients. In general, if an input of shape $s$ was broadcast to a larger output shape, the input's gradient is the output gradient summed over every axis that was stretched. Framework code calls this "unbroadcasting" or "sum to size", and every tensor engine needs it.

import numpy as np

def unbroadcast(grad, shape):
    raise NotImplementedError

Here is a small tensor engine with exactly four operations, broadcasting addition, matrix multiplication, ReLU, and a fused softmax cross-entropy whose backward rule is the $\mathbf{p} - \mathbf{y}$ of §0.2. It uses an iterative topological sort, so deep graphs are fine.

And it trains a real classifier. Three interleaved spirals cannot be separated by any linear model, but a 64-unit hidden layer trained by our engine separates them, reaching 100% training accuracy.

Notice that the engine never needed to know it was training a neural network. It knows four VJP rules and how to order them. Everything PyTorch adds on top (hundreds of operations, GPU kernels, a C++ core, memory management) is engineering around this same idea.

What PyTorch's autograd adds

PyTorch's engine is our engine with some important refinements. Knowing them prevents a whole class of bugs.

• Leaves and requires_grad. Tensors you create directly are leaves. Only leaves with requires_grad=True (parameters, typically) keep their .grad after backward(). Intermediate results get gradients during the pass, but PyTorch frees them unless you call .retain_grad().
• The graph is freed after backward(). Calling backward() twice on the same graph raises an error unless you pass retain_graph=True, as long as any operation in it saved tensors for the backward pass (almost all do), because those saved activations were released to save memory.
• Gradients accumulate. Like our +=, each backward() adds into .grad. That is a feature, since it lets you sum gradients over several small batches before one optimizer step (gradient accumulation, §0.9). It is also why you must zero gradients between steps.
• In-place operations are tracked. If you modify a tensor in place after it was saved for the backward pass, PyTorch notices (each tensor has a version counter) and raises an error instead of silently using the wrong value.

Stopping gradients with detach and no_grad

Sometimes you want a value in the forward pass but no gradient through it. x.detach() returns a tensor that shares x's data but is cut off from the graph, so gradient stops there. You'll use it for targets that should not be trained toward, such as the reference values in some RL losses, and for the stop-gradient trick that keeps self-supervised methods from collapsing, as in the JEPA family of models, where one branch's output is treated as a fixed target. Wrapping code in torch.no_grad() is different. It tells PyTorch not to record a graph at all inside the block, which saves the memory of stored activations. Use it for evaluation and for the parameter update itself. torch.inference_mode() is a stricter, slightly faster version for pure inference.

Higher-order gradients

The backward pass is itself a computation built from differentiable operations. If the engine records it as a graph too (PyTorch does this with create_graph=True), you can differentiate the gradient, which gives second derivatives.

Write $\theta$ for all of a model's parameters, $n$ of them, and $\mathbf{H} = \nabla_\theta^2\mathcal{L}$ for the Hessian, the $n \times n$ matrix of second derivatives. It has $n^2$ entries, hopeless for large $n$. But the Hessian-vector product $\mathbf{H}\mathbf{v}$ is cheap. Since $\mathbf{H}\mathbf{v} = \nabla_\theta\big(\nabla_\theta \mathcal{L}^\top \mathbf{v}\big)$, where $\mathbf{v}$ is held constant, you can compute it with one backward pass through the gradient computation, at a cost of about two to three gradient evaluations (Pearlmutter, 1994). It is also approximately the finite difference of two gradients, $\mathbf{H}\mathbf{v} \approx (\nabla\mathcal{L}(\theta + h\mathbf{v}) - \nabla\mathcal{L}(\theta - h\mathbf{v}))/2h$. Hessian-vector products power curvature estimates (the sharpness of a minimum, or the largest Hessian eigenvalue that limits stable step sizes in §0.6, which power iteration finds from HVPs alone, as §0.10 shows), influence functions, some second-order optimizers, and meta-learning methods like MAML that differentiate through an inner training step.

import numpy as np

def hvp(grad_f, x, v, h=1e-5):
    raise NotImplementedError

def hessian(grad_f, x):
    raise NotImplementedError

In PyTorch

The modern torch.func module offers these as composable functions (grad, vjp, jvp, jacrev, jacfwd, hessian), in the style of JAX. Note that it exposes forward mode too, as jvp and jacfwd, for the cases in the table above where it wins.

Read more

• Karpathy. micrograd. The tiny scalar engine this section's design follows, with a neural-network library on top, in about 150 lines.
• Baydin, Pearlmutter, Radul & Siskind (2018). Automatic Differentiation in Machine Learning: a Survey. The clearest overview of forward and reverse mode, dual numbers, and their costs.
• Griewank & Walther (2008). Evaluating Derivatives: Principles and Techniques of Algorithmic Differentiation. The standard reference on automatic differentiation, including the cheap-gradient bound.
• Rumelhart, Hinton & Williams (1986). Learning representations by back-propagating errors. The paper that made backpropagation the way neural networks are trained.
• Paszke et al. (2019). PyTorch: An Imperative Style, High-Performance Deep Learning Library. How PyTorch records graphs on the fly and runs reverse mode efficiently.
• PyTorch documentation. Autograd mechanics. Leaves, saved tensors, in-place correctness checks, and inference mode.
• Pearlmutter (1994). Fast Exact Multiplication by the Hessian. The Hessian-vector product trick.

Watch next

▶

▶

▶