Stochastic Polyak Step-sizes and Momentum: Convergence Guarantees and Practical Performance
This work addresses the time-consuming tuning process in stochastic optimization for machine learning practitioners, offering adaptive methods that improve robustness and performance.
The paper tackles the problem of tuning step-size and momentum parameters in Stochastic Heavy Ball (SHB) methods for large-scale stochastic optimization by proposing three new Polyak-type step-size variants (MomSPS_max, MomDecSPS, MomAdaSPS). It provides convergence guarantees for convex and smooth problems, with MomSPS_max achieving fast rates matching deterministic HB under interpolation, and MomDecSPS and MomAdaSPS enabling convergence to the exact minimizer without prior knowledge or interpolation assumptions.
Stochastic gradient descent with momentum, also known as Stochastic Heavy Ball method (SHB), is one of the most popular algorithms for solving large-scale stochastic optimization problems in various machine learning tasks. In practical scenarios, tuning the step-size and momentum parameters of the method is a prohibitively expensive and time-consuming process. In this work, inspired by the recent advantages of stochastic Polyak step-size in the performance of stochastic gradient descent (SGD), we propose and explore new Polyak-type variants suitable for the update rule of the SHB method. In particular, using the Iterate Moving Average (IMA) viewpoint of SHB, we propose and analyze three novel step-size selections: MomSPS$_{\max}$, MomDecSPS, and MomAdaSPS. For MomSPS$_{\max}$, we provide convergence guarantees for SHB to a neighborhood of the solution for convex and smooth problems (without assuming interpolation). If interpolation is also satisfied, then using MomSPS$_{\max}$, SHB converges to the true solution at a fast rate matching the deterministic HB. The other two variants, MomDecSPS and MomAdaSPS, are the first adaptive step-size for SHB that guarantee convergence to the exact minimizer - without a priori knowledge of the problem parameters and without assuming interpolation. Our convergence analysis of SHB is tight and obtains the convergence guarantees of stochastic Polyak step-size for SGD as a special case. We supplement our analysis with experiments validating our theory and demonstrating the effectiveness and robustness of our algorithms.