Which Algorithmic Choices Matter at Which Batch Sizes? Insights From a Noisy Quadratic Model

By Guodong Zhang, Lala Li, Zachary Nado, James Martens, Sushant Sachdeva, George E. Dahl, Christopher J. Shallue, Roger Baker Grosse. 2019

192 citations recorded by the source corpus.

Increasing the batch size is a popular way to speed up neural network training, but beyond some critical batch size, larger batch sizes yield diminishing returns. In this work, we study how the critical batch size changes based on properties of the optimization algorithm, including acceleration and preconditioning, through two different lenses: large scale experiments, and analysis of a simple noisy quadratic model (NQM). We experimentally demonstrate that optimization algorithms that employ preconditioning, specifically Adam and K-FAC, result in much larger critical batch sizes than stochastic gradient descent with momentum. We also demonstrate that the NQM captures many of the essential features of real neural network training, despite being drastically simpler to work with. The NQM predicts our results with preconditioned optimizers, previous results with accelerated gradient descent, and other results around optimal learning rates and large batch training, making it a useful tool to generate testable predictions about neural network optimization.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, 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 Vision and Pattern Recognition

In Optimization & Training

Citations in the corpus

Training Compute-Optimal Large Language Models (2022, neutral): “Estimating hyperparameters for large models.”
On Empirical Comparisons of Optimizers for Deep Learning (2019, background): “…work (Shallue et al., 2019; Zhang et al., 2019) has shown that increasing the batch size can increase the gaps between training times for different optimizers, with the gap from SGD to M OMENTUM (Shallue et al., 2019) and from M OMENTUM to A DAM (Zhang et al., 2019) increasing with the batch size.”
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.”
Learning explanations that are hard to vary (2020, background): “Looking at agreement to characterize and improve the representations learned by neural networks is an idea that has been explored in recent work, mostly looking at the cosine similarity between gradients [Du et al., 2018, Eshratifar et al., 2018, Fort et al., 2019, Zhang et al., 2019b].”
Understanding the Role of Training Regimes in Continual Learning (2020, background): “…is further studied by Jastrzebski et al. [32], showing that using a large learning rate in the initial phase of training reduces the variance of the gradient, and improves the conditioning of the covariance of gradients which itself is close to the Hessian in terms of the largest eigenvalues [77].”
Lookahead Optimizer: k steps forward, 1 step back (2019, neutral): “While simple, this model is a proxy for neural network optimization and effectively optimizing it remains a challenging open problem [37, 26, 43, 47].”
Batch Normalization Biases Residual Blocks Towards the Identity Function in Deep Networks (2020, background): “On the other hand, when the batch size is large, the optimal learning rates are independent of batch size [26, 27], and normalized networks use larger learning rates.”
Drawing Multiple Augmentation Samples Per Image During Training Efficiently Decreases Test Error (2021, background): “Since SGD exhibits different behaviour at different batch sizes [Goyal et al., 2017, Ma et al., 2018, Zhang et al., 2019, Smith et al., 2020], we also study augmentation multiplicity in both the small and large batch regimes.”
Pipelined Backpropagation at Scale: Training Large Models without Batches (2020, methodology): “Alternatively initialization methods can be used to enable training without normalization (Zhang et al., 2019; Dauphin & Schoenholz, 2019).”
The Break-Even Point on Optimization Trajectories of Deep Neural Networks (2020, background): “Zhang et al. (2019) argues that in DNNs the Hessian and the covariance of gradients are close in terms of the largest eigenvalues.”
On the Generalization Benefit of Noise in Stochastic Gradient Descent (2020, unknown): “…SGD with and without Momentum achieve similar performance in the noise dominated regime if one maintains a constant effective learning rate (Smith et al., 2019; Zhang et al., 2019), but SGD with Momentum performs better in the curvature dominated regime (Shallue et al., 2018). ar X iv :2 00 6.”
Neural networks with late-phase weights (2020, methodology): “Despite its apparent simplicity, it remains a challenging optimization problem that has important similarities with stochastic minibatch training in deep neural networks (Schaul et al., 2013; Martens, 2016; Wu et al., 2018; Zhang et al., 2019a;b).”
An Empirical Study of Large-Batch Stochastic Gradient Descent with Structured Covariance Noise (2019, background): “The recent work of Zhang et al. (2019) shows that this relationship indeed holds tightly in empirical settings.”
DeepSeek LLM: Scaling Open-Source Language Models with Longtermism (2024, result): “Early works (Goyal et al., 2017; McCandlish et al., 2018; Shallue et al., 2019; Smith et al., 2017; Zhang et al., 2019) provided some empirical observations for setting batch size and learning rate, but we found these observations have limited applicability in our preliminary experiments.”
A Generalizable Approach to Learning Optimizers (2021, methodology): “To do so, we train additional LHOPTs on only MNIST-like datasets (OM) and on only NQMs [54].”