An Empirical Model of Large-Batch Training

TL;DR

Introduces the gradient noise scale — a simple, cheap-to-measure statistic that predicts the largest useful batch size for SGD-style training across supervised learning, RL, and generative modeling. Above the noise scale, doubling batch size wastes compute; below it, batches parallelize almost linearly. Defines the critical batch size $B_\mathrm{crit}$ as the point on the compute-vs-steps Pareto frontier where training takes 2× minimum steps and 2× minimum examples, and shows it can be predicted (to order of magnitude) from the noise scale. Verified across 8 tasks with critical batch sizes ranging from 20 (SVHN autoencoder) to >10M (Dota). The noise scale rises during training as loss falls, so $B_\mathrm{crit}$ grows over the course of a run.

Key claims

• Under a local quadratic approximation, the optimal step size and expected loss improvement per SGD step follow $\Delta L_\mathrm{opt}(B) \propto 1/(1 + B_\mathrm{noise}/B)$; batches below $B_\mathrm{noise}$ scale nearly linearly in speedup, batches above waste compute [§2.2, Eq. 2.7].
• The simplified noise scale $B_\mathrm{simple} = \mathrm{tr}(\Sigma)/|G|^2$ (sum of per-example gradient variances divided by squared gradient norm) is measurable with negligible overhead in data-parallel training and typically differs from the full noise scale by a small constant [§2.2, Eq. 2.9, Appendix A.1].
• Averaging the local prediction over a full training run gives $S/S_\mathrm{min} = 1 + B_\mathrm{crit}/B$ and $E/E_\mathrm{min} = 1 + B/B_\mathrm{crit}$, where $B = B_\mathrm{crit}$ costs exactly 2× minimum steps and 2× minimum examples [§2.3, Eq. 2.11, 2.12].
• Empirically, the simple noise scale predicts the critical batch size at order-of-magnitude level across 8 tasks spanning MNIST, SVHN, CIFAR-10, ImageNet, Billion Word, Atari, Dota, and a variational autoencoder — critical batch sizes vary by 6 orders of magnitude [Fig. 4].
• $B_\mathrm{crit}$ grows over training as the loss falls; the noise scale increases with better-performing models but this dependence is mediated through loss, not model size directly [§2.5, §3, Fig. 7].
• For SGD with momentum, the optimal LR follows the predicted form $\eta^*(B) \propto 1/(1 + B_\mathrm{noise}/B)$; for Adam/RMSProp the optimal LR initially follows a power law with exponent 0.5–1.0 then plateaus, with the plateau scale near the noise scale [§3.1, Fig. 5].
• Difficulty-dependent noise-scale predictions hold: RL tasks with sparse reward and long horizons have larger noise scales than supervised classification; generative models have smaller noise scales than classifiers on the same dataset [§2.5, §3, Fig. 4].

Method

The paper’s central quantity is derived from a local Taylor expansion. Given a stochastic minibatch gradient $\hat{G}_\mathrm{est}$ that estimates the true gradient $G$ with per-example covariance $\Sigma$, the expected loss decrease at optimal step size $\epsilon_\mathrm{opt} = \epsilon_\mathrm{max}/(1 + B_\mathrm{noise}/B)$ is $\Delta L_\mathrm{opt}(B) = \Delta L_\mathrm{max}/(1 + B_\mathrm{noise}/B)$, where $B_\mathrm{noise} = \mathrm{tr}(H\Sigma)/(G^\top H G)$ and $\epsilon_\mathrm{max}, \Delta L_\mathrm{max}$ are the noiseless-gradient values [§2.2]. Under the simplifying assumption of a well-conditioned Hessian, the noise scale reduces to $B_\mathrm{simple} = \mathrm{tr}(\Sigma)/|G|^2$.

Aggregating this local prediction over a full training run to a target loss gives the tradeoff $S/S_\mathrm{min} = 1 + B_\mathrm{crit}/B$ (steps) and $E/E_\mathrm{min} = 1 + B/B_\mathrm{crit}$ (examples), where the critical batch size $B_\mathrm{crit}$ is fit empirically from Pareto curves at multiple batch sizes. At $B = B_\mathrm{crit}$ training uses 2× minimum steps and 2× minimum examples — a natural sweet spot on the compute/time tradeoff.

