How does a language model predict the next word?

A language model predicts the next word given its context. We'll work through a sentence to see how those predictions let us generate text and assign probabilities to complete word sequences. Then we'll estimate the probabilities from word counts and see why the amount of context matters.

What a language model computes

Let's start with a sentence whose last word is missing:

we sat on the bank of the ___

The preceding words give us the context. Several words could complete the sentence, so a language model assigns a probability to each word in its vocabulary. In this example, "river" has the highest probability, followed by alternatives such as "road" and "lake". Together, the probabilities across the vocabulary sum to 1.

wesatonthebankofthe

river

0.41

road

0.09

lake

0.07

canal

0.05

other words

0.38

For a next word $w$ and context $c$, we write the language model's prediction as $P(w \mid c)$. Here, the probability assigned to "river" is:

$${P(\text{river} \mid \text{we sat on the bank of the})} = 0.41$$

Many language models are autoregressive: they generate text by predicting each new word from the words already generated. We select a word using the predicted probabilities, add it to the context, and repeat. After selecting "river", our context becomes "we sat on the bank of the river", and we can predict what follows.

Autoregressive large language models such as GPT-3 use this process with tokens, which can be whole words or parts of words [1]. They learn to predict the next token from its context during training and use those predictions to generate text.

From one word to a sentence

We can also use next-word predictions to calculate the probability of a complete word sequence. We begin with the probability of the first word, then multiply by the probability of each following word given its context. For the first two words of our sentence:

$$P(\text{we}, \text{sat}) = P(\text{we}) \, P(\text{sat} \mid \text{we})$$

Continuing through a sequence of $n$ words gives the chain rule of probability:

$$\begin{aligned} & P(w_1, \dots, w_n) \\ & \quad = P(w_1) \times P(w_2 \mid w_1) \\ & \qquad \times \cdots \times P(w_n \mid w_1, \dots, w_{n-1}) \end{aligned}$$

We can write this more compactly:

$${P(w_1, \dots, w_n)} = {\prod_{i=1}^{n} P(w_i \mid w_1, \dots, w_{i-1})}$$

In our eight-word sentence, each factor predicts one word using all the words before it as context. The first word has no preceding context; by the time we reach "river", the context contains seven words.

$$\begin{aligned} & P(\text{we}) \\[0.4em] {\color{#72746c}\times}\; & P(\text{sat}\mid{\color{#72746c}\text{we}}) \\[0.4em] {\color{#72746c}\times}\; & P(\text{on}\mid{\color{#72746c}\text{we sat}}) \\[0.4em] {\color{#72746c}\times}\; & P(\text{the}\mid{\color{#72746c}\text{we sat on}}) \\[0.2em] & {\color{#72746c}\vdots} \\[0.2em] {\color{#72746c}\times}\; & P({\color{#c9342d}\text{river}}\mid{\color{#72746c}\text{we sat on the bank of the}}) \end{aligned}$$

This gives us a way to compare possible word sequences. In statistical speech recognition, we combine a sequence's language-model probability with how well it matches the recorded audio [2]. For similar-sounding phrases such as "recognize speech" and "wreck a nice beach", the language model favors the more common phrase.

recognize speech

wreck a nice beach

Estimating the probabilities

To calculate a sentence probability, we need estimates for each next-word probability. One approach, used by count-based language models, is to count occurrences in a collection of text. We count how often word $w$ follows context $c$, then divide by the number of times that context appears:

$$\widehat{P}(w \mid c) = \frac{\text{count}(c, w)}{\text{count}(c)}$$

If "of the" appears 1,000,000 times and is followed by "river" 2,000 times, we estimate:

$$\widehat{P}(\text{river} \mid \text{of the}) = \frac{2{,}000}{1{,}000{,}000} = 0.002$$

But our sentence requires the probability of "river" after the full context "we sat on the bank of the". That exact sequence may appear only a few times in the training text, or never at all.

history: 7 words

wesatonthebankofthe

river

seen in training: maybe never

With a vocabulary of 50,000 words, there are $50{,}000^7 \approx 7.8 \times 10^{32}$ possible seven-word contexts. A text collection covers only a small fraction of them. If a context never appears, its count is zero and we can't compute the estimate above.

The Markov assumption

We can make counting more practical by limiting the context to the last few words. Under the Markov assumption, we approximate the next-word probability using this shorter context. Keeping just the last word gives:

$$P(w_i \mid w_1, \dots, w_{i-1}) \approx P(w_i \mid w_{i-1})$$

For our example, we estimate the probability of "river" after "the" rather than after all seven preceding words:

$$\begin{aligned} & P(\text{river} \mid \text{we sat on the bank of the}) \\ & \quad \approx P(\text{river} \mid \text{the}) \end{aligned}$$

Using one preceding word gives us a bigram model, which counts word pairs. Using two gives us a trigram model; for "river", its context would be "of the".

wesatonthebank

trigram

of

thebigram

river

Substituting the bigram approximation into the chain rule gives:

$$P(w_1, \dots, w_n) \approx P(w_1) \prod_{i=2}^{n} P(w_i \mid w_{i-1})$$

We now need counts of individual words and word pairs, so common contexts provide many more examples. When predicting "river", though, a bigram model sees only "the"; the clue in "bank" is lost.

Keeping "bank" in the context lets us use that clue, but leaves us with fewer matching examples to count. For a count-based language model, more context is useful only when we have enough training text to estimate the probabilities reliably.

References

1. GPT-3: Brown, T. B., et al. (2020). Language Models are Few-Shot Learners. NeurIPS 2020. arXiv:2005.14165 ↩
2. Statistical speech recognition: Bahl, L. R., Jelinek, F., & Mercer, R. L. (1983). A Maximum Likelihood Approach to Continuous Speech Recognition. IEEE Transactions on Pattern Analysis and Machine Intelligence 5(2). doi:10.1109/TPAMI.1983.4767370 ↩