5 — Learning to predict

Chapter 4 finished the model and left it knowing nothing. Every one of its weights was random, its next-token scores were noise, and asked to continue “Attention is” it wrote ” retrospective programmer programmer programmer”. The machinery to predict the next token was all there. What it lacked was any pressure to predict it well.

Training supplies that pressure. It needs two things: a single number that says how wrong the model’s predictions are, called the loss, and a procedure that nudges every weight in the direction that makes that number smaller. Repeat the nudge a few hundred times over real text and the predictions improve. This chapter builds the loss on a four-word example, trains a small GPT on one book on this Mac, watches it improve and then overfit, and then looks at the other half of using a language model: turning its probabilities into text.

All runs were on an Apple M3 Pro with 18 GB of memory, with Python 3.12, PyTorch 2.15.0 and tiktoken 0.14.0. Training ran on MPS, the Mac’s GPU, through PyTorch’s mps device. Seeds are given with each run.

How wrong is a prediction?

Take a vocabulary of four words, the, cat, sat and down, and the sentence “the cat sat down”. As chapter 2 set up, the inputs are the first three words and the targets are the same sentence shifted by one: after “the” should come “cat”, after “the cat” should come “sat”, after “the cat sat” should come “down”.

Suppose a model produces these logits, one row per position, one column per word:

Softmax, from chapter 3, turns each row into probabilities: exponentiate every score and divide by the row’s total.

$p_i = \frac{e^{z_i}}{\sum_j e^{z_j}}$

The model did well on the first two and was unsure on the third, where it gave “the” more probability than the right answer.

A good loss should be small when the correct word gets high probability and large when it gets low probability. The standard choice is the negative log probability of the correct word:

$\text{loss} = -\ln p_{\text{correct}}$

If the correct word gets probability 1, the loss is −ln 1 = 0. As the probability falls toward 0, the loss grows without limit. For the three positions: −ln 0.7030 = 0.3524, −ln 0.5851 = 0.5359, and −ln 0.3051 = 1.1870. The loss for the whole sequence is the average, 0.6918.

This is cross-entropy, the loss every GPT is trained with. The logarithm isn’t arbitrary. Multiplying the probabilities of each correct word gives the probability the model assigns to the whole sentence; taking the log turns that product into a sum, and the minus sign turns “make the probability high” into “make the number low”. Minimizing cross-entropy is the same as making the training text as likely as possible under the model.

PyTorch computes it in one call, from the raw logits:

import torch
import torch.nn.functional as F

logits = torch.tensor([[0.5, 2.0, 0.1, -1.0],
                       [0.2, -0.5, 1.5, 0.3],
                       [1.0, 0.0, 0.3, 0.8]])
targets = torch.tensor([1, 2, 3])        # cat, sat, down
print(F.cross_entropy(logits, targets))

tensor(0.6918)

F.cross_entropy takes a 2-D tensor of logits, one row per prediction, and a 1-D tensor of the correct indices, and returns the average negative log probability. It applies softmax internally, in a numerically safer way than doing the two steps separately, so the model’s output goes in as it is.

Perplexity

A loss of 0.6918 is hard to picture. Perplexity turns it into something easier: the exponential of the average loss.

$\text{perplexity} = e^{\text{loss}}$

Here that is $e^{0.6918} = 1.9973$, about 2. A perplexity of 2 means the model is, on average, as uncertain as if it were choosing uniformly between 2 words. A model that knew nothing and spread its probability evenly over the 4 words would have a loss of ln 4 = 1.3863 and a perplexity of exactly 4. A perfect model has perplexity 1.

The same scale applies to GPT-2’s vocabulary. An untrained model spreading probability evenly over 50,257 tokens would have a loss of ln 50,257 = 10.82 and a perplexity of 50,257. Any loss below 10.82 means the model has learned something.

For a sequence model the loss is averaged over every position of every sequence in a batch. The model’s output has shape (batch, tokens, vocabulary), and the targets (batch, tokens). F.cross_entropy wants a list of predictions, so both get flattened:

def batch_loss(model, x, y):
    logits = model(x)                                   # (B, T, V)
    return F.cross_entropy(logits.flatten(0, 1), y.flatten())

flatten(0, 1) merges the first two dimensions, turning (B, T, V) into (B·T, V): one row per position across the batch. y.flatten() turns (B, T) into B·T target IDs that line up with those rows.

The book and the model

