S04 Reading — Backpropagation I | Deep Learning

Module 1 · Reading 1.4 · read before S04: Backpropagation I: Computation Graphs

The PDF posted in Canvas is the assigned version of this reading. It is built from this page, so the text is identical; anything here marked as deep reading is optional and is not in the PDF.

Section 1.3 ended with a promise: give gradient descent the gradient of the loss with respect to every parameter, and it trains anything. This section pays the debt. The obstacle is scale — a modern network has millions of parameters, and computing each derivative separately (say, by nudging each parameter and re-running the network) would cost one full forward pass per parameter. Backpropagation computes them together in one reverse sweep whose cost is a small constant multiple of the forward computation. Its cost still scales with the graph and model size; the key is that it does not require one complete forward pass per parameter.

The algorithm is nothing more than the chain rule, organized well. But “organized well” is the whole game, so we build it in several passes: the chain rule on computation graphs, a full forward/backward pass by hand on a four-parameter network, the other organization of the chain rule (forward mode) so you can see what reverse mode is being compared against, and the vectorized version that scales — plus the finite-difference check that will save you hours of debugging in HW1.

How to read this

Computation graphs and the chain rule
A forward and backward pass by hand
Forward mode, in code
Vectorized backprop: whole layers at a time
Gradient checking as a habit
What frameworks automate — and what they don’t
Going deeper
Common pitfalls
Try it yourself
Further reading

Computation graphs and the chain rule

A computation graph is a bookkeeping device: break a computation into primitive operations (add, multiply, tanh, matrix multiply), make each intermediate result a node, and draw an edge from each input of an operation to its output. Any network-plus-loss becomes a graph whose final node is the scalar loss \(L\).

The payoff is that differentiation becomes local. Each operation knows its own derivative (its local derivative) with respect to each of its inputs, and nothing else. The chain rule then says: the derivative of \(L\) with respect to any node is the sum, over all paths from that node to \(L\), of the product of local derivatives along each path. For a simple chain \(x \to z \to h \to L\):

\[\frac{dL}{dx} = \frac{dL}{dh}\cdot\frac{dh}{dz}\cdot\frac{dz}{dx}\]

Two orders of traversal compute this. Forward mode propagates “how does everything change per unit change of this one input” from inputs toward the loss: one sweep per input, cheap when inputs are few. Reverse mode propagates “how does the loss change per unit change of this node” from the loss back toward the inputs: one sweep total, delivering \(dL/d(\text{node})\) for every node at once. Networks have millions of inputs (the parameters) and one output (the loss), so reverse mode wins by a factor of millions. Backpropagation is reverse-mode automatic differentiation applied to a network’s computation graph.

Historical note. Reverse-mode automatic differentiation was published by Seppo Linnainmaa in his 1970 master’s thesis in Helsinki, with no reference to neural networks; Paul Werbos described its application to training networks in his 1974 Harvard PhD thesis. Neither was noticed by the neural network community of the time. The 1986 Nature paper by Rumelhart, Hinton, and Williams independently derived the method, named the resulting learning procedure, and, decisively, demonstrated that trained hidden units develop meaningful internal representations. The lesson generalizes: in this field, ideas often “arrive” not when first derived but when someone shows them working.

Why did an algorithm this simple take until 1986 to matter, when the mathematics existed by 1970? Partly visibility: Linnainmaa’s thesis was in numerical analysis and Werbos’s in another field’s literature entirely, and the communities that needed the result did not read the communities that had it — a pattern worth remembering whenever “nobody has done X” appears in a research plan. Partly framing: as section 1.2 described, the 1969 Perceptrons critique had pushed neural networks to the margins, so few people were looking for a hidden-layer training algorithm at all. And partly demonstration: the chain rule is not persuasive; a network that learns interesting internal representations is. The 1986 paper supplied the demonstration, on problems small enough for 1986 hardware, and the field reorganized around it within a few years. It is a clean example of a recurring theme from section 1.1: the bottleneck for an idea is rarely its derivation.

The quantity \(dL/d(\text{node})\) is conventionally just called the node’s gradient during backprop, and every deep learning framework’s .grad attribute is exactly this. Two structural rules govern the backward sweep:

Chain (series): a node’s gradient is the downstream gradient times the local derivative.
Fan-out (parallel): if a node feeds several downstream operations, its gradient is the sum of contributions from each path.

