Backpropagation, fully derived

The problem

Training means minimizing $\mathcal{L}(\theta)$ over millions (or trillions) of parameters by gradient descent (1.10 Optimization: from gradient descent to Adam), so we need $\nabla_\theta\mathcal{L}$ at every step. Finite differences — perturb one parameter, re-run the forward pass, repeat — would cost one forward pass per parameter: a 124M-parameter model would need 124 million forward passes per update. Backpropagation (Rumelhart, Hinton & Williams, 1986) computes the entire gradient for about twice the cost of one forward pass, regardless of the number of parameters. It is nothing more than the multivariable chain rule (1.4 Derivatives, gradients and the chain rule), applied in the right order and with intermediate results cached.

Computational graphs and the chain rule

Write the computation as a directed graph: nodes are intermediate values, edges say which values feed which operation. Each node knows its local derivative with respect to its inputs. The chain rule then says: the derivative of the output with respect to any node is the sum, over all paths from that node to the output, of the product of local derivatives along the path. Organizing this sum efficiently is the whole trick.

Take the one-neuron loss $\mathcal{L}=(\sigma(wx+b)-y)^2$ with $w=0.5,\ x=2,\ b=-0.5,\ y=1$.

Forward (left to right, caching every value): $z = wx+b = 0.5$; $a=\sigma(z)=0.6225$; $\mathcal{L}=(a-y)^2 = 0.1425$.

Backward (right to left, each node receives $\partial\mathcal{L}/\partial(\text{node})$ and multiplies by its local derivative):

Two things to notice. The backward pass reuses the forward values ($a$, $x$, $w$) — that is why they must be stored. And the "gradient so far" flows backward through the graph, being multiplied by one local derivative per node; nodes with several outputs sum the gradients arriving from each. The backprop-graph demo below lets you step through exactly this.

Deriving backprop for a 2-layer softmax classifier

Single example $\mathbf{x}\in\mathbb{R}^{d}$, label $y\in\{1,\dots,K\}$, one-hot $\mathbf{y}$; hidden width $m$:

Shapes: $\mathbf{W}^{(1)}\in\mathbb{R}^{m\times d}$, $\mathbf{W}^{(2)}\in\mathbb{R}^{K\times m}$. We go backward one layer at a time, naming $\boldsymbol{\delta}^{(\ell)} = \partial\mathcal{L}/\partial\mathbf{z}^{(\ell)}$ (the "error signal" at layer $\ell$).

Step 1 — the output logits. Using $\mathcal{L} = -z^{(2)}_{y} + \log\sum_j e^{z^{(2)}_j}$,

This is why softmax and cross-entropy are always used together: separately, the softmax Jacobian $p_i(\delta_{ij}-p_j)$ and the $-1/p_y$ from the log are messy; combined, they cancel into predicted minus actual (4.3 Loss functions catalog).

Step 2 — output weights and biases. Since $z^{(2)}_k = \sum_j W^{(2)}_{kj}h_j + b^{(2)}_k$, we have $\partial z^{(2)}_k/\partial W^{(2)}_{kj} = h_j$ and $\partial z^{(2)}_k/\partial b^{(2)}_k = 1$, so

An outer product: entry $(k,j)$ is the error at output $k$ times the activation of hidden unit $j$ — a weight changes in proportion to how active its input was and how wrong its output was (Hebbian learning with a sign).

Step 3 — back into the hidden activations. Each $h_j$ feeds every output, so the gradients sum: $\partial\mathcal{L}/\partial h_j = \sum_k \delta^{(2)}_k W^{(2)}_{kj}$, i.e.

The forward pass multiplied by $\mathbf{W}^{(2)}$; the backward pass multiplies by its transpose. That is always true for linear layers.

Step 4 — through the nonlinearity. $\phi$ acts element-wise, so its Jacobian is diagonal and the chain rule is an element-wise product:

For ReLU, $\phi'$ is a 0/1 mask: the error only flows through units that were active in the forward pass.

Step 5 — first-layer weights and biases, by the same argument as Step 2:

The general recursion. For any depth,

