8.3 Sequences and language models

A language model is a probability distribution over the next word, given the words before it. Generating text is sampling from it, one word at a time. Everything else, from counting word pairs to a large transformer, is a better way to estimate that one distribution.

import collections
import glob
import re
import warnings

import matplotlib.pyplot as plt
import numpy as np
import statsmodels.api as sm
from matplotlib.axes import Axes
from sklearn.linear_model import LinearRegression
from sklearn.neural_network import MLPRegressor
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

from statly import style

style.use()
rng = np.random.default_rng(1)
warnings.filterwarnings("ignore")

sun = sm.datasets.sunspots.load_pandas().data  # sunspots per year, 1700 to 2008


def prose(source: str) -> str:
    """The running text and figure captions of a chapter, without code, formulas or tables."""
    captions = " ".join(re.findall(r'#\| fig-cap: "(.*)"', source))
    text = re.sub(r"```.*?```", "", source, flags=re.S)
    text = re.sub(r"\$\$.*?\$\$", "", text, flags=re.S)
    text = re.sub(r"\$[^$\n]*\$", "", text)
    text = re.sub(r"^(:::|#|\|).*$", "", text, flags=re.M)
    return text + " " + captions


# The training text: every chapter of this book, as lower-case words.
chapters = {path: open(path).read() for path in sorted(glob.glob("[0-9]*.qmd"))}
book_words = {
    path: re.findall(r"[a-z]+(?:'[a-z]+)?", prose(text).lower())
    for path, text in chapters.items()
}
all_words = [w for words in book_words.values() for w in words]

8.3.1 Sequences: recurrent networks

Time series, sound and text come in order. A recurrent network reads a sequence one step at a time and carries a hidden state forward, a running summary of what it has seen:

$$h_t = g\big(W x_t + U h_{t-1} + b\big), \qquad \hat y_t = \beta_0 + \beta^\top h_t$$

The same weights W and U are used at every step, just as a convolution uses the same filter at every position. Unrolled in time, the network is a long chain:

def plot_unrolled(ax: Axes, steps: int = 4) -> None:
    """Boxes for inputs, hidden states and outputs, with arrows between them."""
    for t in range(steps):
        x = t * 1.2
        for y, label, color in (
            (0, f"x{t + 1}", style.BLUE),
            (1, f"h{t + 1}", style.GREEN),
            (2, f"ŷ{t + 1}", style.CORAL),
        ):
            ax.add_patch(plt.Rectangle((x - 0.3, y - 0.25), 0.6, 0.5, color=color, alpha=0.6))
            ax.text(x, y, label, ha="center", va="center", fontsize=10)
        ax.annotate(
            "",
            xy=(x, 0.75),
            xytext=(x, 0.25),
            arrowprops=dict(arrowstyle="->", color=style.INK),
        )
        ax.annotate(
            "",
            xy=(x, 1.75),
            xytext=(x, 1.25),
            arrowprops=dict(arrowstyle="->", color=style.INK),
        )
        if t:
            ax.annotate(
                "",
                xy=(x - 0.3, 1),
                xytext=(x - 0.9, 1),
                arrowprops=dict(arrowstyle="->", color=style.INK),
            )
    ax.set(xlim=(-0.6, steps * 1.2 - 0.6), ylim=(-0.5, 2.5), aspect="equal")
    ax.set_axis_off()


LAGS = 10
series = sun.SUNACTIVITY.to_numpy()
windows = np.column_stack([series[i : len(series) - LAGS + i] for i in range(LAGS)])
next_year = series[LAGS:]
years = sun.YEAR.to_numpy()[LAGS:]
train = years < 1960
network = make_pipeline(
    StandardScaler(),
    MLPRegressor(hidden_layer_sizes=(20,), alpha=1.0, max_iter=5000, random_state=0),
).fit(windows[train], next_year[train])
linear = LinearRegression().fit(windows[train], next_year[train])


