Following the gradient backwards

For the bigram, we derived the gradient by hand: softmax and cross-entropy gave us the short expression $p - \text{onehot}$. That let us move the probabilities towards the corpus frequencies. What happens when we put several more layers of calculations between the input and the prediction?

Now we want a smarter model, one that looks at the previous three letters. A table for three letters would take 97⁴ ≈ 88.5 million cells, and most of them would stay empty. Instead of a table we will use a small neural network of 39,025 numbers. But now each of its numbers is linked to the error not by a two-step hop but by a long chain: multiply, add, bend, multiply again… Sprout has thousands of such chains and 17 million numbers. Deriving every gradient on paper is out of the question.

We need a way to calculate derivatives for every parameter without deriving a separate formula for each one. Backpropagation does this by working through the operations from the result back towards the inputs. It does not choose new weights itself: it finds the gradients that an optimiser uses to update them. We will trace it through a small graph and write a teaching version in Python.

How an input affects the result

We can draw a computation as a graph: input numbers on the left, a result on the right, and nodes for the operations in between. Take $L = (a \cdot b + c) \cdot f$. For each input, we want to know how $L$ responds when we increase just that input slightly. A partial derivative gives the answer; the derivatives together form the gradient. Calling this “blame” is a metaphor. More precisely, we are measuring the result’s sensitivity to its parameters. Here $L$ is a teaching example and can be negative, unlike cross-entropy.

Walk back one step at a time and watch what each node does. No node ever looks at the whole graph. Each one receives a single number from above, the gradient that has reached it, multiplies it by its own local derivative and hands the result on to its inputs. A sum passes the gradient on unchanged. A product gives each factor the gradient times the other factor. That is the whole secret.

After the backward pass, choose an input and move it. The graph shows two quantities: the change in $L$ predicted by the gradient for a shift of 0.01, and the change found by recomputing everything. They agree in this simple example. In general, a derivative gives a local approximation: smaller shifts in a smooth function make the prediction more accurate.

The chain rule

Why multiply the gradient? Picture three gears. The first turns three times faster than the second, the second twice as fast as the third. How many times faster does the first turn than the third? Six: $3 \cdot 2$. A derivative is exactly this kind of ratio: how strongly the output responds when you wiggle the input. When quantities are linked in a chain, the responses multiply.

In formulas: if $L$ depends on $d$, and $d$ depends on $e$, then

$$\frac{\partial L}{\partial e} = \frac{\partial L}{\partial d} \cdot \frac{\partial d}{\partial e}.$$

This is the chain rule. The first factor is the gradient that came from above; the second is the node's local derivative. We need only a handful of local derivatives, which we can examine one at a time:

• sum $d = e + c$: $\frac{\partial d}{\partial e} = \frac{\partial d}{\partial c} = 1$, so the gradient passes straight through and both terms get all of it;
• product $e = a \cdot b$: $\frac{\partial e}{\partial a} = b$ and $\frac{\partial e}{\partial b} = a$, so the factors "swap places";
• tanh $y = \tanh(z)$: $\frac{\partial y}{\partial z} = 1 - y^2$, conveniently computed from the output we already have.

The chain rule is easiest to feel with a nudge. Below is a chain of three links: $u = 1.5x$, then $y = \tanh u$, then $L = y^2$. Nudge $x$ by $\Delta x$ and the nudge runs down the chain, getting multiplied by the local slope at every link. It reaches $L$ multiplied by the product of all the slopes, and that product is $\frac{\partial L}{\partial x}$.

Let's run it on our graph: $a = 2$, $b = -3$, $c = 10$, $f = -2$. Forward: $e = -6$, $d = 4$, $L = -8$. Backward: $\frac{\partial L}{\partial d} = f = -2$. The sum passes the gradient through unchanged, so $\frac{\partial L}{\partial e} = \frac{\partial L}{\partial c} = -2$. The product swaps the factors: $\frac{\partial L}{\partial a} = -2 \cdot b = 6$ and $\frac{\partial L}{\partial b} = -2 \cdot a = -4$. That leaves $\frac{\partial L}{\partial f} = d = 4$. Exactly the numbers the widget showed.

One subtlety. What if the same number is used in two places? Switch the widget to "a used twice", where $L = a \cdot b + a$. Wiggle $a$ and both terms change, and their changes add up. So the gradient that reaches $a$ along two paths adds up as well: $\frac{\partial L}{\partial a} = b + 1 = -1$ (in this graph $b = -2$). In neural networks this happens all the time: hundreds of neurons read the same input, and the same matrix is applied at every position of a text. The rule is always the same: sum over all the paths.

