An Empirical Model of Large-Batch Training

By Sam McCandlish, Jared Kaplan, Jared Kaplan, Dario Amodei, OpenAI Dota Team. 2018

416 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Computer Networks and Communications, Optimization & Training

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Networks and Communications

In Optimization & Training

Citations in the corpus

Language Models are Few-Shot Learners (2020, methodology): “We measure the gradient noise scale during training and use it to guide our choice of batch size [MKAT18].”
Scaling Laws for Neural Language Models (2020, methodology): “B crit is measured empirically from the data shown in Figure 18, but it is also roughly predicted by the gradient noise scale, as in [MKAT18]. and studying it as a function of N, D .”
PaLM: Scaling Language Modeling with Pathways (2022, background): “, better loss as a function of tokens seen) earlier in training, while larger batch sizes are beneficial later in training due to better gradient estimates (Smith et al., 2018; McCandlish et al., 2018), and (2) larger batch sizes result in larger matrix multiplication dimensions, which increases TPU”
Which Algorithmic Choices Matter at Which Batch Sizes? Insights From a Noisy Quadratic Model (2019, background): “…an initial period of perfect scaling where doubling the batch size halves the steps to target, but once the batch size exceeds a problem-dependent critical batch size, there are rapidly diminishing returns, matching the results of [Goyal et al., 2017, McCandlish et al., 2018, Shallue et al., 2018].”
Unified Scaling Laws for Routed Language Models (2022, methodology): “1.6) a unified prediction of the loss achieved by a model with size N training for S steps: L(N,S) = ( Nc N )αN + ( Sc S )αS (28) 3This sidesteps the issue of critical batch size [McCandlish et al., 2018, Kaplan et al., 2020], consideration of which requires a substantially larger sweep of models.”
GLaM: Efficient Scaling of Language Models with Mixture-of-Experts (2021, methodology): “To improve training stability, empirical heuristics and hyperparameter tuning are usually required (McCandlish et al., 2018; Kaplan et al., 2020).”
ZeRO: Memory optimizations Toward Training Trillion Parameter Models (2019, background): “1Prior work [9] shows, very large batch size could slow down convergence.”
GPT-NeoX-20B: An Open-Source Autoregressive Language Model (2022, background): “The critical batch size [McCandlish et al., 2018] for the model may have been lower, possibly resulting in a negative impact on training.”
Improved Denoising Diffusion Probabilistic Models (2021, methodology): “We conﬁrmed this by evaluating the gradient noise scales (McCandlish et al., 2018) for models trained with both objectives, as shown in Figure 7.”
ExT5: Towards Extreme Multi-Task Scaling for Transfer Learning (2021, background): “Another explanation as to why this happens is that large-batch training can offer benefits even for single-task models (Smith et al., 2018) — a trend formalized by McCandlish et al. (2018).”
Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling (2023, background): “It is widely held (McCandlish et al., 2018; Zhang et al., 2019; Kaplan et al., 2020; Brown et al., 2020; Hoffmann et al., 2022) that using larger batch sizes is desirable, but that smaller LLMs require smaller batch sizes to avoid convergence issues.”
Measuring the Algorithmic Efficiency of Neural Networks (2020, background): “Such efﬁciency gains are not generally observed empirically, though they can be calculated with asymptotic analysis in some cases and estimated with empirical scaling laws in others [McCandlish et al., 2018].”
Don't Use Large Mini-Batches, Use Local SGD (2018, background): “…> 10 3 on ImageNet) typically degrades the performance both in terms of training and test error (often denoted as generalization gap ) (Chen & Huo, 2016; Li, 2017; Li et al., 2014; Keskar et al., 2017; Shallue et al., 2018; McCandlish et al., 2018; Golmant et al., 2018; Masters & Luschi, 2018).”
Dota 2 with Large Scale Deep Reinforcement Learning (2019, methodology): “Per results in [28], we hoped to find (in the early part of training) linear speedup from increasing batch size; i.”
Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization (2019, neutral): “However, model parallelism requires access to additional expensive compute accelerators, fast networks, and non-trivial partitioning of model state to balance communication and computation [18, 31, 39].”