The training text is H. G. Wells’s The Time Machine (1895), Project Gutenberg eBook #35, the same book chapters 1 and 2 tokenized. It is in the public domain in the United States. With Project Gutenberg’s header and license text stripped off, it is 179,698 characters, 32,454 words and 45,259 GPT-2 tokens: 40,733 for training and 4,526 for validation. The first 90% of the tokens are for training and the last 10% are held out for validation: text the model never trains on, used only to check whether what it learned carries over to new text.

text = load_book("data/pg35.txt")          # strips the Gutenberg header and footer
ids = enc.encode(text)
split = int(0.9 * len(ids))
train_loader = make_loader(ids[:split], context_length=128, batch_size=16, shuffle=True, seed=123)
val_loader = make_loader(ids[split:], context_length=128, batch_size=16, shuffle=False)

make_loader is chapter 2’s sliding-window Dataset wrapped in a DataLoader, with the window moving 128 tokens at a time so windows don’t overlap. That gives 318 training windows of 128 tokens, or 19 batches of 16 per pass through the book (the last partial batch is dropped), and 35 validation windows.

The model is chapter 4’s GPT at a size that trains in a couple of minutes:

cfg = GPTConfig(context_length=128, emb_dim=256, n_layers=4, n_heads=4, dropout=0.1)
torch.manual_seed(123)
model = GPT(cfg).to("mps")

.to("mps") moves every weight onto the Mac’s GPU; batches are moved the same way before each step. This model has 28,923,904 parameters, and the arithmetic says something about small models: the token embedding is 50,257 × 256 = 12,865,792 and the untied output head is the same again, so 89% of the parameters are the two vocabulary-sized ends. The four blocks hold just 3.16 million.

Before training, its loss on the training and validation text is 10.872 and 10.875, right at ln 50,257 = 10.82. It knows nothing, as expected.

The training loop

A training step is four lines:

optimizer = torch.optim.AdamW(model.parameters(), lr=4e-4, weight_decay=0.1)

