Backpropagation on computational graphs

Before this lesson: lesson 01 (the MLP in matrix form) and lesson 02 (losses, softmax cross-entropy, the training loop). From other tracks: the chain rule as a product of Jacobians (calculus track); JVPs, VJPs and forward vs reverse mode in linear-algebra lesson 15; round-off and the unit roundoff $u$ in linear-algebra lesson 13. Time: about 65 minutes, plus the notebook.

---

Why this matters #

Every call to loss.backward() in PyTorch or jax.grad in JAX runs the algorithm of this lesson. With the explicit matrix formulas you can:

• write a custom backward pass (a PyTorch autograd.Function or a fused GPU kernel) and find bugs where a gradient is silently wrong: a wrong sum over the batch, a missing mask, a missing transpose;
• predict the cost of training. One training step of a model with $N$ parameters costs about $6N$ floating-point operations (FLOPs) per token, and this number sets the compute budgets of large language models;
• understand where training memory goes and what activation checkpointing does;
• trust a new layer, because a gradient check confirms its backward pass.

In ML for science, the same reverse mode is called the adjoint method. It gives the sensitivity of a simulator or a PDE solver to all its parameters at once.

1. Predict first #

2. The picture #

The chain rule, read from the end #

Start with a chain of scalars: $x \to u = g(x) \to \mathcal{L} = f(u)$. The chain rule gives $d\mathcal{L}/dx = f'(u)\, g'(x)$. Backpropagation evaluates this product from the end:

1. Start with $\partial\mathcal{L}/\partial\mathcal{L} = 1$.
2. Multiply by the local derivative $f'(u)$. This gives $\partial \mathcal{L}/\partial u$.
3. Multiply by the local derivative $g'(x)$. This gives $\partial \mathcal{L}/\partial x$.

We write $\bar u = \partial \mathcal{L}/\partial u$ and call it the adjoint of $u$: the derivative of the final scalar loss with respect to that variable. Each step needs only two numbers: the adjoint that arrives from above (the upstream gradient) and the derivative of one operation (the local gradient). The rule to remember is

In this lesson $\mathcal{L}$ is always the scalar loss and $L$ is the number of layers.

Computational graphs #

A computational graph has one node per operation and one edge per value. Here is a tiny network with one hidden ReLU unit, a residual connection and a squared loss:

The value $h$ is used twice: directly in $\hat y = h + v$ (the residual path) and inside $v = w_2 h$. We say that $h$ fans out. Two rules make backprop work on any such graph:

1. Each node maps the adjoint of its output to adjoints of its inputs (upstream times local). For vectors this is a vector-Jacobian product (VJP): $\bar x^\top = \bar u^\top\, \partial u/\partial x$. The full Jacobian $\partial u/\partial x$ is never formed.
2. Fan-out sums. If a variable feeds several nodes, its adjoint is the sum of the contributions from all of them. This is the multivariable chain rule. Residual connections $x + F(x)$ and tied (shared) weights are fan-out.

Interactive figure: loading…

The default numbers are $x = 1.5$, $w_1 = 0.8$, $b_1 = -0.2$, $w_2 = 0.5$, $y = 1$.

Forward pass (top to bottom): $u = 0.8 \cdot 1.5 = 1.2$; $z = 1.2 - 0.2 = 1.0$; $h = \operatorname{relu}(1.0) = 1.0$; $v = 0.5 \cdot 1.0 = 0.5$; $\hat y = 1.0 + 0.5 = 1.5$; $r = 1.5 - 1 = 0.5$; $\mathcal{L} = 0.5^2 = 0.25$.

Backward pass (from the loss down). Each line is "upstream × local":

adjoint	upstream × local	value
$\bar{\mathcal{L}}$	(the seed)	$1$
$\bar r$	$\bar{\mathcal{L}} \cdot 2r$	$1 \cdot 1.0 = 1.0$
$\bar{\hat y}$	$\bar r \cdot 1$	$1.0$
$\bar v$	$\bar{\hat y} \cdot 1$	$1.0$
$\bar w_2$	$\bar v \cdot h$	$1.0 \cdot 1.0 = 1.0$
$\bar h$	$\bar{\hat y} \cdot 1 + \bar v \cdot w_2$ (two paths)	$1.0 + 0.5 = 1.5$
$\bar z$	$\bar h \cdot \operatorname{relu}'(z)$	$1.5 \cdot 1 = 1.5$
$\bar b_1$, $\bar u$	$\bar z \cdot 1$	$1.5$
$\bar w_1$	$\bar u \cdot x$	$1.5 \cdot 1.5 = 2.25$

