Module 1 · Reading 1.5 · read before S05: Backpropagation II

When you call loss.backward() in PyTorch, autograd computes gradients for the parameters that require them and contributed to that loss. Nothing about that line looks like the chain-rule bookkeeping you did by hand in backpropagation, and treating it as magic makes debugging harder: you will eventually hit a silent zero gradient, a doubled gradient, or an exploding memory footprint, and you will need a model of what went wrong.

The cure is to build the machine yourself once. The core of a reverse-mode automatic differentiation engine (the thing PyTorch, JAX, and TensorFlow all contain) fits in about sixty lines of Python. This section builds one, step by step, and then checks it against torch.autograd on the same expression. HW1 (“micro-autograd”) asks you to build a more complete version of exactly this engine, so read this section with an editor open.

How to read this

  • Vectorized backprop and gradient checking, in one page — the first half of class works from here. It is a self-contained recap of the two things section 1.4 ended on: the three matrix rules for a linear layer, and the finite-difference check that tests them. The gradient check you watched in S04; the three matrix rules you have so far only read, since S04 stopped before them and S05 derives them step by step. Read this section properly and run the example; everything after it assumes these rules.

  • What a node must remember — the four core fields every node stores (value, gradient, parents, local backward rule) and reverse mode’s tradeoff: all parameter gradients in one sweep, with forward values saved for backward.
  • A working engine in sixty lines — the worked example to actually run. Slow down on the three details that repay close reading: closures carry the local state, gradients accumulate with +=, and wrappers make constants work.
  • Why topological order matters — the reuse graph a * a + a: no node’s rule may fire until every consumer has deposited its contribution.
  • Checking against torch.autograd — the habit of testing an engine against a reference on the same expression. HW1’s tests do the same job with known values and one central finite difference, needing no PyTorch; the torch comparison in that suite is optional and unscored. Note why agreement is ~1e-8, not exact (float32 versus float64).
  • Where calculus runs out: kinks and infinities — the relu'(0) = 0 convention (PyTorch’s, and HW1’s starter’s) and the where-trap: inf * 0 = nan can contaminate upstream gradients.
  • The graph’s lifecycle: retention, accumulation, detachment — when to clear gradients, and what detach() and no_grad() are for. Differentiating the derivative and From sixty lines to a real framework are safe to skim on first pass.

Four of the six runnable examples below import PyTorch. If you have not installed it yet, read their printed output: the course installs PyTorch at S08, and nothing in HW1 needs it. The two examples to run yourself are the layer check, which needs Python and NumPy, and the sixty-line engine, which uses only Python’s standard library.

Vectorized backprop and gradient checking, in one page

Running the examples. Run cells in order in a fresh kernel. The NumPy layer check and the first Value engine need no PyTorch; the other examples use it. Printed results are reference runs, not a promise of identical digits across PyTorch versions, operating systems, or hardware. The deliberately caught errors and the inf/nan values in the kink example are expected observations.

Everything below builds an engine that automates the chain rule. Before automating something, you should be able to do it by hand at the scale the engine will actually work at: whole layers, not single scalars. That is where section 1.4 ended, and this section restates it so that this page stands on its own.

Scalar-by-scalar backprop would take forever on a real network. The chain rule respects matrix structure, though, so an entire layer can be one node with a matrix-level backward rule. For a linear layer Z = X @ W + b (batch X of shape (N, D), weights W of shape (D, H), incoming gradient dZ of shape (N, H)) there are exactly three rules:

rule expression shape why
weights dW = X.T @ dZ (D, H) — same as W W[i,j] multiplies X[n,i] into Z[n,j] for every one of the N examples, so its gradient sums X[n,i] * dZ[n,j] over n, which is the (i,j) entry of X.T @ dZ
inputs dX = dZ @ W.T (N, D) — same as X the same argument run along the other index; needed to keep the sweep going into earlier layers
bias db = dZ.sum(axis=0) (H,) — same as b b was broadcast across N rows going forward, which is fan-out, so its gradient sums the N contributions

Two habits catch many mistakes. First, shape discipline: a parameter’s gradient always has the parameter’s own shape, and X.T @ dZ combines (D, N) with (N, H) to produce (D, H). Matching shapes is necessary but does not prove the formula is correct. Second, broadcast forward, sum backward: the db rule sums the contributions from every use of the bias, and forgetting the sum is the classic bias-gradient bug.

Here is the whole thing on a two-layer ReLU classifier: forward, backward, the shape check, and the finite-difference check that tests the backward rules. The softmax - onehot line for dZ2 comes from section 1.3; the ReLU backward is an elementwise gate that passes gradient where the pre-activation was positive and blocks it elsewhere. 

import numpy as np

np.random.seed(0)

# The same two-layer ReLU classifier as section 1.4: 8 examples, 3 features,
# 4 hidden ReLU units, 2 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: backward reuses X, Z1 (or its mask), A1, W2, and logp
    Z1 = X @ W1 + b1                                     # (N, H)
    A1 = np.maximum(0, Z1)                               # (N, H)
    Z2 = A1 @ W2 + b2                                    # (N, C)
    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, three rules per linear layer
    dZ2 = np.exp(logp)                                   # softmax probabilities
    dZ2[np.arange(N), y] -= 1                            # the softmax - onehot rule
    dZ2 /= N                                             # the mean's 1/N, exactly once
    grads = {"W2": A1.T @ dZ2, "b2": dZ2.sum(axis=0)}
    dA1 = dZ2 @ W2.T                                     # dX rule, second layer
    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}\n")

# A parameter's gradient always has the parameter's shape. Check it, every time.
print(f"{'param':<7}{'shape':<10}{'gradient shape':<16}{'max |rel err| vs finite diff'}")
eps = 1e-5
for name, p in params.items():
    assert grads[name].shape == p.shape
    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:<7}{str(p.shape):<10}{str(grads[name].shape):<16}{rel.max():.2e}")

Running this prints:

loss = 1.0065

