How Transformers Learn and Generate Text

Here is the whole story in two loops.

The training loop (repeated many times):

1. Take a batch of real text windows, and their targets: the same windows shifted one character left.
2. Run the model once; it predicts the next character at every position.
3. Measure how wrong it was with cross-entropy loss.
4. Backpropagate and update the weights.

The generation loop (repeated once per new character):

1. Run the model on the text so far.
2. Take its probability distribution for the next character.
3. Choose one character from that distribution.
4. Append it to the text and repeat.

Figure 6.6.1. Teacher forcing compares every position with its true next character in one masked pass. Generation samples from the last position, appends the choice, and repeats with fixed weights. ↵ denotes a newline; the sample text is illustrative.

The first half of this lesson (Sections 2–7) is about training; the second half (Sections 8–10) is about generation.

The model and data are exactly as before: 1,003,854 training characters, 111,540 validation characters, a 65-character vocabulary, and a model with 162,561 parameters.

A language model is trained on one deceptively simple task: given the text so far, predict the next token. No labels are needed beyond the text itself; every character in the training data is the answer to the question "what comes after the characters before it?". This makes language modelling self-supervised, like the skip-gram objective in Lesson 5.3, and it is why language models can learn from enormous amounts of raw text.

Why should such a simple task produce anything impressive? Because predicting the next character well requires a lot of knowledge. To predict that "KING RICH" continues with "A", the model must have learned character names. To predict a line break followed by a speaker's name in capitals, it must have learned the layout of a play. To predict the end of a word it has never seen in exactly that context, it must have learned spelling patterns. Every regularity in the text becomes a way to reduce the loss. Large language models are trained on exactly this task, with subword tokens and vastly more text, and the knowledge they show is a by-product of predicting the next token very well.

As we set up in Lesson 6.2, each training example is a window of text x and a target y that is the same window shifted one position to the left, so that y[t] is the character that follows x[t].

The input begins 't, for his purpose in th' and the target begins one character later, ', for his purpose in the'. The check confirms that the target is exactly the input shifted by one (True). Combined with the causal mask, this shift turns one window into many separate questions: at position t, "having seen x[0] to x[t], what is y[t]?".

At each position, the model outputs a probability distribution over the 65 characters. The loss at that position is −log(the probability the model gave to the correct next character), and the overall loss is the average over all positions. This is the same cross-entropy loss you have used for classification since Module 1; a language model is simply a 65-way classifier applied at every position.

Here is a single position worked by hand, with three possible tokens:

Softmax turns the logits [2.0, 1.0, 0.1] into probabilities [0.659, 0.242, 0.099]. The correct token is the first, so the loss is −log(0.659) = 0.417, which matches F.cross_entropy. The table underneath shows how the loss depends on the probability given to the correct token: 0.105 at a probability of 0.9, 0.693 at 0.5, 2.303 at 0.1 and 4.605 at 0.01. The loss punishes confident mistakes severely and rewards confident correct answers.

A loss of, say, 2.0 means nothing on its own. Before training, we compute the validation loss of three very simple predictors, which give us a scale to read our model's loss against:

• Uniform: every character equally likely.
• Character frequencies (unigram): always predict each character with its overall frequency in the training text, ignoring context entirely.
• Previous character only (bigram): predict the next character from the counts of which character follows which in the training text. This uses exactly one character of context.

On the validation text, a uniform guess scores 4.174, character frequencies 3.347, and the bigram model 2.482. The bigram model has already learned some real structure: after q, it predicts u with probability 0.898. Our Transformer sees up to 64 characters of context, so to be worth anything it must beat 2.482 clearly.

Key terms

• Next-token prediction: the task of predicting each token from the tokens before it.
• Self-supervised learning: learning from targets created from the data itself, here the next character.
• Unigram / bigram model: a predictor using no context / only the previous token.

