Year: 2024
Type: article
Abstract: Simple stochastic momentum methods are widely used in machine learning optimization, but their good practical performance is at odds with an absence of theoretical guarantees of acceleration in the literature. In this work, we aim to close the gap between theory and practice by showing that stochastic heavy ball momentum retains the fast linear rate of (deterministic) heavy ball momentum on quadratic optimization problems, at least when minibatching with a sufficiently large batch size. The algorithm we study can be interpreted as an accelerated randomized Kaczmarz algorithm with minibatching and heavy ball momentum. The analysis relies on carefully decomposing the momentum transition matrix, and using new spectral norm concentration bounds for products of independent random matrices. We provide numerical illustrations demonstrating that our bounds are reasonably sharp.
Authors: Raghu Bollapragada, Tyler Chen, Rachel Ward
Institutions: The University of Texas at Austin, New York University
Language: English
Subfield: Artificial Intelligence
Field: Computer Science
Domain: Physical Sciences
Open Access status: bronze
Awards: G1523888516, G3474365249, Graduate Research Fellowship Program (GRFP), FRG: Collaborative Research: Randomized Algorithms for Solving Linear Systems, G6562578720, Co-Design of Neural Operators and Stochastic Optimization Algorithms for Learning Surrogates for PDE-Constrained Optimization Under Uncertainty, AI Institute: Institute for Foundations of Machine Learning, HDR TRIPODS: UT Austin Institute on the Foundations of Data Science
OpenAlex ID: W4401445256