Backpropagation from scratch in NumPy: let the shapes do the work

You can use PyTorch for years without ever having to write a backward pass. Then someone asks you for one in plain NumPy, on a two-layer network, and the chain rule you know on paper turns into a fight with matrix shapes. Here's how I'd write it, along with the trick that'll keep those shapes straight for you.

The network

We'll take a batch X of shape (N, D), one hidden layer of width H with ReLU, and C output classes:

$$\begin{aligned} A &= XW_1 + b_1, & H &= \max(0, A), \cr Z &= HW_2 + b_2, & P &= \operatorname{softmax}(Z), \cr L &= -\frac{1}{N} \sum_n \log P_{n, y_n} \end{aligned}$$

$W_1$ is (D, H), $b_1$ is (H,), $W_2$ is (H, C), $b_2$ is (C,). Broadcasting adds the bias to all of the rows. None of this needs a loop over examples, and if your interview answer's got one, you'll be asked to take it out.

Backpropagation is the chain rule, in reverse order

Write G_X for ∂L/∂X, which always has the same shape as X. We start at the loss and walk back through the forward pass a line at a time:

• Softmax with cross-entropy: $G_Z = (P - Y)/N$ , with Y the one-hot labels. You don't ever have to form the softmax Jacobian.
• A linear layer $Z = HW_2 + b_2$ : $G_{W_2} = H^\top G_Z$ , $G_{b_2}$ is the column sums of $G_Z$ , and $G_H = G_Z W_2^\top$ .
• ReLU: $G_A = G_H \odot [A > 0]$ . The gradient gets through where the unit was active, and stops where it wasn't.
• The first layer: the same rule as the second, $G_{W_1} = X^\top G_A$ .

There's a quick way to get all of these matrix products right without memorising them. The gradient of a weight has to have the weight's shape. $G_{W_2}$ is (H, C); the only way to make an (H, C) matrix from H (N, H) and $G_Z$ (N, C) is $H^\top G_Z$ .

Cache what the backward pass needs

The backward pass reuses values from the forward pass: A for the ReLU mask, H for $G_{W_2}$ , and P for $G_Z$ . Real frameworks store them on the way forward, and that's why training takes several times the memory of inference. If you recompute them on the way back, that's called activation checkpointing, and you're trading memory for time.

Check the gradient numerically

You should check any gradient you've written by hand against central differences before you trust it:

def numerical_grad(f, theta, h=1e-6):
    g = np.zeros_like(theta)
    for i in np.ndindex(theta.shape):
        old = theta[i]
        theta[i] = old + h; up = f()
        theta[i] = old - h; down = f()
        theta[i] = old
        g[i] = (up - down) / (2 * h)
    return g

Central differences are accurate to $O(h^2)$ , so in float64 you'd expect them to agree to $10^{-6}$ or better, relative. If all you get is $10^{-2}$ , that's a bug, and rounding won't explain it. The one fair exception is a ReLU unit whose input sits within h of zero, where the function's got a kink.

Optimisers

Plain gradient descent uses the same learning rate for all of the parameters. Adam keeps two running averages for a parameter, one of its gradient and one of the gradient's square, and divides the first by the square root of the second. So the step it takes is roughly the learning rate times the sign of a smoothed gradient, whatever the gradient's scale. Both of the averages start at zero, which biases them towards zero early on, and Adam corrects for that by dividing by 1 − β^t. If you left that out, the first moment on its own would make the early steps far too small, and the second moment on its own would make them far too large.

If this comes up in an interview

State the shapes before you write any code. Once you can say what shape all of the gradients have, you'll find that most of the code follows. Be ready to explain why P − Y turns up, and that it only holds for softmax with cross-entropy. Bring up gradient checking yourself, before you're asked. And you should expect the follow-up on what happens to a ReLU whose input is always negative.

There's a longer version, with the shapes worked through, here: https://geodesicpartners.co.uk/ml-backpropagation