PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:1907.04164v2cs.LGstat.ML

TL;DR

The paper asks how optimization-algorithm properties change the critical batch size, addressing limits of costly empirical studies and restrictive analyses. It combines large-scale neural-network experiments with a simple noisy quadratic model, finding that preconditioning enables larger critical batch sizes and that the model reproduces many observed training effects.

  • Problem

    Existing empirical scaling studies are expensive, while analytical critical-batch-size results often rely on strong assumptions, are difficult to apply, or focus on plain SGD.

  • Method

    The paper combines large-scale experiments on neural networks with analysis of a noisy quadratic model under convex-quadratic, gradient-noise, and online-training assumptions.

  • Results

    Preconditioning methods such as Adam and K-FAC produce much larger critical batch sizes than momentum SGD, while the NQM captures many essential features of real neural-network training.

  • Takeaways & Limitations

    The NQM provides a fast tool for formulating testable predictions about neural-network optimization and batch-size scaling.

  • Takeaways & Limitations

    The NQM assumes a convex quadratic loss, codiagonalizable noise covariance and Hessian, and online training, rather than modeling overfitting effects.

Abstract

from arXiv · show

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.

1 Introduction

The paper seeks a practical middle ground between expensive empirical scaling studies and restrictive theoretical analyses of batch-size effects. It develops an easy-to-run noisy quadratic model and investigates optimizers beyond momentum SGD, including preconditioned methods.

  • Existing empirical scaling curves are costly to generate, even for one optimizer on one task.
  • Prior analytical critical-batch-size studies rely on strong assumptions, yield difficult-to-apply parameter-dependent results, or focus on plain SGD.
  • The noisy quadratic model closely matches batch-size effects in real neural networks while enabling experiments in seconds and testable predictions.
  • The model predicts that momentum helps at larger but not small batch sizes, while preconditioning extends perfect scaling and helps at small batch sizes.
  • Exponential moving averages can reduce the steps needed for a batch size and achieve equivalent acceleration with smaller batches.

2 Related Work

Related work established batch-size effects, critical-batch-size analyses, momentum’s large-batch advantage, and benefits of averaging, but left important optimizer and mini-batch questions open.

  • Prior experiments found that critical batch size depends on both the dataset and network architecture.
  • Analytical studies derived critical-batch-size notions under varying assumptions, but, to the authors’ knowledge, all focused on SGD.
  • SGD and momentum SGD are equivalent for small learning rates after appropriate rescaling, yet this does not explain momentum’s occasional mini-batch advantage.
  • Smith et al. showed that momentum outperforms plain SGD at large batch sizes.
  • Prior averaging work mainly studied tail averaging, whereas exponential moving averages are favored in deep learning for simplicity and non-convex settings.
  • No prior analysis specifically examined exponential moving averages when mini-batches are used.

3 Analysis of the Noisy Quadratic Model (NQM)

The NQM models stochastic optimization as noisy convex-quadratic dynamics whose analytically tractable behavior explains how batch size interacts with momentum, preconditioning, and averaging. It predicts small-batch linear scaling, large-batch diminishing returns, and optimizer-specific changes in the critical batch size.

  • Model setup: The NQM assumes a convex quadratic loss, noisy gradients, codiagonalizable noise covariance and Hessian, and online i.i.d. observations.Mini-batching averages independent gradient queries, reducing covariance from C to C/B.
  • Model setup: Each coordinate converges exponentially to a steady-state risk, creating a trade-off between faster convergence from larger learning rates and higher residual risk.Increasing batch size lowers steady-state risk proportionally.
  • Momentum: Momentum matches plain SGD at small batch sizes but increases the critical batch size and accelerates training in the large-batch regime.Momentum amplifies steady-state risk by roughly 1/(1−β), while batch size reduces that risk, explaining its larger-batch advantage.
  • Preconditioning: Preconditioning improves convergence by reducing the transformed condition number, with stronger preconditioners offering larger benefits when near-deterministic dynamics dominate.For preconditioning power p, the convergence-rate improvement scales as κ^p; larger batches make steady-state risk less limiting.
  • Exponential moving average: Exponential moving average reduces steady-state risk without sacrificing convergence rate when its averaging coefficient is chosen appropriately.Its benefit is strongest in stochastic settings and becomes redundant near the deterministic large-batch regime.
  • Batch-size predictions: The NQM reproduces two batch-size regimes: perfect linear scaling for small batches and batch-insensitive performance after the critical batch size.Its predictions agree with empirical optimizer behavior, including momentum and preconditioning advantages at larger batches.

4 Neural Network Experiments

