Perfect Parallelization in Mini-Batch SGD with Classical Momentum Acceleration

TL;DR

This paper develops a theoretical framework showing that classical momentum methods can be perfectly parallelized in mini-batch stochastic gradient descent, with acceleration proportional to batch size, applicable to deep learning models.

Contribution

It introduces a general theory for stochastic momentum acceleration in quadratic optimization, encompassing various momentum schemes and mini-batch sizes, with minimal noise assumptions.

Findings

Acceleration scales with mini-batch size up to saturation

Provides a simple, effective momentum parameter choice

Enables perfect parallelization of mini-batch computations

Abstract

Accelerating stochastic gradient methods with classical momentum schemes, such as Polyak's heavy ball, has proven highly successful in training large-scale machine learning models, particularly when combined with the hardware acceleration of large mini-batch computations. Yet, the effect of classical momentum on stochastic mini-batch optimization has been poorly understood theoretically, with prior works requiring strong noise assumptions and extremely large mini-batches. In this work, we develop a general theory of stochastic momentum acceleration for optimizing over quadratics in the interpolation regime, a popular abstraction for studying deep learning dynamics which also includes classical methods such as randomized Kaczmarz and coordinate descent. Our framework encompasses both heavy ball and Nesterov-style momentum, allows for arbitrary mini-batch sizes, and makes minimal assumptions on the stochastic noise. In particular, we show that acceleration from classical momentum is directly proportional to the gradient mini-batch size (up to a natural saturation point), thereby enabling perfect parallelization of mini-batch computations. Our theory also provides a simple choice for the momentum parameter, which is shown to be effective empirically.