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

Authors: Dimitris Oikonomou, Nicolas Loizou
Primary area: optimization
Keywords: Stochastic Polyak Step-size, SGD, Stochastic Heavy Ball, Convergence Analysis, Convex Optimization, Momentum

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: MomSPSmax, MomDecSPS, and MomAdaSPS. For MomSPSmax, 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 MomSPSmax, 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.

Decision: accept (poster)
Reviewer scores: 8, 6, 8, 6

Reviews and discussion:
<thread>
creation date:
1729651268567

thread_type:
official_review

<thread_object>
creator: reviewer_jvkn

<content>
summary:
This paper studies stochastic Polyak step-size for SHB. Three step size selections are proposed and associated convergence guarantees are developed. Specifically, using the MomSPSmax stepsize, the SHB is proved to converge to a neighborhood of the solution, if the objective function is convex and smooth. Moreover, another two stepsizes, I.e., MomDecSPS and MomAdaSPS, are proposed for non-interpolated SHB, which guarantee that SHB iterates converging to the exact solution. Finally, numerical experiments are presented to show the effectiveness of the proposed methods.

strengths:
1.The writing and exposition are clear.

2.The proposed methods are practical to implement and have theoretical improvements compared with existing works.

weaknesses:
Most of techniques are from existing literature. The authors did not explain the difficulty of applying the Polyak step-size to SHB compared with the SGD. This makes the technique used in this paper look like a direct generalization of Loizou et al. (2021)

questions:
None.
</content>
</thread_object>

</thread>
<thread>
creation date:
1729992196967

thread_type:
official_review

<thread_object>
creator: reviewer_aqjo

<content>
summary:
The authors introduce new step-size schedules which are inspired by stochastic Polyak step-size (SPS) for the Stochastic Heavy Ball (SHB) momentum optimization method. The authors present convergence rates for each of their step-size selections showing that SHB can converge to the true solution at a fast rate when interpolation is satisfied and converge to the minimizer adaptively when interpolation is not satisfied.

strengths:
The paper effectively integrates SPS with SHB, an approach previously applied to SGD in the existing literature. The proposed methods are validated both theoretically and empirically by demonstrating their convergence rates and conducting relevant experiments.

weaknesses:
Gaining understanding of SHB within SPS framework and establishing a convergence rate is a valuable contribution. However, identifying the specific settings or problem types where the proposed method is most effective would provide a more comprehensive and practical understanding of its applicability. Based on my current understanding of the methods, from a theoretical perspective, there does not appear to be a clear advantage for incorporating momentum within the SPS framework.

questions:
In the convex setting, the paper indicate that no momentum produce the best convergence rate (line 354), the convergence rate in Theorem 3.5 and Theorem 3.6 deteriorate when $\beta$ increases. Hence, within the SPS framework, which specific settings or problem types momentum should be employed instead of relying solely on SGD?

The assumption of bounded iterates in Section 3.2 constitutes a strong condition, which may not hold even for quadratics without additional constraints. In the absence of this assumption, is it possible to derive meaningful theoretical results for the MomDecSPS and MomAdaSPS methods?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730590649064

thread_type:
official_review

<thread_object>
creator: reviewer_qp47

<content>
summary:
This paper extends the recently developed Stochastic Polyak Stepsize (SPS) to incorporate the Pokyak momentum. In particular, the paper analyzes three SPS variants, vanilla SPS, DesSPS, and AdaSPS, and gets convergence guarantee for their momentum variants: MomSPS, MomDesSPS, and MomAdaSPS. Extensive numerical experiments demonstrate the practical performance of the proposed methods.

strengths:
The paper is well-written and solid, with details and intuitions clearly explained to the readers. Moreover, extensive experiments validate the practical performance of the proposed stepsizes.

weaknesses:
Although the paper is a solid work, my foremost concern is its limited technical novelty. 

1. Limited technical contributions

   SPS [1], DesSPS [2], and AdaSPS [3] considered in this paper are already present in the literature, and the contribution of the paper only lies in showing the feasibility of incorporating momentum into these existing stepsizes.

2. Strong assumptions

   The analysis of MomSPS imposes a constraint on $\beta$, and the analysis of two other stepsizes needs bounded iterates. Although the authors claim these are standard assumptions, they do not usually appear in the literature for SHB methods. 