Unrolling the recursion gives exactly the product of $\mathbf{W}^\top\mathrm{diag}(\phi')$ factors from 4.2 Activation functions — the vanishing-gradient argument is this formula read aloud.

Batched form. With $B$ examples as rows, $\mathbf{H}^{(\ell)}\in\mathbb{R}^{B\times n_\ell}$ and $\boldsymbol{\Delta}^{(\ell)}\in\mathbb{R}^{B\times n_\ell}$, the weight gradient for the mean loss is $\nabla_{\mathbf{W}^{(\ell)}} = \frac{1}{B}\boldsymbol{\Delta}^{(\ell)\top}\mathbf{H}^{(\ell-1)}$ (shape $n_\ell\times n_{\ell-1}$ ✓) and the bias gradient is the column mean of $\boldsymbol{\Delta}^{(\ell)}$. Every per-example outer product collapses into one matrix multiply, which is why the backward pass is as GPU-friendly as the forward pass.

Vector–Jacobian products: why reverse mode

Think of each layer as a function $\mathbf{h}^{(\ell)} = f_\ell(\mathbf{h}^{(\ell-1)})$ with Jacobian $\mathbf{J}_\ell\in\mathbb{R}^{n_\ell\times n_{\ell-1}}$. The full chain rule is $\nabla_{\mathbf{x}}\mathcal{L}^\top = \nabla_{\mathbf{h}^{(L)}}\mathcal{L}^\top\,\mathbf{J}_L\cdots\mathbf{J}_1$. Multiplying from the left, starting with the row vector $\nabla_{\mathbf{h}^{(L)}}\mathcal{L}^\top$, each step is a vector–Jacobian product (VJP) $\mathbf{v}^\top\mathbf{J}_\ell$ — a vector in, a vector out, and the Jacobian itself is never materialized. (A Linear layer's VJP is $\mathbf{v}^\top\mathbf{W}$; ReLU's is a mask. That is all Steps 3–4 were.) This is reverse-mode automatic differentiation. Multiplying from the right (forward mode, Jacobian–vector products) would require one pass per input dimension. For a scalar loss and millions of inputs, reverse mode wins by a factor of millions — and forward mode wins when there are few inputs and many outputs, which is why it appears in some physics simulations but not in training.

What backprop costs

Compute. Each Linear layer does one matmul in the forward pass ($\mathbf{H}\mathbf{W}^\top$) and two in the backward pass: one for the gradient with respect to the input ($\boldsymbol{\Delta}\mathbf{W}$) and one for the gradient with respect to the weights ($\boldsymbol{\Delta}^\top\mathbf{H}$). So backward ≈ 2× forward, and a full training step ≈ 3× a forward pass. Since a forward pass costs about $2N$ FLOPs per token for a model with $N$ parameters (4.12 Compute, hardware and scale basics), training costs about $6N$ per token — the constant in the scaling-law formula $C\approx 6ND$ (7.9 Scaling laws and compute-optimal training).

Memory. Step 2 needed $\mathbf{h}$, Step 4 needed $\mathbf{z}^{(1)}$, Step 5 needed $\mathbf{x}$: every layer's input must be kept alive from the forward pass until its backward step. Activation memory therefore grows with batch size × sequence length × depth × width, and it — not the parameters — is what overflows your GPU first. Activation checkpointing stores only some layers' activations and recomputes the rest during the backward pass, trading ~30% more compute for a large memory saving (4.12 Compute, hardware and scale basics).

Gradient checking

Backprop is easy to get subtly wrong when you implement layers yourself. Verify against the centered finite difference

in float64, and compare with the relative error $\frac{|g_{\text{bp}}-g_{\text{num}}|}{|g_{\text{bp}}|+|g_{\text{num}}|}$: below $10^{-7}$ is right, above $10^{-4}$ is a bug. Check a handful of random parameters, not all of them.

How autograd implements it

PyTorch builds the graph dynamically as your code runs: every tensor operation records its inputs and a function that computes its VJP. Calling loss.backward() sorts the graph topologically and walks it in reverse, accumulating each tensor's .grad (which is why you call optimizer.zero_grad() each step). Karpathy's micrograd is the whole mechanism in 100 lines:

class Value:
    def __init__(self, data, children=(), op=""):
        self.data, self.grad = data, 0.0
        self._backward, self._prev = lambda: None, set(children)

    def __mul__(self, other):
        out = Value(self.data * other.data, (self, other), "*")
        def _backward():                 # local derivative × upstream gradient
            self.grad  += other.data * out.grad
            other.grad += self.data  * out.grad
        out._backward = _backward
        return out

    def backward(self):                  # reverse topological order
        topo, seen = [], set()
        def build(v):
            if v not in seen:
                seen.add(v); [build(c) for c in v._prev]; topo.append(v)
        build(self)
        self.grad = 1.0
        for v in reversed(topo): v._backward()

The += is the "sum over paths" rule; the reversed topological order guarantees every node's upstream gradient is complete before it propagates further. Tensor autograd is the same idea with matrices and a library of hand-written VJPs.

Where this shows up later

• The backward recursion through many layers is the vanishing/exploding-gradient mechanism — 4.2 Activation functions, 4.6 Initialization and normalization; optimizers consume the gradients it produces — 4.5 Optimizers and learning-rate schedules.
• Backprop through time unrolls the same recursion across time steps — 6.4 Recurrent neural networks and BPTT; attention's backward pass and FlashAttention's recomputation trick — 7.2 Scaled dot-product and multi-head attention, fully derived, 7.10 Efficient inference: KV cache, GQA, FlashAttention, speculative decoding, quantization.
• The 6N FLOPs-per-token cost model — 4.12 Compute, hardware and scale basics, 7.9 Scaling laws and compute-optimal training; activation memory and checkpointing — 4.12 Compute, hardware and scale basics; LoRA cuts optimizer memory but not activation memory — 8.2 Fine-tuning and PEFT.
• Debugging with gradient norms and gradient checks — 4.8 Training dynamics and debugging.