Backpropagation

The idea#

At the end of a long production line sits a faulty part. You walk the line backwards and ask at every station: how much of this fault came from here? Whoever contributed most has to correct themselves most.

Backpropagation does exactly that, with computation steps instead of stations.

Rebuilt by hand#

import numpy as np

def relu(x):      return np.maximum(0, x)
def relu_grad(x): return (x > 0).astype(float)

rng = np.random.default_rng(0)
X = rng.normal(size=(64, 8)); y = rng.normal(size=(64, 1))
W1 = rng.normal(0, 0.1, (8, 32)); b1 = np.zeros(32)
W2 = rng.normal(0, 0.1, (32, 1)); b2 = np.zeros(1)

And the loop itself. Every line in the backward part is exactly one link of the chain rule.

for step in range(300):
    # Forward - z1 and a1 have to be kept for the backward pass.
    z1 = X @ W1 + b1; a1 = relu(z1)
    z2 = a1 @ W2 + b2
    loss = float(((z2 - y) ** 2).mean())

    # Backward - each line is one link of the chain rule.
    dz2 = 2 * (z2 - y) / len(X)
    dW2 = a1.T @ dz2;      db2 = dz2.sum(0)
    da1 = dz2 @ W2.T
    dz1 = da1 * relu_grad(z1)
    dW1 = X.T @ dz1;       db1 = dz1.sum(0)

    for p, g in ((W1, dW1), (b1, db1), (W2, dW2), (b2, db2)):
        p -= 0.05 * g
    if step % 100 == 0:
        print(step, round(loss, 4))

Common mistakes#

Why backwards and not forwards#

At n = 7e9 parameters and m = 1 loss value that is the difference between seven billion passes and one. Which is the entire reason it is called backpropagation.

The memory bill#

Worked through for b = 8, s = 4096, h = 4096, a = 32 in bfloat16: the linear part is roughly 268 MB per layer with k₁ ≈ 2, the quadratic part 8 × 32 × 4096² × 2 bytes ≈ 8.6 GB per layer. The second term dominates completely, which is why FlashAttention, which never materialises that matrix in full, is not a detail but the precondition for long contexts. See Context windows in depth.

Gradient checkpointing#

Store only every √L-th layer out of L and recompute the rest, and activation memory falls from O(L) to O(√L) at roughly 30 percent more compute. With 80 layers that means 9 stored instead of 80. It is the difference between "does not fit" and "fits".