Learning rate, batch size and schedules

Why the learning rate matters most. The gradient gives a direction; the learning rate $\eta$ decides the step. Too small and the model crawls: the Decoder's model at a peak of 0.001 ends 2,000 steps at 2.948, still on the plateau of chapter 1. Too large and the steps overshoot: at 1 it ends at 2.906. In between is a band where it trains well: the best mean here is at 0.01 (2.277 ± 0.090), with 0.03 indistinguishable (2.280 ± 0.107). None of these runs diverged outright, though 1.0 is on the edge: recomputed on another computer, one of its three seeds blew up (results.md §6). Gradient clipping and Adam's normalised step (chapter 5) usually keep such a run finite; it just learns little. Plain SGD and momentum have no such safety net: chapter 8 shows one bad batch blowing a momentum run up.

A learning-rate sweep is how that band is found: train the same model, for the same number of steps, at a ladder of learning rates a factor of about 3 apart, and keep the best. The chart below is that sweep for the Decoder's model (AdamW, warmup then cosine, three seeds at each rate, mean ± standard deviation). Every optimiser and every method has its own best rate, so a fair comparison tunes the learning rate for each one separately: chapter 5 compared four optimisers each at the best of its own three rates, and a comparison of distributed-training methods (chapter 13) does the same, because a method that changes what the optimiser sees changes the rate it wants. Because training at the higher rates is chaotic (the same run on another computer can move by about 0.1), the seeds matter as much as the ladder.

You can run this sweep yourself: the trainer below is chapter 1's, with four peak learning rates. Each label quotes the recorded run with the same settings (data seed 1); at the higher rates the run is chaotic, so your browser's may land somewhere else.

Loading the animation…

Warmup. Most runs start with a few hundred to a few thousand steps of linearly rising learning rate. Early on, the gradients are large and change fast, and Adam's running averages (chapter 5) are built from only a few steps, so a full-size first step can throw the weights somewhere bad. Warmup was popularised for large-batch training by Goyal et al. (2017) and is essential for post-LN transformers (Xiong et al. 2020; chapter 7). This tiny pre-LN model at a peak of 0.01 does not need it: the run that starts at the full rate leaves chapter 1's plateau sooner (validation loss 2.382 at step 1,000, against 2.640 with 50 warmup steps) and keeps the lead: 2.159 ± 0.045 at the end, against 2.246 ± 0.025. Warmup is insurance that this model happens not to need.

Decay: cosine and WSD. At the end of a run, a smaller learning rate lets the weights settle into a valley instead of bouncing across it (the jitter of chapter 5's valley). The cosine schedule (Loshchilov & Hutter 2016) decays smoothly from the start; LLaMA, for example, used it down to 10% of the peak (Touvron et al. 2023, §2.3). Its drawback is that the length of the run must be fixed in advance. Warmup–stable–decay (WSD; Hu et al. 2024, studied by Hägele et al. 2024) holds the peak and decays only over the last part, so a run can be extended from any checkpoint of the stable phase. In the animation, WSD is the same run as warmup + constant until step 1,600; then its decay pulls it below: 2.224 ± 0.038 at the end, against 2.246 ± 0.025. The cosine run ends at 2.277 ± 0.090: a 2,000-step run that is still improving at its end gains little from spending its last steps at a small learning rate. The decay pays off in long runs that are close to converged.

Batch size and gradient noise. Each step's gradient is an average over the batch: a noisy estimate of the gradient over all the text. A bigger batch gives a better estimate, so the same progress needs fewer steps, but each step costs more. McCandlish et al. (2018) model the trade-off: doubling the batch halves the steps while the batch is well below a critical batch size, and barely helps above it, where the extra examples only polish an already good gradient. The chart measures it on the Decoder's model, whose examples are windows of 8 characters (its whole context, Decoder, chapter 2): the steps to reach a validation loss of 2.6 at five batch sizes, each at the better of two learning rates, three seeds. Fitting their law $S = S_{\min}(1 + B_{\text{crit}}/B)$ gives a critical batch of about 10 windows here: from 4 to 64 windows the steps fall from 1950 to 625, far less than the 16-fold the batch grew. Their quick estimate of the critical batch, the gradient noise scale $B_{\text{simple}} = \operatorname{tr}(\Sigma)/|G|^2$ (the gradient's variance over its squared size, measured here from 256 single-window gradients), grows as training goes on: 2.8 windows at initialisation, 79 at step 1,000 and 456 at the end of the default run. The two estimates agree only to within about an order of magnitude here, which is the precision McCandlish et al. claim for it. The trend is the useful part: as the loss falls, the signal in the gradient shrinks relative to its noise, so later in training a bigger batch is worth more, which is why some large runs grow their batch during training.

The linear scaling rule (Goyal et al. 2017) says: when you multiply the batch by $k$, multiply the learning rate by $k$ too (up to a limit), since $k$ small steps in about the same direction add up to one $k$-times larger step. Adam's normalised step weakens that relation, which is why large language-model runs tune the batch size and learning rate together.