Backprop from scratch

The puzzle: PyTorch did calculus you never wrote

Something strange has been happening for two weeks. You write y = torch.log(x * x), call y.backward(), and PyTorch announces that the derivative is 0.6667. You never supplied calculus or imported a symbolic math library. log and * are ordinary functions. Where did a derivative come from?

The answer is delightfully mechanical: PyTorch was taking notes. Every tensor operation records its name, inputs, and output, building a graph as the computation runs. Differentiation then walks those operations in reverse and asks each one a local question. The work is bookkeeping plus one multiplication per edge. By the end of the notebook, you will have written the mechanism yourself in plain Python.

This is home turf for coders. The work is a dependency graph, a topological sort, and careful accumulation: problems you already know how to solve.

The chain rule for programmers

You need one piece of calculus, stripped of notation. A derivative is a local multiplier: if b changes a little, how much does c change? If a function has derivative 3 at your current input, then wiggling the input by a tiny h wiggles the output by about 3h. Local multipliers compose by multiplying: if a changes b, and b changes c, then a changes c by the product of those local effects. That composition rule is the chain rule.

You can measure a local multiplier with a nudge:

# 01: a derivative is a measured slope (no calculus required)
import math

def f(x):
    return math.log(x * x)

x = 3.0
h = 0.001
slope = (f(x + h) - f(x)) / h
print(f"measured slope at x=3: {slope:.4f}")
print(f"claimed  slope 2/x:    {2/x:.4f}")

Expected output: measured slope at x=3: 0.6666 and claimed slope 2/x: 0.6667. The measured and claimed values agree to three decimals (the tiny gap is the finite nudge; shrink h and it shrinks). This "nudge test" is also how you'll sanity-check your own backward functions all week: any gradient you compute can be audited by a two-line finite difference.

Choose h with care. A large nudge measures the average slope across a wide interval and misses local curvature. An extremely small nudge makes floating-point subtraction lose precision. Values around 1e-3 or 1e-5 work for these exercises. Gradient checks are slow because they rerun the function once per input, so use them on tiny test arrays. Backprop computes every derivative in one reverse pass.

For y = log(x * x), the computation has two steps: square, then log. Backprop walks those steps backward, asking each operation how upstream change should be routed to its inputs.

# no-run
# 02: PyTorch agrees with the nudge test
import torch

x = torch.tensor(3.0, requires_grad=True)
y = torch.log(x * x)
y.backward()
print(x.grad)
# tensor(0.6667) because d log(x^2) / dx = 2 / x

The requires_grad=True flag tells PyTorch which tensors need notes. x is a leaf tensor, an input to the graph with no producing operation, and PyTorch accumulates gradients in its .grad field.

Autograd records a graph

Every tensor operation creates a node. The node stores enough information to answer a backward question later. Autograd is the system that records this graph and traverses it in reverse topological order.

The graph is directed and acyclic: values flow from inputs toward one output, with no operation depending on its own future result. A topological sort places every producer before its consumers. Reversing that order guarantees that a node has received gradients from all downstream paths before it routes their sum to its inputs. Shared values make that ordering visible because several branches can converge on one leaf.

The entire graph for y = log(x * x) sends forward values down and backward gradients up:

Trace the blue path with actual numbers; this trace is the algorithm. Start at the bottom: the gradient of y with respect to itself is always 1.0. Step up through log. Its local rule multiplies by 1/input, and its input was e = 9.0, so the gradient becomes 1.0 × 1/9 ≈ 0.111. Step up through multiply. Its local rule routes gradient to each input by multiplying by the other input, so both edges back to x carry 0.111 × 3.0 = 0.333. Both edges land on the same x, so their gradients add: 0.333 + 0.333 ≈ 0.667. This matches the 0.6667 from PyTorch in cell 02 and the nudge test in cell 01. Three methods, one number.

For coders, the shape is familiar: build dependency graph, sort dependencies, run callbacks in reverse. The math is local to each operation.

# no-run
# 03: the notebook's backward-function signatures
def log_back(grad_out, out, x):
    # `out` = log(x) is also passed; other operations use it
    return grad_out / x

def multiply_back0(grad_out, out, x, y):
    return grad_out * y      # gradient with respect to x

def multiply_back1(grad_out, out, x, y):
    return grad_out * x      # gradient with respect to y

