Stochastic Polyak Step-sizes and Momentum: Convergence Guarantees and Practical Performance
Authors: Dimitris Oikonomou
Nicolas Loizou
Venue: ICLR 2025 Poster (2025)
Keywords: Stochastic Polyak Step-size
SGD
Stochastic Heavy Ball
Convergence Analysis
Convex Optimization
Momentum
Primary area: optimization
TL;DR: Efficient convergence analysis of SGD with stochastic Polyak step-sizes and heavy-ball 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.

Published as a conference paper at ICLR 2025
STOCHASTIC POLYAK STEP-SIZES AND MOMENTUM: CONVERGENCE GUARANTEES AND PRACTICAL PERFORMANCE
Dimitris Oikonomou CS & MINDS
Nicolas Loizou
AMS & MINDS
Johns Hopkins University dof konolol & edu
Johns Hopkins University nloi zoul jh. edu
ABSTRACT
Stochastic gradient descent with momentum, also known 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 pro- hibitively 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 suit- able 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 maxx 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 the exact minimizer without 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 special case. We supplement our analysis with experiments validating our theory and demonstrating the effectiveness and robustness of our algorithms.
maxx.
INTRODUCTION
We consider the unconstrained finite-sum optimization problem.
1 12 n Σ f(x)
min xERd
f(x)
(1)
where each fa Rd + R is convex, smooth, and lower bounded by &. Let X be the set of minimizers of (1). We assume that X of 0 and we fix x" E X*. This problem is the cornerstone of machine learning tasks, (Hastie et al., 2009), where x corresponds to the model parameters, fe(x) represents the loss on the training point and the aim is minimize the average loss f(x) across training points.
When n is large, stochastic gradient methods are the preferred methods for solving (1) mainly be- cause of their cheap per iteration cost. One of the most popular stochastic algorithms for solving such large-scale machine learning optimization problems is stochastic gradient descent (SGD) with momentum, (Sutskever et al., 2013), a.k.a. stochastic heavy ball method (SHB) given by:
_++1 x² 7: V fs, (x*) + Br -0'=)=''1) =.
(SHB)
where St C [n]] random subset of data-point (mini-batch) with cardinality B sampled indepen- dently at each iteration f, and fs. (x') = 1o2 Eies, ff((2') is the mini-batch gradient. Here Yt is the step-size/learning at iteration l while 3, 2 0 represents the momentum parameter.

---

Published as a conference paper at ICLR 2025
When the momentum parameter 3, = 0, Vt 0. then the update rule SHB is equivalent to the well-studied mini- SGD, x² x=11 = x' % Vf, (x"), (Robbins and Monro, 1951), which has been efficiently analyzed under different properties of problem (1) and different step-size selections "t. (Nemirovski and Yudin, 1983; Nemirovski et al.. 2009; Hardr et al., 2016; Needell et al., 2016; Nguyen et al., 2018; Gower et al., 2019; 2021). Additionally, when the cardinality of the random subset St is B = n. then the update rule of SHB is equivalent to the deterministic heavy ball method (HB) proposed by Polyak (1964), as a way improve the convergence behavior of deterministic Gradient Descent (GD).
There is a rich literature on the convergence of SHB and HB in different scenarios. In Polyak (1964), it was proved that for a specific choice of the step-size and the momentum parameter 3. the HB method enjoys an accelerated linear convergence when minimizing strongly convex quadratic func- tions while more recently, Ghadimi et al. (2015) proved global sublinear convergence guarantee for HB for convex and smooth functions. In the stochastic setting, several works focus on conver- gence guarantees of SHB under constant step-size and momentum parameters (Ma and Yarats, 2019; Kidambi et al., 2018; Yan et al., 2018; Gitman et al., 2019; Liu et al., 2020). However, in practical scenarios, these choices suffer from a prohibitively expensive and time-consuming tuning process. This has motivated a large body of research on the development of adaptive SHB method that adapts their parameters using information collected during the iterative process. Such analysis is challenging, and the current adaptive versions of SHB either focus on the full batch set- ting (deterministic) (Barré et al., 2020; Saab Jr et al., 2022; Wang et al., 2023), or assume that an interpolation condition is satisfied (Schaipp et al., 2023b) or focus on moving averaged gradient (a different form of momentum) (Wang et al., 2023).
parammtr
Previous studies in the fully stochastic (non- mnoninteerpolate) scenario have predominantly concentrated on analyzing adaptive versions of SGD, with limited attention given to developing adaptive variants for the SHB. In this work, we take inspiration from the recently introduced and highly efficient Polyak-type adaptive step-sizes for SGD and investigate its applicability and extension to SHB.
1.1 MAIN CONTRIBUTIONS
Our main contributions are summarized below.
Efficient Polyak Step-sizes via IMA viewpoint. We explain and illustrate by experiment (see Figure 1) why naively using SPS of Loizou et al. (2021) as a step-size "r in the update rule of SHB is not robust, leading to divergence even in simple problems. To resolve this issue, we provide an alternative way of selecting Polyak-type step-sizes for SHB via the Iterate Moving Average view- point from Sebbouh et al. (2021). Through approach, we propose three novel adaptive step-size selections, namely MomSPS max. MomDeeSPS, MomAdaSPS. Each of the proposed step-sizes de- pends the choice of the momentum parameter adding further stability to SHB and comes with specific benefits over their constant step-size counterparts or other adaptive variants of SHB.
MomSPSmact max: Convergence of SHB in non-interpolated setting. Our first step-size selection of SHB is MomSPSmax which has a similar structure to SPS of Loizou et al. (2021) but includes also (1 - 3) in its expression. For this choice, we provide convergence guarantees for SHB to a neighborhood of the solution for convex and smooth problems. Our analysis provides the first convergence guarantees of adaptive SHB using Polyak-type step-size. Previous works on Polyak step-size with momentum in the stochastic setting have guarantees either only under the interpolation setting (Schaipp et al., 2023b) or for a moving averaged gradient momentum (Wang et al., 2023). In addition, as corollary of our main theoretical results, we show that MomSPS max under the interpolation setting and deterministic scenarios (full batch) converges to the true solution at a fast rate matching the deterministic HB.
Convergence of SHB to exact solution via MomDeeSPS and MomAdaSPS. Inspired by two recent Polyak-type step-size selections for SGD, the DecSPS in Orvieto et al. (2022) and AdaSPS in Jiang and Stich (2023), we propose two new ways for tuning the step-size for SHB. These are MoMDecSPS and MomAdaSPS. Our analysis provides the first O(1/VT) convergence guarantees to the exact solution in the non-interpolated regime for Polyak-type adaptive variant of SHB. Our proposed update rules converge for any choice of momentum parameter / E [0, 1), which makes them particularly useful in practical scenarios.
2

---

