Walking downhill

In chapter 1 we built a model without training anything: we counted which letter follows which, and divided. In chapter 2 we measured how good it was: 2.365 nats of surprise per letter on the very text it was counted from (and 2.364 on text it had never seen). For a 97 × 97 table you cannot do better. But counting has a ceiling, and we are about to hit it.

Suppose we want to look at the previous three letters instead of one. There are 97³ = 912,673 different triples, and each one needs a row of 97 probabilities: almost 89 million cells. Our corpus has 41 million letters, so most triples never show up even once, and their rows stay empty. Sprout, meanwhile, looks 512 tokens back. A table for that much context would not fit in the universe.

We need another way to find good numbers, one that works for any model and not just for tables. The idea is simple: start with numbers picked at random, measure the error, and nudge every number a little in the direction that makes the error smaller. Repeat thousands of times. That is training, and at its heart sits gradient descent. In this chapter we try it on a familiar problem: we build the bigram again, this time by descent instead of counting, and check that we arrive at exactly the same place.

Lost in the fog

Picture yourself in the mountains when a fog rolls in so thick you cannot see your own boots. You need to get down to the valley. No map, no visible path. What can you do? Feel with your feet which way the ground tilts right where you stand, and step the way it drops most steeply. Then feel the slope again, and take another step.

Play a couple of rounds from different spots, then lift the fog. You will notice three things, and all three are about how neural networks learn.

1. You don't need the map. To go down, it is enough to know the slope where you are standing. Models learn exactly like this: nobody ever sees the whole "terrain" of the error, but the slope at the current point can always be computed.
2. The steps shrink on their own. The button steps a distance proportional to the steepness: long strides on a steep slope, tiny ones near the bottom. At the very bottom the slope is almost zero, and the walker nearly stops.
3. A low point is not necessarily the lowest. From an unlucky start, walking downhill takes you to a shallow lake where every direction leads up. Such a place is called a local minimum.

Now let's translate the metaphor. Height is the model's error, the loss from chapter 2. Position on the map stands for the model's parameters. The walker has two of them, west–east and south–north. Our bigram has 9,409; Sprout has 17 million. You cannot draw that terrain, but you can still compute its slope and step the same way. What remains is to say what "slope" means in numbers.

A derivative is a slope

Start with a single knob. Say the error depends on one number $x$ and looks like a bowl: $f(x) = \tfrac12 (x-2)^2$. The bottom is at $x = 2$. Standing at some $x$, how do you know which way to go?

Move a tiny distance $h$ away from $x$ and see how much the height changes. Rise over run is the slope:

$$\text{slope} \approx \frac{f(x+h) - f(x)}{h}.$$

The smaller $h$, the more accurate. The value this ratio approaches as $h \to 0$ is called the derivative, written $f'(x)$ or $\frac{df}{dx}$. For our bowl it takes one line:

$$\frac{f(x+h)-f(x)}{h} = \frac{\tfrac12 (x-2+h)^2 - \tfrac12 (x-2)^2}{h} = \frac{(x-2)h + \tfrac12 h^2}{h} = (x-2) + \tfrac12 h.$$

Once $h$ vanishes we are left with $f'(x) = x - 2$. The sign of the derivative tells you which way is uphill: at $x = -1$ the derivative is $-3$, so left is up and right is down. Its size tells you how steep it is. The descent rule falls out by itself: step against the derivative, by a small fraction of it.

$$x \leftarrow x - \eta \, f'(x)$$

The number $\eta$ (Greek "eta") is the learning rate, the step size. It is the only setting gradient descent has, and everything depends on it. See for yourself:

In the bowl everything can be computed exactly. The distance to the bottom is $x - 2$, and after one step it becomes

$$x - \eta(x-2) - 2 = (1-\eta)(x-2).$$