The fan-out rule is the one most often overlooked, so here is the smallest case that needs it. Let \(L = x \cdot x\). As a graph, the multiply node has two inputs, and both are \(x\): the node fans out. Each path contributes the local derivative of a product with respect to one slot, holding the other fixed (\(x\) from the first slot, \(x\) from the second), and the sum gives \(dL/dx = x + x = 2x\), exactly what calculus says. Skip the sum and you get \(x\), silently half the true gradient. Real networks fan out constantly: a weight matrix used at several timesteps of a recurrent network, an embedding shared between input and output layers, a tensor feeding both a main branch and a skip connection. In every case the rule is the same (gradients from all consumers add), and every autograd engine (including the one you build in HW1) implements it as grad += rather than grad =. A surprising fraction of hand-written backward-pass bugs are a = where a += belonged.

Everything else is arithmetic.

A forward and backward pass by hand

Take the smallest network that exercises every rule: one input, one tanh hidden unit, one linear output, squared-error loss. Three parameters, w1, w2, b2 (we omit a first-layer bias to keep the trace short), and the graph

\[z = w_1 x, \qquad h = \tanh(z), \qquad \hat{y} = w_2 h + b_2, \qquad L = \tfrac12(\hat{y} - y)^2\]

The forward pass computes and stores every intermediate value; the backward pass then consumes them in reverse order. Watch the quantity dL_d* walk backward through the graph, one local derivative per step:

import numpy as np

np.random.seed(0)

# A tiny network on one example:
#   z = w1 * x, h = tanh(z), yhat = w2 * h + b2, L = 0.5 * (yhat - y)^2
x, y = 2.0, 1.0
w1, w2, b2 = 0.5, -1.5, 0.5

# ---- forward pass: compute and store every intermediate ----
z = w1 * x
h = np.tanh(z)
yhat = w2 * h + b2
L = 0.5 * (yhat - y) ** 2
print(f"forward:  z = {z:.4f}, h = {h:.4f}, yhat = {yhat:.4f}, L = {L:.4f}")

# ---- backward pass: walk the graph in reverse, one local derivative each ----
dL_dyhat = yhat - y                    # d/dyhat of 0.5*(yhat - y)^2
dL_db2 = dL_dyhat * 1.0                # yhat = w2*h + b2  ->  dyhat/db2 = 1
dL_dw2 = dL_dyhat * h                  #                      dyhat/dw2 = h
dL_dh = dL_dyhat * w2                  #                      dyhat/dh  = w2
dL_dz = dL_dh * (1 - h ** 2)           # h = tanh(z)       ->  dh/dz = 1 - tanh(z)^2
dL_dw1 = dL_dz * x                     # z = w1*x          ->  dz/dw1 = x
print(f"backward: dL/dyhat = {dL_dyhat:.4f}, dL/dw2 = {dL_dw2:.4f}, "
      f"dL/db2 = {dL_db2:.4f}, dL/dw1 = {dL_dw1:.4f}")

# ---- check dL/dw1 against a finite difference ----
def loss_at(w1_val):
    return 0.5 * (w2 * np.tanh(w1_val * x) + b2 - y) ** 2

eps = 1e-6
fd = (loss_at(w1 + eps) - loss_at(w1 - eps)) / (2 * eps)
print(f"finite difference for dL/dw1: {fd:.4f}")

Running this prints:

forward:  z = 1.0000, h = 0.7616, yhat = -0.6424, L = 1.3487
backward: dL/dyhat = -1.6424, dL/dw2 = -1.2508, dL/db2 = -1.6424, dL/dw1 = 2.0693
finite difference for dL/dw1: 2.0693

Trace dL_dw1 yourself with the numbers: \((-1.6424) \times (-1.5) \times (1 - 0.7616^2) \times 2 = 2.0693\). Every factor is either something the forward pass stored (h, yhat) or a parameter, which is why backprop’s memory cost is storing the forward intermediates, and why frameworks keep activations alive until the backward pass finishes. The final line verifies the whole chain against a finite difference: nudge w1 by 1e-6 in each direction, take the slope of the loss, and the analytic gradient matches to four decimals. Also note dL_dw1 is positive while the error yhat - y is negative: gradients tell you which way is uphill for the loss, not which way the output is wrong — gradient descent’s minus sign handles the rest.