During training, the model always predicts each position from the real previous characters, never from its own earlier predictions. If it predicted the wrong character at position 5, position 6 still sees the true character at position 5. This is called teacher forcing: like a teacher correcting a student after every word, the model is always given the right history.

Teacher forcing has a huge practical benefit. Because the inputs at every position are known in advance (they are just the real text), the model does not have to wait for its own predictions, so all positions can be computed at once. Together with the causal mask, one forward pass gives a prediction and a loss term at every position in every window.

A batch of shape (32, 64) gives 2,048 next-character predictions from a single forward pass. A recurrent network would have to process the 64 positions one after another; a Transformer computes them all together. This parallelism is one of the main reasons Transformers could be trained on far more data than their predecessors.

Teacher forcing also has a cost, which we will meet in Section 10: at generation time there is no teacher, and the model must live with its own mistakes.

Key terms

• Teacher forcing: during training, feeding the model the true previous tokens rather than its own predictions.
• Parallel training: computing predictions and losses for all positions in one forward pass.

The training function below brings together four tools you have met in Module 3 and Lesson 6.4:

• AdamW is Adam with decoupled weight decay: a small pull of every weight towards zero, applied separately from the gradient update, which acts as regularisation. It is the standard optimizer for Transformers.
• Learning-rate warmup raises the learning rate linearly from almost zero over the first 100 steps, so the first, most chaotic updates are small. In Lesson 6.4 it was the difference between a post-norm model that collapsed and one that trained.
• Cosine decay then lowers the learning rate smoothly from its peak to zero, following half a cosine wave. Large steps early let the model make fast progress; small steps late let it settle into a good solution.
• Gradient clipping rescales the gradients whenever their overall size (norm) exceeds a threshold, here 1.0, so that one unusual batch cannot produce a huge, damaging update. clip_grad_norm_ also returns the norm before clipping, which we record as a health check.

We also prepare 10 fixed validation batches, so that every validation loss is measured on exactly the same 640 windows.

Here is the learning-rate schedule we will use for the main run: a straight climb to the peak over 100 steps, then a smooth cosine descent to zero at step 1,400.

The peak learning rate is the most important setting to get right, so we choose it with a short experiment: train for 200 steps at three learning rates, with a constant learning rate and no warmup, and compare the validation loss. We never use the validation data to compute gradients, and we use the same random batches for every run.

After 200 steps, a learning rate of 0.001 reaches a validation loss of 2.489, 0.003 reaches 2.330 and 0.01 reaches 2.190. The highest learning rate wins, so we choose 0.01.

A caution about short sweeps: they favour higher learning rates, because a larger step size makes faster progress early on, even if it might cause trouble later. That is one reason we pair the chosen rate with warmup and decay, and why we check the full run's curves rather than trusting the sweep blindly.

Common mistake: choosing settings by looking at the test set, or at the training loss. Training loss rewards memorisation; the test set must stay untouched until the very end. Choose on validation.

Now the main run: 1,400 steps with a peak learning rate of 0.01, 100 warmup steps, cosine decay, AdamW and clipping at 1.0. Every 100 steps, we measure the validation loss. We also keep a copy of the untrained model for comparison.

The validation loss falls steadily, from 2.487 after 100 steps to 1.742 at the end. It is still creeping down at the end, but slowly: the last three measurements are 1.752, 1.744 and 1.742, as the cosine schedule brings the learning rate down to zero. The average training loss over the last 100 steps is 1.547. Let's look at the curves, together with the gradient norms.

Reading the loss curves (left). Start with the baselines. The model passes the character-frequency baseline (3.347) almost immediately, and the bigram baseline (2.482) at about step 100, where its validation loss is 2.487. From there on it is using more than one character of context. It ends at 1.742, far below the bigram's 2.482.

