Compute Efficiency and Serial Runtime Tradeoffs for Stochastic Momentum Methods

2026-06-17

AI Summary

This study investigates the trade-off between serial runtime and computational efficiency in stochastic momentum methods—specifically, stochastic heavy ball (SHB) and accelerated stochastic gradient descent (ASGD). Through finite-dimensional discrete-time theoretical analysis, spectral-dependent modeling, and synthetic linear regression experiments, the work demonstrates that SHB merely extends the effective batch size window of SGD without advancing the computational efficiency frontier. In contrast, ASGD enhances small-batch computational efficiency under rapidly decaying spectra and improves serial runtime at large batch sizes. Both theory and experiments consistently show that under slowly decaying spectra, the two methods perform similarly, whereas under rapidly decaying spectra, a clear trade-off emerges: the effective batch window of SHB can reach up to √κ times the critical batch size of SGD, where κ denotes the condition number of the problem.

Abstract

Stochastic momentum methods such as heavy ball (HB), Nesterov momentum, and variants of Accelerated SGD (ASGD) [Kidambi et al., 2018] are widely used in modern training, but their stochastic benefits depend on two distinct quantities: serial runtime, the number of iterations needed to reach a target accuracy, and compute efficiency (CE), the inverse total gradient-query or FLOP cost. Larger batches reduce serial runtime without hurting CE only when the contraction gap grows linearly with batch size. We study stochastic HB and ASGD for consistent linear regression with Gaussian covariates and prove finite-dimensional, discrete-time lower bounds on their batch-size tradeoffs. Our first result shows that HB does not improve the CE frontier over SGD for arbitrary spectra; rather, it preserves SGD-level CE over a larger batch-size window, allowing larger batches to reduce serial runtime until HB reaches its deterministic accelerated scale. This window can be a factor $\sqrtκ$ larger than the SGD critical batch size. For ASGD, the picture is more spectrum-dependent: for rapidly decaying power-law spectra, ASGD improves small-batch CE over HB/SGD, but as batch size grows it trades this CE advantage for improved serial runtime. Synthetic linear-regression experiments verify these qualitative regimes, including near-overlap of ASGD and HB for slowly decaying spectra and the predicted CE--serial tradeoff for rapidly decaying spectra.

Problem

• stochastic momentum
• compute efficiency
• serial runtime
• batch size
• linear regression

Innovation

• stochastic momentum
• compute efficiency
• serial runtime
• batch-size tradeoff
• linear regression