Learning to Predict

We now have a whole architecture: tokens go in, embeddings pick up position, a stack of attention and MLP blocks mixes them, and a vector comes out for every position. But a freshly built transformer is a random function. It does not know any language. So how does it learn, and what exactly is it learning?

The honest answer is almost embarrassingly simple. We give the model a chunk of text, hide the next token, and ask it to guess. Then we measure how bad the guess was and nudge every weight a tiny bit in the direction that would have made the guess better. Repeat a few hundred billion times.

Everything in this chapter is that sentence unpacked: what "guess" means (a probability distribution over the vocabulary), what "how bad" means (cross-entropy), what "nudge" means (the gradient), and what to do with the guesses once training is over (sampling).

The objective: predict the next token

Here is the problem. Supervised learning usually needs labels, and labels are expensive. Nobody has hand-labeled the internet. So where do the training targets for a language model come from?

The trick is that text labels itself. Take the sentence "the cat sat on the mat". If the model has read "the cat sat on the", then the correct next token is "mat", and we know that because it is right there in the data. Every position in every document is a free training example, with the answer already attached.

One sequence, L training examples

This is worth being precise about, because it is the reason language models are so data-efficient per document. A sequence of $L$ tokens $x_1, x_2, \dots, x_L$ does not give one training example. It gives $L$ of them, all at once, in a single forward pass. At position $t$ the model sees $x_1 \dots x_t$ and must predict $x_{t+1}$.

Mechanically, we build the inputs and the targets from the same array by shifting it one step. If the data is the token id array [464, 3797, 3332, 319, 262, 2603], the input is everything but the last token and the target is everything but the first.

The causal mask from the attention chapter is what makes this legal. Without it, the hidden state at position $t$ could just look ahead and read $x_{t+1}$, and the loss would collapse to zero without the model learning anything.

Why such a dumb objective works

Predicting the next token sounds like autocomplete, and at small scale that is what it is: the model learns which words are frequent and which follow which. But to keep lowering the loss on real text, it has to keep learning more. To predict the next token of "the capital of France is", it has to know geography. To predict the next line of a Python function, it has to model what the function does. Compression of text, done well enough, requires modeling the world that produced the text. The objective is dumb; the data is not.

From hidden state to probabilities: the LM head

The last transformer block hands us a vector $h_t \in \mathbb{R}^{d}$ for each position. That is a point in a 768-dimensional space, not a guess about a token. How do we turn it into "I think the next token is 'mat' with probability 0.31"?

We need one score per vocabulary entry. So we multiply the hidden state by a matrix $W_U \in \mathbb{R}^{V \times d}$ that has one row per vocabulary item. This is the unembedding, also called the language-model head:

Each entry of $z_t$ is the dot product of $h_t$ with one vocabulary row, so it says "how much does this hidden state look like the direction for token $v$". These raw scores are called logits. They can be any real number, positive or negative, and they do not sum to anything in particular.

Then softmax turns the logits into a proper probability distribution: every entry positive, all entries summing to one, larger logits getting exponentially more mass.

That is the model's guess: a full distribution over what comes next. The forward pass of a language model ends here. Everything after this point is either measuring the guess (training) or acting on it (sampling).

Cross-entropy: loss as surprise

Now we have a distribution and we know the true next token. How bad was the guess? We want a single number that is small when the model put lots of probability on the right answer and large when it did not.

The number we use is the negative log of the probability the model gave to the correct token:

Read it as surprise. If the model said the true token had probability 1, the loss is $-\log 1 = 0$: no surprise. If it said probability 0.5, the loss is $\log 2 \approx 0.69$. If it said 0.01, the loss is $\log 100 \approx 4.6$. Probability near zero means surprise near infinity, which is exactly the punishment we want for confidently ruling out the truth.

Averaged over all positions in a sequence (and all sequences in a batch), this is the training loss:

The name "cross-entropy" comes from information theory. If $q$ is the true distribution of next tokens and $p$ is the model's, the cross-entropy is $-\sum_v q(v)\log p(v)$. Our data gives us only one sample from $q$ at each position (the token that actually appeared), so $q$ is a one-hot vector, and the sum collapses to the single term $-\log p(\text{correct})$. Same formula, seen from two angles.