param  shape     gradient shape  max |rel err| vs finite diff
W1     (3, 4)    (3, 4)          4.32e-10
b1     (4,)      (4,)            5.78e-10
W2     (4, 2)    (4, 2)          7.99e-11
b2     (2,)      (2,)            6.01e-12

The check is a central difference, (L(p + eps) - L(p - eps)) / (2 * eps). At a smooth point its truncation error is \(O(\varepsilon^2)\); make eps too small and floating-point cancellation takes over instead. The comparison uses relative error, so a discrepancy is judged against the size of the gradients. The tiny errors printed here mean the backward rules agree with the finite difference on this test; a test is evidence, not a proof. The working rule for a float64 check away from kinks:

  • below 1e-6: agreement, on the entries tested;
  • above 1e-2: look for a bug, in the code or in the check itself;
  • in between: vary eps and look at the discrepancy before deciding.

Two things to check when a relative error is large. When both gradients are near zero, look at the absolute difference too: a tiny absolute gap reads as a large relative one. And check whether the perturbation flips the sign of a ReLU pre-activation: that crosses a kink, where the central-difference argument no longer applies, so shrink eps or test at an input away from the kink. Not every bug shows up as an error near 1e0; a missing transpose usually fails the shape check first.

Run it in float64 and on toy shapes. Float32 does not have the digits: at eps = 1e-6 a correct gradient can report a relative error of 1.00e+00, because L(p + eps) and L(p - eps) are the same float32 number. And the check costs two full forward passes per parameter entry, so it is a debugging instrument, never part of training. That is the state of the art you can reach by hand. It works, and it does not scale: every new architecture means deriving and coding new backward rules, and every one of them is a chance to drop a transpose. The rest of this section builds the machine that does it for you.

What a node must remember

Reverse-mode automatic differentiation rests on one organizing idea: while the forward pass computes values, it also records how each value was computed. Every intermediate result becomes a node in a directed acyclic graph, and every node stores four core pieces of information:

  1. its value — the number computed in the forward pass;
  2. its parents — the nodes it was computed from (we will call them _prev; the constructor parameter that fills them is named _children, micrograd’s word for the same nodes. Both words are in use, in opposite directions: a data-flow graph calls a node’s inputs its parents, an expression tree calls them its children. Recognize the nodes, not the name);
  3. its accumulated gradient (grad) — initially zero, then filled with contributions from its consumers during backward;
  4. a local backward rule — a small function that, given dL/d(this node), adds the correct contribution to each parent’s gradient using that operation’s local derivative.

The backward pass then visits nodes in reverse order, calling each node’s rule. Each rule is one application of the chain rule: for out = a * b, the rule adds b * out.grad to a.grad and a * out.grad to b.grad, because d(a*b)/da = b and d(a*b)/db = a. No rule needs to know anything about the rest of the graph. That locality is why the whole engine stays tiny no matter how complicated the expressions get.

Why reverse mode rather than forward mode? Forward mode pushes derivatives from inputs toward outputs; its cost scales with the number of inputs you differentiate with respect to. Reverse mode pulls derivatives from one output back toward all inputs; its cost scales with the number of outputs. Training a neural network means differentiating one scalar loss with respect to millions of parameters (one output, many inputs), so reverse mode computes every parameter gradient in a single backward sweep that costs a small constant multiple of the forward pass.

For a million parameters, one-sided numerical differentiation needs about a million additional forward passes, or two million for central differences. Both carry truncation error as well as floating-point error. Automatic differentiation applies analytic local rules, avoiding finite-difference truncation error, but both forward and reverse mode still incur floating-point roundoff. Forward mode needs one sweep per input direction to assemble the full gradient; reverse mode gets it in one backward sweep. Forward plus backward usually costs a small multiple of the forward pass alone, with the multiplier depending on the operations. The memory cost comes from saving the forward values needed by those backward rules, unless we choose to recompute them.

Because the graph is recorded while ordinary Python code runs, this style is called define-by-run (or “taping”). You write y = (a * b + c).tanh() and the graph materializes as a side effect. Earlier frameworks such as Theano (developed at the Université de Montréal in the late 2000s) instead made you define a static graph up front and then compile it; Chainer (2015) demonstrated the define-by-run alternative, and PyTorch (released in 2017) made it the mainstream default. The tradeoff: define-by-run is easier to debug (the graph is your program), at the price of rebuilding the record every iteration.

Historical note. Reverse-mode automatic differentiation is older than its fame. It is commonly credited to Seppo Linnainmaa’s 1970 master’s thesis at the University of Helsinki, and Paul Werbos’s 1974 Harvard dissertation discussed applying such gradient methods to network-like models. The idea was rediscovered several times before Rumelhart, Hinton, and Williams’s 1986 Nature paper demonstrated that “backpropagation” (the same reverse-mode computation, specialized to neural networks) could learn useful internal representations, which is when the neural-network community took notice.

A working engine in sixty lines

Here is the whole engine, plus a first test: a single tanh neuron y = tanh(w1*x1 + w2*x2 + b), and a second expression that reuses a variable. The class is deliberately minimal: three operations are enough to expose every important design decision. 

import math