Notice the growing gap between training loss (1.547 at the end) and validation loss (1.742). The model does somewhat better on the plays it trained on than on the held-out ones, but both curves are still falling together, with no sign of the validation loss turning upwards. This is a generalisation gap. There is no observed validation upturn in this run, although that alone cannot rule out all overfitting: in Module 3's terms, "train ↓, validation ↓".

Reading the gradient norms (right). The gradient norm starts high in the first few steps, peaking at 0.84 at step 8, then settles around a median of 0.30. It never reaches the clipping threshold of 1.0: clipping was active on 0 of 1,400 steps. In this run, with pre-norm blocks and warmup, clipping turned out to be unnecessary. That is typical of a healthy run: clipping is insurance, which only matters when a bad batch or an unstable setup produces a spike like the 117.6 in Lesson 6.4's collapsed post-norm run. It costs almost nothing to keep.

Let's compare the untrained and trained models' predictions after a few prompts.

The untrained model's best guesses are random, around 0.04 to 0.05 each. The trained model is confident and usually right:

• After "First Citiz" it predicts 'e' with probability 0.966: it knows the speaker name "First Citizen".
• After "KING RICH" it predicts 'A' with 0.986: it knows "KING RICHARD", written in capitals.
• After "my good lor" it predicts 'd' with 0.993.

But not always. After "What light through yonder windo" its top guesses are 'm' (0.396) and 'n' (0.353), and the correct 'w' gets only 0.069. "Window" is rare in this text, and a small model trained for a couple of minutes has not learned it; "windom" and "window" might be equally plausible to it as spellings. Its knowledge is real but patchy.

The model sees up to 64 characters, but how much does it actually benefit from them? We can measure the validation loss separately at each position in the window: at position 0 the model has seen only one character, at position 63 it has seen 64.

At position 0, with just one character to go on, the validation loss is 2.424, close to the bigram model's 2.482. That makes sense: with one character of context, the Transformer can do little more than a bigram model. The loss falls quickly as context grows: 2.196 with 2 characters, 1.952 with 3, 1.794 with 5 and 1.712 with 9. After that, the curve flattens out at roughly 1.7, with some noise from position to position: 1.663 at position 16, 1.770 at 32 and 1.748 at 63.

Loss improves rapidly over the first few positions and then levels off in this diagnostic. This is consistent with useful short context, but it does not establish that the model ignores everything beyond ten characters: different positions predict different targets and use different learned position vectors. A controlled context-length test should predict the same held-out target tokens with progressively truncated histories. Longer-range ability depends on the data, training, architecture and task.

In Lesson 6.3 we built, by hand, a head that attends to the previous token. Did training produce anything similar? Here are the attention weights of the four heads in the first block for the line "KING RICHARD:\nThe king is dead", with numbers summarising, for each head, how much weight goes to each distance back.

Two heads have learned sharp positional rules entirely on their own:

• Head 1 puts an average weight of 0.97 on the previous character: it is the "previous token" head we designed by hand in Lesson 6.3, discovered by training.
• Head 3 puts 0.91 on the character two positions back.

Together, they let the next block see the last three characters directly. Heads 0 and 2 spread their attention more widely over the recent past (no more than 0.24 on any single distance in the first few), gathering broader context. Compare these sharp, purposeful patterns with the noisy weights of the untrained module in Lesson 6.3.

Common mistake: reading too much into a single attention map. Attention shows where information is gathered from, not what is done with it, and later blocks may change everything. Here we backed the pictures with averages over a whole line; for real conclusions, measure over many inputs.

Now the second half: generation. A trained model gives us a probability distribution for the next character. To write text, we repeat four steps:

1. Run the model on the text so far (at most the last 64 characters, the context window).
2. Take the logits at the last position only: the model's scores for the next character.
3. Choose a character from them. Section 9 is about how.
4. Append it to the text, and go back to step 1.

This is called autoregressive generation: each new token is predicted from the model's own previous outputs. Here is the full function, with the choice step written to support all the strategies of Section 9, followed by 300 characters generated after the prompt "ROMEO:" by sampling from the model's distribution as it is.