Let $L = a \cdot a$ with $a = 3$. What gradient $\partial L / \partial a$ will the backward pass compute?

3: a product gives each factor the other factor 6: $a$ enters the product twice, and the two contributions of 3 add up 9: that is the value of $L$

The multiplication node got the same number on both inputs. Each path brings a gradient contribution $1 \cdot 3 = 3$, together $6$, which is exactly the derivative of $a^2$, that is $2a$. Forget to add the contributions and you get 3: the most common bug when people write their own autograd.

The whole backward pass

Let's assemble the algorithm. It makes two passes over the graph.

1. Forward. Compute the nodes from the inputs to the output and remember every intermediate value: the local derivatives will need them.
2. Backward. The output gets a gradient of 1. Then visit the nodes in reverse order, so that each node is handled only after everything that uses it. Each one multiplies its gradient by its local derivatives and adds the result to the gradient of its inputs.

The order "every node after all of its users" is called reverse topological order. It guarantees that by the time a node passes its gradient on, it has collected contributions from every path.

The main advantage is the amount of computation. For ordinary layers, a backward pass costs the same order of work as a forward pass; often it is roughly twice as expensive, though the ratio depends on the operations. For that cost, we get derivatives for every parameter. Nudging Sprout’s parameters one at a time would require about 17 million additional forward passes per training step. Backpropagation needs one forward pass and one backward pass.

How do you check that a backward pass is right?

The blunt method has not gone away; it has become a test. Take a few parameters, shift each one by a tiny $h$ in both directions and compute $\frac{L(w + h) - L(w - h)}{2h}$. This "central difference" is more accurate than the one-sided version we used in the last chapter. If it matches the backward-pass gradient to several digits, all is well; if not, a factor or a += is missing somewhere. Researchers who write new layers by hand run such a gradient check almost every time. Bugs in a backward pass are sneaky, because a network with a wrong gradient often still learns something, just worse.

An autograd in 50 lines

Let's write all this in Python. The idea: instead of plain numbers, use Value objects that remember where they came from. Every operation creates a new Value and also records a tiny function, _backward, that says how to pass gradients to its inputs. The backward() method lines the nodes up in topological order and calls those functions from the end.

import math class Value: """A number that remembers where it came from, so it can pass blame back.""" def __init__(self, data, children=(), op=''): self.data = data self.grad = 0.0 # dL/d(this value), filled in by backward() self._backward = lambda: None # how to pass my grad on to my inputs self._prev = set(children) self._op = op def __add__(self, other): other = other if isinstance(other, Value) else Value(other) out = Value(self.data + other.data, (self, other), '+') def _backward(): self.grad += out.grad # a sum passes the blame on unchanged other.grad += out.grad out._backward = _backward return out def __mul__(self, other): other = other if isinstance(other, Value) else Value(other) out = Value(self.data * other.data, (self, other), '*') def _backward(): self.grad += other.data * out.grad # a product: each factor gets other.grad += self.data * out.grad # the blame times the other one out._backward = _backward return out def tanh(self): t = math.tanh(self.data) out = Value(t, (self,), 'tanh') def _backward(): self.grad += (1 - t * t) * out.grad out._backward = _backward return out def backward(self): order, seen = [], set() def visit(v): # inputs before outputs: a topological order if v not in seen: seen.add(v) for child in v._prev: visit(child) order.append(v) visit(self) self.grad = 1.0 # dL/dL = 1 for v in reversed(order): # from the end back to the start v._backward() __radd__ = __add__ __rmul__ = __mul__ def __neg__(self): return self * -1 def __sub__(self, other): return self + (-other) def __repr__(self): return f"Value(data={self.data:.4f}, grad={self.grad:.4f})"

Note the += in every _backward: that is the "contributions add up" rule. Now let's test the class on the graph from the widget. The cells on this page work like a notebook: if you run the next one, the cell with the class runs first by itself.

a, b, c, f = Value(2.0), Value(-3.0), Value(10.0), Value(-2.0) e = a * b d = e + c L = d * f L.backward() for name, v in zip('abcfedL', [a, b, c, f, e, d, L]): print(name, v)

The same numbers the widget showed. This class is a simplified version of Andrej Karpathy's micrograd (2020): the original also has powers, ReLU and a few conveniences, but the heart is the same. PyTorch works on the same principle, except that its "values" are whole tensors and it has local derivatives written for hundreds of operations.

Why reset grad to zero before every step?