class Value:
    """A scalar that remembers how it was computed, so gradients can flow back."""

    def __init__(self, data, _children=(), _op=""):
        self.data = float(data)   # the scalar value (forward pass)
        self.grad = 0.0           # dL/d(this node), filled in by backward()
        self._backward = lambda: None   # how to push out.grad into the parents
        self._prev = set(_children)     # its parents: the nodes this one was computed from
        self._op = _op            # label for debugging ("+", "*", "tanh")

    def __add__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data + other.data, (self, other), "+")

        def _backward():
            self.grad += out.grad      # d(a+b)/da = 1
            other.grad += out.grad     # d(a+b)/db = 1
        out._backward = _backward
        return out

    def __mul__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data * other.data, (self, other), "*")

        def _backward():
            self.grad += other.data * out.grad   # d(a*b)/da = b
            other.grad += self.data * out.grad   # d(a*b)/db = a
        out._backward = _backward
        return out

    def tanh(self):
        t = math.tanh(self.data)
        out = Value(t, (self,), "tanh")

        def _backward():
            self.grad += (1.0 - t * t) * out.grad   # d tanh(x)/dx = 1 - tanh(x)^2
        out._backward = _backward
        return out

    def backward(self):
        # 1. Topological sort: dependencies/operands before their result.
        topo, visited = [], set()

        def build(v):
            if v not in visited:
                visited.add(v)
                for parent in v._prev:
                    build(parent)
                topo.append(v)
        build(self)

        # 2. The base case: set the output gradient to 1, then walk the graph in reverse.
        self.grad = 1.0
        for node in reversed(topo):
            node._backward()

    __radd__ = __add__
    __rmul__ = __mul__


# A single tanh neuron: y = tanh(w1*x1 + w2*x2 + b)
x1, x2 = Value(1.5), Value(-2.0)
w1, w2 = Value(0.6), Value(0.4)
b = Value(0.1)

s = w1 * x1 + w2 * x2 + b
y = s.tanh()
y.backward()

print(f"y      = {y.data:.4f}")
print(f"dy/dw1 = {w1.grad:.4f}   (should be x1 * (1 - y^2) = {1.5 * (1 - y.data**2):.4f})")
print(f"dy/dw2 = {w2.grad:.4f}")
print(f"dy/db  = {b.grad:.4f}")
print(f"dy/dx1 = {x1.grad:.4f}")

# One value in three operand slots: gradients must ACCUMULATE.
a = Value(3.0)
f = a * a + a          # f = a^2 + a, so df/da = 2a + 1 = 7
f.backward()
print(f"f      = {f.data:.4f}")
print(f"df/da  = {a.grad:.4f}   (should be 2a + 1 = 7)")

Running this prints:

y      = 0.1974
dy/dw1 = 1.4416   (should be x1 * (1 - y^2) = 1.4416)
dy/dw2 = -1.9221
dy/db  = 0.9610
dy/dx1 = 0.5766
f      = 12.0000
df/da  = 7.0000   (should be 2a + 1 = 7)

Three details in this code repay close reading. The extra _op field is a debugging label, not one of the four fields required for differentiation.

Closures carry the local state. Each operation defines a fresh _backward function that closes over its operands (self and other, or the cached output t for tanh). A closure keeps those references available after the operation returns. Multiplication reads the operands’ .data when backward runs, so those values must remain unchanged between forward and backward; this engine does not copy or protect them. The graph’s “tape” is a web of Python closures.

Gradients accumulate with +=. In f = a * a + a, the node a is used in three operand positions across two consuming operations: twice by *, once by +. Each use contributes to \(\partial f/\partial a\), and the pieces must be summed. The printed 7.0000 (which is \(2a + 1\) at \(a = 3\)) confirms all three contributions arrived. Although _prev is a set and visits a only once, the multiplication closure still adds both operand contributions.

Wrappers make constants work. __radd__ and __rmul__ plus the isinstance check let you write 2 * x + 1 without manually wrapping the literals in Value. HW1 extends this pattern to subtraction, division, and powers.

People. The design here follows micrograd, an autograd engine whose core fits in roughly one hundred lines, that Andrej Karpathy, a founding member of OpenAI and formerly director of AI at Tesla, published on GitHub in 2020 and later built live on camera in his “Neural Networks: Zero to Hero” lecture series. Its point, and ours, is pedagogical: once you have written _backward closures and a topological sort yourself, no framework’s autograd will feel like magic again.

Why topological order matters

backward() cannot just call _backward on nodes in any order. A node’s rule reads its own grad and pushes contributions into its parents, so its rule must not fire until every consumer of that node has already deposited its contribution. Otherwise part of the gradient gets propagated and the rest is stranded.

The fix is a topological sort of the graph: an ordering in which every result node appears after its dependencies (its input operands). The recursive build function produces exactly this — depth-first, appending a node only after all of its inputs have been appended. Processing reversed(topo) then guarantees the invariant: when we call node._backward(), node.grad is final.

Consider the reuse graph in f = a * a + a. The product node p = a * a and the sum node f both touch a. In reverse topological order we process f first (depositing 1.0 into p.grad and 1.0 into a.grad), then p (depositing a * p.grad = 3.0 twice into a.grad). Total: 1 + 3 + 3 = 7. Process p before f and you would propagate p.grad while it is still 0.0.

The other half of the invariant is the base case of the chain rule: self.grad = 1.0 says dL/dL = 1. PyTorch does the same thing when you call .backward() on a scalar; for a non-scalar output it makes you supply that starting gradient explicitly.

Pitfall. The single most common bug in a hand-written engine is writing self.grad = ... instead of self.grad += ... inside _backward. Every test where each variable appears once will pass, because with a single consumer overwrite and accumulate coincide. The bug surfaces when a value is reused, as happens routinely with shared weights, residual branches, recurrent state, and tensors consumed by multiple later operations. Test your engine on expressions like a * a + a where reuse is explicit, not just on straight-line formulas.

Checking against torch.autograd

An autograd engine is easy to get subtly wrong, and the professional habit is to check it against a reference on the same expression. One of HW1’s tests does this with a central finite difference on a composite expression, at a tolerance of 1e-5, and needs no PyTorch; the parity test against torch.autograd in the same suite is optional and carries no points. PyTorch itself is tested both ways (torch.autograd.gradcheck compares analytic gradients to finite differences). Here we run one moderately tangled expression (reuse, nesting, a nonlinearity) through our Value class and through PyTorch, and compare every gradient. The block repeats the class so it runs on its own. When extending Value for the exercises, update this copy too, or reuse your extended class and run only the comparison code below it. 

import math
import torch