Notice what the backward pass never did: it never formed \(dL/dw_1\) as one giant expanded expression. Naive symbolic differentiation can duplicate shared subexpressions and grow exponentially unless the expression is kept as a graph or common subexpressions are reused; numerical differentiation costs one forward pass per parameter. Backprop is the graph-organized third way (derivatives accurate up to floating-point arithmetic, at a cost comparable to a small number of passes), and this distinction is worth internalizing before HW1.

Forward mode, in code

Reverse mode’s claim to fame is a comparison, so the comparison deserves to be concrete. Forward mode has an implementation of almost absurd elegance: replace every number with a pair (value, derivative), a dual number, and teach each primitive operation to propagate both components. The derivative component carries d(anything)/d(seed), where the seed is whichever input you set the derivative to 1 for. One sweep forward, and every quantity in the program knows its sensitivity to that one seed:

import numpy as np

np.random.seed(0)

class Dual:
    """A value paired with a derivative, propagated forward through every op."""
    def __init__(self, val, dot):
        self.val, self.dot = val, dot
    def _wrap(self, o):
        return o if isinstance(o, Dual) else Dual(o, 0.0)
    def __add__(self, o):
        o = self._wrap(o); return Dual(self.val + o.val, self.dot + o.dot)
    __radd__ = __add__
    def __sub__(self, o):
        o = self._wrap(o); return Dual(self.val - o.val, self.dot - o.dot)
    def __mul__(self, o):
        o = self._wrap(o)
        return Dual(self.val * o.val, self.dot * o.val + self.val * o.dot)
    __rmul__ = __mul__
    def __pow__(self, k):
        return Dual(self.val ** k, k * self.val ** (k - 1) * self.dot)

def tanh(a):
    t = np.tanh(a.val)
    return Dual(t, (1 - t ** 2) * a.dot)

# The same tiny network as the by-hand example
x, y = 2.0, 1.0
params = {"w1": 0.5, "w2": -1.5, "b2": 0.5}

def loss(w1, w2, b2):
    z = w1 * x
    h = tanh(z)
    yhat = w2 * h + b2
    return 0.5 * (yhat - y) ** 2

# Forward mode: ONE pass per parameter, seeding that parameter's dot with 1
for name in params:
    args = {k: Dual(v, 1.0 if k == name else 0.0) for k, v in params.items()}
    L = loss(**args)
    print(f"pass seeded at {name}: L = {L.val:.4f}, dL/d{name} = {L.dot:.4f}")

Running this prints:

pass seeded at w1: L = 1.3487, dL/dw1 = 2.0693
pass seeded at w2: L = 1.3487, dL/dw2 = -1.2508
pass seeded at b2: L = 1.3487, dL/db2 = -1.6424

The three derivatives match the backward pass exactly: same chain rule, same arithmetic, opposite orientation. Each Dual multiplication applies the product rule as the computation happens, so nothing needs to be stored: no graph, no saved activations, no separate backward sweep. The price is in the loop: three parameters took three complete forward passes, each recomputing the same L = 1.3487. Scale that to a million parameters and forward mode costs a million passes where reverse mode costs one — that ratio, not any subtlety of the math, is why deep learning is built on reverse mode.

The clean way to summarize the trade uses the Jacobian, the matrix of all output-input derivatives of a function from \(n\) inputs to \(m\) outputs. One forward-mode sweep computes a Jacobian-times-vector product (a linear combination of Jacobian columns); one reverse-mode sweep computes a vector-times-Jacobian product (a combination of rows). To assemble a full Jacobian you need \(n\) forward sweeps or \(m\) reverse sweeps — so the winner is decided by shape alone. Training a network is the extreme case \(n\) in the millions with \(m = 1\), maximally favorable to reverse mode. But forward mode is not a museum piece: with \(n\) tiny it can win because it carries only the requested tangent information and needs no reverse tape or separate backward sweep. It still stores the ordinary values and live tangents required by the forward computation, so it is not literally zero-memory. Modern frameworks expose both modes (PyTorch ships torch.func.jvp alongside reverse mode), and compositions such as forward-over-reverse compute Hessian-vector products.

Vectorized backprop: whole layers at a time

