Choosing the next word

Training is over. Sprout has read 330 million tokens, and given any beginning, it can assign a probability to each of its 8,192 tokens. But a list of probabilities is not yet a text. Somebody has to decide which token actually gets written, and, surprisingly, that decision is not made by the neural network at all. It is made by a short piece of code that runs after the model, has no weights and learns nothing: the sampler.

Swap the sampler, and the same Sprout, with exactly the same 17 million numbers, becomes a bore that repeats itself, a careful storyteller or a babbling madman. In this chapter we take the sampler apart knob by knob. Then we look at the two engineering tricks that make Sprout fast enough to live in a browser tab: the key–value cache and 8-bit weights.

The sampling desk

Here is the whole control panel at once. The box at the top holds a beginning; below it are Sprout's real probabilities for the next token, computed right now on your device. Draw a few tokens, move the sliders and watch which candidates survive. Don't try to make sense of every number yet: each knob gets its own section below, and you can always scroll back and try again.

Two things stand out before we touch a single knob. First, the distribution has a head and a very long tail. A handful of tokens take most of the probability, and thousands of others share the crumbs: on the lower chart the line falls by many orders of magnitude between rank 1 and rank 8,192. Second, the crumbs add up. Under the bars, the widget adds up everything that didn't make the top twelve. After a predictable beginning this tail is thin; after an open one it can hold a sizeable share, and then quite a few draws land somewhere in it.

Everything a sampler does is a decision about the head and the tail: how strongly to favour the head, and whether to allow the tail at all.

Temperature: one dial for boldness

We want a single knob that moves smoothly between "always take the favourite" and "anything goes", without ever changing the model's order of preference. The favourite should stay the favourite; only how much it is favoured should change.

Recall where the probabilities come from (chapter 3): the network outputs a score, a logit $z_i$, for every token, and softmax turns the scores into probabilities. Temperature $T$ simply divides every logit before the softmax:

$$p_i(T) = \frac{e^{z_i/T}}{\sum_{j} e^{z_j/T}}$$

To see what this does, compare two tokens. The denominators cancel:

$$\frac{p_i(T)}{p_j(T)} = \frac{e^{z_i/T}}{e^{z_j/T}} = e^{(z_i - z_j)/T}.$$

At $T = 1$ the same ratio is $e^{z_i - z_j}$. Raising it to the power $1/T$ gives exactly the line above, so

$$\frac{p_i(T)}{p_j(T)} = \left(\frac{p_i(1)}{p_j(1)}\right)^{1/T}.$$

That is the whole secret of temperature: the ratio between any two probabilities is raised to the power $1/T$. Suppose Sprout thinks " girl" is three times as likely as " boy". At $T = 0.5$ the power is 2, and " girl" becomes $3^2 = 9$ times as likely. At $T = 2$ the power is ½, and the gap shrinks to $\sqrt{3} \approx 1.7$. The ranking never changes, only the gaps between neighbours.

Now the two extremes. As $T \to 0$, every gap is raised to a huge power, so the favourite swallows all the probability: that is greedy decoding, always the top token. As $T \to \infty$, every $z_i / T$ shrinks towards zero, every $e^{z_i/T}$ towards $e^0 = 1$, and all 8,192 tokens become equally likely: pure noise. Between them sits $T = 1$, the distribution the model actually learned.

How do we measure how spread out a distribution is? Chapter 2 gave us the tool: the entropy $H = -\sum_i p_i \ln p_i$. It has a friendly reading. A fair die with $n$ faces has entropy $\ln n$, so a distribution with entropy $H$ is as uncertain as a fair die with $e^H$ faces. The sampling desk shows this number as "effective choices". Put the temperature at 0.3 and often only one or two real choices are left; put it at 1.5 and the die grows dozens, often hundreds, of faces.

Here is the same beginning written four times, at four temperatures, with the same random seed:

Look at the numbers under the columns. They measure how surprised Sprout itself is by each text when it reads it back at the normal temperature: the average $-\ln p$ per token, the very quantity it was trained to minimise. The cold text is far more predictable than real stories: people don't always say the most likely thing. The hot text shocks even the model that wrote it. Around $T = 1$ the samples are roughly as surprising as the stories Sprout learned from (on real held-out stories its loss is 1.55), and in practice people often run samplers a little colder than that, trading some liveliness for fewer slips.

At T = 1 Sprout thinks " girl" is 4 times as likely as " boy". How many times as likely is it at T = 0.5?

2 times 4 times: temperature keeps the ratios 8 times 16 times

The ratio is raised to the power 1/T = 2, so 4² = 16. Lower temperature widens every gap; higher temperature narrows them.

Greedy decoding goes round in circles

If the model knows which token is most likely, why not always take it? Greedy decoding sounds like the safest choice. Chapter 0 already hinted at the problem; now let's watch it properly. Below, both columns write greedily, so the output is fully determined by the beginning. Red marks every token that repeats a four-token phrase already present in the text.

Why does a model that has read hundreds of millions of words fall into loops? Because repetition is a real pattern of language. Names come back, refrains come back, " the" comes back every few words, and a transformer has attention heads whose whole job is to find an earlier occurrence of the current token and predict whatever followed it last time. These are induction heads, described by Olsson and colleagues in 2022. Once a phrase appears twice, this copying machinery makes a third time more likely, and the third makes a fourth more likely still. Greedy decoding has no way to step aside, so the text settles into a loop, like a record stuck in a groove. Sampling escapes by luck: sooner or later the dice pick something else.

Not every beginning loops, and a better-trained model loops less. Try the third beginning: greedy Sprout writes a whole story about Lily and her big red ball without getting stuck. But the duck and the mouse show that the tendency never disappears completely, and it is one of the reasons chatbots don't decode greedily.

The repetition penalty

A blunt but effective fix came with the CTRL model (Keskar et al., 2019): before choosing, weaken every token that already appears in the recent text. Our engine looks back 64 tokens and, for each token $i$ it finds there, rescales its logit by a factor $r > 1$:

$$z_i \leftarrow \begin{cases} z_i / r, & z_i > 0, \\ z_i \cdot r, & z_i \le 0. \end{cases}$$

Both branches lower the score: a positive logit shrinks towards zero, a negative one moves further below it. Move the slider in the widget above: a small penalty usually breaks the loop. Then push it to 2 and see the flip side. The penalty cannot tell a boring loop from a necessary repeat, so the text starts avoiding " the", full stops and the hero's own name, and turns strange. That is why the chat on this site uses only a gentle $r = 1.1$.

Cutting off the tail: top-k and top-p

If greedy is too rigid, why not sample from the model's own distribution, $T = 1$ and nothing else? After all, that is exactly the distribution it learned. The trouble is the tail. Every tail token on its own is unlikely, but there are thousands of them, and a story is hundreds of tokens long. Suppose the tail holds just 3% at each step. The chance of never touching it in 200 tokens is $0.97^{200} \approx 0.002$. In other words, a sampled story almost certainly contains a few tokens from the tail.

And the tail is where the model's mistakes live: half-learned guesses, misspellings, junk the model never quite ruled out. Worse, one bad token doesn't stay alone. The model reads its own output as if a careful author had written it and happily continues from nonsense, so errors compound. Hence the idea: sample, but only from the head.

top-k

The first recipe (Fan, Lewis and Dauphin, 2018) keeps the $k$ most likely tokens, sets all others to zero, and rescales the survivors so they add up to one again. With $k = 40$ the tail is gone. But no single $k$ suits every distribution. After "Once upon a" there is essentially one sensible token, " time", and top-40 keeps 39 bad ones next to it. After "For breakfast she ate" dozens of words are fine, and top-5 throws good ones away.

top-p, the nucleus

Holtzman, Buys, Du, Forbes and Choi (2019, published at ICLR 2020) proposed cutting by mass instead of by count. Sort the tokens from most to least likely, $p_{(1)} \ge p_{(2)} \ge \dots$, and keep the smallest head whose probabilities add up to at least $p$:

