Predicting the next word

The objective: predict the next piece

By the chain rule, the probability of a text $x_1,\dots,x_n$ factors into next-piece predictions:

A language model is any model of $P(x_t\mid x_{<t})$. GPT-1 is trained by maximizing $\sum_t \ln P_\theta(x_t\mid x_{t-k},\dots,x_{t-1})$ over a corpus, the “standard language modeling objective” (equation 1 of the paper), where $k$ is the size of the context window. It is equivalent to minimizing the cross-entropy $L=-\tfrac1n\sum_t \ln p_\theta(x_t\mid x_{<t})$. Generating is the reverse trip: $x_t\sim P(\cdot\mid x_{<t})$, one piece at a time (autoregression).

Counting: the n-gram model

With a window of $n-1$ pieces, the maximum-likelihood estimate is the relative frequency:

where $N(c,w)$ counts how many times $w$ followed context $c$. In the lesson ($n=2$, bigrams), “is” appears $N=8$ times and $N(\text{is},\text{blue})=4$, so $P=4/8=0.5$. Shannon (1948) already generated text this way: his “second-order word approximation” picks each word according to the transition frequencies from the previous one. The limit is sparsity: with a vocabulary of size $V$ there are up to $V^{\,n-1}$ contexts, almost all rare or never seen, and the table does not generalize. Bengio et al. (2003) proposed replacing it with a network that learns a vector per word: if “dog” and “cat” have similar vectors, what is learned from one sentence carries over to the other.

Softmax and temperature

A network produces a score $z_i$ (a logit) for every piece $i$ of the vocabulary. The probability is

With $T=1$ it is the ordinary softmax. Adding the same constant to every score changes nothing; that is why the lesson’s scores, $z_i=\ln c_i-\ln 2$ (counts $4,2,1,1$; rounded on screen to $0.7,\,0,\,-0.7,\,-0.7$), give exactly count ÷ total. In general, with $z_i=\ln c_i$ we get $p_i=c_i^{1/T}/\sum_j c_j^{1/T}$: with $T=0.5$, blue gets $16/22\approx0.727$; with $T=2$, $2/(2+\sqrt2+1+1)\approx0.369$; with $T=0.25$, $256/274\approx0.934$. As $T\to0$ all the mass goes to the highest-scoring piece (greedy); as $T\to\infty$ the distribution tends to uniform. The name comes from statistical physics: the Boltzmann distribution has this same form, with the temperature dividing. It plays this role in, for example, knowledge distillation (Hinton et al., 2015).

Top-k and top-p

After the softmax one can truncate. Top-k (Fan et al., 2018, who sample among the $k=10$ most likely) keeps the $k$ most likely pieces. Top-p, or nucleus sampling (Holtzman et al., 2020), keeps the smallest set $V^{(p)}$ whose probability adds up to at least $p$:

Then it shares the probability out again, $q_i=p_i/\sum_{j\in V^{(p)}}p_j$, and samples with $q$. Holtzman et al. observed that decoding methods that maximize probability produce text that is “bland, incoherent, or gets stuck in repetitive loops” (the greedy loop in the lesson), and that truncating the unreliable tail, without going all the way to greedy, yields text closer to human writing.

In code

The lesson’s model: count pairs, softmax with temperature, top-p, and greedy writing, with the output it prints:

import math
from collections import Counter, defaultdict

text = ("the sky is blue . the sea is blue . the sky is gray . the sun is big . "
        "at night the sky is dark . the sky is blue and the sun shines . "
        "the sea is gray . today the sky is blue .")
pieces = text.split()

# 1. Count which piece follows each piece (a bigram model)
table = defaultdict(Counter)
for a, b in zip(pieces, pieces[1:]):
    table[a][b] += 1

def softmax(scores, T=1.0):
    z = [s / T for s in scores]
    m = max(z)
    e = [math.exp(v - m) for v in z]    # subtracting the max does not change the result
    return [v / sum(e) for v in e]

def distribution(context, T=1.0, top_p=1.0):
    words = [w for w, _ in table[context].most_common()]
    scores = [math.log(c) for _, c in table[context].most_common()]   # score = ln(count)
    p = softmax(scores, T)
    nucleus, acc = [], 0.0               # top-p: the most likely words until they reach top_p
    for w, q in zip(words, p):
        nucleus.append((w, q)); acc += q
        if acc >= top_p - 1e-9: break
    total = sum(q for _, q in nucleus)
    return [(w, round(q / total, 3)) for w, q in nucleus]   # share it out again

print(dict(table["is"]))
for T in (1.0, 0.5, 2.0):
    print(T, distribution("is", T))
print(distribution("is", top_p=0.7))

# 2. Write: always take the most likely piece (greedy) and repeat
written = ["the", "sky", "is"]
for _ in range(9):
    written.append(table[written[-1]].most_common(1)[0][0])
print(" ".join(written))
# {'blue': 4, 'gray': 2, 'big': 1, 'dark': 1}
# 1.0 [('blue', 0.5), ('gray', 0.25), ('big', 0.125), ('dark', 0.125)]
# 0.5 [('blue', 0.727), ('gray', 0.182), ('big', 0.045), ('dark', 0.045)]
# 2.0 [('blue', 0.369), ('gray', 0.261), ('big', 0.185), ('dark', 0.185)]
# [('blue', 0.667), ('gray', 0.333)]
# the sky is blue . the sky is blue . the sky

The text was written for the lesson, and its Spanish and English versions are parallel word for word, so the counts match in both languages. The scene uses this same model (lm.ts), and verify-data.ts checks it; both live in the lesson folder.

Scale: GPT-1, GPT-2, and GPT-3

All three are decoder-only transformers trained with the objective above. GPT-1 (2018) read BooksCorpus, over 7,000 unpublished books; its size, 117 million parameters, appears in table 2 of the GPT-2 paper (“the smallest model is equivalent to the original GPT”). GPT-2 (2019) reached 1,542 million and read WebText, slightly over 8 million documents (40 GB of text). GPT-3 (2020) has 175,000 million (175 billion). A curiosity: the GPT-3 paper counts 125 million for a model with the same shape as GPT-1 (12 layers of 768). The gap fits the embedding table, which grows with the vocabulary: GPT-1 used 40,000 BPE merges and GPT-2 expanded the vocabulary to 50,257 pieces.

What comes next

We said “words,” but the model works with pieces: tokens (lesson 20). The probability is computed by a network that looks at the whole context with attention (lessons 21 and 22), and what it means to train it on almost the whole internet is the subject of lesson 23.