Published as a conference paper at ICLR 2025
Step-size
Assumptions Setting
Adaptive Exact Convergence Rate
Constant (Liu et al., 2020)
Knowledge of L
0 + +
IMA (Sehbouh et al., 2021)
Knowledge of I
X
0
+
ALR-SMAG (Wang al., 2023)
Moving Averaged Gradient
X
0 +
MomSPS (Thm 3.2)
Restriction on / Bounded Iterates
X
0 +
MomDecSPS (Thm 3.6)
0
MomAdaSPS (Thm 3.7)
Bounded Iterates
0
Table 1: Summary of the considered step-sizes and the corresponding theoretical results in the stochastic setting. All the rates are given for convex and smooth functions. The quantity of conver- gence in our rates is [*f("") - [()")), where x² - x!. Here 6² = E | [=<x")] < oo and o² = [ Ej] < oc. The "Exact Convergence" column refers to convergence to the exact solution with no interpolation assumption.
Tight Convergence Guarantees. All of our convergence guarantees are true generalizations of the theoretical analysis of SGD using SPS DecSPS, and AdaSPS. That is, if 3 0 (no momentum) in the update rule of SHB, our theorems obtain as a special case the best-known convergence rates of Polyak-type step-size for SGD, highlighting the tightness of our analysis. See also Table for summary of our main complexity results and comparison with closely related works.
Further Convergence Results. As byproduct of our theoretical analysis, we provide two in- teresting corollaries: novel analysis of constant step-size SHB (as corollary of our Theorem on MomSPS maxx) and the first robust convergence of SHB via our theorem on MomAdaSPS. For constant step-size, the Corollary 3.5 of our Theorem 3.2 allows larger step-sizes than the analysis of SHB in Liu et al. (2020) and provides convergence without assuming the restrictive bounded variance condition (there exist 0 such that E fi f=))) q). In addition. via The- orem 3.7 we provide the first robust convergence of adaptive SHB that guarantees convergence to the exact solution and automatically adapts whether if interpolation is assumed, then the rate of SHB with MomAdaSPS is the same as the rate of SHB with MomSPS max (or constant step-size), and if no interpolation is assumed, then it matches the rate of SHB with MomDecSPS. The analysis achieves the best-known rates in both settings.
problem is interpolated or not. That is,
Numerical Evaluation. In Section we verify our theoretical results via numerical ex- periments on various problems, demonstrating the effectiveness and practicality of our ap- proach. An open-source implementation of our method is available https: / thub. com/ dimitris -oik k McmSPSS
EXPLORING THE INTERPLAY OF SPS AND HEAVY BALL MOMENTUM
In this section, present the expression of Stochastic Polyak Step-size (SPS) and its different variants. We illustrate that naive combination of SPS with momentum is not robust, leading to divergence in simple problems. We explain how we resolve this issue using the Iterate Moving Average viewpoint of SHB and propose three adaptive Polyak-type step-sizes for SHB.
2.1 BACKGROUND ON STOCHASTIC POLYAK STEP-SIZE
The deterministic Polyak step-size (PS) was first introduced by Polyak (1987) as an efficient step- size selection for GD for solving convex optimization problems. It has the following expression:
T: f(x")-f(ar)) which is obtained by minimizing an upper bound of the quantity ||æ - x x²1² Mr
in the analysis of GD. Since its original proposal, the PS has been successfully used in the analysis of deterministic subgradient methods in different settings with favorable convergence guarantees (Boyd al., 2003; Davis al., 2018; Hazan and Kakade, 2019). PS requires the prior knowledge of f(x*), which might look like a strong assumption. However, as shown in Boyd et al. (2003), this is known in several applications, including finding point in the intersection of convex and positive semi-definite matrix completion.
Inspired by the convergence of PS in the deterministic setting, Loizou et al. (2021) has effectively modified the Polyak step-size for the stochastic setting, achieving convergence rates comparable to those of finetuned SGD. The proposed stochastic Polyak step-size (SPS) has several benefits,
3

---

Published as a conference paper at ICLR 2025
including independence on parameters of the problem (e.g., L-smoothness μ strong convexity) and competitive performance in over- weer-paametrizedd models. More specifically, Loizou et al. (2021) proposed the SPSmax given below":
Tt min { Ss, (x') S. cl Vss, (x²)
mSPS
Here, 76 0 is bound that restricts SPS from being very large and is essential to ensure conver- gence to a small neighborhood around the solution and с > 0 is a positive constant that depends on the function class the objective ʃ belongs to.
As mentioned in Orvieto et al. (2022), the SPS comes with strong convergence guarantees and competitive performance; however, it has main drawback when used in -araovee-prrammetrriieedd regimes: It can guarantee convergence only to a neighborhood of the solution. For this reason, Orvieto et al. (2022) suggests decreasing variant of the original SPS named DecSPS, given by:
fs,(z") - Es. V ct a0-77-11
7
min
(DecSPS)
Ct
where is an increasing sequence of positive real numbers and c-1 == co and 7-1 76 >> 0. The authors proved that SGD with DecSPS and c vt + 1, converges with sublinear rate O(1/VT) for convex and smooth functions with bounded iterates (i.e., D == max || - x= 11 < oo).
More recently, in Jiang and Stich (2023), another decreasing variant of SPS was introduced, named AdaSPS:
fs. (a"") all (I)) Viizo (s,")))" 7-11
Y = min
(AdaSPS)
where 7-1 +00 and c > O. For convex and L-smooth functions with bounded iterates, it can be shown that SGD with AdaSPS converges to exact solution with rate O _ VT) where
2e + -. The interesting aspect of the convergence of AdaSPS is that it provides robust result for SGD, meaning that the method recovers the best bounds for both the interpolated (a = 0) and non-interpolated regimes.
2.2 NAIVE SPS IN SHB
All of the above variants of SPS were proposed and analyzed for SGD (SHB with momentum). With the increased popularity of momentum in machine learning, one can naturally ask the following question: Is it possible combine SPS with momentum? This question was partially answered in Wang et al. (2023) when they proposed a Polyak-type step-size for a moving averaged gradient (MAG) momentum but not for the SHB. Their MAG framework is not equivalent to SHB when their stepsizes are adaptive and while their proposed stepsize reduces to the original SPS stepsize when 8 0, their guarantees do not reduce to the original SPS max guarantees for convex functions.
The most straightforward approach is to directly apply the SPSmax in the SHB up- date rule. Let us call this update rule SPSmax with naive momentum. Unfortu- nately, this approach does not necessar- ily lead to convergence for natural choices of momentum parameter 3 E (0, 1), as shown in Figure 1.2 In this experiment, even for simple convex and smooth prob- lems like a logistic regression with syn- thetic data, the naive rule fails to converge when 3 gets larger (a typical choice for 0.9). This
Hond
Figure 1: Comparison of MomSPSman max versus SPSmax with naive momentum for different momentum param- eters 3 on a logistic regression problem. Left: 3 = 0.2, Right: 3 0.3
momentum parameter is indicates that more careful step-size se- lection is needed, which may also depend
Originally Loizou et al. (2021) use fö, instead the lower bound f The lower bound is due to Orvieto et al. (2022), which proves that the more relaxed lower bound can still lead to the same convergence guarantees. "See also Figure 9 in Appendix for more values of 3.
4

---

Published as a conference paper at ICLR 2025
on the momentum parameter. In Figure 1, compare the SPSmax with naive momentum with one of our proposed and analyzed step-size selection MomSPSmax As seen in Figure 1, when / is small, then both SPSmax with naive momentum and SHB with MomSPS have similar behavior. When 3 = O, the two methods have identical performance as both are reduced to the same method: SGD with SPS However, B gets larger, the performance of SPS with naive momentum gets worse and less stable, and at some point, it diverges. On the other hand, in this example, SHB with MomSPS max converges for any of the selected
2.3 ITERATIVE MOVING AVERAGE: BALANCE BETWEEN SPS AND MOMENTUM
Having explained how naively combining SPS with heavy ball momentum can lead diver- gence, in this section, we leverage the iterate moving-averagg (IMA) viewpoint of SHB from Seb- bouh et al. (2021) to propose SPS-type adaptations that depend momentum parameter well for SHB. As we will see later, this viewpoint provides stability and robustness to the proposed update rules.
Iterate Moving Average (IMA). Sebbouh et al. (2021) provide the IMA as an alternative way of expressing the update rule of SHB and explain how the IMA formulation is crucial in comparing SHB and SGD as it allows to establish connections between the step-sizes of the two methods. The Iterate Moving Average method is given by the following update rule:
Ae+1
1+1 =² T/e (x),
x: 1+1
1+1
A441
A+t1 +
where 20 x0 7t 0 and X, 2 As proved in Sebbouh al. (2021), if for any E N, holds A and 7): (1 + % then the 2; iterates of the IMA method are equal the 2t iterates produced by the SHB method. For completeness, we include the proof of this statement in Appendix C.3.
+ Ac+1 3:
In our proposed methods, we select the most common setting of constant momentum (B, B)- Considering the equivalence between IMA and SHB under the correct parameter selection, let us present the following proposition that will allow us to obtain convergence guarantees for SHB via convergence analysis of IMA.
Proposition 2.1. If X₂ = X 2 0 then assuming that B I+x and "tt = (1 - B)T the 2, iterates of the IMA method are equal to the z, iterates produced by the SHB method.
Using the above proposition, let us provide the following corollary that explains the derivation of proposed step-size selections for SHB.
Corollary 2.2. Let the step-size Tt. IMA be one of the previously proposed Polyak step-sizes: SPSmax. DecSPS AdaSPS and let x Then via Proposition 2.1, the SHB has constant momentum parameter 3 = and the step-sizes are MomSPSmaxs MomDecSPS or MomAdaSPS respectively.
Three New Adaptive Step-sizes for SHB. Using the IMA viewpoint and Proposition 2.1 and Corol- lary 2.2 let present three Polyak-type step-size for SHB.
(i) MomSPSmax- Let us start with the variant associated with the SPS That is, via Corollary 2.2 use SPS step-size for IMA, then this is equivalent using the following step-size in the update rule of SHB:
Ss, (x+) f cl/V (1'))
it = (1 3) min
(MomSPS
Here, the parameter "t > has the same purpose as the original SPS and it a bound that re- stricts MomSPS max from being very large and is essential to ensure convergence to a neighborhood around the solution. is clear from its expression that MomSPS has the same form as SPSmas but multiplied by correcting momentum factor 1 3. This follows from Corollary 2.2. In practice, this small change allows the SHB to be more "stable" for different momentum parameters, as we show in Figure 1. Using MomSPS max and restriction the momentum parameter 3 in Section 3, we establish O(1/T) convergence for SHB up to neighborhood of the solution for convex and smooth functions.
5

---

Published as a conference paper at ICLR 2025
Our following two proposed step-size selections variants of the two decreasing variants of Polyak step-sizes DecSPS and AdaSPS. As such, we can prove SHB's 's convergence to the exact solution instead of neighborhood. More importantly, using MomDecSPS MomAdaSPS, convergence guarantees for any choice of the momentum coefficients 3 E [0, 1) making SHB fully adaptive (no tuning necessary).
provide
(ii) MomDecSPS. Firstly, we propose an adaptation of DecSPS with momentum. We call this step- size selection MomDecSPS, and it is given by:
(1 3)lfs, fl
1--10-11
7t min
(MomDecSPS)
<ell [*sss(())]'
Ce
76 0 is step-size bound and C cvt+1 for 2 0 with C-1 cp be a constant to regulate the step-size. Note that "/ is indeed decreasing. This step-size is adaptive and, in its definition, does not require knowledge of any function properties (e.g., smoothness/strong convexity constants). In Section 3 we prove that SHB with MomDecSPS and constant momentum
where 7-1
| E [0, 1) converges to exact solution with a rate O(1/TT).
(iii) MomAdaSPS. Similar to the previous two steps, let propose the following adaptation of AdaSPS with momentum. We call this MomAdaSPS, and is given by:
(("))a 7/ = min ( (w') 11-1
(MomAdaSPSS)
By definition, this decreasing step-size (% S 7-1)) In Section 3, we prove that SHB with MomAdaSPS also converges to the exact solution for any choice of B E [0, 1). More specifically, we show that are in the interpolation regime (i.e., o² 0), the convergence of SHB with MomAdaSPS matches SHB with MomSPS mmx while when the setting does not satisfy interpolation condition, this became equivalent to the O(I/VT) of MomDecSPS. Following terminology of Jiang and Stich (2023), our results are the first robust adaptive step-size selection for SHB, it can automatically adapt to the optimization setting (interpolation vs. non-interpolation)
converges to the exact solution with rate O (+ + Thus, when we
Let us close this section by mentioning two remarks related to the above three step-size selections:
Remark 2.3. The "correcting factor" - /3 is outside the minimum in the MomSPS while for the decreasing variants MomDecSPS and MomAdaSPS, it only appears in the first term of the minimum. This follows from Proposition 2.1 and the derivation of Corollary 2.2 and is explained in detail in Appendix C.3. In Appendix F.3, via experiments, we also test other variants of the above step-sizes with the correcting factor outside the minimum. We observe that the other choices do not have good practical performance.
Remark 2.4. For no momentum, i.c.. 3 all three proposed step-size choices, MomSPS max: MomDecSPS and MomAdaSPS are reduced to SPS DecSPS and AdaSPS and our conver- gence analysis recovers the convergence guarantees for SGD showing the tightness of our ap- proach.
CONVERGENCE ANALYSIS
In this section, we present the convergence results for SHB with all of the proposed Polyak step- sizes. For the formal definitions and helpful lemmas, see Appendix C. The proofs can be found in Appendix For all our results, we make the following assumption:
Assumption 3.1 (Finite optimal objective difference).
d² = (s,, - (s. = f(x") - Es, [C', < oo
This assumption first introduced in Loizou et al. (2021) and Orvieto et al. (2022). Note that o² < oc when all fa are lower bounded like we have assumed. We say that problem (1) is in- terpolated if o² 0. If the interpolation condition is satisfied then there exists x² X such

---

Published as a conference paper at ICLR 2025
that f(x*) = fz, = E3s for all St [n]. Many modern machine learning models satisfy this condition. Examples include non-parametric regression (Liang and Rakhlin, 2020) and over- parameterized deep neural networks (Zhang et al., 2021; Ma et al., 2018).
3.1 CONVERGENCE TO A NEIGHBORHOOD OF THE SOLUTION
We start with the analysis of SHB with MomSPS
Theorem 3.2. Assume that each Si is convex and Li-smooth. Then, the iterates of SHB with
MomSPS max with = and 3 E 0, (ū-aa where a = min 1 max; {Li), converge as
TLmax 71 and Lmax =
B[J(T") f(x))] Cj C11100 xz"
+ C202
T
where x²a Et= x² and the constants C1
1-8 a,44 2376
and C2
2n -00--00-
2/7%%
Firstly, let note that both constants C1 and C2 in Theorem 3.2 positive since the denominator of both is positive because 3 E [0. 1) and the numerator of C2 is positive because % 2 > 0. The above result shows that MomSPSmax has sublinear convergence to a neighborhood matching the best-known rate for SHB in the convex setting, see Liu et al. (2020). Moreover, when 8 0, the update rule of SHB with MomSPSmax becomes equivalent to SGD with SPS (no momentum) and the result of Theorem 3.2 reduces to convergence guarantees provided in Loizou et al. (2021), only with slightly better neighborhood of convergence (here when B = we have C2 = (2% aa/a while in the original it is C2 2n//a).
In addition, note that in Theorem 3.2, there a restriction in the momentum coefficient. This stems from the technique used in the proof (which forces C₁ and C2 to be positive). However, this restriction is not vital as in practical scenarios, the method converges even when 3 is selected outside the given interval. Moreover, let us highlight that the constants C1 and C2 are increasing when viewed as functions of B in the given interval (see also Figure 14). This means that our theory suggests that 3 (no momentum) is the best theoretical choice, which typical in many works on stochastic methods with momentum (Wang et al., 2023; Loizou and Richtárik, 2020). This is not ideal but does undermine the importance of Theorem 3.2, as this is the first result showing the convergence of SHB with SPS in non-interpolated setting.
From Theorem 3.2, we can also deduce rates for the deterministic and interpolated regimes.
Corollary 3.3 (Deterministic Heavy Ball Method). Assume that f is convex and T-smooth. Then HB with 3 E [0, a22% a)) and
{
f(x')-
f(x') - f(x²) }
tt (1 B) min
|| .77
(MomPS
where a min 27. nh converges as miinen)i(((>)) - f(x*)} s Taßta 27
(1-0)]>*-
Note that in the setting of Corollary 3.3, we have used fx** instead of lower bound t". This is in accordance with previous work in the deterministic setting (Polyak, 1987). We highlight that in Hazan and Kakade (2019) (PS) and Wang et al. (2023) (ALR-MAG), there is no upper bound in the step-size for the deterministic setting. This is due to the fact that in these works, the step-size is derived as a minimizer of an upper bound of the quantity ||2- x I|2 Our step-size MomPS max is not a minimizer of such bound when 8 >> O. Nevertheless, when 8 = 0 (no momentum), by setting no oo we recover the corresponding result in Hazan and Kakade (2019). In our setting, we cannot set 7h when 3 > 0 because the term B%, appears in the denominator of our convergence result.
Corollary 3.4 (Interpolation). Assume interpolation (o² = 0) and let all assumptions of The- [0, a2% oi) where min 21.max 76 converges as Bff f(x*)] (ap+a 24%)7 where If = ET- -1 x..
orem 3.2 be satisfied. Then SHB with MomSPS max and [1-8)] =="

---

Published as a conference paper at ICLR 2025
In the closely related works Loizou et al. (2021) (SPSmax) and Wang et al. (2023) (ALR-SMACH
there is no upper bound in the proposed step-sizes for the interpolated setting. In our analysis, we focus on the fully stochastic setting (interpolation is only corollary), and the result of Corollary 3.4 shows a saddle connection between the momentum parameter and the parameter " (we cannot select V6 when momentum is used). Nevertheless, via corollary, when / 0 by setting 7/6 oc, we can recover the corresponding result in Loizou et al. (2021).
Another Corollary of Theorem 3.2 is novel analysis for SHB with constant step-size.
Corollary 3.5 (Constant Step-size). Let all assumptions of Theorem 3.2 be satisfied. If 7o I
2lomaxx then SHB with 7 S 1-8 and [0, 1) converges ([ < ||x--x'1 -x'1" +0?
2hmaxx where IJ Ei= x².
T7
Note that in Corollary 3.5, there no restriction that depends on the smoothness parameter in the
momentum parameter (3. We allow have B [0, 1). A similar result is obtained in Liu et al. (2020), where the authors establish the convergence rate O [(29) f(x")
Lyd²
for the quantity 1-p
+ E fs. (x*) [[3] when (1-2)² T. min T-B1.8 2 23422
Here 8² is assumed to be
bound of the variance E, ( S 0². In comparison, in analysis, we provide the same asymptotic rate O(1/T) for the common in the convex setting quantity E[f(=") f(x")) By comparing the two results, our step-size is larger when 3 2 v5 2 0.236. However, in our result, the neighborhood of convergence is constant and does not depend on the step-size y or the momentum coefficient. For numerical comparison of SHB with the two different constant step-sizes, see Appendix F.2.
3.2 CONVERGENCE TO THE EXACT SOLUTION
In this section, we provide theoretical results for the adaptive decreasing step-sizes, MomDecSPS and MomAdaSPS. The main advantage of these decreasing step-sizes is that can guarantee vergence to the exact solution while keeping the main adaptiveness properties. Due to the nature of the decreasing step-sizes, there is restriction in the momentum parameter, as was the case with Theorem 3.2. For the results of this section, we make the extra assumption of bounded ilerates, ie., we assume that 2 max [T] ||.= 112 is finite. This is a standard assumption for several adaptive step-sizes, see: (Reddi et al., 2018; Ward et al., 2020; Orvieto et al., 2022; Jiang and Stich, 2023).
Theorem 3.6. Assume that each fs is convex and L;-smooth. Let म.. = T ET L' a = min 2I } and D2 = maxte[1] ||xd x=1 Then, the iterates of SHB with MomDecSPS with c = and 3 E [0, 1) converge as:
Elf(=") f(x))] v 28[f(x0) ff(x-)
+ (1 (1+ +
202
(1 (1- ()-
(1 (1-8)ark ()-B)aVT VT
Notice that when there is no momentum, i.e., for = 0, Theorem 3.6 recovers the original result of SGD with DecSPS from Orvieto al. (2022). To knowledge, the only other result that guarantees convergence to the exact solution for SHB in the non-interpolated regime appears in Sebbouh et al. (2021). In particular, Sebbouh et al. (2021) proves that for 7t
I+A.41 and 3,
d where nh - with 0 Sn T and X = Σ THA, as E[f(x') f(x)) < O(lott/V)) This result only holds for specific choice of 3t, it needs knowledge of Lmax. and only guarantees (logt/ (/t) rate. However, does not make bounded iteration assumption, and it shows convergence of the last iterate. In contrast, our convergence of SHB with MomDeeSPS holds for any choice of E [0, 1), the siep is adaptive and does not depend on L and it achieves a rate of O(1/VR).
2+1 with Ao 0, SHB converges
Next, we present the theoretical guarantees for MomAdaSPS.
8

---

Published as a conference paper at ICLR 2025
Theorem 3.7. Assume that each Ji is convex and Li-smooth. Then, the iterates of SHB with MomAdaSPS and 3 E [0, 1) converge as
7² TO
+
T
VT
where H² Er= x', x T= 8/ (((")) f(x)
(1+0)cLD²
1 2c
and D² max || x²I|-
1-9
()
As in our previous theorems, when there is no momentum (8 = 0). Theorem 3.7 recovers the con- vergence guarantees of AdaSPS from Jiang and Stich (2023). Let us highlight that our convergence guarantees of SHB with MomAdaSPS offers the first robust step-size selection for SHB, in the sense that it can automatically adapt to the optimization setting. More specifically, for no interpolation, our result has rate of O(I/VT), which matches the best-known rate for SHB while we assume interpolation (a² 0), then we are able to achieve rate of O(1TT) which matches the best-known rate for SHB in the convex setting (see Liu et al. (2020) and Theorem 3.2). Furthermore, under interpolation, Theorem 3.7 improves the convergence of Corollary 3.4. In particular, when o² 0, Theorem 3.7 reaches the rate (1/T") for any E [0, 1), while in Corollary 3.4 this is possible only under tighter restrictions B.
4 NUMERICAL EXPERIMENTS
In this section, we test our proposed algorithms in deterministic and stochastic convex problems as well as in training popular deep neural networks (DNNs). Our experiments are designed to highlight the benefits of our momentum variants over the vanilla (no-momentum) SPS.
Deterministic Setting. For the de- terministic setting, focus the least squares problem. The loss func- tion in this is given by f(x)= :IAA - bll?, where A E Rxx and b E R In experiments, we fol- low the setting of Wang et al. (2023) and choose n matrix A has been generated accord- ing Lenard and Minkoff (1984)
1000 while the
Audtacc All
ALTMAG
such that the condition number of AT A 10a. We test the following algorithms: Gradient Descent (GD), Nesterov's Accelerated Gradient De- scent (AGD) (Nesterov, 2018), 8), Gra- dient Descent with Polyak's 's step-size (Polyak, 1987), Adaptive Heavy Ball (AHB) (Saab Jr al., 2022). Polyak's 's Heavy Ball (HB) (Polyak, 1964), Accelerated Gradient Method (AGM) (Barré et al., 2020), ALR-MAG (Wang al., 2023) and finally our step-size from Corollary 3.3, called (MomPSmax). which guarantees O(1/T) rate. For the least squares problem, there is an established theory from Polyak (1964) that guarantees acceleration of HB with optimal step-size and momentum coefficient given by /VVL + vn)² and 7* = + //*)/L. For ALR-MAG and SHB with MomPS max. we have used the optimal momentum 9 for direct comparison with optimal HB. More- over, for GD Polyak, ALR-MAG, and MomPS we have used the precomputed f(x") (via GD), and for MomPS we have chosen 7 = 100 (for other choices, see Appendix F.4). The results Figure 2 show that MomPSmax has performance comparable to AGM and ALR-MAG which are faster than vanilla Polyak step-size after the initial iterations. Furthermore, to explore the behavior of the various adaptive step-sizes, we plot the Polyak, the ALR-SMAG, the MomPS and the optimal HB step-sizes. We can see that the original Polyak step-size oscillates around the optimal HB step-size while MomPS converges from below to the optimal.
1500
Figure 2: Comparison of deterministic algorithms for least
squares. Left: Relative error, Right: Step-sizes.
3* (VL
Stochastic Setting: Convex Problems. Here, compare step-sizes with previous works in the stochastic setting. As noticed initially for SPS (Loizou et al., 2021), the value of the upper bound 7o that results in good convergence depends the problem and requires careful parameter

---

Published as a conference paper at ICLR 2025
tuning. To alleviate this problem, Loizou et al. (2021) uses smoothing procedure that prevents large fluctuations in the step-size across iterations. That is, 7 = - for each iteration & where r 2, b is the batch-size, and n is the number of examples. We use the same smoothing trick for 76 in the implementation of the proposed methods. Similar smoothing procedures have been used in Tan et al. (2016); Vaswani et al. (2019b).
We considered multi-class logistic regression applied to commonly used benchmark datasets from the LIBSVM repository (Chang and Lin, 2011). We separate our experiments between two classes of step-sizes: non-decreasing step-sizes (which might guarantee convergence the solution) and decreasing step-sizes (which guarantee convergence to the exact solution). For the first class of step-sizes
neighborhood of
test the following algorithms: SGD (Nesterov, 2018), SHB with the constant step-size and /3 = 0.9, ADAM as described in Kingma and Ba 5), the warm-up version of ALR-SMAC with the hyper-parameters suggested in Wang et al. (2023), the SPS the SPS with naive momentum as described in Section 2.2 and our proposed MomSPS max- For the Polyak- based step-sizes such as ALR-SMAG, SPS and MomSPS we select E and 1. For SHB, ALR- All the convex experiments are run for 100 epochs and for trials. We plot the average of the trials and the standard deviation. Since there are no standard train/test splits, and due to the small sizes of the datasets, we present training loss and accuracy curves only. We present the comparison of these methods in Figure 3. In all experiments, we observe that the performance of proposed MomSPS max significantly outperforms the other methods in both training loss and test accuracy. For the decreasing variants, we test AdaGrad-Norm (Duchi al., 2011), SGD with DecSPS, SGD with AdaSPS, and our proposed SHB with MomDecSPS and SHB with MomAdaSPS. We present the outcome of this comparison in Figure 4. MomDecSPS and MomAdaSPS outperform their no- momentum counterparts in both training loss and test accuracy.
MomSPS SPS with naive momentum, we use momentum / = 0.9.
Si6 Mivortear
Aidds Ner dccass
Miceiciorua
Pretectss
Figure 3: LibSVM datasct: vowel, Batch size: Figure 4: LibSVM dataset: letter, Batch size: 52
1500
MomSPS for DNNs. In our final experiment, we go beyond the convex setting of conver- gence guarantees, and we test MomSPS on the training of DNNs. We consider non-convex minimization for multi-class classification using deep network models the CIFAR 10 and CI- FAR 100 datasets (Krizhevsky et al., 2009). We use two standard image-classffication architectures: ResNet18 and ResNet34, (He al., 2016). For space concerns, we report only the ResNet34 exper- iments in the main paper and relegate the ResNet18 to the Appendix. In these experiments we run MomSPS max with c 0.4. We present the results for CIFAR 10 in Figure 5 and for CIFAR 100 in Figure 6. We observe that MomSPS consistently outperforms its no-momentum counterpart SPS max and has competitive generalization performance compared to other popular optimizers.
2.0
E.50
381
B.A. Ite
Clous
Offnce
Offnuce Siie Momenttin
189
180
20
50
Figure 5: Resnet 34 on CIFAR 10
Figure 6: ResNet 34 on CIFAR 100
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Farshed Abdukhakimov, Chulu Xiang, Dmitry Kamzolov, Robert Gower, and Martin Takáč. Sania: Polyak-type optimization framework leads to scale invariant stochastic algorithms. arXiv preprint arXiv:2312. 173669, 2023.
Hilal Asi and John Duchi. The importance of better models in stochastic optimization. Proceed- ings of the National Academy of Sciences, (16(46):22924-2299) 2019a.
Hilal Asi and John Duchi. Stochastic (approximate) proximal point methods: Convergence, opti- mality, and adaptivity. SIAM Journal on Optimization, 29(3):2257-2290, 2019b.
Mahmoud Assran, Nicolas Loizou, Nicolas Ballas, and Mike Rabbat. Stochastic gradient push for distributed deep learning. In International Conference on Machine Learning, 2019.
Mathieu Barré, Adrien Taylor, and Alexandre d' Aspremont. Complexity guarantees for polyak steps with momentum. In COLT, 2020.
Leonard Berrada, Andrew Zisserman, and M Pawan Kumar. Training neural networks for and by interpolation. In ICML,, 2020.
Dimitri Bertsekas and John N Tsitsiklis. Neuro-dynamie Programming. Athena Scientific, 1996.
Stephen Boyd, Lin Xiao, and Almir Mutapcic. Subgradient methods. Lecture Notes of EE392o, Stanford University, Autumn Quarter, 2003.
Karan Chadha, Gary Cheng, and John Duchi. Accelerated, optimal and parallel: Some results on model-based stochastic optimization. In ICML, 2022.
Chih- Salh-Chung Chang and Chih-Jen Lin. Libsvm: library for support vector machines. ACM Trans- actions Intelligent Systems and Technology, 2-3): 2011.
Sayantan Choudhury, Eduard Gorbunov, and Nicolas Loizou. Single-call stochastic extragradient methods for structured non-monotone variational inequalities: Improved analysis under weaker conditions. NeurIPS, 2024a.
Sayantan Choudhury, Nazarii Tupitsa, Nicolas Loizou, Samuel Horvath, Martin Takac, and Eduard Gorbunov. Remove that square root: new efficient scale-invariant version of adagrad. arXiv preprint arXiv:2403.026648 2024b.
Damek Davis, Dmitriy Drusvyatskiy, Kellie MacPhee, and Courtney Paquette. Subgradient meth- ods for sharp weakly convex functions. Journal of Optimization Theory and Applications, 179 (3):962-998) 2018
Aaron Defazio and Samy Jelassi, A momentumized, adaptive, dual averaged gradient method.
Journal of Machine Learning Research, 23(144):1-34 2022.
Alexandre Défossez, Léon Bottou, Francis R. Bach, and Nicolas Usunier. A simple convergence proof of adam and adagrad. Transactions on Machine Learning Research, 2022, 2022.
Ryan O'Orazio, Nicolas Loizou, Issam H Laradji, and loannis Mitliagkas. Stochastic mirror descent: Convergence analysis and adaptive variants via the mirror stochastic polyak stepsize. Transactions on Machine Learning Research, 2023.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(7), 2011.
Guillaume Garrigos and Robert M Gower. Handbook of convergence theorems for (stochastic) gradient methods. arXiv preprint arXiv:2301.11233 2023.
Guillaume Garrigos, Robert M Gower, and Fabian Schaipp. Function value learning: Adap- tive learning rates based on the polyak stepsize and function splitting in erm. arXiv preprint arXiv:2307.14522 2023.
11

---

Published as a conference paper at ICLR 2025
Euhanna Ghadimi, Hamid Reza Feyzmahdavian, and Mikael Johansson. Global convergence of the heavy-ball method for convex optimization. European Control Conference (ECC), 2015.
Igor Gitman, Hunter Lang, Pengchuan Zhang, and Lin Xiao. Understanding the role of momentum in stochastic gradient methods. In NewrIPS, 2019.
Eduard Gorbunov, Hugo Berard, Gauthier Gidel, and Nicolas Loizou. Stochastic extragradient: General analysis and improved rates. In AISTA7S, 2022.
Robert Gower, Othmane Sebbouh, and Nicolas Loizou. SGD for structured nonconvex functions: Learning rates, minibatching and interpolation. In AISTATS, 2021.
Robert M Gower, Mathieu Blondel, Nidham Gazagnadou, and Fabian Pedregosa. Cutting some slack for SGD with adaptive polyak stepsizes. arXiv preprint arXiv:2202. 12328, 2022.
Robert Mansel Gower, Nicolas Loizou, Xun Qian, Alibek Sailanbayex, Egor Shulgin, and Peter Richtárik. SGD: General analysis and improved rates. ICMI., 2019.
Moritz Hardt, Benjamin Recht, and Yoram Singer. Train faster, generalize better: Stability of stochastic gradient descent. In ICML,, 2016.
Trevor Hastie, Robert Tibshirani, and Jerome H Friedman. The Elements of Statistical Learning: Data Mining, Inference, and Prediction, volume 2. Springer, 2009.
Elad Hazan and Sham Kakade. Revisiting the polyak step size. arXiv preprint 33 2019.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. CVPR, 2016.
Geoffrey Hinton, Nitish Srivastava, and Kevin Swersky. Neural networks for machine learning: Lecture 6a overview of mini-batch gradient descent. Lecture Notes, page 2012.
Xiaowen Jiang and Sebastian U Stich. Adaptive SGD with polyak stepsize and line-search: Robust convergence and variance reduction. In NewrlPS, 2023.
Ahmed Khaled, Othmane Sebbouh. Nicolas Loizou, Robert Gower, and Peter Richtárik. Unified analysis of stochastic gradient methods for composite convex and smooth optimization. Journal of Optimization Theory and Applications, 199(2):499-544 2023.
Rahul Kidambi, Praneeth Netrapalli, Prateek Jain, and Sham Kakade. On the insufficiency of ex- isting momentum schemes for stochastic optimization. In Information Theory and Applications Workshop (ITA), pages 1-9. IEEE, 2018.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR, 2015.
Anastasia Koloskova, Nicolas Loizou, Sadra Boreiri, Martin Jaggi, and Sebastian Stich. A unified theory of decentralized sgd with changing topology and local updates. In ICML, 2020.
Alex Krizhevsky al. Learning Multiple Layers of Features from Tiny Images. Technical report, University of Toronto, 2009.
Melanie L Lenard and Michael Minkoff. Randomly generated test problems for positive definite quadratic programming. ACM Transactions on Mathematical Software (TOMS), 10(1):86-96, 1984.
Shuang Li, William J. Swartworth, Martin Takác, Deanna Needell, and Robert M. Gower. SP2 A second order stochastic polyak method. In ICLR, 2023.
Xiaoyu Li and Francesco Orabona. On the convergence of stochastic gradient descent with adaptive stepsizes. In AISTATS, 2019.
Tengyuan Liang and Alexander Rakhlin. Just iterpolate: Kernel "ridgeless" regression can general- ize. The Annals of Statistics, 122-)1133911347, 2020.
12

---

Published as a conference paper at ICLR 2025
Yanli Liu, Yuan Gao, and Wotao Yin. An improved analysis of stochastic gradient descent with momentum. In NeurIPS, 2020.
Nicolas Loizou and Peter Richtárik. Momentum and stochastic momentum for stochastic gradi- ent, newton, proximal point and subspace descent methods. Computational Optimization and Applications, 77(3):653-710, 2020.
Nicolas Loizou, Hugo Berard, Alexia Dolicoeur-Martinee Pascal Vincent, Simon Lacoste-Julien, and Ioannis Mitliagkas. Stochastic hamiltonian gradient methods for smooth games. In ICML, 2020.
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 AISTATS, 2021.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization, In ICLR, 2019.
Jerry Ma and Denis Yarats. Qussi-hyperbolic momentum and adam for deep learning. In ICLR, 2019.
Siyuan Ma, Raef Bassily, and Mikhail Belkin. The power of interpolation: Understanding the effectiveness of SGD in modern per-parametrized learning. In ICML, 2018.
Sohom Mukherjee, Nicolas Loizou, and Sebastian U Stich. Locally adaptive federated learning via stochastic polyak stepsizes. arXiv preprint urXiv:2307.06304 2023.
D. Needell, N. Srebro, and R. Ward. Stochastic gradient descent, weighted sampling, and the ran- domized kaczmarz algorithm. Mathematical Programming, Series A, 155(1):549-573, 2016.
Arkadi Nemirovski and David B. Yudin. Problem Complexity and Method Efficiency in Optimiza- tion. Wiley Interscience, 1983.
Arkadi Nemirovski, Anatoli Juditsky, Guanghui Lan, and Alexander Shapiro. Robust stochastic approximation approach to stochastic programming. SIAM Journal on Optimization, 19(4):1574- 1609, 2009.
Yurii Nesterov. A method of solving a convex programming problem with convergence rate ((/k?²) In Doklady Akademii Nauk, volume 269, pages 543 547. Russian Academy of Sciences, 1983.
Yurii Nesterov. Lectures on Convex Optimization, volume 137. Springer, 2018.
Lam Nguyen, Phuong Ha Nguyen, Marten Dijk, Peter Richtárik, Katya Scheinberg, and Martin Takác. Sgd and hogwild! convergence without the bounded gradients assumption. ICML, 2018.
Antonio Orvieto, Simon Lacoste-Julien, and Nicolas Loizou. Dynamics of SGD with stochastic polyak stepsizes: Truly adaptive variants and convergence exact solution. In NeurIPS, 2022.
Boris T Polyak. Some methods of speeding up the convergence of iteration methods. USSR Com- putational Mathematics and Mathematical Physics, 4(5):1-17, 1964.
Boris T Polyak. Introduction to Optimization. New York, Optimization Software, 1987.
Mariana Oliveira Prazeres and Adam M. Oberman. Stochastic gradient descent with polyak's 's learn- ing rate. Journal Scientific Computing. 89(1):25, 2021.
Alexander Rakhlin, Ohad Shamir, and Karthik Sridharan. Making gradient descent optimal for strongly convex stochastic optimization. In ICML, 2012.
Benjamin Recht, Christopher Re, Stephen Wright, and Feng Niu. Hogwild!: A lock-free approach to parallelizing stochastic gradient descent. In NeurlPS, 2011.
Sashank J. Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence adam and beyond. In ICLR, 2018.
13

---

Published as a conference paper at ICLR 2025
Herbert Robbins and Sutton Monro. A stochastic approximation method. The Annals of Mathemat- ical Statistics, pages 400-407, 1951.
Michal Rolinek and Georg Martius. L4: Practical loss-based stepsize adaptation for deep learning. In NeurlPS, 2018
Samer Saab Jr, Shashi Phoha, Minghui Zhu, and Asok Ray. An adaptive polyak heavy-ball method. Machine Learning, ((1)))324--377) 2022.
Fabian Schaipp, Robert M. Gower, and Michael Ulbrich. A stochastic proximal polyak step size. Transactions Machine Learning Research, 2023a.
Fabian Schaipp, Ruben Ohana, Michael Eickenberg, Aaron Defazio, and Robert M Gower. Momo: Momentum models for adaptive learning rates. arXiv preprint arXiv:2305.077 07538, 2023b.
Mark Schmidt and Nicolas Le Roux. Fast convergence of stochastic gradient descent under a strong growth condition. arXiv preprint 1rXiv:1308. 6370, 2013.
Othmane Sebbouh, Robert M Gower, and Aaron Defazio. Almost sure convergence rates for stochas- tic gradient descent and stochastic heavy ball. In COLT, 2021.
Ohad Shamir and Tong Zhang. Stochastic gradient descent for non-smooth optimization: Conver- gence results and optimal averaging schemes. In ICML, 2013.
Zheng Shi, Abdurakhmon Sadiev, Nicolas Loizou, Peter Richtárik, and Martin Takáč. Ai-sarah: Adaptive and implicit stochastic recursive gradient methods. Transactions on Machine Learning Research, 2022.
Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initializa- tion and momentum in deep learning. In ICML, 2013.
Conghui Tan, Shiqian Ma. Yu-Hong Dai, and Yuqiu Qian. Murzilai-borweit step size for stochastic gradient descent. In NeurlPS, 2016.
Sharan Vaswani, Francis Bach, and Mark Schmidt. Fast and faster convergence of SGD for over- parameterized models and an accelerated perceptron. In AISTATS, 2019a.
Sharan Vaswani, Aaron Mishkin, Issam Laradji, Mark Schmidt, Gauthier Gidel, and Simon Lacoste- Julien. Painless stochastic gradient: Interpolation, line-search, and convergence rates. In NeurlPS, 2019b.
Sharan Vaswani, Issam Laradji, Frederik Kunstner, Si Yi Meng, Mark Schmidt, and Simon Lacoste- Julien. Adaptive gradient methods converge faster with over-parameterization (but you should do line-search). arXiv preprint arXiv:2006.0000 06835, 2020.
Jun-Kun Wang, Chi-Heng Lin, Andre Wibisono, and Bin Hu. Provable acceleration of heavy ball beyond quadratics for a class of polyak-lojasiewiez functions when the non-convexity is averaged- out. In ICML, 2022.
Xiaoyu Wang, Mikael Johansson, and Tong Zhang. Generalized polyak step size for first order optimization with momentum. In ICMI,, 2023.
Rachel Ward, Xiaoxia Wu, and Leon Bottou. Adagrad stepsizes: Sharp convergence over nonconvex landscapes. The Journal of Machine Learning Research, 21(1)99047-9076 2020.
Yan Yan, Tianbao Yang, Zhe Li, Qihang Lin, and Yi Yang. A unified analysis of stochastic mo- mentum methods for deep learning. In Proceedings of the Twenth-Seventh International Joint Conference on Artificial Intelligence, IJCAI, 2018
Yun Zeng, Deren Han, Yansheng Su, and Jiaxin Xic. On adaptive stochastic heavy ball momentum for solving linear systems. arXiv preprint arXiv:2305.05482 2023.
Chiyuan Zhang, Samy Bengio, Moritz Hardt. Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM, 64(3):107- 115, 2021.
14

---

Published as a conference paper at ICLR 2025
Supplementary Material
The Supplementary Material is organized as follows: Appendix A, we have more details on related work adaptive and momentum methods. In Appendix B, include the main pseudo- codes for our algorithms. In Appendix C, we give the basic definitions and lemmas as well as the basic theory of IMA. Appendix D presents the proofs of the theoretical guarantees from the main paper. In Appendix F, we describe in detail our experimental setup and provide additional experiments.
CONTENTS
1 Introduction
Main Contributions
Exploring the Interplay of SPS and Heavy Ball Momentum
2.1 Background on Stochastic Polyak Step-size
2.2 Naive SPS in SHB
4
2.3 Iterative Moving Average: Balance between SPS and Momentum
5
3
Convergence Analysis
6
3.1 Convergence neighborhood of the solution
3.2 Convergence the exact solution
Numerical Experiments
Further Related Work
17
Pseudo-codes
19
Technical Preliminaries
C.I Basic Definitions
21
C.2 Basic Lemmas
21
C.3 Iterate Moving Average
22
C.4 Connection of MomSPS and SPS
23
Proofs of the main results
25
D.I Proof of Theorem 3.2
25
D.2 Proof of Theorem 3.6
28
D.3 Proof of Theorem 3.7
Beyond the bounded iterates assumption
33
Additional Experiments
35
E.1 Extra Deterministic Experiment
35
F2 Stochastic Heavy Ball with Constant Step-size
35
E3 Other Choices
35
F4 Is % needed?
37
15

---

Published as a conference paper at ICLR 2025
E5
Plots for C1 and C2
38
F6
More convex experiments and parameter settings
39
More deep learning experiments and parameter settings
39
16

---

Published as a conference paper at ICLR 2025
A FURTHER RELATED WORK
Adaptive methods. Two of the first adaptive algorithms are AdaGrad (Duchi et al., 2011) and RMSProp (Hinton et al., 2012). The very popular algorithm Adam (Kingma and Ba, 2015) was in- troduced as momentum extension of RMSProp, and its more recent weight-decay variant Adam W (Loshchilov and Hutter, 2019) is the de facto optimizer for deep neural networks. Investigating the convergence guarantees of adaptive methods across various settings continues to be an active area of research (Vaswani et al., 2020; Ward et al., 2020; Li and Orabona, 2019; Shi et al., 2022; Défossez et al., 2022; Choudhury et al., 2024b; Defazio and Jelassi, 2022).
More recently, a new line of work for adaptive algorithms has appeared, inspired by Polyak step- sizes. Some attempts for efficiently generalizing the Polyak step-size from the deterministic setting to the stochastic were made in Rolinek and Martius 20; Prazeres and Oberman (2021); Berrada et al. (2020). Loizou al. (2021) was the first work proposing SPS and providing strong conver- gence guarantees in different settings, including strongly convex, convex, and non-convex functions. Further extensions with decreasing variants were proposed in Orvieto et al. (2022); Jiang and Stich (2023). Many recent works propose SPS-type update rules for solving optimization problems in different settings, For example, extensions of SPS to proximal setting analyzed in Schaipp et al. (2023a), SPS variants for mirror descent presented in D'Orazio et al. (2023), and SPS for federated learning proposed in Mukherjee et al. (2023). There are also strong connections between SPS and Model-Basee -iaeed approaches (Asi and Duchi, 20;9;;;; Chadha et al., 2022). For further results related to SPS, see Gower et al. (2022); Li et al. (2023); Garrigos et al. (2023) and Abdukhakimov et al. (2023).
Momentum with constant step-size. Polyak (1964) introduced the concept of momentum for the Gradient Descent (GD) method. He showed that for strongly convex quadratic problems. the momentum method provably accelerates GD with optimal momentum coefficient 3 = (VI /m)²/ + vT)² and 7* = -1 + / where μs is the strong convexity constant and L the smoothness constant. Ghadimi et al. (2015) proved a global sublinear convergence guarantee for HB for convex and smooth functions and global linear convergence when the function is also smooth and strongly convex. Moreover, Nesterov (1983) introduced Nesterov's Accelerated Gradient Descent (AGD) method, where he shows acceleration over GD for strongly convex and smooth as well as convex and smooth objectives. Note that both of these methods require knowledge of both μ and L. In the stochastic setting, Liu et al. (2020) provides new convergence guarantees for SHB with constant step-size. In particular, it shows that for both strongly convex and non-convex objectives, SHB enjoys the same convergence bound as SGD. Sebbouh et al. (2021), among other things, shows that in the smooth and convex setting SHB converges in expectation at the last iterate with a rate of O1/TT) to a neighborhood of the minimum and at a Olog / T/VT) rate to the minimum exactly. For these guarantees, one needs knowledge of the problem parameters (such the smoothness constant), and their results hold for specific selection of the momentum parameter. Momentum has also been successfully applied in linear systems, see Loizou and Richtárik (2020) where it provides convergence guarantees of SGD, stochastic Newton, stochastic proximal point, and stochastic dual subspace ascent with momentum for consistent linear systems.
Adaptive Momentum. In this paragraph, we review the bibliography for adaptive methods with momentum. Momentum can be either SHB or Nesterov's acceleration, and in this case, the adaptiv- ity can be either the step-size or the momentum coefficient. In the deterministic setting, Barré et al. (2020) proposed the Accelerated Gradient Method (AGM), an adaptive algorithm built upon AGD that guarantees acceleration over GD. AGM approximates the strong convexity constant μ by the inverse of the classical Polyak Step-size. However, their algorithm still requires the knowledge of L (but not μ1). Furthermore, Saab Jr et al. (2022) suggests an Adaptive Heavy Ball (AHB) algorithm where it approximates the constants l and L iteratively in each step based on the formula of the optimal constants of HB, but it does not show any acceleration over HB or GD. In the stochastic set- ting, Wang et al. (2023) propose adaptive algorithm for a moving averaged gradient momentum, named ALR-SMAG. It shows that ALR-SMAG enjoys linear convergence rate for semi-strongly convex and smooth functions. Schaipp al. (2023b) proposes new adaptive learning rate that can be combined with any momentum- based method. Finally, Zeng et al. (2023) proposes adaptive variants for SHB in linear systems solvers
17

---

Published as a conference paper at ICLR 2025
On Technical Assumptions for Convergence. In the literature of stochastic optimization prob- lems, there are several assumptions on the noise of the stochastic estimators that typically are made on top of smoothness and convexity to prove the convergence of stochastic optimization algorithms.
For example, many works (Recht et al., 2011; Rakhlin et al., 2012; Shamir and Zhang, 2013; Nguyen et al., 2018) assume bounded gradients, i.e., that is, there exists M R such that E V fx) M. While this might look like a natural assumption, in the unconstrained setting, it contradicts the assumption of strong convexity leading to convergence guarantees that hold for an empty set of problems (Nguyen et al., 2018; Gower et al., 2019; 2021). A much more relaxed assumption used in the literature is the growth condition the stochastic gradients, (Bertsekas and Tsitsiklis, 1996; Schmidt and Roux, 2013; Vaswani et al., 2019a). It states that there exist constants E R and e I such that E || V f=())1 S f(=] Based on strong growth condition (6 = 0), Schmidt and Roux (2013) were the first to establish linear convergence of SGD, with Vaswani et al. (2019a) later showing that SGD can find first-order stationary point as efficiently as full gradient descent non- convex settings. Similar conditions have also been proved and used in the analysis of decentralized variants SGD (Koloskova al., 2020; Assran et al., 2019). More recently, a line of work that uses smoothness (via expected smoothness/residual conditions) to provide closed-form expressions for the values of p and of growth condition was able to provide tight convergence analysis of several stochastic algorithms, including SGD (Gower et al., 2019; 2021), variance reduced methods (Khaled et al., 2023), stochastic algorithms for min-max optimization (Loizou et al., 2021; 2020; Gorbunov al., 2022; Choudhury et al., 2024a).
In the convex regime, the original analysis SGD with SPS of Loizou et al. (2021) was one of the first papers that did not require any additional assumptions to guarantee convergence for SGD. Fol- lowing the convergence guarantees of Loizou et al. (2021), we highlight that our proposed analysis of SHB with Polyak step-size does not require any additional assumptions for guaranteeing conver- gence. To the best of our knowledge, as mentioned in the main paper, our approach provides the first analysis of SHB without the restrictive bounded variance and growth conditions.
For proving convergence of SHB, Liu et al. (2020) makes the strong assumption of bounded variance and assumes knowledge of the smoothness L for tuning the parameters. Another work theoretical guarantees for SHB is Sebbouh et al. (2021), where last iterate convergence is proved for a specific choice of "t and 3r, but also requires knowledge of L and only achieves rate Oloog T / /vT) for the decreasing step-size variant. Furthermore, the authors of Wang et al. (2022) propose Polyak-based stepsize for a moving averaged gradient momentum variant, different from SHB, and provide guar- antees for the deterministic and stochastic setting for strongly convex objectives. Finally, Schaipp et al. (2023b) assumes stochastic convex and interpolated regime for their results.
18

---

Published as a conference paper at ICLR 2025
B PSEUDO-CODES
In this section, we include the pseudo-codes of the SHB with the proposed stochastic Polyak step- sizes for better understanding of the methods and easier comparison with other algorithms.
SHB with MomSPSmax. We start with the pseudo-code for SHB with MomSPS max. In The- orem 3.2 we established convergence to neighborhood for B [ where a
ميا - and for c 1. However, in practice, it seems to have convergence even for different choices of even if 3 is chosen outside of this bound, for a general 3 E [0, 1) and c. Thus, the user must provide the momentum coefficient 33, the upper bound 76. the constant c, and the lower bounds fz. In practice, the usual choices are 3 0.9, = L, and Es. = O.
min 2L
provide the pseudo-code
Algorithm SHB with MomSPS
1: Parameters: 3 E [0, 1), >> O, 0, f3, lower bounds
2: Initialization: _0 E Rd x- x00
3: for = 0, 2.... do
4: Choose uniformly at random S. S nI,...,n
5: Y1 (1 -1-18 3) min Js, (='))(s, ds, (2)))22
6:
x+41 = x² Ss (x*) B(x² 11
7:
end for
SHB with MomDecSPS. For decreasing step-sizes, Theorems 3.6 and 3.7 make assumption on E [0, 1). For SHB with MomDecSPS, the user needs to provide the momentum coefficient 3, the upper bound n6 the constant and the lower bounds fs., like the previous algorithm. In practice, the usual choices are e and f.s = 0. Note that when = 0, the step-size is equal to
%o = min (1 3fss (I f',
T6e
fs. (2 Es
)
(1 B) min
cllV/ Ss. (x") 7b3.
V (x:) 112
which is equal to MomSPS max for =
Algorithm 2 SHB with MonDecSPS
1: Parameters: 3 [0,1), = 76 0 bounds
for 20, 0, E3, lower
2: Initialization: x0 E Rdd x- = σ
3: for t. 0, do
4: Choose uniformly at random S₁ S {I,...,,.) n}
(1 (=) ---10-11 [a,
5: 7tt = min
6:
r11 xt
Jt.
V.s (x*) + 3(x² x0 -11
7:
end for
SHB with MomAdaSPS. Finally, for SHB with MomAdaSPS the requirements are less since the user only needs to provide the momentum coefficient B3, the constant and the lower bounds f Furthermore, according Jiang and Stich (2023) can set [=)) after randomly
choosing So S [n] in the first iteration. Moreover, the minimum with respect to 7 is to ensure that the step-size is decreasing.
19

---

Published as a conference paper at ICLR 2025
Algorithm 3 SHB with MomAdaSPS
1: Parameters: 3 [0, 1), 7-1 t00, 0, f3 lower bounds
2: Initialization: x0 E Rd x= = x0
3: for t = 1, do
4: Choose uniformly at random St E fI,....,1) mn (x'
5: 7t min
(x')), Eizu (x-) 7T-1
6: x141 x² SS. (x 3(x² =x
7:
end for
Note that for all the above algorithms, one needs only lower bounds Es,, which can be chosen equal to f.. 0 in practice. Of course, if the true infima fú, inf fs, are known, then one should use them because will lead to better convergence in terms of the neighborhood. Furthermore, notice that start by initializing x e-
æº, where xº is the given starting point. This is because the first step of SHB (i.e. for f = 0) is just an iteration of SGD because we have no previous iterate
x-1 that is x! x 7o V Iso (x*) B(x 1) x 7o V VSs (x²).
20

---

Published as a conference paper at ICLR 2025
C TECHNICAL PRELIMINARIES
C.I BASIC DEFINITIONS
In this section, we present some basic definitions use throughout the paper.
Definition C.1 (Convexity). A differentiable function f IR + I. is convex if f(x) 2 f(y) + (Vf(()) y).
for all x; y Rdd
Definition C.2 (L-smooth). differentiable function F IR is L-smooth if there exists constant such that
[V]())-V/([) - S L|aa - -yl,
or equivalently
If(z) - f(v) - (Vf(ww -- 4>1 _ 1| -111
L
2
for all v. y E Rdd
C.2
BASIC LEMMAS
Here have gathered useful lemmas for various Polyak related step-sizes. These lemmas are used repeatedly in the proofs of the main results.
Lemma C.3 ((Loizou et al., 2021)). Suppose each f; is L4-smooth, then the step-size of SPS satisfies:
a min 2eImmxx 7o S 71 S To and 7f V.s )1 s mnlss,('") - E.y,).
(2)
where Lmax maxie[n] {Li).
Lemma C.4 (Lemma 1, (Orvieto et al., 2022)). Suppose each Ji is Li-smooth, then the step-size of DecSPS satisfies
min
12 Ct
co7b
7t
Co
and 7- S 7/t and Vs.
< 7t [ss. (x²)
C
a
(3)
where Lmax maxie (n) {L,).
Lemma C.5 (Lemma 12, (Jiang and Stich, 2023)). For any non-negative sequence ....... OT the following holds:
T
T
04
<
2
Σ
04
(4)
1=0
1=0
Es=0
t=0
If ag 2 1, then the following holds:
T
at
< log
Σ at
+
(5)
Σ1=
21

---

Published as a conference paper at ICLR 2025
Lemma C.6 (Lemma 13, (Jiang and Stich, 2023)). If I a(x + b) 2 and 0, then holds that a S a + vbb.
Lemma C.7 (Lemma 16, (Jiang and Stich, 2023)). Suppose each fi is Lm-oth,th, then the step- size of AdaSPS satisfies
T1 Th S fs, (x²) - f.s.
(6)
2cLmax v V (x²( E's
cl Vs.. (x')11
V V fs,(x*) C...
-
where T.max maxie(n) {1,}-
C.3 ITERATE MOVING AVERAGE
As mentioned in the main paper, the iterates of the SHB method are equal to the 2.3 method. The following proposition is from Sebbouh al. (2021). We include the proof for completeness.
Proposition C.8 (Proposition 1.6, (Sebbouh et al., 2021)). Let 1ht > O and x 2 0. Consider the Iterate Moving Average (IMA) method given by the following update rule:
x11 2' 111 (s.
(7)
x 1+
d4t1
D' +
1
24
(8)
Aet1
+
d4t1 +
where 2 = x.. If the following equations hold for any L E N
1 + det1 =
(9)
Be
T) (1 \et1)a
(10)
then the xa iterates of the 2.3 method are equal to the iterates produced by the SHB method.
Proof. We will start from the 2.3 update rule and will prove that its (x*) iterates satisfy the SHB update rule. Suppose that we have the (z++1) and (x*+1) iterates as well as the (Ax) and (17) sequences. Define the sequences (3.) and (1x) using eqs. and (10). We will show that (x*) iterates satisfy the SHB update rule. Solving for 2 2+1 in eq. (8) we get
++ ++t + 1 =²')
(11)
Now substituting eqs. (7) and (11) in (8) we have
A4t1
1
x. +
1-1
d411 +
A+1
+1
(
det1
(2" T1 V... (x,)))))
Ac+1 +
Ae+1 +
(11) Aet1
1
I²
(s' Ae -xt-1) m1 Vss, (x¹)))
A1t1 +
d
fis (x') + X4o+1 + (x² (x'- -1--1)
A +
(9.10) L V. (x') +B((x' - -1
as wanted.
The above proposition essentially states that momentum is convex combination of SGD under a certain transformation. this paper, we work only with a constant momentum coefficient, i.c., B B. For this case, also assume constant A and have Proposition 2.1 and Corollary 2.2.
22

---

Published as a conference paper at ICLR 2025
Proof of Proposition 2.1. Setting X = X in Equation (9) we have
A
X
7
Br
=
=:
8.
51
1+
Then solving for À we get X =
0 1-3
Hence
11s
78 1 +
+ I-3
= (1 B)mT
Proof of Corollary 2.2.
Let fpt SPS Then
Tt (1 3)
fs, (x¹) - E
(1 3) min
cl [s. (x²)| 1/b
Sss (x²) f'ss
= (1 3) min
cl f., (z²) 76
Here %/ 176.
Let 77 DecSPS. Then
Tt (1 - B)
fs. (x')
Es. 19-10-11
= (1 - B) min
cell V Ss, (xt) 12
Ct
(1
B)(fs, (x*)
(1
3)
min
c. <1| | fs, (x²)||
(1 B)(fs. (x²)
-Gs.] 1-
min
CsIVSs, (1²)
Ct
since N-11 (1 B)m from Proposition 2.1.
Let 1/1 AdaSPS. Then
7. = (1 - -B)
fs. (x) E3.
= (1 - 8) min
cl fs. (x*)) v E'=o (s) - -
(1 - B)ss, (x*)) -[3,]
min
(1 .3 33000-2
c| cl fs. fx (22
E'= fs, (xx) - -
(1 B)ss) (x²) E
min
---11
cl Vss E' Ss, (x*) - -P'.
C.4 CONNECTION OF MOMSPS max AND SPS
max
In this section, we explore the connection between the MomSPS and the SPS step-sizes. In particular, view the MomSPS step-size as the SPS step-size with 1/(1
23

---

Published as a conference paper at ICLR 2025
8). Here, explain why these are two different step-sizes. Recall from Proposition 2.1 and Corollary 2.2 that our proposed step-sizes for the SHB setting are scaled versions of Polyak step- sizes for SGD. In particular, let's assume that
Ss. (a') fs.
SPSmaz min
SpSaa
πt
CSPSmar
Vis, (x+)]12
and
fs. (2') f3.
min
M MomSPSmer
CMomSPSmaz V fss, (xt) (2
we have PSPSaa
MMos if and only if
1-
and 76
SPSmaz = (1 B)7b MoMSPSPPaa
CMomSPSmar CarSmar
Without loss of generalization, we can assume that CMomSPSmes (this also what Theo- rem 3.2 guarantees). In this case we get f = carSmar
Hence, we can restate Theorem 3.2 as
follows:
Theorem C.9. Assume that each Is is convex and L--smooth. Then, the iterates of SHB with 1 where min 2bmax 76 and Lmax
SPSmax with 1, 2n--2 ) and 270-a
maxxein] Li), converge
B[S(x)) f(x") Ci C1 x0112
+ C202
T
where H² = T T- of is the Cesaro average and the constants C1
1-8
and C2
a0+a-201
27h-a0 29--
auta-237s
This means that MomSPS max can be viewed special case of SPS SPS However, note that the update rules are different since in SHB we have the extra term +3(x" x²-1) ) so a different analysis needed.
The same line of thought can be applied to other works related to HB. For example, in Ghadimi al. (2015), the first global guarantees for HB were established. The authors show that for a convex and L-smooth function f HB with any 8 E [0. 1) and 0 < < L O(1/T). Also, recall that GD for convex and L-smooth function converges for < < Z with rate O(1/T). Hence, we can restate the result of Ghadimi al. (2015) follows: Let f be convex and L-smooth function. Then HB with y = f where c (0, 2), and B = 1 converges with rate O(1/T).
2(1-8) converges with a rate of
Similarly, for SGD in the convex and smooth setting, we are guaranteed convergence O(1/T') to neighborhood of the solution when ? S In Liu al. (2020), the authors guarantee that SHB
with any E (0, 1) and S (1= min 1-9+ 2 42" converges with rate O(1/T) to neighborhood. This can be restated as follows: SHB with 7 f. where (0, 1), and given by
(1-0)3 (1-8)²
the solution of min
converges with rate O(1/T) to neighborhood.
4-848
2 22/20420
24

---

Published as a conference paper at ICLR 2025
D PROOFS OF THE MAIN RESULTS
In this section we present the proofs of the main theoretical results presented in the main paper, i.e., the convergence analysis of SHB with MomSPS MomDecSPS and MomAdaSPS for convex and smooth functions fi and f of Problem 1. The idea for all the proof' is that we firstly show the corresponding result in the 2.3 setting, and then we transfer the SHB setting with the help of Proposition 2.1.
Compared to the analysis of SGD with SPS, the SHB update rule requires taking into consideration
extra terms related to the previous iterate x'-1. This needs new algebraic tricks in order to use convexity, and after that, one needs to deal with the two quantities f(x*) and f(x==1) at the same time. We deal with this using the 2.3 framework, similarly with Sebbouh et al. (2021). However, we highlight that we finish our proofs with a different approach. More specifically, in the proof of Theorem G.1 in Sebbouh et al. (2021) the authors choose the momentum coefficients xe 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 X This more prominent 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 E[ f((')) * f( 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 clas- sical analysis of SGD and constant step-size SHB. This makes the convex proof more challenging than the standard analysis of SHB.
D.1 PROOF OF THEOREM 3.2
First, we state and prove the corresponding theorem for SPSmax in the 2.3 setting.
Theorem D.1. Assume that each fi is convex and Lm-shh.hh. Then, the iterates of 2.3 with
n min { ss, (x²) Es mo } and d x
[o [ 2(n a)
[x')((('')1
where a min Tmmme mh } and Lmax = maxig(n] Li), converge as
1:00 - ==11|
(2nh + 2npn - 22x a))o
E[f(*T) - f(x)) (2a) + a 21/22T
2a + a 2n/b
where xJ T- x' and 0² = Elfi(z*) - fi1
Proof. We have
|| - x² 1|2
= x=11 2n, (V 2f n² V fz,((")))
(12) ||2"- 2n V
2n (1 -x-11)
+ V (12
convesity
112 212 112 21 (fs - - -fss (x (x0
+ V (s,')(1'
2 || 1011==21100
- -21 [s, (x²) [s, (x') - 2n - Ss. (crtt))
+ (x²) - -(5.]
|| - xx x=1 -2ne (x²)- [ss] + (ve(ss((x') -65,) -2mpdelSs (0²) ((s.
25

---

Published as a conference paper at ICLR 2025
+ 2ffs. (z²-1) - -fš,] + n:lfs (x') -fs.]
)
|122 - x=112= - 27/e (1+ 2/2 Ifss - [.,l- + 22 2f - [3,]]
+ 22/1/s (a*)) - f
(2)
S 11 x² e'Ip- -2a \e + 2) [Fs)(t') - fb,,]+ + 270 -
+2m fossss.. (1²) - Es11
5 $11=' - x²112- - 2ar (>r (Ai+ +2++2 1 (ss - [s((+""]- - 22 -2cr (x
1
2 [ss, [ss((²)) - - -(y.]
+ 2. - fs. }| + (m)r' - -fs,1 + 2m -//.1.
Now take expectation conditional on x.' to get
Es, || - x²112 S Es, ||2l x²11 2a \1 + 2) Ls., - ( (x')
2a (>:
1
Es, lfs. (xx) [ + 21 Us. 1fs, (xt - (x')
2)
+ 2m Es, [fs.
{s,] 1 + 2m Es, (fs, -5,,]
Ax (x++ 2 [Mxx') - [ax)|| - -2a (ar + 11)
1
|| -x1112 -
-2a -2x
2) 0²
+ def - -[(x') + +21oo10²+ + +210022
Ax (^+ 2 [f(x*) - 1f((11()) + ( - -((=))))))))(())1
11 - =x²11 -2a
+2 7b (1+ Ae) -a Ax + 2 1) 0².
(12)
Note that in eq. (12) the coefficient 2 (m (1 + A1) - a (de + 2)) of 0² is positive since nh 2 a.
Taking expectation again and using the tower property we have
1 2
2x
Av
B[s(x') - [(x)]] - 2r B[](x*1))-1 1 [(x'))
S 1112 x²1 $12 - -B 111 - x² 112 +2 1hb (1+ - a (M + 22)
) 0².
(13)
: T and telescoping we get Summing eq. (13) for 1. 0,...;
T-1
1
2
d +
2 B[J(x') - -s(x")] - Ae E|5(x*1))- - [(x01)
T
E || x 1|2 - E E1 - x x²1
+
T-1
Σ
2
16 (1+ -a
-a(0)aa222 ( (Ar+ + 2 0²
1=0
1=
T-1
17 [b+ Nx) - a ( + 0²
= E1120 x²1 EE 12
10122 + Σ
2 0z. The LHS of eq. (14) can be bounded as follows (using x = x0)
2=z00
1 1| - 2²112-
T-1
Σ 2
no (1+ - a (Ar )
(14)
1=0
+ 1 2
2a >x (Ar
Blf(x')- f(x") - 27hhAr B[f(x'-)) - --=-)))()(x)))
0=
26

---

Published as a conference paper at ICLR 2025
T-1 1
Σ
20
A +
E[f(x') - [(x=)]]
2 B[
- fxx))
2
1=0
T-1
1
Σ 1=0
20
A +
BIs(a*) -[(x)] -
Σ
2ch B(f(xt) [(x))
2
t=
T-2
20
A
1
B[ -s(x)) + Σ
2a
dec +
2/2 - -((')
T-2
Σ
21th 10+1 [l)(x') - [xx))
T-2
Ar-1 + 1/2 Eis(x - -s(x)) + Σ
de + -12
20
2x
210
B[s(xt) - [xx)]
2
2
1=0
Ar-1 + 5/ - ~m B[ - f(x')]
2a
T-2
2a Ax 1 2 Blf(x*) f(x')1
t=0
T-1
\e (++z 1
Σ
2a
21
[B[[xt") ((x2))
(15)
Combining eq. (14) and eq. (15) we have
T
2a
Ae Net
+ 2)
2n/
Elf(x') - f(x')1
T-1
17b (1+ h1) -a (x+ 1 (+++)) 0².
$ ||00 x²112 +
Σ
(16)
t=
In eq. (16), setting d = X, Yt
0 with 0 < x <
2 aaj (in order to ensure that quantity
2x (X+ + )) 21 in the LHS is positive) we get
2a A+ 1 20au
T
Blf(x*) - f(z))) )1 S 11 - x=1 + 2T To (1+ - (A+ - 0².
Σ
2
(17)
Now dividing eq. (17) by T 22x (A+ 2)- - 2m/6a] and using Jensen's inequality we get
T-1
Blf(T") f(x")l < T
Σ
B[f(x') - =f(x")]
1=0
1000 x²1 112 + TT + 2np 20d - a)
(2ad aa 2-11///T
1 -zº112
2np + 21/6. - -20d a 0²
(2ad + a -176oiT
2ad + a -216ed
as wanted.
Now the proof of Theorem 3.2 follows immediately from Theorem D.1 setting X 8 (from
1-0
Proposition 2.1) and 76 77 (from Corollary 2.2). The bound on 1+ follows from the bound on À and from the fact that the function f(x) I+ is strietly increasing for x 2 0.
A
27

---

Published as a conference paper at ICLR 2025
D.2 PROOF OF THEOREM 3.6
Here we state and prove the corresponding theorem for DecSPS in the 2.3 setting.
Theorem D.2. Assume that each fi is convex and smooth. Then, the iterates of 2.3 with
1 min Ct
fs, (x') Cg.
[ss. (z²)|2 : and X 2 0,
where (ct) an increasing sequences with e₁ 2 1, C-1 cu. 971 17 and (De) decreasing sequence, converge as
E[f(=") - f(x))] 2x f(()))
Cr
o² T-1
Σ
T
Ta
T
t=0
where :" 7 T-1 x', = min ( mbco) } and D = max 1| - xx212
Proof. We have
| x=1 (7) 1| - x² x-11=- 21p. V Vis (x'),
-xx) + V (v,(("")) (x²)
Rearranging and dividing by 1/2 > 0 get
2(V -x)) ||_ x=1 ||
xx11
Mel
V (x,''')
17
178
Summing the above for t = 0,...,T and telescoping we have
2
[s. (x²)) (==)),=' -x)
|| 14-rra x²1 112
1 x01
+ T/e
V fs. (x
1/
|| x=1 11
1
101
T-1 Σ
+
m
(2')11
t=
t=
122
x= 11
x= 11?
1
x=
Mel Vel (x'))1
170
7/
1)t
120 x=11" 112
|| 1=1 1
|
x0 x0112
ne|| Vsl (x')1²
1/0
M0+1
1/t
120
112
||
x=1
| x0112
||
x01
Th (x')))
770
Mc+1
1/c
7T
t=0
1100
x=11
112
T
1
11 - x= 11?
T-1
+
1
x=
Mc V ff,s((''))
170
7h+1
1/
t=
1=0
1
T-2
1
D3 +
nel
V (x,)))1
110
M/++1
7
D2
Tell
nellV
fs, (2
(18)
11-1
Taking expectation in eq. (18) we have
D2
T-1 Σ
2 E
2'''
SE
mell
[
(19)
11--1
1=U
t=0
28

---

Published as a conference paper at ICLR 2025
Now the LHS of eq. (19) can be bounded as follows
T-1
-----""11 x"]
t=
T-1 2''
(+"=*")
2 E f( ) -2')] +2
Σ
E
1=1
(11) 2 (Vf())) ) -20) Σ
E 12(V (V -
2 [k[f(=)) f(x')] +
E (2(f(="'- -
-f((--11)]
T-1
=2E[f(x9) - f(x*)] +
Σ
[2(1 + - Bif - [[x)]]
-
22
T-1
=2E[f(x - [|xx)] +
2(1 + N1) B(f(t()) - f(x')] - Σ
2d -) -(((=))
T-2
=2B[ - [(x))] + - 2(1 + Xx) B(( [(x')] -
Σ
2d2+1 Blf(x*) f(x'))
1=0
T-1
T-1
=2 - f(x*)] +2 Σ B(f(z')- f(xx)]++++++]])))))))+))) + Σ
24, B[f(x') f(x")]
T-2
Σ 21++1 Elf(x')- - f((x))
= -2 - f(x')] +2 T-1
B[f(x*) - B10 f(x=))
1=1
T-2
T-2
24. (( - Σ
-
2d++1 B(F(2t) - 221 -[((=9)-$(55
2
Σ B[f("') f(x')( Elf(2--- -f(x'=))))
T-2
+2 2
Σ (Ne ((f - f(x*)( - 221 Bif(z9) - f(x*)]
t=1
T-1
22
Σ f(=") 2
(20)
1=0
using the fact that (Ae) is decreasing and the fact that Elf(x*) - f(x*)] 2 0. Combining eq. (19)
and eq. (20) we have
2 T-1
D2
T-1
Σ Elf(x*) - f(x*)) - -(=)(
+
Σ 0
Em Vi
(21)
11-
Now for the RHS of eq. (21) we have
D2
T-1
(3)
D2
+
Thl ([')
[ss, (xx( - E',,]
1--1
1--
C
D²
T-
T-1
Σ
Uss (x²) Ss. fs. (2
+
Σ
-
[s,(x") - [(,,]
W/-1
1=0
29

---

Published as a conference paper at ICLR 2025
CT CT-1Da D
T-1
- (
Σ t=
[fs, (x*)
(22)
where the last inequality follows from the fact that 1/t is decreasing and 1 2 an with a min Z.maa moco), from Lemma C.4. Thus, combining eq. (21) and eq. (22) we have
2 Σ - -x:1f(=)))= f(x") cr T
T-1
T-1 2
Σ
E[f(x') - f(x")]
+
Σ
t=
1=0
(23)
Rearranging eq. (23) we get
T-1 Σ t=0
(2- 2 - 1 -((=))(((x)))) 2\) - f(x))] + CT-1D
T-1
0²
Σ
(24)
r=0
and using the fact that
2 1>> 0 eq. (24) reduces
B[((2') - f(x")| S 22 [(x))+ Cr-1 2
T-1
Σ
(25)
2=
1=0
Dividing eq. (25) by I' and using Jensen's 's inequality we get
E[S(=") [(x))
S(x'))
-
CT-1D Ta
o²
Ct
1=
wanted.
If o² > 0 like in the standard SGD analysis under decreasing step-sizes, the choice leads to the asymptotic rate O(I/Vt).
vl +
Corollary D.3. Under the setting of Theorem D.2, for C: = vt+1 T
1- = Co) we have
Elf(=") - f(x')] ~ -
D2 20
T
avT VT
Proof. It follows from the well known inequality
Σ
S2VT.
Now the proof of Theorem 3.6 follows immediately from Corollary D.3 setting d = A (from Proposition 2.1). Moreover we have
1-B
||
12
[[ ))x² -a--1 x²112
< + A)||=*
x x=1 + A|| - x11
< (1 + 2A) max I|,." 101
(26)
thus D3 < (1 + 2X) max, || 1² = 1+0 D2 where D2 max || - x²11.
30

---

Published as a conference paper at ICLR 2025
D.3 PROOF OF THEOREM 3.7
Here we state and prove the corresponding theorem for AdaSPS in the 2.3 setting.
Theorem D.4. Assume that each fi is convex and smooth. Then, the iterates of 2.3 with
Js, (x.(x²) Es.
T/e min
71-1 and x 2 0,
V V fs, (2*)
F.
where 77-1 +00, c and (De) a decreasing sequence, converge as
-² + TO - T VT
where xJ Et= x', T= (A: ( + cLD + 1
and D = max; || x²1|
2c
Proof.
Using Equation (21) from the proof of Theorem 3.6, we have
D2
2 Σ (
Σ
Emel EmllVfss((
(27)
1T-1
t=0
t=
By Lemma C.7 we have
D 1-
T-1
< cLD?
Σ (xt - -3.. ==0
(28)
and by Lemma C,5 we have
T
Ss. (x*) Pss.
VrlVss
Ss. (xx) -E'.-
(29)
20 Eizo fs. (x*) E.s
0=
0=0
Now combining Equations (27) to (29) and using Jensen's inequality we get
T-1
2
Elf(x') f(x'((())))))))))))) - 2x1lf(x²) f(x")]
1=
T-1
<
cLD?
Σ fs. fs.. c
Σ
fs. (xr) -83.
T-1
(CLD2 + 2
E
(x*) EGs, -b,
s=0
(LDD + :)
E
(xs) fs. (xx) + fs. (x*) -(..)
<
CLD? +
Σ ((x)))
(30)
Rearranging eq. (30) we have
T-1
CLD2 + 1
E[f(x')- - + - +
Bl - [(xx)) + 02.
1=0
(31)
31

---

Published as a conference paper at ICLR 2025
Now let us choose cy -V((2²)- - f(x=). Then
2\11((=0) -((x)) - 2A1 ( - f(x*)
- 2
Σ - -
(32)
Combining eq. (31) and eq. (32) we get
T-l - f(x)] (xcy +
cLD² 1
Σ f(xt))
(33)
2c
t=0
==0
Squaring both sides of eq. (33) we have
T-1 Σ
BLS( B((()))
[()')]
S²²
E [(x)) To2)
(34)
where T = (Xice + CLD'
+
1 2c
Now we use Lemma C.6 in eq. (33) to get
2
Σ
(35)
Finally by Jensen's 's inequality and eq. (35) we have
12 TO
Elf(F') - +
T VT'
as wanted.
Now the proof Theorem 3.7 follows immediately from Theorem D.4 setting Proposition 2.1). Moreover, using eq. (26) have D
(from
I-Ø
1+3 D².
T-8
32

---

Published as a conference paper at ICLR 2025
E BEYOND THE BOUNDED ITERATES ASSUMPTION
The two main theorems on the convergence of MonDeeSPS/MomAdaS// require the bounded iterates assumption. As mentioned in the main paper, this is standard assumption for several adaptive step-sizes, see: (Reddi et al., 2018; Ward et al., 2020; Orvieto et al., 2022; Jiang and Stich, 2023). However, one can artificially remove this assumption by adding a projection step in the update rule of IMA onto compact and convex subset D C Rd as described in this section.
Consider the constrained finite-sum optimization problem,
min RED
f(x) = ((())
(36)
where each fi [R + IR is convex, smooth, and lower bounded by E and D S I is compact and convex subset. Let X S D be the set of minimizers of (36). We assume that X F 0 and we fix x² E EXD-
The new update rule of IMA, takes the following form (Projected IMA):
V.s. (x²)) x 1+1 A1+1 x²
2t1 = projo [2'
(37)
1
(38)
Aet1
A
with x0 = 20 E D, where projp(x) E arg minder ||d - z=1². Note that the projection step is only needed for the update of z'41 because r²+1 already convex combination of elements the convex set D, namely 2++1 and xf by induction. We highlight the fact that Projected IMA is not necessarily equivalent to Projected SHB. Now the proofs of Theorems D.2 and D.4 will go through using the non-expansiveness of the projection operator, explained in the following lemma:
Lemma E.1 (Non-expansiveness). For all x. y E TR it holds
[ - projo (v) < || - yl
Proof. See Example 8.14 and Lemma 8.16 in Garrigos and Gower (2023).
Now we have:
Theorem E.2 (Pojected IMA version of Theorem D.2), Assume that each Ji is convex and Li- smooth. Then, the iterates of Projected IMA with
1 min
fs. (x') - -3 [[Ss, (x') (11 and X 2 O,
where (ct) is an increasing sequences with e 2 L, c-1 co. 7-1 17b and (Ne) is decreasing sequence, converge as
-(((z)))((E)) - [(x))) -
CT eT
0² T
T
Ta
T
1=
where :" F Et= x1, min ( mboo } and D = diam D.
Proof. We have
|| x² 1²1
[Iproj
7/ V.s,, (x¹)) (=)))
Lem E.1
<
|| n/ v fs, (x') - x01
|| x² 1| 1| 2n fs, (x*), 2'- x²) (f
Now we continue exactly like the rest of the proof of Theorem D.2.
33

---

Published as a conference paper at ICLR 2025
Theorem E.3 (Pojected IMA version of Theorem D.4), Assume that each Sa is convex and Li- smooth. Then, the iterates of Projected IMA with
Js, (2:') Es. c ( V V
7i min
76-1
and d 0,
(x*) -ss
f.s.
where 1/ +00, >> 0 and (Ax) a decreasing sequence, converge
72 +
-(f(F")-(==))
TO
T VT
where x" 7 Et=o x², -= (e:r-((1v))) (
cLD²
and D diam D.
2 2c
Proof.
Same as above.
34

---

Published as a conference paper at ICLR 2025
F ADDITIONAL EXPERIMENTS
F.1 EXTRA DETERMINISTIC EXPERIMENT
Here we have included an extra experiment in the deterministic setting, this one for logistic regres- sion. We have performed a grid search find the best for MomPSmax since no optimal choices for 3 are known for the general convex case. We use the same 3 for HB and ALR-MAG for direct com- parison. For the step-size of HB we used 7 = 2(1 3) as recomended in Ghadimi et (2015). The results presented in Figure 7. In this problem the fastest, having similar performace with ALR-MAG and GD Polyak. For the deterministic setting have used 12 core AMD Ryzen 5 5600H CPU to run the experiments.
observe that our step-size (MomPSmas is
MAC
(a) Relative Error Plot
(b) Step-size Plot
Figure 7: Comparison of various deterministic algorithms for the logistic regression problem on synthetic data.
F.2 STOCHASTIC HEAVY BALL WITH CONSTANT STEP-SIZE
As we saw in Corollary 3.5, when %o with the constant step-size 7 2bmax
S 2bmax then we have new theoretical guarantees for SHB The most recent analysis of constant SHB is from Liu al.
1-8
(2020) with step-size = (1-8)2 min
As mentioned in the main text our
1-333 2 23+234
step-size is larger when B 2 2 ~ 0.236. In Figure 8 we provide numerical comparison of these results for logistic regression with synthetic data. We observe that SHB with the Liu et al. (2020) step-size has faster convergence for the first iterations but reaches plateau much earlier. Moreover, note that as 3 + both step-sizes have similar performance.
F.3 OTHER CHOICES
In this paper, we have provided convergence guarantees for SHB update rule x++1
x
7 (xt, + 3(x² 1111) with the MomSPS step-size given by % (1
3) min (x)))
However, as discussed in the main paper a more natural step-size would
(a'))
be to choose the well known SPS the SHB rule. We call this new update rule SPS with naive momentum. Here we numerically compare these two updates. As we see Figure 9, when is "small" and close to 0 then both SPSmax with naive momentum and MomSPS are close in performance. Of course, when B = they are both equal to standard SPS However, as gets larger the performance of SPS with naive momentum gets worse and less stable and at some point it diverges. For 3 2 0.7 there is no SPS with naive momentum, since it cannot be numerically computed.
35

---

Published as a conference paper at ICLR 2025
(a) 0.1
(b) B 0.2 PErOA
(c) 3 0.3
(d) 8 0.4
(e) 0.5
( B ===0. 0.6
(g) 8 ==0.7 0.7
(h) 3 = 0.8
(i) B = 0.9
Figure 8: Comparison of SHB with constant step-size. (Liu et al., 2020) vs Corollary 3.5 for various momentum coefficients 3 on logistic regression with synthetic data.
(a) 0.0
(b) B 0.1
(c) 3 0.2
(d) 8 0.3 IPflr
MandPha
1
i
(e) 0.4
( 3 = 0.5
(g) 0.6
(h) 3 0.7
Figure 9: Comparison of MomSPS versus SPS with naive momentum for various momentum coefficients 3 on logistic regression with synthetic data.
Furthermore, recall that our decreasing variants MomDecSPS and MomAdaSPS are respectively given by
- (x') 13g1 2--V
7t- = min
(MomDecSPPS)
cvt+ (())(())) vt+11
(1-B)[[s, (2x²) [z,]]
7 = min
(MomAdaSPS)
f(t / E'zo (x. -3'.
36

---

Published as a conference paper at ICLR 2025
The above step-sizes theoretically inspired by IMA, as explained in Appendix C.3, but again natural question is what happens we take the "correcting factor" 1 - 3 outside the minimum. We call these step-sizes Alternative MomDecSPS and Alternative MomAdaSPS and given by
fs, E3. M-IvE
7t = (1 (1-3) min
(Alt MomDecSPS)
cvt V V ss. (x9)1²² Vt+IJ
Ss. (x²) fss. cll fa.
= 1 8) min
Yt-1
(Alt MomAdaSPS)
fx,(x*)
E..
In Figure 10 and Figure 11
that when > O then our proposed step-sizes are significantly better. A possible reason is that for both Alt
have similiar performances, however when
MomDecSPS and Alt MomadaSPS holds ? < (1 7-- for all t. Thus inductively we get < (1 r0 So when 3 > 0, (1 - g) + as + oc which means that the step-size is too
small and barely updates x².
1000 0m.1)
(a) 3 = 0.0
(b) B 0.3
(c) 3 0.6
(d) 3 = 0.9
Figure 10: Comparison of MoMDecSPS versus Alternative MomDecSPS for various momentum coefficients 3 on logistic regression with synthetic data.
(a) B 0.0
(b) B 0.3
(C) 0.6
(d) 3
Figure 11: Comparison of MomAdaSPS versus Alternative MomAdaSPS for various momentum coefficients / on logistic regression with synthetic data.
F.4 IS 7bh NEEDED?
Recall the definition of MomPS "/t = (1- min f(a')_f(x") [V 76 }. The more natural choice would be without the upper bound "6- Moreover, in Appendix F.3, we can have the naive version of both of the above step-sizes, i.e. with no correcting factor 1 3. We compare the following step- sizes on both the least squares problem, in Figure 12, and logistic regression problem, in Figure 13, in the deterministic setting synthetic data. The method is SHB with 0.97 for the least squares and B = 0.3 for the logistic regression.
= (1- [(x²) ((=)) f(x*)
(MomPS)
[Vf(2*)11
f(x²) - f(x))
(Alt MomPS)
|VV((>']12
f(x*) - f(x²) [f(x')] 776
= (1 - 3) min
MomPS
37

---

Published as a conference paper at ICLR 2025
= min {
f(xx)) [V](2+)||1 76
(Alt MomPS max)
We see that " chosen large enough, all these step-sizes have similar performance. However, we highlight the fact that convergence guarantees are known only for MomPS max-
(a) 0.1
(b)
0.5
(c)
(d) 76b 10
Figure 12: Comparison of various upper bounds for MomPS and other alternatives on least squares with synthetic data.
(a) 76
(b) 76 2
(c) 76 10
76
100
Figure 13: Comparison of various upper bounds for MomPSmax and other alternatives on logistic regression with synthetic data.
F.5 PLOTS FOR C AND C₂
Below it plot of the constants C₁ and C2 functions of with %6- 2 and (recall % 2 a) from Theorem 3.2. Both functions are hyperbolas and the vertical line is the constant 2m--a"
We that in the interval
00, 2%6 2x-0a both functions are increasing.
Figure 14: Plots for the constants C₁ and C2 as functions of 3. Here we have chosen % = 2 and
a=1 1.
38

---

Published as a conference paper at ICLR 2025
F.6 MORE CONVEX EXPERIMENTS AND PARAMETER SETTINGS
In this section, we list the parameters, architectures and hardware that we used for the deep learning experiments. The information is collected in Table We also include some extra experiments in Figures 15 to 18.
Hyper-parameter Value
Architecture
Logistic Regression
GPUs
Ix NVIDIA GeForce RTX 3050
Batch-size
See caption each plot
Epochs
100
Trials
5
Weight Decay
0.0
Table 2: Logistic regression experiment
2115
125
5-g
SHB
511
Alknn
Noke Navertar
SiSrac saka Monemmu
Soocks
Spachs
6pcds
Figure 15: LibSVM dataset: vehicle, Batch size:
Figure 16: LibSVM datasct: letter, Batch size: 256
Decohs
AdSSSS
AecEPS
Morrees''s
MOTDNUSFS
AEEDS
Addiis Nom
toads
Spacha
Epoda
Figure 17: LibSVM dataset: vehicle, Batch size: 85
Figure 18: LibSVM dataset: glass, Batch size: 32
F.7 MORE DEEP LEARNING EXPERIMENTS AND PARAMETER SETTINGS
In this section, we list the parameters, architectures and hardware that we used for the deep learning experiments. The information is collected in Table 3. We also include some extra experiments in Figures 19 and 20. For ALR-SMAG we use 0.1 the experiments.
Hyper-parameter
Value
Architecture
ResNet 10/34 (He et al., 2016)
GPUs
1x Nvidia RTX 6000 Ada Generation
Batch-size
256
Epochs
100
Trials
5
Weight Decay
0.0
Table 3: CIFAR10 experiment
39

---

Published as a conference paper at ICLR 2025
157
2.4
LAME
60
22.3
1R Parnd
1750
100
Figure 19: Resnet 18 on CIFAR 10
Figure 20: Resnet 18 on CIFAR 100
40


Reviews:
[{"review_id": "5tVpwjKYOz", "reviewer": "Reviewer_SmTn", "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.", "questions": "## Questions\n1. Am I correct, that by \"interpolated regime\" you mean that noise variance is zero and we have determenistic case?\n1. 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.\n1. Lines 382-383. What do you mean by \"saddle connection\"?\n\n## Suggestions\n1. 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.\n1. Please, increase font size in the figures.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "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. \nOverall, 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Q6yAsaCpkQ", "reviewer": "Reviewer_QP47", "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.", "questions": "**Questions**\n\n1. 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?\n2. 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?\n\n**Minor issues**\n\n1. Line 55\n\n   has efficiently analyzed => has been efficiently analyzed.\n\n2. Line 400, 410\n\n   Please be consistent with step size and step-size.\n\n3. Line 1098\n\n   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.\n\n**References**\n\n[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.\n\n[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.\n\n[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.\n\n[4] Cutkosky, Ashok, and Francesco Orabona. Momentum-based variance reduction in non-convex sgd. *Advances in neural information processing systems* 32, 2019.\n\n[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.\n\n[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.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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. \n\n1. Limited technical contributions\n\n   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.\n\n2. Strong assumptions\n\n   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. \n\nOverall, 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "GvMAmf6CFU", "reviewer": "Reviewer_aQjo", "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.", "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?\n\nThe 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?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "GjtPZbaUBk", "reviewer": "Reviewer_jVKN", "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.", "questions": "None.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.The writing and exposition are clear.\n\n2.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)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)