Experiments across neural networks and tasks show that batch-size scaling depends strongly on the optimizer, with preconditioning producing the largest critical batch sizes. The observed learning-rate behavior also follows the noisy quadratic model’s predictions.

  • Experimental setup: Experiments compare SGD, momentum SGD, Adam, and K-FAC across five architectures, three image-classification tasks, and one language-modeling task.Learning rates, schedules, and optimizer-specific metaparameters were tuned independently for each model, task, optimizer, and batch size.
  • Measurement: The experiments include steps to validation-set accuracy or cross entropy, with divergent trials discarded when the learning rate was too high.The target metric and computational budget were selected using initial experiments.
  • Critical batch size: K-FAC has the largest critical batch size in all cases, while momentum SGD extends perfect scaling beyond plain SGD.Below plain SGD’s critical batch size, momentum SGD requires as many steps as plain SGD; Adam and K-FAC can help even at the smallest batch sizes, though not universally.
  • Critical batch size: As batch size increases, steps initially scale perfectly, but larger batches yield rapidly diminishing returns after a problem-dependent critical batch size.Perfect scaling means doubling the batch size halves the steps to target.
  • Optimal learning rate: The optimal learning rate initially scales linearly with batch size, then plain SGD saturates while momentum SGD’s optimal effective learning rate continues increasing.The two methods appear to deviate at the same batch size in both optimal-learning-rate and steps-to-target plots.

5 Conclusion

The paper combines neural-network experiments with a noisy quadratic model and develops analytical treatments of preconditioning-related optimization dynamics. Its results connect optimizer properties to critical batch sizes, convergence, and computational implementation.

  • Conclusion: The noisy quadratic model agrees with neural-network phenomena including learning-rate scaling, critical batch sizes, momentum, preconditioning, and averaging.The model is presented as a simple framework for analyzing interactions between batch size and optimization algorithms.
  • K-FAC: K-FAC approximates the curvature matrix with Kronecker factors, avoiding exact-Fisher quadratic storage costs and enabling tractable approximate natural-gradient computation.For a layer, the factors are A = E[aa⊤] and S = E[{∇sL}{∇sL}⊤].
  • K-FAC: K-FAC computes natural-gradient updates using matrix transformations comparable in size to the weight matrix, making the method efficient.The paper also extends K-FAC to convolutional layers under spatial homogeneity and spatially uncorrelated-derivative assumptions.
  • K-FAC for Transformer: For Transformers, K-FAC handles the shared embedding and pre-softmax weight matrix by swapping the positions of its two Kronecker factors.The shared matrix is transposed at the pre-softmax layer, changing how input and output factors receive contributions.
  • Momentum analysis: The momentum analysis tracks parameter expectations, variances, and their covariance through a transition matrix whose characteristic polynomial determines convergence.The analysis treats each dimension independently and incorporates the correlation between parameters and momentum variables.
  • Momentum analysis: At the fastest-convergence setting, all transition-matrix eigenvalues equal β.This condition is identified from the transition-matrix eigenvalue analysis.

C Proof of Theorem 1

The proof derives closed-form dynamics for the stochastic and deterministic components of a linear system associated with the analyzed optimizer. A second-order recurrence and Z-transform provide the analytical solution.

  • Linear-system solution: The proof represents the relevant dynamics as a linear system and analyzes a single term of its infinite sum.A zero vector is appended for convenience when forming the system representation.
  • Recurrence reduction: Merging updates eliminates the variable that does not determine the loss, yielding a second-order difference equation.The proof applies this reduction to the component controlling the loss.
  • Z-transform: The second-order recurrence is solved with a Z-transform, producing a Z-domain function whose roots are r1 and r2.The roots satisfy z2 −(1 −αh + β)z + β = 0.
  • Z-transform: The proof uses r1r2 = β and r1 + r2 = 1 −αh + β to simplify the analytical expression.These root identities are substituted into the inverse-Z-transform solution.
  • Deterministic term: The deterministic term follows the same second-order recurrence as the stochastic term, but with different initial conditions.The proof then combines this case with the preceding recurrence analysis.

D Proof of Theorem 2

This proof analyzes moving-average dynamics by tracking parameter expectations, variances, and covariance in a linear system. It reduces the loss-relevant component to a second-order recurrence and solves it using characteristic roots.

  • Moving-average dynamics: The moving-average analysis tracks the parameter and average together, including their correlation, before deriving variance and covariance dynamics.The expected risk is expressed through the squared expectation plus variance of the averaged parameter.
  • Moving-average dynamics: The covariance recurrence includes contributions from parameter variance, parameter-average covariance, and gradient-noise variance.The update is Cov(t + 1) = (1 −γ)(1 −αh)2V[θ(t)] + (1 −αγ)Cov(t) + (1 −γ)α2c.
  • Linear-system formulation: The resulting quantities are organized as a linear dynamical system with a transition matrix and an infinite-sum representation.The proof introduces notation to simplify the closed-form derivation.
  • Closed-form solution: Eliminating the variable irrelevant to the loss produces a second-order difference equation, which is solved using its characteristic roots.The roots are r1 = 1 −αh and r2 = γ.
  • Closed-form solution: The proof obtains an explicit expression for yp involving α, c, h, γ, and the initial-condition factor B.The displayed expression is derived after substituting the two roots into the inverse-transform solution.
  • Bound: An upper bound is then stated for the remaining contribution to the solution.This bound follows the analysis of the second term involving the infinite sum.