Compare this with the random characters of the untrained model in Lesson 6.5. The trained model has learned the form of a play: speaker names in capitals followed by a colon and a line break ("JULIET:", "QUEEN MERSET:", "KINGHAM:"), lines of verse that start with capitals, punctuation in sensible places, and many real words ("I would", "to thee", "that's your", "a man is hither"). It has not learned meaning: the sentences do not make sense, and it invents words ("bleepelies", "excheeksiard") and names ("BUCKINGHRMOP"). For a 162,561-parameter model trained for a couple of minutes on one CPU core, reading one character at a time, that is about what we should expect.

Visual walkthrough: Follow the autoregressive loop.

Play or step through three character choices. Each chosen character becomes part of the next input; model weights stay fixed.

These supplied numbers illustrate the calculation. The web app adds a paused, step-by-step animation of the same example.

At step 3 of the generation loop, the model gives us a probability distribution, and we must turn it into a single choice. This choice, called the decoding strategy, has a big effect on the text.

• Greedy decoding always picks the single most likely token. It is deterministic: the same prompt always gives the same text.
• Sampling draws the next token at random according to the model's probabilities, so a token with probability 0.3 is chosen about 30% of the time.

Sampling is usually modified in one of three ways:

• Temperature. Divide the logits by a number T before softmax. T below 1 makes the distribution sharper (more confident); T above 1 makes it flatter (more adventurous); T = 1 leaves it unchanged; and T approaching 0 becomes greedy.
• Top-k. Keep only the k most likely tokens and set the rest to zero probability before sampling.
• Top-p (nucleus sampling). Keep the smallest set of most likely tokens whose probabilities add up to at least p, and drop the rest. Unlike top-k, the number of candidates adapts to how confident the model is.

Let's see temperature's effect on one real distribution: the model's prediction for the character after "...What light through yonder ".

At temperature 1.0 the most likely next character is 't' (as in "the" or "thy") with probability 0.12, with an entropy of 3.07. At temperature 0.5 the same distribution becomes much sharper, with 't' at 0.25 and the entropy down to 2.44. At 1.5 it flattens to 0.08 and 3.39. Temperature never changes the order of the candidates, only how strongly the model prefers the top ones. Here, top-p with p = 0.9 keeps 18 candidates, while top-k with k = 5 would keep exactly 5.

Now let's measure what each strategy does to the generated text. For each strategy we generate 300 characters after "ROMEO:", three times with different random seeds, and measure two things:

• Training-vocabulary words: the share of generated words that appear somewhere in the training text. This measures vocabulary coverage, not factuality, grammaticality or meaning.
• Distinct 8-character pieces: the share of all 8-character stretches of the output that are different from each other, a measure of variety. Text that repeats itself scores low.

The trade-off is very clear:

Greedy decoding uses almost only real words (0.984), but only 0.188 of its 8-character pieces are distinct: it repeats itself. Raising the temperature trades training-vocabulary word coverage for variety in these samples: at 1.5, fewer than half the words are real (0.450). Top-k and top-p improve on plain sampling at temperature 1.0 (0.869 and 0.830 real words, against 0.801) without losing variety, because they cut off the long tail of unlikely characters, where most invented words come from. Here are some of the samples behind those numbers:

The greedy output falls into a loop almost immediately: "the shall the shall the shall..." over and over. Once the model has produced "the shall", the most likely continuation leads straight back to it, and greedy decoding has no way to escape. Temperature 0.5 reads most like English, with plausible phrases such as "serve the contents". Temperature 1.5 dissolves into invented words and stray capitals ("LeWToHat nexhand"). Top-p keeps the variety while producing a more convincing play layout.

