LM5-T1 — The Next-Token Objective and Cross-Entropy

Learning objectives

After this lesson, you should be able to:

• Derive the cross-entropy loss from the likelihood of the training text.
• Compute the loss and the perplexity of a model from the probabilities it gave to the correct tokens.
• State the loss of a uniform guess and use it as a sanity check.
• Write the loss in PyTorch with the correct tensor shapes, and explain why it is computed from logits.

The idea: reward the model for not being surprised

Take a piece of training text and a model. At every position the model states a probability for each possible next token. Look at the probability it gave to the token that really came next. If it was high, the model was not surprised and did well; if it was near zero, the model was badly wrong. The training goal is to make the model, on average, as unsurprised as possible by real text.

From likelihood to loss

The probability a model gives to a whole text is the product of the probabilities of its tokens, each conditioned on the earlier ones (the chain rule, Lesson LM1-T2). Training by maximum likelihood chooses the parameters that make the probability of the training text largest. A product of thousands of numbers below 1 is vanishingly small, so we maximise its logarithm instead; the logarithm turns the product into a sum and does not change where the maximum is. We then flip the sign, so that bigger is worse and the quantity can be minimised, and divide by the number of tokens, so that texts of different length can be compared. The result is the cross-entropy loss.

Why the logarithm is the right penalty

The curve above shows −ln p. A prediction that gave the right token probability 0.5 costs 0.69; 0.1 costs 2.30; 0.01 costs 4.61. The penalty grows without bound as the probability approaches zero, so a confident wrong answer is punished much more than an unsure one. That is the behaviour we want: a model should never say “impossible” about something that then happens.

Perplexity

A loss of 1.15 nats is hard to feel. Perplexity is its exponential, e^1.15 ≈ 3.2, and has a direct reading: it is the number of equally likely tokens the model is effectively choosing between at each step. A model with perplexity 3 is as uncertain as a fair three-sided die. A model that knows nothing about a vocabulary of 50,257 tokens has perplexity 50,257.

Units: nats and bits

With the natural logarithm the loss is in nats; with base 2 it is in bits. One nat is 1/ln 2 ≈ 1.44 bits. Papers on compression and text report bits per character or per byte; PyTorch reports nats. They differ by a constant factor, and a perplexity is the same whichever base is used if the exponential matches.

The mathematics: likelihood, cross-entropy, perplexity

Let the text be tokens t₁, …, t_T and let p_θ(t_i | t_<i) be the probability the model with parameters θ assigns to the true token given the tokens before it. The log-likelihood is

and the loss we minimise is its negative average:

T

the number of positions at which a prediction is scored (all positions of all sequences in the batch).