def plot_forecasts(ax: Axes) -> None:
    """One-step-ahead forecasts on the test years."""
    test = ~train
    ax.plot(years[test], next_year[test], color=style.GREY, lw=1.5, label="observed")
    for model, color, name in (
        (linear, style.BLUE, "linear"),
        (network, style.CORAL, "network"),
    ):
        forecast = model.predict(windows[test])
        rmse = np.sqrt(np.mean((next_year[test] - forecast) ** 2))
        ax.plot(years[test], forecast, color=color, lw=1.4, label=f"{name}: error {rmse:.1f}")
    ax.legend(fontsize=8, loc="upper right")
    ax.set(xlabel="year", ylabel="sunspots")


fig = plt.figure(figsize=(10, 3.1))
grid = fig.add_gridspec(1, 2, width_ratios=[1, 1.6])
plot_unrolled(fig.add_subplot(grid[0]))
plot_forecasts(fig.add_subplot(grid[1]))
fig.tight_layout()
plt.show()

Figure 8.3.1: Left: a recurrent network unrolled over four steps. Each step takes the new input (bottom) and the previous hidden state (left arrow), and gives a new hidden state and a prediction (top). The same weights are used at every step. Right: forecasting next year’s sunspots from the previous ten years, tested on 1960 to 2008, which the models never saw. A network on the ten-year window (red) and a linear autoregression (blue, Chapter 9.2) both follow the cycle; the network’s error is a little lower.

(The network in the figure reads a fixed window instead of carrying a state, the simplest stand-in for a recurrent network.) Plain recurrent networks forget quickly; LSTM and GRU units add gates that decide what to keep and what to forget, and can carry information over long stretches. For many time series, though, the models of Chapter 9.2 are as accurate and far easier to check.

8.3.2 Text as numbers: the bag of words

A computer needs numbers, so text must be turned into them. The simplest way is the bag of words: one column per word of the vocabulary, counting how often each word occurs in a document, ignoring the order. A plain network or a lasso logistic regression (Chapter 6.4) on these columns is a strong baseline for sorting documents.

SENTENCES = [
    "a good film, a really good cast",
    "not good, a bad film",
    "bad cast, bad film, not really bad",
]
words = sorted({w.strip(",") for s in SENTENCES for w in s.split()})
counts = np.array(
    [[sum(w.strip(",") == word for w in s.split()) for word in words] for s in SENTENCES]
)


def plot_bag(ax: Axes) -> None:
    """The count table as a grid of numbers."""
    ax.imshow(counts > 0, cmap="Blues", alpha=0.4, aspect="auto")
    for i in range(counts.shape[0]):
        for j in range(counts.shape[1]):
            ax.text(j, i, str(counts[i, j]), ha="center", va="center", fontsize=10)
    ax.set_xticks(range(len(words)), words)
    ax.set_yticks(range(len(SENTENCES)), [f'"{s}"' for s in SENTENCES], fontsize=8.5)
    ax.grid(False)
    ax.tick_params(length=0)


fig, ax = plt.subplots(figsize=(9, 2))
plot_bag(ax)
plt.show()

Figure 8.3.2: Three made-up review sentences turned into a bag of words: one column per word, the count of each word in each sentence. The order of the words is lost, which is the method’s weakness (‘not good, bad’ and ‘not bad, good’ look the same).

The rest of this chapter follows the road from such counts to large language models, using one small dataset throughout: the text of this book.

8.3.3 Tokens

A language model works on tokens, the units it reads and writes. Whole words make a huge vocabulary and cannot handle a word never seen before; single characters make every text very long. Modern models use subword tokens, found by byte pair encoding: start from single characters, and repeatedly merge the pair of symbols that occurs most often next to each other into a new symbol. Common words end up as one token, rare words as a few pieces.

SENTENCE = "the bootstrap resamples households to estimate uncertainty"
word_counts = collections.Counter(all_words)


def learn_merges(rounds: int) -> list[tuple[str, str]]:
    """Byte pair encoding: the most frequent adjacent pair is merged, `rounds` times."""
    spelled = {tuple(word) + ("_",): count for word, count in word_counts.items()}
    merges = []
    for _ in range(rounds):
        pairs = collections.Counter()
        for symbols, count in spelled.items():
            for pair in zip(symbols, symbols[1:]):
                pairs[pair] += count
        best = max(pairs, key=pairs.get)
        merges.append(best)
        spelled = {merge_pair(symbols, best): count for symbols, count in spelled.items()}
    return merges


