[Paper Review] An Empirical Model of Large-Batch Training

Sam McCandlish, Jared Kaplan+2 more|arXiv (Cornell University)|Dec 14, 2018

Optimization and Search Problems36 references139 citations

TL;DR

The paper introduces the gradient noise scale as a simple statistic that predicts the largest useful batch size across supervised, RL, and generative modeling tasks, and analyzes the tradeoff between compute-efficiency and time-efficiency. It tests across multiple domains and shows the noise scale increases as training progresses and with task difficulty.

ABSTRACT

In an increasing number of domains it has been demonstrated that deep learning models can be trained using relatively large batch sizes without sacrificing data efficiency. However the limits of this massive data parallelism seem to differ from domain to domain, ranging from batches of tens of thousands in ImageNet to batches of millions in RL agents that play the game Dota 2. To our knowledge there is limited conceptual understanding of why these limits to batch size differ or how we might choose the correct batch size in a new domain. In this paper, we demonstrate that a simple and easy-to-measure statistic called the gradient noise scale predicts the largest useful batch size across many domains and applications, including a number of supervised learning datasets (MNIST, SVHN, CIFAR-10, ImageNet, Billion Word), reinforcement learning domains (Atari and Dota), and even generative model training (autoencoders on SVHN). We find that the noise scale increases as the loss decreases over a training run and depends on the model size primarily through improved model performance. Our empirically-motivated theory also describes the tradeoff between compute-efficiency and time-efficiency, and provides a rough model of the benefits of adaptive batch-size training.

Motivation & Objective

• Motivate and understand why batch size limits differ across domains and datasets.
• Introduce the gradient noise scale as a practical predictor of optimal batch size.
• Develop a simple theory linking batch size, gradient noise, and training efficiency.
• Empirically validate predictions across diverse tasks including ImageNet, CIFAR-10, SVHN, MNIST, BillionWord, Atari, and Dota.

Proposed method

• Define and derive the gradient noise scale as B_noise = tr(H Σ) / (G^T H G).
• Relate optimal step size to batch size with ε_opt(B) = ε_max / (1 + B_noise/B).
• Define the simplified noise scale B_simple = tr(Σ) / |G|^2 for practical measurement.
• Predict a Pareto-like tradeoff between training time and compute via a hyperbola around B_crit ~ B_noise.
• Measure B_simple, B_noise, and B_crit across tasks and track how they evolve during training.
• Fit empirical Pareto fronts to assess the alignment with the model’s predictions.

Experimental results

Research questions

• RQ1 What is the gradient noise scale and how does it relate to optimal batch size across tasks?
• RQ2 Can B_simple/B_noise predict the critical batch size where compute efficiency gains degrade?
• RQ3 How does the noise scale evolve during training and across task types (supervised, RL, generative)?
• RQ4 What is the impact of learning rate and conditioning on the observed batch-size tradeoffs?
• RQ5 Do dynamic batch-size adjustments offer efficiency benefits as predicted by the theory?

Key findings

• The gradient noise scale roughly predicts the largest usable batch size at the order-of-magnitude level across tasks.
• Training efficiency follows a Pareto frontier; gains from larger batches diminish beyond the noise scale.
• The noise scale increases during training as the model achieves lower loss.
• B_simple provides a practical estimate of B_crit in many tasks, with B_noise offering closer predictions in some cases.
• The noise scale is larger for more complex tasks (e.g., RL/Dota) and grows as training progresses.
• Dynamic batch-size tuning, guided by the noise scale, is predicted to improve efficiency.