Walking downhill

For our first model, we found the parameters by counting character pairs and normalising the rows of a table. The loss was about 2.365 nats per character on training text and 2.364 on held-out text. Now we will solve the same problem differently: start with poor probabilities and improve them a little at a time. This introduces the training method we will need for a neural network.

With three preceding characters, there are $97^3 = 912\,673$ possible contexts. Each needs 97 probabilities, for almost 89 million cells. Even in 41 million characters, many combinations may never appear and rare ones will have too few observations. The table grows exponentially as we extend the context. The finished Sprout uses up to 512 tokens of context; listing every possible sequence of that length is entirely impractical.

We still want to find parameters from data, but a neural network does not give us a simple counting formula for the answer. Instead, start with some parameter values, measure the loss, and calculate how a small change to each parameter would affect it. Then take a step towards a lower loss. This is the basis of gradient descent. We will first try it on the familiar bigram, where counting gives us a reference to compare against.

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}.$$

In exact arithmetic, this ratio approaches the derivative as $h$ shrinks, provided the derivative exists. We write it as $f'(x)$ or $\frac{df}{dx}$. On a computer, an excessively small $h$ can make the estimate worse because subtracting nearly equal values magnifies rounding error. For our bowl, we can find the derivative exactly by expanding the brackets:

$$\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. It scales the update: the actual movement is the learning rate times the derivative. This is the setting we vary in the next experiment. See how it changes the motion:

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.

Start from random logits instead of zeros. The initial loss is higher, around 5.0, but a suitable learning rate again brings it close to the counted result. This model’s loss is convex in its logits: it has no local minima with a worse value than the global minimum. That does not mean there is one unique set of parameters. Adding a constant to a whole row leaves softmax unchanged, and zero observed frequencies require some logits to tend towards minus infinity. Convexity also cannot rescue an excessively large step, as the previous experiment showed.

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 and regularisation

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.

Another way to discourage overly confident predictions is to add a penalty for large logits to the loss. It plays a similar role to Laplace smoothing, but it is a different calculation and generally gives different probabilities:

$$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 has a similar dependence on the number of observations: 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

A stronger penalty keeps the model from fitting the training frequencies as closely. The rare “Q” row makes this especially visible. This is regularisation: we limit the fit in the hope of improving predictions on new data. Whether it helps must be checked on held-out data; there is no automatic guarantee. Here we use L2 regularisation, adding a sum of squared weights to the loss. With ordinary gradient descent, that is equivalent to shrinking the weights at every step. With adaptive optimisers the distinction matters: AdamW applies weight decay separately from the loss gradient.

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.

In the next chapter we will add a hidden layer and feed the model several preceding characters. Deriving every gradient by hand is still possible, but it soon becomes tedious. We will learn how a program can do that work through backpropagation.