Scalar-by-scalar backprop would take forever on real networks. Fortunately the chain rule respects matrix structure, and each layer becomes one node with a matrix-level local derivative. For a linear layer \(Z = XW + b\) (batch \(X\) of shape (N, D), weights \(W\) of shape (D, H)) with incoming gradient \(dZ = \partial L/\partial Z\) of shape (N, H), the three backward rules are:

\(dW = X^\top dZ\): shape (D, H), same as \(W\)

\(dX = dZ\, W^\top\): shape (N, D), same as \(X\)

\(db_j = \sum_n dZ_{nj}\) (sum over batch rows): shape (H,), same as \(b\)

The derivation of the first rule is worth doing once in your life, so here is its skeleton. Entry \(W_{ij}\) touches the forward pass in exactly one place per example: it multiplies feature \(X_{ni}\) and contributes to pre-activation \(Z_{nj}\). So by the chain rule, \(\partial L/\partial W_{ij}\) sums over everything downstream that \(W_{ij}\) influenced (all \(N\) examples’ \(Z_{nj}\)), giving \(\sum_n X_{ni}\, dZ_{nj}\). Stare at that sum: it is the \((i, j)\) entry of \(X^\top dZ\). The matrix form is not a new fact; it is \(N \cdot D \cdot H\) scalar chain rules with the summation packed into a matrix multiply, which is precisely why backprop runs at matrix-multiply speed on hardware built for matrix multiplies.

You can derive the other rules the same way, but you can also remember them by shape discipline: a parameter’s gradient always has the same shape as the parameter, and \(X^\top dZ\) is the only way to combine (N, D) and (N, H) into (D, H). The db rule is the fan-out rule in disguise — b was broadcast across N rows in the forward pass, so its gradient sums the N per-row contributions. Broadcast forward, sum backward; forget the sum and your bias gradient has the wrong shape or, worse, silently broadcasts into wrongness.

Here is a complete two-layer ReLU classifier — forward, backward, and the check that certifies it. The first backward line uses the \(\operatorname{softmax} - \operatorname{onehot}\) gradient derived in section 1.3; the ReLU backward is an elementwise gate, passing gradient where the pre-activation was positive and blocking it elsewhere:

import numpy as np

np.random.seed(0)

# Batch of 8 examples, 3 features, 4 hidden ReLU units, 2 output classes
N, D, H, C = 8, 3, 4, 2
X = np.random.randn(N, D)
y = np.random.randint(0, C, size=N)

params = {
    "W1": 0.5 * np.random.randn(D, H), "b1": np.zeros(H),
    "W2": 0.5 * np.random.randn(H, C), "b2": np.zeros(C),
}

def forward_backward(params, X, y):
    W1, b1, W2, b2 = params["W1"], params["b1"], params["W2"], params["b2"]
    N = X.shape[0]

    # forward
    Z1 = X @ W1 + b1                         # (N, H) pre-activations
    A1 = np.maximum(0, Z1)                   # (N, H) ReLU
    Z2 = A1 @ W2 + b2                        # (N, C) logits
    Z2s = Z2 - Z2.max(axis=1, keepdims=True)             # stable log-softmax
    logp = Z2s - np.log(np.exp(Z2s).sum(axis=1, keepdims=True))
    loss = -logp[np.arange(N), y].mean()

    # backward — same graph, reverse order
    dZ2 = np.exp(logp)                       # softmax probabilities, (N, C)
    dZ2[np.arange(N), y] -= 1                # subtract 1 at the true class
    dZ2 /= N                                 # mean over the batch
    grads = {
        "W2": A1.T @ dZ2, "b2": dZ2.sum(axis=0),
    }
    dA1 = dZ2 @ W2.T                         # (N, H)
    dZ1 = dA1 * (Z1 > 0)                     # ReLU gate: pass or block
    grads["W1"] = X.T @ dZ1
    grads["b1"] = dZ1.sum(axis=0)
    return loss, grads

loss, grads = forward_backward(params, X, y)
print(f"loss = {loss:.4f}")

# Gradient check: compare every analytic gradient to a central finite difference
eps = 1e-5
for name, p in params.items():
    num = np.zeros_like(p)
    for idx in np.ndindex(p.shape):
        old = p[idx]
        p[idx] = old + eps; lp, _ = forward_backward(params, X, y)
        p[idx] = old - eps; lm, _ = forward_backward(params, X, y)
        p[idx] = old
        num[idx] = (lp - lm) / (2 * eps)
    rel = np.abs(num - grads[name]) / (np.abs(num) + np.abs(grads[name]) + 1e-12)
    print(f"{name}: max relative error vs finite difference = {rel.max():.2e}")

