Calculating Cross-Entropy Loss and Perplexity

Hello! In the previous lesson, you produced a vocabulary-sized vector of logits at each position. Those scores are useful only if we can compare them with what actually came next. This lesson makes that comparison: you will calculate cross-entropy loss for individual next-token predictions, average it across a short sequence, and convert that average into perplexity.

Allow about 40 minutes for the short video, two focused readings, and the PyTorch example. We will use a five-token vocabulary from the previous lesson, so you can concentrate on the measurement rather than a new model architecture.

Which prediction gets scored?

The image below shows the basic idea: a language model assigns a probability to the next item given what precedes it. It uses whole words for illustration; a Llama model makes predictions over tokens, which may be words, pieces of words, or special tokens.

For our tiny sequence, suppose the token IDs represent I cast heal <eos>. We score three predictions:

The target is shifted one position ahead of the input. The score produced at cast, for example, is compared with heal, not with cast itself. During training, many such comparisons can be made from one sequence.

Watch the two selected passages from Andrej Karpathy’s Let’s build GPT. His tokens are characters, but the target alignment and loss calculation work the same way for other tokenizers.

From logits to cross-entropy

A logit is not a probability. For one position, let ${\ell }_j$ be the logit for vocabulary token $j$. Softmax turns all $V$ logits into a probability distribution:

To score that position, we look up the probability assigned to the actual next token, $y$. Its cross-entropy loss, when the target is a single token ID, is

Here $\log$ is the natural logarithm. A confident correct prediction has a small loss: if $p_y=1$, its loss is $0$. Assigning a small probability to the actual token produces a large loss, even if some other token received the highest probability. This is also called negative log-likelihood for a single correct target.

Suppose our three predictions assign the following probabilities to their actual targets. These are illustrative, hand-set probabilities—not outputs from a trained model.

The middle prediction contributes the most loss because it assigned the least probability to what happened. To summarize the sequence, average over the three scored positions:

The product $0.10$ is the probability assigned to these three successive targets under the stated predictions. Using the average negative log probability rather than the total makes sequences with different numbers of scored tokens more comparable.

What perplexity adds

Perplexity puts that mean loss on another scale:

Read Section 9.3.2, “Perplexity,” in Dive into Deep Learning for the connection between average cross-entropy and this metric.

A perplexity of $1$ is the ideal limit: every actual next token receives probability $1$. If predictions are uniform over our five-token vocabulary, every target receives probability $1/5$; mean loss is $\log 5$, and perplexity is $5$. Our illustrative $2.1544$ is lower than that uniform baseline.

Perplexity is sometimes described as an effective number of choices the model faces, but it is not literally the number of plausible words at each position. Lower perplexity means higher assigned probability to the observed tokens on the evaluation text being scored. Compare perplexities only with care: changing the tokenizer changes what counts as a token, and changing the evaluation text changes the predictions being tested.

Calculate both metrics in PyTorch

PyTorch’s cross-entropy function accepts raw logits and integer target IDs. It performs the log-softmax calculation internally; do not apply softmax to the logits before passing them in. The documentation also specifies a shape convention: for a two-dimensional call, logits have shape (number_of_predictions, vocabulary_size) and targets have shape (number_of_predictions,).

This standalone CPU script reproduces our hand calculation. We specify full probability distributions first, then take their logarithms to obtain convenient illustrative logits: softmax of those logits recovers the specified probabilities. A real model would output logits directly.

import torch
import torch.nn.functional as F

vocab = ["<eos>", "I", "cast", "heal", "fireball"]

# One sequence with three prediction positions; last axis follows vocab.
probs = torch.tensor([[
    [0.05, 0.05, 0.50, 0.20, 0.20],  # after "I"
    [0.05, 0.05, 0.25, 0.25, 0.40],  # after "I cast"
    [0.80, 0.05, 0.05, 0.05, 0.05],  # after "I cast heal"
]], dtype=torch.float64)

sequence = torch.tensor([[1, 2, 3, 0]], dtype=torch.long)
inputs = sequence[:, :-1]   # I, cast, heal
targets = sequence[:, 1:]   # cast, heal, <eos>

logits = probs.log()         # Hand-constructed logits, shape (1, 3, 5)
B, T, V = logits.shape

# Treat each batch/position pair as one next-token prediction.
per_token = F.cross_entropy(
    logits.reshape(B * T, V),
    targets.reshape(B * T),
    reduction="none",
)
mean_loss = per_token.mean()
perplexity = mean_loss.exp()

print("inputs:", inputs.tolist())
print("targets:", targets.tolist())
print("losses:", per_token.tolist())
print(f"mean loss:  {mean_loss.item():.6f}")
print(f"perplexity: {perplexity.item():.6f}")

# Independently verify which probabilities the targets selected.
target_probs = probs.reshape(B * T, V)[
    torch.arange(B * T), targets.reshape(B * T)
]
assert torch.allclose(target_probs, torch.tensor([0.50, 0.25, 0.80]))
assert torch.allclose(per_token, -target_probs.log())

The final two printed values should be approximately 0.767528 and 2.154435. Notice the second probability row: fireball has the highest probability, but heal is the actual target. Cross-entropy scores the probability of heal, not the probability of the model’s preferred answer.

The reshape preserves the pairing between each position’s five logits and its target ID. For larger batches, a model commonly produces logits of shape (B, T, V) and targets of shape (B, T); flattening both in the same order gives PyTorch (B*T, V) and (B*T,). If a batch contains padding or positions you deliberately do not score, exclude those targets from the mean—for example, with ignore_index=-100. Perplexity must be the exponential of the mean over scored tokens, not over padded positions.

Takeaways

For each next-token prediction, cross-entropy is the negative natural log of the probability assigned to the actual next token. Average those losses over the scored positions, then exponentiate the average to obtain perplexity. In PyTorch, supply logits and integer targets directly to cross-entropy, with the vocabulary on the class axis; do not softmax first.

Next, you will use this scalar loss to train a small model. The key change is that instead of merely measuring how surprising its predictions are, you will adjust its weights to make the observed targets less surprising.