Backpropagation, by Hand: The Chain Rule, Run Backward

The router of Learning Is Optimization and Measuring a Classifier Honestly is a linear model: a weighted sum of features, passed through a sigmoid. Features and Overfitting showed where that runs out. It can only weigh the features it’s given, it can’t combine them, and on our data one hand-picked feature matched 379 learned ones.

The way out is to let the model build its own features: put a layer of computed quantities between the input and the output, each a nonlinear function of the inputs, and learn what they should compute. That’s a neural network. Training one needs the same thing Learning Is Optimization needed, the gradient of the loss with respect to every weight. The difference is that the weights are now buried under layers of other computation. Backpropagation is how we get those gradients, all of them, in one backward pass. This post builds it from the chain rule. First the smallest problem no straight line can solve, then one step of training on it with every number shown, then the matrix rules for a whole layer, a numerical check that they are right, and the cost of all this in FLOPs. The last part is where the rule that training a model costs about six times its parameter count in FLOPs per token comes from. Cross-Entropy and the Softmax Gradient then replaces the sigmoid with a softmax and routes to all four parts of the blog.

Table of contents

What one neuron computes

A neuron takes a vector $x$, computes a weighted sum plus a bias, and passes the result through a fixed nonlinear function $\phi$, the activation. With $x$ a row vector, as everywhere in the series:

Logistic regression is exactly one neuron with $\phi = \sigma$. The bias is easy to overlook but does a specific job. $x w$ alone is always zero at the origin, so without $b$ every neuron’s decision boundary, the set where $z = 0$, would have to pass through the origin. With it, the boundary $x w = -b$ can sit anywhere. Read $-b$ as a threshold: the neuron “fires” when the weighted evidence exceeds it. In the router from Learning Is Optimization, $b = -1.46$ was what made one serving word necessary before the probability passed 0.5.

The activation used in hidden layers today is usually the ReLU, $\phi(z) = \max(0, z)$. It passes positive values unchanged and blocks negative ones. It’s cheap, and its gradient is exactly 1 wherever it passes, which matters later when gradients have to travel through many layers.

Why a nonlinearity at all? Two linear layers stacked without one compute $(xW_1)W_2 = x(W_1 W_2)$, a single linear layer with the weight matrix $W_1 W_2$. Any number of stacked linear layers computes a function that one linear layer could compute. Depth only adds power when something nonlinear sits between the layers.

The smallest problem a line can’t solve

The classic example is XOR: output 1 when exactly one of two binary inputs is 1.

No single neuron can compute it. Suppose one did, with logit $z = w_1 x_1 + w_2 x_2 + b$, positive exactly on the two middle rows. The rows give four inequalities:

• $(0,0) \mapsto 0$ needs $b < 0$;
• $(0,1) \mapsto 1$ needs $w_2 + b > 0$;
• $(1,0) \mapsto 1$ needs $w_1 + b > 0$;
• $(1,1) \mapsto 0$ needs $w_1 + w_2 + b < 0$.

Adding the middle two gives $w_1 + w_2 + 2b > 0$, so $w_1 + w_2 + b > -b > 0$, which contradicts the last one. No line separates the two classes. Minsky and Papert’s 1969 Perceptrons made limits like this precise. They knew that networks with more layers could compute such functions. What nobody had then was a practical way to train the layers in between, and that’s what this post is about.

One hidden layer fixes it. Here is a solution by hand, with two ReLU units:

• At $(0,0)$: $h = (0, 0)$, so $z = -1$.
• At $(0,1)$ or $(1,0)$: $h = (1, 0)$, so $z = 1$.
• At $(1,1)$: $h = (2, 1)$, so $z = 4 - 4 - 1 = -1$.

The signs are exactly XOR. The hidden units built two new features, “at least one input is on” and “both are on”, and in those features the problem is linear. That’s the whole idea of a hidden layer, and it’s the same move a router would need to learn that “memory” plus “training” means something different from “memory” plus “serving”.

A network can represent XOR. Whether gradient descent can find those weights is a separate question, and we’ll see it has a surprising answer.

The chain rule, run backward

Write the two-layer network out as a chain of simple steps. For one input $x$ with target $y$:

(In the notation of AI Engineering, from the Inside, $z$ is a logit. Here $z_2$ is the logit and $z_1$, the hidden layer’s pre-activation, is local to this post.)

