08

Training

How a stack of random weights becomes a language model: one loss, one gradient, one careful step at a time.

Key terms

TermMeaning
ObjectiveThe single number training tries to make small (the loss)
Cross-entropyLoss that measures how surprised the model is by the true next token
Surprise-log(probability) the model assigned to what actually happened
GradientThe direction in weight-space that increases the loss fastest
Gradient descentRepeatedly stepping the weights opposite the gradient
Learning rateHow big each step is (the lr)
MinibatchA small group of examples averaged into one gradient
AdamAn 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:

ExampleTrue tokenp(true)loss = -ln(p)
1B0.600.511
2A0.301.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:

lrmultiplier (1 − 2·lr)x0x1x2x3behavior
0.10.810.80.640.512converges smoothly
0.50.01000jumps to the minimum
1.0−1.01−11−1oscillates forever
1.5−2.01−24−8diverges (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.

IdeaPlain gradient descentAdam
Step directionRaw gradientMomentum-smoothed gradient
Step sizeOne fixed lr for allPer-weight, scaled by recent gradient size
Robust to bad lrFragileMore 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 is x·(1−2·lr), so lr>1 diverges and lr=1 oscillates — 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 the f(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?