def merge_pair(symbols: tuple[str, ...], pair: tuple[str, str]) -> tuple[str, ...]:
    """Replace every occurrence of `pair` in `symbols` by one joined symbol."""
    out, i = [], 0
    while i < len(symbols):
        if symbols[i : i + 2] == pair:
            out.append(symbols[i] + symbols[i + 1])
            i += 2
        else:
            out.append(symbols[i])
            i += 1
    return tuple(out)


def encode(word: str, merges: list[tuple[str, str]]) -> list[str]:
    """Split a word into subword tokens by applying the merges in order."""
    symbols = tuple(word) + ("_",)
    for pair in merges:
        symbols = merge_pair(symbols, pair)
    return [s.replace("_", "") for s in symbols if s != "_"]


def draw_tokens(ax: Axes, tokens: list[str], y: float, label: str) -> None:
    """Tokens as coloured boxes in a row, sized by their length."""
    x = 0.0
    palette = [style.BLUE, style.CORAL, style.GREEN, style.GOLD, style.PURPLE]
    for i, token in enumerate(tokens):
        width = 0.16 * max(len(token), 1) + 0.12
        ax.add_patch(
            plt.Rectangle((x, y - 0.3), width - 0.04, 0.6, color=palette[i % 5], alpha=0.35)
        )
        ax.text(
            x + width / 2 - 0.02,
            y,
            token if token != " " else "␣",
            ha="center",
            va="center",
            fontsize=8.5,
        )
        x += width
    ax.text(-0.15, y, f"{label}\n{len(tokens)} tokens", ha="right", va="center", fontsize=8.5)


merges = learn_merges(400)
subwords = [token for word in SENTENCE.split() for token in encode(word, merges)]
fig, ax = plt.subplots(figsize=(11, 2.6))
draw_tokens(ax, list(SENTENCE), 2, "characters")
draw_tokens(ax, subwords, 1, "subwords")
draw_tokens(ax, SENTENCE.split(), 0, "words")
ax.set(xlim=(-1.6, 17.2), ylim=(-0.5, 2.5))
ax.set_axis_off()
plt.show()

Figure 8.3.3: One sentence split three ways. Top: characters. Middle: subword tokens, from 400 rounds of byte pair encoding learned on the text of this book; frequent words are single tokens, rarer words are built from pieces. Bottom: whole words.

Each token gets a number, its position in the vocabulary, and from then on the model only sees numbers. Large models have vocabularies of 50,000 to 200,000 tokens.

8.3.4 Predicting the next word

Any text can be built one word at a time, and so can its probability. By the multiplication rule of probability (Chapter 2.4),

$$P(w_1, w_2, \dots, w_n) = P(w_1)\, P(w_2 \mid w_1)\, P(w_3 \mid w_1, w_2) \cdots P(w_n \mid w_1, \dots, w_{n-1})$$

so a model of the next word, given everything before it, is a model of all text. The simplest estimate counts. An n-gram model looks only at the last n − 1 words and uses the share of times each word followed them in the training text:

$$\hat P(w \mid \text{previous words}) = \frac{\text{count}(\text{previous words}, w)}{\text{count}(\text{previous words})}$$

def next_word_counts(context: tuple[str, ...]) -> collections.Counter:
    """How often each word follows `context` anywhere in the book."""
    n = len(context)
    following = collections.Counter()
    for words in book_words.values():
        for i in range(n, len(words)):
            if tuple(words[i - n : i]) == context:
                following[words[i]] += 1
    return following


def plot_next(ax: Axes, context: tuple[str, ...], color: str) -> None:
    """The eight most likely next words after `context`, as probabilities."""
    following = next_word_counts(context)
    total = sum(following.values())
    top = following.most_common(8)
    ax.barh(range(len(top)), [c / total for _, c in top], color=color, alpha=0.7)
    ax.set_yticks(range(len(top)), [w for w, _ in top])
    ax.invert_yaxis()
    ax.xaxis.set_major_formatter(style.PERCENT)
    ax.set(
        title=f"after “{' '.join(context)}” ({total} times)",
        xlabel="probability of the next word",
    )