Each backward function receives the gradient of the output (plus the forward output and inputs) and returns the gradient for one input. That is why the notebook has multiply_back0 and multiply_back1: one function per argument. (The notebook's real versions also pass each result through unbroadcast; the next section explains why.)

grad_out carries everything that happened after the current operation. If e = x * x feeds y = 2 * e, the multiply node receives grad_out = 2, then multiplies it by each local slope. This split keeps every rule reusable. multiply_back0 behaves the same whether its output feeds a logarithm, ten later layers, or the final loss directly.

Each function contains one local rule. log_back knows how log routes a gradient and has no knowledge of the surrounding graph. Composing these small rules in reverse dependency order computes the global derivative.

Run the whole trace yourself:

# 04: backprop by hand, the full algorithm in ten lines
import math

x = 3.0
e = x * x                 # forward step 1
y = math.log(e)           # forward step 2
print(f"forward: e = {e}, y = {y:.4f}")

grad_y = 1.0              # dy/dy, where every backward pass starts
grad_e = grad_y / e       # log_back: route through log
grad_x = grad_e * x       # multiply_back0: path through the first x
grad_x += grad_e * x      # multiply_back1: second path through x; accumulate!
print(f"backward: x.grad = {grad_x:.4f}")

Expected output: forward: e = 9.0, y = 2.1972, then backward: x.grad = 0.6667. That is the diagram, executed. The notebook's backprop() exercise generalises this cell by topologically sorting an arbitrary graph and running the right *_back function at each node.

The two gotchas: reuse and broadcasting

The engine is simple; two details make it subtle, and the notebook tests you on both.

Reuse means accumulate

If a tensor is used twice, gradients accumulate. The graph has two paths back to the same value, and the answer is their sum. You already met this in cell 04: += gives 0.667, while = overwrites the first path and gives 0.333. Every "write a gradient to a node" in your engine must use +=. Real PyTorch training loops call zero_grad() because .grad fields also accumulate across backward passes; clear them between steps.

Broadcasting forward means summing backward

Broadcasting reverses through a sum. If the forward pass stretched a tensor from (3,) to (2, 3) (week 1's implicit copy), the backward pass sums over the stretched axis to return to (3,). This is the reuse rule in disguise: the forward pass used a broadcast value once per copy, so its gradient sums the contributions from all copies.

# no-run
# 05: broadcast forward, sum backward (PyTorch's view)
import torch

x = torch.ones(2, 3, requires_grad=True)
b = torch.arange(3.0, requires_grad=True)
y = (x + b).sum()
y.backward()
print(b.grad)
# tensor([2., 2., 2.]) because b was used once per row

The notebook's unbroadcast exercise asks you to implement that reverse move in numpy. Work through both cases before the exercise:

# 06: unbroadcast both cases (numpy, runs in your browser)
import numpy as np

# case 1: forward stretched (3,) -> (2, 3) by adding a new leading axis.
# backward: sum over that axis, and the axis itself disappears.
grad_out = np.ones((2, 3))
print(grad_out.sum(axis=0))            # -> shape (3,)

# case 2: forward stretched (2, 1) -> (2, 3) along an existing size-1 axis.
# backward: sum over it and KEEP the axis, returning to shape (2, 1).
grad_out2 = np.ones((2, 3))
print(grad_out2.sum(axis=1, keepdims=True))   # -> shape (2, 1)

Expected output: [2. 2. 2.], then [[3.] [3.]]. Case 1 (new axes prepended) sums without keepdims; case 2 (existing size-1 axes stretched) sums with keepdims=True. Every unbroadcast bug in the session will be one of these two cases handled with the wrong rule.

A general unbroadcast implementation follows the shapes. First, sum away leading dimensions that the forward pass added. Then inspect the remaining axes: whenever the original input had size 1 and the output had a larger size, sum that axis with keepdims=True. Assert that the result now matches the input shape. That final assertion turns a vague downstream failure into a local shape error at the backward rule that caused it.

From autograd to modules

The notebook finishes by rebuilding Parameter, Module, and Linear, then training MNIST on the tiny engine you wrote. That closes the loop from week 2: the abstractions package graph bookkeeping into reusable classes.

The circle is complete. Week 2's Linear was "a matrix multiply plus a bias, wrapped in a class." Week 3's optimizer was "fifteen lines that nudge parameters using .grad." This week supplies the last missing organ, the code that fills in .grad: a reverse graph walk that calls functions such as multiply_back0. Stack all three weeks and you have a working deep-learning framework, written end to end by you and training a real model on real data. PyTorch uses the same machinery with better performance engineering.

Pair-session guide

Core work is ARENA 0.4 sections 1 and 2: primitive backward functions and the backprop engine. Stretch work is rebuilding the neural-network library and training MNIST. Have one person draw the graph while the other writes code for the backward pass; switch after each exercise.

What you should see

Your backprop() should agree with PyTorch .backward() on the same small graphs. In the stretch section, MNIST should train using your own mini-autograd engine. The win condition is being able to say where each gradient came from.

Where to go next

Act II starts next week: you build GPT-2 itself. Two short readings cement this week first:

• Calculus on Computational Graphs: Backpropagation by Chris Olah, the ARENA notebook's recommended written version of this week's mental model.
• 3Blue1Brown's What is backpropagation really doing? if you prefer the animated version.