Training

What is being minimised

The model outputs, for each position, a vector of V logits — one unnormalised score per vocabulary token. Softmax turns those into a probability distribution over what comes next. The loss is the negative log of the probability the model assigned to the token that actually came next:

This objective has a property worth appreciating: it is not "get it right", it is "be confident and right, and be uncertain when uncertain". Assigning 0.9 to the correct token costs 0.105 nats; assigning 0.5 costs 0.69; assigning 0.01 costs 4.6. The penalty for confident wrongness is unbounded, which is what stops the model from bluffing.

The lab exposes the loss function itself as editable code. The default calls a fused softmax-and-cross-entropy op, and the comment in it explains why: computed separately, the gradient involves log of a probability that may have already rounded to zero, giving -Infinity. Fused, the gradient is just (p − onehot), which is always finite. Numerical stability in deep learning is largely a story of finding forms like this one.

The loop

for step in range(max_steps):
    batch      = sample_random_windows(data)   # [B, T] tokens
    logits     = model(batch[:, :-1])          # forward
    loss       = cross_entropy(logits, batch[:, 1:])
    grads      = backward(loss)                # every parameter's derivative
    grads      = clip_by_norm(grads, 1.0)      # tame the outliers
    params     = optimizer.step(params, grads) # move downhill

Forward and backward

The forward pass records every operation it performs into a graph. The backward pass walks that graph in reverse, applying the chain rule at each node, and arrives at ∂loss/∂p for every parameter p in the model — the direction in which increasing that parameter would increase the loss. Backpropagation is exactly the chain rule with the intermediate results cached so nothing is computed twice. It costs roughly twice a forward pass, which is why the standard estimate for training compute is about 3× the forward cost per token.

Why batches

Computing the true gradient would mean a forward and backward pass over the entire dataset for one parameter update. Instead each step uses a random batch, giving a noisy but unbiased estimate — stochastic gradient descent. Bigger batches mean less noisy gradients and better hardware utilisation; smaller batches mean more updates for the same compute, and the noise itself acts as a mild regulariser. The trade is one of the more studied and less settled questions in the field.

Adam, and why nobody uses plain SGD here

Plain gradient descent — p ← p − lr · g — has one learning rate for every parameter in the model. That is a poor fit for a transformer, where an embedding row for a rare token and a weight in the output projection see gradients differing by orders of magnitude. Adam gives each parameter its own effective step size, derived from its own gradient history.

m = β₁·m + (1 − β₁)·g          # momentum: running mean of the gradient
v = β₂·v + (1 − β₂)·g²         # running mean of the squared gradient
m̂ = m / (1 − β₁^t)             # bias correction — m starts at 0
v̂ = v / (1 − β₂^t)
p = p − lr · m̂ / (√v̂ + ε)      # step, normalised by recent magnitude

• m (momentum, β₁ ≈ 0.9) smooths the direction over roughly the last ten steps, damping the noise from batch sampling.
• v (β₂ ≈ 0.999) tracks typical gradient magnitude over roughly the last thousand steps. Dividing by its square root means a parameter with consistently tiny gradients still takes meaningful steps, and one with huge gradients does not overshoot.
• Bias correction matters because both averages start at zero; without it the first few dozen steps would be far too small.
• ε (≈ 1e-8) stops division by zero for a parameter whose gradient has been zero throughout.

The cost is memory: Adam stores two extra floats per parameter, so optimiser state is twice the model size. For a 7B-parameter model in fp32 that is 56 GB of optimiser state on top of 28 GB of weights — one of the main reasons training a model needs far more memory than running one.

The learning rate is the knob that matters

If you can tune exactly one hyperparameter, tune this one. It sets how far each step moves, and both failure directions are recognisable once you have seen them.

The "too low" case is the expensive one, because nothing looks wrong. The curve is clean and descending; it will simply still be descending when your budget is gone. Both failure modes are one button each in the lab, and the second is worth staring at precisely because it looks so healthy.

Schedules