Because of that same +=. Gradients accumulate, and if you don't zero them, the gradient from the previous step gets added to the new one. PyTorch has optimizer.zero_grad() for this, and forgetting to call it is a classic beginner's bug. Sometimes the accumulation is used on purpose, though: compute gradients over several small batches of data in a row and only then take a step. That is how people imitate a big batch that does not fit in memory.

A neuron

We can now find gradients through a chain of supported operations. Let’s build a model from them. One basic component is an artificial neuron: it multiplies inputs by weights, adds the results and a bias $b$, then applies an activation function. A nonlinear activation changes the shape of the input-output relationship. Our example uses the hyperbolic tangent:

$$y = \tanh(w_1 x_1 + w_2 x_2 + b).$$

The weights and the bias are parameters, and descent picks their values. By the way, the third graph in the widget above, "a neuron", is exactly this formula taken apart into nodes.

Play with the weights. The sum $w_1 x_1 + w_2 x_2 + b$ is constant along straight lines, so a single neuron always splits the plane with a straight line: "yes" on one side, "no" on the other. The weights turn the line, the bias shifts it, and the bigger the weights, the sharper the transition.

There are many activation functions; we will start with two common examples.

• tanh, a smooth step from −1 to 1. It is the one inside mlp-char-3, the network we take apart at the end of this chapter.
• ReLU, $\max(0, z)$: "let positives through, cut negatives off". It could not be simpler, and it works very well in deep networks; around 2010 it became a popular alternative to sigmoid and tanh (Nair and Hinton, 2010; Glorot, Bordes and Bengio, 2011). Transformers, Sprout included, use its smooth relatives; we will meet them in chapter 8.

Now look at the slope under the activation curve. For the backward pass this is the neuron's local derivative: how much blame it lets through. For tanh the slope is largest near zero and all but vanishes at the edges: at $|z| = 3$ it is already below 0.01. A neuron pinned against ±1 lets almost no blame through and hardly learns at all: this is the plateau from the last chapter. ReLU's slope is 1 for any positive $z$ and 0 for any negative one: all or nothing.

Layers and bends

One neuron only draws a straight line. Put several neurons side by side and you get a layer: each neuron looks at all the inputs, each has its own weights, and the layer outputs as many numbers as it has neurons. The outputs of one layer become the inputs of the next. Such a network is called a multilayer perceptron, or MLP.

But why the bend? Let's try without it. A layer with no activation is just $W x + b$. Two such layers in a row give

$$W_2 (W_1 x + b_1) + b_2 = (W_2 W_1)\, x + (W_2 b_1 + b_2).$$

One matrix and one bias again, which is to say one layer again. However many linear layers you stack, they collapse into one and draw a single straight line. A bend between the layers stops the collapse: each layer bends and folds the space produced by the one before it.

Try all four datasets. The rings and xor usually take the network just a few dozen steps, the moons a couple of hundred, and the spiral well over a thousand; with only two or four neurons per layer it never quite untangles the spiral. Then switch the bend off ("none"): the boundary straightens and never bends again, however long you train. For the moons a straight line almost works (about 93% of the points end up on the correct side), but for the rings and xor it is no better than a coin toss.

The code behind this widget is the same forward and backward pass, only written for a whole layer at once: instead of separate $a \cdot b$ nodes there is a matrix times a vector. And here is the same idea built on our Value: a three-layer network learning to give the right answers for four examples.

import random random.seed(1) class Neuron: def __init__(self, n_in): self.w = [Value(random.uniform(-1, 1)) for _ in range(n_in)] self.b = Value(0.0) def __call__(self, x): z = sum((wi * xi for wi, xi in zip(self.w, x)), self.b) return z.tanh() def parameters(self): return self.w + [self.b] class Layer: def __init__(self, n_in, n_out): self.neurons = [Neuron(n_in) for _ in range(n_out)] def __call__(self, x): out = [n(x) for n in self.neurons] return out[0] if len(out) == 1 else out def parameters(self): return [p for n in self.neurons for p in n.parameters()] class MLP: def __init__(self, n_in, sizes): dims = [n_in] + sizes self.layers = [Layer(a, b) for a, b in zip(dims, dims[1:])] def __call__(self, x): for layer in self.layers: x = layer(x) return x def parameters(self): return [p for layer in self.layers for p in layer.parameters()] xs = [[2.0, 3.0, -1.0], [3.0, -1.0, 0.5], [0.5, 1.0, 1.0], [1.0, 1.0, -1.0]] ys = [1.0, -1.0, -1.0, 1.0] # the answers we want net = MLP(3, [4, 4, 1]) print(len(net.parameters()), "parameters") for step in range(41): preds = [net(x) for x in xs] loss = sum(((p - y) * (p - y) for p, y in zip(preds, ys)), Value(0.0)) for p in net.parameters(): p.grad = 0.0 # forget the blame from the last step loss.backward() for p in net.parameters(): p.data -= 0.05 * p.grad # one step of gradient descent if step % 5 == 0: print(f"step {step:2d} loss {loss.data:.4f}") print("answers:", [round(net(x).data, 2) for x in xs])