$$k^{*} = \min\Big\{k : \sum_{i=1}^{k} p_{(i)} \ge p\Big\}, \qquad p'_{(i)} = \frac{p_{(i)}}{\sum_{j=1}^{k^{*}} p_{(j)}} \ \text{ for } i \le k^{*}.$$

They called this head the nucleus. It adapts by itself: when the model is confident, the first token alone may reach $p$ and the nucleus shrinks to one token; when many continuations are plausible, it grows to hundreds. Try it on the sampling desk: set top-p to 0.9, switch between the ready-made beginnings and watch the number of surviving tokens jump. Then switch the chart to "running total": top-p is the horizontal line where the running total crosses $p$, while top-k would be a vertical line at a fixed rank.

Here are all three knobs in a few lines of numpy, applied to a made-up distribution over eight words, with ten thousand draws each. You can run it right here:

import numpy as np logits = np.array([3.0, 2.2, 1.5, 0.3, -0.5, -1.0, -2.0, -4.0]) words = ["girl", "boy", "dog", "cat", "bird", "king", "spoon", "the"] def next_token_probs(logits, temperature=1.0, top_k=0, top_p=1.0): z = logits / max(temperature, 1e-5) order = np.argsort(-z) # most likely first if top_k: z = z.copy() z[order[top_k:]] = -np.inf # everything below the k-th is gone p = np.exp(z - z.max()) p /= p.sum() if top_p < 1.0: sorted_p = p[order] drop = np.cumsum(sorted_p) - sorted_p > top_p p[order[drop]] = 0 # keep the head up to top_p p /= p.sum() return p settings = [(1.0, 0, 1.0), (0.5, 0, 1.0), (2.0, 0, 1.0), (1.0, 3, 1.0), (1.0, 0, 0.8)] rng = np.random.default_rng(0) for T, k, top_p in settings: p = next_token_probs(logits, T, k, top_p) draws = rng.choice(len(p), size=10_000, p=p) counts = np.bincount(draws, minlength=len(p)) shown = " ".join(f"{w} {c}" for w, c in zip(words, counts) if c) print(f"T={T} top_k={k} top_p={top_p}: {shown}")

Look at the line for top_p = 0.8: only "girl" and "boy" survive, because "girl" alone holds about 55% and the two together pass 80%. The line drop = np.cumsum(sorted_p) - sorted_p > top_p is a neat trick: it drops a token only if the tokens before it already exceed $p$, so the token that crosses the line stays in.

And here is the real thing, the generation loop from Sprout's model.py. The order matters: temperature first, then top-k, then softmax, then top-p, then a weighted draw with torch.multinomial.

@torch.no_grad() def generate(self, idx, max_new, temperature=0.8, top_k=None, top_p=0.95, stop=None): for _ in range(max_new): logits, _ = self(idx[:, -self.cfg.context:]) logits = logits[:, -1, :].float() / max(temperature, 1e-5) if top_k: kth = torch.topk(logits, top_k).values[:, -1, None] logits[logits < kth] = -float('inf') probs = F.softmax(logits, dim=-1) if top_p and top_p < 1.0: sorted_p, order = probs.sort(descending=True) drop = sorted_p.cumsum(-1) - sorted_p > top_p sorted_p[drop] = 0 probs = torch.zeros_like(probs).scatter(-1, order, sorted_p) probs /= probs.sum(-1, keepdim=True) nxt = torch.multinomial(probs, 1) idx = torch.cat((idx, nxt), dim=1) if stop is not None and nxt.item() in stop: break return idx

Notice one wasteful line: self(idx[:, -self.cfg.context:]) runs the whole model over the whole text again for every new token. It is short, correct and fine for a terminal script. The browser engine does better, and we will see how in a moment.

top-p = 0.9. Sprout is 95% sure the next token is " time". How many tokens can the sampler pick from?

Just one: " time" 90% of the vocabulary, about 7,400 tokens The top ten, whatever their probabilities

The first token alone already reaches 0.9, so the nucleus is a single token and this step is effectively greedy. When the model is sure, top-p lets it be sure; when it hesitates, the nucleus widens.

Where the randomness comes from, and where it stops

