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

We would like to thank the reviewer for their time and positive evaluation. Below, we address the concern raised by the reviewer.

Stating that our approach or techniques are the same as [Loizou, 2021] is comparable to claiming that all analyses of SGD with momentum are equivalent to those of vanilla SGD. Our work builds on previous techniques but introduces significant differences that result in new theoretical contributions and practical implications.

On proof techniques:

Compared to the analysis of SGD with SPS from [Loizou, 2021], the SHB update rule requires taking into consideration extra terms related to the previous iterate $x^{t-1}$. This needs new algebraic techniques in order to use convexity, and after that, one needs to deal with the two quantities $f(x^t)$ and $f(x^{t-1})$ at the same time. We deal with this using the IMA framework, similar to Sebbouh [2021]. However, we highlight that we finish our proofs with a different approach. More specifically, in the proof of Theorem G.1 in Sebbouh [2021] the authors choose the momentum coefficients $\lambda_t$ in such a way that the main sum telescopes and the final result only has the expectation of the last iterate. In our proof, we make no such simplifications, and instead, we use Jensen’s inequality to finish. This approach allows us to have more freedom in the selection of $\lambda_t$. This is more prominent in the proofs of the decreasing variants.

Compared to the classical analysis of constant step-size SHB, the use of SPS requires an adaptive step-size that uses the loss and stochastic gradient estimates at an iterate, resulting in correlations of the step-size with the gradient, i.e. we might have $\mathbb{E}[\gamma_t\nabla f_i(x^t)]\neq\gamma_t\mathbb{E}[\nabla f_i(x^t)]$. One of the technical challenges in the proofs is to carefully analyze the SHB iterates, taking these correlations into account. Moreover, since we try to be adaptive to the Lipschitz constant, we can not use any standard descent lemmas (implied by the smoothness and the SHB update) that are used in the classical analysis of SGD and constant step-size SHB. This makes the convex proof more challenging than the standard analysis of SHB.

Building on the analysis of previous works should not be considered a drawback or a reason to question the novelty of a paper. In the end, a paper should be valuable if it answers open questions in the literature, and we believe that our work does that. If you agree that we managed to address all issues, please consider raising your mark to support our work. If you believe this is not the case, please let us know so that we have a chance to respond.

The authors' rebuttal basically addresses my concern. I increase the contribution score to 3 and maintain the overall score.

We thank the reviewer for taking the time to review our work. We appreciate that the reviewer finds our convergence rates “a valuable contribution”. Below, we address all the concerns and questions raised.

On using momentum:

The main advantage of momentum lies in increasing the generalization performance of an algorithm in applications. In all of the experiments of the paper, we demonstrate that our MomSPS variant provides significant benefits compared to the no-momentum variants, which is not achieved by simply applying standard momentum to SPS-type methods. We agree with the reviewer that theoretically, our results do not show an advantage of using momentum on top of Polyak-type stepsizes. However, to the best of our knowledge, this is the case with all existing analyses of momentum on top of vanilla SGD. The benefit of the momentum in all existing works is practical and typically leads to better generalization. This should not undermine the main contribution of our work which is the understanding that the naive use of momentum on top of Polyak-type stepsizes might fail to converge and the proposal of the IMA viewpoint that allows us to resolve the issue.

On bounded iterates assumption and projection:

Indeed our two main theorems on the convergence of MomDecSPS/MomAdaSPS require the bounded iterates assumption, however, the results are meaningful. This is exactly the assumption used in the original papers by [Orvieto, 2022] and [Jiang, 2023] that proposed the analyses of the no-momentum variants (DecSPS/AdaSPS). In these papers, it is possible to remove the bounded iterates assumption by assuming strong convexity of individual f_i. This is a strong alternative assumption that holds for a much more restrictive setting than the general convex regime, which is the main focus of our work. We speculate that with this extra assumption, a similar result might be possible for the MomDecSPS/MomAdaSPS too.

Another alternative to removing this assumption is by projection: Our results can be extended to the constrained regime by adding a projection step in the update rule of IMA onto a convex and bounded domain D. That is $z^{t+1}=proj_D(z^t-\eta_t\nabla f_{S_t}(x^t))$ with $x^0=z^0\in D$. Note that the projection step is only needed for the update of $z^{t+1}$ because $x^{t+1}$ is already a convex combination of elements in the convex set D. Our current proofs will then go through using the non-expansiveness of the projection operator.

If you agree that we managed to address all issues, please consider raising your mark. If you believe this is not the case, please let us know so that we have a chance to respond.

We thank the reviewer for taking the time to review our work. We appreciate that the reviewer finds our paper “solid work”. Below, we address all the concerns and questions raised.

Limited technical contributions:

Building on the analysis of previous works should not be considered a drawback or a reason to question the novelty of a paper. In the end, a paper should be valuable if it answers open questions in the literature and we believe that our work does that. That being said, in this work, we do not just blindly take the previous stepsizes and we prove that they also work in the SHB setting. This is the naive approach we discuss in Section 2.2 and we demonstrate that it might lead to divergence even in simple scenarios. We use the IMA framework to derive our stepsizes and to make sure they depend on the momentum coefficient β. Furthermore, the proofs are not straightforward adaptations from the SGD case. Compared to the analysis of SGD with SPS, the SHB update rule requires taking into consideration extra terms related to the previous iterate $x^{t-1}$. This needs new algebraic tricks in order to use convexity, and after that, one needs to deal with the two quantities $f(x^t)$ and $f(x^{t-1})$ at the same time.