Real runs vary the learning rate over time. Warmup ramps it up from near zero over the first few hundred steps, because a large step taken while Adam's variance estimates are still unreliable is how runs die in the first minute. Cosine decay then anneals it smoothly toward zero, so the model takes large exploratory steps early and fine ones at the end. The combination is close to universal.

How runs die

Gradient clipping, and the spikes it saves you from

Occasionally a batch produces an enormous gradient — an unusual sequence, an unlucky interaction of weights. One such step can undo thousands of good ones. Clipping computes the global norm of all gradients and, if it exceeds a threshold (1.0 is the usual choice), scales the whole vector down to that norm. Direction is preserved; only magnitude is capped. It costs nothing when nothing is wrong, which is why essentially every serious run uses it.

Watch the grad norm stat during the run. A healthy run shows it high at the start and settling into a stable band. If it climbs steadily, the run is heading for trouble.

The NaN cascade

NaN is not a random glitch; it has a specific and repeatable life cycle. An oversized step produces huge activations; exp() of a large logit overflows float32 to Infinity; Infinity − Infinity or 0 × Infinity produces NaN; NaN propagates through every arithmetic operation it touches; one backward pass later every parameter in the model is NaN. There is no recovery — the information is gone. Real training runs checkpoint frequently for exactly this reason, and restarting from the last checkpoint with a lower learning rate is a normal operational event at scale.

Loss is a proxy. Score the thing you actually want.

Falling loss proves the optimiser is working. It does not prove the model learned what you think. Switch the corpus to Arithmetic and the distinction becomes concrete: most of the available loss reduction comes from learning the format of an addition problem — digits, then a plus, then digits, then an equals — and a model can do all of that while being unable to add.

So the lab scores exact-match accuracy on freshly generated sums: greedy-decoded, compared character for character, no partial credit. That number cannot be argued with. Single-digit addition is reachable at 100% by a model this small — and it typically arrives suddenly, over a few dozen steps, after a long plateau where the loss was improving but accuracy was flat. That shape is the small-scale version of the "emergent capability" phenomenon reported in large models: a smooth loss curve hiding a sharp transition in a capability metric.

Two-digit addition requires carrying and has ten thousand distinct problems, so memorisation is unavailable. Expect roughly a quarter correct at these settings. That gap — between a task the model can genuinely learn and one it can only half-learn at this scale — is the honest lesson of the chapter.

Vocabulary

Cross-entropy loss

Negative log probability of the correct next token, averaged over positions. Measured in nats.

Logits

The model's raw per-token scores before softmax.

Backpropagation

Reverse-mode application of the chain rule to obtain every parameter's gradient.

Adam / AdamW

Optimiser with per-parameter step sizes from running estimates of gradient mean and variance.

Learning rate

How far each step moves. The single most consequential hyperparameter.

Gradient clipping

Capping the global gradient norm so one bad batch cannot destroy the run.

Warmup / cosine decay

Ramping the learning rate up early and annealing it to zero late.

Check yourself

If you can answer these without re-reading, the lab below will make sense. If you cannot, the relevant section is worth a second pass — that is a better use of your time than clicking buttons.

1. Where do the training labels come from, and why is that called self-supervised?
2. What does dividing by √v̂ accomplish in Adam that a single global learning rate cannot?
3. Describe the loss curve of a learning rate that is slightly too high, and of one that is far too low.
4. Walk through the chain of events from one oversized step to every parameter being NaN.
5. Gradient clipping changes the magnitude of the update but not its direction. Why is that the right trade?
6. Loss on the arithmetic corpus is falling steadily but exact-match accuracy is 0%. What is the model learning?

Where this comes from

• Kingma & Ba, "Adam: A Method for Stochastic Optimization" (2014) — The optimiser, in six pages.
• Loshchilov & Hutter, "Decoupled Weight Decay Regularization" (2017) — Why AdamW replaced Adam.
• Karpathy, "A Recipe for Training Neural Networks" — The practical debugging methodology, including overfitting a single batch.