Each step is a function whose derivative we know. We want $\partial\mathcal{L}/\partial W_1$, but $W_1$ affects $\mathcal{L}$ only through $z_1$, which acts through $a_1$, then $z_2$, then $p$. The chain rule multiplies the local derivatives along that path:

This is schematic. Each factor is really a matrix of partial derivatives (a Jacobian), and with row vectors the matrix products run in the other order, which the rules below get right. The structure is what matters here.

The key observation is that the first factors are shared. Every parameter’s gradient starts with $\partial\mathcal{L}/\partial z_2$. Every first-layer parameter’s gradient continues through $\partial z_2/\partial a_1$ and $\partial a_1/\partial z_1$. So compute the derivatives from the loss backward, one step at a time, and reuse each result for everything below it. That is all backpropagation is: the chain rule, evaluated from the output toward the input, caching each intermediate derivative. Call $\delta_v = \partial\mathcal{L}/\partial v$ the gradient arriving at quantity $v$. Each step then applies one local rule:

• Loss and sigmoid, together: $\delta_{z_2} = p - y$. That’s the result in Learning Is Optimization: the sigmoid’s derivative cancels against the log-loss’s.
• Through a weighted sum $z_2 = a_1 w_2 + b_2$: $\delta_{w_2} = \delta_{z_2}\, a_1$, $\delta_{b_2} = \delta_{z_2}$, and $\delta_{a_1} = \delta_{z_2}\, w_2$. The weight’s gradient is the arriving gradient times the weight’s input. The input’s gradient is the arriving gradient times the weight.
• Through ReLU: $\delta_{z_1} = \delta_{a_1} \odot \mathbb{1}[z_1 > 0]$. The gradient passes where the unit was active and is blocked where it wasn’t. ReLU has no derivative at exactly 0; the code uses 0 there, and the choice rarely matters because $z_1$ is almost never exactly 0. (One place it is: the input $(0, 0)$ with zero biases, which comes up below.)
• Through the first weighted sum: $\delta_{W_1} = x^\top \delta_{z_1}$ and $\delta_{b_1} = \delta_{z_1}$.

One step, every number

Take the XOR network with two hidden units and these starting weights. The rows of $W_1$ belong to the inputs, and the columns to the hidden units:

Train on one example: $x = (1, 0)$, target $y = 1$.

Forward.

• $z_1 = (1 \cdot 0.5 + 0 \cdot 0.3 + 0,\ \ 1 \cdot (-0.4) + 0 \cdot 0.8 + 0.1) = (0.5,\ -0.3)$.
• $a_1 = \operatorname{ReLU}(z_1) = (0.5,\ 0)$. The second hidden unit is off.
• $z_2 = 0.5 \times 0.7 + 0 \times (-0.6) + 0.1 = 0.45$.
• $p = \sigma(0.45) = 0.6106$.
• $\mathcal{L} = -\log 0.6106 = 0.4932$.

Backward, one local rule per step:

• $\delta_{z_2} = p - y = 0.6106 - 1 = -0.3894$. The prediction is too low, so the gradient says “raise $z_2$”.
• $\delta_{w_2} = \delta_{z_2}\, a_1 = -0.3894 \times (0.5,\ 0) = (-0.1947,\ 0)$, and $\delta_{b_2} = -0.3894$. The second output weight gets no gradient, because the unit feeding it was off.
• $\delta_{a_1} = \delta_{z_2}\, w_2 = -0.3894 \times (0.7,\ -0.6) = (-0.2726,\ 0.2336)$.
• $\delta_{z_1} = \delta_{a_1} \odot \mathbb{1}[z_1 > 0] = (-0.2726,\ 0)$. The ReLU blocks the second unit’s gradient, since $z_1$ was negative there.
• $\delta_{W_1} = x^\top \delta_{z_1} = \begin{pmatrix} -0.2726 & 0 \\ 0 & 0 \end{pmatrix}$, and $\delta_{b_1} = (-0.2726,\ 0)$. Only the weight from input 1 to hidden unit 1 moves, because it’s the only path that was active end to end.

Update, with $\eta = 0.5$: every parameter moves by $-0.5\,\delta$. On the same example, $p$ rises from 0.6106 to 0.7131, and the loss falls from 0.4932 to 0.3381.

