Vol. XI · Optimisation at Scale Plate 2 of 10batch size, gradient noise, critical batch size, gradient noise scale, data parallelism, steps versus examplesLC

Batch Size.

Twice the batch, half the steps — until the gradient is already as clean as it needs to be.

A batch of B sequences gives a gradient whose noise shrinks as 1/√B. A cleaner gradient allows bigger, surer steps, so larger batches reach a target in fewer steps — which matters because steps are sequential while the examples within a step can be spread over thousands of accelerators (Vol. XIV). The question is how far this goes.

McCandlish and colleagues (2018) found a simple law across many tasks: the steps S and examples E needed for a target obey (S/S_min − 1)(E/E_min − 1) = 1. Small batches are example-efficient but slow; huge batches are fast but waste examples. Between them sits the critical batch size B_crit ≈ E_min / S_min, where doubling the batch still nearly halves the steps. It grows during training, as the gradient’s signal shrinks relative to its noise. The workshop measures it on the noisy quadratic model, tuning the learning rate at each batch size.

Below: left, steps to reach the target loss against batch size. Right, examples consumed against steps — the trade-off curve — with the critical batch size marked.

What you will learn
Workshop
Noisy quadratic model, momentum SGD with a constant learning rate tuned from a grid at each batch size; steps until the loss first falls below 0.02. Left: steps against batch size (log–log); the dashed line is perfect scaling, steps ∝ 1/B. Right: examples (steps × batch) against steps; E_min and S_min are read from the smallest and largest batches.

Fewer steps, more examples

With the learning rate tuned at each size, a batch of 4 reaches the target 3.4 times faster than a batch of 1, against an ideal of 4; a batch of 16 is 5.9 times faster, already well short of 16; and a batch of 256 only 18.2 times faster and uses 14.0 times the minimum number of examples. The trade-off curve is the hyperbola McCandlish described, and its corner gives a critical batch of about 18. Raise the gradient noise to 3 and the critical batch grows to 71; lower it to 0.5 and it shrinks to 6: the noisier the gradient, the more a large batch helps.

In real training the noise grows relative to the signal as the loss falls, so the critical batch grows too — which is why many large runs ramp the batch size up during training, and why the largest runs use batches of millions of tokens. Llama 3 405B, for example, started at 4 million tokens per batch and doubled it twice. Past the critical batch, more accelerators buy little speed: the step count stops falling, and the extra examples are spent for nothing.

Learning rate and batch size are linked: a larger batch tolerates a larger learning rate, roughly in proportion for SGD at small batches and more weakly for Adam. Comparing batch sizes without retuning the learning rate understates the larger batches.

Terms

Batch size
The number of sequences (or tokens) whose gradients are averaged in one step.
Gradient noise
The difference between a batch’s gradient and the true gradient; shrinks as 1/√B.
Critical batch size
The batch beyond which doubling it no longer nearly halves the steps.
Steps–examples trade-off
(S/S_min − 1)(E/E_min − 1) = 1.
Larger batches need fewer steps but more examples; the critical batch size marks where the trade stops paying, and it grows as training proceeds.

Sources

  1. McCandlish, Kaplan, Amodei and the OpenAI Dota Team, “An Empirical Model of Large-Batch Training”, arXiv:1812.06162, 2018.
  2. Zhang et al., “Which Algorithmic Choices Matter at Which Batch Sizes?”, NeurIPS 2019.