Every step multiplies the distance to the bottom by $1 - \eta$. With $\eta = 0.1$ that is $0.9$, and the ball crawls. With $\eta = 1$ the factor is zero: straight to the bottom in one step. With $\eta = 1.5$ the factor is $-0.5$: the ball jumps over the bottom, but each jump is half as long as the last. With $\eta = 2$ it swings back and forth forever, and with $\eta > 2$ each jump is longer than the one before and the ball flies off. Training is then said to have "diverged": instead of falling, the error grows without bound.

Now switch the curve to "Two dips" and "Plateau". With two dips, the ball rolls into whichever is closer, even when the other one is deeper. On the plateau the slope is almost zero, so the steps are almost zero too: the ball barely moves, although a deep dip lies ahead. Both problems are real, and we will meet both again, the plateau as early as the next chapter.

All of this fits in a few lines of Python. We compute the derivative the blunt way, straight from the definition: rise over a tiny run.

def f(x): return 0.5 * (x - 2) ** 2 def slope(f, x, h=1e-6): return (f(x + h) - f(x)) / h # rise over run, on a tiny step x = -1.0 # where the ball starts lr = 0.5 # the step size, η for step in range(8): g = slope(f, x) print(f"step {step}: x = {x:7.4f} f(x) = {f(x):7.4f} slope = {g:7.4f}") x = x - lr * g # walk against the slope

With $\eta = 0.5$ the slope, and with it the distance to the bottom, halves at every step: 3, 1.5, 0.75… Change lr to 2.5 and run it again.

Many knobs: the gradient

The walker in the fog has two coordinates, and the height depends on both: $f(x, y)$. How do you measure a slope when there are infinitely many directions? The trick is to turn the knobs one at a time first.

Freeze $y$ and move only $x$. What you get is an ordinary function of one variable, with an ordinary derivative. It is called the partial derivative with respect to $x$ and written $\frac{\partial f}{\partial x}$; the curly "d" reminds you that every other variable is frozen for the moment. In the same way you get $\frac{\partial f}{\partial y}$. The two numbers together are the gradient:

$$\nabla f = \left( \frac{\partial f}{\partial x},\ \frac{\partial f}{\partial y} \right).$$

The gradient is an arrow. It points in the direction where the function grows fastest, and its length says how fast. We want to go down, so we step against the arrow, and the rule looks just like the one-knob version:

$$(x, y) \leftarrow (x, y) - \eta \, \nabla f(x, y).$$

Look at the arrows at the ball: the step really is the sum of two independent pieces. Each partial derivative tells you how to turn its own knob, and together they give the direction of steepest descent. This scales up without a single change: the bigram has 9,409 knobs, so its gradient is 9,409 numbers, one per knob. Sprout's has 17 million. One step of descent moves all the knobs at once.

Raise the step η to 0.1 and the ball starts bouncing between the shores of a lake, or even gets thrown out of it, just like the bowl with $\eta > 2$. Narrow valleys make a big step especially dangerous: across the valley the walls are steep, so the ball keeps jumping over the bottom while hardly moving along it.

Why does the gradient point straight up the steepest slope?

Take a tiny step in the direction of a unit vector $u = (u_x, u_y)$. Close up the surface is nearly flat, so the height changes by about $\frac{\partial f}{\partial x} u_x + \frac{\partial f}{\partial y} u_y$. That is the dot product $\nabla f \cdot u$, which equals $|\nabla f| \cdot |u| \cdot \cos\alpha$, where $\alpha$ is the angle between the two vectors. It is largest when $\cos\alpha = 1$, meaning $u$ points the same way as the gradient, and smallest when it points exactly the other way. Across the gradient ($\cos\alpha = 0$) the height does not change at all: that is a level line, a contour on the map. This is why the downhill arrows always cross the contours at right angles.

Probabilities from any numbers

Back to the bigram. To train it by descent we need numbers that can be pushed anywhere. Probabilities cannot: they are never negative, and every row has to add up to one. A descent step knows nothing about those rules and will happily break them.