The pattern in those numbers generalizes. In a ReLU network, a weight gets gradient from an example only if the input feeding it was nonzero and at least one path from it to the output ran through active units. (A deeper network has many paths, and the gradient sums over them.) The second hidden unit, off for this example, learned nothing from it. With smooth activations like the GELU of later posts, the gradient through an “off” unit is small rather than exactly zero, and an L2 penalty still reaches every weight.

Step through the same example below, one rule at a time. Then apply the update and keep going, or switch to another input. On $(0, 0)$ both inputs are zero, so no first-layer weight gets a gradient (only the biases can). On $(1, 1)$ both hidden units are on and every edge gets one.

import numpy as np

sigmoid = lambda z: 1 / (1 + np.exp(-z))
relu = lambda z: np.maximum(z, 0.0)

def init(n_in, n_hidden, rng):
    """He initialization for the ReLU layer; small random output weights."""
    return {"W1": rng.normal(0, np.sqrt(2 / n_in), (n_in, n_hidden)),
            "b1": np.zeros(n_hidden),
            "w2": rng.normal(0, np.sqrt(1 / n_hidden), n_hidden),
            "b2": 0.0}

def forward(P, X):
    z1 = X @ P["W1"] + P["b1"]            # (n, hidden)
    a1 = relu(z1)
    z2 = a1 @ P["w2"] + P["b2"]           # (n,)  one logit per example
    return z1, a1, z2

def backward(P, X, y, lam=0.0):
    """Gradients of the mean log-loss (+ L2 on weights) for every parameter."""
    n = len(y)
    z1, a1, z2 = forward(P, X)
    p = sigmoid(z2)
    dz2 = (p - y) / n                     # the p - y of "Learning Is Optimization", per example
    grads = {"w2": a1.T @ dz2 + lam * P["w2"],
             "b2": dz2.sum()}
    da1 = np.outer(dz2, P["w2"])          # back through the output weights
    dz1 = da1 * (z1 > 0)                  # back through ReLU: pass or block
    grads["W1"] = X.T @ dz1 + lam * P["W1"]
    grads["b1"] = dz1.sum(axis=0)
    p = np.clip(p, 1e-12, 1 - 1e-12)
    loss = -(y * np.log(p) + (1 - y) * np.log(1 - p)).mean()
    loss += 0.5 * lam * ((P["W1"] ** 2).sum() + (P["w2"] ** 2).sum())
    return grads, loss

def train(X, y, n_hidden, lam=0.0, steps=3000, eta=0.5, seed=0):
    P = init(X.shape[1], n_hidden, np.random.default_rng(seed))
    for _ in range(steps):
        grads, _ = backward(P, X, y, lam)
        for k in grads:
            P[k] = P[k] - eta * grads[k]
    return P

Full script: mlp.py

The backward function is the four local rules above, written for a whole batch at once. The next sections explain the matrix form. xor.py reproduces the hand calculation, and asks how often training finds XOR at all:

import numpy as np
from mlp import sigmoid, relu, forward, backward, train

X = np.array([[0, 0], [0, 1], [1, 0], [1, 1]], dtype=float)
y = np.array([0, 1, 1, 0], dtype=float)          # XOR

# One example, one step, every number.
P = {"W1": np.array([[0.5, -0.4], [0.3, 0.8]]), "b1": np.array([0.0, 0.1]),
     "w2": np.array([0.7, -0.6]), "b2": 0.1}
x, t = X[2:3], y[2:3]                             # x = (1, 0), target 1
z1, a1, z2 = forward(P, x)
grads, loss = backward(P, x, t)
print(f"z1 {z1[0]}  a1 {a1[0]}  z2 {z2[0]:.4f}  p {sigmoid(z2[0]):.4f}  loss {loss:.4f}")
for k in ["w2", "b2", "W1", "b1"]:
    print(f"dL/d{k} = {np.round(grads[k], 4).tolist()}")
P = {k: v - 0.5 * grads[k] for k, v in P.items()}
print(f"after one step: p {sigmoid(forward(P, x)[2][0]):.4f}  loss {backward(P, x, t)[1]:.4f}")

# How often does training find XOR, by hidden-layer width?
for h in [2, 3, 4, 8, 16]:
    solved = 0
    for seed in range(100):
        P = train(X, y, h, steps=5000, eta=1.0, seed=seed)
        solved += ((sigmoid(forward(P, X)[2]) >= 0.5) == y).all()
    print(f"{h:2d} hidden units: XOR solved from {solved}/100 random starts")

