Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Stochastic Polyak Step-sizes and Momentum: Convergence Guarantees and Practical Performance

Dimitris Oikonomou, Nicolas Loizou

TL;DR

This work tackles the practical challenge of tuning step-sizes and momentum in stochastic optimization with momentum (SHB) by introducing Polyak-type adaptive step-sizes via an Iterate Moving Average (IMA) view. It proposes three SHB variants—$\text{MomSPS}_{\max}$, $\text{MomDecSPS}$, and $\text{MomAdaSPS}$—that adaptively set the learning rate $\gamma_t$ in concert with momentum $\beta$, achieving convergence to a neighborhood without interpolation and, in two variants, convergence to the exact minimizer without prior problem-parameter knowledge. The analysis shows that naive SPS with momentum can be unstable, but the IMA-based formulations yield tight, generation-spanning guarantees that reduce to known SPS results for $\beta=0$ and extend them to SHB; corollaries include robust constant-step behavior and a projection-based extension. Empirically, the proposed methods demonstrate robust performance across convex and nonconvex tasks, including deep networks, with open-source code provided, underscoring practical impact for scalable training where hyperparameter tuning is costly.

Abstract

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.

Stochastic Polyak Step-sizes and Momentum: Convergence Guarantees and Practical Performance

TL;DR

This work tackles the practical challenge of tuning step-sizes and momentum in stochastic optimization with momentum (SHB) by introducing Polyak-type adaptive step-sizes via an Iterate Moving Average (IMA) view. It proposes three SHB variants—, , and —that adaptively set the learning rate in concert with momentum , achieving convergence to a neighborhood without interpolation and, in two variants, convergence to the exact minimizer without prior problem-parameter knowledge. The analysis shows that naive SPS with momentum can be unstable, but the IMA-based formulations yield tight, generation-spanning guarantees that reduce to known SPS results for and extend them to SHB; corollaries include robust constant-step behavior and a projection-based extension. Empirically, the proposed methods demonstrate robust performance across convex and nonconvex tasks, including deep networks, with open-source code provided, underscoring practical impact for scalable training where hyperparameter tuning is costly.

Abstract

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, MomDecSPS, and MomAdaSPS. For MomSPS, 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, 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.
Paper Structure (37 sections, 22 theorems, 82 equations, 22 figures, 3 tables, 3 algorithms)

This paper contains 37 sections, 22 theorems, 82 equations, 22 figures, 3 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Main Contributions
  3. Exploring the Interplay of SPS and Heavy Ball Momentum
  4. Background on Stochastic Polyak Step-size
  5. Naive SPS in SHB
  6. Iterative Moving Average: Balance between SPS and Momentum
  7. Convergence Analysis
  8. Convergence to a neighborhood of the solution
  9. Convergence to the exact solution
  10. Numerical Experiments
  11. Further Related Work
  12. Adaptive methods.
  13. Momentum with constant step-size.
  14. Adaptive Momentum.
  15. On Technical Assumptions for Convergence.
  16. ...and 22 more sections

Key Result

Proposition 2.1

If $\lambda_t=\lambda\geq0$ then assuming that $\beta=\frac{\lambda}{1+\lambda}$ and $\gamma_t=(1-\beta)\eta_t$ the $x_t$ iterates of the IMA method are equal to the $x_t$ iterates produced by the eq:shb method.

Figures (22)

  • Figure 1: Comparison of \ref{['eq:mospsmax']} versus SPS$_{\max}$ with naive momentum for different momentum parameters $\beta$ on a logistic regression problem. Left: $\beta=0.2$, Right: $\beta=0.5$
  • Figure 2: Comparison of deterministic algorithms for least squares. Left: Relative error, Right: Step-sizes.
  • Figure 3: LibSVM dataset: vowel, Batch size: 52
  • Figure 4: LibSVM dataset: letter, Batch size: 1500
  • Figure 5: Resnet 34 on CIFAR 10
  • ...and 17 more figures

Theorems & Definitions (36)

  • Proposition 2.1
  • Corollary 2.2
  • Remark 2.3
  • Remark 2.4
  • Theorem 3.2
  • Corollary 3.3: Deterministic Heavy Ball Method
  • Corollary 3.4: Interpolation
  • Corollary 3.5: Constant Step-size
  • Theorem 3.6
  • Theorem 3.7
  • ...and 26 more