What the loss looks like at initialization

Here is a useful sanity check that you will use every time you train a model. Before any training, the weights are random and small, so all the logits are close to zero, so softmax gives every token roughly the same probability: $p(v) \approx 1/V$. The loss is then

For GPT-2's vocabulary of $V = 50{,}257$ that is $\ln 50257 \approx 10.82$. For our 5-token toy vocabulary it is $\ln 5 \approx 1.61$. If your model reports a loss of 10.8 at step 0, the plumbing is right. If it reports 25, something (usually the initialization scale) is off, because the model is starting out confidently wrong, which is worse than clueless.

What does "loss = 3.0 nats" mean?

Loss is measured in nats when we use the natural log, or in bits when we use $\log_2$ (one nat is about 1.44 bits). A loss of 3.0 nats means that, on average, the correct token was assigned probability $e^{-3.0} \approx 0.05$. Not that every token got 5%; rather, the geometric mean of the probabilities assigned to the correct tokens was 5%. Some tokens (the "the" after "of") get 90%; some (the first word of a new sentence) get 0.1%; the average of the logs lands at 3.0.

For reference, GPT-2 small reaches roughly 3.3 nats per token on web text, and modern frontier models are reported around 2 nats or below on similar data, though the numbers depend heavily on the tokenizer and dataset, so never compare losses across different tokenizers.

Perplexity: the loss, exponentiated

Cross-entropy in nats is a bit abstract. Perplexity is the same number made intuitive: it is $e$ raised to the loss.

The interpretation: a perplexity of $N$ means the model is, on average, as uncertain as if it were choosing uniformly among $N$ tokens. Loss 3.0 nats is perplexity $e^{3} \approx 20$: the model is "effectively choosing among 20 tokens" at each step. At initialization, perplexity is $e^{\ln V} = V$: choosing uniformly among the whole vocabulary, which is exactly what it is doing.

The companion code/lumen/eval.py computes perplexity properly: it sums the per-token losses over a held-out text, divides by the number of predicted tokens, and exponentiates once at the end (never average the per-sequence perplexities; average the losses).

The gradient signal

We have a loss. Training means computing how the loss changes when each weight changes, and moving the weights the other way. Backpropagation does this through the whole network, but the very first step, the gradient of the loss with respect to the logits, has a form so clean it is worth memorizing.

If $y$ is the one-hot vector for the correct token (a 1 at the correct index, 0 elsewhere), then the gradient of the cross-entropy loss with respect to the logit vector is simply:

That is it. The gradient on each logit is the probability the model assigned to that token, minus 1 if it was the correct token. Gradient descent moves logits in the negative gradient direction, so: the correct token's logit gets pushed up by $1 - p(\text{correct})$, and every wrong token's logit gets pushed down by $p(\text{wrong})$, proportional to how much probability it was wrongly given.

This is the mechanism in one sentence: the model is told, at every position, "more of what happened, less of what you expected instead, in proportion to how wrong you were". When the model is already right ($p(\text{correct}) \to 1$) the gradient goes to zero and it stops learning from that token. The loss self-regulates: easy tokens stop contributing and hard tokens dominate.

From the logits, backprop continues into $W_U$, then the last block, and so on down to the embedding table. Every parameter gets a gradient. We never write those by hand: PyTorch's autograd does it when we call loss.backward(). But it helps to know that everything upstream is being told, indirectly, the same thing: make the correct token more likely.

The training loop end to end

Let us zoom out and see one complete step of training. The question is: what happens between "here is a batch of text" and "the weights are slightly better"?

In PyTorch the whole thing is a few lines. The companion code/lumen/train.py adds a learning-rate schedule, gradient clipping and mixed precision on top of this skeleton, but the skeleton is the important part:

import torch, torch.nn.functional as F

def train_step(model, optimizer, x, y):
    # x, y: (B, L) integer token ids; y is x shifted one to the right
    logits = model(x)                            # (B, L, V)
    loss = F.cross_entropy(                      # softmax + -log p(correct), averaged
        logits.view(-1, logits.size(-1)),        # (B*L, V)
        y.view(-1),                              # (B*L,)
    )
    optimizer.zero_grad(set_to_none=True)
    loss.backward()                              # gradients for every parameter
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()
    return loss.item()