Full script: xor.py

z1 [ 0.5 -0.3]  a1 [0.5 0. ]  z2 0.4500  p 0.6106  loss 0.4932
dL/dw2 = [-0.1947, 0.0]
dL/db2 = -0.3894
dL/dW1 = [[-0.2726, 0.0], [0.0, 0.0]]
dL/db1 = [-0.2726, 0.0]
after one step: p 0.7131  loss 0.3381
 2 hidden units: XOR solved from 18/100 random starts
 3 hidden units: XOR solved from 52/100 random starts
 4 hidden units: XOR solved from 73/100 random starts
 8 hidden units: XOR solved from 96/100 random starts
16 hidden units: XOR solved from 100/100 random starts

Can training find it?

The last five lines are the surprise. Two hidden units are enough to represent XOR, since the solution by hand above uses exactly two. But gradient descent from a random start finds a solution only 18 times in 100 (give or take about 4, from 100 starts, at this fixed $\eta = 1$ and 5,000 steps). What happens to the other 82?

import numpy as np
from collections import Counter
from mlp import sigmoid, forward, backward, init, train

X = np.array([[0, 0], [0, 1], [1, 0], [1, 1]], dtype=float)
y = np.array([0, 1, 1, 0], dtype=float)

# Why do 82 of 100 two-unit starts fail? Classify every failure.
outcome, dead_start, solved_if_dead = Counter(), 0, 0
for seed in range(100):
    W1 = init(2, 2, np.random.default_rng(seed))["W1"]
    dead = ((W1 <= 0).all(axis=0)).any()          # a unit with both weights <= 0 never fires
    dead_start += dead
    P = train(X, y, 2, steps=5000, eta=1.0, seed=seed)
    right = int(((sigmoid(forward(P, X)[2]) >= 0.5) == y).sum())
    solved_if_dead += dead and right == 4
    if right < 4:
        grads, loss = backward(P, X, y)
        size = np.sqrt(sum((g ** 2).sum() for g in grads.values()))
        outcome[(right, round(loss, 3), "dead unit at start" if dead else "both units alive",
                 "gradient < 1e-3" if size < 1e-3 else "gradient >= 1e-3")] += 1
print(f"starts with a unit dead from step 0: {dead_start}/100, of which solved: {solved_if_dead}")
for (right, loss, dead, grad), n in sorted(outcome.items(), key=lambda kv: -kv[1]):
    print(f"{n:3d} failures: {right}/4 right, loss {loss:.3f}, {dead}, {grad}")

Full script: xor_failures.py

starts with a unit dead from step 0: 42/100, of which solved: 0
 27 failures: 3/4 right, loss 0.477, dead unit at start, gradient < 1e-3
 22 failures: 3/4 right, loss 0.347, both units alive, gradient < 1e-3
 15 failures: 2/4 right, loss 0.693, dead unit at start, gradient < 1e-3
 14 failures: 3/4 right, loss 0.477, both units alive, gradient < 1e-3
  2 failures: 2/4 right, loss 0.693, both units alive, gradient < 1e-3
  1 failures: 3/4 right, loss 0.347, both units alive, gradient >= 1e-3
  1 failures: 3/4 right, loss 0.693, both units alive, gradient < 1e-3

• They’re stuck, not slow. In 81 of the 82 failures the gradient has shrunk below $10^{-3}$: training has stopped at a point it can’t leave. Most get three of the four points right; the rest sit near $\log 2 = 0.693$, predicting about 0.5.
• Units dead from the first step are the biggest cause. The biases start at zero and the inputs are 0 or 1, so a hidden unit whose two incoming weights are both negative has $z_1 \le 0$ on every input. It never fires, so by the ReLU rule it never gets a gradient, and it never recovers. That happens to each unit with probability 1/4, so about $1 - (3/4)^2 = 44\%$ of two-unit starts have one. Here 42 did, and not one of them solved XOR: one unit alone can’t.
• The other 40 failures have both units alive, and 37 of them fit three of the four points (loss 0.347 or 0.477); all but one of those stall on a flat region of the loss where the gradient toward the fourth has vanished. The other 3 sit near $\log 2$.

