The single most important number in a run, and how it is found: a learning-rate sweep, warmup, cosine and warmup–stable–decay schedules, and the batch size, with the critical batch and the gradient noise scale measured on the model.
This deck follows chapter 6, learning rate, batch size and schedules, where you can run the sweep yourself. Measured numbers: results.md §6a–6d, every setting three data seeds, mean ± standard deviation.
The gradient gives a direction; the learning rate \(\eta\) decides the step. Too small and the model crawls; too large and the steps overshoot. In between is a band where it trains well.
Ends 2,000 steps at 2.948, still on deck 01’s plateau.
2.277 ± 0.090; 0.03 is indistinguishable (2.280 ± 0.107).
Ends at 2.906; recomputed on another computer, one of its seeds blew up.
None of these runs diverged outright in the recorded results: gradient clipping and Adam’s normalised step usually keep such a run finite; it just learns little. Plain SGD and momentum have no such safety net (deck 07 blows one up with a single bad batch).
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. AdamW, warmup 50 then cosine to a tenth of the peak, 2,000 steps, clipping at 1.0, three seeds each (results.md §6a):
Every optimiser and every method has its own best rate, so a fair comparison tunes the learning rate for each one separately: deck 04 compared optimisers each at the best of its own three rates, and deck 12 does the same for distributed-training methods, because a method that changes what the optimiser sees changes the rate it wants.
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 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; deck 06).
This tiny pre-LN model at a peak of 0.01 trains fine without warmup: the run that starts at the full rate is at 2.382 at step 1,000 against 2.640 with 50 warmup steps, and ends at 2.159 ± 0.045 against 2.246 ± 0.025. Warmup is insurance that this model happens not to need.
At the end of a run a smaller learning rate lets the weights settle into a valley instead of bouncing across it. The cosine schedule (Loshchilov & Hutter 2016) decays smoothly from the start; LLaMA used it down to 10% of the peak (Touvron et al. 2023, §2.3). Its drawback: the run’s length must be fixed in advance. WSD (Hu et al. 2024, studied by Hägele et al. 2024) holds the peak and decays only at the end, so a run can be extended from any checkpoint of the stable phase.
AdamW at peak learning rate 0.01, 2,000 steps, three seeds, warmup 50 where used, WSD decaying over the last 20% to a tenth of the peak (results.md §6b):
| Schedule | Final validation loss (mean ± sd) | Seeds 1, 2, 3 |
|---|---|---|
| constant | 2.159 ± 0.045 | 2.111, 2.165, 2.201 |
| warmup + constant | 2.246 ± 0.025 | 2.275, 2.238, 2.227 |
| warmup + cosine | 2.277 ± 0.090 | 2.299, 2.354, 2.178 |
| WSD | 2.224 ± 0.038 | 2.264, 2.220, 2.189 |
WSD is the same run as warmup + constant until step 1,600; then its decay pulls it below. The cosine run ends where it does because a 2,000-step run that is still improving at its end gains little from spending its last steps at a small learning rate. Decay pays off in long runs that are close to converged; this toy isn’t, and the spreads overlap.
Each step’s gradient is an average over the batch: a noisy estimate of the true gradient. 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 well below a critical batch size, and barely helps above it.
$$S = S_{\min}\Big(1 + \frac{B_{\text{crit}}}{B}\Big)$$Steps to a validation loss of 2.6 at five batch sizes, each at the better of two learning rates, three seeds (results.md §6c). From 4 to 64 windows the steps fall from 1,950 to 625, far less than the 16-fold the batch grew. The fit gives \(B_{\text{crit}}\) = 10.4 windows, \(S_{\min}\) = 592 steps.
McCandlish et al.’s quick estimate of the critical batch is the gradient noise scale: the gradient’s variance over its squared size (their eq. 2.9), the batch at which noise is as large as signal.
$$B_{\text{simple}} = \frac{\operatorname{tr}(\Sigma)}{|G|^2}$$| Step | tr(Σ) | |G|2 | Bsimple (windows) |
|---|---|---|---|
| 0 | 1.6e+00 | 5.8e-01 | 2.8 |
| 250 | 1.9e+00 | 3.1e-02 | 62.8 |
| 500 | 1.3e+00 | 1.0e-02 | 126.4 |
| 1,000 | 2.3e+00 | 2.8e-02 | 79.5 |
| 1,500 | 9.0e+00 | 5.6e-02 | 161.4 |
| 2,000 | 1.2e+01 | 2.6e-02 | 455.7 |
Measured from 256 single-window gradients along the default run (results.md §6d). It agrees with the fitted 10 windows only to within about an order of magnitude, which is the precision McCandlish et al. claim for it. The trend is the useful part: as the loss falls, the gradient’s signal 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): 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.
SGD, below the critical batch, after warmup (Goyal et al. trained ImageNet with large batches this way, with a warmup).
Adam’s normalised step: the step size no longer scales with the gradient. Large language-model runs therefore tune the batch size and learning rate together, and the measured best rate here changed between batch sizes (results.md §6c).
Batch size is also where parallelism meets optimisation: more data-parallel workers means a bigger batch per step (deck 10), and the critical batch says how far that helps.
Animated: chapter 6, learning rate, batch size and schedules · next: deck 06, initialisation, normalisation and µP