fig, axes = style.subplots(2, height=3, width=10)
plot_next(axes[0], ("standard",), style.BLUE)
plot_next(axes[1], ("the", "standard"), style.CORAL)
fig.tight_layout()
plt.show()

Figure 8.3.4: Next-word probabilities estimated by counting in the text of this book. Left: after the one word ‘standard’. Right: after the two words ‘the standard’. Each is a probability distribution over the vocabulary; only the eight most likely words are shown.

8.3.5 More context, and where counting breaks down

A longer context should predict better: “the standard” says more than “standard”. But the longer the context, the more often it has never been seen at all, and then counting has nothing to go on. Models are judged on held-out text by their perplexity, the exponential of the average negative log probability they gave to each actual next word:

$$\text{perplexity} = \exp\Big(-\frac{1}{N} \sum_{t=1}^{N} \log \hat P(w_t \mid w_1, \dots, w_{t-1})\Big)$$

It reads as “the model was as unsure as if it chose among this many words”. Lower is better. The average inside is the cross-entropy loss, the same negative log-likelihood that logistic regression minimises (Chapter 4.4).

paths = sorted(book_words)
held_out = [p for i, p in enumerate(paths) if i % 7 == 3]
train_text = [w for p in paths if p not in held_out for w in book_words[p]]
test_text = [w for p in held_out for w in book_words[p]]
known = {w for w, c in collections.Counter(train_text).items() if c >= 2}


def known_only(words: list[str]) -> list[str]:
    """Words outside the training vocabulary become one shared 'unknown' token."""
    return [w if w in known else "<unknown>" for w in words]


class NGram:
    """An n-gram model with Witten-Bell smoothing: unseen contexts fall back to shorter ones."""

    def __init__(self, words: list[str], n: int) -> None:
        self.n, self.size = n, len(words)
        self.vocabulary = len(set(words)) + 1
        self.counts = [collections.Counter() for _ in range(n + 1)]
        self.context = [collections.Counter() for _ in range(n + 1)]
        self.types = [collections.Counter() for _ in range(n + 1)]
        for k in range(1, n + 1):
            for i in range(k - 1, len(words)):
                history = tuple(words[i - k + 1 : i])
                if self.counts[k][history + (words[i],)] == 0:
                    self.types[k][history] += 1  # one more distinct follower of this history
                self.counts[k][history + (words[i],)] += 1
                self.context[k][history] += 1

    def prob(self, history: tuple[str, ...], word: str, k: int | None = None) -> float:
        """P(word | the last k − 1 words of history), mixed with the (k − 1)-gram estimate."""
        k = self.n if k is None else k
        if k == 1:
            return (self.counts[1][(word,)] + 1) / (self.size + self.vocabulary)
        shorter = self.prob(history, word, k - 1)
        h = tuple(history[len(history) - k + 1 :])
        seen, distinct = self.context[k][h], self.types[k][h]
        if seen == 0:
            return shorter
        return (self.counts[k][h + (word,)] + distinct * shorter) / (seen + distinct)

    def perplexity(self, words: list[str]) -> float:
        """exp of the mean negative log probability of each word given the ones before it."""
        logs = [
            np.log(self.prob(tuple(words[max(0, i - 4) : i]), words[i]))
            for i in range(4, len(words))
        ]
        return float(np.exp(-np.mean(logs)))


ORDERS = range(1, 6)
train_known, test_known = known_only(train_text), known_only(test_text)
models = {n: NGram(train_known, n) for n in ORDERS}


def seen_share(n: int) -> float:
    """Share of held-out positions whose n − 1 word context occurred in training."""
    if n == 1:
        return 1.0
    hits = [
        models[n].context[n][tuple(test_known[i - n + 1 : i])] > 0
        for i in range(4, len(test_known))
    ]
    return float(np.mean(hits))


