Language Modeling Basics

A language model does one thing: given the tokens of a text so far, it gives a probability to every possible next token. Everything else in this roadmap, from attention to RLHF (reinforcement learning from human feedback), is a better way of computing that one distribution.

Example. After "the cat sat on the", an imaginary model might give "mat" probability 0.4, "floor" 0.2, "sofa" 0.1, and spread the remaining 0.3 over every other token of the vocabulary. The probabilities of all the tokens add up to 1.

This page builds the two oldest models that still teach the most, a counting model and a small neural one, and then shows how to turn their probabilities into text.

The chain rule of probability

A language model only predicts one token at a time, yet it gives a probability to whole texts. The link is the chain rule of probability: the probability of a sequence is the product of the next-token probabilities.

where:

• $x_1, \dots, x_T$ are the $T$ tokens of the text;
• $p(x_t \mid x_1, \dots, x_{t-1})$, read "the probability of $x_t$ given the tokens before it", is the model's next-token prediction at position $t$;
• $\prod$ means "multiply together", for $t$ from 1 to $T$.

Example. For the text "the cat sat": if $p(\text{the}) = 0.1$, $p(\text{cat} \mid \text{the}) = 0.2$ and $p(\text{sat} \mid \text{the cat}) = 0.5$, then $p(\text{the cat sat}) = 0.1 \cdot 0.2 \cdot 0.5 = 0.01$.

This is not an approximation: it is just the definition of conditional probability, applied $T$ times. So a model that predicts the next token well models whole texts well, and training it means minimizing the cross-entropy of the previous section, averaged over every position.

Counting: the bigram model

The simplest model forgets everything except the last token: to predict $x_t$ it looks only at $x_{t-1}$. This is a bigram model ("bi" for the two tokens involved). It is a table with one row per previous token and one column per next token, and it is built by counting.

1. Count. For every pair of neighbours in the training text, add 1 to the cell (previous token, next token).
2. Normalize. Divide each row by its total, so that every row adds up to 1.

where $\mathrm{count}(i \to j)$ is the number of times token $j$ comes right after token $i$, and the sum below counts every token that comes after $i$. In words: of all the times we saw $i$, the fraction followed by $j$. This is the maximum-likelihood estimate: the table under which the training text is most probable.

Example. The text "a b a b a c", with vocabulary $\{a, b, c\}$. Its pairs are a→b, b→a, a→b, b→a, a→c.

previous	→ a	→ b	→ c	probabilities
a	0	2	1	$(0, \frac{2}{3}, \frac{1}{3})$
b	2	0	0	$(1, 0, 0)$
c	0	0	0	undefined

The middle columns are the counts: how many times "a", "b" or "c" came next. Row "c" is $0/0$, undefined: "c" is never followed by anything.

After "a", the model predicts "b" with probability $2/3$ and "c" with $1/3$.

Smoothing: no zeros allowed

The problem is zeros. In the example, $p(a \mid a) = 0$: the model is sure that "a a" never happens. If a test text contains "a a" even once, the probability of the whole text is 0 (the chain rule multiplies by 0) and its loss is infinite. And the row of "c" is not even defined, because "c" was never followed by anything.

Add-$\alpha$ smoothing (called Laplace smoothing when $\alpha = 1$) pretends every pair was seen $\alpha$ extra times:

where $V$ is the vocabulary size: $\alpha$ is added to each of the $V$ cells of the row, so the total grows by $\alpha V$ and the row still adds up to 1.

Example. Same counts, $\alpha = 1$, $V = 3$.

• Row a: counts $(0, 2, 1)$ become $(1, 3, 2)$, total 6: probabilities $(0.17, 0.5, 0.33)$. "a a" is now unlikely but possible.
• Row b: $(2, 0, 0)$ becomes $(3, 1, 1)$, total 5: $(0.6, 0.2, 0.2)$.
• Row c: $(0, 0, 0)$ becomes $(1, 1, 1)$: uniform, $(1/3, 1/3, 1/3)$. With no data, the model has no preference.

Perplexity

Models are compared by perplexity, $e^{\mathcal{L}}$, where $\mathcal{L}$ is the mean negative log-likelihood per token, measured on text the model was not trained on. It reads as "the number of tokens the model is hesitating between, on average": a model that guesses uniformly among $V$ tokens has perplexity exactly $V$, a perfect model has perplexity 1.

Example. The smoothed table on its own training text "a b a b a c": the five transitions get probabilities $0.5, 0.6, 0.5, 0.6, 0.33$. The mean of $-\log p$ is $0.70$, so the perplexity is $e^{0.70} \approx 2.0$: the model hesitates between about 2 of the 3 tokens.

Why counting does not scale

Longer contexts help: a model that looks at the last 2 tokens (a trigram) knows that after "on the" comes a noun. But the table needs one row per possible context. With a vocabulary of $V = 50{,}000$ tokens, a bigram table has $V^2 = 2.5$ billion cells, and a trigram table $V^3 \approx 1.25 \times 10^{14}$. Almost all of them are never observed in any training text, so the model has no information about most contexts it meets. This is the curse of dimensionality that pushed language modeling towards neural networks.

A neural language model

Bengio et al. (2003) replaced the table with a function that can generalize: if it has seen "the cat sat", it can say something sensible about "the dog sat", because it knows that "cat" and "dog" are similar. The model reads the last $n$ tokens and works in four steps:

1. Embed. Look up each of the $n$ tokens in an embedding table $C$ of shape $(V, D)$: one vector of $D$ numbers per token (section 4).
2. Concatenate. Put the $n$ vectors side by side, oldest first, into one vector $e$ of length $nD$.
3. Hidden layer. A linear layer followed by tanh: $h = \tanh(e W_1 + b_1)$.
4. Logits. Another linear layer gives one score per token of the vocabulary: $z = h W_2 + b_2$. Softmax turns them into the next-token probabilities.

where:

• $e = [C_{x_{t-n}}, \dots, C_{x_{t-1}}]$ is the concatenation of the embeddings of the last $n$ tokens;
• $W_1, b_1$ and $W_2, b_2$ are the weights and biases of the two layers, and $H$ is the size of the hidden layer $h$;
• $z$ are the logits, $V$ numbers.

Example. $n = 3$ tokens of context, embeddings of $D = 64$ numbers, a hidden layer of $H = 256$, and $V = 50{,}000$.

• Shapes: $e$ has $3 \cdot 64 = 192$ numbers, $W_1$ is $192 \times 256$, $W_2$ is $256 \times 50{,}000$.
• Parameters: $3.2$ million in $C$, about $49{,}000$ in $W_1$ and $b_1$, $12.85$ million in $W_2$ and $b_2$: about $16$ million in total.

A table for the same context of 3 tokens would need $V^4 \approx 6 \times 10^{18}$ cells (one row per context of 3 tokens, one column per next token).

The number of parameters grows linearly with $n$ instead of exponentially, and similar tokens get similar embeddings, so the model generalizes to contexts it has never seen. It still has a fixed window of $n$ tokens; attention (section 7) removes that limit.

From logits to text: sampling

A trained model gives logits $z$ for the next token. Turning them into text is a loop: turn the logits into probabilities, choose one token, append it to the text, and ask the model again. How to choose?

• Greedy decoding always picks the most likely token. It is deterministic, and tends to fall into repetitive loops.
• Sampling draws a token at random, each with its probability: a token with probability 0.3 is chosen 30% of the time. Before drawing, the distribution is usually reshaped, in one of the three ways below.

Temperature

Temperature divides the logits by a number $T > 0$ before the softmax:

where $q_i$ is the probability of token $i$ (here $T$ is the temperature, not the length of the text).

• $T < 1$ makes the differences between logits bigger, so the distribution is sharper: the favourite token gets even more probability. As $T \to 0$ it becomes greedy decoding.
• $T > 1$ makes the differences smaller, so the distribution is flatter. As $T \to \infty$ it becomes uniform.
• $T = 1$ leaves the distribution unchanged.

Example. Logits $(2, 1, 0)$:

temperature	probabilities
$T = 0.5$	$(0.867, 0.117, 0.016)$
$T = 1$	$(0.665, 0.245, 0.090)$
$T = 2$	$(0.506, 0.307, 0.186)$

Top-k and top-p

Even a good model gives a little probability to thousands of unlikely tokens (the tail of the distribution), and together they add up: sampling picks one of them now and then, and one bad token can derail the rest of the text. The two filters below cut the tail off before sampling.

Top-k sampling (used by Fan et al., 2018, for story generation) keeps only the $k$ most likely tokens and renormalizes them so they add up to 1. Its weakness: a fixed $k$ is too many when the model is confident, and too few when it is not.

Top-p (nucleus) sampling (Holtzman et al., 2020) keeps the smallest set of most likely tokens whose total probability reaches $p$: a handful of tokens when the distribution is peaked, many when it is flat.

Example. Probabilities $(0.5, 0.3, 0.15, 0.05)$.

• Top-k with $k = 2$ keeps the first two, $0.5$ and $0.3$, and renormalizes them by their total $0.8$: $(0.625, 0.375)$.
• Top-p with $p = 0.9$ adds tokens from the most likely until the total reaches $0.9$: $0.5$, then $0.8$ (not yet), then $0.95$ (reached). It keeps three, renormalized by $0.95$: about $(0.526, 0.316, 0.158)$.

Holtzman et al. observe that maximizing likelihood at decoding time (greedy or beam search) gives bland, strangely repetitive text, and that the long tail of unlikely tokens is unreliable; sampling only from the "nucleus" avoids both problems.

In practice these filters are applied to the logits, by setting the removed tokens to $-\infty$: since $e^{-\infty} = 0$, the same softmax and sampling code then works for all of them.

References

• C. E. Shannon, A Mathematical Theory of Communication, Bell System Technical Journal 27(3), 379–423, 1948. Includes the first n-gram approximations of English text.
• D. Jurafsky, J. H. Martin, Speech and Language Processing, 3rd ed. draft, Chapter 3: N-gram Language Models. Counting, perplexity and smoothing (section 3.6).
• Y. Bengio, R. Ducharme, P. Vincent, C. Jauvin, A Neural Probabilistic Language Model, JMLR 3, 1137–1155, 2003.
• A. Fan, M. Lewis, Y. Dauphin, Hierarchical Neural Story Generation, ACL, 2018. Top-k sampling.
• A. Holtzman, J. Buys, L. Du, M. Forbes, Y. Choi, The Curious Case of Neural Text Degeneration, ICLR, 2020. Nucleus (top-p) sampling.
• A. Karpathy, makemore: bigram and MLP character-level language models, step by step.