"Draw a token" means the same weighted dice as in chapter 1: take a random number $u$ between 0 and 1, walk along the kept tokens adding up their probabilities, and stop at the first token where the running total passes $u$. A token with probability 0.3 owns a 0.3-wide slice of the interval, so it is hit 30% of the time.

But computers don't roll dice. The random numbers come from a pseudo-random generator: a small formula that turns one number, the seed, into a long sequence that looks random. Sprout's engine uses a 32-bit generator of this kind. Same seed, same model, same knobs, and you get the same text, token for token. The four-temperature widget above relies on this: keep the seed, press the button again, and every column repeats itself exactly; let the seed change and each run is new. Seeds make experiments reproducible and comparisons fair. They are also why "regenerate" in a chatbot gives a different answer: it simply uses a new seed.

Why the four columns often start with the same words

The columns that draw at random all consume the same stream of random numbers (the greedy one needs none). At the first step the favourite often owns a large slice of the interval, so the same $u$ lands on it in several columns at once. As soon as one column picks a different token, the texts diverge, and from then on they are unrelated. Lower temperatures make the favourite's slice wider, which is why the cold columns stay together longest.

And how does generation end? In one of three ways. The model can produce a special token: during pre-training every document ended with <|endoftext|>, so the base model learned to emit it when a story is finished, and our engine stops right there. We can choose our own stop tokens: in the next chapter the chat model will end every reply with <|end|>. Or we hit a length limit that we set ourselves. There is also a hard limit inside the model: Sprout can see at most 512 tokens. When a text grows longer, our engine keeps the last 256 tokens, re-reads them and carries on, forgetting the beginning.

Why it's fast: the key–value cache

Look again at generate: to choose token number 101 it runs the model over all 100 tokens; for token 102, over all 101, and so on. Almost all of that work is repeated. Can we avoid it?

Think back to attention (chapter 7). Every token, in every block, produces a query, a key and a value. Because of the causal mask, the key and the value of token $j$ depend only on tokens $1 \dots j$. Adding a new token at the end changes nothing about the past: all the old keys and values stay exactly as they were. So we can compute them once and keep them. That store is the key–value cache. For each new token the model computes only that token's own query, key and value, adds the key and value to the cache, and lets the query look over everything stored.

How much does it save? Say the prompt has $P$ tokens and we generate $n$ more. Without a cache, the step that produces token number $t+1$ has to push all $t$ tokens through the model, so the total is

$$\sum_{t=P}^{P+n-1} t \;=\; nP + \frac{n(n-1)}{2}\ \text{ token passes},$$

which grows like $n^2$. With the cache each token goes through the model once: $P + n$ passes. For a 14-token prompt and 30 new tokens that is 855 passes against 44. A pass through Sprout costs about 17.3 million multiply-adds (every weight is used once), plus the attention over the stored tokens. Here is the race, run for real in your browser:

The cache is not free: it costs memory. For every token Sprout stores a key and a value of 384 numbers each in every one of its 8 blocks: $2 \times 8 \times 384 = 6{,}144$ numbers per token. A full 512-token context takes 3.1 million numbers, 12.6 MB in float32, about three quarters of the 17.5 MB the weights themselves take. For large models with long contexts the cache runs into gigabytes, and a whole line of research exists just to shrink it.

Fitting in a browser: 8-bit weights and SIMD

Sprout has 17.31 million parameters. Stored as ordinary 32-bit floats, that is 69.2 MB, a heavy download for a web page. The file your browser loaded is 17.5 MB: every weight matrix is stored as 8-bit integers, with one scale factor per row. This is quantisation.

The recipe for one row of a matrix: find its largest absolute value and map it to 127. A signed byte holds the integers from −128 to 127, and we use the symmetric part of that range.

$$s_r = \frac{\max_c |W_{rc}|}{127}, \qquad q_{rc} = \operatorname{round}\!\left(\frac{W_{rc}}{s_r}\right) \in [-127, 127], \qquad W_{rc} \approx q_{rc}\, s_r.$$