# --- the same Value class as before, condensed ---
class Value:
    def __init__(self, data, _children=()):
        self.data, self.grad = float(data), 0.0
        self._backward, self._prev = lambda: None, set(_children)  # _prev: its parents

    def __add__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data + other.data, (self, other))
        def _backward():
            self.grad += out.grad
            other.grad += out.grad
        out._backward = _backward
        return out

    def __mul__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data * other.data, (self, other))
        def _backward():
            self.grad += other.data * out.grad
            other.grad += self.data * out.grad
        out._backward = _backward
        return out

    def tanh(self):
        t = math.tanh(self.data)
        out = Value(t, (self,))
        def _backward():
            self.grad += (1.0 - t * t) * out.grad
        out._backward = _backward
        return out

    def backward(self):
        topo, visited = [], set()
        def build(v):
            if v not in visited:
                visited.add(v)
                for c in v._prev:
                    build(c)
                topo.append(v)
        build(self)
        self.grad = 1.0
        for node in reversed(topo):
            node._backward()

    __radd__ = __add__
    __rmul__ = __mul__


def expression(a, b, c):
    """Same formula for both engines: reuse, nesting, and a nonlinearity."""
    d = a * b + c
    e = (d + a).tanh()
    return e * d + b


# Ours
a, b, c = Value(0.7), Value(-1.3), Value(2.0)
out = expression(a, b, c)
out.backward()

# PyTorch
at = torch.tensor(0.7, dtype=torch.float32, requires_grad=True)
bt = torch.tensor(-1.3, dtype=torch.float32, requires_grad=True)
ct = torch.tensor(2.0, dtype=torch.float32, requires_grad=True)
outt = expression(at, bt, ct)
outt.backward()

print(f"{'':8s}{'micro':>12s}{'torch':>12s}{'|diff|':>12s}")
print(f"{'out':8s}{out.data:12.6f}{outt.item():12.6f}{abs(out.data - outt.item()):12.2e}")
for name, v, t in [("dL/da", a, at), ("dL/db", b, bt), ("dL/dc", c, ct)]:
    print(f"{name:8s}{v.grad:12.6f}{t.grad.item():12.6f}{abs(v.grad - t.grad.item()):12.2e}")

Running this prints:

               micro       torch      |diff|
out        -0.269121   -0.269121    5.19e-08
dL/da      -1.263999   -1.263999    2.42e-08
dL/db       1.742557    1.742557    4.19e-08
dL/dc       1.060796    1.060796    5.93e-08

Note that expression is written once and fed both Value objects and PyTorch tensors: Python duck typing lets us reuse the same formula. Differences around 1e-8 remain because our engine computes in Python’s 64-bit floats while the PyTorch tensors here are float32, whose relative precision is about 1e-7. HW1’s optional parity test uses float64 in both engines and still compares with a tolerance, since operation order and math libraries can produce small rounding differences. Exact agreement on a particular expression is also possible and does not by itself indicate a faulty test.

Where calculus runs out: kinks and infinities

The ordinary chain rule assumes each local derivative exists. Several workhorse functions violate that assumption at isolated points: \(\operatorname{relu}(x) = \max(0, x)\) has a kink at \(x = 0\) where the left derivative is 0 and the right derivative is 1; abs also has a kink there, with one-sided derivatives of -1 and 1. For sqrt, the derivative \(1/(2\sqrt{x})\) grows without bound as \(x\) approaches 0 from above, so there is no finite derivative at 0. Frameworks specify how backward behaves at such points: they may return a chosen value, a non-finite value, or raise an error. The examples below show why those choices matter. 

import torch

torch.manual_seed(0)

# relu is not differentiable at 0. What does autograd return there?
x = torch.tensor([-1.0, 0.0, 1.0], requires_grad=True)
torch.relu(x).sum().backward()
print(f"relu'(x) at x = -1, 0, 1 : {x.grad.tolist()}")

# abs has the same kink at 0
xa = torch.tensor([-1.0, 0.0, 1.0], requires_grad=True)
xa.abs().sum().backward()
print(f"abs'(x)  at x = -1, 0, 1 : {xa.grad.tolist()}")

# sqrt at 0: the one-sided derivative is genuinely infinite
xs = torch.tensor([0.0, 1.0], requires_grad=True)
xs.sqrt().sum().backward()
print(f"sqrt'(x) at x = 0, 1     : {xs.grad.tolist()}")

# The where-trap: masking the VALUE does not mask the GRADIENT.
xw = torch.tensor([0.0, 4.0], requires_grad=True)
safe_looking = torch.where(xw > 0, torch.sqrt(xw), torch.zeros_like(xw))
safe_looking.sum().backward()
print(f"where(x>0, sqrt(x), 0)'  : {xw.grad.tolist()}")

# Keep the same piecewise values, but evaluate sqrt only on positive inputs.
xf = torch.tensor([0.0, 4.0], requires_grad=True)
positive = xf > 0
safe = torch.zeros_like(xf)
safe[positive] = torch.sqrt(xf[positive])
safe.sum().backward()
print(f"mask BEFORE sqrt       : {xf.grad.tolist()}")

# In-place mutation of a tensor that backward still needs is an error.
a = torch.ones(3, requires_grad=True)
b = a.exp()          # exp's backward rule reuses its OUTPUT b
b.add_(1.0)          # in-place: modifies b after it was saved
try:
    b.sum().backward()
except RuntimeError as e:
    print(f"in-place then backward -> RuntimeError: {str(e)[:75]}...")

Running this prints:

relu'(x) at x = -1, 0, 1 : [0.0, 0.0, 1.0]
abs'(x)  at x = -1, 0, 1 : [-1.0, 0.0, 1.0]
sqrt'(x) at x = 0, 1     : [inf, 0.5]
where(x>0, sqrt(x), 0)'  : [nan, 0.25]
mask BEFORE sqrt       : [0.0, 0.25]
in-place then backward -> RuntimeError: one of the variables needed for gradient computation has been modified by a...