fig, axes = style.subplots(2, height=3, width=10)
axes[0].plot(
    ORDERS, [models[n].perplexity(test_known) for n in ORDERS], "o-", color=style.CORAL, lw=2
)
axes[0].set(
    xticks=list(ORDERS),
    xlabel="n (words in the n-gram)",
    ylabel="perplexity on held-out text",
    title="Prediction",
)
axes[1].plot(ORDERS, [seen_share(n) for n in ORDERS], "o-", color=style.BLUE, lw=2)
axes[1].yaxis.set_major_formatter(style.PERCENT)
axes[1].set(
    xticks=list(ORDERS),
    xlabel="n (words in the n-gram)",
    ylabel="contexts seen in training",
    title="Coverage",
    ylim=(0, 1.05),
)
fig.tight_layout()
plt.show()

Figure 8.3.5: N-gram models trained on most chapters of this book and tested on the rest. Left: perplexity on the held-out chapters. Going from one word (no context) to two (one word of context) helps a lot; longer contexts make it worse, because most of them never occurred in training (right: the share of held-out contexts seen in training), and the model falls back on shorter ones. Unseen words of the held-out text are counted as one ‘unknown word’ token.

This is the curse of dimensionality (Chapter 1.3) in a new form: the number of possible contexts grows explosively with their length, and no text is long enough to see them all. Counting treats “the standard deviation” and “the standard error” as unrelated contexts. A neural model instead maps contexts to vectors, so that similar contexts give similar predictions, and it can learn from one what to expect in the other.

8.3.6 Embeddings: words as points

Words get their meaning from the company they keep. Count how often each pair of words occurs near each other, turn the counts into a table of how much more often than chance they meet, and compress each word’s row of that table to a few dozen numbers with the singular value decomposition that underlies principal components (Chapter 7.1). The result is an embedding: one vector per word, with words used in similar ways close together.

frequent = [w for w, _ in word_counts.most_common(1500)]
index = {w: i for i, w in enumerate(frequent)}
together = np.zeros((len(frequent), len(frequent)))
for words in book_words.values():
    for i, w in enumerate(words):
        if w in index:
            for j in range(max(0, i - 4), min(len(words), i + 5)):
                if j != i and words[j] in index:
                    together[index[w], index[words[j]]] += 1
total = together.sum()
p_word = together.sum(axis=1, keepdims=True) / total
p_context = (together.sum(axis=0, keepdims=True) / total) ** 0.75  # dampens very common words
p_context /= p_context.sum()
with np.errstate(divide="ignore"):
    association = np.maximum(np.log(together / total / (p_word * p_context)), 0)
U, singular, _ = np.linalg.svd(association)
vectors = U[:, :50] * np.sqrt(singular[:50])
vectors /= np.linalg.norm(vectors, axis=1, keepdims=True)

GROUPS = {
    "colours": (["red", "blue", "green", "grey"], style.CORAL),
    "models": (
        ["regression", "logistic", "linear", "tree", "forest", "network", "boosting"],
        style.BLUE,
    ),
    "spread": (["variance", "deviation", "spread", "standard", "error"], style.GREEN),
    "households": (["income", "food", "households"], style.GOLD),
}
NEIGHBOURS = ["network", "regression", "variance", "income", "red"]


def plot_map(ax: Axes) -> None:
    """Chosen words placed by the first two principal components of their vectors."""
    chosen = [w for words, _ in GROUPS.values() for w in words if w in index]
    points = vectors[[index[w] for w in chosen]]
    centred = points - points.mean(axis=0)
    xy = centred @ np.linalg.svd(centred, full_matrices=False)[2][:2].T
    colour = {w: c for words, c in GROUPS.values() for w in words}
    for (x, y), w in zip(xy, chosen):
        ax.scatter(x, y, s=25, color=colour[w])
        ax.annotate(
            w, (x, y), xytext=(4, 3), textcoords="offset points", fontsize=8.5, color=colour[w]
        )
    ax.set(
        xticks=[],
        yticks=[],
        xlabel="component 1",
        ylabel="component 2",
        title="Word vectors in 2D",
    )


