This note gives next-token prediction — introduced conceptually in GPT Architecture and used throughout LLM Pretraining — its full mathematical and numerical treatment.
Formula
This is exactly categorical cross-entropy (see Categorical Cross-Entropy), averaged over every position \(t\) in a sequence of length \(T\) — at each position, the "true label" is simply whatever token actually came next in the training text.
The Parallel-Supervision Trick, Enabled by Causal Masking
Here's the elegant part: because of causal masking (see Masked Self-Attention), position \(t\)'s prediction only ever depends on positions \(1,\ldots,t-1\) — so a single forward pass through the whole sequence simultaneously produces a valid next-token prediction at every position, all supervised at once against the actual next tokens in that same sequence. There's no need to run \(T\) separate forward passes for a sequence of length \(T\) — this parallel-supervision property is a major reason next-token pretraining is so efficient at scale.
Numerical Example
For the sequence "the cat sat," with true next tokens ["cat", "sat", "<end>"], and the model assigning probabilities \(P(\text{"cat"}\mid\text{"the"})=0.4\), \(P(\text{"sat"}\mid\text{"the cat"})=0.3\), \(P(\text{"<end>"}\mid\text{"the cat sat"})=0.6\):
Code
import torch
import torch.nn.functional as F
# logits: (seq_len, vocab_size), one prediction distribution per position
logits = torch.tensor([[2.0, 0.5, -1.0], # predicting the token after "the"
[0.3, 1.8, 0.1], # predicting the token after "the cat"
[-0.5, 0.2, 2.1]]) # predicting the token after "the cat sat"
true_next_tokens = torch.tensor([0, 1, 2]) # the ACTUAL next token id at each position
loss = F.cross_entropy(logits, true_next_tokens)
print(loss.item()) # a single scalar, averaged across all 3 positions -- computed in ONE callCommon Mistakes
- Assuming next-token prediction requires a separate forward pass per position during training — causal masking is exactly what allows one single forward pass to supervise every position in a training sequence simultaneously.
- Confusing this training-time parallelism with inference-time generation — generation genuinely must happen one token at a time (each new token depends on the actual previous tokens generated), unlike training, which can process a full known sequence in parallel.
Interview Relevance
Q: "Why can an LLM be trained on a full sequence in one forward pass, even though it generates text one token at a time at inference?" During training, the full target sequence is already known, and causal masking ensures each position's prediction only depends on genuinely earlier positions — so a single forward pass can compute valid next-token predictions, and their losses, at every position simultaneously. At inference, future tokens don't exist yet (they haven't been generated), so generation must proceed one step at a time, feeding each new token back in autoregressively.
Practice Question
For a training sequence of length 500, how many next-token predictions does one forward pass produce and supervise simultaneously?