41 parameters, 40 steps of descent, the loss drops from 4.5 to 0.02, and not a single derivative worked out by hand. Value and the chain rule did it all.

The same in PyTorch

PyTorch has all of this built in. A tensor with requires_grad=True is our Value, and backward() is our backward pass (this code is for your own computer; PyTorch does not run in the browser):

import torch a = torch.tensor(2.0, requires_grad=True) b = torch.tensor(-3.0, requires_grad=True) c = torch.tensor(10.0, requires_grad=True) f = torch.tensor(-2.0, requires_grad=True) L = (a * b + c) * f L.backward() # the same backward pass, done by PyTorch print(L.item(), a.grad.item(), b.grad.item(), c.grad.item(), f.grad.item())

And here is the real code that trained the network in the next section. An embedding layer, hidden layers with tanh, an output layer, and not a single derivative: loss.backward() finds them all.

class MLP(nn.Module): def __init__(self, vocab, context, d_embed, hidden): super().__init__() self.context = context self.embed = nn.Embedding(vocab, d_embed) dims = [context * d_embed] + hidden self.hidden = nn.ModuleList(nn.Linear(a, b) for a, b in zip(dims, dims[1:])) self.out = nn.Linear(dims[-1], vocab) def forward(self, x): # x: (batch, context) h = self.embed(x).flatten(1) for layer in self.hidden: h = torch.tanh(layer(h)) return self.out(h)

The mlp-char-3 network is MLP(97, 3, 16, [256]): an alphabet of 97 symbols, a window of 3 letters, embeddings of 16 numbers, and one hidden layer of 256 neurons. It was trained with the AdamW optimizer for 20,000 steps of 512 windows each, with the weight decay knob from the last chapter set to 0.01.

A three-letter window

We now have everything we need to take apart the first real neural language model, the one Bengio, Ducharme, Vincent and Jauvin described in "A Neural Probabilistic Language Model" (2003). They used words and we use letters, but the design is the same.

1. Each of the last three letters turns into a list of 16 numbers, its embedding. That is simply a row of a trainable 97 × 16 table: every letter has its own.
2. The three embeddings are glued into 48 numbers and fed into a hidden layer of 256 tanh neurons.
3. The output layer turns the 256 numbers into 97 logits, and softmax turns those into probabilities for the next letter.

Compare the two columns. After "ti" the bigram sees only the "i" and guesses broadly, with "n" on top at 27%. The network sees " ti" with the space before it and is 73% sure the next letter is "m": "time", of course. Try "Lily wanted to pl": the network bets 94% on "a", while the bigram after "l" wavers between "e", "l" and "i". Or "Tom and his mo": the network is torn between "r" and "m", "mother" or "mom", and that is honest doubt, not ignorance.

What does all this cost? Let's count the parameters:

part	shape	numbers
embeddings	97 × 16	1,552
hidden layer	48 × 256 + 256	12,544
output layer	256 × 97 + 97	24,929
total		39,025

39 thousand numbers instead of 88.5 million table cells: over two thousand times fewer. And the network is not just thriftier, it is better. On the same held-out text, which none of the models ever saw, the bigram is surprised by 2.364 nats per letter and the three-letter network by 1.479. In terms of chapter 2's "number of equally likely options", that is $e^{2.364} \approx 10.6$ against $e^{1.479} \approx 4.4$: the network picks the next letter as if it had well under half as many options to choose from.

Where do the savings come from? The table stores every triple separately and knows nothing about how they resemble each other. The network shares its knowledge: the same 256 neurons serve every triple, and what it learned about "th" can help with "Th" if the embeddings of "t" and "T" end up similar. That kind of reuse is called generalisation.

How embeddings are organised, and why similar letters end up close together in them, is the subject of the next chapter.

Sprout right now

Sprout now has a network with a hidden layer: embeddings, 256 hidden neurons and an output layer. Backpropagation calculated the gradients, and AdamW updated the parameters. Three characters of context improve prediction over the bigram, and familiar fragments of words appear more often in the generated text. Long sentences are still beyond it: by the end of a long word, the beginning has left its window.

Now for more memory. In chapter 5 we look at what the embeddings have actually learned, learn to write such networks properly in PyTorch, and widen the window to eight letters.