Choosing the next word

Sprout has finished training on 330 million tokens and can now score possible continuations. After “The little”, for example, it may assign one probability to “ girl”, another to “ boy” and another to “ cat”. Which one appears on the page? A decoding algorithm makes that choice. It can always take the highest-scoring token, or draw at random according to the probabilities. The second approach is sampling, and the code that performs it is a sampler.

With exactly the same weights, Sprout can repeat familiar phrases, try less expected continuations or lose its thread. The token-selection settings make a substantial difference. We will take them one at a time, then look at how a key–value cache and 8-bit weights make it practical to run the model in a browser.

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.

At the edges of the scale, as $T \to 0^+$, probability concentrates on the token with the highest logit. If several tokens tie for the maximum, the limiting distribution remains uniform over them. In software, “temperature 0” usually selects one maximum directly, without dividing by zero: this is greedy decoding. As $T \to \infty$, every finite logit divided by $T$ tends to zero; without further filters, the distribution approaches uniform over the vocabulary. At $T = 1$ we retain the model's original distribution.

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 the effective count often approaches one or two; this does not literally count nonzero probabilities; 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.

Repetition is part of language: names, phrases and topics recur. A model learns to use that pattern, and its own repeated output can make another repetition increasingly likely. Greedy decoding always follows the current leader, leaving it particularly vulnerable to this feedback. Researchers have studied copying circuits such as induction heads in transformers, but attributing a particular Sprout loop to them would need a separate investigation. Random sampling offers a chance to take another branch, though it does not guarantee an escape.

Greedy decoding does not inevitably loop. Try the third beginning: Sprout can finish the story about Lily and her red ball without an obvious repetition cycle. The duck and mouse examples show other possibilities. Whether to use greedy decoding or sampling depends on the task: stories often benefit from variety, while some tasks favour reproducible output. There is no universal rule against the greedy option.

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.

The tail can contain poorly learned continuations and obvious mistakes, but also rare, perfectly appropriate words. Truncation cannot tell them apart: probability is only a rough guide. Why use it, then? A badly chosen token becomes part of the next input, and the model may continue the nonsense it has begun. Cutting the tail can reduce that risk at the cost of some variety.

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$, only the first forty candidates remain. But no single $k$ suits every distribution. After "Once upon a" there is essentially one sensible token, " time", and top-40 keeps another 39 candidates, many of them unlikely here. 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.

Sprout gets its random numbers from a pseudo-random generator: an algorithm that produces a long sequence from an initial state, the seed. With the same prompt, weights, settings and computations, the same seed reproduces the same token choices. The four-temperature comparison uses this property: keep the seed and repeat the run, then change it. In other engines or on other devices, a seed alone may not suffice, because numerical differences and the order of calculations can also affect the result. “Regenerate” usually starts a fresh random draw, but does not promise a different answer.

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: the pre-training stream separated documents 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 quadratically with the number of new tokens. A cache lets us process the prompt once, then feed in one new token at a time. Producing $n$ tokens takes $P+n-1$ position passes; processing the final emitted token as well makes that $P+n$. Our widget counts the latter: 44 passes for a 14-token prompt and 30 new tokens, rather than 855. The new token's attention still reads the stored context, so caching does not make the entire computation constant in sequence length. Try the race 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.$$

The engine can therefore take the row's scale factor outside the inner sum. A small WebAssembly kernel does the main work with SIMD (single instruction, multiple data): it loads 16 eight-bit weights, converts them to floats and multiplies four at a time. A recorded Node benchmark on a Mac with an M4 Pro reached about 14 billion multiply-adds per second in the kernel and roughly 500 generated tokens per second, compared with 90 for plain JavaScript. Those figures belong to that device and test. Your browser, context length and phone hardware affect the speed you see.

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 has not yet been fine-tuned for the assistant role: 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.