Training
How a stack of random weights becomes a language model: one loss, one gradient, one careful step at a time.
Key terms
| Term | Meaning |
|---|---|
| Objective | The single number training tries to make small (the loss) |
| Cross-entropy | Loss that measures how surprised the model is by the true next token |
| Surprise | -log(probability) the model assigned to what actually happened |
| Gradient | The direction in weight-space that increases the loss fastest |
| Gradient descent | Repeatedly stepping the weights opposite the gradient |
| Learning rate | How big each step is (the lr) |
| Minibatch | A small group of examples averaged into one gradient |
| Adam | An optimizer that adds momentum and per-weight step sizing |
The objective: predict the next token
Training needs one number to push down. For a language model that number is simple: given the tokens so far, how well did the model predict the actual next token? The whole architecture from day 7 ends in a logit per vocabulary word; softmax (day 3) turns those logits into probabilities.
There are no human labels here — the text is the label. For the sentence "the cat sat", the model is quizzed three times: predict "cat" after "the", predict "sat" after "the cat", and so on. Every position is a next-token prediction problem.
Cross-entropy = average surprise
Once the model outputs a probability for the true next token, we score it. A confident-and-right prediction should cost almost nothing; a confident-and-wrong one should cost a lot. The function with exactly that shape is cross-entropy: the loss for one prediction is -log(p_true), the surprise of the token that actually occurred.
Work a tiny vocabulary {A, B, C, D}. Two predictions in a minibatch:
| Example | True token | p(true) | loss = -ln(p) |
|---|---|---|---|
| 1 | B | 0.60 | 0.511 |
| 2 | A | 0.30 | 1.204 |
The batch loss is the average: (0.511 + 1.204) / 2 = 0.858 nats. Note the pattern — example 2 was less confident about the right answer, so it was punished more (1.204 > 0.511). Perfect confidence p = 1.0 gives loss -ln(1) = 0; a coin-flip-bad guess gives a large loss. Cross-entropy is literally the average surprise across the batch.
Gradient descent
The loss depends on millions of weights. The gradient is the vector of slopes — for each weight, "if I nudge you up, does the loss go up or down, and how steeply?" To reduce the loss we step every weight a little in the downhill direction: weight_new = weight_old - lr × gradient.
Strip it to one weight so the arithmetic is visible. Let the loss be f(x) = x², a simple bowl with its minimum at x = 0. Its gradient is f'(x) = 2x. The update rule becomes:
x_new = x - lr * f'(x)
= x - lr * 2x
= x * (1 - 2*lr)
So each step just multiplies x by the constant (1 - 2*lr). Whether training converges depends entirely on that multiplier — which brings us to the learning rate.
Learning rate (converge or diverge)
The learning rate is the single most important training dial. Too small and learning crawls; too large and it explodes. Using the same f(x) = x² bowl starting at x = 1, watch what the multiplier (1 - 2*lr) does over a few steps:
| lr | multiplier (1 − 2·lr) | x0 | x1 | x2 | x3 | behavior |
|---|---|---|---|---|---|---|
| 0.1 | 0.8 | 1 | 0.8 | 0.64 | 0.512 | converges smoothly |
| 0.5 | 0.0 | 1 | 0 | 0 | 0 | jumps to the minimum |
| 1.0 | −1.0 | 1 | −1 | 1 | −1 | oscillates forever |
| 1.5 | −2.0 | 1 | −2 | 4 | −8 | diverges (blows up) |
The tipping point is exactly lr = 1: at or below it the size of x never grows; above it the size grows every step and the loss races to infinity. Real networks are far more complex, but the failure mode is identical — a too-large learning rate makes the loss increase instead of decrease. Try it yourself: drag the learning rate and step through the bowl.
Play with the slider until you can predict, before pressing Step, whether the next dot lands closer to the bottom or further away.
Minibatches + Adam
Two practical upgrades make training usable at scale. First, minibatches: instead of one example at a time, average the gradient over a small group (say 32 sequences). The average is a less noisy estimate of the true downhill direction, and modern hardware computes a batch in parallel about as fast as a single example.
Second, a smarter optimizer. Plain gradient descent uses one fixed lr for every weight. Adam keeps two running averages per weight: a momentum term (a smoothed gradient, so steps keep rolling in a consistent direction) and a scale term (how large that weight's recent gradients were, so noisy weights take smaller steps). The effect is a per-weight, self-tuning step size, which is why Adam trains transformers far more reliably than raw gradient descent.
| Idea | Plain gradient descent | Adam |
|---|---|---|
| Step direction | Raw gradient | Momentum-smoothed gradient |
| Step size | One fixed lr for all | Per-weight, scaled by recent gradient size |
| Robust to bad lr | Fragile | More forgiving |
Key takeaways
- The training objective is next-token prediction; the text is its own label.
- Cross-entropy loss for one token is
-log(p_true)— the surprise of what actually happened — averaged over the batch. - Gradient descent steps weights opposite the gradient:
weight - lr × gradient. - For
f(x)=x²the update isx·(1−2·lr), solr>1diverges andlr=1oscillates — the learning rate decides everything. - Minibatches denoise the gradient; Adam adds momentum and per-weight step sizing.
Checklist
- [ ] Can you compute the cross-entropy loss for a given probability with
-ln(p)? - [ ] Can you explain why a confident-wrong prediction costs more than an unsure one?
- [ ] Can you derive
x·(1−2·lr)from thef(x)=x²update rule? - [ ] Can you say what happens at lr = 0.1, 1.0, and 1.5, and why?
- [ ] Can you name the two things Adam tracks per weight?