for epoch in range(30):
    for x, y in train_loader:
        x, y = x.to("mps"), y.to("mps")
        loss = batch_loss(model, x, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

loss.backward() is where PyTorch earns its keep. Every tensor operation in the forward pass was recorded, and backward walks that record in reverse to compute, for every one of the 28.9 million parameters, the gradient: how much the loss would change if that parameter were nudged up slightly. The gradients land in each parameter’s .grad. optimizer.zero_grad() clears the previous step’s gradients first, because PyTorch adds new gradients to old ones by default.

optimizer.step() then moves every parameter a little in the direction that lowers the loss. Plain gradient descent would subtract the gradient times a learning rate. AdamW is smarter: it keeps a running average of each parameter’s gradient and of its square, and scales each parameter’s step by them, so parameters with consistently small gradients still move and noisy ones move cautiously. The learning rate lr=4e-4 sets the overall step size. weight_decay=0.1 pulls every weight slightly toward zero at each step, a mild brake on the model leaning too hard on any one weight.

One pass over all the training batches is an epoch. Every 10 steps the loop also measures the average loss over the whole training set and the whole validation set, with the model in evaluation mode (dropout off) and gradients off. Those are the numbers below.

Watching it learn

The run took 148.6 seconds on MPS for 30 epochs, 570 steps, and 1,167,360 tokens seen (19 batches × 16 windows × 128 tokens × 30 passes).

The greedy continuation of “The Time Traveller” at a few points along the way, 30 new tokens each with newlines written as \n, shows what the numbers mean:

step  50: The Time Traveller\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
step 150: The Time Traveller.\n““““““““““““““
step 250: The Time Traveller, to\n“I’s,” said the Time Traveller.\n“I” said the Time Traveller,
step 400: The Time Traveller. The\n” said the Time Traveller, and his\n\n“I think. “I” said the Time Traveller
step 570: The Time Traveller\n his glass for will be to be to be more. Then, and it was as it. It was as I\nthe thing that the world

At step 50 the model has learned that the newline is the most common token in the book (the Gutenberg text breaks its lines at no more than 82 characters) and predicts nothing else. A model that only knows token frequencies has a loss of 6.35 on this book, the entropy of its token counts, as chapter 4 measured, and the training loss is passing through that level right here. At step 150 it has moved on to the opening quotation mark, and is just as stuck on it. By step 250 it has learned “said the Time Traveller”, a phrase the book uses 11 times, along with commas, periods and dialogue punctuation. By the end it strings together clauses that read almost like sentences.

Overfitting

The two curves tell different stories. Training loss falls the whole way, to 1.572. Validation loss bottoms out at 5.700 at step 250 and then climbs, ending at 6.587, higher than it was at step 100.

This is overfitting. The training set is only 40,733 tokens, and the model has 28.9 million parameters, about 700 per training token. For the first dozen epochs it learns what generalizes: the common words, the shape of a sentence, the book’s vocabulary. After that, the easiest way to keep lowering training loss is to learn the particular sequences it has seen, and those lessons don’t transfer. By the end, its perplexity on the training text is $e^{1.572} \approx 4.8$, as if it were choosing among about 5 tokens at each step. On the held-out chapter it is $e^{6.587} \approx 726$, worse than it was at step 100. It has become confidently wrong in ways specific to the training text.

Counting how often the model’s top prediction is the actual next token makes the gap concrete. Over every position of the training and validation windows:

At the end, the model guesses the next token of a training window correctly seven times in ten, and of an unseen window about one time in eight, worse than it managed at step 250. It has stored a great deal about the training windows, in the exact contexts it saw them.

That memory is brittle, though. Prompted with 32 tokens from the middle of the training text (” me, and I was led to make a further remark, which puzzled me still more: that aged and infirm among this people there were none”), the final model continues ”. I could not no more of my / “I could see the Morlocks. The Morlocks”, while the book goes on ”. “I must confess that my satisfaction with my first theories…”. Only the first token matches. The windows it trained on started at fixed places, every 128 tokens, and this prompt starts somewhere else, so the contexts it memorized never appear exactly. Its greedy text is a pastiche of the book’s phrases, not a copy: the longest stretch of its 60-token continuation of “The Time Traveller” that appears word for word in the book is 25 characters, " the Time Machine, and I ".

The usual defenses are more data, a smaller model, stronger regularization, and stopping early. The cheapest is early stopping: keep a copy of the weights from the step with the lowest validation loss and use those. The loop above saved the step-250 weights, and the rest of this chapter samples from them.

From probabilities to text

The model’s output at the last position is 50,257 scores. Generating text means choosing one token from them, again and again. Chapter 4’s generate always took the highest-scoring token. That is one choice among several, and it changes the text a lot.

Everything in this section uses the step-250 weights, moved to the CPU so that seeded sampling repeats exactly. For “The Time Traveller”, the model’s eight most likely next tokens are:

Greedy

Greedy decoding takes the top token every time. It is deterministic: the same prompt always gives the same text. Its weakness shows up quickly in small models. The most likely continuation of a phrase often leads back to the same phrase, and greedy decoding has no way out of the loop:

The Time Traveller, to
“I’s,” said the Time Traveller.
“I” said the Time Traveller,’s
’s
“

Temperature

Sampling draws the next token at random, in proportion to its probability. That breaks loops, but it also gives every one of the 50,257 tokens some chance, including the ones the model rightly thinks are absurd.

Temperature controls how adventurous the draw is. The logits are divided by a number $T$ before softmax:

$p_i = \frac{e^{z_i / T}}{\sum_j e^{z_j / T}}$

Dividing by a number below 1 stretches the differences between scores, so the top token’s share grows. Dividing by a number above 1 squeezes them, so the distribution flattens. On three logits [2.0, 1.0, 0.1]:

As $T$ approaches 0 the top token’s probability approaches 1 and sampling becomes greedy. At $T = 1$ the model’s own probabilities are used unchanged. On the real model, after “The Time Traveller”:

At 0.5 the two punctuation marks take 89% of the probability between them. At 1.5 the distribution is so flat that the top token gets under 3%, and the entropy, 8.53 nats, is not far from the 10.82 of a uniform guess over the whole vocabulary.

Two samples at each temperature, with seeds 1 and 2:

[T=0.5, seed 1] The Time Traveller, a shape
the length. I had been no vast at a little people, I had been
the Time Machine, I was a very
in the Morlocks, and my mind. I

[T=0.5, seed 2] The Time Traveller, to
’s that it,“I have been
” said the story to the Time Traveller.
I
“It“Iaster“I is

[T=1.0, seed 1] The Time Traveller up down in the time in the Amateur
buildPir I rareomedposeotelising drew the balance from the heat, I was dislressed
in the speech of the judge from utter move in

[T=1.0, seed 2] The Time Traveller people, there down that angry increase. The garments else found
Shell pervasiveaming at have have story to spweet people stood another not broader and of
the Specific delight,asterness of
now

[T=1.5, seed 1] The Time Traveller up yet shape Nemesis hungnell UNDER Amateur granite wouldPir hysterical rare AUTposeotelising drew Detection were soations above Profession Iology dislressed carving altogetheraturation said pointing semblance judge balloon Apparently those Tid

[T=1.5, seed 2] The Time Traveller peopleo distress down between angry increase dug and garments else Clicker peopleShell pervasiveaming corro Starts were storyply spweet people stood acid plunged broader coefficient of shaft realised Specific delightchedasterness touppernow

At 0.5 the text stays inside the book’s vocabulary but wanders. At 1.0 fragments like “buildPir” and “rareomedposeotelising” appear: these are runs of rare subword tokens. The book uses only 5,724 of GPT-2’s 50,257 tokens, and the other 44,000-odd were never the right answer during training, so each got pushed down to a tiny probability. Tiny times 44,000 is not tiny, though, and with sampling at full temperature the model lands in that tail every few tokens. At 1.5 the tail takes over almost completely. Seed 1 at 1.0 and seed 1 at 1.5 share whole stretches (“Amateur”, “drew”, “dislressed”, “judge”): both draws use the same sequence of random numbers, and only the probabilities they are applied to differ.

Top-k

Top-k sampling keeps only the k highest-scoring tokens, sets every other logit to −∞ so its probability becomes zero, and samples among the survivors. It cuts off the long tail of unlikely tokens that makes high-temperature text fall apart, while keeping variety among the plausible ones. It combines naturally with temperature:

@torch.no_grad()
def generate(model, idx, max_new_tokens, temperature=0.0, top_k=None, generator=None):
    model.eval()
    ctx = model.cfg.context_length
    for _ in range(max_new_tokens):
        logits = model(idx[:, -ctx:])[:, -1, :]
        if temperature == 0.0:
            nxt = logits.argmax(dim=-1, keepdim=True)
        else:
            logits = logits / temperature
            if top_k is not None:
                kth = torch.topk(logits, top_k).values[:, -1:]
                logits = logits.masked_fill(logits < kth, float("-inf"))
            probs = torch.softmax(logits, dim=-1)
            nxt = torch.multinomial(probs, 1, generator=generator)
        idx = torch.cat([idx, nxt], dim=1)
    return idx

torch.topk(logits, k) returns the k largest values in descending order; [:, -1:] is the smallest of them, the cut-off. masked_fill replaces every logit below the cut-off with −∞, which softmax turns into exactly zero. torch.multinomial(probs, 1) draws one index at random with the given probabilities, and passing a seeded torch.Generator makes the draw repeatable.

[T=1.0, top-k=10, seed 1] The Time Traveller, a kind of the Morlocks
the little Traveller, I had been very
the Psych were so, and it. It’s up for the Time Machine and I found, in

[T=1.0, top-k=10, seed 2] The Time Traveller? And at a thing. The little people were
of-“I felt in the world—” said the
’s of
the Time Traveller,“I was

[T=1.5, top-k=10, seed 1] The Time Traveller
”�
to you the Psych—’t a little Machine, I had you of his
in the Psych of the Time Traveller. I had no very machine. Then I

[T=1.5, top-k=10, seed 2] The Time Traveller? And at a thing. The little people were-—that was
had made me. It was a kind of the
and, with the Time Machine, in their
the Morller was

None of these settings is “correct”. Greedy is right when there is one right answer and repetition isn’t a risk. Low temperature or small k keeps text safe and plain; higher values make it more varied and more error-prone. Most text generation uses some mix of temperature and a cut-off like top-k, or the related top-p, which keeps the smallest set of tokens whose probabilities add up to p.

Where this leaves things

The loop works. Cross-entropy measured how wrong the model was, backpropagation found how each of 28.9 million weights contributed, and AdamW moved them. In two and a half minutes on this Mac, the model went from a loss of 10.87 to a best validation loss of 5.700, from noise to text with the book’s words and something like its sentences. And then, with nothing new to learn from, it began memorizing.

The limit here isn’t the method. It is the amount of text. One short novel is about 45,000 tokens. GPT-2 was trained on WebText, which the GPT-2 paper describes as slightly over 8 million documents and 40 GB of text, more than 200,000 times the size of this book, with a model three times wider and three times deeper than this one, on hardware far beyond one laptop. The loss this small model reached on unseen text, 5.700, is a perplexity of about 299. A model that has read that much should do far better on the same chapter.

Training GPT-2 again on this Mac is out of reach. But OpenAI published the result: every one of GPT-2’s 124 million trained weights. The model from chapter 4 has exactly the right shape to hold them. Chapter 6 downloads them, loads them in, checks that the result matches a reference implementation, and measures it on the same held-out chapter of The Time Machine.