Next-token prediction

Next-token prediction is the task of giving each token of the vocabulary (see Section 5.1) the probability that it comes next in a text. A language model is trained on this task, and it writes a text by repeating it, one token at a time. A written token never changes, so each later token has to continue from it.

7.1Logits

The logits of a prediction are one number for each token of the vocabulary, and a token with a larger logit gets a larger probability. For example, after the text The ride of 3 km cost $, a model can give the logits 2, 1.5 and 0 to the tokens 9, 8 and 10.

The model calculates the logits as follows. First, each token id becomes its embedding, a vector of $d$ numbers (see Section 6.2). After this, the network turns the embeddings into new vectors of $d$ numbers, one for each position (see Chapter 8). Finally, the logit of a token is the dot product of the last vector with the row of the token in an output matrix (see Section 3.1).

The following picture shows this path for the above text, where the last vector is shaded:

The vector at a position depends only on the tokens up to that position, and its logits predict the token after it. Thus, the shaded vector predicts the token after ␣$, and the picture shows its logits for 9, 8 and 10.

The output matrix has one row for each token of the vocabulary. For example, with 128,256 tokens and $d = 4096$, the output matrix has 128,256 rows of 4096 numbers, the same size as the embedding matrix (see Section 6.2)¹.

7.2Softmax

The softmax² turns the logits into probabilities. Let $z(t)$ denote the logit of the token $t$. The softmax gives $t$ the following probability: $$p(t) = \frac{e^{z(t)}}{\sum_{u} e^{z(u)}}$$ where the sum goes over all tokens $u$ of the vocabulary. Each $e^{z(u)}$ is positive, so each probability is positive, and the probabilities add up to 1. A larger logit gives a larger probability.

For example, suppose that we leave out all tokens except 9, 8 and 10. Their logits 2, 1.5 and 0 give the sum $e^2 + e^{1.5} + e^0 \approx 7.39 + 4.48 + 1 = 12.87$ and the following probabilities:

For example, the probability of 9 is $7.39 / 12.87 \approx 0.57$.

The training uses the same probabilities. At each position of a training text, the loss is $-\ln p(t)$, where $t$ is the token that actually comes next (see Section 1.3). The training lowers this loss by making the logit of $t$ larger compared with the others. The network gives a vector for every position, so one pass through a training text gives the loss at all of its positions (see Chapter 9).

7.3Writing a text

The prompt is the text that a model reads before it writes. The program that runs the model turns the prompt into token ids. After this, it repeats one step: it gets the probabilities of the next token from the model, chooses one token and adds it to the text.

The end token is a special token that marks the end of a text or of a turn in a conversation (see Section 5.1). The program stops when the model writes it³.

In greedy decoding, the program always chooses the token with the highest probability. As an example, consider the prompt The ride of 3 km cost $. The following picture shows the three steps, where the new token is shaded, its probability follows the row, and <|end|> stands for the end token:

After the third step, the text is The ride of 3 km cost $9., and the program stops. The model gives this answer the probability $0.57 \cdot 0.9 \cdot 0.8 \approx 0.41$, the product of the probabilities of its tokens.

Each step goes once through every layer of the network, so an answer of 100 tokens takes 100 passes. Chapter 10 shows what a pass costs, and why it does not have to calculate the vectors of the prompt again.

The program also stops after a maximum number of tokens, which the request sets. A text that stops there has no end token, and it can end in the middle of a sentence⁴.

7.4Sampling

Greedy decoding writes the same text each time for the same prompt. In practice, even greedy decoding can write different texts, because the logits of a request change a little with the number of other requests that the program runs at the same time (see Chapter 10)⁵.

Long texts written by greedy decoding often repeat a phrase again and again⁶. Sampling chooses each token at random, with the probabilities of the model. For example, after the above prompt, sampling writes 9 in about 57 of 100 runs, 8 in about 35 and 10 in about 8. When it writes 8, the next steps continue from $8.

The temperature is a positive number $T$ by which the program divides every logit before the softmax. For example, the logits 2, 1.5 and 0 give the following probabilities at three temperatures:

When $T = 0.5$, the logits become 4, 3 and 0. A temperature below 1 moves probability to the tokens with the largest logits, and a temperature above 1 spreads it more evenly. A lower temperature makes the texts better but less varied⁷.

As $T$ goes toward 0, the probability of 9 goes toward 1. Thus, when one token has the largest logit, sampling becomes greedy decoding. A program that runs models can take the temperature 0 to mean greedy decoding⁸, and the APIs of two large providers use the temperature 1 by default⁹.

Top-p sampling¹⁰ samples only among the most probable tokens. It keeps the fewest of them whose probabilities add up to at least a given value, and it scales their probabilities to add up to 1. For example, with the value 0.9, it keeps 9 and 8, whose probabilities add up to 0.92, and it writes them with the probabilities 0.62 and 0.38. A smaller value removes more unlikely tokens, and it also makes the texts less varied.

Without this limit, tens of thousands of unlikely tokens together get too much probability, and sampling often writes text that makes no sense.

7.5Restricting the format

A program that reads the output of a model often needs a fixed format, for example JSON with given fields. We can describe the format in the prompt, but the model can still write a token that breaks it. Constrained decoding¹¹ removes such tokens. In each step, the program that runs the model sets the probability of every token that would break the format to 0, and it scales the other probabilities to add up to 1.

As an example, suppose that the format allows only {"refund": true} and {"refund": false}, and that the text so far is {"refund":. The model gives the following probabilities to the next token:

The token ␣yes would break the format, so its probability becomes 0, and the other two probabilities are doubled to add up to 1. Thus, a text that ends with the end token is always valid JSON.

Note that the probability of ␣yes does not go to ␣true, which means the same here. Before the restriction, the model gave a refund the probability 0.5 + 0.3 = 0.8, and after it, only 0.6. The restriction changes the answer less when the model already gives little probability to the tokens that break the format.

A straightforward way to find the allowed tokens is to check every token of the vocabulary in each step. With the vocabulary of Section 5.1, this takes 200,000 checks for each new token.

However, when a regular expression can describe the format, the allowed tokens depend only on the place in the format, such as inside a string or after a colon. The idea is to find the allowed tokens of each place once, before the generation starts¹². A step then needs one lookup instead of 200,000 checks.

A format that nests to any depth, such as JSON objects inside objects, also needs a stack of the open brackets. In practice, an API can take longer for the first request with a new format, because it first turns the format into a grammar and prepares it for the steps¹³.

Note that a text that stops at the maximum number of tokens can still be incomplete JSON. The restriction only removes tokens, and it cannot add the missing end.