Lessons From Implementing a Neural Network From Scratch

I’d worked through the math of backpropagation on paper — the chain rule applied layer by layer, nothing conceptually mysterious about it — and assumed that meant implementing it would be mostly transcription. It wasn’t. The gap between understanding the derivation and getting a working implementation turned out to be full of small, specific mistakes that the math itself doesn’t warn you about, and most of what I actually learned came from finding and fixing those mistakes, not from the derivation.

The network

Nothing exotic: two layers, a ReLU nonlinearity in the hidden layer, and mean squared error as the loss, trained on a small synthetic regression problem. The whole point was keeping the architecture boring enough that any bugs would have to be in my implementation of the mechanics, not hidden in some more complicated piece of machinery.

import numpy as np

def init_params(input_dim, hidden_dim, output_dim, rng):
    return {
        "W1": rng.normal(0, 0.1, size=(input_dim, hidden_dim)),
        "b1": np.zeros(hidden_dim),
        "W2": rng.normal(0, 0.1, size=(hidden_dim, output_dim)),
        "b2": np.zeros(output_dim),
    }

def forward(params, X):
    z1 = X @ params["W1"] + params["b1"]
    a1 = np.maximum(z1, 0)  # ReLU
    z2 = a1 @ params["W2"] + params["b2"]
    cache = {"X": X, "z1": z1, "a1": a1}
    return z2, cache

def loss_fn(pred, y):
    return np.mean((pred - y) ** 2)

The forward pass is straightforward enough that it’s hard to get wrong. The backward pass is where things fell apart the first several times.

Where the first bugs actually were

My first implementation of the backward pass produced a network that trained, in the sense that the loss went down — which is exactly what makes this category of bug dangerous. A completely broken backward pass usually produces a loss that doesn’t move at all, or explodes to NaN, both of which are loud enough to notice immediately. A subtly wrong backward pass often still reduces the loss, just less effectively than the correct gradient would, and there’s no obvious signal telling you the numbers you’re computing are wrong rather than merely slow to converge.

def backward(params, cache, pred, y):
    m = cache["X"].shape[0]
    dz2 = 2 * (pred - y) / m  # dL/dz2

    dW2 = cache["a1"].T @ dz2
    db2 = dz2.sum(axis=0)

    da1 = dz2 @ params["W2"].T
    dz1 = da1 * (cache["z1"] > 0)  # ReLU derivative

    dW1 = cache["X"].T @ dz1
    db1 = dz1.sum(axis=0)

    return {"W1": dW1, "b1": db1, "W2": dW2, "b2": db2}

That’s the corrected version. What I actually had at first differed in two small ways, each of which cost me real time to find.

The ReLU derivative masked the wrong tensor. My first pass computed the ReLU gradient mask from a1 (the post-activation values) instead of z1 (the pre-activation values). For a ReLU, a1 > 0 and z1 > 0 are actually equivalent conditions, so this particular mistake happened to work — but only because I’d chosen the one activation function where pre- and post-activation signs coincide. It was a bug that got lucky, not a bug that was absent, and it would have silently broken the moment I switched to an activation function like tanh where that coincidence doesn’t hold.

I forgot to average the gradient over the batch in one place but not the others. I’d divided by m when computing dz2, correctly, but then also divided by m a second time inside the weight gradient computation, effectively dividing the gradient for W2 by the batch size twice. The network still trained — a gradient that’s systematically too small in one layer just makes that layer’s parameters update more slowly, which looks like slightly worse convergence, not like an obvious bug. I only found this by comparing the relative magnitudes of the gradients for W1 versus W2 and noticing they were oddly imbalanced given how the two layers were being used.

Neither of these produced an error message. Both produced a network that trained, just worse than it should have. That’s the specific danger of backpropagation bugs: the failure mode isn’t “broken,” it’s “quietly suboptimal,” and quietly suboptimal is much harder to notice than a crash.

The tool that actually found these: gradient checking

What eventually caught both bugs wasn’t careful re-reading of the derivation — I’d done that several times and it looked fine each time, because the bugs weren’t conceptual, they were implementation slips that a mental re-derivation doesn’t naturally catch. What caught them was gradient checking: numerically approximating the gradient at each parameter using finite differences, and comparing that numerical estimate against what the backward pass computed analytically.

def numerical_gradient(params, X, y, key, idx, eps=1e-5):
    original = params[key][idx]

    params[key][idx] = original + eps
    pred_plus, _ = forward(params, X)
    loss_plus = loss_fn(pred_plus, y)

    params[key][idx] = original - eps
    pred_minus, _ = forward(params, X)
    loss_minus = loss_fn(pred_minus, y)

    params[key][idx] = original  # restore
    return (loss_plus - loss_minus) / (2 * eps)

def check_gradients(params, X, y, n_checks=20, rng=None):
    pred, cache = forward(params, X)
    analytic = backward(params, cache, pred, y)

    rng = rng or np.random.default_rng()
    max_rel_error = 0.0
    for key in params:
        flat_shape = params[key].shape
        for _ in range(n_checks):
            idx = tuple(rng.integers(0, dim) for dim in flat_shape)
            numeric = numerical_gradient(params, X, y, key, idx)
            analytic_val = analytic[key][idx]
            rel_error = abs(numeric - analytic_val) / max(abs(numeric) + abs(analytic_val), 1e-8)
            max_rel_error = max(max_rel_error, rel_error)

    return max_rel_error

The idea is almost embarrassingly simple: the definition of a derivative already gives you a way to check any gradient computation, independent of how the analytic gradient was derived. Nudge a single parameter by a tiny amount in each direction, see how much the loss moves, and divide — that’s a numerical approximation of the same partial derivative your backward pass is supposed to be computing analytically. If the two disagree by more than a small numerical tolerance, the backward pass has a bug, full stop, regardless of how plausible the code looks or how well the network still seems to train.

Both of my bugs showed up immediately as large relative errors on specific parameters the moment I ran this check — the W2 gradient consistently off by close to a factor of two, exactly matching the double-division bug once I went looking for it. Neither bug was visible from watching the loss curve. Both were glaringly obvious the moment I had a second, independently-computed answer to compare against.

What this changed about how I trust ML code generally

The broader lesson wasn’t really about backpropagation specifically. It was that a training loss going down is much weaker evidence of correctness than it feels like in the moment, because gradient descent is remarkably good at making progress even with a moderately wrong gradient — it’s a stochastic, iterative, self-correcting process, which is exactly the property that makes it robust to noise and exactly the property that makes it dangerously tolerant of certain classes of bugs.

I now treat “the loss is decreasing” as necessary but nowhere near sufficient evidence that a custom training implementation is correct, and I reach for some form of independent verification — gradient checking, comparing against a known-correct reference implementation on a toy problem, overfitting a tiny dataset and confirming the loss actually reaches close to zero — any time I’m implementing the mechanics of training rather than just calling a library that’s already had this verification done for me. It’s a small amount of extra work that would have saved me several hours of confused, inconclusive staring at a loss curve that was technically doing its job, just not as well as it should have been.