Running this prints:

loss = 1.0065
W1: max relative error vs finite difference = 4.32e-10
b1: max relative error vs finite difference = 5.78e-10
W2: max relative error vs finite difference = 7.99e-11
b2: max relative error vs finite difference = 6.01e-12

Relative errors around 1e-10 are the signature of a correct implementation checked with central differences at eps = 1e-5: agreement down to floating-point noise. A genuinely wrong gradient (a missing transpose, a dropped ReLU mask, a forgotten 1/N) shows up as relative error near 1e0, orders of magnitude away; there is no ambiguity. As a working rule: below 1e-6, trust the gradient; above 1e-2, it is wrong; in between, suspect numerical issues (an unlucky eps, float32 instead of float64, or a ReLU kink; see the pitfall below) before suspecting the math.

Gradient checking as a habit

The finite-difference check deserves its own habits, because it is the single highest-value debugging tool in this course:

Use central differences. (L(p + eps) - L(p - eps)) / (2 * eps) has error of order eps^2; the one-sided version has error of order eps and will muddy your verdicts. With eps around 1e-5 and float64, expect agreement to ~1e-8 or better.
Use relative error, |num - ana| / (|num| + |ana| + tiny) — absolute error is meaningless across gradients whose scales differ by orders of magnitude.
Check small. The check costs two forward passes per parameter entry (the very cost backprop exists to avoid), so run it on toy shapes (like our 3 -> 4 -> 2 net), never on the real model. Correctness established at small scale transfers; the code is the same.
Check after every change, not once at the end. A gradient check wired into your tests (HW1’s suite does exactly this) converts silent training bugs into loud red failures.

The float64 advice in the working rule above is not superstition, and it is worth one experiment to see why. A central difference computes L(p + eps) - L(p - eps), a subtraction of two nearly equal numbers, and dividing their difference by 2 * eps amplifies whatever rounding error the subtraction left behind. The smaller eps gets, the truer the formula but the noisier the arithmetic: truncation error falls as eps^2 while cancellation error grows as (machine precision)/eps. Where those curves cross depends entirely on how many digits your floats carry:

import numpy as np

np.random.seed(0)

def loss_and_grad(w, X, y):
    """Logistic regression loss and analytic gradient, in w's dtype."""
    p = 1 / (1 + np.exp(-(X @ w)))
    tiny = np.finfo(w.dtype).tiny
    L = -np.mean(y * np.log(p + tiny) + (1 - y) * np.log(1 - p + tiny))
    g = X.T @ (p - y) / len(y)
    return L, g

for dtype in [np.float64, np.float32]:
    X = np.random.default_rng(0).standard_normal((64, 10)).astype(dtype)
    w = (0.1 * np.random.default_rng(1).standard_normal(10)).astype(dtype)
    yv = (np.random.default_rng(2).uniform(size=64) > 0.5).astype(dtype)
    _, g = loss_and_grad(w, X, yv)
    for eps in [1e-2, 1e-4, 1e-6]:
        num = np.zeros_like(g)
        for i in range(len(w)):
            wp = w.copy(); wp[i] += dtype(eps)
            wm = w.copy(); wm[i] -= dtype(eps)
            num[i] = (loss_and_grad(wp, X, yv)[0] - loss_and_grad(wm, X, yv)[0]) / (2 * eps)
        rel = np.max(np.abs(num - g) / (np.abs(num) + np.abs(g) + 1e-12))
        print(f"{np.dtype(dtype).name}, eps = {eps:.0e}: max relative error = {rel:.2e}")

Running this prints:

float64, eps = 1e-02: max relative error = 3.10e-05
float64, eps = 1e-04: max relative error = 3.10e-09
float64, eps = 1e-06: max relative error = 1.77e-09
float32, eps = 1e-02: max relative error = 7.75e-05
float32, eps = 1e-04: max relative error = 9.75e-03
float32, eps = 1e-06: max relative error = 1.00e+00