The way out is to store free numbers instead of probabilities and turn them into probabilities at the last moment. The free numbers are called logits (we'll write them $z$), and the conversion is called softmax:

$$p_i = \frac{e^{z_i}}{\sum_j e^{z_j}}.$$

It works in two moves. The exponential $e^{z}$ makes any number positive: large logits become very large, negative ones tiny but never zero. Dividing by the sum normalises them, so that together they add up to exactly 1. Play with the logits of the five most common letters after "t":

Press "+1 to all": the logits go up, but the probabilities don't move a hair. Softmax only cares about the differences between logits: $e^{z_i + c} = e^c \cdot e^{z_i}$, and the factor $e^c$ cancels between the top and the bottom. "×2" stretches the differences, and the distribution becomes sharper, more confident. (Temperature works the same way when generating text; we'll come back to it in chapter 11.)

A neural bigram is just a 97 × 97 table $W$ holding logits instead of probabilities. Row $a$ holds the logits for the letter that comes after $a$. To predict the next letter, take the row and pass it through softmax. Why call that a neural network? Write the current letter as a vector of 97 zeros with a single 1 in that letter's position, a so-called one-hot vector. Then "take row $a$" is the same as multiplying that vector by the matrix $W$. One matrix, one layer, no hidden ones: the smallest neural network there is. In chapter 4 we will give it more layers.

Which way to pull the logits

Descent needs the gradient of the error with respect to every logit. The error on one example is the surprise at the right letter $y$, from chapter 2: $L = -\ln p_y$. Let's write it in terms of the logits by plugging in softmax:

$$L = -\ln \frac{e^{z_y}}{\sum_j e^{z_j}} = -z_y + \ln \sum_j e^{z_j}.$$

Take the partial derivative with respect to some logit $z_i$. For the first term it is $-1$ when $i = y$ and zero otherwise. For the second one:

$$\frac{\partial}{\partial z_i} \ln \sum_j e^{z_j} = \frac{e^{z_i}}{\sum_j e^{z_j}} = p_i.$$

Add them up:

$$\frac{\partial L}{\partial z_i} = p_i - [i = y].$$

The square brackets mean 1 when the condition holds and 0 when it doesn't. The whole gradient is the vector of probabilities minus the one-hot vector of the right answer: $p - \text{onehot}(y)$. Machine learning has few formulas simpler than this one, and it is worth remembering.

Now read the formula as descent would, stepping against the gradient. The right letter's logit gets a push up of size $1 - p_y$: the less the model believed the right answer, the harder the push. Every wrong letter's logit gets pushed down by its own $p_i$: the more confidently the model was wrong, the harder it is pushed down. Go back to the widget above, choose the right letter and press "Descent step" a few times. The last column is exactly $p - \text{onehot}$, and the surprise melts with every step.

A formula like this deserves an honest check: wiggle each logit and watch how the error changes. It is the same trick as with the ball, only with five knobs now.

import numpy as np def softmax(z): e = np.exp(z - z.max()) # subtracting the max changes nothing, but keeps exp() from overflowing return e / e.sum() z = np.array([2.0, 1.0, 0.5, 0.0, -1.0]) # logits for h, space, o, e, i y = 0 # the right answer is 'h' p = softmax(z) loss = -np.log(p[y]) grad = p.copy() grad[y] -= 1 # our formula: p - onehot # the slow, honest check: nudge each logit a little and watch the loss nudged = np.zeros_like(z) for i in range(len(z)): dz = np.zeros_like(z) dz[i] = 1e-6 nudged[i] = (-np.log(softmax(z + dz)[y]) - loss) / 1e-6 print("p ", p.round(4)) print("formula ", grad.round(4)) print("nudging ", nudged.round(4))

Now from one example to the whole corpus. The bigram's error is the average surprise over all $N$ = 41.3 million pairs of letters. The pair "$a$, then $b$" occurs $n_{ab}$ times, and every occurrence pushes row $a$ by the formula above. Add up all the pushes and divide by $N$:

$$\frac{\partial L}{\partial W_{ab}} = \frac{n_a}{N}\left(p_{ab} - \frac{n_{ab}}{n_a}\right),$$