Add width and the problem gets easier: 52 of 100 with three units, 73 with four, 96 with eight, and all 100 with sixteen. More units means more random starting directions, so a dead unit costs less and some units are likely to start well placed. In small networks like this one, extra width makes the loss much easier to descend, a real and studied effect. It isn’t why language models are large, though. Their size is set by how the loss improves with parameters and data, the scaling laws of How LLM Architecture Evolved, 2017–2026.

The start must also be random. Two hidden units with identical incoming weights, bias and outgoing weight compute the same thing and, by the rules above, receive identical gradients, so they stay identical forever and act as one unit. Starting every weight at zero is worse still: with $w_2 = 0$ no gradient reaches the first layer, with $a_1 = 0$ none reaches $w_2$, and on XOR, whose labels are balanced, not even the output bias moves. The logistic regression from Learning Is Optimization could start at zero because it had no hidden units to tell apart.

Backprop for a whole layer

The worked example used one input. Training uses a batch: the inputs $X$ are an $n \times d$ matrix, one example per row. A layer computes $Z = XW + b$ with $W$ of shape $d \times k$. Given the gradient $\delta_Z$ arriving from above ($n \times k$, one row per example), the three gradients the layer must produce follow from the same local rules, now summed over examples:

Checking the shapes is the fastest way to remember them. $\delta_W$ must have $W$‘s shape $d \times k$, and $X^\top \delta_Z$ is $(d \times n)(n \times k)$. The sum over the $n$ inside that product is the sum of each example’s contribution. $\delta_X$ is passed further down, to the layer below, as its arriving gradient. An elementwise activation contributes a mask or a scale: $\delta_Z = \delta_A \odot \phi'(Z)$.

Those three lines, plus one per activation, are all of backpropagation for a feed-forward network. Every framework’s automatic differentiation is a general engine for applying rules like these to any graph of operations. PyTorch from the Inside builds a small one.

Checking the gradients

Backprop code is easy to get subtly wrong: a missing transpose, a sum over the wrong axis, a forgotten $1/n$. A wrong gradient often still makes the loss go down, just worse, so the bug can hide for a long time. The defense is to compare against the definition of a derivative. Nudge one parameter by $\pm\varepsilon$ and measure the change in the loss:

This central difference has error proportional to $\varepsilon^2$, where the one-sided version has error proportional to $\varepsilon$. With $\varepsilon = 10^{-6}$ in double precision, a correct gradient agrees to roughly seven or more significant digits.

import numpy as np
from mlp import init, backward

# Compare backprop to finite differences on a small random problem.
rng = np.random.default_rng(1)
X, y = rng.normal(size=(5, 4)), rng.integers(0, 2, 5).astype(float)
P = init(4, 3, rng)
grads, _ = backward(P, X, y, lam=0.01)
eps, worst = 1e-6, 0.0
for k in ["W1", "b1", "w2"]:
    for i in np.ndindex(P[k].shape):
        old = P[k][i]
        P[k][i] = old + eps; up = backward(P, X, y, 0.01)[1]
        P[k][i] = old - eps; down = backward(P, X, y, 0.01)[1]
        P[k][i] = old
        numeric = (up - down) / (2 * eps)                 # central difference
        worst = max(worst, abs(numeric - grads[k][i]) / max(1e-12, abs(numeric) + abs(grads[k][i])))
print(f"largest relative error, backprop vs finite differences: {worst:.1e}")

Full script: grad_check.py

largest relative error, backprop vs finite differences: 9.7e-08

A relative error around $10^{-7}$ means the gradients are right. A bug typically shows up as $10^{-2}$ or worse. (The loop skips b2 only because it’s a plain number in this code rather than an array; its rule is the same as b1’s. And finite differences are unreliable within $\varepsilon$ of a ReLU’s kink, which random inputs like these avoid.) The check costs two forward passes per parameter, so it’s only for testing on tiny problems. Which raises the real question: why is backpropagation so much cheaper?

Why backward, and what it costs

The finite-difference check needs two loss evaluations for each of the $N$ parameters, so $2N$ forward passes to get the gradient. Backpropagation gets every one of the $N$ partial derivatives from one forward pass and one backward pass. The reason is the shape of the problem. The loss is one number and the parameters are many. Going backward from that one number, each intermediate derivative is computed once and reused by everything below it. Each backward rule is a vector–Jacobian product: it takes the arriving gradient and returns the one for the layer below, without ever forming the Jacobian matrix itself. For ReLU that Jacobian is diagonal, which is why its rule is an elementwise mask. (Going forward instead, the approach of forward-mode differentiation, a separate derivative would have to be pushed through the network for each parameter.) This is reverse-mode automatic differentiation, and its efficiency for one-output, many-input functions is why it trains every neural network.

