Training and Generation in Practice

You now have the whole architecture: tokens → embeddings → stacked blocks of (causal self-attention + MLP) → a final linear layer. This lesson connects that machine to the two things it actually does: learn from text and produce text. They're the same forward pass used in two different loops.

1. The task: next-token prediction

The model's entire job is one function:

$\text{tokens}_{1..t} \;\longrightarrow\; \text{a probability distribution over the vocabulary for token } t+1$

The final layer (often called the unembedding or LM head) is a matrix $W_U \in \mathbb{R}^{d_{\text{model}} \times V}$ , where $V$ is the vocabulary size (say 50,000). It maps the residual-stream vector at position $t$ into $V$ raw scores called logits:

$z = h_t W_U \in \mathbb{R}^{V}$

Softmax turns logits into probabilities:

$p_i = \frac{e^{z_i}}{\sum_{j=1}^{V} e^{z_j}}$

That's it. There is no separate "understanding" module. Everything the model knows about grammar, facts, and style exists because it improves this one prediction.

The trick that makes training cheap

Here's the part people usually miss when explaining transformers. Because of the causal mask, position $t$ can only attend to positions $\le t$ . So a single forward pass over a 1,024-token document produces 1,024 predictions simultaneously, and every one of them is a valid training example — position 5 predicts token 6, position 6 predicts token 7, and so on. No leakage, because position 5 physically cannot see token 6.

In code, the setup is a one-line shift:

input_ids  = tokens[:-1]   # what the model sees
target_ids = tokens[1:]    # what it should predict at each position
logits = model(input_ids)  # shape: (seq_len, vocab_size)

This is called teacher forcing: during training the model always conditions on the true prefix, never on its own guesses. It's why training is one big parallel matrix computation while generation (next section) is stubbornly sequential.

2. Cross-entropy loss

For one position with target token $y$ , the loss is the negative log-probability the model assigned to the correct token:

$\mathcal{L} = -\log p_y = -\log \frac{e^{z_y}}{\sum_j e^{z_j}}$

Read it as: "How surprised was the model by what actually came next?" Perfect prediction ( $p_y = 1$ ) gives loss $0$ . Assigning $p_y = 0.01$ gives loss $\approx 4.6$ . The total loss for a batch is the mean over all positions and all sequences.

Worked example

Suppose the context is "the cat sat on the" and, for simplicity, the vocabulary has 6 tokens. The model produces these logits:

| token | mat | rug | floor | cat | banana | the | |---|---|---|---|---|---|---| | logit $z_i$ | 2.0 | 1.0 | 0.5 | 0.0 | −1.0 | −2.0 |

Exponentiate: $e^{2.0}=7.389$ , $e^{1.0}=2.718$ , $e^{0.5}=1.649$ , $e^{0}=1.000$ , $e^{-1}=0.368$ , $e^{-2}=0.135$ . Sum $= 13.259$ .

$p = [0.557,\; 0.205,\; 0.124,\; 0.075,\; 0.028,\; 0.010]$

If the true next token is mat:

$\mathcal{L} = -\log(0.557) = 0.585$

If it had been rug:

$\mathcal{L} = -\log(0.205) = 1.585$

Note that the loss depends only on the probability of the correct token — but that probability depends on all the logits, because of the normalizing sum.

Why the gradient is so clean

Differentiate cross-entropy with respect to the logits and almost everything cancels:

$\frac{\partial \mathcal{L}}{\partial z_i} = p_i - \mathbb{1}[i = y]$

In our example, with target mat, the gradient is $[0.557 - 1,\; 0.205,\; 0.124,\; 0.075,\; 0.028,\; 0.010] = [-0.443,\; 0.205,\; \ldots]$ . Gradient descent pushes the correct logit up and pushes every wrong logit down in proportion to how much probability it wrongly claimed. banana barely gets touched; rug gets a real shove. This signal then flows backward through the unembedding, the MLPs, and the attention weights — which is how attention heads end up learning to do things like "copy the token that followed this same bigram earlier in the document."

Perplexity

Loss values are usually reported as perplexity, $\text{PPL} = e^{\mathcal{L}}$ , which you can read as "the effective number of tokens the model is choosing between." Our example loss of $0.585$ is a perplexity of $1.79$ — the model was effectively picking between two options. A good modern LM on ordinary web text sits around loss $2$ – $3$ , i.e. perplexity $7$ – $20$ .

3. Generation: the autoregressive loop

Training was parallel. Generation is not, because the model's output at step $t$ becomes its input at step $t+1$ .

tokens = tokenizer.encode(prompt)