Check one of them directly: $\mathcal{L} = \big((1 + w_2)\operatorname{relu}(w_1 x + b_1) - y\big)^2$, so $\partial\mathcal{L}/\partial w_1 = 2r\,(1 + w_2)\,x = 2 \cdot 0.5 \cdot 1.5 \cdot 1.5 = 2.25$ (the ReLU is active, so its derivative is 1). The same number, without ever writing the whole formula.

This answers question 3 of section 1: the contributions add, so $\bar h = 1.0 + 0.5 = 1.5$.

Why reverse order? #

The loss is a map $\mathcal{L}: \mathbb{R}^n \to \mathbb{R}$ from $n$ parameters to one number. Its gradient is one row of length $n$. The chain rule writes it as a product of Jacobians: $\nabla\mathcal{L}^\top = J_k J_{k-1} \cdots J_1$, where $J_k$ (the last one) has only one row.

• Reverse mode multiplies from the left. It starts with the $1 \times 1$ seed and every intermediate result is a single row. Each step is a VJP and costs about as much as the forward operation. One backward pass gives all $n$ derivatives.
• Forward mode multiplies from the right. It pushes one input direction through the graph per pass (a Jacobian-vector product, JVP). It needs $n$ passes for $n$ parameters.

For a scalar loss and millions of parameters, reverse mode wins by a factor of millions. Linear-algebra lesson 15 treats JVPs and VJPs in detail.

Interactive figure: loading…

3. Definitions and main results #

The setting #

A batch has $B$ examples, one per row (the row convention of lesson 01). A linear layer computes

where $\mathbf{1} \in \mathbb{R}^{B \times 1}$ is a column of ones, so $\mathbf{1}b$ adds the row $b$ to every row. Write $G = \partial \mathcal{L}/\partial Z$. It has the same shape as $Z$: entry $G_{bj}$ is the derivative of the loss with respect to entry $Z_{bj}$. In the same way, $\bar W$ has the shape of $W$ and $\bar A$ the shape of $A$.

Rule 1: the linear layer #

First, entry by entry. Entry $W_{ij}$ appears in column $j$ of $Z$, in every row: $Z_{bj} = \sum_k A_{bk} W_{kj} + b_j$, so $\partial Z_{bj}/\partial W_{ij} = A_{bi}$. By the chain rule, summed over every entry of $Z$ that $W_{ij}$ touches (one per example):

In the same way, $A_{bi}$ appears in row $b$ of $Z$, in every column, with $\partial Z_{bj}/\partial A_{bi} = W_{ij}$. So $\bar A_{bi} = \sum_j G_{bj} W_{ij} = (G W^\top)_{bi}$. And $b_j$ appears in every row of column $j$, so $\bar b_j = \sum_b G_{bj}$.

Theorem (backward pass of a linear layer). For $Z = AW + \mathbf{1}b$ and $G = \partial\mathcal{L}/\partial Z$:

Shapes as a memory aid. $\bar W$ must have the shape of $W$, $n \times m$. From $A$ ($B \times n$) and $G$ ($B \times m$), the only product with that shape is $A^\top G$. From $G$ ($B \times m$) and $W$ ($n \times m$), the only product with the shape of $A$ is $GW^\top$. When you forget a formula, check the shapes.

A compact proof with differentials

For a scalar $\mathcal{L}$, $d\mathcal{L} = \sum_{b,j} G_{bj}\, dZ_{bj} = \langle G, dZ \rangle$, the Frobenius inner product $\langle X, Y \rangle = \operatorname{tr}(X^\top Y)$. If we can write $d\mathcal{L} = \langle M, dW \rangle$, then $\bar W = M$. Now $dZ = dA\, W + A\, dW + \mathbf{1}\, db$. Move each factor to the other side of the inner product (cyclic property of the trace):

