Backpropagation in NumPy: Gradients of a Two-Layer MLP

For the network A = XW1 + b1, H = max(0, A), Z = HW2 + b2 and the loss L = −(1/N) Σ_n log softmax(Z_n)_{y_n}, return [dW1, db1, dW2, db2], the gradients of L with respect to each parameter, each with the same shape as the parameter. y holds integer labels in 0…C−1.

The tests are generated so that no hidden pre-activation is exactly zero.

Examples

Input X = [[1, 2]], y = [0], W1 = [[1, 0], [0, 1]], b1 = [0, 0], W2 = [[1, 0], [0, 1]], b2 = [0, 0]

Output [[[-0.7311, 0.7311], [-1.4621, 1.4621]], [-0.7311, 0.7311], [[-0.7311, 0.7311], [-1.4621, 1.4621]], [-0.7311, 0.7311]]

One example. dZ = P − one_hot(0); everything else follows from it.

Input X = [[-1, -1]], y = [1], W1 = [[1, 0], [0, 1]], b1 = [0, 0], W2 = [[1, 0], [0, 1]], b2 = [0, 0]

Output [[[0, 0], [0, 0]], [0, 0], [[0, 0], [0, 0]], [0.5, -0.5]]

Both hidden units are off, so no gradient reaches W1 or b1.

Input X = [[1, 0], [0, 1]], y = [0, 1], W1 = [[2, -1], [-1, 2]], b1 = [0.5, 0.5], W2 = [[1, -1], [-1, 1]], b2 = [0, 0.1]

Output [[[-0.0074, 0], [0, -0.0061]], [-0.0074, -0.0061], [[-0.0092, 0.0092], [0.0076, -0.0076]], [-0.0007, 0.0007]]

Two examples: the 1/N averages their contributions.

Constraints

1 ≤ N ≤ 256, 1 ≤ D, H ≤ 128, 2 ≤ C ≤ 32
Checked against central finite differences to a tolerance of 10⁻⁶

Hints

Hint 1

Run the forward pass and keep A, H and P. The backward pass needs all three.

Hint 2

Start with dZ = (P − Y)/N. Then dW2 = Hᵀ dZ, db2 = dZ summed over rows, and dH = dZ W2ᵀ.

Hint 3

Through the ReLU, dA = dH where A > 0 and 0 elsewhere. Then dW1 = Xᵀ dA and db1 = dA summed over rows.

Where candidates go wrong

Masking with H > 0 instead of A > 0. They agree here, but only because ReLU(A) > 0 exactly when A > 0; for leaky ReLU they do not.
Forgetting the 1/N, so that every gradient is N times too large.
Summing db over the wrong axis and returning a vector of length N.
Transposing the wrong factor: dW2 must have W2's shape (H, C). Hᵀ dZ is the only product that does.

The idea

Work backwards through the forward pass, one line at a time, carrying the gradient of L with respect to each intermediate value.

Loss and softmax together: dZ = (P − Y)/N, where Y is the one-hot label matrix.
Z = HW2 + b2: dW2 = HᵀdZ (H × C), db2 = Σ_n dZ_n (C), dH = dZ W2ᵀ (N × H).
H = max(0, A): dA = dH ⊙ 1[A > 0].
A = XW1 + b1: dW1 = XᵀdA (D × H), db1 = Σ_n dA_n (H).

Each product is forced by the shapes, which is the fastest way to get the transposes right under pressure. The whole backward pass costs about twice the forward pass, which is the general result for reverse-mode differentiation.

Time O(N(DH + HC)) · Space O(N(H + C))

Reference solution

import numpy as np

def mlp_grads(X, y, W1, b1, W2, b2):
    X, W1, b1, W2, b2 = (np.asarray(a, dtype=float) for a in (X, W1, b1, W2, b2))
    y = np.asarray(y)
    n = X.shape[0]
    # forward, keeping what the backward pass needs
    A = X @ W1 + b1                           # (N, H) pre-activations
    H = np.maximum(0.0, A)
    Z = H @ W2 + b2                           # (N, C) logits
    Z = Z - Z.max(axis=1, keepdims=True)
    P = np.exp(Z)
    P /= P.sum(axis=1, keepdims=True)
    # backward
    dZ = P.copy()
    dZ[np.arange(n), y] -= 1.0
    dZ /= n                                   # softmax + cross-entropy: (P - Y) / N
    dW2 = H.T @ dZ                            # (H, C)
    db2 = dZ.sum(axis=0)                      # (C,)
    dH = dZ @ W2.T                            # (N, H)
    dA = dH * (A > 0)                         # ReLU passes gradient only where it was active
    dW1 = X.T @ dA                            # (D, H)
    db1 = dA.sum(axis=0)                      # (H,)
    return [dW1, db1, dW2, db2]

Console⌘/Ctrl + Enter runs

Python runs in your browser with NumPy. The first run downloads the runtime, which takes a few seconds.

Why interviewers ask it

The defining question of ML research interviews. Interviewers want each gradient stated with its shape before it is coded, and they will ask how you know your answer is right. The answer they hope for is a gradient check.

The theory behind it is in the lesson Backpropagation from scratch.

Keep going

Easy10 min

Backpropagation through a two-layer network