Rounding moves each weight by at most half a step, $s_r / 2$. Why a scale per row and not one for the whole matrix? Because weights are not all alike: a single large weight would stretch a shared scale, and every other number in the matrix would be rounded coarsely. With one scale per row, an outlier only hurts its own row. This is the function from export.py that packed Sprout for the browser:

def quantize_rows(w): scale = np.abs(w).max(axis=1) / 127.0 scale[scale == 0] = 1.0 q = np.clip(np.round(w / scale[:, None]), -127, 127).astype(np.int8) return q, scale.astype(np.float32)

Let's test it on a random 384 × 384 matrix, the size of the output projection in Sprout's attention, with one planted outlier, and compare it with a single scale for the whole matrix:

rng = np.random.default_rng(1) W = rng.normal(0, 0.02, size=(384, 384)).astype(np.float32) W[7, 42] = 0.5 # one planted outlier q, s = quantize_rows(W) err_rows = np.abs(q * s[:, None] - W) s_all = np.abs(W).max() / 127.0 # one scale for the whole matrix err_all = np.abs(np.round(W / s_all) * s_all - W) print("bytes:", W.nbytes, "->", q.nbytes + s.nbytes) print(f"average error, a scale per row: {err_rows.mean():.1e}") print(f"average error, one scale for all: {err_all.mean():.1e}") print(f"in the row with the outlier: {err_rows[7].mean():.1e}")

A quarter of the bytes, and with a scale per row the error is several times smaller than with one shared scale: only the outlier's own row pays. On the real model the damage is tiny. We measured it on a checkpoint from Sprout's training run (step 3,000 of 10,070): validation loss 1.8388 nats in float32 and 1.8394 after 8-bit rounding with a scale per row, a difference of well under a thousandth. At 8 bits even one shared scale per matrix does as well; the scale per row earns its keep when bits get scarce. At 4 bits, rounding with a scale per row costs 0.06 nats, while one shared scale per matrix costs half a nat.

There is a second gift hidden in this format. When a row is multiplied by the input vector $x$, the scale doesn't depend on $c$ and can be taken out of the sum:

$$\sum_c W_{rc}\, x_c \;\approx\; \sum_c s_r\, q_{rc}\, x_c \;=\; s_r \sum_c q_{rc}\, x_c.$$

So the engine multiplies small integers by the input and applies one scale per row at the end. That inner sum is where nearly all of Sprout's time goes: 17.3 million multiply-adds per token. Our engine runs it in a tiny WebAssembly kernel that uses SIMD (single instruction, multiple data): it loads 16 weights at once into a 128-bit register, widens them to floats and multiplies four at a time. On a Mac with an M4 Pro (measured in Node) the kernel reaches about 14 billion multiply-adds per second, and Sprout writes some 500 tokens per second, almost six times faster than the same loop in plain JavaScript (about 90). A phone is several times slower, but still fast enough to read along.

How far can we squeeze? Try it yourself. The widget below loads Sprout's weights into this page and re-rounds all 17 million of them to fewer bits, then lets the squeezed model write and measures its surprise on a short story it has never seen.

On the final Sprout, six, five and even four bits change nothing you can see on this story: the surprise stays at about 1.77 nats per token, and the greedy text barely moves. (Over the whole validation text the step-3,000 checkpoint did lose 0.06 nats at 4 bits, so the damage is real, just small.) Three bits hurt a lot: the surprise jumps by more than half a nat, and the story loses its thread. At two bits, where every weight becomes −1, 0 or +1 times a scale, the model falls apart: the surprise shoots up to about 12 nats, and on the step-3,000 checkpoint 2-bit rounding pushed the validation loss to 10.6. Both are worse than guessing uniformly among all 8,192 tokens ($\ln 8192 \approx 9.0$).

Sprout right now

This is the same base model as at the end of chapter 10, but now you control how it chooses. Turn the temperature down for a calm, predictable story, turn it up for a wild one, and remember what happens at the edges. Sprout tells a decent story, but it still cannot hold a conversation: ask it a question and it will simply carry on writing, as if your question were a line from a story. In the next chapter we teach it to talk.