where $n_a$ is how many times letter $a$ occurred, and $\frac{n_{ab}}{n_a}$ is the share of $b$ among the letters after $a$: exactly the counted probability from chapter 1. When does the gradient become zero? When $p_{ab} = \frac{n_{ab}}{n_a}$ for every pair. Descent comes to rest exactly where the network's probabilities match the counts.

A model gave the right letter probability 0.3 and one of the wrong letters 0.5. What are the gradients of their logits?

+0.3 for the right one, −0.5 for the wrong one −0.7 for the right one, +0.5 for the wrong one Zero for both: the model already knows something

The gradient is $p - \text{onehot}$: for the right letter $0.3 - 1 = -0.7$, for the wrong one $0.5 - 0 = 0.5$. Descent steps against the gradient, so the right letter's logit goes up and the wrong one's goes down.

Training the bigram

Everything is in place. A table $W$ of 9,409 logits, all zero at the start, so each of the 97 symbols gets probability 1/97 and the surprise is $\ln 97 \approx 4.575$ nats. The gradient comes from the formula above, and the counts $n_{ab}$ from the same table of counts as in chapter 1. Press Train.

The curve drops steeply at first, then creeps more and more slowly towards the green dashed line. In the first 100 steps the surprise falls from 4.575 to 2.49; by step 1,000 it is 2.381, and by step 20,000 it is 2.366, a thousandth away from counting. The table below the chart comes into focus too: in the "q" row a bright "u" lights up, after a space come "t", "a" and "s". These are the same patterns as in chapter 1, except that this time the network found them by itself.

Tap the "q" row. Even when the curve has almost settled on the dashed line, this row is still visibly behind. After 300 steps less than 1% of the probability in the space row is in the wrong place, but in the "q" row it is a whopping 60%. The culprit is the factor $\frac{n_a}{N}$: the rarer the letter, the weaker the push its row gets. Spaces are 18% of all letters; "q" is three hundredths of a percent. Common letters are learned first, rare ones last. That unfairness is one of the reasons real models are trained not with bare descent but with optimizers like Adam, which choose a step size for every number separately (chapter 10).

Why are the last thousandths so hard to get? The table has 7,051 cells with a zero count: pairs that never occurred, to which counting gives a probability of exactly 0. The network cannot do that: for $e^{z}$ to become zero, the logit would have to reach minus infinity. Descent dutifully keeps pushing those logits down, but the push is proportional to the probability $p$ itself, and that shrinks with every step. The network gets ever closer to counting without ever reaching it.

Now play with the step size. With η = 5 descent works, but sluggishly: after 1,000 steps the surprise is still 2.49. With η = 150 the curve falls at first, then starts to shake around 2.5 and goes no lower: for the rows of the most common letters this step is too big (like $\eta > 2$ in the bowl), and they start to swing. The space row alone receives almost a fifth of all the pushes, and "e", "t" and "h" are not far behind. With η = 500 training never settles at all: the surprise jumps around somewhere between 6 and 9.5, worse than guessing blindly. It is the ball in the bowl all over again, only in 9,409 dimensions.

One last experiment: start from random logits instead of zeros. The surprise starts higher, around 5.0, because a random model is confidently wrong, but it ends up at the same point. This problem has a single valley with no false dips (mathematicians call such a loss convex), so wherever you start, descent finds its bottom. Real neural networks have many valleys, and yet in practice descent keeps finding good places. Why it works out so well is still not fully understood, even by researchers.

Here is the same descent in numpy: 600 steps, a couple of seconds right in your browser.

