A neural network from scratch in NumPy

This is where the pieces meet. You write a two-layer network with nothing but NumPy: the forward pass that turns inputs into predictions, the backward pass that works out how much each weight contributed to the error, and the loop that repeats both until the thing learns. Nobody builds production models this way. The point is that after doing it once, a PyTorch model stops being magic and becomes a stack of things you have already written yourself.

Three spirals#

import numpy as np
rng = np.random.default_rng(0)

N, K = 100, 3                                    # points per class, classes
X = np.zeros((N * K, 2)); y = np.zeros(N * K, dtype=int)
for k in range(K):
    r = np.linspace(0.0, 1.0, N)
    angle = np.linspace(k * 4, (k + 1) * 4, N) + rng.normal(0, 0.2, N)
    X[k * N:(k + 1) * N] = np.c_[r * np.sin(angle), r * np.cos(angle)]
    y[k * N:(k + 1) * N] = k

That claim is checkable in two lines, and worth checking, because the whole exercise is pointless on data a straight line already handles. Fitting scikit-learn's logistic regression to these 300 points scores 0.533. The network below reaches 0.993 on the same data with no feature engineering at all; the hidden layer is doing the work that a person would otherwise do by inventing polar coordinates.

Forward#

H = 100
W1 = rng.normal(0, 0.01, (2, H)); b1 = np.zeros(H)
W2 = rng.normal(0, 0.01, (H, K)); b2 = np.zeros(K)

def forward(X):
    h = np.maximum(0, X @ W1 + b1)                      # [n, H]  hidden layer, ReLU
    logits = h @ W2 + b2                                # [n, K]
    e = np.exp(logits - logits.max(axis=1, keepdims=True))
    p = e / e.sum(axis=1, keepdims=True)                # softmax
    return h, p

def loss_of(p, y):
    return -np.log(p[np.arange(len(y)), y]).mean()      # cross-entropy

Two details in there are not decoration. Subtracting the row maximum before the exponential changes nothing mathematically, because the constant cancels in the division, but without it a logit of 800 becomes inf and the softmax returns NaN; every framework does this internally and it is why you are told to hand the loss raw logits. And the indexing in loss_of picks out the probability of the correct class for each row in one operation, which is what cross-entropy is: the negative log of the probability you assigned to the right answer.

Before training anything, check the loss at step 0. The weights are near zero, so all three classes get roughly equal probability, so the loss should be the negative log of one third, which is 1.0986. It prints 1.0986. If it had printed 7, the labels or the initialisation would be wrong, and that is a thirty-second check rather than an hour-long one.

Backward#

def backward(X, y, h, p):
    n = len(y)
    d_logits = p.copy()
    d_logits[np.arange(n), y] -= 1                      # p - y, from the calculus section
    d_logits /= n
    dW2 = h.T @ d_logits;  db2 = d_logits.sum(axis=0)
    dh = d_logits @ W2.T                                # send the error back through W2
    dh[h <= 0] = 0                                      # ReLU passes gradient only where it was active
    dW1 = X.T @ dh;        db1 = dh.sum(axis=0)
    return dW1, db1, dW2, db2

Each line is the chain rule applied to one operation, in reverse order. The shape of each gradient matches the shape of the thing it is the gradient of, which is the quickest check that a line is right.

Three of those lines carry the whole idea and are worth saying out loud. The first, that the gradient of cross-entropy with respect to the logits is simply the predicted probabilities minus the one-hot target, is the single most useful fact in this section: softmax and cross-entropy are derived together precisely so that everything cancels into a subtraction. It is also why frameworks fuse them into one operation and refuse to take a softmax you applied yourself.

The second is dh, which sends the error backwards through W2 by multiplying by its transpose. Forward, the weights map hidden units to outputs; backward, the same numbers map output errors to hidden units, and the transpose is nothing more cunning than that. The third is the ReLU line: a unit that was inactive contributed nothing to the output, so it receives no blame. That one line is where the next section's bug lives.

Check the gradients#

A hand-written backward pass is a derivation, and derivations have errors. The check does not need any calculus: nudge one parameter up by a tiny amount, nudge it down, see how much the loss actually moved, and compare that with what your formula claimed. If they disagree, the formula is wrong.

def numerical(param, eps=1e-5):
    g = np.zeros_like(param)
    for i in np.ndindex(param.shape):
        old = param[i]
        param[i] = old + eps; plus = loss_of(forward(X)[1], y)
        param[i] = old - eps; minus = loss_of(forward(X)[1], y)
        param[i] = old
        g[i] = (plus - minus) / (2 * eps)
    return g

h, p = forward(X)
analytic = backward(X, y, h, p)[2]                      # dW2
num = numerical(W2)
print(np.abs(analytic - num).max() / (np.abs(analytic).max() + 1e-12))   # 9.03e-09

Relative error rather than absolute, because a gradient of 1e-7 and a gradient of 1e7 are both legitimate and an absolute threshold would pass one and fail the other. Below about 1e-7 the derivation is right. Above 1e-2 it is wrong. In between, suspect the kink described below before suspecting yourself.

Now delete the ReLU masking line and run the same check against dW1.

dW1, ReLU mask present    relative error 5.02e-03
dW1, ReLU mask removed    relative error 5.10e-01     <- a hundred times worse

Then train both for 5,000 steps anyway:
  with the mask       final loss 0.028   accuracy 0.993
  without the mask    final loss 0.572   accuracy 0.793     <- it still learns

That is why this check exists. The broken network does not crash, does not produce NaN, and does not sit at chance; it trains to 79 per cent, which on an unfamiliar dataset is entirely plausible and would be shipped. Nothing but the gradient check distinguishes it from the correct one.

Train#

No optimiser, no momentum, no schedule: subtract the learning rate times the gradient from every parameter, five thousand times. A learning rate of 1.0 would be wild for a deep network and is fine here because the network is two layers on 300 points, and there is no mini-batching either, since the whole dataset fits in one matrix.

lr = 1.0
for step in range(5001):
    h, p = forward(X)
    grads = backward(X, y, h, p)
    for param, g in zip((W1, b1, W2, b2), grads):
        param -= lr * g
    if step % 1000 == 0:
        print(step, round(loss_of(p, y), 3), round((p.argmax(1) == y).mean(), 3))

# step     loss    accuracy
#    0    1.099      0.463      log(3): the model knows nothing
# 1000    0.134      0.947
# 2000    0.070      0.983
# 3000    0.047      0.987
# 4000    0.035      0.993
# 5000    0.028      0.993      logistic regression on the same data: 0.533

The curve is the ordinary shape: most of the progress in the first thousand steps, then a long tail buying decreasing amounts. Accuracy stops at 0.993 while the loss keeps falling, which is the model becoming more confident about points it already had right rather than fixing the two it has wrong. Those two sit where the spiral arms meet and are genuinely ambiguous.