Compute Efficiency and Serial Runtime Tradeoffs for Stochastic Momentum Methods

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.

• The paper shows that Heavy Ball momentum preserves SGD-level compute efficiency while enabling larger batch sizes to reduce serial runtime.
• It derives explicit lower bounds and identifies critical batch size transitions for both HB and ASGD under different data covariance spectra.
• The findings guide optimizer selection by clarifying the tradeoffs between runtime reduction and compute efficiency in large-scale machine learning.

Compute Efficiency and Serial Runtime Tradeoffs for Stochastic Momentum Methods

Motivation and Problem Formulation

The paper presents a rigorous study of stochastic momentum methods—specifically, Heavy Ball (HB) momentum and variants of Accelerated Stochastic Gradient Descent (ASGD)—in terms of their compute efficiency (CE) and serial runtime as a function of problem instance (covariate spectrum) and batch size in consistent linear regression. The study is motivated by practical concerns in large-scale machine learning, where not only sample complexity but also the translation of larger batch sizes into actual reductions in wall-clock training time are critical factors in algorithm selection.

Momentum methods are essential in modern optimizers, either directly or via equivalents in adaptive schemes, but their benefits in stochastic regimes have remained inadequately characterized, especially concerning how they trade off between compute efficiency (total gradient queries/FLOPs to reach a target accuracy) and serial training time (number of sequential steps). This tradeoff is especially subtle when optimization is constrained by noisy gradients and high-dimensional spectra typical of large-scale learning, such as LLM pretraining.

Theoretical Results: Lower Bounds on Compute Efficiency

Heavy Ball (HB) Momentum

For arbitrary data covariance spectra, the paper proves that HB does not improve the best-achievable compute efficiency over SGD at any batch size; the CE frontier of HB matches that of SGD. However, HB's key benefit is to preserve this optimal CE over a larger window of batch sizes—up to batch sizes on the order of $$B_{\text{HB}} \sim \sqrt{\kappa} \, \text{Tr}(H)/\lambda_{\max}$$, where $\kappa$ is the condition number of the covariance. This extension allows practitioners to employ larger batches—hence reducing serial runtime—without incurring additional compute costs, all the way until deterministic acceleration saturates and no further improvement in serial runtime is possible without sacrificing efficiency.

For power-law spectra $\lambda_i \sim i^{-a}$, the HB-accelerated regime transition occurs at $B_{\text{HB}} \sim d^{a/2}$, and beyond this, the serial runtime reduction saturates at the optimal accelerated complexity.

Accelerated Stochastic Gradient Descent (ASGD)

ASGD, parameterically broader due to the presence of an additional buffer, achieves strictly better small-batch CE than SGD or HB in the regime of rapidly decaying spectra. In contrast to HB, the CE-preserving linear scaling window of ASGD is shorter: as batch size increases, the marginal serial runtime gain comes at the expense of spending ASGD's initial CE advantage. For spectra with power-law decay, ASGD saturates its accelerated scaling at smaller batch sizes—specifically, for $\lambda_i \sim i^{-a}$ the CE-optimal batch size is $B_{\text{ASGD}} \sim d^{1/2}$, which is always less than $B_{\text{HB}}$ for $a > 1$.

In summary:

• HB: Enlarges the batch size window preserving SGD-level CE, allowing runtime reduction up to the accelerated regime.
• ASGD: Outperforms HB/SGD CE at small batch sizes, but loses this advantage at large batches due to an earlier saturation of the scaling window.

Empirical results validate these theoretical distinctions with synthetic regression experiments, demonstrating clear qualitative and quantitative separations between the regimes of operation as the spectrum decay and batch size vary.

Implications for Algorithm Selection and Large-Scale Learning

The results have major implications for optimizer selection in large-batch training, particularly in the context of deep learning and LLM pretraining:

• Momentum is only beneficial for serial runtime in large-batch regimes; at small batch sizes, it does not improve overall compute efficiency.
• ASGD-type methods should be preferred if maximizing compute efficiency at small batch sizes is desired, especially when the spectral decay is rapid (fatter tails).
• The batch size should be increased up to the method-specific critical batch size (CBS) to minimize wall-clock time without crossing into regimes where compute is wasted.
• For practical large-scale pretraining (with very large token counts), batch sizes typically exceed those for which ASGD’s CE advantage is realized. In these cases, HB or standard momentum preserves SGD-level efficiency but allows runtime reduction at batch sizes inaccessible to plain SGD.

These findings inform batch scaling strategies and the design of optimizer hyper-parameterization for new architectures and tasks, highlighting the importance of problem spectrum in effective optimizer selection.

Technical Approach

The analysis hinges on a careful spectral treatment of the stochastic covariance recursion for the error, focusing on augmented states that render the dynamics first-order linear. For both HB and ASGD, the contraction rates (spectral gaps) as a function of batch size and spectrum are sharply characterized with lower bounds that hold for any feasible choice of algorithm parameters, not relying on pathologically constructed cases.

These results are substantially stronger and more general than previous lower bounds, subsuming high-dimensional asymptotics and concretely identifying the batch-sizing transitions for acceleration in finite dimension.

Limitations and Future Directions

The presented lower bounds are non-asymptotic and explicit for power-law spectra, but rely on the linear regression/Gaussian covariate setting. Extension to nonlinear models, or to general function classes with non-quadratic loss, is a major direction for future work. Additionally, the methods may be leveraged to obtain sharp upper bounds, yielding tight characterizations of optimal stepsize and momentum parameter ranges.

Moreover, the interplay between optimizer-induced stability, deviation from edge-of-stability, and the emergence of scaling laws in practical neural networks (outside linear regimes) is left for ongoing investigation.

Conclusion

The paper systematically delineates the regimes and tradeoffs for stochastic momentum methods, conclusively demonstrating that the principal benefit of momentum in the stochastic regime is an expanded batch size window for runtime reduction, not increased compute efficiency. ASGD offers improved compute efficiency at small batch sizes for rapidly decaying spectra but only up to an earlier saturation point. These theoretical results provide essential guidance for the principled deployment and tuning of optimizers in large-scale machine learning systems.

Reference:

"Compute Efficiency and Serial Runtime Tradeoffs for Stochastic Momentum Methods" (2606.19179)