Two details that matter. First, F.cross_entropy takes raw logits, not probabilities; it computes the softmax and the log together in a numerically stable way, so never apply softmax yourself before it. Second, the loss is averaged over all $B \times L$ positions, so a batch of 8 sequences of 1024 tokens is 8,192 next-token predictions per step.

Reading a loss curve

The one plot every practitioner stares at is loss versus step. What does a healthy one look like, and how do you tell healthy from sick?

The fast drop

In the first few hundred steps the loss falls off a cliff, from $\ln V$ down to maybe 6 or 7 nats. This is the model learning unigram and bigram statistics: which tokens are common at all, and which tokens tend to follow which. That is cheap knowledge with a huge payoff, so gradient descent finds it first.

The long tail

After that the curve bends and flattens into a slow, grinding decline that never quite stops. Each additional nat of loss now requires something more expensive: syntax, then facts, then reasoning. On a log-log plot of loss against steps this tail is roughly a straight line, which is the empirical basis of scaling laws (Kaplan et al., 2020; you will meet them properly in Scaling Laws and Optimization).

Spikes and divergence

Sometimes the loss jumps up by a nat or two in a single step and then recovers over the next few hundred. That is a loss spike: a bad batch, a numerically unlucky attention logit, or a learning rate that is a touch too high. Occasional spikes that recover are tolerable. If the loss jumps and keeps climbing, or turns into a not-a-number, the run has diverged, and you lower the learning rate, tighten gradient clipping, or restore from a checkpoint before the spike.

Sampling: turning probabilities into text

Training is over. The model can now produce, for any prefix, a distribution over the next token. But a distribution is not text. To generate, we have to pick a token, append it to the prefix, and run the model again. The question is how to pick, and it turns out the obvious answers are both bad.

Greedy decoding: always take the best

The obvious choice is to take the highest-probability token every time, $\hat{x} = \arg\max_v p(v)$. This is deterministic and cheap, and for short factual answers it works fine. For anything longer it is boring and, worse, it gets stuck: the model wanders into a phrase, the most likely continuation of that phrase is more of the same phrase, and you get "I don't know. I don't know. I don't know." Holtzman et al. (2019) showed that greedy and beam-search text from GPT-2 is repetitive and has far higher per-token probability than real human text. Humans do not always say the most likely thing.

Pure sampling: roll the dice

The opposite choice is to sample from $p$ exactly. This fixes repetition and is what the model "really believes", but it is chaotic. The tail of the distribution is enormous: 50,000 tokens each at probability 0.0001 add up to a lot of probability, so every few tokens the sample lands on something weird, and once one weird token is in the context, the next distribution is conditioned on nonsense and things unravel. Everything below is a way to trim that tail while keeping some randomness.

Temperature: sharpen or flatten

Divide the logits by a temperature $T$ before the softmax:

With $T = 1$ nothing changes. With $T < 1$ the differences between logits get magnified, so the distribution sharpens toward the top choices; as $T \to 0$ it becomes greedy. With $T > 1$ it flattens toward uniform. A neat way to see it: $p_T(v) \propto p(v)^{1/T}$, so $T = 0.5$ squares the probabilities and renormalizes.

Top-k: keep only the k best

Zero out every token except the $k$ highest-probability ones, then renormalize the survivors (Fan et al., 2018). With $k = 50$ the model can still be creative among 50 plausible tokens but can never emit one of the 50,000 junk tokens. The weakness: $k$ is fixed, but the number of sensible continuations is not. After "the Eiffel Tower is in" there is one good answer; after "she opened the door and saw" there are hundreds. A fixed $k$ is too large for the first and too small for the second.

Top-p (nucleus): keep the smallest set that covers p

Sort tokens by probability and keep the smallest prefix whose cumulative probability reaches $p$, typically $0.9$ or $0.95$ (Holtzman et al., 2019). Then renormalize. When the model is confident, the nucleus is one or two tokens; when it is uncertain, the nucleus grows to hundreds. The cutoff adapts to the shape of the distribution, which is exactly what top-k could not do.

