Backpropagation for a One-Hidden-Layer MLP

By Better Wise

Why a one-hidden-layer net breaks the closed-form habit

Two things you already own do most of the work here: the finite-difference harness you built for gradient checking, and the softmax cross-entropy gradient $\partial L/\partial Z = (P - Y)/N$. A one-hidden-layer MLP is those two ideas stacked with a ReLU in between. What is genuinely new is bookkeeping: the chain rule now runs backward through a nonlinearity and two matrix transposes, and every one of those steps is a place to ship a silent transpose bug.

By the end of this lesson you should be able to:

• write down the shape of every intermediate before opening an editor
• derive dW1, db1, dW2, db2 by hand from the loss
• implement the forward and backward pass in roughly thirty lines of NumPy
• flatten the whole parameter vector into your existing finite-difference check
• read a gradient-check report and name the bug it points at

Draw the shapes before you touch NumPy

Fixed notation, because shape confusion is 90% of the debugging hours in this module. Batch size $N$, input dimension $D$, hidden width $H$, classes $K$.

Parameter count: $DH + H + HK + K$. The invariant that catches almost everything: every gradient must have exactly the shape of the parameter it updates. If dW1 comes out (H, D), you transposed, and the optimizer will either crash or silently broadcast garbage.

Note: Convention for the whole module: rows are examples, X @ W maps features forward. Pick it once and never mix it with the W @ x convention from a tutorial, or you will spend an evening on transposes.

Exercise 1. For the capstone setup $D = 64$, $H = 32$, $K = 10$, $N = 128$: write the shape of H1, A1, Z2, P, and the total parameter count.

Answer

H1 and A1 are (128, 32). Z2 and P are (128, 10). Parameters: $64 \cdot 32 + 32 + 32 \cdot 10 + 10 = 2410$.

Forward pass: two affines, one ReLU, one softmax

$$ H_1 = XW_1 + b_1,\quad A_1 = \max(H_1, 0),\quad Z_2 = A_1W_2 + b_2,\quad P = \mathrm{softmax}(Z_2) $$

Keep the logits and do the softmax inside a log-sum-exp that subtracts the row maximum. Subtracting the max does not change the probabilities, and it is the difference between nan and a working model on your first real run.

def forward(params, X):
    H1 = X @ params["W1"] + params["b1"]        # (N, H)
    A1 = np.maximum(H1, 0.0)                    # (N, H)
    Z2 = A1 @ params["W2"] + params["b2"]       # (N, K)
    return H1, A1, Z2

def softmax_xent(Z2, y):
    Z2 = Z2 - Z2.max(axis=1, keepdims=True)    # stability shift
    logP = Z2 - np.log(np.exp(Z2).sum(axis=1, keepdims=True))
    N = Z2.shape[0]
    loss = -logP[np.arange(N), y].mean()
    return loss, np.exp(logP)

Exercise 2. At initialization with balanced classes ($K = 10$), your forward pass reports a loss of 2.3026. Pass or fail, and why does that number look familiar?

Answer

Pass. $\ln 10 \approx 2.3026$ is exactly the loss of a uniform prediction, which is what a freshly initialized network with zeroed biases should produce. A loss of 0.31 or 47.0 at step zero means the forward pass or the label indexing is broken, not that the model is bad.

The backward pass on paper

Derive it once by hand, then code it. The loss is the mean cross-entropy, so the batch factor $1/N$ belongs in dZ2; carrying it there keeps every downstream expression clean and consistent.

$$ \frac{\partial L}{\partial Z_2} = \frac{1}{N}(P - Y), \qquad \frac{\partial L}{\partial W_2} = A_1^\top \frac{\partial L}{\partial Z_2}, \qquad \frac{\partial L}{\partial H_1} = \left(\frac{\partial L}{\partial Z_2} W_2^\top\right) \odot \mathbb{1}[H_1 > 0] $$

Then dW1 = X.T @ dH1 and both bias gradients are column sums. Three rules that make the derivation self-checking:

• Walk the shapes backward: (N, K) to (N, H) to (D, H). Each step consumes one stored tensor.
• A matrix that appears transposed in the forward pass appears transposed again on the way back: $W_2$ multiplies on the right forward, on the right-transposed backward.
• The ReLU mask is built from the pre-activation H1 > 0. Using A1 > 0 happens to agree almost everywhere, but H1 is what the derivative actually is.

def backward(params, X, y, H1, A1, P):
    N = X.shape[0]
    dZ2 = P.copy()
    dZ2[np.arange(N), y] -= 1.0
    dZ2 /= N                                    # (N, K)
    grads = {"W2": A1.T @ dZ2, "b2": dZ2.sum(axis=0)}
    dH1 = (dZ2 @ params["W2"].T) * (H1 > 0.0)    # (N, H)
    grads["W1"] = X.T @ dH1                     # (D, H)
    grads["b1"] = dH1.sum(axis=0)               # (H,)
    return grads

Exercise 3. Row sums of dZ2 should be exactly zero for unregularized mean cross-entropy. Prove it in one line and turn it into an assertion.

Answer

Each row of $P$ sums to 1 and each row of the one-hot $Y$ sums to 1, so each row of $P - Y$ sums to 0. Assert np.allclose(dZ2.sum(axis=1), 0) in the backward pass. It catches label-indexing bugs and double-counted samples immediately.