Strong assumptions:

Indeed our two main theorems on the convergence of MomDecSPS/MomAdaSPS require the bounded iterates assumption. This is exactly the assumption used in the original papers by [Orvieto, 2022] and [Jiang, 2023] that proposed the analyses of the no-momentum variants (DecSPS/AdaSPS). In these papers, it is possible to remove the bounded iterates assumption by assuming strong convexity of individual f_i. This is a strong alternative assumption that holds for a much more restrictive setting than the general convex regime, which is the main focus of our work. We speculate that with this extra assumption, a similar result might be possible for the MomDecSPS/MomAdaSPS too.

Another alternative to removing the bounded iterates assumption is via projection: Our results can be extended to the constrained regime by adding a projection step in the update rule of IMA onto a convex and bounded domain D. That is $z^{t+1}=proj_D(z^t-\eta_t\nabla f_{S_t}(x^t))$ with $x^0=z^0\in D$. Note that the projection step is only needed for the update of $z^{t+1}$ because $x^{t+1}$ is already a convex combination of elements in the convex set D. Our current proofs will then go through using the non-expansiveness of the projection operator.

[Q1]: Let’s focus firstly on the no momentum case (ie β=0) where SHB reduces to SGD. SGD with constant stepsizes converges up to a neighborhood while SGD with decreasing stepsizes converges to exact solution with rate $O(1/\sqrt{K})$. A similar thing is true for SPS. SGD with SPSmax converges up to a neighborhood while SGD with DecSPS converges to an exact solution with rate $O(1/\sqrt{K})$. All of these can be transferred to SHB. As you mentioned SHB with decreasing stepsizes ( $a=O(1/\sqrt{K})$) converges to exact solution with rate $a=O(1/\sqrt{K})$. In Thm 3.5, we show that SHB with MomDecSPS (and MomAdaSPS) converges to the exact solution again with rate $a=O(1/\sqrt{K})$.

[Q2]: This is an intriguing question. We do not know the answer to this, however, it seems like an interesting future research direction.

Minor issues: Thanks for catching them. All of them will be fixed.

We are glad that the reviewer recognizes our contribution to the SPS literature and would like to reiterate that our work addresses broader regimes with no stronger assumptions than prior studies, making it a significant theoretical advancement. We hope that with our response, we will clarify the theoretical novelty of our results further and explain the assumptions used (no stronger than previous works on adaptive methods). If you agree that we managed to address all your questions, please consider raising your mark to support our work. If you believe this is not the case, please let us know so that we have a chance to respond.

We thank the reviewer for taking the time to review our work. Below, we address all the questions raised.

We respectfully disagree with the reviewer’s statement that the paper is “not ready yet for publication” due to its presentation. This is a theory/optimization paper, and our structure aligns with the standard conventions for papers in this field, emphasizing precise and rigorous statements in the main body. Moreover, two other reviewers (QP47 and jVKN) have explicitly commented on the clarity of the exposition and noted that the paper is well-written. Furthermore, the reviewer suggested adding tables; however, we already include a table (Table 1) that provides informal summaries of our theorems alongside comparisons to related works. Additionally, regarding the suggestion for a “Preliminaries” section, we do have such a section in Appendix C, titled “Technical Preliminaries,” where all details are presented (due to space limitation, such details are typically included in the appendix)

About the reviewer’s Questions:

[Q1]: You are correct that interpolation means that the noise variance ($\sigma^2$) is zero. However, this does not necessarily imply that we are in the deterministic case. The deterministic case is only a special case of a setting where the variance is zero. Besides that, there are plenty of stochastic interpolated examples. As mentioned in the paper, non-parametric regression (Liang and Rakhlin, 2020) and over-parameterized deep neural networks (Zhang et al., 2021; Ma et al., 2018) satisfy the interpolation condition.

[Q2]: Thanks for catching these. These can be easily fixed.

[Q3]: The word saddle can be removed.

As we highlighted above, the current structure of our paper and its presentation effectively communicates the content. We believe that the presentation is a personal style of the authors of each paper, and one should not force their preferences on them. Providing feedback is, of course, welcome for improving details on the exposition, and we appreciate the comments. However, claiming that a theoretically focused paper should not provide details of the main theoretical statements as corollaries and remarks (a common practice in the majority of optimization/theory-oriented papers) after the main Theorems in the main paper is not constructive feedback.

We agree with the reviewer that the main result is introducing the adaptive Polyak step size to SHB, and we devote section 2 (more than 3 pages of the main paper) to exactly the introduction of the main ideas and the connection with what exists. However, just the introduction of the new ideas is not enough; one needs a detailed exposition of the convergence guarantees of the method in different variants and settings. This is when the main theoretical results thrive compared to what is known in the literature. At first glance, Section 3 might look overwhelming for the reviewer, but it is divided into sections with the main theorems highlighting the importance of each result and the discussion after explaining the benefits and limitations of the proposed analysis. This is done by design so that experts in the area can read and understand exactly the importance of each statement.

The only concern raised as a weakness of our work is the presentation of Section 3. As you argue, this is a personal taste of the authors of each paper and should not be considered a reason to suggest rejection. Please see the other reviewers that highlight the quality of this section and how well-written it is. If you agree that we managed to address your concers, please consider raising your mark. If you believe this is not the case, please let us know so that we have a chance to respond.

The authors adressed all of my concerns. I increase my rating to 8, and presentation score - to 2.