for _ in range(max_new_tokens):
    logits = model(tokens)          # (len(tokens), vocab_size)
    next_logits = logits[-1]        # only the LAST position matters now
    probs = softmax(next_logits / temperature)
    next_token = sample(probs)      # or argmax, for greedy decoding
    tokens.append(next_token)
    if next_token == EOS:
        break

print(tokenizer.decode(tokens))

Three things worth saying out loud to a colleague:

The model still computes logits at every position; we throw away all but the last. The earlier ones would just re-predict tokens we already have.
This is where the KV cache comes in. Recomputing keys and values for the whole prefix at every step is wasteful, and because of the causal mask, the keys and values for token 5 never change when token 200 arrives. So implementations cache them: each new step computes a query for one new position and attends against the cached $K$ and $V$ . That turns per-step cost from $O(t^2)$ to $O(t)$ and is the single biggest reason inference is affordable.
The model has no plan. It has no representation of the sentence it's "about to write." Coherence over a paragraph is an emergent consequence of each step conditioning on everything already emitted.

4. Temperature

If you always take argmax (greedy decoding), the model is deterministic and tends to loop: "the best way to do that is to do that is to do that…" Sampling from $p$ fixes this but sometimes picks genuinely bad tokens. Temperature $T$ is the dial between those failure modes. It divides the logits before the softmax:

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

Using the same six logits from before:

| token | $T = 0.5$ | $T = 1.0$ | $T = 2.0$ | |---|---|---|---| | mat | 0.829 | 0.557 | 0.357 | | rug | 0.112 | 0.205 | 0.216 | | floor | 0.041 | 0.124 | 0.168 | | cat | 0.015 | 0.075 | 0.131 | | banana | 0.002 | 0.028 | 0.080 | | the | 0.000 | 0.010 | 0.048 |

$T < 1$ sharpens: gaps between logits get amplified, the top token dominates. As $T \to 0$ this becomes greedy decoding.
$T = 1$ is the model's honest, calibrated distribution — the one it was trained to produce.
$T > 1$ flattens: rare tokens become plausible. As $T \to \infty$ it approaches uniform sampling over the entire vocabulary, i.e. noise. Notice banana went from a 1-in-500 shot to a 1-in-12 shot.

An important nuance: temperature is a monotonic transformation. It never reorders the tokens — mat is the most likely candidate at every temperature. It only changes how often you gamble on the rest of the tail.

Because the tail of a 50,000-token vocabulary contains a lot of garbage, temperature is usually combined with a truncation step:

Top- $k$ : keep only the $k$ highest-probability tokens, renormalize, sample.
Top- $p$ (nucleus): keep the smallest set of tokens whose probabilities sum to $p$ (e.g. 0.9), renormalize, sample. Adaptive — a confident step keeps 2 tokens, an ambiguous one keeps 40.

Typical practice: $T \approx 0.7$ with top- $p = 0.9$ for prose; $T \to 0$ for code, extraction, or anything where you want reproducibility.

5. The 90-second version for your colleague

The model reads a sequence and, at every position, outputs a score for every token in the vocabulary. Softmax turns those scores into probabilities. Training just means: show it real text, and penalize it by the negative log of the probability it assigned to the token that actually came next — that's cross-entropy. Because attention is masked so each position can only look leftward, one pass over a document gives you thousands of these predictions at once, which is why transformers train so efficiently.

To generate, you run the same forward pass, look only at the last position's distribution, pick a token from it, append it, and repeat. Temperature is a knob applied to the scores before the softmax: divide by a small number and the distribution sharpens toward the single best guess; divide by a large number and it flattens toward random. Nothing else about the model changes — same weights, same attention — you're just deciding how much you want it to gamble.

Try it yourself

Open a Python REPL — no ML framework needed, just numpy:

Write softmax(logits, T=1.0) that returns probabilities. (Watch out for overflow: subtract logits.max() before exponentiating. This is exactly what production implementations do, and it's a nice thing to be able to explain.)
Reproduce the table in Section 4 with the logits [2.0, 1.0, 0.5, 0.0, -1.0, -2.0] at $T \in \{0.1, 0.5, 1, 2, 10\}$ . At what temperature does banana cross 5%?
Write cross_entropy(logits, target_index) and verify that it equals $-\log p_{\text{target}}$ from your softmax. Then confirm numerically that the gradient of the loss with respect to the logits is $p - \text{onehot}(y)$ , by nudging one logit by $\pm 10^{-5}$ and taking a finite difference.
Finally, say the Section 5 summary out loud without reading it, and notice which sentence you stumble on. That's the part to re-read.