At the ReLU kink, PyTorch returns 0. That is a choice, not a theorem: any value between 0 and 1 is a legitimate “subgradient” (the convex-analysis generalization of the derivative at a kink), and a framework could defensibly pick 0, 1, or 0.5. For abs, PyTorch’s answer at 0 is likewise 0. In theory the choice almost never matters (a continuous input has probability zero of landing exactly on the kink), but in float32 practice, exact zeros are common (ReLU itself manufactures them), and Bertoin and colleagues (NeurIPS 2021) report that the choice of relu'(0) can measurably change training outcomes at 16- and 32-bit precision, while being negligible at 64-bit. You do not need to lose sleep over the convention; you do need to know it exists, because HW1’s starter, like PyTorch, defines relu'(0) = 0.

The sqrt and where rows are the ones that cost real debugging hours. PyTorch’s inf for the backward rule at 0 reflects the unbounded right-hand limit; it is not a finite derivative. The expression torch.where(x > 0, torch.sqrt(x), 0) looks protective because its value at x = 0 is a clean 0. But both branch tensors are computed before where runs, so backward still reaches sqrt(0). Its infinite local derivative combines with a zero incoming gradient to produce nan, which can contaminate earlier gradients on the affected paths.

The added row fixes the computation by selecting positive inputs before evaluating sqrt. It preserves the intended piecewise forward values and returns 0 on the excluded branch, including at 0 by convention; the piecewise function still has no finite derivative at that boundary. Clamping or adding an epsilon can also stabilize some computations, but changes the function near zero; adding an epsilon alone does not protect arbitrary negative inputs.

Pitfall. Computing an unsafe expression and then masking its output can leave NaNs in backward even when the loss is finite. To debug this, run both forward and backward inside with torch.autograd.detect_anomaly():. It checks backward for NaNs and can report the forward traceback of the failing operation; it does not automatically rerun backward or detect every infinity. Also inspect inputs, losses, and gradients with torch.isfinite. These checks add overhead, so use them while debugging.

The last row demonstrates the guard rail PyTorch built for in-place operations. Methods ending in _ (add_, relu_) and Python’s += overwrite a tensor’s storage rather than allocating a new one. That is sometimes worth doing for memory, but it collides with reverse mode’s core requirement: backward rules reuse forward values. Here exp saves its own output (because d exp(x)/dx = exp(x), the cheapest form of the derivative), the add_ then rewrites that output in place, and backward would silently compute with the wrong number. PyTorch tracks a version counter on every tensor, records the version at save time, and refuses to run the backward rule if the tensor changed since — an error instead of a wrong gradient. Our sixty-line engine has no such protection: mutate some_node.data between forward and backward and you simply get corrupted gradients. When you write your HW1 engine, notice how much of a production framework is exactly this kind of defensive bookkeeping.

The graph’s lifecycle: retention, accumulation, detachment

Keeping the output of our Value expression keeps its reachable graph alive. The toy engine does not release saved state or reset intermediate .grad fields after backward. Calling it again on the same graph therefore reuses stale gradients: for f = a * a + a at a = 3, a second f.backward() leaves a.grad at 20, not 14. For a clean repeat, reset all reachable nodes’ gradients before reseeding the output. In training, rebuild the forward expression each iteration and clear the reused parameter gradients before accumulating the next update.

PyTorch distinguishes the graph structure from the tensors saved for backward. By default, backward releases its saved references as their rules finish; tensors still referenced elsewhere remain alive. Graph nodes may remain reachable through an output’s grad_fn, but another backward fails if it needs saved tensors that have been released. Simple operations that save no tensors can be exceptions. The first experiment below uses multiplication that needs saved values, so its second backward does fail.

There is also a distinction between computing an intermediate gradient and storing it in .grad. By default, PyTorch accumulates .grad on leaf tensors that require gradients, such as model parameters. An intermediate non-leaf tensor can participate in backward while its .grad remains None; call its retain_grad() before backward if you need to inspect that field. retain_grad() preserves that gradient, whereas retain_graph=True preserves saved state for another backward call. Our toy engine keeps a .grad field on every node. 

import torch

torch.manual_seed(0)

# 1. This backward releases saved values; repeating it needs those values.
w = torch.tensor(2.0, dtype=torch.float32, requires_grad=True)
loss = (w * w) * 3.0                       # dL/dw = 6w = 12
loss.backward()
print(f"after 1st backward: w.grad = {w.grad.item():.1f}")
try:
    loss.backward()
except RuntimeError as e:
    print(f"2nd backward -> RuntimeError: {str(e)[:60]}...")

# 2. retain_graph=True keeps the buffers; gradients then ACCUMULATE.
w.grad = None
loss = (w * w) * 3.0
loss.backward(retain_graph=True)
loss.backward()
print(f"two backwards with retain_graph: w.grad = {w.grad.item():.1f}  (12 + 12)")

# 3. detach() cuts the recording: b is a constant as far as autograd knows.
a = torch.tensor(3.0, requires_grad=True)
b = (a * a).detach()                       # value 9.0, no history
c = b * a
c.backward()
print(f"c = detach(a*a) * a: dc/da = {a.grad.item():.1f}  (only the last factor counts)")

# 4. This multiplication in no_grad() creates no reverse-mode history.
with torch.no_grad():
    d = a * a
print(f"inside no_grad: (a*a).requires_grad = {d.requires_grad}")

Running this prints:

after 1st backward: w.grad = 12.0
2nd backward -> RuntimeError: Trying to backward through the graph a second time (or direc...
two backwards with retain_graph: w.grad = 24.0  (12 + 12)
c = detach(a*a) * a: dc/da = 9.0  (only the last factor counts)
inside no_grad: (a*a).requires_grad = False

The first block demonstrates why ordinary training builds a fresh graph each iteration. retain_graph=True preserves saved state when several backward calls need the same forward computation, such as computing separate gradients for two losses before an update. It does not recompute values after parameters change, so do not use retention as a substitute for a fresh forward pass after opt.step(). If only the combined gradient is needed, call backward once on the sum of the losses.

The second block shows that w.grad becomes 24: parameter gradients accumulate across calls. Gradient accumulation turns this into a tool: run backward() on each of k small minibatches, each on its own fresh graph, then one opt.step(). For equal-sized minibatches, fixed parameters throughout the group, and a loss that is a mean of independent per-example losses, this matches the combined batch gradient: backpropagate loss / k for each. For unequal sizes, multiply each minibatch mean by batch_size / total_examples instead. Batch-dependent operations such as training-mode BatchNorm can break the equivalence. Clear the parameter gradients once before the group, not between its minibatches. PyTorch’s accumulation example uses exactly this scaling.

The third and fourth blocks control which computations are differentiated. detach() returns a tensor sharing the same storage but disconnected from the preceding history; later operations can still build a graph through other inputs. It is the mechanism behind every stop-gradient trick you will meet: freezing a pretrained encoder while training a head, or target networks whose weights must not receive gradients. It does not undo the work or memory already spent recording the original computation, and modifying detached storage can still affect the original tensor. torch.no_grad() prevents reverse-mode recording for operations performed inside the block. Use it for inference when gradients are unnecessary; tensors are freed when no longer referenced, not necessarily on the next line. It is separate from model.eval(), which controls behaviors such as dropout and BatchNorm.

Differentiating the derivative

In our engine, grad fields are plain floats — dead ends. Nothing stops you from asking a sharper question: what is the derivative of the derivative? If the backward pass is itself executed with recorded operations, the gradient becomes a differentiable function like any other, and reverse mode can run on it again. PyTorch exposes this through create_graph=True, most cleanly via the functional interface torch.autograd.grad, which returns gradients as tensors instead of stashing them in .grad: 

import torch

torch.manual_seed(0)

# Higher-order derivatives of tanh at x = 0.5, checked against closed forms.
x = torch.tensor(0.5, requires_grad=True)
y = torch.tanh(x)

(g1,) = torch.autograd.grad(y, x, create_graph=True)   # graph OF the gradient
(g2,) = torch.autograd.grad(g1, x, create_graph=True)
(g3,) = torch.autograd.grad(g2, x)

t = y.item()
print(f"f   = tanh(0.5)          = {t:.6f}")
print(f"f'  autograd = {g1.item():9.6f}   analytic 1 - t^2          = {1 - t*t:9.6f}")
print(f"f'' autograd = {g2.item():9.6f}   analytic -2 t (1 - t^2)   = {-2*t*(1 - t*t):9.6f}")
print(f"f'''autograd = {g3.item():9.6f}   analytic (1 - t^2)(6t^2-2) = {(1 - t*t)*(6*t*t - 2):9.6f}")

# A full Hessian, one call: f(w) = w0^2 * w1 at w = (1.5, -2.0)
def f(w):
    return w[0] ** 2 * w[1]

w = torch.tensor([1.5, -2.0])
H = torch.autograd.functional.hessian(f, w)
print(f"\nHessian of w0^2 * w1 at (1.5, -2.0):")
print(f"  autograd: [[{H[0,0]:.1f}, {H[0,1]:.1f}], [{H[1,0]:.1f}, {H[1,1]:.1f}]]")
print(f"  analytic: [[2*w1, 2*w0], [2*w0, 0]] = [[-4.0, 3.0], [3.0, 0.0]]")

Running this prints:

f   = tanh(0.5)          = 0.462117
f'  autograd =  0.786448   analytic 1 - t^2          =  0.786448
f'' autograd = -0.726862   analytic -2 t (1 - t^2)   = -0.726862
f'''autograd = -0.565209   analytic (1 - t^2)(6t^2-2) = -0.565209

Hessian of w0^2 * w1 at (1.5, -2.0):
  autograd: [[-4.0, 3.0], [3.0, 0.0]]
  analytic: [[2*w1, 2*w0], [2*w0, 0]] = [[-4.0, 3.0], [3.0, 0.0]]

Three derivatives of tanh evaluated by autodiff, each obtained by differentiating the graph that computed the previous one, and a two-by-two Hessian that matches the hand calculation entry for entry. Think about what the third line required: a graph recording the backward pass of the backward pass of the forward pass. The machinery is the same sixty-line idea, applied to its own output.

The costs deserve respect. Each additional order multiplies running time by a constant factor and (the reverse-mode signature again) retains yet more intermediate state. Worse, a full Hessian of a network with \(n\) parameters is an \(n \times n\) object requiring \(n\) backward sweeps to materialize row by row: fine at \(n = 2\), out of the question at \(n = 10^7\). The affordable object is the Hessian-vector product \(H\mathbf{v}\), which for a twice continuously differentiable loss is the gradient of the scalar \(\mathbf{g}^\top \mathbf{v}\), with \(\mathbf{g} = \nabla L\) and \(\mathbf{v}\) held constant: one extra backward sweep, a trick worked out by Pearlmutter in 1994, decades before it became a standard tool for curvature diagnostics and second-order optimizer research.

Where will you actually meet double backward? Gradient-penalty losses regularize the norm of a gradient (the WGAN-GP objective in Chapter 4 is the canonical example), which requires differentiating through a backward pass; meta-learning methods differentiate through an entire inner training loop; physics-informed networks put d(output)/d(input) terms directly in the loss. Whenever you see create_graph=True in real code, this section is what is happening.

From sixty lines to a real framework

What separates this toy from torch.autograd? Mostly engineering, not concept.

Tensors, not scalars. PyTorch nodes hold whole tensors, and each backward rule is a vectorized kernel: the local derivative of a matrix multiply \(Z = XW\) is itself expressed with matrix multiplies (\(\partial L/\partial X = (\partial L/\partial Z)\, W^\top\) and \(\partial L/\partial W = X^\top\, \partial L/\partial Z\)), as you saw in backpropagation. One graph node can thus represent millions of scalar nodes from our engine — vectorization and optimized kernels account for much of the speedup.

Broadcasting rules. When a bias b of shape (n,) is added to a batch of shape (B, n), the forward pass copies b across B rows implicitly; the backward pass must therefore sum the incoming gradient over the broadcast dimension. Getting these reductions right is a classic source of shape bugs in hand-rolled engines, and the opening NumPy example shows the required sum.

Graph lifecycle and memory. Backward rules save the forward values they need, which is why training takes more memory than inference; parameter gradients and optimizer state add to it. PyTorch normally releases saved references as backward proceeds; retain_graph=True delays that release. torch.no_grad() skips reverse-mode recording when gradients are unnecessary. Our toy engine has no explicit release policy and relies on Python to collect unreachable nodes and their closure cycles.

Robustness. A real engine handles in-place mutation detection, mixed dtypes and devices, double-backward (differentiating the backward pass itself), and an iterative graph traversal — our recursive build would hit Python’s recursion limit on a chain approaching the default limit of about 1000 recursive calls, a refinement worth making once your HW1 engine passes its tests.

Tape, graph, and back again. The define-by-run tape also sits at a particular point in a design argument that has now swung twice. The static-graph frameworks that preceded it (Theano, and TensorFlow 1.x after 2015) asked you to declare the whole computation up front, in effect writing a program in a small graph language, and then compiled and ran it. That bought real advantages: with the entire graph in hand, the framework could fuse operations, plan memory, and ship the result to production without Python attached. It also made everyday work painful: errors surfaced far from the line that caused them, and a simple if statement on a tensor’s value required special graph constructs. Chainer and then PyTorch bet that researcher iteration speed mattered more, and won: by the late 2010s define-by-run was the dominant research interface. But the tape leaves performance on the table — every operation dispatches through Python, and no cross-operation optimization is possible when the framework only ever sees one op at a time. So the field moved partway back: JAX’s jit traces Python functions for compilation with XLA, and PyTorch 2.0 (2023) added torch.compile, which captures graphs out of running bytecode and compiles them, falling back to eager execution where capture fails. The modern synthesis is not the old static world (you still write, run, and debug imperative code), but it quietly rebuilds the graph the compilers always wanted. PyTorch in Practice returns to what this means day to day.

None of this changes the picture you should carry forward: loss.backward() walks a recorded DAG in reverse topological order, calling one small local-derivative rule per node and summing contributions into .grad. Frameworks add policies for saved tensors, gradient buffers, and recording controls around this same reverse sweep.

Going deeper

The memory wall, and checkpointing. Saving forward activations can become a major memory constraint as depth, batch size, or sequence length grows; parameters, gradients, and optimizer state also consume memory. Gradient checkpointing (also called activation recomputation or rematerialization) saves selected activations and recomputes others during backward. Chen and colleagues’ 2016 paper, “Training Deep Nets with Sublinear Memory Cost,” describes a square-root-of-depth storage scheme with the compute cost of an extra forward pass under its assumptions. Actual savings depend on the model and checkpoint placement; an order-of-magnitude reduction is not guaranteed. torch.utils.checkpoint provides this tradeoff in PyTorch. Our toy engine would need changes to release and recompute the values its closures retain.

A cleaner vocabulary: VJPs and JVPs. The literature describes what _backward does as a vector-Jacobian product (VJP): each node’s rule takes the incoming gradient vector and multiplies it by that operation’s local Jacobian, without ever building the Jacobian as a matrix. Reverse mode is a chain of VJPs from the output back; forward mode is the mirror-image chain of Jacobian-vector products (JVPs) from an input forward. The composition view has become an interface: JAX, and PyTorch’s torch.func namespace, expose grad, vjp, jvp, and vmap as composable function transforms, so “gradient of a function” is itself a function you can transform again. Higher-order tricks drop out of composing the two modes — for example, forward-over-reverse computes a Hessian-vector product by applying a JVP to the gradient function. It is often more memory-efficient than reverse-over-reverse, but the best choice depends on the operations and implementation.

What backprop computes at kinks is still being formalized. The relu'(0) convention from earlier is a small window onto an active research question: when a program contains nondifferentiable operations, the thing autograd returns is not always a gradient or even a subgradient of the composite function, and optimization theorists (Bolte, Pauwels, and others, in work from around 2020 onward) have developed notions such as “conservative gradients” to characterize what backpropagation-as-implemented actually produces and why SGD still converges with it. For daily practice the takeaway is the modest one from this section (know the conventions, keep unsafe expressions out of the computation), but it is worth knowing that the tidy chain-rule story has genuinely open edges.

How real engines schedule the sweep. Our backward() is a recursive depth-first topological sort, which is correct but fragile (recursion depth) and sequential. PyTorch’s autograd engine instead does the moral equivalent of a reference-counted reverse sweep: it counts, for each node, how many consumers must deposit gradient contributions before the node’s own rule may fire, and processes nodes whose counts reach zero — a form of Kahn’s algorithm that needs no recursion and exposes parallelism (independent branches can become ready together, though actual concurrency also depends on device and stream dependencies). The invariant being protected is exactly the one from the reuse example: no rule fires until its node’s gradient is final.

Accumulation is not always a bigger batch. The loss / k equivalence assumes equal-sized minibatches and a loss that is a mean over examples at fixed parameters. Unequal minibatches need weighting by their share of the examples, and batch-dependent operations such as BatchNorm in training mode (S07) break the equivalence outright.

Common pitfalls

  • Symptom: gradients are exactly right for simple formulas but wrong when a variable appears twice. Cause: _backward overwrites with = instead of accumulating with +=. Fix: accumulate, and add a regression test such as a * a + a whose true derivative you know.
  • Symptom: gradients include contributions from earlier updates. Cause: parameter .grad fields persist; reusing the toy graph can also reuse stale intermediate gradients. Fix: clear parameter gradients before each update’s accumulation group and build a fresh forward graph for each minibatch. For a clean repeat on the same toy graph, reset every reachable node before reseeding the output.
  • Symptom: RecursionError: maximum recursion depth exceeded on a long chain of operations. Cause: the depth-first topological sort is recursive, and Python’s default recursion limit is about 1000 frames. Fix: rewrite the sort iteratively with an explicit stack.
  • Symptom: every gradient comes back 0.0. Cause: either the output’s gradient was never set to 1.0, or somewhere in the forward pass you dropped out of the graph — for example calling math.tanh(x.data) and wrapping the result in a fresh Value with no parents. Fix: set the output’s gradient to 1.0 before the walk, and make sure every operation goes through methods that record _prev and _backward.
  • Symptom: the comparison differs in the last few digits and you loosen the tests to 1e-2. Cause: mixed float32/float64 may explain small discrepancies, but is not proof that the implementation is correct. Fix: compare both engines in float64 with suitable relative and absolute tolerances, especially near zero. Investigate remaining discrepancies instead of weakening HW1’s tests.

Try it yourself

  1. Extend Value with __pow__ (for constant exponents), __neg__, __sub__, __truediv__, and the reflected forms __rsub__ and __rtruediv__ (define a / b as a * b**-1). Verify each new operation against torch.autograd using the harness in Checking against torch.autograd. If PyTorch is not installed, use hand-derived values and central differences at smooth points.
  2. Add exp(), then implement sigmoid two ways: composed from your primitives (1 / (1 + (-x).exp())) and as a single fused operation with local derivative s * (1 - s). Compare their gradients within floating-point tolerance on several moderate inputs. The literal 1 / ... needs the __rtruediv__ method from exercise 1.
  3. Deliberately change one += to = in __mul__. Find the smallest expression whose gradient becomes wrong, and one nontrivial expression whose gradient stays right. Explain both.
  4. Using only your Value class, fit \(y = 2x - 1\) from five data points by gradient descent on two parameters w and b with squared-error loss. You now have end-to-end training with no framework at all.
  5. Explain why your Value engine cannot compute second derivatives as written (what type does grad have, and what would it need to be?). Then sketch (or implement for + and * only) a version whose _backward builds new Value graphs instead of accumulating floats, and check \(\mathrm{d}^2(a^3)/\mathrm{d}a^2 = 6a\) against torch.autograd.grad with create_graph=True.

Further reading

  • Rumelhart, Hinton, and Williams, “Learning representations by back-propagating errors,” Nature, 1986. nature.com/articles/323533a0 — the four-page paper that made reverse-mode differentiation famous in neural networks; worth reading for how little it needed to say.
  • Baydin, Pearlmutter, Radul, and Siskind, “Automatic Differentiation in Machine Learning: a Survey,” JMLR, 2018. arxiv.org/abs/1502.05767 — the clearest map of forward versus reverse mode, and of where “backprop” sits inside the broader AD landscape.
  • Karpathy, micrograd (2020). github.com/karpathy/micrograd — the roughly hundred-line engine this section and HW1 are modeled on; read its tests to see how a reference check against PyTorch is organized.
  • Griewank and Walther, Evaluating Derivatives: Principles and Techniques of Algorithmic Differentiation, 2nd ed., SIAM, 2008. epubs.siam.org/doi/book/10.1137/1.9780898717761 — the standard reference on AD theory, including the checkpointing trade-offs sketched in Going deeper.
  • Pearlmutter, “Fast Exact Multiplication by the Hessian,” Neural Computation, 1994. doi.org/10.1162/neco.1994.6.1.147 — the Hessian-vector-product trick: exact curvature information for one extra pass, no Hessian ever materialized.
  • Bertoin, Bolte, Gerchinovitz, and Pauwels, “Numerical influence of ReLU’(0) on backpropagation,” NeurIPS 2021. arxiv.org/abs/2106.12915 — what the kink convention actually does to real training runs at 16-, 32-, and 64-bit precision.
  • Paszke et al., “Automatic differentiation in PyTorch,” NIPS 2017 Autodiff Workshop. openreview.net/forum?id=BJJsrmfCZ — a short design paper on torch.autograd: the tape, in-place versioning, and the engineering around them.
  • Tokui et al., “Chainer: A Deep Learning Framework for Accelerating the Research Cycle,” KDD 2019. arxiv.org/abs/1908.00213 — the define-by-run originators’ own account of the design and why it displaced static graphs for research.
  • Dive into Deep Learning, section 2.5, “Automatic Differentiation.” d2l.ai/chapter_preliminaries/autograd.html — the same machinery at tensor level, with exercises that pair well with HW1.
  • PyTorch documentation, “Autograd mechanics.” docs.pytorch.org/docs/stable/notes/autograd.html — the authoritative page on graph retention, in-place rules, and double backward.

Before class you should be able to

  • List the four core fields every graph node stores, and explain how closures carry forward-pass values into the backward pass.
  • Explain why gradients must accumulate with +=, and why a * a + a is the test that exposes the overwrite bug.
  • Explain, using the reuse example, why the backward sweep must run in reverse topological order.
  • State the relu'(0) convention and explain how a masked sqrt still produces NaN gradients (the where-trap).
  • Explain when a second backward() needs retained saved tensors, why parameter gradients accumulate, how detach() cuts history, and how no_grad() prevents reverse-mode recording.
  • Explain when to clear parameter gradients and how to scale minibatch mean losses for gradient accumulation.
  • Reproduce the vectorized layer rules from the S04 reading and pass a gradient check on a small MLP.

Connects to

  • This session is covered by Quiz 1 (Module 1: S01–S08), given at S09.
  • This section is the primary reference for HW1 (micro-autograd, out at S04, due at S08), which grades your engine on known gradient values and one finite-difference check like the first example’s; the comparison against torch.autograd is an optional, unscored extra.
  • The PyTorch-facing consequences (zero_grad(), no_grad(), detach()) return in 1.8 (S08) and in every assignment after it.

The PDF posted in Canvas is the assigned version of this reading; it is built from this page.

← S04 · Backpropagation I · All readings · S06 · Optimizers →