Backpropagation from scratch

Autograd is a convenience. Interviewers for research roles want to see that you could do without it.

The network

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

A = XW₁ + b₁, H = max(0, A), Z = HW₂ + b₂, P = softmax(Z), L = −(1/N) Σ_n log P_{n,y_n}

W₁ is (D, H), b₁ is (H,), W₂ is (H, C), b₂ is (C,). Broadcasting adds each bias to every row. Nothing here needs a loop over examples, and an interview answer that has one will be asked to remove it.

Backpropagation is the chain rule, in reverse order

Write G_X for ∂L/∂X, which always has the same shape as X. Start at the loss and walk back through each line of the forward pass:

Softmax with cross-entropy: G_Z = (P − Y)/N, with Y the one-hot labels. The softmax Jacobian never has to be formed.
A linear layer Z = HW₂ + b₂: G_W₂ = HᵀG_Z, G_b₂ = the column sums of G_Z, and G_H = G_ZW₂ᵀ.
ReLU: G_A = G_H ⊙ [A > 0]. The gradient passes where the unit was active and stops where it was not.
The first layer: the same rule as the second, G_W₁ = XᵀG_A.

There is a quick way to get every matrix product right without memorising it: the gradient of a weight must have the weight's shape. G_W₂ is (H, C); the only way to make an (H, C) matrix from H (N, H) and G_Z (N, C) is HᵀG_Z.

A two-weight network: forward, then the chain rule backwards.

Step 1 / 10

Forward pass. The input x = 2; the weights are w₁ = 1.5 and w₂ = -0.5; the target is 1.
Step 2 / 10

z = w₁ · x = 1.5 × 2 = 3.
Step 3 / 10

h = relu(z) = 3. z is positive, so the unit is active; cache that fact for the backward pass.
Step 4 / 10

y = w₂ · h = -0.5 × 3 = -1.5.
Step 5 / 10

loss = ½ (y − t)² = ½ (-1.5 − 1)² = 3.125.
Step 6 / 10

Backward pass, from the end. ∂loss/∂y = y − t = -2.5.
Step 7 / 10

∂loss/∂w₂ = ∂loss/∂y × h = -2.5 × 3 = -7.5. That is the gradient for w₂.
Step 8 / 10

∂loss/∂h = ∂loss/∂y × w₂ = -2.5 × -0.5 = 1.25.
Step 9 / 10

relu passes the gradient through because z > 0: ∂loss/∂z = 1.25.
Step 10 / 10

∂loss/∂w₁ = ∂loss/∂z × x = 1.25 × 2 = 2.5. Every gradient used a value the forward pass cached.

Cache what the backward pass needs

The backward pass reuses the forward pass's intermediate values: A for the ReLU mask, H for G_W₂, P for G_Z. Real frameworks store them during the forward pass, and that is why training uses several times the memory of inference. Recomputing them instead is called activation checkpointing; it trades memory for time.

Check the gradient numerically

Every hand-written gradient should be checked against central differences before it is trusted:

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²). Expect relative agreement of 10⁻⁶ or better in float64. Agreement of 10⁻² means a bug, not rounding. The one legitimate exception is a ReLU unit whose input sits within h of zero, where the function has a kink.

Optimisers

Plain gradient descent uses one learning rate for every parameter. Adam keeps two running averages per parameter, of the gradient and of its square, and divides one by the square root of the other. The step for each parameter is then roughly the learning rate times the sign of a smoothed gradient, whatever the gradient's scale. Both averages start at zero, so they are biased towards zero early on; Adam corrects that by dividing by 1 − β^t. Without the correction, the first steps are far too small, or, for the second moment, far too large.

What interviewers listen for

Shapes stated before code. If you can say what shape every gradient has, the code follows.
Knowing why P − Y appears, and that it holds only for softmax paired with cross-entropy.
Mentioning gradient checking without being asked.
Knowing what dies when a ReLU's input is always negative, and what that does to training.

Practise it

Easy10 min