1. $\langle G, A\,dW \rangle = \langle A^\top G, dW \rangle$, so $\bar W = A^\top G$.
2. $\langle G, dA\,W \rangle = \langle G W^\top, dA \rangle$, so $\bar A = G W^\top$.
3. $\langle G, \mathbf{1}\,db \rangle = \langle \mathbf{1}^\top G, db \rangle$, so $\bar b = \mathbf{1}^\top G$. $\blacksquare$

Rule 2: an elementwise activation #

For $A = \sigma(Z)$, applied entry by entry, each $A_{bj}$ depends only on $Z_{bj}$. So the Jacobian is diagonal, and the VJP is an elementwise product ($\odot$):

For ReLU, $\sigma'(Z)$ is the mask $\mathbf{1}[Z > 0]$: the gradient passes where the unit was active and is blocked where it was off.

Rule 3: softmax cross-entropy #

Take one example with logits $z \in \mathbb{R}^K$, probabilities $p = \operatorname{softmax}(z)$ and a target $y$ with $\sum_c y_c = 1$ (one-hot, or smoothed). The loss is $\ell = -\sum_c y_c \log p_c$, and $\log p_c = z_c - \log \sum_j e^{z_j}$. So

Lesson 02 derived this too. For the batch mean $\mathcal{L} = \frac1B \sum_b \ell_b$, each example's loss is divided by $B$, so

The softmax Jacobian and its VJP

The softmax Jacobian is $J = \operatorname{diag}(p) - pp^\top$ ($K \times K$). Its VJP with a row $v^\top$ is $v^\top J = \big(p \odot (v - (v^\top p)\mathbf{1})\big)^\top$, computed in $O(K)$ without forming $J$. With $v = \partial\ell/\partial p = -y/p$ (entrywise), $v^\top p = -\sum_c y_c = -1$, and the formula gives $p \odot (-y/p + \mathbf{1}) = p - y$ again. Attention (lesson 09) applies a softmax to every row of a score matrix and uses this VJP.

Theorem: the MLP backward pass #

Forward pass. $A_0 = X$; $Z_l = A_{l-1}W_l + \mathbf{1}b_l$ for $l = 1, \dots, L$; $A_l = \sigma(Z_l)$ for $l < L$; $P = \operatorname{softmax}(Z_L)$ row by row; and $\mathcal{L} = -\frac1B \sum_{b,c} Y_{bc} \log P_{bc}$ with one-hot rows in $Y$.

Theorem. With $G_l := \partial \mathcal{L}/\partial Z_l$, for $l = L, L-1, \dots, 1$:

1. $G_L = (P - Y)/B$.
2. $\bar W_l = A_{l-1}^\top G_l$ and $\bar b_l = \mathbf{1}^\top G_l$.
3. $\bar A_{l-1} = G_l W_l^\top$.
4. $G_{l-1} = (G_l W_l^\top) \odot \sigma'(Z_{l-1})$ (for $l \ge 2$).

Proof (induction from the top layer down).

1. Step 1 is rule 3.
2. Suppose $G_l$ is correct. The graph is a chain: $\mathcal{L}$ depends on $W_l$, $b_l$ and $A_{l-1}$ only through $Z_l$. Rule 1 with $A = A_{l-1}$ and $G = G_l$ gives steps 2 and 3.
3. $A_{l-1} = \sigma(Z_{l-1})$ elementwise, and $\mathcal{L}$ depends on $Z_{l-1}$ only through $A_{l-1}$. Rule 2 gives step 4, which is $G_{l-1}$. So the claim holds for $l - 1$. $\blacksquare$

Note where $1/B$ enters: once, in $G_L$. The products $A_{l-1}^\top G_l$ and $\mathbf{1}^\top G_l$ then sum over the batch, so every gradient is a mean over the batch.

Interactive figure: loading…

Cost: time and memory #

Time. A forward matmul $AW$ with $A \in \mathbb{R}^{B \times n}$, $W \in \mathbb{R}^{n \times m}$ costs about $2Bnm$ FLOPs (one multiply and one add per term). The backward pass needs two matmuls of the same size, $A^\top G$ and $G W^\top$. So the backward pass costs about 2× the forward pass, and one training step costs about 3× the forward pass. Elementwise steps cost $O(Bm)$, which is small next to the matmuls.