import json import numpy as np from pyodide.http import open_url # in the browser; on your computer: open('bigram-counts.json') data = json.load(open_url('/llm/data/bigram-counts.json')) C = np.array(data['counts'], dtype=float) # C[a, b]: how often b came right after a N = C.sum() # 41.3 million pairs n = C.sum(axis=1, keepdims=True) # how often each letter a occurs W = np.zeros((97, 97)) # logits; all equal = every symbol gets 1/97 lr = 50 for step in range(601): P = np.exp(W - W.max(axis=1, keepdims=True)) P /= P.sum(axis=1, keepdims=True) # softmax, row by row loss = -(C * np.log(P)).sum() / N # average surprise per letter if step % 100 == 0: print(f"step {step:4d} loss {loss:.4f}") W -= lr * (n * P - C) / N # the gradient: (n_a / N) * (p - counted) q, u = data['alphabet'].index('q') + 1, data['alphabet'].index('u') + 1 print("p(u | q): network", round(P[q, u], 3), " counting", round(C[q, u] / n[q, 0], 3))

In real projects nobody derives gradients by hand: a library such as PyTorch finds them by itself. Let's check that it arrives at exactly our formula (this code is for your own computer; PyTorch does not run in the browser):

import torch import torch.nn.functional as F C = torch.tensor(data['counts'], dtype=torch.float64) N, n = C.sum(), C.sum(1, keepdim=True) W = torch.zeros(97, 97, dtype=torch.float64, requires_grad=True) loss = -(C * F.log_softmax(W, dim=1)).sum() / N loss.backward() # PyTorch works out the gradient by itself P = F.softmax(W, dim=1).detach() print(f"loss {loss.item():.4f}") print(torch.allclose(W.grad, (n * P - C) / N)) # the same as our formula?

How loss.backward() manages this for any formula, however tangled, is the subject of the next chapter.

Smoothing is a penalty

The counted table has one nasty habit: pairs that never appeared in the corpus get probability zero. And the surprise at a zero-probability event is infinite: one unforeseen pair and the score for the whole text collapses. That is why people usually add a little to every count, say one. This is smoothing, and the bigram-char model that babbled at the end of chapter 1 is smoothed exactly this way.

A neural network has its own way to do the same thing: add a penalty for large logits to the error.

$$L_{\text{total}} = L + \lambda \sum_{a,b} W_{ab}^2.$$

The penalty's gradient is $2\lambda W_{ab}$: it pulls every logit towards zero, and so every row towards the uniform distribution. The data pull towards counting, the penalty towards "all letters are equal", and the network settles somewhere in between. Where there is plenty of data, the data win; where there is little, the penalty does. Smoothing behaves in just the same way: an extra one hardly changes anything in the space row with its 7.5 million pairs, but it blurs the row of a rare letter a lot.

Try the penalty λ in the widget above (after "Start over"). Here is what you get after 20,000 steps:

λ	surprise	p(u | q)	p(u | Q)
0	2.366	0.994	0.93
10⁻⁶	2.371	0.949	0.57
10⁻⁵	2.403	0.667	0.03
10⁻⁴	2.577	0.045	0.01

The stronger the penalty, the higher the surprise on the training text: the network deliberately refuses to fit the table to the data down to the last digit. The "Q" row (958 occurrences in the corpus) blurs even under a weak penalty, while the "q" row (13,061 occurrences) holds out longer. Why would you want that if the surprise goes up? Because what we care about is not the training text but text the model has never seen. A bigram trained on 41 million letters has almost no room to overfit, that is, to learn quirks of its training text that do not carry over to new text: on unseen text the counted table scores 2.364 nats, practically the same as on the old one. Models with millions of parameters are another matter. They can memorise the accidents of their training data, and the penalty helps keep them from doing so. The technique is called regularisation, and this particular kind is L2 regularisation, or weight decay. The AdamW optimizer used to train large models has a dedicated knob for it.

Sprout right now

Sprout still sees only the previous letter, so it babbles just as it did in chapter 1. In spirit, though, this is a different model: its numbers were not counted but found by descent, the very method that trained all 17 million numbers of the real Sprout. Below are three texts written with the same rolls of the dice: before training, by your network, and by the counted table. If you haven't trained the network yet, you can do it right here.

To see further back than one letter, a model needs hidden layers, and the gradient through them can no longer be worked out on a napkin like $p - \text{onehot}$. In the next chapter we learn to find it automatically, with backpropagation.