How much does the backward pass cost? Count FLOPs for one linear layer, $Z = XW$, with $n$ rows and a $d \times k$ weight matrix ($n$, $d$ and $k$ are local to this section):

• Forward: $Z = XW$ is $n \cdot d \cdot k$ multiply-adds, which is $2ndk$ FLOPs.
• Backward: $\delta_X = \delta_Z W^\top$ is another $2ndk$, and $\delta_W = X^\top \delta_Z$ is another $2ndk$.

So backward costs twice forward, and one training step costs three times a forward pass: $6ndk$ FLOPs. The layer has $dk$ parameters, and in a language model each row of $X$ is one token. So per token, a training step costs about $6$ FLOPs per parameter:

for a model with $N$ parameters trained on $D$ tokens. Here, and only here, $N$ and $D$ have their scaling-law meanings, parameters and training tokens, not the query heads and model width of the notation table in AI Engineering, from the Inside. It’s an approximation, and for large models a good one. The backbone of every training-compute estimate, it leaves out three things:

• Attention scores add about $2 \cdot L \cdot S \cdot D_\text{model}$ FLOPs per token in the forward pass, which depends on the context length $S$ rather than on parameters. For Llama 3 8B at an 8,192-token context that’s $2 \times 32 \times 8{,}192 \times 4{,}096 \approx 2.1 \times 10^9$, about 13% on top of the $2N = 16 \times 10^9$.
• Embedding lookups cost no multiplications, so counting the input embedding table ($128{,}256 \times 4{,}096 \approx 0.5$ billion parameters) in $N$ overcounts by about 7%.
• Recomputation. Backpropagation needs every layer’s forward activations, so they must be kept in memory until the backward pass reaches them. That memory grows with batch size, sequence length and depth, and it’s often what limits training, more than FLOPs. The common fix, activation checkpointing, stores only some of them and recomputes the rest during backward, which costs about one more forward pass, $8ND$ in total. Utilization figures are still quoted against $6ND$.

Put numbers in. Meta reports pretraining Llama 3 8B on over 15 trillion tokens:

An H100 does about 989 TFLOPS of dense BF16 math. Real training keeps it busy only part of the time; assume 40%. That’s $7.2 \times 10^{23} / (0.4 \times 989 \times 10^{12})$ seconds, about 21,000 GPU-days, or 58 GPU-years. Meta’s model card reports 1.3 million H100-hours for pretraining the 8B, about 54,000 GPU-days. That’s 2.6 times this estimate, which would mean about 16% of peak if all of it went to the one run. Real training is less efficient than a 40% guess, or includes more than the final run, or both. The estimate gives the order of magnitude, not the bill. Inference, by contrast, needs only the forward pass, $2N$ per token. That’s the ”$2 \times$ parameters” in the cost sketch of AI Engineering, from the Inside, and the reason training a model costs so much more than running it on the same number of tokens.

The router, with a hidden layer

Does the router benefit? Give it 16 hidden ReLU units on the same features as the best linear router from Features and Overfitting, a bag of words plus the keyword count, and compare on the same 20 splits. Following the method of Features and Overfitting, every difference is paired split by split with the corrected standard error, the choice is made on validation, and only the chosen model is scored on test:

def features(rs, X):
    return np.c_[X, keyword_count(rs)]            # bag of words + keyword count

models = {"linear, L2 0.003": None, "linear, L2 0.01": None}
models.update({f"16 hidden, L2 {lam}": lam for lam in [0.0, 0.003, 0.01, 0.03, 0.1]})
val, test, train_loss = {m: [] for m in models}, {m: [] for m in models}, {m: [] for m in models}
for seed in range(20):
    tr, va, te = split(seed)
    ((Xtr, ytr), (Xva, yva), (Xte, yte)), _ = featurize(tr, va, te)
    A, B, C = features(tr, Xtr), features(va, Xva), features(te, Xte)
    for name, lam in models.items():
        if name.startswith("linear"):
            w, b = train_linear(A, ytr, float(name.split()[-1]))
            prob = lambda M: sigmoid(M @ w + b)
        else:
            P = train(A, ytr, n_hidden=16, lam=lam, seed=seed)
            prob = lambda M: sigmoid(forward(P, M)[2])
        for rows, M, t in [(val, B, yva), (test, C, yte), (train_loss, A, ytr)]:
            p = prob(M)
            rows[name].append((log_loss(p, t), ((p >= .5) == t).mean()))

