[Paper Review] Which Algorithmic Choices Matter at Which Batch Sizes? Insights From a Noisy Quadratic Model

Guodong Zhang, Lala Li+6 more|arXiv (Cornell University)|Jul 9, 2019

Stochastic Gradient Optimization Techniques49 references15 citations

TL;DR

This paper introduces a noisy quadratic model (NQM) to study how algorithmic choices—such as preconditioning, momentum, and exponential moving averages—affect the critical batch size in neural network training. The NQM accurately predicts that preconditioned optimizers like Adam and K-FAC extend perfect batch size scaling to much larger batch sizes than SGD with momentum, while also improving performance at small batch sizes.

ABSTRACT

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.

Motivation & Objective

• To understand how different optimization algorithms affect the critical batch size in neural network training.
• To investigate whether preconditioning, momentum, and exponential moving averages extend the benefits of large-batch training beyond what is achievable with standard SGD.
• To validate the predictive power of a simple noisy quadratic model (NQM) for real-world neural network optimization dynamics.
• To provide a scalable, analytically tractable alternative to expensive empirical batch size studies.
• To generate testable predictions about optimal learning rate schedules and algorithmic hyperparameters in large-batch training.

Proposed method

• Develops a noisy quadratic model (NQM) that captures key aspects of neural network optimization, including gradient noise and curvature.
• Uses the NQM to analytically derive predictions about optimal learning rates, critical batch sizes, and the impact of algorithmic components.
• Employs large-scale experiments on multiple models (CNNs, ResNets, VGG, Transformers) and datasets (MNIST, CIFAR10, ImageNet, LM1B) to validate NQM predictions.
• Applies preconditioned optimizers (Adam, K-FAC) and compares their performance against SGD with momentum across varying batch sizes.
• Implements learning rate schedules (constant and linear decay) and analyzes their interaction with batch size and optimizer choice.
• Uses techniques like ghost batch normalization, label smoothing, and channel-wise weight normalization to stabilize training across batch sizes.

Experimental results

Research questions

• RQ1 How does the critical batch size depend on the choice of optimization algorithm, particularly when comparing SGD with momentum to preconditioned methods like Adam and K-FAC?
• RQ2 To what extent can the noisy quadratic model (NQM) predict real-world batch size scaling behavior in neural network training?
• RQ3 Does exponential moving average (EMA) reduce the number of training steps required for a given batch size, and can it enable equivalent performance with smaller batch sizes?
• RQ4 How do learning rate schedules interact with batch size and optimizer type to affect training efficiency?
• RQ5 What is the relative impact of preconditioning versus momentum in extending the range of effective large-batch training?

Key findings

• The noisy quadratic model (NQM) accurately predicts the batch size scaling behavior observed in real neural network training, including the effects of momentum and preconditioning.
• Preconditioned optimizers like Adam and K-FAC extend perfect batch size scaling to significantly larger batch sizes than SGD with momentum, with K-FAC achieving the largest critical batch size.
• Unlike momentum, preconditioning provides performance benefits even at small batch sizes, improving training efficiency across the entire batch size spectrum.
• Exponential moving averages reduce the number of training steps required for a given batch size, enabling equivalent performance with smaller batch sizes and saving computational resources.
• The optimal learning rate scales linearly with batch size for both constant and linear decay schedules, consistent with NQM predictions.
• The NQM successfully reproduces known phenomena such as the diminishing returns of large batches and the existence of a critical batch size, validating its use as a predictive tool for optimization.