There is no single best setting: it depends on the use. Lower temperature can make outputs more repeatable; more sampling can increase variety. Neither greedy decoding nor a low temperature guarantees factual accuracy. Our "real words" score only checks membership in the training vocabulary: repeated familiar words can still form meaningless or false sentences. Choose a strategy using evaluations appropriate to the task, including factuality when it matters.

Common mistake: concluding that a model is "broken" because greedy decoding produces endless repetition. Repetition loops are a known property of greedy decoding, not necessarily of the model. Try sampling with temperature, top-k or top-p before drawing conclusions.

Key terms

• Autoregressive generation: producing a sequence one token at a time, each predicted from the tokens generated so far.
• Decoding strategy: the rule for choosing a token from the model's distribution.
• Greedy decoding: always choosing the most likely token.
• Temperature: dividing logits by T before softmax; T < 1 sharpens, T > 1 flattens.
• Top-k / top-p sampling: sampling only from the k most likely tokens / from the smallest set covering probability p.

Training and generation use the same model, but they differ in important ways:

Two of these differences show up directly in code:

model.train() and model.eval() set the model's training flag (True and False). Our model has no dropout or batch normalisation, so the flag changes nothing here, but in models that have them it changes their behaviour, so the correct mode is a habit worth keeping (Module 2). Inside torch.no_grad(), the loss has requires_grad equal to False: PyTorch does not store anything for backpropagation, which saves memory and time during generation.

The last line points to the main cost of generation: 300 new characters need 300 separate forward passes, one after another. Training processes many positions in parallel, but generation is inherently sequential. Production systems speed this up with a KV cache, which stores the keys and values already computed for earlier tokens so that each new step only has to process the newest token. You will meet it again in Module 7.

The other important difference is exposure bias: during training the model always saw real text, but during generation it sees its own output, mistakes included. One invented word, such as "BUCKINGHRMOP", becomes part of the context for everything that follows, and the model was never trained to recover from its own errors. Sampling strategies that avoid unlikely tokens reduce this problem; they do not remove it.

Key terms

• Inference: using a trained model to make predictions, without updating its weights.
• KV cache: stored keys and values of earlier tokens, reused so each generation step only processes the new token.
• Exposure bias: the mismatch between training on real text and generating from the model's own outputs.

A language model learns by next-token prediction. Each training window x comes with a target y shifted one position left, and with teacher forcing and the causal mask, one forward pass makes a prediction at every position: 2,048 per batch. The loss is cross-entropy, −log(probability of the true next character), averaged over all positions, and we read it against baselines: 4.174 (uniform), 3.347 (character frequencies) and 2.482 (bigram).

We trained our 162,561-parameter model with AdamW, a 100-step warmup, cosine decay and gradient clipping, choosing the peak learning rate of 0.01 on validation (2.190 after 200 steps, against 2.330 and 2.489). After 1,400 steps it reached a validation loss of 1.742, with a healthy gap to the training loss (1.547). Clipping never activated: insurance, not a necessity, in a stable run. Its loss-by-position diagnostic improves most over the first few positions; this does not establish a ten-character effective-context limit. One of its attention heads attends to the previous character with weight 0.97 and another to the character two back with 0.91, the kind of head we designed by hand in Lesson 6.3.

To generate, the model runs autoregressively: predict, choose, append, repeat. The decoding strategy controls the result. Greedy decoding used real words (0.984) but looped ("the shall the shall..."); higher temperatures added variety at the cost of real words (0.450 at temperature 1.5); top-k and top-p cut the unlikely tail and kept both reasonably high. The model learned the form of a play, but not its meaning. Generation is sequential and suffers from exposure bias, because the model must build on its own mistakes.

That completes Module 6: you have built a Transformer from its parts, trained it, and generated text with it. In Module 7 you will use Transformers that others have pretrained at enormous scale, and in Module 8 you will adapt them to your own tasks.

Figure 6.6.2. The training loop predicts real next tokens, measures loss and updates model weights. Generation branches from the trained model and feeds chosen tokens back into its context; it never runs the optimizer.