In float64, shrinking eps helps and then plateaus around 1e-9, comfortably inside “trust the gradient” territory. In float32 the story inverts: eps = 1e-6 produces relative error of 1.00e+00 — the check has failed completely on a correct gradient, because L(p + eps) and L(p - eps) are identical to all seven of float32’s decimal digits, and the difference is pure rounding noise. Even the best float32 row sits near 1e-4, in the ambiguous zone where you cannot distinguish “correct but noisy” from “subtly wrong.” This is the origin of the working rule: gradient checks run in float64 or they run in the dark. The same cancellation arithmetic, incidentally, is why training in low precision requires care with small gradient values, a thread Going deeper picks up.

Pitfall. ReLU can fool the gradient check itself. If some pre-activation sits within eps of zero, the finite-difference probe steps across the kink: the two evaluations straddle different linear pieces, the numerical estimate is garbage there, and a correct implementation reports a large error in that one entry. With random inputs this is rare (our run shows no sign of it), but when a check fails in isolated entries of a ReLU layer, inspect the pre-activations before tearing your math apart — or rerun the check with a tanh activation, which has no kink. The same non-differentiability is harmlessly resolved in training by picking either subgradient at exactly zero; it matters only to the checker’s probes.

What frameworks automate — and what they don’t

Everything in this section is mechanical: record operations into a graph during the forward pass, then sweep backward applying each operation’s stored local-derivative rule. Mechanical means automatable, and that automation is called autograd — PyTorch’s defining feature and the subject of section 1.5, where you build a working autograd engine in about a hundred lines and then reimplement it as HW1. After this course you will rarely write a backward pass by hand again.

How the graph gets built was, for a few years, the field’s liveliest engineering argument. Early frameworks (TensorFlow 1.x, Theano) were static-graph: you declared the entire computation up front as a data structure, the framework compiled and optimized it, and then you pushed batches through — efficient, but debugging meant staring at a compiled artifact, and control flow (loops, branches that depend on data) was painful to express. PyTorch popularized the define-by-run alternative: the graph is simply a byproduct of executing ordinary Python, rebuilt fresh on every forward pass, so a print statement or a data-dependent if works exactly as it would anywhere else. Researchers voted decisively with their feet for define-by-run, and TensorFlow 2 switched its default to match. The trade never disappeared, though (a graph known in advance can be optimized in ways a graph discovered on the fly cannot), and modern PyTorch’s torch.compile recovers static-style optimization by tracing the dynamic graph and compiling what it sees. You will feel this history directly in HW1: your micro-autograd engine will be define-by-run, because that is the version a hundred lines can express clearly.

People. The tradition of teaching backprop by building a tiny engine owes much to Andrej Karpathy, a founding member of OpenAI and formerly director of AI at Tesla, whose open-source micrograd (2020), a scalar autograd engine of roughly one hundred lines, and the accompanying lecture video walking through it line by line have become a standard on-ramp to the field. HW1’s micro-autograd assignment is in this lineage: the fastest way to trust an automated backward pass is to have written one yourself.

What the framework does not automate is everything this section taught you to see: that memory scales with stored activations (why big models hit out-of-memory in the backward pass), that gradients flow only through the graph you actually built (why an accidental detach or an integer cast silently stops learning), and that saturated activations multiply small local derivatives into vanishing gradients (section 1.2). When training misbehaves, you debug it with this section’s mental model, not the framework’s source code.

Going deeper

Backprop’s compute cost gets all the attention, but its memory cost often limits the models you can train. Reverse mode normally retains forward intermediates until their backward rules consume them, so activation memory grows with batch or sequence size, width, and depth. Activations can dominate for large batches or long sequences; weights, gradients, and optimizer state can dominate in other regimes. Forward mode avoids a reverse tape but still carries primal values and tangents through the live forward computation.

The classic escape is gradient checkpointing: store only every \(k\)-th activation during the forward pass, and when the backward sweep needs a segment you did not keep, recompute it from the nearest checkpoint. You pay a second forward pass for a segment; you save the memory of everything between checkpoints. Here it is on a 16-layer chain, with the exactness of the result checked directly:

import numpy as np

np.random.seed(0)

# A chain of 16 tanh layers; loss = mean of the final output.
L_layers, width = 16, 32
Ws = [np.random.randn(width, width) / np.sqrt(width) for _ in range(L_layers)]
x0 = np.random.randn(8, width)

