Softmax and Probability Distributions over a Vocabulary

A practical question comes up as soon as a language model produces text: how do you put a number on how good that text is? The answer is a loss value, but a loss can only be computed after the model has produced its output. So the first job is to run the model on a small batch and look closely at what it actually generates. Doing this on a small batch makes every step visible before any scoring happens.

The sequence that runs the model, converts its raw outputs into probabilities, picks token IDs, and decodes them back into words is the same sequence that the generate_text_simple function performs internally. This article walks through that sequence on a concrete batch, then shows how the loss is built from the probabilities the model assigned to the correct tokens.

The running batch contains two inputs, 'every effort moves' and 'I really like'. Each has already been converted to token IDs. Each input has three tokens, so the batch holds six token positions in total. That count of six comes back later when the loss is averaged.

Text generation is a five-step procedure. First, the input text is mapped to token IDs using a vocabulary. Second, the model produces a probability score for every token in the vocabulary at every input position. Third, for each position, the index with the highest probability is located. Fourth, those highest-probability indices are collected as the predicted token IDs. Fifth, the predicted token IDs are mapped back into text using an inverse vocabulary that goes from IDs to words.

The same five steps apply whether the vocabulary is tiny or large. The pipeline diagram uses a seven-word vocabulary so that every probability vector fits on one page. The GPT model used in this work has a far larger vocabulary, and that difference matters when you interpret the numbers, as the next section explains.

The output of the model is a probability tensor, and its shape encodes three things at once. The first dimension is the batch size, which is the number of input examples. The second dimension is the number of tokens in each input. The last dimension matches the vocabulary size, because each input position receives one probability for every possible next token. With the two-example batch here and a vocabulary of 50,257 tokens, the probability tensor has shape torch.Size([2, 3, 50257]).

Read the diagram as a lookup table. Each slot in one probability vector is an index, and each index corresponds to a vocabulary token ID. Token IDs for this model range from 0 to 50,256. To choose the predicted token at a position, the third step of the pipeline locates the index whose probability is highest in that vector. The index itself is the token ID, which is why the predicted IDs can be decoded back into words through the inverse vocabulary.

The targets are the token IDs the model should produce. They are built from the inputs by shifting every sequence one position forward, so that each input token is paired with the token that follows it. The input 'every effort moves' has the target 'effort moves you'. The input 'I really like' has the target 'really like chocolate'. Each target token is the next word after the corresponding input token.

This shifting is what teaches a language model to predict the next token in a sequence. Without it, the model would be asked to reproduce its input, which is a different task.

The loss does not look at the single most likely token. It looks at how much probability the model gave to the correct token at each position. For every input position, the probability vector has one entry per vocabulary token. The entry at the index of the target token is the value that matters. The training goal is to raise that entry, so that the correct token becomes more likely relative to the others.

Before any training, the model's vectors are essentially random. Each input token produces a vector, each vector index corresponds to a vocabulary word, and training is meant to maximize the value at the index of the target token. Because the vectors start random, the values at the target indices are not special yet.

Applied to the two-example batch, indexing the probability tensor with the target IDs gives the untrained probabilities assigned to the correct tokens. For the first input these are roughly 7.45e-05, 3.11e-05, and 1.16e-05. For the second input they are roughly 1.03e-05, 5.68e-05, and 4.76e-06. Every value is tiny, which matches the expectation for a large vocabulary.

The loss is computed in a fixed sequence that starts from the six target probabilities collected above. First, a logarithm is applied to each probability, which makes the values easier to optimize than the raw probabilities. Second, the six log probabilities are averaged into a single number. Third, that average is multiplied by -1. The result is the cross entropy loss, which is the quantity training reduces.

In the formula, $p_{i}$ is the probability the model assigned to the correct target token at position $i$, and $n$ is the number of positions, which is six for this batch. The source does not fix a base for the logarithm, so the worked numbers below use the natural logarithm.

Notice how the steps respond to the numbers. Every target probability is far below 1, so every logarithm is strongly negative. Averaging and negating turn those negative values into a positive loss. A larger probability on the correct token pulls its logarithm toward zero, which lowers the loss. That is the direction training pushes.

When generated text is compared with its target, the first useful step is to decode both sides and read them position by position. Decoding the first batch's targets gives ' effort moves you'. Decoding the model's predicted token IDs for the same batch gives ' Armed heNetflix'. The two strings differ throughout, which is consistent with an untrained model.

A divergence report should identify the first position where the predicted token differs from the target, because that is where the sequence departs from what was expected. The loss value for that position is then the place to investigate. In this batch every decoded output differs from its target, so the investigation starts at the first position of each sequence.

A trained model would be expected to match the target at more positions, and the first mismatch would move further along the sequence. Finding that point is the debugging step: it tells you where the model's behaviour first breaks from the target.