Experiments train each task with grid-searched constant learning rates at a wide range of batch sizes, measure Pareto frontiers of steps vs examples to reach target loss, and fit $B_\mathrm{crit}$. The simple noise scale is measured with a low-overhead estimator that computes per-worker gradient statistics in data-parallel training. Predictions are compared to fitted $B_\mathrm{crit}$ across MNIST, SVHN, CIFAR-10, ImageNet, Billion Word, Atari (Alien, Beam Rider, Breakout, Pong, Space Invaders, Q*bert), Dota 5v5, and an SVHN VAE.

Results

• Order-of-magnitude prediction across 6 decades: measured $B_\mathrm{crit}$ ranges from ~20 (SVHN autoencoder) to ~10⁷ (Dota); the simple noise scale predicts within a small constant factor across all 8 tasks [Fig. 4].
• Growth during training: $B_\mathrm{crit}$ typically increases by an order of magnitude or more over a single training run, tracking the loss decrease [§3.2, Fig. 7].
• Model-size dependence is indirect: at fixed loss, $B_\mathrm{noise}$ is nearly independent of parameter count; larger models get larger noise scales only via the mechanism of achieving lower loss [§2.5, empirical confirmation §3].
• Optimizer dependence: SGD/momentum optimal LRs follow the predicted $\epsilon_\mathrm{opt}(B)$ shape; Adam/RMSProp show an initial power-law LR-vs-B regime with exponent 0.5–1.0 that transitions to constant-LR at a scale somewhat below $B_\mathrm{noise}$ [§3.1, Fig. 5].
• Task-difficulty prediction confirmed: Dota’s noise scale (~10⁷) reflects the sparse-reward long-horizon setting; VAE generative-model noise scales are consistently smaller than classifier noise scales on matched datasets [§2.5, Fig. 4].

Why it’s interesting

This paper is the foundational reference for the “critical batch size” framing that appears throughout the wiki’s Hyperparameter scaling laws entries — it’s cited as [MKAT18] in Scaling Laws for Neural Language Models (co-filed here) and re-invoked by The Design Space of Tri-Modal Masked Diffusion Models‘s “SDE reparam eliminates the batch-size search up to $B_\mathrm{crit}$” and How to Allocate Your Tokens? Scaling Laws with Training Steps and Batch Size‘s three-term $L(N, T, B)$ scaling law that keeps steps and batch size as separate axes. The concrete implication for Luma’s video/image sweeps is that $B_\mathrm{crit}$ is loss-dependent — it grows as loss falls — so a fixed target batch size is only optimal at one point in training, and the CBS calculus that LLM teams use (measure gradient noise scale, size the run near it) transfers directly if the noise scale can be measured on our diffusion / flow objectives. Contrasts with Lost in Backpropagation: The LM Head is a Gradient Bottleneck and other single-optimizer papers by giving a task-invariant statistic rather than an optimizer- or architecture-specific one.

See also

• Scaling Laws for Neural Language Models — the Kaplan et al. paper that operationalizes $B_\mathrm{crit}$ into an LLM scaling law and shows it depends on the loss alone
• Hyperparameter scaling laws — the concept page where critical batch size sits alongside Fitting/Transfer/Enforce optimizer scaling
• The Design Space of Tri-Modal Masked Diffusion Models — Apple’s tri-modal MDM design-space study uses SDE reparam to make training loss batch-size-invariant up to $B_\mathrm{crit}$, extending this paper’s local-quadratic framework
• How to Allocate Your Tokens? Scaling Laws with Training Steps and Batch Size — three-term $L(N, T, B)$ scaling law that recovers optimal-batch-size scaling as a consequence of a single fit
• Training stability at scale — where the LR-vs-B optimal-tradeoff shape and its interaction with warmup / weight decay is debated
• Distributed training parallelism — data-parallel batch scaling depends critically on knowing $B_\mathrm{crit}$ to avoid wasting communication bandwidth on above-CBS batches