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

TL;DR

The work analyzes how batch size interacts with optimization algorithms by combining a tractable Noisy Quadratic Model with extensive neural network experiments. It demonstrates that momentum helps primarily in large-batch regimes, while preconditioning (e.g., Adam, K-FAC) extends perfect scaling to even larger batch sizes, albeit with trade-offs in steady-state risk. Exponential moving averages (EMA) reduce steady-state error and can lower the required batch size for the same speed, with EMA benefits diminishing at very large batches. Across multiple datasets and architectures, the results align with NQM predictions, offering a fast, testable framework for predicting optimizer performance under varying 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.

Figures (15)

• Figure 1: Cartoon of the evolution of risk for different coordinates with and without learning rate decay.
• Figure 2: Convergence rate and steady state risk (SSK) as a function of momentum for a single dimension with $\alpha h = 0.0005$ and batch size $B = 1$.
• Figure 3: (a) Effects of momentum and preconditioning. Steps required to reach target loss as a function of batch size under different preconditioning power. Solid lines are momentum SGD while dashed lines are plain SGD. The black dashed line is the information theoretic lower bound. (b) Effect of learning rate decay. The solid lines use the optimized piecewise constant scheme, which are shown in (c) for power $0$. The dashed curves in (b) are plain SGD for comparison. We observe that learning rate schedules close most of the gap between the fixed learning rate performance and the information theoretic lower bound.
• Figure 4: Effects of exponential moving average (EMA). Solid lines are SGD with EMA while dashed lines are plain SGD.
• Figure 5: Empirical relationship between batch size and steps to result. Key observations: 1) momentum SGD has no benefit over plain SGD at small batch sizes, but extends the perfect scaling to larger batch sizes; 2) preconditioning also extends perfect scaling to larger batch sizes, i.e. K-FAC > Adam > momentum SGD. This is most noticeable in the Transformer model; 3) preconditioning (particularly K-FAC) reduces the number of steps needed to reach the target even for small batch sizes. All of these agree with the predictions by NQM.

Theorems & Definitions (2)

• Theorem 1
• Theorem 2