Overall, despite the limited theoretical novelty due to the strong assumptions, I believe this paper is a solid contribution to the SPS literature and I recommend weak acceptance.

questions:
**Questions**

1. The comparison after Corollary 3.4 looks unclear to me. In particular, the results in [6] can achieve $O(1/\sqrt{K})$ convergence in the presence of noise by taking $\alpha =O( 1/\sqrt{K})$ (Theorem 1, [6]) , while SPS cannot achieve exact convergence. Could you elaborate more on this comparison?
2. Momentum in stochastic optimization is often shown to achieve certain variance reduction effects [4, 5]. Do you think this can help improve dependence on $\sigma^2$ in the convergence analysis of MomSPS?

**Minor issues**

1. Line 55

   has efficiently analyzed => has been efficiently analyzed.

2. Line 400, 410

   Please be consistent with step size and step-size.

3. Line 1098

   Definition C.2. Smoothness requires a two-sided bound $|f(x) - f(y) - \langle \nabla f(y), x- y\rangle | \leq \frac{L}{2}\\|x - y\\|^2 $. Concave functions without Lipschitz continuous gradient can satisfy the given one-sided bound.

**References**

[1] Nicolas Loizou, Sharan Vaswani, Issam Hadj Laradji, and Simon Lacoste-Julien. Stochastic polyak step-size for sgd: an adaptive learning rate for fast convergence. In *International Conference on Artificial Intelligence and Statistics*, pages 13061314. PMLR, 2021.

[2] Antonio Orvieto, Simon Lacoste-Julien, and Nicolas Loizou. Dynamics of sgd with stochastic polyak stepsizes: truly adaptive variants and convergence to exact solution. *Advances in Neural Information Pro- cessing Systems*, 35:2694326954, 2022.

[3] Xiaowen Jiang and Sebastian U Stich. Adaptive sgd with polyak stepsize and line-search: robust convergence and variance reduction. *Advances in Neural Information Processing Systems*, 36, 2024.

[4] Cutkosky, Ashok, and Francesco Orabona. Momentum-based variance reduction in non-convex sgd. *Advances in neural information processing systems* 32, 2019.

[5] Gao, Yuan, Anton Rodomanov, and Sebastian U. Stich. Non-Convex Stochastic Composite Optimization with Polyak Momentum. In *Forty-First International Conference on Machine Learning*, 2024.

[6] Liu, Yanli, Yuan Gao, and Wotao Yin. An improved analysis of stochastic gradient descent with momentum. *Advances in Neural Information Processing Systems* 33, 2020.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730713906243

thread_type:
official_review

<thread_object>
creator: reviewer_smtn

<content>
summary:
This paper introduces Polyak step size to Stochastic Heavy Ball method. To do this, authors consider iterate moving-avergage viewpoint of stochastic heavy ball method. With slight modification of the step size adjustment technique, authors propose three new step size selection: one is based on regular Polyak step size for SGD, and the rest - on decreasing versions of Polyak step size for SGD. The resulting algorithms do not need knowledge about any of problem parameters and converge either to exact solution in non-interpolated regime or to the area of solution in interpolated regime. Convergence rates match the existing state-of-the-art results.

strengths:
Authors propose a connection between both Polyak step size and adaptive Polyak step size with stochastic heavy ball method, which seems like an important result for optimization community. The proposed technique does not need any knowledge about problem parameters, which makes methods easier to implement. Theoretical convergence results show that proposed methods either outperform or perform on the same level with existing state-of-the-art algorithms while sometimes mitigating some of the issues of these algorithms. Finally, experimental results show either similar or better performance on both synthetic and real problems in comparison with existing analogs.

weaknesses:
The Section 3 is written in rather overwhelming way. After each theorem goes a paragraph, mentioning some special cases and how the provided result matches the results from other existing works, which is rather hard to comprehend for a new reader. You could change these paragraphs to small tables with small supporting captions. This should take about the same size as the paragraph in text, although it would be much more evident. Additionally, after every theorem follows a lot of corollaries, that again consider some special cases. It is easy to get lost in so many special cases. I think, some of these corollaries can also be moved to appendix, leaving some mention about them in the main part. Since the main result is introduction of Polyak step size and adaptive Polyak step size to SHB, it should be the central point of main part of the paper. 
Overall, despite the seemingly good results, considering everything above and some suggestions in Question part, I think the main text needs polishing and not ready yet for publication.