Full script: router_mlp.py

The rest of the script prints the table: each model’s mean training and validation loss, its difference from the linear router paired split by split with the corrected standard error, and the test score of the network chosen on validation.

                      train  val loss  vs linear 0.003  val acc
linear, L2 0.003      0.065     0.347                     0.866
linear, L2 0.01       0.141     0.346   -0.001 ± 0.021    0.861
16 hidden, L2 0.0     0.000     0.689   +0.343 ± 0.143    0.872
16 hidden, L2 0.003   0.010     0.416   +0.069 ± 0.040    0.868
16 hidden, L2 0.01    0.032     0.365   +0.018 ± 0.015    0.869
16 hidden, L2 0.03    0.098     0.342   -0.005 ± 0.011    0.866
16 hidden, L2 0.1     0.314     0.399   +0.053 ± 0.047    0.834

chosen on validation: 16 hidden, L2 0.03. test loss 0.332 vs linear 0.333, difference -0.002 ± 0.008

(This uses router_data.py and overfit.py from Features and Overfitting and router_models.py from Measuring a Classifier Honestly.)

The honest answer is no, not yet.

• Unpenalized, it memorizes. Its training loss is 0.000. Its accuracy doesn’t suffer, 0.872 is the highest in the table, but the questions it gets wrong it gets wrong with near-certainty, and that doubles the validation loss (+0.343 ± 0.143 against the linear router). It has about 6,000 weights for 120 questions.
• Penalized, it ties. The penalty sweep has a clear bottom at $\lambda = 0.03$: weaker penalties overfit, and $\lambda = 0.1$ underfits. The best network is $-0.005 \pm 0.011$ from the linear router on validation, and on test, $0.332$ against $0.333$, a difference of $-0.002 \pm 0.008$. That’s a tie.
• The comparison is, if anything, tilted toward the network. It had five penalty settings to choose from, and the linear router had two. The linear router’s second setting, $\lambda = 0.01$, ties its first.

Two things could still hide an advantage, and neither does:

import numpy as np
from mlp import sigmoid, forward, train
from overfit import log_loss
from router_data import split, featurize
from router_models import keyword_count

# Does the answer depend on the step count or on the initialization seed?
runs = {}
for seed in range(20):
    tr, va, te = split(seed)
    ((Xtr, ytr), (Xva, yva), _), _ = featurize(tr, va, te)
    A, B = np.c_[Xtr, keyword_count(tr)], np.c_[Xva, keyword_count(va)]
    for lam in [0.0, 0.03]:
        for steps in [3000, 12000]:
            for init_seed in range(5):
                P = train(A, ytr, n_hidden=16, lam=lam, steps=steps, seed=100 * init_seed + seed)
                runs.setdefault((lam, steps, init_seed), []).append(
                    log_loss(sigmoid(forward(P, B)[2]), yva))
for lam in [0.0, 0.03]:
    for steps in [3000, 12000]:
        means = [np.mean(runs[(lam, steps, s)]) for s in range(5)]
        print(f"L2 {lam:<4} {steps:6,d} steps: val loss {min(means):.3f}-{max(means):.3f} "
              f"across 5 initialization seeds")

Full script: steps_and_seeds.py

L2 0.0   3,000 steps: val loss 0.677-0.692 across 5 initialization seeds
L2 0.0  12,000 steps: val loss 0.778-0.795 across 5 initialization seeds
L2 0.03  3,000 steps: val loss 0.342-0.343 across 5 initialization seeds
L2 0.03 12,000 steps: val loss 0.342-0.342 across 5 initialization seeds

• The step count. The network trains for a fixed 3,000 steps, which is a form of early stopping. With $\lambda = 0.03$ it has converged: 12,000 steps give the same 0.342. Without a penalty, more steps make it worse.
• The random start. Five different initializations per split give validation losses within 0.001 of each other with the penalty and within 0.02 without it, so these results aren’t the luck of one seed.

That’s a real result, not a failure of the method. A hidden layer adds the ability to learn interactions, but 120 examples don’t contain enough evidence to learn them reliably, and the strongest signal in this data, which words appear, is already linear. When the data is the bottleneck, capacity doesn’t help. The next gains come from elsewhere:

• More information per label. Cross-Entropy and the Softmax Gradient routes to all four parts of the blog. The same 200 questions then tell the model which of four routes each belongs to, not just “serving or not”.
• Features that already carry meaning. Embeddings and Contrastive Learning gives each word a vector learned from far more text than 120 questions.

Where this leads

• Cross-Entropy and the Softmax Gradient replaces the sigmoid with the softmax and the log-loss with cross-entropy for four classes. It derives why their gradient is again “prediction minus target”, then builds the four-way router and its 4×4 confusion matrix.
• Initialization and Optimizers takes up what the XOR experiment exposed: how the starting weights are chosen, why the scale matters for deep networks, and the optimizers (momentum, Adam) that make steps better than plain gradient descent.
• PyTorch from the Inside turns the four local rules into a general engine, a small autograd, and shows how PyTorch records a graph of operations and runs backward through it.
• Deep networks. With many layers, the backward pass multiplies many factors like $w$ and $\phi'(z)$. If they’re consistently below 1 the gradient vanishes, and above 1 it explodes. That problem stopped deep networks for years. Residual connections (The Residual Stream), normalization (Dropout and Normalization) and careful initialization (Initialization and Optimizers) exist to solve it.

Takeaways

• A neuron is a weighted sum plus a bias, passed through a nonlinearity. The bias lets the decision boundary sit anywhere. Without nonlinearities, stacked layers collapse to one linear layer.
• XOR is the smallest problem no single neuron can solve (four inequalities, one contradiction). One hidden layer builds features in which it becomes linear.
• Backpropagation is the chain rule evaluated from the loss backward, with each intermediate derivative computed once and reused. Per step it applies local rules: $\delta_W = X^\top\delta_Z$, $\delta_X = \delta_Z W^\top$, and a mask for ReLU.
• In a ReLU network, a weight learns from an example only if its input was nonzero and at least one path from it to the output ran through active units.
• Representing a solution isn’t finding it. Two hidden units can compute XOR, yet training found it from only 18 of 100 random starts, against 100 of 100 with 16 units. The failures are stuck, mostly on a unit that was dead from the first step. Initialization must be random to break symmetry.
• Check gradients with central differences. A relative error near $10^{-7}$ is correct, and $10^{-2}$ is a bug.
• Backward costs about twice forward, so training costs $\approx 6ND$ FLOPs: $7.2 \times 10^{23}$ for an 8B model on 15T tokens, about 21,000 H100-days at 40% utilization. Inference costs $2N$ per token.
• More capacity doesn’t help when data is the bottleneck. On 120 questions, the network tied the linear router.

References

Start with Rumelhart, Hinton and Williams: four pages, and the source of the idea that hidden layers learn representations. Baydin et al. is the clearest account of forward against reverse mode. Karpathy’s micrograd is a 100-line autograd worth reading alongside PyTorch from the Inside.

• Rumelhart, Hinton and Williams, Learning Representations by Back-Propagating Errors (Nature, 1986): backpropagation for neural networks, and the hidden layer as a learned representation.
• Griewank, Who Invented the Reverse Mode of Differentiation? (Documenta Mathematica, 2012): reverse-mode differentiation before neural networks, from Linnainmaa (1970) onward.
• Baydin, Pearlmutter, Radul and Siskind, Automatic Differentiation in Machine Learning: a Survey (JMLR, 2018): forward vs reverse mode and their costs.
• Minsky and Papert, Perceptrons (MIT Press, 1969): the limits of single-layer networks, including XOR.
• Nair and Hinton, Rectified Linear Units Improve Restricted Boltzmann Machines (ICML, 2010): the ReLU in deep learning.
• He, Zhang, Ren and Sun, Delving Deep into Rectifiers (ICCV, 2015): the initialization used in mlp.py.
• Kaplan et al., Scaling Laws for Neural Language Models (2020): the $C \approx 6ND$ training-compute approximation.
• Meta AI, Introducing Meta Llama 3 (launch blog, 2024): Llama 3 pretrained on over 15 trillion tokens.
• Meta AI, Llama 3 model card (2024): the 1.3 million H100-hours of pretraining compute reported for the 8B model.
• Goodfellow, Bengio and Courville, Deep Learning (2016): chapter 6, deep feed-forward networks and back-propagation.