Repetition penalties

Even with good truncation, models like to repeat themselves. A repetition penalty (Keskar et al., 2019) divides the logit of any token that already appears in the context by a factor $\rho > 1$ (or multiplies it if the logit is negative), making reappearance less likely. A frequency penalty subtracts a constant times the count of previous occurrences; a presence penalty subtracts a constant if the token appeared at all. All of these are heuristics with no principled justification; they exist because they help.

The companion code/lumen/sampling.py implements all of these on a batch of logits, in the same order as above: temperature, then top-k, then top-p, then a torch.multinomial draw. The core is about ten lines:

def sample_next(logits, temperature=1.0, top_k=None, top_p=None):
    logits = logits / max(temperature, 1e-6)
    if top_k is not None:
        kth = torch.topk(logits, k=min(top_k, logits.size(-1))).values[..., -1, None]
        logits = logits.masked_fill(logits < kth, float("-inf"))
    if top_p is not None:
        sorted_logits, idx = torch.sort(logits, descending=True)
        cum = torch.softmax(sorted_logits, dim=-1).cumsum(-1)
        remove = cum - torch.softmax(sorted_logits, dim=-1) >= top_p   # keep the token that crosses p
        sorted_logits = sorted_logits.masked_fill(remove, float("-inf"))
        logits = torch.full_like(logits, float("-inf")).scatter(-1, idx, sorted_logits)
    probs = torch.softmax(logits, dim=-1)
    return torch.multinomial(probs, num_samples=1)

Beam search, and why it is rarely used for open-ended text

Beam search keeps the $b$ most probable partial sequences at each step instead of committing to one token, and at the end returns the whole sequence with the highest total log-probability. It was the workhorse of machine translation, where there is a short, nearly-correct answer to find. For open-ended generation it fails for the reason greedy fails, amplified: it finds sequences that are too probable. Highly probable text is generic, repetitive, and short, and beam search is very good at finding it. Stahlberg and Byrne (2019) even showed that for many translation models the single most probable output is the empty string. So in chat and creative settings, everyone samples.

Teacher forcing and exposure bias

One last subtlety about how training and generation differ. During training, at every position the model is conditioned on the true previous tokens from the data, regardless of what it would have predicted. This is called teacher forcing, and it is what makes the parallel, shift-by-one training possible. At inference, the model is conditioned on its own previous outputs. So a model that makes one mistake finds itself in a context it never saw during training, and can compound the error; this mismatch is called exposure bias (Ranzato et al., 2015). Remedies such as scheduled sampling (Bengio et al., 2015) exist, but in practice large models are trained with plain teacher forcing and the problem is mostly handled by the sheer breadth of contexts they have seen, plus the post-training methods of the next chapter, which do train on model-generated text.

Practice

Further reading

• Bengio, Ducharme, Vincent, Jauvin (2003). A Neural Probabilistic Language Model. The origin of neural next-word prediction with a softmax over the vocabulary.
• Radford et al. (2019). Language Models are Unsupervised Multitask Learners. The GPT-2 paper: the argument that next-token prediction on web text learns many tasks at once.
• Holtzman, Buys, Du, Forbes, Choi (2019). The Curious Case of Neural Text Degeneration. Introduces nucleus sampling and shows why greedy and beam search produce repetitive text.
• Fan, Lewis, Dauphin (2018). Hierarchical Neural Story Generation. Where top-k sampling was introduced.
• Keskar et al. (2019). CTRL: A Conditional Transformer Language Model for Controllable Generation. Introduces the repetition penalty.
• Stahlberg and Byrne (2019). On NMT Search Errors and Model Errors: Cat Got Your Tongue? Exact search often returns the empty string, a striking demonstration of the "too probable" problem.
• Ranzato, Chopra, Auli, Zaremba (2015). Sequence Level Training with Recurrent Neural Networks. Names and analyzes exposure bias.
• Bengio, Vinyals, Jaitly, Shazeer (2015). Scheduled Sampling for Sequence Prediction with Recurrent Neural Networks. One classic remedy for exposure bias.
• Kaplan et al. (2020). Scaling Laws for Neural Language Models. The power-law shape of the long tail, quantified.