Exercise 4. A hidden unit has H1[:, j] < 0 for all 128 examples in the batch. What are dW1[:, j] and db1[j]?

Answer

Both are exactly zero: the mask zeroes its column of dH1, so neither the incoming weights nor the bias receive any signal. The unit is dead. It cannot recover on its own, so check the fraction of dead units after a few epochs; too many means your W1 initialization scale is wrong (use $\sqrt{2/D}$ for ReLU, as in He initialization).

Gradient-checking the whole parameter vector

Your harness takes one flat vector, so flatten the parameters, perturb one coordinate at a time with the central difference, and compare against the flattened analytic gradient.

def flat(params):  return np.concatenate([params[k].ravel() for k in params])
def unflat(v, shapes):  # restore the dict from a flat vector
    out, i = {}, 0
    for k, s in shapes.items():
        n = int(np.prod(s)); out[k] = v[i:i+n].reshape(s); i += n
    return out

def rel_err(a, g):  return abs(a - g) / max(1e-8, abs(a) + abs(g))

Check on a tiny model, N = 5, D = 4, H = 6, K = 3, in float64, with h = 1e-5 and no regularizer unless you also differentiate it. A relative error around $10^{-7}$ or below is a pass; $10^{-4}$ means a real bug, usually a missing $1/N$ or a stale forward pass.

Exercise 5. The check reports $4\times10^{-9}$ for W1 and b1 but $2\times10^{-3}$ for W2 and b2. Where is the bug?

Answer

It is localized to the logits branch, downstream of A1. Typical causes: dZ2 divided by N twice or not at all, b2 summed over the wrong axis, or the numerical check calling a forward pass that applies an L2 penalty the analytic gradient ignores. When one branch passes and another fails, stop reading the code top to bottom and bisect along the chain.

Traps that cost you a day of debugging

• Derive before you type. If you cannot write dH1 on paper from the forward equations, no amount of print-statement debugging will save the implementation.
• Gradient-check in float64. In float32, finite differences are dominated by rounding error and the check lies to you.
• Keep the logits. Never feed probabilities into the loss; log-softmax from raw logits is both correct and stable.
• Do not re-run the forward pass inside backward. Pass H1, A1, P in, so the mask and probabilities are guaranteed to come from the same forward call.
• Regularize explicitly. If you add $\lambda \lVert W \rVert^2$ to the loss, add $2\lambda W$ to the gradient, or the check will flag a "bug" that is really a design decision.
• If only one parameter's check fails, suspect a transpose; if all fail together, suspect the loss scaling.

What the gradient check buys you

You now have a forward pass, a hand-derived backward pass, and a check that either passes or points at a branch. Those three artifacts are the entire correctness argument for the model you are about to benchmark, and they are worth more than any accuracy number. The next lesson, on cross-validation and a fair baseline comparison protocol, takes a model that is provably correct and asks the harder question: is it actually better than the library baseline, on splits it has never seen?

Front: X is (N, D), W1 is (D, H), W2 is (H, K). Shapes of H1, A1, Z2, P? Back: H1 and A1 are (N, H); Z2 and P are (N, K).

Front: Parameter count of a one-hidden-layer MLP in terms of D, H, K? Back: DH + H + HK + K.

Front: Gradient of mean cross-entropy with respect to the logits? Back: (P - Y) / N, where P is the softmax output and Y is one-hot.

Front: Formula and shape of dW2? Back: A1.T @ dZ2, shape (H, K).

Front: Formula for db1? Back: dH1.sum(axis=0), shape (H,).

Front: How does the gradient cross the ReLU? Back: Elementwise multiply by the mask (H1 > 0), built from the pre-activation, not the output.

Front: Why subtract the row max before exponentiating logits? Back: It prevents overflow in exp and leaves the probabilities unchanged.

Front: What loss should a K-class model report at initialization? Back: About ln(K), the uniform-prediction loss.

Front: Which dtype and step size for finite-difference checking? Back: float64 with a central difference and h = 1e-5.

Front: Relative error formula used by the gradient checker? Back: |analytic - numeric| / max(1e-8, |analytic| + |numeric|).

Front: Why should the row sums of dZ2 be zero? Back: Both P and the one-hot Y have rows summing to 1, so P - Y does.

Front: What happens to a ReLU unit whose pre-activation is negative for every example? Back: It gets zero gradient forever; it is dead and needs better initialization to avoid.

Front: Check passes for W1 and b1 but fails for W2 and b2. What does that localize? Back: The bug is in the logits branch: dZ2 scaling, b2 axis, or a loss term the analytic gradient ignores.

Front: What must hold between a gradient and its parameter? Back: Identical shapes; a mismatch means a transpose bug.

With logits Z2 and one-hot targets Y, the gradient of mean cross-entropy with respect to the logits is _____, where P is the softmax output. The backward pass through ReLU multiplies the incoming gradient by the mask _____, built from the pre-activation. For W1 of shape (D, H), the gradient is _____, which has the same shape as W1. Gradient checking with a central difference uses the estimate _____. A ReLU unit with negative pre-activation on every example receives _____ gradient, so no update ever revives it.