E More results on the NQM

These analyses compare neural-network curvature and gradient statistics with assumptions used by the noisy quadratic model, finding structured Fisher spectra and support for the model’s alignment assumptions.

  • Loss-surface analysis: The study uses the K-FAC approximate Fisher eigenspectrum as a tractable proxy for the true Fisher and, in turn, the Hessian.The true Fisher is treated as equivalent to the generalized Gauss–Newton Hessian.
  • Loss-surface analysis: The approximate Fisher spectrum has only a few large and small eigenvalues, while most eigenvalues lie in the middle range.After 200 training iterations, the eigenspectrum remains mostly unchanged.
  • Assumption checks: Gradient variance and curvature are compared in a Kronecker-factored eigenbasis to test whether the Hessian and noise covariance are codiagonalizable.This basis is used because obtaining all true-Fisher eigenvectors is typically infeasible.

E.3 Plots for the Evolution of the First Term in Eqn. (6)

The NQM comparison shows that momentum SGD and plain SGD with an adjusted learning rate have closely matching convergence dynamics when their effective learning rates are aligned.

  • Convergence comparison: Figure 11 compares the convergence term as a function of (1 −β), the scaling between momentum SGD’s effective and true learning rates.Red curves represent momentum and blue curves represent plain SGD using the effective learning rate.
  • Convergence comparison: The convergence dynamics of momentum SGD and plain SGD with the effective learning rate closely match in the NQM.The comparison plots the convergence-determining first term of the total risk, while steady-state risk remains unchanged.

E.4 Verification of Eigenspectrum

The eigenspectrum verification compares the NQM’s assumed Hessian spectrum with K-FAC spectra from a real neural network and reports close agreement.

  • Eigenspectrum comparison: The NQM eigenvalue distribution closely matches the K-FAC eigenspectra of ResNet8, supporting the model’s Hessian-spectrum assumption.The comparison is presented in Figure 12.
  • Batch-size experiment: The experiment examines how the number of training steps needed to reach a target loss varies with batch size and loss threshold.For each batch size, minimal steps are obtained by grid-searching over constant learning rates.

E.6 Results of Optimal Learning Rate on NQM

These NQM experiments examine how optimal learning rates and training steps vary with batch size, including fixed and piecewise-constant learning-rate schemes.

  • Training steps: The number of training steps required to reach a target loss is evaluated across batch sizes and multiple loss thresholds.The analysis uses constant learning rates selected by grid search for each batch size.
  • Optimal learning rate: For momentum-free training, the optimal learning rate increases with batch size until reaching the maximum stable learning rate.Larger preconditioning powers lower the optimal learning rate at a given batch size and extend the leveling-off range.
  • Piecewise-constant learning rate: The final learning rate in the piecewise-constant scheme also scales linearly with batch size.This extends the observed linear scaling from fixed learning rates to the final rate under learning-rate decay.

F.1 Data Sets

The experiments use datasets and model architectures largely aligned with prior work, with several stated modifications. The section also specifies learning-rate schedules and optimizer-specific hyperparameter tuning.

  • F.1 Data Sets: Figure 15 reports the final learning rate of the piecewise-constant learning-rate scheme versus batch size.The passage identifies the plotted quantities but does not state the figure’s outcome.
  • F.1 Data Sets: The experiments use MNIST, Fashion MNIST, CIFAR10, ImageNet, and LM1B, with data augmentation applied to CIFAR10.These datasets match Shallue et al. [2018], although CIFAR10 augmentation differed between the studies.
  • F.1 Data Sets: The evaluated models include Simple CNN, ResNet8, ResNet32, VGG11, and a shallow Transformer.The architectures are described as similar or identical to models used in prior work, with modifications such as omitted or ghost batch normalization and altered Transformer depth.
  • F.1 Data Sets: The study compares constant and linear-decay learning-rate schedules.The linear-decay schedule tunes the initial learning rate α0, decay rate γ, and step count T to reach the final learning rate.
  • F.1 Data Sets: Optimizer tuning covers momentum for SGD, β1, β2, and ϵ for Adam, and damping, trust-region, momentum, and damping-average parameters for K-FAC.For K-FAC on Transformer, momentum and the damping moving-average parameter were fixed while damping and the trust-region constraint were tuned.
Loading 1907.04164v2…