• For a transformer, the forward pass costs about $2N$ FLOPs per token for $N$ parameters (one multiply-add per weight). So a training step costs about $6N$ FLOPs per token (attention scores are not counted in this estimate).
• The cheap gradient principle: reverse mode computes the gradient of any scalar program for a small constant times the cost of the program, whatever the number of inputs. Baur and Strassen (1983) prove a constant between 3 and 5 for arithmetic circuits, depending on how operations are counted.

Memory. Steps 2–4 for layer $l$ need $A_{l-1}$ and $\sigma'(Z_{l-1})$. The forward pass produces them in the order $1, \dots, L$; the backward pass uses them in the order $L, \dots, 1$. So all of them are in memory when the backward pass starts: $O(L \cdot B \cdot \text{width})$ numbers, on top of the weights, gradients and optimizer state. For large models with long sequences, these stored activations are the largest part of training memory.

What each operation saves for the backward pass: a matmul saves its input; ReLU needs only the sign (PyTorch saves its output); GELU saves its input; dropout saves its mask; softmax saves its output.

Activation checkpointing #

Activation checkpointing (also called rematerialization) trades compute for memory. Split the $L$ layers into $s$ segments.

1. Forward pass: store only the input of each segment. Do not keep the activations inside a segment.
2. Backward pass, segment by segment from the last one: run the segment's forward pass again from its stored input (now keeping its activations), backpropagate through it, then free them.

• Memory: about $s + L/s$ layer activations ($s$ segment inputs, plus the $L/s$ activations of one segment). This is smallest at $s = \sqrt{L}$, where it equals $2\sqrt{L}$.
• Compute: at most one extra forward pass. A step costs about $4F$ instead of $3F$, about 33% more.

Interactive figure: loading…