questions:
## Questions
1. Am I correct, that by "interpolated regime" you mean that noise variance is zero and we have determenistic case?
1. Figures 8 and 9. In first figure you have orange color for your method, in the second - blue. This is confusing, please, fix the order of the plots in all the figures for consistency. Also in Figure 8 you do not use the name of your particular algorithm. Please also check, that these issues are fixed for all other figures.
1. Lines 382-383. What do you mean by "saddle connection"?

## Suggestions
1. Please, add explicit assumptions in the text about interpolation and smoothness etc., that you use in the theorems in some section like "Preliminaries". Sometimes it is hard to follow, whether you mean interpolation or bounded variance and it is hard to find them in the text.
1. Please, increase font size in the figures.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732158146654

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response to all reviewers

comment:
We thank the reviewers for their feedback and time. 

In particular, we appreciate that the reviewers acknowledged the following strengths of our work:
* Reviewer **SmTn** acknowledges that our work is an important result for the optimization community and recognizes that our stepsizes do not need any knowledge about problem parameters, which makes methods easier to implement. 
* Reviewer **QP47** finds our well-written, with details and intuitions clearly explained to the readers as well as acknowledges that it is solid work. 
* Reviewer **aQjo** appreciates that our proposed stepsizes are validated both theoretically and empirically and finds our convergence rates a valuable contribution. 
* Reviewer **jVKN** highlights that our writing and exposition are clear and that our proposed stepsizes have theoretical improvements compared with existing works. 

With our rebuttal, we address all raised issues. **Here we highlight again that with our work:**

* We explain why naively using SPSmax as a step-size $\gamma_t$ in the update rule of the Stochastic Heavy Ball method (SHB) **leads to divergence** even in simple problems (see also Figure 1). 
* To resolve this issue, we provide an alternative way of selecting Polyak-type step-sizes for SHB **via the Iterate Moving Average viewpoint, and we propose three novel adaptive step-size selections** (MomSPSmax, MomDecSPS, MomAdaSPS).
* MomSPSmax is an adaptive step-size selection that depends on the momentum parameter $\beta$. Most importantly, **SHB with MomSPSmax is the first adaptive SHB with convergence guarantees for the fully stochastic setting (no interpolation assumed).**
* Our decreasing variants, MomDecSPS and MomAdaSPS, are the **first adaptive step-sizes for SHB with convergence guarantees to the exact solution (in the non-interpolated regime)**. Furthermore, via MomAdaSPS we provide the first robust convergence of adaptive SHB that guarantees convergence to the exact solution and automatically adapts to whether our problem is interpolated or not.
* All the provided convergence guarantees are **tight** in the following sense: If $\beta=0$ (no momentum) in the update rule of SHB, our theorems recover as a special case the best-known convergence rates of Polyak-type step-size for SGD.
* Finally, we have extensive numerical evaluations both in the convex and non-convex regimes. 

**We hope that you will engage with us in a back-and-forth discussion and we will be most happy to answer any remaining questions.**
</content>
</thread_object>

</thread>


Metareview:
metareview:
Summary:
This paper extends Stochastic Polyak Stepsize (SPS) methods to incorporate momentum through the Stochastic Heavy Ball (SHB) method. The authors propose three novel adaptive stepsize selections (MomSPSmax, MomDecSPS, MomAdaSPS) and provide theoretical convergence guarantees. Most notably, MomSPSmax is the first adaptive SHB with convergence guarantees for the fully stochastic setting, while the decreasing variants are the first adaptive stepsizes for SHB that guarantee convergence to the exact solution in the non-interpolated regime.

Main Strengths:
The work provides a significant theoretical contribution by developing convergence guarantees that are tight and recover best-known rates for SGD as special cases. The proposed methods are practical to implement as they require no knowledge of problem parameters. The paper is well-written with clear exposition of technical details and intuition, supported by extensive numerical evaluations in both convex and non-convex settings.

Main Weaknesses:
Initial concerns were raised about the technical novelty given that the work builds on existing SPS methods. There were also questions about the strength of assumptions required, particularly the bounded iterates assumption. Some reviewers suggested improvements to the presentation of theoretical results to enhance readability.