Optional Foundation — Build a Small Language Model
How the weights actually move
You have a loss; now watch it drive the weights. This chapter is the training loop itself — backprop to find which direction lowers the loss, an optimizer to take the step, and the practical machinery (epochs, learning-rate schedule, validation curves, gradient accumulation) that turns random numbers into a trained model.
You can now measure how wrong the model is with a single number, the loss. Training is the loop that uses that number to move every weight a little in the direction that lowers it, over and over, millions of times. This chapter is that loop — the mechanism (backprop + an optimizer) and the practical machinery (epochs, learning-rate schedule, validation curves, gradient accumulation) that actually turns random matrices into a working model.
What you will understand by the end
- The four-step loop every training run repeats: forward → loss → backward → update.
- What backpropagation and gradient descent do, in intuition, not calculus.
- Why the learning rate (and its schedule) is the most important training knob.
- How to read train vs validation loss, and what epochs and gradient accumulation are for.
The loop, in four steps
Every training run — from a toy model to a frontier one — repeats the same four steps on batch after batch of data:
┌─────────────────────────────────────────────┐
│ 1. FORWARD run the batch → logits → loss │
│ 2. LOSS one number: how wrong we are │
│ 3. BACKWARD backprop: ∂loss/∂each weight │ ← the gradient
│ 4. UPDATE optimizer nudges each weight │
└───────────────┬─────────────────────────────┘
└── repeat on the next batch, millions of times
Steps 1–2 you already have. The new machinery is steps 3 and 4: figuring out which way to move each weight, and taking the step.
Backpropagation: which way is downhill
The loss is a function of every parameter — change any weight and the loss changes a little. Backpropagation computes, for each weight, the gradient: the direction and rate at which nudging that weight would change the loss. It does this efficiently by applying the chain rule backward through the network, from the loss to every parameter, in one pass.
Backprop answers one question for every weight at once: "if I nudge you up a hair, does the loss go up or down, and how fast?" That per-weight answer is the gradient. It does not change anything — it just points downhill for every parameter simultaneously.
Gradient descent and the optimizer
Once you know the downhill direction, you step the weights a small amount that way:
weight ← weight − learning_rate × gradient. Repeat, and the loss walks downhill. That is
gradient descent.
In practice nobody uses plain gradient descent; they use an optimizer like Adam / AdamW that adapts the step size per parameter using the recent history of its gradients (momentum and scaling). It converges faster and more stably than raw SGD, which is why it is the default for training language models. The idea is unchanged — step downhill — the optimizer just chooses a smarter step.
The learning rate is the master knob
The learning rate sets how big each step is, and it is the single most consequential training hyperparameter:
- Too high and the steps overshoot the valley — the loss bounces or diverges to
NaN. - Too low and training crawls, wasting compute and possibly getting stuck.
Real runs use a schedule: a short warmup (start tiny, ramp up) so the early, chaotic steps don't blow up, followed by a decay (shrink the rate over time) so later steps settle into a good minimum. Think explore-then-refine.
The learning rate is where most training runs live or die. A loss that spikes to NaN early is almost always too high a rate (or missing warmup); a loss that barely moves is often too low. Tune this before almost anything else — architecture rarely rescues a bad learning rate.
Reading the curves: epochs, train vs validation
- An epoch is one full pass over the training data. Classic ML often runs hundreds; large language-model pre-training typically makes only a handful of passes (sometimes less than one) over an enormous corpus — there is simply so much data that the model rarely needs to see it many times.
- Plot training loss and validation loss (on the held-out split from the data chapter) together. Both falling → learning. Training loss falling while validation loss rises → overfitting: the model is memorising the training set instead of the language. The validation curve is your signal for when to stop.
Training loss alone can always be driven down (memorise harder). The validation loss is what tells you the model is learning something that generalises. Watch the gap between the two — it is the same "did it learn the pattern or just the answers?" question the evaluation split asks.
Gradient accumulation: big batches on small GPUs
Large batches make training more stable, but a big batch may not fit in GPU memory alongside the weights and activations. Gradient accumulation is the workaround: run several micro-batches, add up their gradients without updating, then take one optimizer step on the sum. You get the statistics of a large batch with the memory footprint of a small one — the training-time echo of the same memory pressure you met on the inference side.
Mental model
Loop forever: forward, loss, backprop (which way is downhill for every weight), step (optimizer nudges each weight that way). The learning rate sets the step size and is the knob that makes or breaks it; validation loss tells you whether you're learning or memorising; epochs and gradient accumulation are how you pace and fit the run.
Common mistakes
- Mis-setting the learning rate. Too high diverges, too low crawls — the first thing to get right, with warmup and decay.
- Watching only training loss. It can always go down; without the validation curve you can't see overfitting.
- Expecting many epochs. LM pre-training often needs only a few passes over a large corpus, not hundreds.
- Ignoring memory when picking batch size. If the batch won't fit, use gradient accumulation rather than a tiny, noisy batch.
Practical guidance
- Set a learning-rate schedule (warmup + decay) from the start; treat the peak rate as the hyperparameter you tune first.
- Log train and validation loss every so often and stop (or checkpoint) when validation stops improving.
- Use AdamW as a sane default optimizer; reach for exotic ones only with a measured reason.
- When batches don't fit, reach for gradient accumulation before shrinking the batch — keep the large-batch stability, pay it in steps.
- Checkpoint regularly so a diverged run or a crash doesn't cost the whole thing.
Summary
- Training repeats forward → loss → backward → update on batch after batch.
- Backprop gives the gradient (downhill direction for every weight); an optimizer (Adam/AdamW) takes the step; the learning rate sizes it and is the master knob.
- Validation loss vs training loss reveals learning vs overfitting; LM pre-training needs only a few epochs over a big corpus.
- Gradient accumulation buys large-batch stability within a small memory budget.
Knowledge check
Your training loss falls steadily but validation loss bottoms out and starts rising. What is happening, and what do you do?
The model is overfitting — it is now memorising the training set rather than learning patterns that generalise, so it does better on seen data and worse on held-out data. Stop at (or roll back to) the point of lowest validation loss — that checkpoint is your best model. Options to push further: more or cleaner data, regularisation, or a smaller model. The validation curve, not the training curve, decides when you're done.
Backprop just finished for a batch. Has anything about the model changed yet? What actually moves the weights?
No — backprop only computes the gradients (the downhill direction for each weight); it
changes nothing. The optimizer's update step is what moves the weights, taking
weight − learning_rate × gradient (Adam/AdamW adapting the step per parameter). Gradient =
direction; optimizer step = the actual move.
Related chapters
- The forward pass and the loss — the number this loop minimises
- The model as parameters to be learned — the weights that these steps move
- Development, validation and hidden sets — the held-out data the validation loss reads
- From trained weights to generation — what you do once the loop converges