def backward_full(x):
    """Standard backprop: store every activation, then sweep backward."""
    acts = [x]
    for W in Ws:
        acts.append(np.tanh(acts[-1] @ W))
    d = np.ones_like(acts[-1]) / acts[-1].size          # dL/d(output) for L = mean
    grads = [None] * L_layers
    for i in reversed(range(L_layers)):
        dz = d * (1 - acts[i + 1] ** 2)                 # through tanh
        grads[i] = acts[i].T @ dz
        d = dz @ Ws[i].T
    return grads, len(acts)

def backward_checkpointed(x, every=4):
    """Keep every 4th activation; recompute each segment during the backward pass."""
    ckpts, h = {0: x}, x
    for i, W in enumerate(Ws):
        h = np.tanh(h @ W)
        if (i + 1) % every == 0:
            ckpts[i + 1] = h
    d = np.ones_like(h) / h.size
    grads = [None] * L_layers
    for seg_end in range(L_layers, 0, -every):
        seg_start = seg_end - every
        seg = [ckpts[seg_start]]                        # recompute this segment only
        for i in range(seg_start, seg_end):
            seg.append(np.tanh(seg[-1] @ Ws[i]))
        for i in reversed(range(seg_start, seg_end)):
            dz = d * (1 - seg[i - seg_start + 1] ** 2)
            grads[i] = seg[i - seg_start].T @ dz
            d = dz @ Ws[i].T
    return grads, len(ckpts) + every + 1                # checkpoints + one live segment

g_full, mem_full = backward_full(x0)
g_ck, mem_ck = backward_checkpointed(x0)
diff = max(np.max(np.abs(a - b)) for a, b in zip(g_full, g_ck))
print(f"activations stored, standard backprop:  {mem_full}")
print(f"activations stored, checkpointed:       {mem_ck} (5 checkpoints + 1 segment of 4 + input)")
print(f"max |gradient difference|:              {diff:.2e}")

Running this prints:

activations stored, standard backprop:  17
activations stored, checkpointed:       10 (5 checkpoints + 1 segment of 4 + input)
max |gradient difference|:              0.00e+00

In this deterministic NumPy example the maximum displayed difference is zero, and the gradients are bitwise identical on the tested run. More generally, checkpointed and ordinary execution are mathematically equivalent, but bitwise identity requires the same kernels, floating-point order, RNG state, and other stateful behavior. Framework checkpoint utilities preserve common RNG state; custom randomness or mutable state still deserves a test. The bookkeeping here (10 stored versus 17) is modest because the chain is short, but the asymptotics are what matter: checkpointing every \(\sqrt{L}\) layers stores only \(O(\sqrt{L})\) activations at the price of roughly one extra forward pass.

The second economics story is precision. Our gradient-check experiment showed float32 arithmetic destroying a finite-difference estimate; training itself faces a version of the same hazard. Modern accelerators run half-precision (16-bit) arithmetic several times faster than 32-bit, so production training wants 16-bit everywhere — but FP16 can represent nothing smaller than about 6e-8, and real gradient distributions have a long tail of values below that, which silently flush to zero. The now-standard recipe, mixed-precision training (Micikevicius et al., ICLR 2018), threads the needle with three moves: compute forward and backward in FP16; keep a master copy of weights in FP32 so millions of tiny updates are not individually rounded away; and multiply the loss by a scale factor before the backward pass (dividing it back out at the update), shifting the whole gradient distribution up into FP16’s representable range. That last trick, loss scaling, is pure exploitation of this section’s structure: gradients are linear in the incoming \(dL/dL = 1\), so scaling the loss scales every gradient in the graph by the same known constant. PyTorch automates the recipe as torch.amp; the newer bfloat16 format, with FP32’s exponent range and less mantissa, often removes the need for loss scaling entirely. Either way, the numerical texture of your gradients is a design surface, not a given.

A third lever follows from linearity: gradient accumulation. A sum or correctly weighted mean of per-example gradients can be computed in installments — run four forward/backward passes on quarter-sized micro-batches, accumulate into .grad, and step once. This matches the full-batch gradient when loss scaling handles unequal or final partial micro-batches correctly and the model has no batch-dependent behavior that changes between the full batch and micro-batches (ordinary BatchNorm is the classic counterexample). Floating-point summation order can also create tiny numerical differences. Subject to those conditions, peak activation memory falls because only one micro-batch’s activations are alive at a time.