def neighbour_table(ax: Axes) -> None:
    """The five nearest words to each chosen word, as a small table."""
    ax.set_axis_off()
    ax.set_title("Nearest words")
    for row, w in enumerate(NEIGHBOURS):
        similarity = vectors @ vectors[index[w]]
        nearest = [frequent[i] for i in np.argsort(-similarity)[1:6]]
        ax.text(0, 1 - row * 0.2, w, fontsize=9.5, fontweight="bold", va="top")
        ax.text(0.28, 1 - row * 0.2, ", ".join(nearest), fontsize=9, va="top")


fig, axes = style.subplots(2, height=3.6, width=10.5)
plot_map(axes[0])
neighbour_table(axes[1])
fig.tight_layout()
plt.show()

Figure 8.3.6: Left: word vectors learned from the text of this book alone (counts within four words of each other, positive pointwise mutual information, compressed to 50 dimensions), shown in 2D by their first two principal components. Words of a kind gather: colours, models, measures of spread, the households data. Right: the nearest words to a few words, by the cosine of the angle between their vectors.

A network learns its embeddings as part of training, as the first layer, and they come out far richer than these, trained on billions of words. Directions in the space can carry meaning: the difference between the vectors for two related words often points the same way for other pairs with the same relation.

8.3.7 Attention

A recurrent network squeezes everything it has read into one hidden state. Attention lets each position look back at every earlier position directly and decide how much each one matters. Each token’s vector is turned into a query (what am I looking for?), a key (what do I offer?) and a value (what do I pass on?), by three learned matrices. The weights are a softmax of how well each query matches each key:

$$\text{attention}(Q, K, V) = \operatorname{softmax}\!\Big(\frac{Q K^\top}{\sqrt{d}}\Big) V$$

Row i of the softmax says how much token i draws on each other token; the output for token i is the weighted average of their values. In a language model, a token may only look backwards, so the weights above the diagonal are set to zero.

PHRASE = "the network learned filters from the digits and the regression learned coefficients from the households"
tokens = [w for w in PHRASE.split()]
X = np.array([vectors[index[w]] if w in index else np.zeros(50) for w in tokens])
scores = X @ X.T / np.sqrt(X.shape[1]) * 8  # scaled up so the softmax is not flat
scores[np.triu_indices(len(tokens), 1)] = -np.inf  # no looking ahead
weights = np.exp(scores - scores.max(axis=1, keepdims=True))
weights /= weights.sum(axis=1, keepdims=True)

fig, ax = plt.subplots(figsize=(6.4, 5.6))
ax.imshow(weights, cmap="Blues", vmin=0, vmax=1)
ax.set_xticks(range(len(tokens)), tokens, rotation=60, ha="right", fontsize=8.5)
ax.set_yticks(range(len(tokens)), tokens, fontsize=8.5)
ax.set(xlabel="looks at", ylabel="token")
ax.grid(False)
plt.show()

Figure 8.3.7: Attention weights for a sentence, each row a softmax over the words before it (and itself). Here the queries and keys are simply the word vectors of Figure 2.4.3.6, so a word attends to earlier words used in similar ways; a trained transformer learns separate query and key matrices, and many heads, each picking out a different kind of relation.

8.3.8 Transformers and large language models

A transformer stacks blocks of attention and small feed-forward networks (Chapter 8.1), each with a shortcut that adds the block’s input to its output. The token embeddings, plus a signal for each token’s position, enter at the bottom; at the top a softmax gives the probability of every token in the vocabulary as the next one.

def box(ax: Axes, x: float, y: float, w: float, h: float, text: str, color: str) -> None:
    """A labelled, filled box."""
    ax.add_patch(plt.Rectangle((x - w / 2, y - h / 2), w, h, color=color, alpha=0.35))
    ax.text(x, y, text, ha="center", va="center", fontsize=8.5)


def arrow(ax: Axes, start: tuple[float, float], end: tuple[float, float]) -> None:
    """An arrow between two points."""
    ax.annotate("", xy=end, xytext=start, arrowprops=dict(arrowstyle="->", color=style.INK))


