Sprint 1 of this series took two posts: in Part 1 we built a tokenizer, the part of the stack that turns text into numbers, and in Part 2 we made it faster. This is Sprint 2, and it is about learning. Once a network makes a guess and the guess is wrong, how does it know what to change?
Here is the whole idea, in plain terms first. A network is a chain of simple arithmetic. Numbers come in, get multiplied and added in stages, and one final number comes out: the guess. We score the guess against the right answer, and that score is the error (the loss). Training is the work of shrinking that error by adjusting the hundreds of little numbers inside the network, called weights.
The question is which weights to adjust, and by how much. The process that answers it is called backpropagation, and the way I ended up picturing it is simple: share out the blame for the error. Every part of the network gets told how much of the final error it is responsible for, in proportion to how much it contributed. Parts that contributed a lot get adjusted a lot. Parts that contributed nothing get left alone.
The technical name for a share of blame is a gradient, and that is the word I will use from here. It is all backpropagation computes: one gradient for every weight, in a single backward pass.
By the end of this sprint, the engine trains a small network to 100% accuracy on a toy dataset of 200 points arranged in two moon shapes. Getting there took three bugs, and the third one was invisible to every test we had.
The engine is built around one class, Value, which wraps a single number. Each Value remembers which values produced it, and carries a small function (a closure) describing how to pass gradient back to them. That rule is specific to each operation. When two numbers are added, each one takes the full gradient that arrives. When two numbers are multiplied, each one takes the gradient multiplied by the other number, because that is how much it actually influenced the result.
Here is the multiplication rule, exactly as it sits in the engine. Look at the last part of each line inside the rule, the + self.grad: every share is added to whatever is already there. That is the collecting, happening in code:
def __mul__(self, other):
out = Value(self.data * other.data, (self, other), "mul")
out._backward = lambda: [
setattr(self, "grad", out.grad * other.data + self.grad),
setattr(other, "grad", out.grad * self.data + other.grad),
]
return out
Nothing flows backward while the guess is being made. The forward pass just builds the chain and leaves these rules behind, one per step. Afterwards, backward() walks the chain in reverse, starting from the loss, and each rule splits the gradient between the numbers that fed into it. If a number was used in two places, it collects a share from both. That collecting is the whole trick.
A tiny example, with real numbers. Take a = 2 and b = 3. Then c = a × b = 6, and d = c + a = 8. Working backward from d: d's gradient is 1. c fed into one addition, so c gets 1. b sits inside the multiplication, so b gets 1 × a = 2. And a is used twice: once in the addition (share 1) and once in the multiplication (share 1 × b = 3), so a ends with 4.
At real scale, the backward walk first lines the chain up in topological order, meaning every value is visited only after everything that depends on it. We built that ordering with a loop and an explicit stack rather than recursion, because Python gives up around 1,000 nested calls and real chains run far deeper than that. Here is the walk itself. Keep an eye on the visited set: it becomes important in the third bug below.
def backward(self):
topo = []
worklist = [(self, False)]
visited = set()
while worklist:
node, processed = worklist.pop()
if not processed:
node.grad = 0
worklist.append((node, True))
worklist.extend((x, False) for x in node._prev)
if processed and node not in visited:
topo.append(node)
visited.add(node)
topo.reverse()
Three small building blocks sit on top. A Neuron takes its inputs, multiplies each by a weight, adds them up, and squashes the result with tanh so the numbers stay in a tame range. A Layer is a row of neurons all looking at the same inputs. An MLP stacks layers, so MLP(2, [4, 4, 1]) is a network with 2 inputs, two hidden layers of 4, and 1 output: 37 weights in total.
Training is then just a loop:
backward() to share out the blame.p.data -= lr * p.grad.
And here is that same loop as it actually runs in the training script. Four things happen per round: guess, wipe, blame, nudge.
for step in range(steps):
loss, acc = loss_and_acc() # forward pass + score the guess
mlp.zero_grad() # wipe last round's gradients
loss.backward() # share out the blame
for p in mlp.parameters():
p.data -= lr * p.grad # nudge each weight against its gradient
Step 3 looks like housekeeping but it matters. Gradients accumulate on purpose, because a value used in several places has to collect every share. So if you do not zero them between rounds, last round's gradients leak into this round's adjustments.
In step 5, lr is the learning rate, the step size: the gradient says which direction is downhill, the learning rate says how far to trust it. Ours was 0.1.
Before the bugs, the evidence that the engine is correct. We built a chain by repeating one operation N times and compared the engine's gradient against the answer calculus gives for it, at every size:
== graph-size sweep: y = x; repeat N: y = y * w + b ==
N forward (s) backward (s) dy/dx
100 0.000 0.000 1.11
1000 0.002 0.002 2.72
5000 0.011 0.007 148
20000 0.049 0.038 4.8e+08
The last column matches the analytic answer (1.001^N) at every row, and a 20,000-deep chain runs without the recursion problem described above. One full forward-plus-backward pass of the moons network (two hidden layers of 16, so 337 weights) takes 2.88 ms. All timings are from my Mac mini.
Three bugs are worth writing about, because each one passed the tests before it failed somewhere more expensive.
The share that got counted twice. The first bug showed up in the simplest possible case: one number used twice in the same step, added to itself. With a = 2, b = a + a = 4, and c = b × b = 16, the gradient for a should end at 16. Mine returned 72. My addition rule was written like this:
a.grad = out.grad * (a.grad + 1) # wrong
a.grad = a.grad + out.grad # right
The first version re-read the gradient it had just written and multiplied it again on the second slot. The fix is one line and completely boring, which is usually the sign of the right fix.
17 weights that were secretly 4. The second bug was stranger. Training kept crashing inside the power operation, with numbers so large the program gave up. The power operation was not the bug; it was where values that had already grown enormous finally exploded. A quick count told the real story: the network claimed to have 17 weights, but only 4 of them were actually different objects. The cause was one line:
[Value(random.uniform(-1, 1))] * n_in
Multiplying a list copies the reference, not the object, so every weight in the neuron was literally the same Value, doing several jobs at once: collecting gradient from every place it was used, then getting nudged once per list entry. A few hundred rounds of that and the values ran away. Building each Value inside the loop fixed it. A quieter bug sat in the same file: a neuron computed its tanh() result and threw it away instead of returning it, so that step was silently doing nothing.
The plateau. With those fixed, all 38 tests passed, and training still did not work. On moons, the loss fell from 2.3830 and then flatlined at 0.4376, accuracy stuck at 0.880, and tripling the steps changed nothing.
The bug was in backward() itself: my topological walk kept no record of which nodes it had already visited (no visited set), so a value that fans out into two different operations could enter the walk twice and fire its rule before all of its gradient had arrived. Part of the blame was simply never delivered. A tiny check exposes it: with a = 2, h = a + 1, and y = h × 2 + h × 3, the gradient for a should be exactly 5. Mine returned 7 on some runs and 8 on others, depending on set ordering.
No test had caught it because every test reused a value inside one operation, and real networks fan out constantly: every hidden unit feeds several outputs. Adding the visited set, so each node is processed exactly once and only after its full gradient has landed, fixed the check to exactly 5.0, and the moons run went from the plateau to this:
step 0 loss 2.3830 accuracy 0.295
step 10 loss 0.1309 accuracy 0.980
step 20 loss 0.0836 accuracy 0.990
step 49 loss 0.0433 accuracy 1.000
The lesson from that one stayed with me: the training loop will faithfully follow an incomplete gradient for hundreds of steps. It has no way to know the answer it was given is missing a piece.
PyTorch, the tool production training actually runs on, is the same algorithm. Same backward walk, same topological order, same gradients that accumulate until you zero them. Forgetting the zeroing step is a classic production bug too, for exactly the reason it appears in our loop.
The difference is representation, not the idea. PyTorch operates on tensors, whole tables of numbers at once, with the arithmetic in compiled code on fast hardware, so the cost of an operation is paid once per table instead of once per number. Our engine pays Python overhead per number, which is why a few hundred weights is its comfortable size. What production adds on top (graphics cards, lower-precision arithmetic, spreading the gradient computation across machines) is engineering around the same reverse pass.
One design choice is worth naming. Building the graph as the forward pass runs, the way our Value does, is what PyTorch calls define-by-run. The alternative is to capture the whole graph before running it, which is JAX's approach: less flexible to debug, but the whole graph can be optimized at once, fusing operations together. After a weekend of debugging a hand-built graph, both sides of that trade-off make sense to me.
Next in the series: the transformer block, where this MLP comes back as the feed-forward half of every layer, and the network starts generating text instead of just sorting points into two moons.