All three stories generalize into one habit: read your training loop as systems code. The backward pass is twice the forward FLOPs, activations are the memory bill, recomputation and precision are the levers, and every one of those facts follows from the mechanics this section built by hand. Section 4.4 returns to these levers at datacenter scale.

Common pitfalls

Symptom: shapes explode or a dimension error appears in the backward pass only. Cause: a transposed matrix in \(dW = X^\top dZ\) or \(dX = dZ\, W^\top\). Fix: apply shape discipline — write every gradient’s required shape next to it and check that each product is the only one producing that shape.
Symptom: bias gradient has shape (N, H) instead of (H,), or training diverges only when batch size changes. Cause: missing sum over the batch axis for a broadcast parameter. Fix: whatever was broadcast in the forward pass gets summed over those axes in the backward pass, without exception.
Symptom: loss decreases but much slower than expected, gradients all correct except scaled. Cause: dropped or doubled 1/N — the mean in the loss must appear exactly once in the backward pass. Fix: put the 1/N in one place (we fold it into dZ2) and gradient-check, which catches constant-factor errors immediately.
Symptom: gradient check reports ~1e-4 relative error everywhere, not 1e-10 and not 1e0. Cause: usually float32 arithmetic or a too-large/too-small eps, not wrong math — exactly the float32 rows of our precision experiment. Fix: run checks in float64 with central differences and eps near 1e-5; only then judge.
Symptom: out-of-memory during training, but inference on the same batch size works fine. Cause: the backward pass’s stored activations — inference discards intermediates, training cannot. Fix: smaller batches, gradient checkpointing, or mixed precision; and know that the memory bill scales with depth times batch size times width.
Symptom: some parameters never change during training. Cause: gradient genuinely zero (dead ReLU path) or the parameter disconnected from the graph (created after the forward pass, or gradient never applied). Fix: print gradient norms per parameter; a hard zero for a whole tensor means a graph problem, not an optimization problem.

Try it yourself

Extend the by-hand example with a first-layer bias b1 (so z = w1 * x + b1). Derive dL/db1 on paper, add it to the script, and verify it with a finite difference.
Add a sin operation to the Dual class (values sin(v), derivative cos(v) * dot), change the tiny network’s activation from tanh to sin, and verify all three forward-mode gradients against finite differences.
In the vectorized example, deliberately break the backward pass three ways (remove the ReLU mask (Z1 > 0), drop the /= N, transpose W2.T to W2) and record the gradient check’s report for each. Learn what each class of bug “looks like.”
Add a second hidden layer (shapes 3 -> 4 -> 4 -> 2) to the vectorized network and make the gradient check pass again. How many lines did the backward pass grow by, and which lines are pure repetition?
In the checkpointing example, vary every over 1, 2, 4, 8, 16 and tabulate stored activations against the number of extra forward-layer computations. Which setting minimizes total storage for this 16-layer chain, and how does that compare to \(\sqrt{16}\)?
Measure the cost claim: time 1000 forward passes and 1000 forward+backward passes of the vectorized network (increase sizes to N, D, H, C = 256, 100, 100, 10 so timings are stable). Is forward+backward within a small factor of twice the forward-only time?

Before class you should be able to

Walk a small computation graph backward, one local derivative per step, and verify the result with a central finite difference.
Apply the fan-out rule (gradients from multiple consumers add) and explain what goes wrong with grad = instead of grad +=.
Explain why reverse mode beats forward mode for network training, in terms of inputs versus outputs.
Write the three vectorized backward rules for a linear layer and justify each by shape discipline.
State the “broadcast forward, sum backward” rule for bias gradients.
Interpret a gradient check’s output: ~1e-10 means correct, ~1e0 means wrong, in between means suspect the numerics first.
Explain why training runs out of memory in the backward pass when inference on the same batch does not.

Connects to

This session is covered by Quiz 1 (Module 1: S01–S08), given at S09.
HW1 (micro-autograd) goes out today, due at S08 — its test suite is largely gradient checks of the kind shown here.
Continue with 1.5 (the S05 reading) before starting HW1.

← S03 · Losses and Gradient Descent · All readings · S05 · Backpropagation II →

How to read this

Computation graphs and the chain rule

A forward and backward pass by hand

Forward mode, in code

Vectorized backprop: whole layers at a time

Gradient checking as a habit

What frameworks automate — and what they don’t

Going deeper

Common pitfalls

Try it yourself

Further reading

Before class you should be able to

Connects to