For $L = 16$ layers (PyTorch's checkpoint_sequential, counted in section 5):

segments $s$	layer forward calls	stored activations
1 (no checkpointing)	16	17
2	24	10
4	28	8
16	31	17

With $s = 4 = \sqrt{16}$ only 8 activations are stored, for 12 extra layer calls. With 16 segments nothing is gained: every layer input is a segment input.

More memory-compute trades

• Recursive schedules reach $O(\log L)$ memory for $O(L \log L)$ forward compute (Chen et al. 2016, §4).
• In practice, large models checkpoint each transformer block and keep only the block inputs. Megatron's selective recomputation (Korthikanti et al. 2022) recomputes only the attention-score part, which is large in memory but cheap in FLOPs.
• FlashAttention's backward pass recomputes the attention matrix from $Q$, $K$, $V$ and the softmax statistics instead of storing it: the same trade inside one GPU kernel.

Gradient checking #

A gradient check compares the gradient from your backward pass with a numerical estimate. The central difference with step $h$ in coordinate $i$ is

where $e_i$ is the $i$-th unit vector. (PyTorch's gradcheck calls the step eps. We write $h$ because $\varepsilon$ means machine epsilon in linear-algebra lesson 13.) Two errors compete.

Truncation error. By Taylor's theorem, $f(x \pm h) = f \pm h f' + \tfrac{h^2}{2} f'' \pm \tfrac{h^3}{6} f''' + O(h^4)$. In the difference, the even terms cancel:

So the error is $O(h^2)$. The forward difference $(f(x+h) - f(x))/h$ keeps the $f''$ term and has error $O(h)$, which is much worse.

Round-off error. Each computed value of $f$ has a relative error of about the unit roundoff $u$ ($u = 2^{-53} \approx 1.1 \cdot 10^{-16}$ in float64, $u = 2^{-24} \approx 6 \cdot 10^{-8}$ in float32). The difference of two nearly equal numbers keeps this absolute error, about $u|f|$, and the division by $2h$ makes it about $u|f|/h$. It grows when $h$ gets small.

The best step. Add the two errors and minimize over $h$:

When $|f| \approx |f'''|$, this gives $h^* \approx (3u)^{1/3} \approx 7 \cdot 10^{-6}$ in float64, with an error near $2 \cdot 10^{-11}$. In float32, $h^* \approx 6 \cdot 10^{-3}$ and the best error is only about $10^{-5}$: too coarse to separate a small bug from round-off. Check gradients in float64.

Interactive figure: loading…

How to run a gradient check in practice.

• Use the relative error $|g_{\text{num}} - g| / \max(|g_{\text{num}}|, |g|, \text{tiny})$. Gradients of different parameters differ by orders of magnitude, and the relative error does not depend on scale. For a smooth function in float64, expect $10^{-7}$ or less. A relative error of $10^{-2}$ or more almost always means a bug (the rule of thumb of Stanford CS231n).
• Avoid kinks. If $x \pm h$ crosses a ReLU joint or a tie in a max, the numerical derivative is wrong. Use smooth activations for the check, or check that no pre-activation is within $h$ of 0.
• Save time on big models. The check needs two evaluations of $f$ per coordinate. Check a few random coordinates, or one random direction $v$: $(f(x + hv) - f(x - hv))/(2h) \approx \nabla f(x) \cdot v$.

Answers to section 1. (1) About $2F$: the backward pass does two matmuls of the same size for every forward matmul. (2) $10^{-5}$, close to $(3u)^{1/3}$. At $10^{-2}$ the truncation error dominates; at $10^{-10}$ and $10^{-15}$ the round-off error dominates.

4. Worked example #

A batch of $B = 2$ examples, two inputs, one hidden ReLU layer of width 2, $K = 2$ classes, all biases zero:

Forward pass.

1. $Z_1 = XW_1$. Row 1: $(1 \cdot 1 + (-1)(-1),\; 1 \cdot 2 + (-1) \cdot 1) = (2, 1)$. Row 2: $((-1) \cdot 1 + 0,\; (-1) \cdot 2 + 0) = (-1, -2)$. So $Z_1 = \begin{pmatrix} 2 & 1 \\ -1 & -2 \end{pmatrix}$ and $A_1 = \operatorname{relu}(Z_1) = \begin{pmatrix} 2 & 1 \\ 0 & 0 \end{pmatrix}$. Both hidden units are off for example 2.
2. $Z_2 = A_1 W_2$. Row 1: $(2 \cdot 1 + 1 \cdot (-2),\; 2 \cdot (-1) + 1 \cdot 2) = (0, 0)$. Row 2 is $(0, 0)$ because $A_1$ has a zero row.
3. All logits are 0, so every row of $P$ is $(\tfrac12, \tfrac12)$, and $\mathcal{L} = -\tfrac12(\log\tfrac12 + \log\tfrac12) = \log 2 \approx 0.6931$.

Backward pass.

1. $G_2 = (P - Y)/2 = \tfrac12\begin{pmatrix} -0.5 & 0.5 \\ 0.5 & -0.5 \end{pmatrix} = \begin{pmatrix} -0.25 & 0.25 \\ 0.25 & -0.25 \end{pmatrix}$.
2. $\bar W_2 = A_1^\top G_2 = \begin{pmatrix} 2 & 0 \\ 1 & 0 \end{pmatrix}\begin{pmatrix} -0.25 & 0.25 \\ 0.25 & -0.25 \end{pmatrix} = \begin{pmatrix} -0.5 & 0.5 \\ -0.25 & 0.25 \end{pmatrix}$, and $\bar b_2 = \mathbf{1}^\top G_2 = (0, 0)$.
3. $\bar A_1 = G_2 W_2^\top = \begin{pmatrix} -0.25 & 0.25 \\ 0.25 & -0.25 \end{pmatrix}\begin{pmatrix} 1 & -2 \\ -1 & 2 \end{pmatrix} = \begin{pmatrix} -0.5 & 1 \\ 0.5 & -1 \end{pmatrix}$.
4. $G_1 = \bar A_1 \odot \mathbf{1}[Z_1 > 0] = \begin{pmatrix} -0.5 & 1 \\ 0.5 & -1 \end{pmatrix} \odot \begin{pmatrix} 1 & 1 \\ 0 & 0 \end{pmatrix} = \begin{pmatrix} -0.5 & 1 \\ 0 & 0 \end{pmatrix}$. The mask removes example 2.
5. $\bar W_1 = X^\top G_1 = \begin{pmatrix} 1 & -1 \\ -1 & 0 \end{pmatrix}\begin{pmatrix} -0.5 & 1 \\ 0 & 0 \end{pmatrix} = \begin{pmatrix} -0.5 & 1 \\ 0.5 & -1 \end{pmatrix}$, and $\bar b_1 = (-0.5, 1)$.

Shape check. $\bar W_2$ is $2 \times 2$ like $W_2$; $\bar A_1$ is $B \times 2$ like $A_1$. Every gradient has the shape of its variable.

Read the result. Both hidden units are off for example 2. Its row of $A_1$ is zero, so it adds nothing to $\bar W_2$, and the ReLU mask removes it from $G_1$. So example 2 reaches only the output bias $b_2$, and there its gradient cancels the gradient of example 1. The first-layer weights $W_1$ learn from example 1 alone.

Finite differences on a known function. For $f = \sin$ at $x = 1$ (so $f' = \cos 1$), the central difference has these errors in float64:

$h$	$10^{-1}$	$10^{-3}$	$10^{-5}$	$10^{-7}$	$10^{-9}$	$10^{-11}$
error	$9.0 \cdot 10^{-4}$	$9.0 \cdot 10^{-8}$	$1.1 \cdot 10^{-11}$	$1.9 \cdot 10^{-10}$	$3.0 \cdot 10^{-9}$	$1.2 \cdot 10^{-6}$

For $h \ge 10^{-3}$ the error equals the truncation term $h^2 \cos(1)/6$: it falls by $10^4$ when $h$ falls by 100. Near $10^{-5}$ the two sources have the same size. Below that the error grows like $u/h$ (round-off).

5. In code #

NumPy checks the worked example, compares it with central differences, and reproduces the table of errors. The notebook has these cells and the exercises.

import numpy as np

X = np.array([[1., -1.], [-1., 0.]])
W1 = np.array([[1., 2.], [-1., 1.]])
W2 = np.array([[1., -1.], [-2., 2.]])
Y = np.eye(2)                                  # labels (0, 1) as one-hot rows
B = X.shape[0]

def loss(W1, W2):
    Z1 = X @ W1; A1 = np.maximum(Z1, 0); Z2 = A1 @ W2
    Z2 = Z2 - Z2.max(axis=1, keepdims=True)    # stable softmax
    P = np.exp(Z2) / np.exp(Z2).sum(axis=1, keepdims=True)
    return -np.mean(np.log((P * Y).sum(axis=1))), (Z1, A1, P)

L, (Z1, A1, P) = loss(W1, W2)
G2 = (P - Y) / B
dW2 = A1.T @ G2
G1 = (G2 @ W2.T) * (Z1 > 0)
dW1 = X.T @ G1
print(round(L, 4))                             # 0.6931
print(dW2.tolist(), G2.sum(axis=0).tolist())   # [[-0.5, 0.5], [-0.25, 0.25]] [0.0, 0.0]
print(dW1.tolist(), G1.sum(axis=0).tolist())   # [[-0.5, 1.0], [0.5, -1.0]] [-0.5, 1.0]

# central differences for every entry of W1
# (no pre-activation is near 0, so no ReLU kink is crossed)
h = 1e-6
num = np.zeros_like(W1)
for i in range(2):
    for j in range(2):
        E = np.zeros_like(W1); E[i, j] = h
        num[i, j] = (loss(W1 + E, W2)[0] - loss(W1 - E, W2)[0]) / (2 * h)
rel = np.abs(num - dW1).max() / np.abs(dW1).max()
print(rel < 1e-8)                              # True

# error of the central difference for sin'(1) against h
hs = 10.0 ** -np.arange(1, 12, 2)
err = np.abs((np.sin(1 + hs) - np.sin(1 - hs)) / (2 * hs) - np.cos(1))
print(" ".join(f"{a:.0e}:{e:.1e}" for a, e in zip(hs, err)))
# 1e-01:9.0e-04 1e-03:9.0e-08 1e-05:1.1e-11 1e-07:1.9e-10 1e-09:3.0e-09 1e-11:1.2e-06

PyTorch's autograd gives the same numbers, and it can count what checkpointing stores (16 Linear+ReLU layers, width 256):

import torch, torch.nn.functional as F

Xt, W1t, W2t = (torch.tensor(M, dtype=torch.float64) for M in (X, W1, W2))
W1t.requires_grad_(); W2t.requires_grad_()
F.cross_entropy(torch.relu(Xt @ W1t) @ W2t, torch.tensor([0, 1])).backward()
print(W1t.grad.tolist())                       # [[-0.5, 1.0], [0.5, -1.0]]
Xg = Xt.clone().requires_grad_()
print(torch.autograd.gradcheck(lambda t: torch.tanh(t @ W2t).sum(), (Xg,)))  # True
# checkpoint_sequential on 16 layers (the counting code is in the notebook)
# segments, forward calls, saved activations: 1 16 17 | 2 24 10 | 4 28 8 | 16 31 17

6. Where it shows up #

• Deep-learning frameworks. Autograd in PyTorch and JAX is rules 1–3 written for hundreds of operations. Each operation has a forward function and a VJP. Fused GPU kernels come with a hand-written backward pass, and a gradient check is how their authors test it.
• Large-scale training. The estimate of $6N$ FLOPs per token sets compute budgets and the model FLOPs utilization (MFU) of a training run. Activation memory decides the micro-batch size, where to checkpoint, and when to split sequences across GPUs.
• Normalization and residual connections (lesson 05) are fan-out: the residual stream $x + F(x)$ receives the adjoint from both paths, so gradients reach early layers directly.
• ML for science. The adjoint method for ODE and PDE solvers is reverse mode applied to the solver's steps; neural ODEs use its continuous-time form. Variational data assimilation in weather prediction (4D-Var) computes its gradient with an adjoint model, and long adjoint runs use checkpointing schedules such as Griewank's revolve.
• Physics-informed losses. A PINN (physics-informed neural network) loss contains derivatives such as $\partial u/\partial x$ computed by autograd. Training then backpropagates through that derivative (double backward, create_graph=True in PyTorch). This is one reason smooth activations matter (lesson 04).

7. Common mistakes #

• Dividing by $B$ twice, or not at all. Divide once, in $G_L$. Otherwise the gradient has the wrong scale and depends on $B$, and the learning rate seems to change when $B$ changes.
• A wrong order or a missing transpose. $\bar W = A^\top G$, not $G^\top A$ (that is $m \times n$, the transposed shape). Check shapes at every step.
• The mask from the wrong layer. Step 4 uses $\sigma'(Z_{l-1})$, the pre-activation below $W_l$, not $\sigma'(Z_l)$.
• Not summing at fan-out. Residual connections and shared weights get the sum of all contributions. In PyTorch, the opposite bug is forgetting optimizer.zero_grad(), so the gradients of different steps add up.
• A bad gradient check: float32, $h = 10^{-10}$, a step across a ReLU joint, or absolute instead of relative errors.
• In-place operations that overwrite a tensor saved for the backward pass. PyTorch stops with an error about a modified variable version.

8. Check yourself #

1. $A$ has shape $B \times n$ and $G = \partial \mathcal{L}/\partial Z$ has shape $B \times m$, for $Z = AW + \mathbf{1}b$. Give $\bar W$, $\bar A$, $\bar b$ and their shapes.
   Answer

   $\bar W = A^\top G$ ($n \times m$), $\bar A = GW^\top$ ($B \times n$), $\bar b = \mathbf{1}^\top G$ ($1 \times m$, the column sums of $G$ over the batch).
2. Why does $\partial \ell/\partial z = p - q$ also hold for a label-smoothed target $q = (1-\varepsilon)\,y + \varepsilon/K$ (one-hot $y$, smoothing $\varepsilon$, $K$ classes)?
   Answer

   The derivation used only $\sum_c q_c = 1$, and the smoothed target also sums to 1: $(1 - \varepsilon) + K \cdot \varepsilon/K = 1$.
3. You switch from the mean loss over the batch to the sum. What changes in the backward pass?
   Answer

   Only $G_L$ loses the factor $1/B$, so every gradient is $B$ times larger. With SGD that is the same as a $B$ times larger learning rate. Adam is almost invariant to it (except through its small constant in the denominator).
4. A 64-layer network uses $\sqrt{L}$ checkpointing. About how many layer activations are stored, and what is the extra compute?
   Answer

   $s = 8$ segments: about $s + L/s = 16$ activations instead of 64, for at most one extra forward pass (a step of about $4F$ instead of $3F$).
5. What is the truncation error of the central difference for $f = e^x$ at $x = 0$ with $h = 10^{-4}$?
   Answer

   $h^2 f'''(0)/6 = 10^{-8}/6 \approx 1.7 \cdot 10^{-9}$. (The quotient is exactly $\sinh(h)/h = 1 + h^2/6 + \dots$)
6. Reverse mode is right for training. When is forward mode better?
   Answer

   When there are few inputs and many outputs, for example the Jacobian of $\mathbb{R}^3 \to \mathbb{R}^{10^6}$ (the sensitivity of a field to 3 physical parameters): 3 forward passes (JVPs) against $10^6$ reverse passes (VJPs).

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. What is the adjoint $\bar u$ of a variable $u$?

   Show answer

2. How does backprop compute the adjoint of a node's input from the adjoint of its output?

   Show answer

3. A variable feeds two operations that send it adjoints $a$ and $b$. What is its adjoint?

   Show answer

4. For $Z = AW + \mathbf{1}b$ with $G = \partial\mathcal{L}/\partial Z$, what are $\bar W$ and $\bar A$?

   Show answer

5. What is $G_L$ for softmax cross-entropy averaged over a batch of $B$ examples?

   Show answer

6. Write the backward recursion from $G_l$ to $G_{l-1}$ in an MLP.

   Show answer

7. Why does training use reverse mode and not forward mode?

   Show answer

8. About how much does one training step cost compared with a forward pass, and how many FLOPs per token for $N$ parameters?

   Show answer

9. Which step $h$ is best for a central-difference check in float64, and why?

   Show answer

10. What does activation checkpointing trade, and what is the best number of segments?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "Derive the gradient of softmax cross-entropy." $p - y$, from $\log p_c = z_c - \log\sum_j e^{z_j}$ and $\sum_c y_c = 1$.
• "Write the backward pass of a two-layer MLP in NumPy." The four steps of the theorem in section 3, with shapes: G2 = (P - Y) / B, dW2 = A1.T @ G2, db2 = G2.sum(0), G1 = (G2 @ W2.T) * (Z1 > 0), dW1 = X.T @ G1, db1 = G1.sum(0).
• "Where does training memory go, and how do you reduce activation memory?" Weights, gradients, optimizer state and stored activations. Reduce activations with checkpointing ($\sqrt{L}$ segments, or selective recomputation), smaller micro-batches, or by splitting sequences across devices.
• "How do you know your custom backward pass is correct?" gradcheck in float64, relative error, $h \approx 10^{-6}$, smooth functions or no kinks near the test point, or a random-direction check.
• "Why does a training step cost about $6N$ FLOPs per token?" $2N$ for the forward pass (one multiply-add per weight), and two matmuls of the same size in the backward pass for every forward matmul.

10. Practice #

The notebook for this lesson runs in Colab. Its core exercises use NumPy only; PyTorch is used as a reference when it is installed.

1. Write a tiny scalar autograd engine: a Value class with +, *, relu and backward(), and check it on the graph of section 2.
2. Write mlp_backward for a 2-layer MLP with softmax cross-entropy in matrix form (no loop over examples), and check it on the worked example.
3. Write numerical_gradient (central differences) and rel_error, and use them to check exercise 2 on random data.
4. Plot the error of central and forward differences against $h$, and find the best $h$.
5. Write checkpoint_counts(L, s) that returns the forward calls and stored activations of checkpoint_sequential, and find the best $s$ for $L = 64$.

11. Watch and read #

Backpropagation, intuitively | Deep Learning Chapter 33Blue1Brown · 13 min

The idea before the formulas: how each weight's change is pushed back from the loss.

Backpropagation calculus | Deep Learning Chapter 43Blue1Brown · 10 min

The chain rule on a chain of single neurons, and how sums appear when a neuron feeds many others.

Lecture 4 | Introduction to Neural NetworksStanford CS231n (2017) · 74 min

Backprop on computational graphs with numbers at each node, and the add, multiply and max gates.

Building makemore Part 4: Becoming a Backprop NinjaAndrej Karpathy · 115 min

A full manual backward pass of an MLP in PyTorch, checked tensor by tensor against autograd.

• Simon Prince, Understanding Deep Learning (free PDF), chapter 7, §7.1–7.4: the toy example, the backpropagation algorithm and algorithmic differentiation.
• Zhang, Lipton, Li and Smola, Dive into Deep Learning, §5.3 Forward propagation, backward propagation, and computational graphs.
• Stanford CS231n notes, Neural networks part 3, section "Gradient checks".
• Chen, Xu, Zhang and Guestrin, Training deep nets with sublinear memory cost (2016): the $\sqrt{L}$ and recursive checkpointing schedules.
• Korthikanti et al., Reducing activation recomputation in large transformer models (2022): selective recomputation.