p_θ(t_i | t_

the probability of the correct next token at position i; the softmax of the logits at that position, picked at the index of the target.

ln

the natural logarithm; the loss is then in nats per token.

L

a single number; lower is better; 0 only if every correct token got probability 1.

With one-hot targets this is the cross-entropy between the true distribution (all weight on the actual next token) and the model’s distribution. In terms of logits z ∈ ℝ^V at a position with target y:

The second term is the log-sum-exp; subtracting the maximum logit before exponentiating keeps it from overflowing (Lab LM5-L1, Stage 3). Frameworks compute exactly this inside cross_entropy, which is why you pass logits and never probabilities you computed yourself.

Perplexity and baselines.

A model that gives every token the same probability 1/V has loss ln V and perplexity V. For V = 50,257, ln V = 10.825 nats = 15.62 bits. Every trained model must beat this; an untrained model should start close to it. In bits per token, L_bits = L / ln 2.

Shapes. The model returns logits of shape (B, n, V), the targets have shape (B, n). F.cross_entropy expects (N, V) and (N), so flatten the batch and sequence dimensions: N = B · n. The result is the mean over all N positions.

A worked example: four positions

A model gives the correct next token the probabilities 0.5, 0.1, 0.8 and 0.25 at four positions.

Mean loss: (0.693 + 2.303 + 0.223 + 1.386) / 4 = 1.151 nats. Perplexity: e^1.151 = 3.16. In bits: 1.151 / 0.693 = 1.66 bits per token. Compare a model that gives every one of 50,257 tokens the same probability: 10.82 nats, perplexity 50,257.

The same four numbers are also the geometric-mean story: the geometric mean of the probabilities (0.5 · 0.1 · 0.8 · 0.25)^1/4 = 0.316 is 1/3.16, the inverse of the perplexity.

import math
import torch
import torch.nn.functional as F

p = [0.5, 0.1, 0.8, 0.25]
loss = sum(-math.log(v) for v in p) / len(p)
print(round(loss, 3), round(math.exp(loss), 3), round(loss / math.log(2), 2))   # 1.151 3.162 1.66

# the same through logits: logits whose softmax gives exactly these target probabilities at 4 positions with V = 3
V = 3
probs = torch.tensor([[0.5, 0.3, 0.2], [0.1, 0.6, 0.3], [0.8, 0.1, 0.1], [0.25, 0.25, 0.5]])
targets = torch.tensor([0, 0, 0, 0])
logits = probs.log()                                   # any logits with these softmax values; here ln p
print(round(F.cross_entropy(logits, targets).item(), 3))   # 1.151

# the batch layout: logits (B, n, V) and targets (B, n) must be flattened
B, n, V = 2, 5, 11
logits = torch.randn(B, n, V)
targets = torch.randint(0, V, (B, n))
l = F.cross_entropy(logits.flatten(0, 1), targets.flatten())
print(tuple(logits.flatten(0, 1).shape), tuple(targets.flatten().shape), l.item() > 0)   # (10, 11) (10,) True

# the loss of an untrained (near uniform) model is about ln V
uniform = torch.zeros(B * n, V)
print(round(F.cross_entropy(uniform, targets.flatten()).item(), 4), round(math.log(V), 4))   # 2.3979 2.3979

# padding positions can be skipped with ignore_index
padded = targets.flatten().clone(); padded[-3:] = -100
print(F.cross_entropy(logits.flatten(0, 1), padded, ignore_index=-100).item() > 0)   # True

Note the shapes in the middle block. A common error is to pass the logits as (B, n, V) with the targets (B, n): PyTorch then treats the second dimension as the classes and either fails or silently computes something wrong when the sizes happen to fit.

Engineering implications and limitations

Give the loss logits, not probabilities. Computing a softmax yourself and then a logarithm is less stable and slower than cross_entropy, which fuses the two.

Flatten correctly. logits.flatten(0, 1) and targets.flatten() keep positions aligned. Flattening only one of them is a typical bug and often fails only for some shapes.

The first loss is a test. It should be close to ln V. If it is much lower, targets may leak into inputs; if it is much higher, the initial weights or logit scale are off. Check it before every long run.

Compare losses only on the same tokenizer. A loss per token depends on how text is cut into tokens. A model with a larger vocabulary has fewer tokens per word and a higher loss per token for the same quality. Compare models by bits per character or per byte if tokenizers differ.

Validation, not training, loss. The training loss measures fit to seen text. Report validation loss and perplexity on held-out text (Lesson LM2-T6), and say which text.

Mean versus sum. Averaging over tokens makes the loss independent of the batch size and sequence length; that keeps the learning rate meaningful across settings. Summing would scale the gradient with the batch.

Perplexity is not quality. It measures how well the model predicts the reference text. Two models with similar perplexity can differ greatly in usefulness for a task; evaluation (Module 9) goes beyond it.

Common misconceptions

Engineering takeaways

1. Training maximises the likelihood of the text; equivalently it minimises L = −(1/T) Σ ln p(correct token).
2. Perplexity = exp(L), the effective number of choices per token; a uniform guess has loss ln V and perplexity V (10.82 nats and 50,257 for GPT-2’s vocabulary).
3. Nats and bits differ by the factor ln 2: 1 nat = 1.44 bits.
4. In PyTorch: F.cross_entropy(logits.flatten(0, 1), targets.flatten()), with logits of shape (B, n, V) and targets (B, n).
5. Check that the first loss is near ln V, and compare models by held-out loss on the same tokenizer.

References and further reading

See Chapter 5 of the book (calculating the text generation loss: cross-entropy and perplexity) and the PyTorch documentation for torch.nn.functional.cross_entropy.

• Shannon, C. E. (1948). A Mathematical Theory of Communication. Bell System Technical Journal (the source of entropy and cross-entropy in information theory).
• Build a Large Language Model (From Scratch), Chapter 5: Pretraining on Unlabeled Data (primary book reference)
• Build a Large Language Model (From Scratch), Appendix D: Adding Bells and Whistles to the Training Loop (primary book reference)

Knowledge check

Three conceptual questions. Write an answer, then reveal the explanation. Your answers stay in your browser.

What’s next

Next: how the loss is turned into parameter updates, with backpropagation and the AdamW optimizer.

Related content

Nothing below is required to finish this lesson.

• Before
  LM4-T5: Output Logits, Parameter Counts and Shapes

  Project hidden states to vocabulary logits, count parameters and follow tensor shapes end to end.

• Lab
  LM5-L1: Cross-Entropy and Perplexity

  Compute the loss and perplexity of a model on a text and see what a better model changes.