def plot_transformer(ax: Axes) -> None:
    """The stack from tokens to next-token probabilities."""
    box(ax, 0, 0, 3, 0.5, "tokens: “the standard”", style.GREY)
    box(ax, 0, 0.9, 3, 0.5, "embeddings + positions", style.GREEN)
    ax.add_patch(plt.Rectangle((-1.8, 1.45), 3.6, 1.9, fill=False, ls="--", color=style.INK))
    ax.text(1.95, 2.4, "× N\nblocks", fontsize=8.5, va="center")
    box(ax, 0, 1.85, 3, 0.5, "attention (looks back)", style.BLUE)
    box(ax, 0, 2.85, 3, 0.5, "feed-forward network", style.BLUE)
    box(ax, 0, 3.9, 3, 0.5, "softmax over the vocabulary", style.CORAL)
    box(ax, 0, 4.8, 3, 0.5, "P(next = “deviation”) = …", style.GOLD)
    for y0, y1 in ((0.25, 0.65), (1.15, 1.6), (2.1, 2.6), (3.1, 3.65), (4.15, 4.55)):
        arrow(ax, (0, y0), (0, y1))
    ax.set(xlim=(-2.2, 2.9), ylim=(-0.4, 5.2))
    ax.set_axis_off()


def plot_stages(ax: Axes) -> None:
    """Pretraining, instruction tuning and preference tuning as three boxes."""
    stages = [
        (
            "1  Pretraining",
            "predict the next token in\ntrillions of tokens of text",
            style.BLUE,
        ),
        (
            "2  Instruction tuning",
            "continue training on examples\nof questions and good answers",
            style.GREEN,
        ),
        (
            "3  Preference tuning",
            "people (or models) rank answers;\nthe model is pushed towards\nthe preferred ones",
            style.CORAL,
        ),
    ]
    for i, (title, text, color) in enumerate(stages):
        y = 4 - i * 1.7
        box(ax, 0, y, 3.6, 1.2, "", color)
        ax.text(-1.7, y + 0.35, title, fontsize=9, fontweight="bold", va="center")
        ax.text(-1.7, y - 0.15, text, fontsize=8, va="center")
        if i:
            arrow(ax, (0, y + 1.1), (0, y + 0.6))
    ax.set(xlim=(-2, 2), ylim=(0, 4.8))
    ax.set_axis_off()


fig, axes = style.subplots(2, height=4.2, width=10)
plot_transformer(axes[0])
plot_stages(axes[1])
fig.tight_layout()
plt.show()

Figure 8.3.8: Left: a transformer language model. Tokens become vectors, pass through a stack of blocks (attention, then a feed-forward network, each with a shortcut), and the top gives next-token probabilities. Large models stack around 100 such blocks. Right: the three stages that turn it into an assistant.

A large language model is such a transformer with billions of weights, pretrained by minimising the cross-entropy of the next token over trillions of tokens of text. Nothing in that objective mentions facts, grammar or reasoning; they are learned because they help predict the next token. Instruction tuning and preference tuning (often called RLHF, reinforcement learning from human feedback) then turn the text predictor into an assistant that answers questions.

Two statistical patterns govern how well this works. More data and more weights lower the held-out loss smoothly, close to a power law, which is what made ever larger models worth building. Even the counting models of this chapter show the first half of that pattern:

SHARES = np.array([0.1, 0.2, 0.35, 0.5, 0.7, 1.0])

fig, ax = plt.subplots(figsize=(7.5, 3.2))
for n, color in ((2, style.BLUE), (3, style.CORAL)):
    sizes = (SHARES * len(train_known)).astype(int)
    scores_n = [NGram(train_known[:size], n).perplexity(test_known) for size in sizes]
    ax.plot(sizes, scores_n, "o-", color=color, lw=2, label=f"{n}-gram")
ax.set(xscale="log", xlabel="training words (log scale)", ylabel="held-out perplexity")
ax.legend(fontsize=8)
plt.show()

Figure 8.3.9: Perplexity on the held-out chapters of a two-word and a three-word model trained on growing shares of the training chapters (log scale across). More text steadily lowers the error. The vocabulary is fixed from the full training text so that every model is scored on the same tokens.

8.3.9 Sampling, temperature and confident nonsense

To write, the model draws the next token from its distribution, appends it, and repeats. A temperature T reshapes the distribution before drawing:

$$P_T(w) = \frac{P(w)^{1/T}}{\sum_v P(v)^{1/T}}$$

T below 1 sharpens it towards the most likely tokens (safer, more repetitive); T above 1 flattens it (more varied, more errors); T near 0 always takes the top token.

full = NGram(known_only(all_words), 3)
context = ("the", "standard")
next_probs = {w: full.prob(context, w) for w in known}
candidates = sorted(next_probs, key=next_probs.get, reverse=True)[:8]
base = np.array([full.prob(context, w) for w in candidates])

fig, axes = style.subplots(3, height=2.8, width=10.5)
for ax, temperature, color in zip(
    axes, (0.5, 1.0, 2.0), (style.BLUE, style.GREEN, style.CORAL)
):
    reshaped = base ** (1 / temperature)
    ax.barh(range(len(candidates)), reshaped / reshaped.sum(), color=color, alpha=0.7)
    ax.set_yticks(range(len(candidates)), candidates, fontsize=8.5)
    ax.invert_yaxis()
    ax.xaxis.set_major_formatter(style.PERCENT)
    ax.set(title=f"T = {temperature}", xlabel="probability")
fig.tight_layout()
plt.show()

Figure 8.3.10: The next-word distribution after “the standard” from the three-word model, at three temperatures. Low temperature puts nearly everything on the favourite; high temperature spreads probability onto unlikely words.

Generated text from the three-word model of this book, started from the same two words:

vocabulary = sorted(known) + ["<unknown>"]


def generate(start: tuple[str, str], temperature: float, length: int, seed: int) -> str:
    """Sample `length` words, one at a time, from the smoothed three-word model."""
    local = np.random.default_rng(seed)
    words = list(start)
    for _ in range(length):
        p = np.array([full.prob(tuple(words[-2:]), w) for w in vocabulary]) ** (1 / temperature)
        nxt = vocabulary[local.choice(len(vocabulary), p=p / p.sum())]
        words.append(nxt if nxt != "<unknown>" else "…")
    return " ".join(words)


for temperature in (0.5, 1.0):
    print(
        f"> **T = {temperature}:** {generate(('the', 'standard'), temperature, 30, seed=3)}\n"
    )

T = 0.5: the standard deviation has variance so four times the area under it is the residual sum of their own rate of a and badly on hospital b whose patients lie where a

T = 1.0: the standard deviation describes read the bayes error rate chapter line are pulled more towards the prior within each part the final a … for error smaller so the standard deviation of

The output has the local texture of the book and no meaning beyond a few words. A large model’s output is fluent over pages, but it is produced the same way: by sampling plausible continuations. That is why it can state a wrong fact as confidently as a right one, a hallucination: the model learned what answers look like, and a plausible wrong answer can look exactly like a right one. The statistical habits of this book apply: check answers against a source, evaluate on held-out questions the model cannot have seen in training (test data leak easily into web-scale training text), and treat a fluent answer as a prediction, not a measurement.

8.3.10 In one picture

• A language model gives $P(\text{next token} \mid \text{previous tokens})$; the chain rule turns that into a probability for any text.
• Text is split into tokens, usually subwords from byte pair encoding.
• N-gram models count; longer contexts run out of data, which neural models solve by mapping similar contexts to similar vectors.
• Embeddings place words used alike close together.
• Attention: $\operatorname{softmax}(QK^\top/\sqrt{d})V$, each token looking back at all earlier ones; transformers stack attention and feed-forward blocks.
• LLMs are pretrained on next-token cross-entropy, then instruction and preference tuned; they are judged by held-out perplexity and held-out tasks.
• Generation is sampling, shaped by the temperature; fluent output is not checked fact.

Next: models that create new data, from autoencoders to diffusion.