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

## Abstract

Abstract Stochastic gradient descent with momentum, also known as Stochastic Heavy Ball method (SHB), is one of the most popular algorithms for solving large-scale stochastic optimization problems in various machine learning tasks. In practical scenarios, tuning the step-size and momentum parameters of the method is a prohibitively expensive and time-consuming process. In this work, inspired by the recent advantages of stochastic Polyak step-size in the performance of stochastic gradient descent (SGD), we propose and explore new Polyak-type variants suitable for the update rule of the SHB method. In particular, using the Iterate Moving Average (IMA) viewpoint of SHB, we propose and analyze three novel step-size selections: MomSPS max , MomDecSPS, and MomAdaSPS. For MomSPS max , we provide convergence guarantees for SHB to a neighborhood of the solution for convex and smooth problems (without assuming interpolation). If interpolation is also satisfied, then using MomSPS max , SHB converges to the true solution at a fast rate matching the deterministic HB. The other two variants, MomDecSPS and MomAdaSPS, are the first adaptive step-sizes for SHB that guarantee convergence to the exact minimizer without prior knowledge of the problem parameters and without assuming interpolation. The convergence analysis of SHB is tight and obtains the convergence guarantees of SGD with stochastic Polyak step-sizes as a special case. We supplement our analysis with experiments that validate the theory and demonstrate the effectiveness and robustness of the new algorithms.

### 1 Introduction

We consider the unconstrained finite-sum optimization problem,

$$ $\min_{x\in\mathbb{R}^{d}}\left[f(x)=\frac{1}{n}\sum_{i=1}^{n}f_{i}(x)\right],$ (1) $$

where each $f_{i}:\mathbb{R}^{d}\to\mathbb{R}$ is convex, smooth, and lower bounded by $\ell_{i}^{*}$. Let $X^{*}$ be the set of minimizers of (1). We assume that $X^{*}\neq\emptyset$ and we fix $x^{*}\in X^{*}$. This problem is the cornerstone of machine learning tasks (Hastie et al., ), where $x$ corresponds to the model parameters, $f_{i}(x)$ represents the loss on the training point $i$, and the aim is to minimize the average loss $f(x)$ across training points.

When $n$ is large, stochastic gradient methods are the preferred methods for solving (1) mainly because 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., ), a.k.a. stochastic heavy ball method (SHB) given by:

$$ $x^{t+1}=x^{t}-\gamma_{t}\nabla f_{S_{t}}(x^{t})+\beta_{t}(x^{t}-x^{t-1}).$ (SHB) $$

where $S_{t}\subseteq[n]$ a random subset of data-points (mini-batch) with cardinality $B$ sampled independently at each iteration $t$, and $\nabla f_{S_{t}}(x^{t})=\frac{1}{B}\sum_{i\in S_{t}}\nabla f_{i}(x^{t})$ is the mini-batch gradient. Here $\gamma_{t}>0$ is the step-size/learning rate at iteration $t$ while $\beta_{t}\geq 0$ represents the momentum parameter. When the momentum parameter $\beta_{t}=0,\forall t\geq 0$, then the update rule SHB is equivalent to the well-studied mini-batch SGD, $x^{t+1}=x^{t}-\gamma_{t}\nabla f_{S_{t}}(x^{t})$ (Robbins and Monro, ), which has efficiently analyzed under different properties of problem (1) and different step-size selections $\gamma_{t}$ (Nemirovski and Yudin, ; Nemirovski et al., ; Hardt et al., ; Needell et al., ; Nguyen et al., ; Gower et al., , ). Additionally, when the cardinality of the random subset $S_{t}$ is $B=n$, then the update rule of SHB is equivalent to the deterministic heavy ball method (HB) proposed by Polyak (), as a way to 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 (), it was proved that for a specific choice of the step-size $\gamma$ and the momentum parameter $\beta$, the HB method enjoys an accelerated linear convergence when minimizing strongly convex quadratic functions while more recently, Ghadimi et al. () proved a global sublinear convergence guarantee for HB for convex and smooth functions. In the stochastic setting, several works focus on convergence guarantees of SHB under constant step-size and momentum parameters (Ma and Yarats, ; Kidambi et al., ; Yan et al., ; Gitman et al., ; Liu et al., ). However, in practical scenarios, these choices suffer from a prohibitively expensive and time-consuming hyper-parameter tuning process. This has motivated a large body of research on the development of adaptive SHB - a 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 setting (deterministic) (Barré et al., ; Saab Jr et al., ), or assume that an interpolation condition is satisfied (Wang et al., ; Schaipp et al., ).

Previous studies in the fully stochastic (non-interpolated) 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.

##### ⋄ ⋄ \diamond ⋄ Efficient Polyak Step-sizes via IMA viewpoint.

We explain and illustrate by experiment (see Figure 1) why naively using SPSmax of Loizou et al. () as a step-size $\gamma_{t}$ 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 viewpoint from Sebbouh et al. (). Through our approach, we propose three novel adaptive step-size selections, namely MomSPSmax, MomDecSPS, MomAdaSPS. Each of the proposed step-sizes depends on the choice of the momentum parameter $\beta$ adding further stability to SHB and comes with specific benefits over their constant step-size counterparts or other adaptive variants of SHB.

##### ⋄ ⋄ \diamond ⋄ MomSPS max : Convergence of SHB in non-interpolated setting.

Our first step-size selection of SHB is MomSPSmax, which has a similar structure to SPSmax of Loizou et al. () but includes also $(1-\beta)$ 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 (Wang et al., ; Schaipp et al., ) on Polyak step-size with momentum have guaranteed only under the interpolation setting. In addition, as a corollary of our main theoretical results, we show that MomSPSmax under the interpolation setting and in deterministic scenarios (full batch) converges to the true solution at a fast rate matching the deterministic HB.

##### ⋄ ⋄ \diamond ⋄ Convergence of SHB to exact solution via MomDecSPS and MomAdaSPS.

Inspired by two recent Polyak-type step-size selections for SGD, the DecSPS in Orvieto et al. () and AdaSPS in Jiang and Stich (), we propose two new ways for tuning the step-size for SHB. These are MomDecSPS and MomAdaSPS. Our analysis provides the first $O(1/\sqrt{T})$ convergence guarantees to the exact solution in the non-interpolated regime for a Polyak-type adaptive variant of SHB. Our proposed update rules converge for any choice of momentum parameter $\beta\in[0,1)$, which makes them particularly useful in practical scenarios.

##### ⋄ ⋄ \diamond ⋄ Tight Convergence Guarantees.

All of our convergence guarantees are *true* generalizations of the theoretical analysis of SGD using SPSmax, DecSPS, and AdaSPS. That is, if $\beta=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 1 for a summary of our main complexity results and a comparison with closely related works.

##### ⋄ ⋄ \diamond ⋄ Further Convergence Results.

As a byproduct of our theoretical analysis, we provide two interesting corollaries: a novel analysis of constant step-size SHB (as a corollary of our Theorem on MomSPSmax) and the first robust convergence of SHB via our theorem on MomAdaSPS. For a constant step-size, the Corollary 3.4 of our Theorem 3.1 allows larger step-sizes than the analysis of SHB in Liu et al. () and provides convergence without assuming the restrictive bounded variance condition (there exist $q>0$ such that $\operatorname{\mathbb{E}}\|\nabla f_{i}(x)-\nabla f(x)\|^{2}<q$). In addition, via Theorem 3.6, we provide the first robust convergence of adaptive SHB that guarantees convergence to the exact solution and automatically adapts to whether our problem is interpolated or not. That is, if interpolation is assumed, then the rate of SHB with MomAdaSPS is the same as the rate of SHB with MomSPSmax (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.

##### ⋄ ⋄ \diamond ⋄ Numerical Evaluation.

In Section 4, we verify our theoretical results via numerical experiments on various problems, demonstrating the effectiveness and practicality of our approach. An open-source implementation of our method is available at [https://github.com/dimitris-oik/MomSPS](https://github.com/dimitris-oik/MomSPS).

**Table 1: Summary of the considered step-sizes and the corresponding theoretical results. All the rates are given for convex and smooth functions. The quantity of convergence in our rates is $\operatorname{\mathbb{E}}[f(\overline{x}^{T})-f(x^{*})]$, where $\overline{x}^{T}=\frac{1}{T}\sum_{t=0}^{T}x^{t}$. Here $\hat{\sigma}^{2}=\operatorname{\mathbb{E}}\|\nabla f_{i}(x^{*})\|^{2}<\infty$ and $\sigma^{2}=\operatorname{\mathbb{E}}[f_{i}(x^{*})-\ell_{i}^{*}]<\infty$. The “Exact Convergence” column refers to convergence to the exact solution $x^{*}$ with no interpolation assumption.**
| \CT@drsc@<br>Step-size | Assumptions | Adaptive | Exact Convergence | Rate |
| --- | --- | --- | --- | --- |
| Constant (Liu et al., ) | Knowledge of $L$ | ✗ | ✗ | $O(\frac{1}{T}+\hat{\sigma}^{2})$ |
| IMA (Sebbouh et al., ) | Knowledge of $L$ | ✗ | ✓ | $O(\frac{1}{T}+\hat{\sigma}^{2})$ |
| ALR-SMAG (Wang et al., ) | Interpolation | ✓ | ✗ | $O(\frac{1}{T}+\sigma^{2})$ |
| MomSPSmax (Thm 3.1) | Restriction on $\beta$ | ✓ | ✗ | $O(\frac{1}{T}+\sigma^{2})$ |
| MomDecSPS (Thm 3.5) | Bounded Iterates | ✓ | ✓ | $O(\frac{1}{\sqrt{T}})$ |
| MomAdaSPS (Thm 3.6) | Bounded Iterates | ✓ | ✓ | $O(\frac{1}{T}+\frac{\sigma}{\sqrt{T}})$ |
| \CT@drsc@ |  |  |  |  |

### 2 Exploring the Interplay of SPS and Heavy Ball Momentum

In this section, we present the expression of Stochastic Polyak Step-size (SPS) and its different variants. We illustrate that a 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 () as an efficient step-size selection for GD for solving convex optimization problems. It has the following expression: $\gamma_{t}=\frac{f(x^{t})-f(x^{*})}{\|\nabla f(x^{t})\|^{2}}$ which is obtained by minimizing an upper bound of the quantity $\|x^{t+1}-x^{*}\|^{2}$ 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 et al., ; Davis et al., ; Hazan and Kakade, ). PS requires the prior knowledge of $f(x^{*})$, which might look like a strong assumption. However, as shown in Boyd et al. (), this is known in several applications, including finding a point in the intersection of convex sets and positive semi-definite matrix completion.

Inspired by the convergence of PS in the deterministic setting, Loizou et al. () 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, including independence on parameters of the problem (e.g., $L$-smoothness or $\mu$ strong convexity) and competitive performance in over-parametrized models. More specifically, Loizou et al. () proposed the SPSmax given below(^1^11Originally Loizou et al. () use $f_{S_{t}}^{*}$ instead of the lower bound $\ell_{S_{t}}^{*}$. The lower bound is due to Orvieto et al. (), which proves that the more relaxed lower bound can still lead to the same convergence guarantees.):

$$ $\gamma_{t}=\min\left\{\frac{f_{S_{t}}(x^{t})-\ell_{S_{t}}^{*}}{c\|\nabla f_{S_ {t}}(x^{t})\|^{2}},\gamma_{b}\right\}.$ (SPS max ) $$

Here, $\gamma_{b}>0$ is a bound that restricts SPS from being very large and is essential to ensure convergence to a small neighborhood around the solution and $c>0$ is a positive constant that depends on the function class the objective $f$ belongs to.

As mentioned in Orvieto et al. (), the SPSmax comes with strong convergence guarantees and competitive performance; however, it has one main drawback when used in non-over-parameterized regimes: It can guarantee convergence only to a neighborhood of the solution. For this reason, Orvieto et al. () suggests a decreasing variant of the original SPSmax named DecSPS, given by:

$$ $\gamma_{t}=\frac{1}{c_{t}}\min\left\{\frac{f_{S_{t}}(x^{t})-\ell_{S_{t}}^{*}}{ \|\nabla f_{S_{t}}(x^{t})\|^{2}},c_{t-1}\gamma_{t-1}\right\},$ (DecSPS) $$

where $c_{t}$ is an increasing sequence of positive real numbers and $c_{-1}:=c_{0}$ and $\gamma_{-1}=\gamma_{b}>0$. The authors proved that SGD with DecSPS and $c_{t}=\sqrt{t+1}$, converges with a sublinear rate $O(1/\sqrt{T})$ for convex and smooth functions with bounded iterates (i.e., $D^{2}:=\max_{t\in[T-1]}\|x^{t}-x^{*}\|^{2}<\infty$).

More recently, in Jiang and Stich (), another decreasing variant of SPS was introduced, named AdaSPS:

$$ $\gamma_{t}=\min\left\{\frac{f_{S_{t}}(x^{t})-\ell_{S_{t}}^{*}}{c\|\nabla f_{S_ {t}}(x^{t})\|^{2}}\frac{1}{\sqrt{\sum_{s=0}^{t}f_{S_{s}}(x^{s})-\ell_{S_{s}}^{ *}}},\gamma_{t-1}\right\},$ (AdaSPS) $$

where $\gamma_{-1}=+\infty$ and $c>0$. For convex and $L$-smooth functions with bounded iterates, it can be shown that SGD with AdaSPS converges to an exact solution with a rate $O\left(\frac{\tau^{2}}{T}+\frac{\tau\sigma}{\sqrt{T}}\right)$, where $\tau=2cLD^{2}+\frac{1}{c}$. The interesting aspect of the convergence of AdaSPS is that it provides a robust result for SGD, meaning that the method recovers the best bounds for both the interpolated ($\sigma=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 no momentum). With the increased popularity of momentum in machine learning, one can naturally ask the following question:

$$ Is it possible to combine SPS with momentum? $$

This question was partially answered in Wang et al. () when a momentum variant of SGD with Polyak-type step-size was proposed to solve the easier-to-analyze interpolated problems. In this scenario, the method behaves like its deterministic counterpart. However, several stochastic optimization problems do not satisfy the interpolation condition. In our work, we focus on this more general scenario.

The most straightforward approach is to directly apply the SPSmax in the SHB update rule. Let us call this update rule SPSmax with naive momentum. Unfortunately, this approach does not necessarily lead to convergence for natural choices of momentum parameter $\beta\in(0,1)$, as shown in Figure 1.(^2^22See also Figure 9 in Appendix for more values of $\beta$.) In this experiment, even for simple convex and smooth problems like a logistic regression with synthetic data, the naive rule fails to converge when $\beta$ gets larger (a typical choice for momentum parameter is $\beta=0.9$). This indicates that a more careful step-size selection is needed, which may also depend on the momentum parameter. In Figure 1, we compare the SPSmax with naive momentum with one of our proposed and analyzed step-size selection MomSPSmax. As seen in Figure 1, when $\beta$ is small, then both SPSmax with naive momentum and SHB with MomSPSmax have similar behavior. When $\beta=0$, the two methods have identical performance as both are reduced to the same method: SGD with SPSmax. However, as $\beta$ gets larger, the performance of SPSmax with naive momentum gets worse and less stable, and at some point, it diverges. On the other hand, in this example, SHB with MomSPSmax converges for any of the selected $\beta$.

Figure: (a) $\beta=0.0$
Refer to caption: x1.png

#### 2.3 Iterative Moving Average: Balance between SPS and Momentum

Having explained how naively combining SPSmax with heavy ball momentum can lead to divergence, in this section, we leverage the iterate moving-average (IMA) viewpoint of SHB from Sebbouh et al. () to propose SPS-type adaptations that depend on momentum parameter as 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. () 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:

$$ $z^{t+1}=z^{t}-\eta_{t}\nabla f_{S_{t}}(x^{t}),\quad x^{t+1}=\frac{\lambda_{t+1 }}{\lambda_{t+1}+1}x^{t}+\frac{1}{\lambda_{t+1}+1}z^{t+1},$ (IMA) $$

where $z^{0}=x^{0}$, $\eta_{t}>0$ and $\lambda_{t}\geq 0$.
As proved in Sebbouh et al. (), if for any $t\in\mathbb{N}$, holds $1+\lambda_{t+1}=\frac{\lambda_{t}}{\beta_{t}}\text{ and }\eta_{t}=(1+\lambda_{
t+1})\gamma_{t}$ then the $x_{t}$ iterates of the IMA method are equal to the $x_{t}$ iterates produced by the SHB method. For completeness, we include the proof of this statement in Section C.3.

In our proposed methods, we select the most common setting of constant momentum ($\beta_{t}=\beta$). 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 a convergence analysis of IMA.

###### Proposition 2.1 .

If $\lambda_{t}=\lambda\geq 0$ then assuming that $\beta=\frac{\lambda}{1+\lambda}$ and $\gamma_{t}=(1-\beta)\eta_{t}$ the $x_{t}$ iterates of the IMA method are equal to the $x_{t}$ iterates produced by the SHB method.

Using the above proposition, let us provide the following corollary that explains the derivation of our proposed step-size selections for SHB.

###### Corollary 2.2 .

Let the step-size $\eta_{t}$ in IMA be one of the previously proposed Polyak step-sizes: SPSmax, DecSPS or AdaSPS and let $\lambda_{t}=\lambda$. Then via Proposition 2.1, the SHB has constant momentum parameter $\beta_{t}=\beta$ and the step-sizes are MomSPSmax, MomDecSPS or MomAdaSPS respectively.

###### Proposition 2.1 .

###### Corollary 2.2 .

##### Three New Adaptive Step-sizes for SHB.

Using the IMA viewpoint and Proposition 2.1 and Corollary 2.2 let us present three Polyak-type step-size for SHB.

(i) MomSPSmax.
Let us start with the variant associated with the SPSmax. That is, via Corollary 2.2 if we use SPSmax as a step-size for IMA, then this is equivalent to using the following step-size in the update rule of SHB:

$$ $\gamma_{t}=(1-\beta)\min\left\{\frac{f_{S_{t}}(x^{t})-\ell_{S_{t}}^{*}}{c\| \nabla f_{S_{t}}(x^{t})\|^{2}},\gamma_{b}\right\}.$ (MomSPS max ) $$

Here, the parameter $\gamma_{b}>0$ has the same purpose as in the original SPSmax, and it is a bound that restricts MomSPSmax from being very large and is essential to ensure convergence to a neighborhood around the solution. It is clear from its expression that MomSPSmax has the same form as SPSmax but multiplied by a *correcting momentum factor* $1-\beta$. 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 MomSPSmax and a restriction on the momentum parameter $\beta$ in Section 3, we establish $O(1/T)$ convergence for SHB up to a neighborhood of the solution for convex and smooth functions.

Our following two proposed step-size selections are variants of the two decreasing variants of Polyak step-sizes DecSPS and AdaSPS. As such, we can prove SHB’s convergence to the exact solution instead of a neighborhood. More importantly, using MomDecSPS or MomAdaSPS, we can provide convergence guarantees for any choice of the momentum coefficients $\beta\in[0,1)$ making SHB fully adaptive (no tuning necessary).

(ii) MomDecSPS.
Firstly, we propose an adaptation of DecSPS with momentum. We call this step-size selection MomDecSPS, and it is given by:

$$ $\gamma_{t}=\min\left\{\frac{(1-\beta)[f_{S_{t}}(x^{t})-\ell_{S_{t}}^{*}]}{c_{t }\|\nabla f_{S_{t}}(x^{t})\|^{2}},\frac{\gamma_{t-1}c_{t-1}}{c_{t}}\right\},$ (MomDecSPS) $$

where $\gamma_{-1}:=\gamma_{b}>0$ is a step-size bound and $c_{t}=c\sqrt{t+1}$ for $t\geq 0$ with $c_{-1}:=c_{0}=c>0$ be a constant to regulate the step-size. In this scenario we note that $\gamma_{t}$ is indeed decreasing since $\gamma_{t}\leq\frac{\gamma_{t-1}\sqrt{t}}{\sqrt{t+1}}\leq\gamma_{t-1}$. 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 $\beta\in[0,1)$ converges to exact solution with a rate $O(1/\sqrt{T})$.

(iii) MomAdaSPS.
Similar to the previous two steps, let us propose the following adaptation of AdaSPS with momentum. We call this MomAdaSPS, and is given by:

$$ $\gamma_{t}=\min\left\{\frac{(1-\beta)[f_{S_{t}}(x^{t})-\ell_{S_{t}}^{*}]}{c\| \nabla f_{S_{t}}(x^{t})\|^{2}\sqrt{\sum_{s=0}^{t}f_{S_{s}}(x^{s})-\ell_{S_{s}} ^{*}}},\gamma_{t-1}\right\}.$ (MomAdaSPS) $$

By definition, this is a decreasing step-size ($\gamma_{t}\leq\gamma_{t-1}$). In Section 3, we prove that SHB with MomAdaSPS also converges to the exact solution for any choice of $\beta\in[0,1)$. More specifically, we show that MomAdaSPS converges to the exact solution with a rate $O\left(\frac{1}{T}+\frac{\sigma}{\sqrt{T}}\right)$. Thus, when we are in the interpolation regime (i.e., $\sigma^{2}=0$), the convergence of SHB with MomAdaSPS matches SHB with MomSPSmax while when the setting does not satisfy interpolation condition, this became equivalent to the $O(1/\sqrt{T})$ of MomDecSPS. Following the terminology of Jiang and Stich (), our results are the first *robust* adaptive step-size selection for SHB, as it can automatically adapt to the optimization setting (interpolation vs. non-interpolation).

Let us close this section by mentioning two remarks that related to the above three step-size selections:

###### Remark 2.3 .

The “correcting factor” $1-\beta$ is outside the minimum in the MomSPSmax 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 Section C.3. In Section E.3, via experiments, we also test other variants of the above step-sizes with the correcting factor outside of the minimum. We observe that the other choices do not have good practical performance.

###### Remark 2.4 .

For no momentum, i.e., $\beta=0$, all three proposed step-size choices, MomSPSmax, MomDecSPS and MomAdaSPS are reduced to SPSmax, DecSPS and AdaSPS and our convergence analysis recovers the convergence guarantees for SGD showing the tightness of our approach.

###### Remark 2.3 .

###### Remark 2.4 .

### 3 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 D. For all of our results, we make the assumption $\sigma^{2}:=\operatorname{\mathbb{E}}_{S_{t}}[f_{S_{t}}(x^{*})-\ell_{S_{t}}^{*
}]=f(x^{*})-\operatorname{\mathbb{E}}_{S_{t}}[\ell_{S_{t}}^{*}]<\infty$. This assumption was first introduced in Loizou et al. () and Orvieto et al. (). Note that $\sigma^{2}<\infty$ when all $f_{i}$ are lower bounded like we have assumed. We say that problem (1) is *interpolated* if $\sigma^{2}=0$. If the interpolation condition is satisfied then there exists $x^{*}\in X^{*}$ such that $f(x^{*})=f_{S_{t}}(x^{*})=\ell_{S_{t}}^{*}$ for all $S_{t}\subseteq[n]$. Many modern machine learning models satisfy this condition. Examples include non-parametric regression Liang and Rakhlin () and over-parameterized deep neural networks (Zhang et al., ; Ma et al., ).

#### 3.1 Convergence to a neighborhood of the solution

We start with the analysis of SHB with MomSPSmax.

###### Theorem 3.1 .

Assume that each $f_{i}$ is convex and $L_{i}$-smooth. Then, the iterates of SHB with MomSPSmax with $c=1$ and $\beta\in\left[0,\frac{\alpha}{2\gamma_{b}-\alpha}\right)$ where $\alpha=\min\left\{\frac{1}{2L_{\max}},\gamma_{b}\right\}$ and $L_{\max}=\max_{i\in[n]}\{L_{i}\}$, converge as

$$ $\operatorname{\mathbb{E}}[f(\overline{x}^{T})-f(x^{*})]\leq\frac{C_{1}\|x^{0}- x^{*}\|^{2}}{T}+C_{2}\sigma^{2},$ $$

where $\overline{x}^{T}=\frac{1}{T}\sum_{t=0}^{T-1}x^{t}$ and the constants $C_{1}=\frac{1-\beta}{\alpha\beta+\alpha-2\beta\gamma_{b}}$ and $C_{2}=\frac{2\gamma_{b}-\alpha\beta-\alpha}{\alpha\beta+\alpha-2\beta\gamma_{b}}$.

Firstly, let us note that both constants $C_{1}$ and $C_{2}$ in Theorem 3.1 are positive since the denominator of both is positive because $\beta\in[0,1)$ and the numerator of $C_{2}$ is positive because $\gamma_{b}\geq\alpha>0$. The above result shows that MomSPSmax has a sublinear convergence to a neighborhood matching the best-known rate for SHB in the convex setting, see Liu et al. (). Moreover, when $\beta=0$, the update rule of SHB with MomSPSmax becomes equivalent to SGD with SPSmax (no momentum) and the result of Theorem 3.1 reduces to convergence guarantees provided in Loizou et al. (), only with a slightly better neighborhood of convergence (here when $\beta=0$ we have $C_{2}=(2\gamma_{b}-\alpha)/\alpha$ while in the original it is $C_{2}=2\gamma_{b}/\alpha$).

In addition, note that in Theorem 3.1, there is a restriction in the momentum coefficient. This stems from the technique used in the proof (which forces $C_{1}$ and $C_{2}$ to be positive). However, this restriction is not vital as in practical scenarios, the method converges even when $\beta$ is selected outside the given interval. Moreover, let us highlight that the constants $C_{1}$ and $C_{2}$ are increasing when viewed as functions of $\beta$ in the given interval (see also Figure 14). This means that our theory suggests that $\beta=0$ (no momentum) is the best theoretical choice, which is typical in many works on stochastic methods with momentum (Wang et al., ; Loizou and Richtárik, ). This is not ideal but does not undermine the importance of Theorem 3.1, as this is the first result showing the convergence of SHB with SPS in a non-interpolated setting.

From Theorem 3.1, we can also deduce rates for the deterministic and interpolated regimes.

###### Corollary 3.2 .

Assume that $f$ is convex and $L$-smooth. Then HB with $\beta\in[0,\alpha/(2\gamma_{b}-\alpha))$ and

$$ $\gamma_{t}=(1-\beta)\min\left\{\frac{f(x^{t})-f(x^{*})}{\|\nabla f(x^{t})\|^{2 }},\gamma_{b}\right\},$ (MomPS max ) $$

where $\alpha=\min\left\{\frac{1}{2L},\gamma_{b}\right\}$, converges as $\min_{t\in[T]}\{f(x^{t})-f(x^{*})\}\leq\frac{(1-\beta)\|x^{0}-x^{*}\|^{2}}{(
\alpha\beta+\alpha-2\beta\gamma_{b})T}$.

Note that in the setting of Corollary 3.2, we have used $f(x^{*})$ instead of a lower bound $\ell^{*}$. This is in accordance with previous work in the deterministic setting, (Polyak, ). We highlight that in Hazan and Kakade () (PS) and Wang et al. () (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 $\|x^{t}-x^{*}\|^{2}$. Our step-size MomPSmax is *not* a minimizer of such bound when $\beta>0$. Nevertheless, when $\beta=0$ (no momentum), by setting $\gamma_{b}=\infty$ we recover the corresponding result in Hazan and Kakade (). In our setting, we cannot set $\gamma_{b}=\infty$ when $\beta>0$ because the term $\beta\gamma_{b}$ appears in the denominator of our convergence result.

###### Corollary 3.3 .

Assume interpolation ($\sigma^{2}=0$) and let all assumptions of Theorem 3.1 be satisfied. Then SHB with MomSPSmax and $\beta\in[0,\alpha/(2\gamma_{b}-\alpha))$ where $\alpha=\min\left\{\frac{1}{2L_{\max}},\gamma_{b}\right\}$ converges as

$$ $\operatorname{\mathbb{E}}[f(\overline{x}^{T})-f(x^{*})]\leq\frac{(1-\beta)\|x^ {0}-x^{*}\|^{2}}{(\alpha\beta+\alpha-2\beta\gamma_{b})T},$ $$

where $\overline{x}^{T}=\frac{1}{T}\sum_{t=0}^{T-1}x^{t}$.

In the closely related works Loizou et al. () (SPSmax) and Wang et al. () (ALR-SMAG), there is no upper bound in the proposed step-sizes for the interpolated setting. This is because these works treat the interpolated regime the same as the deterministic regime (the proofs of the results in these two settings are almost identical). In our analysis, we focus on the fully stochastic setting (interpolation is only a corollary), and the result of Corollary 3.3 shows a saddle connection between the momentum parameter and the parameter $\gamma_{b}$ (we cannot select $\gamma_{b}=\infty$ when momentum is used). Nevertheless, via our corollary, when $\beta=0$ by setting $\gamma_{b}=\infty$, we can recover the corresponding result in Loizou et al. ().

Another interesting outcome of Theorem 3.1 is a novel analysis for SHB with a constant step size.

###### Corollary 3.4 .

Let all assumptions of Theorem 3.1 be satisfied. If $\gamma_{b}\leq\frac{1}{2L_{\max}}$, then SHB with $\gamma\leq\frac{1-\beta}{2L_{\max}}$ and $\beta\in[0,1)$ converges as

$$ $\operatorname{\mathbb{E}}[f(\overline{x}^{T})-f(x^{*})]\leq\frac{\|x^{0}-x^{*} \|^{2}}{T\gamma_{b}}+\sigma^{2},$ $$

where $\overline{x}^{T}=\frac{1}{T}\sum_{t=0}^{T-1}x^{t}$.

Note that in Corollary 3.4, there is no restriction that depends on the smoothness parameter in the momentum parameter $\beta$. We allow to have $\beta\in[0,1)$. A similar result is obtained in Liu et al. (), where the authors establish the convergence rate $O\left(\frac{f(x^{0})-f(x^{*})}{T}+\frac{L\gamma\hat{\sigma}^{2}}{1-\beta}\right)$ for the quantity $\frac{1}{T}\sum_{t=0}^{T-1}\operatorname{\mathbb{E}}[\|\nabla f_{S_{t}}(x^{t})
\|^{2}]$ when $\gamma\leq\frac{(1-\beta)^{2}}{L}\min\left\{\frac{1}{4-\beta+\beta^{2}},\frac{
1}{2\sqrt{2\beta+2\beta^{2}}}\right\}$. Here $\hat{\sigma}^{2}$ is assumed to be a bound of the variance $\operatorname{\mathbb{E}}_{S_{t}}\|\nabla f_{S_{t}}(x^{t})-\nabla f(x^{t})\|^{
2}\leq\hat{\sigma}^{2}$. In comparison, in our analysis, we provide the same asymptotic rate $O(1/T)$ for the common in the convex setting quantity $\operatorname{\mathbb{E}}[f(\overline{x}^{T})-f(x^{*})]$. By comparing the two results, our step size is larger when $\beta\geq\sqrt{5}-2\approx 0.236$. However, in our result, the neighborhood of convergence is constant and does not depend on the step-size $\gamma$ or the momentum coefficient. For a numerical comparison of SHB with the two different constant step-sizes, see Section E.2.

###### Theorem 3.1 .

###### Corollary 3.2 .

###### Corollary 3.3 .

###### Corollary 3.4 .

#### 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 we can guarantee convergence to the exact solution while keeping the main adaptiveness properties. Due to the nature of the decreasing step sizes, there is no restriction in the momentum parameter, as was the case with Theorem 3.1. For the results of this section, we make the extra assumption of *bounded iterates*, i.e., we assume that $D^{2}=\max_{t\in[T]}\|x^{t}-x^{*}\|^{2}$ is finite. This is a standard assumption for several adaptive step-sizes, see: (Reddi et al., ; Ward et al., ; Orvieto et al., ; Jiang and Stich, ).

###### Theorem 3.5 .

Assume that each $f_{i}$ is convex and $L_{i}$-smooth. Then, the iterates of SHB with MomDecSPS with $c=1$ and $\beta\in[0,1)$ converge as:

$$ $\displaystyle\operatorname{\mathbb{E}}[f(\overline{x}^{T})-f(x^{*})]\leq\frac{ 2\beta[f(x^{0})-f(x^{*})]}{(1-\beta)T}+\frac{(1+\beta)D^{2}}{(1-\beta)\alpha \sqrt{T}}+\frac{2\sigma^{2}}{\sqrt{T}},$ $$

where $\overline{x}^{T}=\frac{1}{T}\sum_{t=0}^{T-1}x^{t}$, $\alpha=\min\left\{\frac{1}{2L},\gamma_{b}\right\}$ and $D^{2}=\max_{t\in[T]}\|x^{t}-x^{*}\|^{2}$.

Notice that when there is no momentum, i.e., for $\beta=0$, Theorem 3.5 recovers the original result of SGD with DecSPS from Orvieto et al. (). To our knowledge, the only other result that guarantees convergence to the exact solution for SHB in the non-interpolated regime appears in Sebbouh et al. (). In particular, Sebbouh et al. () proves that for $\gamma_{t}=\frac{\eta_{t}}{1+\lambda_{t+1}}$ and $\beta_{t}=\frac{\lambda_{t}}{1+\lambda_{t}}$ where $\eta_{t}=\frac{\eta}{\sqrt{t+1}}$ with $0\leq\eta\leq\frac{1}{4L_{\max}}$ and $\lambda_{t}=\frac{1}{2}\sum_{i=0}^{t-1}\sqrt{\frac{t+1}{i+1}}$ with $\lambda_{0}=0$, SHB converges as $\operatorname{\mathbb{E}}[f(x^{t})-f(x^{*})]\leq O(\log t/\sqrt{t})$. This result only holds for a specific choice of $\beta_{t}$, it needs knowledge of $L_{\max}$, and only guarantees $O(\log t/\sqrt{t})$ rate. However, it does not make a bounded iteration assumption, and it shows convergence of the last iterate. In contrast, our convergence of SHB with MomDecSPS holds for any choice of $\beta\in[0,1)$, the step-size is adaptive and does not depend on $L_{\max}$, and it achieves a rate of $O(1/\sqrt{k})$.

Next, we present the theoretical guarantees for MomAdaSPS.

###### Theorem 3.6 .

Assume that each $f_{i}$ is convex and $L_{i}$-smooth. Then, the iterates of SHB with MomAdaSPS and $\beta\in[0,1)$ converge as

$$ $\operatorname{\mathbb{E}}[f(\overline{x}^{T})-f(x^{*})]\leq\frac{\tau^{2}}{T}+ \frac{\tau\sigma}{\sqrt{T}},$ $$

where $\overline{x}^{T}=\frac{1}{T}\sum_{t=0}^{T-1}x^{t}$, $\tau=\frac{\beta\sqrt{f(x^{0})-f(x^{*})}}{1-\beta}+\frac{(1+\beta)cLD^{2}}{2(1
-\beta)}+\frac{1}{2c}$ and $D^{2}=\max_{t\in[T]}\|x^{t}-x^{*}\|$.

As in our previous theorems, when there is no momentum ($\beta=0$), Theorem 3.6 recovers the convergence guarantees of AdaSPS from Jiang and Stich (). 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 a rate of $O(1/\sqrt{T})$, which matches the best-known rate for SHB while if we assume interpolation ($\sigma^{2}=0$), then we are able to achieve a rate of $O(1/T)$ which matches the best-known rate for SHB in the convex setting (see Liu et al. () and Theorem 3.1). Furthermore, under interpolation, Theorem 3.6 improves the convergence of Corollary 3.3. In particular, when $\sigma^{2}=0$, Theorem 3.6 reaches the rate $O(1/T)$ for *any* $\beta\in[0,1)$, while in Corollary 3.3 this is possible only under tighter restrictions on $\beta$.

###### Theorem 3.5 .

###### Theorem 3.6 .

### 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.

#### 4.1 Deterministic Setting

For the deterministic setting, we focus on the least squares problem. The loss function in this case is given by $f(x)=\frac{1}{2}\|Ax-b\|^{2}$, where $A\in\mathbb{R}^{n\times d}$ and $b\in\mathbb{R}^{n}$. In our experiments, we follow the setting of Wang et al. () and choose $n=d=1000$ while the matrix $A$ has been generated according to Lenard and Minkoff () such that the condition number of $A^{T}A$ is $10^{4}$. We test the following algorithms: Gradient Descent (GD), Nesterov’s Accelerated Gradient Descent (AGD) (Nesterov, ), Gradient Descent with Polyak’s step-size (Polyak, ), Adaptive Heavy Ball (AHB) (Saab Jr et al., ), Polyak’s Heavy Ball (HB) (Polyak, ), Accelerated Gradient Method (AGM) (Barré et al., ), ALR-MAG (Wang et al., ) and finally our step-size from Corollary 3.2, called (MomPSmax), which guarantees $O(1/T)$ rate. For the least squares problem, there is an established theory from Polyak () that guarantees acceleration of HB with optimal step-size and momentum coefficient given by $\beta^{*}=(\sqrt{L}-\sqrt{\mu})^{2}/(\sqrt{L}+\sqrt{\mu})^{2}$ and $\gamma^{*}=(1+\sqrt{\beta^{*}})/L$. For ALR-MAG and SHB with MomPSmax, we have used the optimal momentum $\beta^{*}$ for direct comparison with optimal HB. Moreover, for GD Polyak, ALR-MAG, and MomPSmax, we have used the precomputed $f(x^{*})$ (via GD), and for MomPSmax, we have chosen $\gamma_{b}=100$ (for other choices, see Section E.4). The results in 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 MomPSmax, and the optimal HB step-sizes. We can see that the original Polyak step-size oscillates around the optimal HB step-size while MomPSmax converges from below to the optimal.

Figure: Figure 2: Comparison of MomPSmax against various deterministic algorithms for consistent least squares. Left: Relative error, Right: Step-sizes.
Refer to caption: x5.png

#### 4.2 Stochastic Setting: Convex Problems

Here, we compare our step-sizes with previous works in the stochastic setting. As noticed initially for SPSmax in Loizou et al. (), the value of the upper bound $\gamma_{b}$ that results in good convergence depends on the problem and requires careful parameter tuning. To alleviate this problem, Loizou et al. () uses a smoothing procedure that prevents large fluctuations in the step-size across iterations. That is, $\gamma_{b}^{t}=\tau^{b/n}\gamma_{t-1}$ for each iteration $t$ where $\tau=2$, $b$ is the batch-size, and $n$ is the number of examples. We use the same smoothing trick for $\gamma_{b}$ in the implementation of the proposed methods. Similar smoothing procedures have been used in (Tan et al., ; Vaswani et al., ).

We considered multi-class logistic regression applied to commonly used benchmark datasets from the LIBSVM repository, (Chang and Lin, ). We separate our experiments between two classes of step-sizes: non-decreasing step-sizes (which might guarantee convergence to a neighborhood of the solution) and decreasing step-sizes (which guarantee convergence to the exact solution). For the first class of step-sizes we test the following algorithms: SGD (Nesterov, ), SHB with the constant step-size and $\beta=0.9$, ADAM as described in Kingma and Ba (), the warm-up version of ALR-SMAG with the hyper-parameters suggested in Wang et al. (), the SPSmax, the SPSmax with naive momentum as described in Section 2.2 and our proposed MomSPSmax. For the Polyak-based step-sizes such as ALR-SMAG, SPSmax, and MomSPSmax, we select $\ell_{i}^{*}=0$ and $c=1$. For SHB, ALR-SMAG, MomSPSmax, SPS${}_{\max},$ with naive momentum, we use momentum $\beta=0.9$. All the convex experiments are run for 100 epochs and for 5 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 our proposed MomSPSmax significantly outperforms the other methods in both training loss and test accuracy. For the decreasing variants, we test AdaGrad-Norm (Duchi et al., ), 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.

Figure: Figure 3: Comparison between our proposed MomSPSmax, other momentum methods, SPSmax (no momentum) and constant step-size SGD. Setting: Unregularized logistic regression, LIBSVM dataset: vowel, Batch-size for all methods: 52
Refer to caption: x6.png

Figure: Figure 4: Comparison of our proposed MomDecSPS and MomAdaSPS between DecSPS and AdaSPS as well as AdaGrad-Norm. Setting: Unregularized logistic regression, LIBSVM Dataset: letter, Batch-size for all methods: 1500
Refer to caption: x7.png

#### 4.3 MomSPS max for DNNs

In our final experiment, we go beyond the convex setting of our convergence guarantees, and we test MomSPSmax on the training of DNNs. We consider non-convex minimization for multi-class classification using deep network models on the CIFAR-10 and CIFAR-100 datasets, (Krizhevsky et al., ). We use two standard image-classification architectures: ResNet34 (He et al., ) and DenseNet121 (Huang et al., ). For space concerns, we report only the ResNet experiments in the main paper and relegate the DenseNet to the Appendix. We present the results for CIFAR-10 in Figure 5 and for CIFAR-100 in Figure 6. We observe that MomSPSmax consistently outperforms its no-momentum counterpart SPSmax and has competitive generalization performance compared to other popular optimizers. An important observation is that for these experiments (CIFAR-10 and CIFAR-100 with ResNet34), the closely related methods that combine stochastic Polyak step-size with momentum (ALR-SMAG from Wang et al. () and SPSmax with naive momentum from Section 2.2) diverge and not generalize well. This further highlights the importance of balancing the stochastic Polyak step size and the momentum parameter we provide in Section 2.3.

Figure: Figure 5: Comparison between our proposed MomSPSmax, other momentum methods, SPSmax (no momentum) and constant step-size SGD. Model: ResNet 34, Dataset: CIFAR-10, Batch-size for all methods: 256
Refer to caption: x8.png

Figure: Figure 6: Comparison between our proposed MomSPSmax, other momentum methods, SPSmax (no momentum) and constant step-size SGD. Model: ResNet 34, Dataset: CIFAR-100, Batch-size for all methods: 256
Refer to caption: x9.png

### 5 Conclusion and Future Work

We proposed and theoretically analyzed three novel adaptive Polyak step-size selections for SHB: MomSPSmax, MomDecSPS and MomAdaSPS. We provided convergence guarantees for the three variants for solving convex and smooth optimization problems and verified our theory via numerical experiments. There are many possible interesting extensions of our work: explore the convergence of the proposed methods for non-convex or non-smooth problems, work on adaptive Polyak-type momentum algorithms that can guarantee faster convergence over the no-momentum counterparts, and study extensions in distributed and decentralized settings.

## Supplementary Material

The Supplementary Material is organized as follows: In Appendix A, we have more details on related work on adaptive and momentum methods. In Appendix B, we 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 E, we describe in detail our experimental setup and provide additional experiments.

### Appendix A Further Related Work

##### Adaptive methods.

Two of the first adaptive algorithms are AdaGrad [Duchi et al., ] and RMSProp [Hinton et al., ]. The very popular algorithm Adam [Kingma and Ba, ] was introduced as a momentum extension of RMSProp, and its more recent weight-decay variant AdamW [Loshchilov and Hutter, ] 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., , Ward et al., , Li and Orabona, , Shi et al., , Défossez et al., , Choudhury et al., , Defazio and Jelassi, ].

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, , Prazeres and Oberman, , Berrada et al., ]. Loizou et al. [] was the first work proposing SPSmax and providing strong convergence guarantees in different settings, including strongly convex, convex, and non-convex functions. Further extensions with decreasing variants were proposed in [Orvieto et al., , Jiang and Stich, ]. 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. [], SPS variants for mirror descent presented in D’Orazio et al. [], and SPS for federated learning proposed in Mukherjee et al. []. There are also strong connections between SPS and Model-Based approaches [Asi and Duchi, , , Chadha et al., ]. For further results related to SPS, see Gower et al. [], Li et al. [], Garrigos et al. [] and Abdukhakimov et al. [].

##### Momentum with constant step-size.

Polyak [] 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 $\beta^{*}=(\sqrt{L}-\sqrt{\mu})^{2}/(\sqrt{L}+\sqrt{\mu})^{2}$ and $\gamma^{*}=(1+\sqrt{\beta^{*}})^{2}/L$ where $\mu$ is the strong convexity constant and $L$ the smoothness constant. Ghadimi et al. [] 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 [] 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 $\mu$ and $L$. In the stochastic setting, Liu et al. [] 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. [], among other things, shows that in the smooth and convex setting SHB converges in expectation at the last iterate with a rate of $O(1/T)$ to a neighborhood of the minimum and at a $O(\log T/\sqrt{T})$ rate to the minimum exactly. For these guarantees, one needs knowledge of the problem parameters (such as the smoothness constant), and their results hold for a specific selection of the momentum parameter. Momentum has also been successfully applied in linear systems, see Loizou and Richtárik [] 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 adaptivity can be either the step size or the momentum coefficient.
In the deterministic setting, Barré et al. [] proposed the Accelerated Gradient Method (AGM), an adaptive algorithm built upon AGD that guarantees acceleration over GD. AGM approximates the strong convexity constant $\mu$ by the inverse of the classical Polyak Step-size. However, their algorithm still requires the knowledge of $L$ (but not $\mu$). Furthermore, Saab Jr et al. [] suggests an Adaptive Heavy Ball (AHB) algorithm where it approximates the constants $\mu$ 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 setting, Wang et al. [] propose an adaptive algorithm for a variant of SHB, named ALR-SMAG. It shows that under interpolation, ALR-SMAG enjoys a linear convergence rate for semi-strongly convex and smooth functions. Schaipp et al. [] proposes a new adaptive learning rate that can be combined with any momentum-based method. Finally, Zeng et al. [] proposes adaptive variants for SHB in linear systems solvers.

##### On Technical Assumptions for Convergence.

In the literature of stochastic optimization problems, 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., , Rakhlin et al., , Shamir and Zhang, , Nguyen et al., ] assume bounded gradients, i.e., that is, there exists a $M\in\mathbb{R}$ such that $\operatorname{\mathbb{E}}\|\nabla f_{i}(x)\|^{2}\leq 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., , Gower et al., , ]. A much more relaxed assumption used in the literature is the growth condition on the stochastic gradients, [Bertsekas and Tsitsiklis, , Schmidt and Roux, , Vaswani et al., ]. It states that there exist constants $\rho\in\mathbb{R}$ and $\delta\in\mathbb{R}$ such that $\operatorname{\mathbb{E}}\|\nabla f_{i}(x)\|^{2}\leq\rho\|\nabla f(x)\|^{2}+\delta$. Based on a strong growth condition $(\delta=0)$, Schmidt and Roux [] were the first to establish linear convergence of SGD, with Vaswani et al. [] later showing that SGD can find a first-order stationary point as efficiently as full gradient descent in non-convex settings. Similar conditions have also been proved and used in the analysis of decentralized variants of SGD [Koloskova et al., , Assran et al., ]. More recently, a line of work that uses smoothness (via expected smoothness/residual conditions) to provide closed-form expressions for the values of $\rho$ and $\delta$ of growth condition was able to provide tight convergence analysis of several stochastic algorithms, including SGD [Gower et al., , ], variance reduced methods [Khaled et al., ], stochastic algorithms for min-max optimization [Loizou et al., , , Gorbunov et al., , Choudhury et al., ].

In the convex regime, the original analysis of SGD with SPS of Loizou et al. [] was one of the first papers that did not require any additional assumptions to guarantee convergence for SGD. Following the convergence guarantees of Loizou et al. [], we highlight that our proposed analysis of SHB with Polyak step-size does not require any additional assumptions for guaranteeing convergence. 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. [] makes the strong assumption of bounded variance and assumes knowledge of the smoothness $L$ for tuning the parameters. Another work on theoretical guarantees for SHB is Sebbouh et al. [], where a last iterate convergence is proved for a specific choice of $\gamma_{t}$ and $\beta_{t}$, but also requires knowledge of $L$ and only achieves rate $O(\log T/\sqrt{T})$ for the decreasing step-size variant. Furthermore, the authors of Wang et al. [] provide guarantees for a Polyak-based adaptive SHB method for the deterministic setting and the stochastic interpolated regime for convex and strongly convex objectives. Finally, Schaipp et al. [] also assumes a stochastic convex and interpolated regime for their results.

### Appendix B Pseudo-codes

In this section, we include the pseudo-codes of the SHB with the proposed stochastic Polyak step-sizes for a better understanding of the methods and easier comparison with other algorithms.

##### SHB with MomSPSmax.

We start with the pseudo-code for SHB with MomSPSmax. In Theorem 3.1 we established convergence to a neighborhood for $\beta\in\left[0,\frac{\alpha}{2\gamma_{b}-\alpha}\right)$ where $\alpha=\min\left\{\frac{1}{2L_{\max}},\gamma_{b}\right\}$ and for $c=1$. However, in practice, it seems to have convergence even for different choices of $c$ or even if $\beta$ is chosen outside of this bound, so we provide the pseudo-code for a general $\beta\in[0,1)$ and $c$. Thus, the user must provide the momentum coefficient $\beta$, the upper bound $\gamma_{b}$, the constant $c$, and the lower bounds $\ell_{S_{t}}^{*}$. In practice, the usual choices are $\beta=0.9$, $c=1$, and $\ell_{S_{t}}^{*}=0$.

Figure: Algorithm 1 SHB with MomSPSmax

##### SHB with MomDecSPS.

For the decreasing step-sizes, Theorems 3.5 and 3.6 make no assumption on $\beta\in[0,1)$. For SHB with MomDecSPS, the user needs to provide the momentum coefficient $\beta$, the upper bound $\gamma_{b}$, the constant $c$ and the lower bounds $\ell_{S_{t}}$, like the previous algorithm. In practice, the usual choices are $c=1$ and $\ell_{S_{t}}^{*}=0$. Note that when $t=0$, the step-size is equal to

$$ $\displaystyle\gamma_{0}=\min\left\{\frac{(1-\beta)[f_{S_{0}}(x^{0})-\ell_{S_{0 }}^{*}]}{c\|\nabla f_{S_{0}}(x^{0})\|^{2}},\frac{\gamma_{b}c}{c}\right\}=(1- \beta)\min\left\{\frac{f_{S_{0}}(x^{0})-\ell_{S_{0}}^{*}}{c\|\nabla f_{S_{0}}( x^{0})\|^{2}},\gamma_{b}\right\},$ $$

which is equal to MomSPSmax for $t=0$.

Figure: Algorithm 2 SHB with MomDecSPS

##### SHB with MomAdaSPS.

Finally, for SHB with MomAdaSPS the requirements are less since the user only needs to provide the momentum coefficient $\beta$, the constant $c$ and the lower bounds $\ell_{S_{t}}^{*}$. Furthermore, according to Jiang and Stich [] one can set $c=\frac{1}{\sqrt{f_{S_{0}}(x^{0})-\ell_{S_{0}}^{*}}}$ after randomly choosing $S_{0}\subseteq[n]$ in the first iteration. Moreover, the minimum with respect to $\gamma_{t-1}$ is to ensure that the step-size is decreasing.

Figure: Algorithm 3 SHB with MomAdaSPS

Note that for all the above algorithms, one needs only lower bounds $\ell_{S_{t}}^{*}$, which can be chosen equal to $\ell_{S_{t}}^{*}=0$ in practice. Of course, if the true infima $f_{S_{t}}^{*}=\inf_{x\in\mathbb{R}^{d}}f_{S_{t}}$ are known, then one should use them because it will lead to better convergence in terms of the neighborhood. Furthermore, notice that we start by initializing $x^{-1}=x^{0}$, where $x^{0}$ is the given starting point. This is because the first step of SHB (i.e. for $t=0$) is just an iteration of SGD because we have no previous iterate $x^{-1}$, that is $x^{1}=x^{0}-\gamma_{0}\nabla f_{S_{0}}(x^{0})+\beta(x^{0}-x^{-1})=x^{0}-\gamma
_{0}\nabla f_{S_{0}}(x^{0})$.

### Appendix C Technical Preliminaries

#### C.1 Basic Definitions

In this section, we present some basic definitions we use throughout the paper.

###### Definition C.1 (Convexity) .

A differentiable function $f:\mathbb{R}^{d}\to\mathbb{R}$ is convex if

$$ $\displaystyle f(x)\geq f(y)+\langle\nabla f(y),x-y\rangle,$ $$

for all $x,y\in\mathbb{R}^{d}$.

###### Definition C.2 ( $L$ -smooth) .

A differentiable function $f:\mathbb{R}^{d}\to\mathbb{R}$ is $L$-smooth if there exists a constant $L>0$ such that

$$ $\displaystyle\|\nabla f(x)-\nabla f(y)\|\leq L\|x-y\|,$ $$

or equivalently

$$ $\displaystyle f(x)-f(y)\leq\langle\nabla f(y),x-y\rangle+\frac{L}{2}\|x-y\|^{2}$ $$

for all $x,y\in\mathbb{R}^{d}$.

###### Definition C.1 (Convexity) .

###### Definition C.2 ( L 𝐿 L italic_L -smooth) .

#### C.2 Basic Lemmas

Here we 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_{i}$ is $L_{i}$-smooth, then the step-size of SPSmax satisfies:

$$ $\displaystyle\alpha=\min\left\{\frac{1}{2cL_{\max}},\gamma_{b}\right\}\leq \gamma_{t}\leq\gamma_{b}\text{ and }\gamma_{t}^{2}\|\nabla f_{S_{t}}(x^{t})\|^ {2}\leq\gamma_{t}[f_{S_{t}}(x^{t})-\ell_{S_{t}}^{*}],$ (2) $$

where $L_{\max}=\max_{i\in[n]}\{L_{i}\}$.

###### Lemma C.4 (Lemma 1, [Orvieto et al., 2022 ] ) .

Suppose each $f_{i}$ is $L_{i}$-smooth, then the step-size of DecSPS satisfies

$$ $\displaystyle\min\left\{\frac{1}{2c_{t}L_{\max}},\frac{c_{0}\gamma_{b}}{c_{t}} \right\}\leq\gamma_{t}\leq\frac{c_{0}\gamma_{b}}{c_{t}}\text{ and }\gamma_{t-1 }\leq\gamma_{t}\text{ and }\gamma_{t}^{2}\|\nabla f_{S_{t}}(x^{t})\|^{2}\leq \frac{\gamma_{t}}{c_{t}}[f_{S_{t}}(x^{t})-\ell_{S_{t}}^{*}],$ (3) $$

where $L_{\max}=\max_{i\in[n]}\{L_{i}\}$.

###### Lemma C.5 (Lemma 12, [Jiang and Stich, 2023 ] ) .

For any non-negative sequence $a_{0},\dots,a_{T}$ , the following holds:

$$ $\displaystyle\sqrt{\sum_{t=0}^{T}a_{t}}\leq\sum_{t=0}^{T}\frac{a_{t}}{\sqrt{ \sum_{s=0}^{t}a_{s}}}\leq 2\sqrt{\sum_{t=0}^{T}a_{t}}.$ (4) $$

If $a_{0}\geq 1$, then the following holds:

$$ $\displaystyle\sum_{t=0}^{T}\frac{a_{t}}{\sqrt{\sum_{s=0}^{t}a_{s}}}\leq\log \left(\sum_{t=0}^{T}a_{t}\right)+1.$ (5) $$

###### Lemma C.6 (Lemma 13, [Jiang and Stich, 2023 ] ) .

If $x^{2}\leq a(x+b)$ for $a\geq 0$ and $b\geq 0$, then it holds that $x\leq a+\sqrt{ab}$.

###### Lemma C.7 (Lemma 16, [Jiang and Stich, 2023 ] ) .

Suppose each $f_{i}$ is $L_{i}$-smooth, then the step-size of AdaSPS satisfies

$$ $\displaystyle\frac{1}{2cL_{\max}}\frac{1}{\sqrt{\sum_{s=0}^{t}f_{S_{s}}(x^{s}) -\ell_{S_{s}}^{*}}}\leq\eta_{t}\leq\frac{f_{S_{t}}(x^{t})-\ell_{S_{t}}^{*}}{c \|\nabla f_{S_{t}}(x^{t})\|^{2}}\frac{1}{\sqrt{\sum_{s=0}^{t}f_{S_{s}}(x^{s})- \ell_{S_{s}}^{*}}},$ (6) $$

where $L_{\max}=\max_{i\in[n]}\{L_{i}\}$.

###### Lemma C.3 ( Loizou et al. [ 2021 ] ) .

###### Lemma C.4 (Lemma 1, [Orvieto et al., 2022 ] ) .

###### Lemma C.5 (Lemma 12, [Jiang and Stich, 2023 ] ) .

###### Lemma C.6 (Lemma 13, [Jiang and Stich, 2023 ] ) .

###### Lemma C.7 (Lemma 16, [Jiang and Stich, 2023 ] ) .

#### C.3 Iterate Moving Average

As mentioned in the main paper, the iterates of the SHB method are equal to the IMA method. The following proposition is from Sebbouh et al. []. We include the proof for completeness.

###### Proposition C.8 (Proposition 1.6, [Sebbouh et al., 2021 ] ) .

Let $\eta_{t}>0$ and $\lambda_{t}\geq 0$. Consider the Iterate Moving Average (IMA) method given by the following update rule:

$$ $\displaystyle z^{t+1}=z^{t}-\eta_{t}\nabla f_{S_{t}}(x^{t})$ (7) $\displaystyle x^{t+1}=\frac{\lambda_{t+1}}{\lambda_{t+1}+1}x^{t}+\frac{1}{ \lambda_{t+1}+1}z^{t+1},$ (8) $$

where $z^{0}=x^{0}$. If the following equations hold for any $t\in\mathbb{N}$

$$ $\displaystyle 1+\lambda_{t+1}=\frac{\lambda_{t}}{\beta_{t}}$ (9) $\displaystyle\eta_{t}=(1+\lambda_{t+1})\gamma_{t},$ (10) $$

then the $x_{t}$ iterates of the IMA method are equal to the $x_{t}$ iterates produced by the SHB method.

###### Proof.

We will start from the IMA update rule and we will prove that its $(x^{t})$ iterates satisfy the SHB update rule. Suppose that we have the $(z^{t+1})$ and $(x^{t+1})$ iterates as well as the $(\lambda_{t})$ and $(\eta_{t})$ sequences. Define the sequences $(\beta_{t})$ and $(\gamma_{t})$ using eqs. 9 and 10. We will show that $(x^{t})$ iterates satisfy the SHB update rule.
Solving for $z^{t+1}$ in eq. 8 we get

$$ $\displaystyle z^{t+1}:=x^{t+1}+\lambda_{t+1}(x^{t+1}-x^{t})$ (11) $$

Now substituting eqs. 7 and 11 in eq. 8 we have

$$ $\displaystyle x^{t+1}$ $\displaystyle=\frac{\lambda_{t+1}}{\lambda_{t+1}+1}x^{t}+\frac{1}{\lambda_{t+1 }+1}z^{t+1}$ $\displaystyle\overset{(\ref{eq:ima-1})}{=}\frac{\lambda_{t+1}}{\lambda_{t+1}+1 }x^{t}+\frac{1}{\lambda_{t+1}+1}\left(z^{t}-\eta_{t}\nabla f_{S_{t}}(x^{t})\right)$ $\displaystyle\overset{(\ref{eq:z-def})}{=}\frac{\lambda_{t+1}}{\lambda_{t+1}+1 }x^{t}+\frac{1}{\lambda_{t+1}+1}\left(x^{t}+\lambda_{t}(x^{t}-x^{t-1})-\eta_{t }\nabla f_{S_{t}}(x^{t}))\right)$ $\displaystyle=x^{t}-\frac{\eta_{t}}{\lambda_{t+1}+1}\nabla f_{S_{t}}(x^{t})+ \frac{\lambda_{t}}{\lambda_{+t+1}+1}(x^{t}-x^{t-1})$ $\displaystyle\overset{(\ref{eq:ima-hb-equiv-1},\ref{eq:ima-hb-equiv-2})}{=}x^{ t}-\gamma_{t}\nabla f_{S_{t}}(x^{t})+\beta_{t}(x^{t}-x^{t-1}),$ $$

as wanted.
∎

The above proposition essentially states that momentum is a convex combination of SGD under a certain transformation. In this paper, we work only with a constant momentum coefficient, i.e., $\beta_{t}=\beta$. For this case, we also assume constant $\lambda_{t}$, and we have Proposition 2.1 and Corollary 2.2.

###### Proof of Proposition 2.1 .

Setting $\lambda_{t}=\lambda$ in Equation 9 we have

$$ $\displaystyle 1+\lambda=\frac{\lambda}{\beta_{t}}\Rightarrow\beta_{t}=\frac{ \lambda}{1+\lambda}=:\beta.$ $$

Then solving for $\lambda$ we get $\lambda=\frac{\beta}{1-\beta}$. Hence

$$ $\displaystyle\gamma_{t}$ $\displaystyle=\frac{\eta_{t}}{1+\lambda}$ $\displaystyle=\frac{\eta_{t}}{1+\frac{\beta}{1-\beta}}$ $\displaystyle=(1-\beta)\eta_{t}$ $$

∎

###### Proof of Corollary 2.2 .

- •
Let $\eta_{t}=\textnormal{SPS}_{\max}$. Then
$\displaystyle\gamma_{t}$
$\displaystyle=(1-\beta)\eta_{t}$
$\displaystyle=(1-\beta)\min\left\{\frac{f_{S_{t}}(x^{t})-\ell_{S_{t}}^{*}}{c\|
\nabla f_{S_{t}}(x^{t})\|^{2}},\eta_{b}\right\}$
$\displaystyle=(1-\beta)\min\left\{\frac{f_{S_{t}}(x^{t})-\ell_{S_{t}}^{*}}{c\|
\nabla f_{S_{t}}(x^{t})\|^{2}},\gamma_{b}\right\}.$
Here $\gamma_{b}=\eta_{b}$.
- •
Let $\eta_{t}=\textnormal{DecSPS}$. Then
$\displaystyle\gamma_{t}$
$\displaystyle=(1-\beta)\eta_{t}$
$\displaystyle=(1-\beta)\min\left\{\frac{f_{S_{t}}(x^{t})-\ell_{S_{t}}^{*}}{c_{
t}\|\nabla f_{S_{t}}(x^{t})\|^{2}},\frac{\eta_{t-1}c_{t-1}}{c_{t}}\right\}$
$\displaystyle=\min\left\{\frac{(1-\beta)[f_{S_{t}}(x^{t})-\ell_{S_{t}}^{*}]}{c
_{t}\|\nabla f_{S_{t}}(x^{t})\|^{2}},\frac{(1-\beta)\eta_{t-1}c_{t-1}}{c_{t}}\right\}$
$\displaystyle=\min\left\{\frac{(1-\beta)[f_{S_{t}}(x^{t})-\ell_{S_{t}}^{*}]}{c
_{t}\|\nabla f_{S_{t}}(x^{t})\|^{2}},\frac{\gamma_{t-1}c_{t-1}}{c_{t}}\right\},$
since $\gamma_{t-1}=(1-\beta)\eta_{t-1}$ from Proposition 2.1.
- •
Let $\eta_{t}=\textnormal{AdaSPS}$. Then
$\displaystyle\gamma_{t}$
$\displaystyle=(1-\beta)\eta_{t}$
$\displaystyle=(1-\beta)\min\left\{\frac{f_{S_{t}}(x^{t})-\ell_{S_{t}}^{*}}{c\|
\nabla f_{S_{t}}(x^{t})\|^{2}\sqrt{\sum_{s=0}^{t}f_{S_{s}}(x^{s})-\ell_{S_{s}}
^{*}}},\eta_{t-1}\right\}$
$\displaystyle=\min\left\{\frac{(1-\beta)[f_{S_{t}}(x^{t})-\ell_{S_{t}}^{*}]}{c
\|\nabla f_{S_{t}}(x^{t})\|^{2}\sqrt{\sum_{s=0}^{t}f_{S_{s}}(x^{s})-\ell_{S_{s
}}^{*}}},(1-\beta)\eta_{t-1}\right\}$
$\displaystyle=\min\left\{\frac{(1-\beta)[f_{S_{t}}(x^{t})-\ell_{S_{t}}^{*}]}{c
\|\nabla f_{S_{t}}(x^{t})\|^{2}\sqrt{\sum_{s=0}^{t}f_{S_{s}}(x^{s})-\ell_{S_{s
}}^{*}}},\gamma_{t-1}\right\}.$

∎

###### Proposition C.8 (Proposition 1.6, [Sebbouh et al., 2021 ] ) .

###### Proof.

###### Proof of Proposition 2.1 .

###### Proof of Corollary 2.2 .

#### C.4 Connection of MomSPS max and SPS max

In this section, we explore the connection between the MomSPSmax and the SPSmax step-sizes. In particular, one can view the MomSPSmax step-size as the SPSmax step-size with $c=1/(1-\beta)$. Here, we 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

$$ $\displaystyle\gamma_{t}^{SPSmax}=\min\left\{\frac{f_{S_{t}}(x^{t})-f_{S_{t}}^{ *}}{c_{SPSmax}\|\nabla f_{S_{t}}(x^{t})\|^{2}},\gamma_{b}^{SPSmax}\right\}$ $$

and

$$ $\displaystyle\gamma_{t}^{MomSPSmax}=(1-\beta)\min\left\{\frac{f_{S_{t}}(x^{t}) -f_{S_{t}}^{*}}{c_{MomSPSmax}\|\nabla f_{S_{t}}(x^{t})\|^{2}},\gamma_{b}^{ MomSPSmax}\right\}$ $$

so we have $\gamma^{t}_{SPSmax}=\gamma^{t}_{MomSPSmax}$ if and only if

$$ $\displaystyle\frac{1-\beta}{c_{MomSPSmax}}=\frac{1}{c_{SPSmax}}\text{ and } \gamma_{b}^{SPSmax}=(1-\beta)\gamma_{b}^{MomSPSmax}.$ $$

Without loss of generalization, we can assume that $c_{MomSPSmax}=1$ (this is also what Theorem 3.1 guarantees). In this case we get $\beta=1-\frac{1}{c_{SPSmax}}$. Hence, we can restate Theorem 3.1 as follows:

###### Theorem C.9 .

Assume that each $f_{i}$ is convex and $L_{i}$-smooth. Then, the iterates of SHB with SPSmax with $c\in\left[1,\frac{2\gamma_{b}-2\alpha}{2\gamma_{b}-\alpha}\right)$ and $\beta=1-\frac{1}{c}$ where $\alpha=\min\left\{\frac{1}{2L_{\max}},\gamma_{b}\right\}$ and $L_{\max}=\max_{i\in[n]}\{L_{i}\}$, converge as

$$ $\operatorname{\mathbb{E}}[f(\overline{x}^{T})-f(x^{*})]\leq\frac{C_{1}\|x^{0}- x^{*}\|^{2}}{T}+C_{2}\sigma^{2},$ $$

where $\overline{x}^{T}=\frac{1}{T}\sum_{t=0}^{T-1}x^{t}$ is the Cesaro average and the constants $C_{1}=\frac{1-\beta}{\alpha\beta+\alpha-2\beta\gamma_{b}}$ and $C_{2}=\frac{2\gamma_{b}-\alpha\beta-\alpha}{\alpha\beta+\alpha-2\beta\gamma_{b}}$.

This means that MomSPSmax can be viewed as a special case of SPSmax. However, note that the update rules are different since in SHB we have the extra term $+\beta(x^{t}-x^{t-1})$ so a different analysis is needed.

The same line of thought can be applied to other works related to HB. For example, in Ghadimi et al. [], the first global guarantees for HB were established. The authors show that for a convex and $L$-smooth function $f$ HB with any $\beta\in[0,1)$ and $0<\gamma<\frac{2(1-\beta)}{L}$ converges with a rate of $O(1/T)$. Also, recall that GD for convex and $L$-smooth function converges for $0<\gamma<\frac{2}{L}$ with rate $O(1/T)$. Hence, we can restate the result of Ghadimi et al. [] as follows: Let $f$ be a convex and $L$-smooth function. Then HB with $\gamma=\frac{c}{L}$, where $c\in(0,2)$, and $\beta=1-\frac{c}{2}$ converges with rate $O(1/T)$.

Similarly, for SGD in the convex and smooth setting, we are guaranteed convergence $O(1/T)$ to a neighborhood of the solution when $\gamma\leq\frac{1}{L}$. In Liu et al. [], the authors guarantee that SHB with any $\beta\in(0,1)$ and $\gamma\leq\frac{(1-\beta)^{2}}{L}\min\left\{\frac{1}{4-\beta+\beta^{2}},\frac{
1}{2\sqrt{2\beta^{2}+2\beta}}\right\}$ converges with rate $O(1/T)$ to a neighborhood. This can be restated as follows: SHB with $\gamma=\frac{c}{L}$, where $c\in(0,1)$, and $\beta$ given by the solution of $c=\min\left\{\frac{(1-\beta)^{2}}{4-\beta+\beta^{2}},\frac{(1-\beta)^{2}}{2
\sqrt{2\beta^{2}+2\beta}}\right\}$, converges with rate $O(1/T)$ to a neighborhood.

###### Theorem C.9 .

### Appendix 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 MomSPSmax, MomDecSPS and MomAdaSPS for convex and smooth functions $f_{i}$ and $f$ of Problem 1. The idea for all the proofs is that we firstly show the corresponding result in the IMA setting, and then we transfer it to 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^{t-1}$. This needs new algebraic tricks in order to use convexity, and after that, one needs to deal with the two quantities $f(x^{t})$ and $f(x^{t-1})$ at the same time. We deal with this using the IMA framework, similarly with Sebbouh et al. []. 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. [] the authors choose the momentum coefficients $\lambda_{t}$ in such a way that the main sum telescopes and the final result only has the expectation of the last iterate. In our proof, we make no such simplifications, and instead, we use Jensen’s inequality to finish. This approach allows us to have more freedom in the selection of $\lambda_{t}$. This is more prominent in the proofs of the decreasing variants.

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

#### D.1 Proof of Theorem 3.1

First, we state and prove the corresponding theorem for SPSmax in the IMA setting.

###### Theorem D.1 .

Assume that each $f_{i}$ is convex and $L_{i}$-smooth. Then, the iterates of IMA with

$$ $\displaystyle\eta_{t}=\min\left\{\frac{f_{S_{t}}(x^{t})-\ell_{S_{t}}^{*}}{\| \nabla f_{S_{t}}(x^{t})\|^{2}},\eta_{b}\right\}\text{ and }\lambda_{t}=\lambda \in\left[0,\frac{\alpha}{2(\eta_{b}-\alpha)}\right),$ $$

where $\alpha=\min\left\{\frac{1}{2L_{\max}},\eta_{b}\right\}$ and $L_{\max}=\max_{i\in[n]}\{L_{i}\}$, converge as

$$ $\displaystyle\operatorname{\mathbb{E}}[f(\overline{x}^{T})-f(x^{*})]\leq\frac{ \|x^{0}-x^{*}\|^{2}}{(2\alpha\lambda+\alpha-2\eta_{b}\lambda)T}+\frac{(2\eta_{ b}+2\eta_{b}\lambda-2\alpha\lambda-\alpha)\sigma^{2}}{2\alpha\lambda+\alpha-2 \eta_{b}\lambda},$ $$

where $\overline{x}^{T}=\frac{1}{T}\sum_{t=0}^{T-1}x^{t}$ and $\sigma^{2}=\operatorname{\mathbb{E}}[f_{i}(x^{*})-f_{i}^{*}]$.

###### Proof.

We have

$$ $\displaystyle\|z^{t+1}-x^{*}\|^{2}$ $\displaystyle=\|z^{t}-x^{*}\|^{2}-2\eta_{t}\langle\nabla f_{S_{t}}(x^{t}),z^{t }-x^{*}\rangle+\eta_{t}^{2}\|\nabla f_{S_{t}}(x^{t})\|^{2}$ $\displaystyle\overset{(\ref{eq:z-def})}{=\joinrel=}\|z^{t}-x^{*}\|^{2}-2\eta_{ t}\langle\nabla f_{S_{t}}(x^{t}),x^{t}-x^{*}\rangle-2\eta_{t}\lambda_{t} \langle\nabla f_{S_{t}}(x^{t}),x^{t}-x^{t-1}\rangle$ $\displaystyle\quad+\eta_{t}^{2}\|\nabla f_{S_{t}}(x^{t})\|^{2}$ $\displaystyle\overset{convexity}{\leq}\|z^{t}-x^{*}\|^{2}-2\eta_{t}[f_{S_{t}}( x^{t})-f_{S_{t}}(x^{*})]-2\eta_{t}\lambda_{t}[f_{S_{t}}(x^{t})-f_{S_{t}}(x^{t- 1})]$ $\displaystyle\quad+\eta_{t}^{2}\|\nabla f_{S_{t}}(x^{t})\|^{2}$ $\displaystyle\overset{(\ref{eq:lem:sps-bound})}{\leq}\|z^{t}-x^{*}\|^{2}-2\eta _{t}[f_{S_{t}}(x^{t})-f_{S_{t}}(x^{*})]-2\eta_{t}\lambda_{t}[f_{S_{t}}(x^{t})- f_{S_{t}}(x^{t-1})]$ $\displaystyle\quad+\eta_{t}[f_{S_{t}}(x^{t})-\ell_{S_{t}}^{*}]$ $\displaystyle=\|z^{t}-x^{*}\|^{2}-2\eta_{t}[f_{S_{t}}(x^{t})-\ell_{S_{t}}^{*}] +2\eta_{t}[f_{S_{t}}(x^{*})-\ell_{S_{t}}^{*}]-2\eta_{t}\lambda_{t}[f_{S_{t}}(x ^{t})-\ell_{S_{t}}^{*}]$ $\displaystyle\quad+2\eta_{t}\lambda_{t}[f_{S_{t}}(x^{t-1})-\ell_{S_{t}}^{*}]+ \eta_{t}[f_{S_{t}}(x^{t})-\ell_{S_{t}}^{*}]$ $\displaystyle=\|z^{t}-x^{*}\|^{2}-2\eta_{t}\left(1+\lambda_{t}-\frac{1}{2} \right)[f_{S_{t}}(x^{t})-\ell_{S_{t}}^{*}]+2\eta_{t}\lambda_{t}[f_{S_{t}}(x^{t -1})-\ell_{S_{t}}^{*}]$ $\displaystyle\quad+2\eta_{t}[f_{S_{t}}(x^{*})-\ell_{S_{t}}^{*}]$ $\displaystyle\overset{(\ref{eq:lem:sps-bound})}{\leq}\|z^{t}-x^{*}\|^{2}-2 \alpha\left(\lambda_{t}+\frac{1}{2}\right)[f_{S_{t}}(x^{t})-\ell_{S_{t}}^{*}]+ 2\eta_{b}\lambda_{t}[f_{S_{t}}(x^{t-1})-\ell_{S_{t}}^{*}]$ $\displaystyle\quad+2\eta_{b}[f_{S_{t}}(x^{*})-\ell_{S_{t}}^{*}]$ $\displaystyle\leq\|z^{t}-x^{*}\|^{2}-2\alpha\left(\lambda_{t}+\frac{1}{2} \right)[f_{S_{t}}(x^{t})-f_{S_{t}}(x^{*})]-2\alpha\left(\lambda_{t}+\frac{1}{2 }\right)[f_{S_{t}}(x^{*})-\ell_{S_{t}}^{*}]$ $\displaystyle\quad+2\eta_{b}\lambda_{t}[f_{S_{t}}(x^{t-1})-f_{S_{t}}(x^{*})]+2 \eta_{b}\lambda_{t}[f_{S_{t}}(x^{*})-\ell_{S_{t}}^{*}]+2\eta_{b}[f_{S_{t}}(x^{ *})-\ell_{S_{t}}^{*}].$ $$

Now take expectation conditional on $x^{t}$ to get

$$ $\displaystyle\operatorname{\mathbb{E}}_{S_{t}}\|z^{t+1}-x^{*}\|^{2}$ $\displaystyle\leq\operatorname{\mathbb{E}}_{S_{t}}\|z^{t}-x^{*}\|^{2}-2\alpha \left(\lambda_{t}+\frac{1}{2}\right)\operatorname{\mathbb{E}}_{S_{t}}[f_{S_{t} }(x^{t})-f_{S_{t}}(x^{*})]$ $\displaystyle\quad-2\alpha\left(\lambda_{t}+\frac{1}{2}\right)\operatorname{ \mathbb{E}}_{S_{t}}[f_{S_{t}}(x^{*})-\ell_{S_{t}}^{*}]+2\eta_{b}\lambda_{t} \operatorname{\mathbb{E}}_{S_{t}}[f_{S_{t}}(x^{t-1})-f_{S_{t}}(x^{*})]$ $\displaystyle\quad+2\eta_{b}\lambda_{t}\operatorname{\mathbb{E}}_{S_{t}}[f_{S_ {t}}(x^{*})-\ell_{S_{t}}^{*}]+2\eta_{b}\operatorname{\mathbb{E}}_{S_{t}}[f_{S_ {t}}(x^{*})-\ell_{S_{t}}^{*}]$ $\displaystyle=\|z^{t}-x^{*}\|^{2}-2\alpha\left(\lambda_{t}+\frac{1}{2}\right)[ f(x^{t})-f(x^{*})]-2\alpha\left(\lambda_{t}+\frac{1}{2}\right)\sigma^{2}$ $\displaystyle\quad+2\eta_{b}\lambda_{t}[f(x^{t-1})-f(x^{*})]+2\eta_{b}\lambda_ {t}\sigma^{2}+2\eta_{b}\sigma^{2}$ $\displaystyle=\|z^{t}-x^{*}\|^{2}-2\alpha\left(\lambda_{t}+\frac{1}{2}\right)[ f(x^{t})-f(x^{*})]+2\eta_{b}\lambda_{t}[f(x^{t-1})-f(x^{*})]$ $\displaystyle\quad+2\left(\eta_{b}\left(1+\lambda_{t}\right)-\alpha\left( \lambda_{t}+\frac{1}{2}\right)\right)\sigma^{2}.$ (12) $$

Note that in eq. 12 the coefficient $2\left(\eta_{b}\left(1+\lambda_{t}\right)-\alpha\left(\lambda_{t}+\frac{1}{2}
\right)\right)$ of $\sigma^{2}$ is positive since $\eta_{b}\geq\alpha$. Taking expectation again and using the tower property we have

$$ $\displaystyle 2\alpha\left(\lambda_{t}+\frac{1}{2}\right)\operatorname{\mathbb {E}}[f(x^{t})-f(x^{*})]-2\eta_{b}\lambda_{t}\operatorname{\mathbb{E}}[f(x^{t-1 })-f(x^{*})]$ $\displaystyle\leq\operatorname{\mathbb{E}}\|z^{t}-x^{*}\|^{2}-\operatorname{ \mathbb{E}}\|z^{t+1}-x^{*}\|^{2}+2\left(\eta_{b}\left(1+\lambda_{t}\right)- \alpha\left(\lambda_{t}+\frac{1}{2}\right)\right)\sigma^{2}.$ (13) $$

Summing eq. 13 for $t=0,\dots,T-1$ and telescoping we get

$$ $\displaystyle\sum_{t=0}^{T-1}\left[2\alpha\left(\lambda_{t}+\frac{1}{2}\right) \operatorname{\mathbb{E}}[f(x^{t})-f(x^{*})]-2\eta_{b}\lambda_{t}\operatorname {\mathbb{E}}[f(x^{t-1})-f(x^{*})]\right]$ $\displaystyle\leq\sum_{t=0}^{T-1}\left[\operatorname{\mathbb{E}}\|z^{t}-x^{*} \|^{2}-\operatorname{\mathbb{E}}\|z^{t+1}-x^{*}\|^{2}\right]+\sum_{t=0}^{T-1}2 \left(\eta_{b}\left(1+\lambda_{t}\right)-\alpha\left(\lambda_{t}+\frac{1}{2} \right)\right)\sigma^{2}$ $\displaystyle=\operatorname{\mathbb{E}}\|z^{0}-x^{*}\|^{2}-\operatorname{ \mathbb{E}}\|z^{T}-x^{*}\|^{2}+\sum_{t=0}^{T-1}2\left(\eta_{b}\left(1+\lambda_ {t}\right)-\alpha\left(\lambda_{t}+\frac{1}{2}\right)\right)\sigma^{2}$ $\displaystyle\overset{z^{0}=x^{0}}{\leq}\|x^{0}-x^{*}\|^{2}+\sum_{t=0}^{T-1}2 \left(\eta_{b}\left(1+\lambda_{t}\right)-\alpha\left(\lambda_{t}+\frac{1}{2} \right)\right)\sigma^{2}.$ (14) $$

The LHS of eq. 14 can be bounded as follows (using $x^{-1}=x^{0}$)

$$ $\displaystyle\sum_{t=0}^{T-1}\left[2\alpha\left(\lambda_{t}+\frac{1}{2}\right) \operatorname{\mathbb{E}}[f(x^{t})-f(x^{*})]-2\eta_{b}\lambda_{t}\operatorname {\mathbb{E}}[f(x^{t-1})-f(x^{*})]\right]$ $\displaystyle=\sum_{t=0}^{T-1}2\alpha\left(\lambda_{t}+\frac{1}{2}\right) \operatorname{\mathbb{E}}[f(x^{t})-f(x^{*})]-\sum_{t=0}^{T-1}2\eta_{b}\lambda_ {t}\operatorname{\mathbb{E}}[f(x^{t-1})-f(x^{*})]$ $\displaystyle=\sum_{t=0}^{T-1}2\alpha\left(\lambda_{t}+\frac{1}{2}\right) \operatorname{\mathbb{E}}[f(x^{t})-f(x^{*})]-\sum_{t=0}^{T-2}2\eta_{b}\lambda_ {t+1}\operatorname{\mathbb{E}}[f(x^{t})-f(x^{*})]$ $\displaystyle=2\alpha\left(\lambda_{T-1}+\frac{1}{2}\right)\operatorname{ \mathbb{E}}[f(x^{T-1})-f(x^{*})]+\sum_{t=0}^{T-2}2\alpha\left(\lambda_{t}+ \frac{1}{2}\right)\operatorname{\mathbb{E}}[f(x^{t})-f(x^{*})]-\sum_{t=0}^{T-2 }2\eta_{b}\lambda_{t+1}\operatorname{\mathbb{E}}[f(x^{t})-f(x^{*})]$ $\displaystyle=2\alpha\left(\lambda_{T-1}+\frac{1}{2}\right)\operatorname{ \mathbb{E}}[f(x^{T-1})-f(x^{*})]+\sum_{t=0}^{T-2}\left[2\alpha\left(\lambda_{t }+\frac{1}{2}\right)-2\eta_{b}\lambda_{t+1}\right]\operatorname{\mathbb{E}}[f( x^{t})-f(x^{*})]$ $\displaystyle\geq\left[2\alpha\left(\lambda_{T-1}+\frac{1}{2}\right)-2\eta_{b} \lambda_{T}\right]\operatorname{\mathbb{E}}[f(x^{T-1})-f(x^{*})]+\sum_{t=0}^{T -2}\left[2\alpha\left(\lambda_{t}+\frac{1}{2}\right)-2\eta_{b}\lambda_{t+1} \right]\operatorname{\mathbb{E}}[f(x^{t})-f(x^{*})]$ $\displaystyle=\sum_{t=0}^{T-1}\left[2\alpha\left(\lambda_{t}+\frac{1}{2}\right )-2\eta_{b}\lambda_{t+1}\right]\operatorname{\mathbb{E}}[f(x^{t})-f(x^{*})].$ (15) $$

Combining eq. 14 and eq. 15 we have

$$ $\displaystyle\sum_{t=0}^{T-1}\left[2\alpha\left(\lambda_{t}+\frac{1}{2}\right) -2\eta_{b}\lambda_{t+1}\right]\operatorname{\mathbb{E}}[f(x^{t})-f(x^{*})]\leq \|x^{0}-x^{*}\|^{2}+\sum_{t=0}^{T-1}2\left(\eta_{b}\left(1+\lambda_{t}\right)- \alpha\left(\lambda_{t}+\frac{1}{2}\right)\right)\sigma^{2}.$ (16) $$

In eq. 16, setting $\lambda_{t}=\lambda,\forall t>0$ with $0\leq\lambda<\frac{\alpha}{2(\eta_{b}-\alpha)}$ (in order to ensure that quantity $2\alpha\left(\lambda+\frac{1}{2}\right)-2\eta_{b}\lambda$ in the LHS is positive) we get

$$ $\displaystyle\left[2\alpha\left(\lambda+\frac{1}{2}\right)-2\eta_{b}\lambda \right]\sum_{t=0}^{T-1}\operatorname{\mathbb{E}}[f(x^{t})-f(x^{*})]\leq\|x^{0} -x^{*}\|^{2}+2T\left(\eta_{b}\left(1+\lambda\right)-\alpha\left(\lambda+\frac{ 1}{2}\right)\right)\sigma^{2}.$ (17) $$

Now dividing eq. 17 by $T\left[2\alpha\left(\lambda+\frac{1}{2}\right)-2\eta_{b}\lambda\right]$ and using Jensen’s inequality we get

$$ $\displaystyle\operatorname{\mathbb{E}}[f(\overline{x}^{T})-f(x^{*})]$ $\displaystyle\leq\frac{1}{T}\sum_{t=0}^{T-1}\operatorname{\mathbb{E}}[f(x^{t}) -f(x^{*})]$ $\displaystyle\leq\frac{\|x^{0}-x^{*}\|^{2}+T\left(2\eta_{b}+2\eta_{b}\lambda-2 \alpha\lambda-\alpha\right)\sigma^{2}}{(2\alpha\lambda+\alpha-2\eta_{b}\lambda )T}$ $\displaystyle=\frac{\|x^{0}-x^{*}\|^{2}}{(2\alpha\lambda+\alpha-2\eta_{b} \lambda)T}+\frac{2\eta_{b}+2\eta_{b}\lambda-2\alpha\lambda-\alpha}{2\alpha \lambda+\alpha-2\eta_{b}\lambda}\sigma^{2},$ $$

as wanted.
∎

Now the proof of Theorem 3.1 follows immediately from Theorem D.1 setting $\lambda=\frac{\beta}{1-\beta}$ (from Proposition 2.1) and $\gamma_{b}=\eta_{b}$ (from Corollary 2.2). The bound on $\beta=\frac{\lambda}{1+\lambda}$ follows from the bound on $\lambda$ and from the fact that the function $f(x)=\frac{x}{1+x}$ is strictly increasing for $x\geq 0$.

###### Theorem D.1 .

###### Proof.

#### D.2 Proof of Theorem 3.5

Here we state and prove the corresponding theorem for DecSPS in the IMA setting.

###### Theorem D.2 .

Assume that each $f_{i}$ is convex and $L_{i}$-smooth. Then, the iterates of IMA with

$$ $\displaystyle\eta_{t}=\frac{1}{c_{t}}\min\left\{\frac{f_{S_{t}}(x^{t})-\ell_{S _{t}}^{*}}{\|\nabla f_{S_{t}}(x^{t})\|^{2}},c_{t-1}\eta_{t-1}\right\}\text{ and }\lambda_{t}\geq 0,$ $$

where $(c_{t})$ is an increasing sequences with $c_{t}\geq 1$, $c_{-1}=c_{0}$, $\eta_{-1}=\eta_{b}$ and $(\lambda_{t})$ is a decreasing sequence, converge as

$$ $\displaystyle\operatorname{\mathbb{E}}[f(\overline{x}^{T})-f(x^{*})]\leq\frac{ 2\lambda_{1}[f(x^{0})-f(x^{*})]}{T}+\frac{c_{T-1}D^{2}_{z}}{T\alpha}+\frac{ \sigma^{2}}{T}\sum_{t=0}^{T-1}\frac{1}{c_{t}},$ $$

where $\overline{x}^{T}=\frac{1}{T}\sum_{t=0}^{T-1}x^{t}$, $\alpha=\min\left\{\frac{1}{2L},\eta_{b}c_{0}\right\}$ and $D^{2}_{z}=\max_{t\in[T]}\|z^{t}-x^{*}\|^{2}$.

###### Proof.

We have

$$ $\displaystyle\|z^{t+1}-x^{*}\|^{2}\overset{(\ref{eq:ima-1})}{=}\|z^{t}-x^{*}\| ^{2}-2\eta_{t}\langle\nabla f_{S_{t}}(x^{t}),z^{t}-x^{*}\rangle+\eta_{t}^{2}\| \nabla f_{S_{t}}(x^{t})\|^{2}.$ $$

Rearranging and dividing by $\eta_{t}>0$ we get

$$ $\displaystyle 2\langle\nabla f_{S_{t}}(x^{t}),z^{t}-x^{*}\rangle\leq\frac{\|z^ {t}-x^{*}\|^{2}}{\eta_{t}}-\frac{\|z^{t+1}-x^{*}\|^{2}}{\eta_{t}}+\eta_{t}\| \nabla f_{S_{t}}(x^{t})\|^{2}.$ $$

Summing the above for $t=0,\dots,T-1$ and telescoping we have

$$ $\displaystyle 2\sum_{t=0}^{T-1}\langle\nabla f_{S_{t}}(x^{t}),z^{t}-x^{*}\rangle$ $\displaystyle\leq\sum_{t=0}^{T-1}\left[\frac{\|z^{t}-x^{*}\|^{2}}{\eta_{t}}- \frac{\|z^{t+1}-x^{*}\|^{2}}{\eta_{t}}+\eta_{t}\|\nabla f_{S_{t}}(x^{t})\|^{2}\right]$ $\displaystyle=\sum_{t=0}^{T-1}\frac{\|z^{t}-x^{*}\|^{2}}{\eta_{t}}-\sum_{t=0}^ {T-1}\frac{\|z^{t+1}-x^{*}\|^{2}}{\eta_{t}}+\sum_{t=0}^{T-1}\eta_{t}\|\nabla f _{S_{t}}(x^{t})\|^{2}$ $\displaystyle=\frac{\|z^{0}-x^{*}\|^{2}}{\eta_{0}}+\sum_{t=1}^{T-1}\frac{\|z^{ t}-x^{*}\|^{2}}{\eta_{t}}-\sum_{t=0}^{T-1}\frac{\|z^{t+1}-x^{*}\|^{2}}{\eta_{t }}+\sum_{t=0}^{T-1}\eta_{t}\|\nabla f_{S_{t}}(x^{t})\|^{2}$ $\displaystyle=\frac{\|z^{0}-x^{*}\|^{2}}{\eta_{0}}+\sum_{t=0}^{T-2}\frac{\|z^{ t+1}-x^{*}\|^{2}}{\eta_{t+1}}-\sum_{t=0}^{T-1}\frac{\|z^{t+1}-x^{*}\|^{2}}{ \eta_{t}}+\sum_{t=0}^{T-1}\eta_{t}\|\nabla f_{S_{t}}(x^{t})\|^{2}$ $\displaystyle=\frac{\|z^{0}-x^{*}\|^{2}}{\eta_{0}}+\sum_{t=0}^{T-2}\frac{\|z^{ t+1}-x^{*}\|^{2}}{\eta_{t+1}}-\sum_{t=0}^{T-2}\frac{\|z^{t+1}-x^{*}\|^{2}}{ \eta_{t}}-\frac{\|z^{T}-x^{*}\|^{2}}{\eta_{T-1}}+\sum_{t=0}^{T-1}\eta_{t}\| \nabla f_{S_{t}}(x^{t})\|^{2}$ $\displaystyle=\frac{\|z^{0}-x^{*}\|^{2}}{\eta_{0}}+\sum_{t=0}^{T-2}\left(\frac {1}{\eta_{t+1}}-\frac{1}{\eta_{t}}\right)\|z^{t+1}-x^{*}\|^{2}-\frac{\|z^{T}-x ^{*}\|^{2}}{\eta_{T-1}}+\sum_{t=0}^{T-1}\eta_{t}\|\nabla f_{S_{t}}(x^{t})\|^{2}$ $\displaystyle\leq\left(\frac{1}{\eta_{0}}+\sum_{t=0}^{T-2}\left[\frac{1}{\eta_ {t+1}}-\frac{1}{\eta_{t}}\right]\right)D^{2}_{z}+\sum_{t=0}^{T-1}\eta_{t}\| \nabla f_{S_{t}}(x^{t})\|^{2}$ $\displaystyle=\frac{D^{2}_{z}}{\eta_{T-1}}+\sum_{t=0}^{T-1}\eta_{t}\|\nabla f_ {S_{t}}(x^{t})\|^{2}.$ (18) $$

Taking expectation in eq. 18 we have

$$ $\displaystyle 2\sum_{t=0}^{T-1}\operatorname{\mathbb{E}}\left[\langle\nabla f( x^{t}),z^{t}-x^{*}\rangle\right]\leq\operatorname{\mathbb{E}}\left[\frac{D^{2} _{z}}{\eta_{T-1}}+\sum_{t=0}^{T-1}\eta_{t}\|\nabla f_{S_{t}}(x^{t})\|^{2} \right].$ (19) $$

Now the LHS of eq. 19 can be bounded as follows

$$ $\displaystyle 2\sum_{t=0}^{T-1}\operatorname{\mathbb{E}}\left[\langle\nabla f( x^{t}),z^{t}-x^{*}\rangle\right]$ $\displaystyle\overset{(z^{0}=x^{0})}{=}2\operatorname{\mathbb{E}}\left[\langle \nabla f(x^{0}),x^{0}-x^{*}\rangle\right]+2\sum_{t=1}^{T-1}\operatorname{ \mathbb{E}}\left[\langle\nabla f(x^{t}),z^{t}-x^{*}\rangle\right]$ $\displaystyle\overset{(\ref{eq:z-def})}{=}2\operatorname{\mathbb{E}}\left[ \langle\nabla f(x^{0}),x^{0}-x^{*}\rangle\right]+\sum_{t=1}^{T-1}\operatorname {\mathbb{E}}\left[2\langle\nabla f(x^{t}),x^{t}-x^{*}\rangle+2\lambda_{t} \langle\nabla f(x^{t}),x^{t}-x^{t-1}\rangle\right]$ $\displaystyle\overset{convexity}{\geq}2\operatorname{\mathbb{E}}\left[f(x^{0}) -f(x^{*})\right]+\sum_{t=1}^{T-1}\operatorname{\mathbb{E}}\left[2(f(x^{t})-f(x ^{*}))+2\lambda_{t}(f(x^{t})-f(x^{t-1}))\right]$ $\displaystyle=2\operatorname{\mathbb{E}}\left[f(x^{0})-f(x^{*})\right]+\sum_{t =1}^{T-1}\left[2(1+\lambda_{t})\operatorname{\mathbb{E}}[f(x^{t})-f(x^{*})]-2 \lambda_{t}\operatorname{\mathbb{E}}[f(x^{t-1})-f(x^{*})]\right]$ $\displaystyle=2\operatorname{\mathbb{E}}\left[f(x^{0})-f(x^{*})\right]+\sum_{t =1}^{T-1}2(1+\lambda_{t})\operatorname{\mathbb{E}}[f(x^{t})-f(x^{*})]-\sum_{t= 1}^{T-1}2\lambda_{t}\operatorname{\mathbb{E}}[f(x^{t-1})-f(x^{*})]$ $\displaystyle=2\operatorname{\mathbb{E}}\left[f(x^{0})-f(x^{*})\right]+\sum_{t =1}^{T-1}2(1+\lambda_{t})\operatorname{\mathbb{E}}[f(x^{t})-f(x^{*})]-\sum_{t= 0}^{T-2}2\lambda_{t+1}\operatorname{\mathbb{E}}[f(x^{t})-f(x^{*})]$ $\displaystyle=2\operatorname{\mathbb{E}}\left[f(x^{0})-f(x^{*})\right]+2\sum_{ t=1}^{T-1}\operatorname{\mathbb{E}}[f(x^{t})-f(x^{*})]+\sum_{t=1}^{T-1}2 \lambda_{t}\operatorname{\mathbb{E}}[f(x^{t})-f(x^{*})]$ $\displaystyle\quad-\sum_{t=1}^{T-2}2\lambda_{t+1}\operatorname{\mathbb{E}}[f(x ^{t})-f(x^{*})]$ $\displaystyle=2\operatorname{\mathbb{E}}\left[f(x^{0})-f(x^{*})\right]+2\sum_{ t=1}^{T-1}\operatorname{\mathbb{E}}[f(x^{t})-f(x^{*})]+2\lambda_{T-1} \operatorname{\mathbb{E}}[f(x^{T-1})-f(x^{*})]$ $\displaystyle\quad+\sum_{t=1}^{T-2}2\lambda_{t}\operatorname{\mathbb{E}}[f(x^{ t})-f(x^{*})]-\sum_{t=1}^{T-2}2\lambda_{t+1}\operatorname{\mathbb{E}}[f(x^{t}) -f(x^{*})]$ $\displaystyle=2\sum_{t=0}^{T-1}\operatorname{\mathbb{E}}[f(x^{t})-f(x^{*})]+2 \lambda_{T-1}\operatorname{\mathbb{E}}[f(x^{T-1})-f(x^{*})]$ $\displaystyle\quad+2\sum_{t=1}^{T-2}(\lambda_{t}-\lambda_{t+1})\operatorname{ \mathbb{E}}[f(x^{t})-f(x^{*})]-2\lambda_{1}\operatorname{\mathbb{E}}[f(x^{0})- f(x^{*})]$ $\displaystyle\geq 2\sum_{t=0}^{T-1}\operatorname{\mathbb{E}}[f(x^{t})-f(x^{*}) ]-2\lambda_{1}[f(x^{0})-f(x^{*})],$ (20) $$

using the fact that $(\lambda_{t})$ is decreasing and the fact that $\operatorname{\mathbb{E}}[f(x^{t})-f(x^{*})]\geq 0$. Combining eq. 19 and eq. 20 we have

$$ $\displaystyle 2\sum_{t=0}^{T-1}\operatorname{\mathbb{E}}[f(x^{t})-f(x^{*})]-2 \lambda_{1}[f(x^{0})-f(x^{*})]\leq\operatorname{\mathbb{E}}\left[\frac{D^{2}_{ z}}{\eta_{T-1}}+\sum_{t=0}^{T-1}\operatorname{\mathbb{E}}\eta_{t}\|\nabla f_{S _{t}}(x^{t})\|^{2}\right].$ (21) $$

Now for the RHS of eq. 21 we have

$$ $\displaystyle\frac{D^{2}_{z}}{\eta_{T-1}}+\sum_{t=0}^{T-1}\eta_{t}\|\nabla f_{ S_{t}}(x^{t})\|^{2}$ $\displaystyle\overset{(\ref{eq:lem:decsps-bound})}{\leq}\frac{D^{2}_{z}}{\eta_ {T-1}}+\sum_{t=0}^{T-1}\frac{1}{c_{t}}[f_{S_{t}}(x^{t})-\ell_{S_{t}}^{*}]$ $\displaystyle=\frac{D^{2}_{z}}{\eta_{T-1}}+\sum_{t=0}^{T-1}\frac{1}{c_{t}}[f_{ S_{t}}(x^{t})-f_{S_{t}}(x^{*})]+\sum_{t=0}^{T-1}\frac{1}{c_{t}}[f_{S_{t}}(x^{* })-\ell_{S_{t}}^{*}]$ $\displaystyle\leq\frac{c_{T-1}D^{2}_{z}}{\alpha}+\sum_{t=0}^{T-1}\frac{1}{c_{t }}[f_{S_{t}}(x^{t})-f_{S_{t}}(x^{*})]+\sum_{t=0}^{T-1}\frac{1}{c_{t}}[f_{S_{t} }(x^{*})-\ell_{S_{t}}^{*}],$ (22) $$

where the last inequality follows from the fact that $\eta_{t}$ is decreasing and $\eta_{t}\geq\frac{\alpha}{c_{t}}$ with $\alpha=\min\{\frac{1}{2L_{\max}},\eta_{b}c_{0}\}$, from Lemma C.4. Thus, combining eq. 21 and eq. 22 we have

$$ $\displaystyle 2\sum_{t=0}^{T-1}\operatorname{\mathbb{E}}[f(x^{t})-f(x^{*})]-2 \lambda_{1}[f(x^{0})-f(x^{*})]\leq\frac{c_{T-1}D^{2}_{z}}{\alpha}+\sum_{t=0}^{ T-1}\frac{1}{c_{t}}\operatorname{\mathbb{E}}[f(x^{t})-f(x^{*})]+\sum_{t=0}^{T- 1}\frac{\sigma^{2}}{c_{t}}.$ (23) $$

Rearranging eq. 23 we get

$$ $\displaystyle\sum_{t=0}^{T-1}\left(2-\frac{1}{c_{t}}\right)\operatorname{ \mathbb{E}}[f(x^{t})-f(x^{*})]\leq 2\lambda_{1}[f(x^{0})-f(x^{*})]+\frac{c_{T- 1}D^{2}_{z}}{\alpha}+\sigma^{2}\sum_{t=0}^{T-1}\frac{1}{c_{t}}$ (24) $$

and using the fact that $2-\frac{1}{c_{t}}\geq 1>0$ eq. 24 reduces to

$$ $\displaystyle\sum_{t=0}^{T-1}\operatorname{\mathbb{E}}[f(x^{t})-f(x^{*})]\leq 2 \lambda_{1}[f(x^{0})-f(x^{*})]+\frac{c_{T-1}D^{2}_{z}}{\alpha}+\sigma^{2}\sum_ {t=0}^{T-1}\frac{1}{c_{t}}$ (25) $$

Dividing eq. 25 by $T$ and using Jensen’s inequality we get

$$ $\displaystyle\operatorname{\mathbb{E}}[f(\overline{x}^{T})-f(x^{*})]$ $\displaystyle\leq\frac{1}{T}\sum_{t=0}^{T-1}\operatorname{\mathbb{E}}[f(x^{t}) -f(x^{*})]$ $\displaystyle\leq\frac{2\lambda_{1}[f(x^{0})-f(x^{*})]}{T}+\frac{c_{T-1}D^{2}_ {z}}{T\alpha}+\frac{\sigma^{2}}{T}\sum_{t=0}^{T-1}\frac{1}{c_{t}},$ $$

as wanted.
∎

If $\sigma^{2}>0$ like in the standard SGD analysis under decreasing step-sizes, the choice $c_{t}=\sqrt{t+1}$ leads to the asymptotic rate $O(1/\sqrt{t})$.

###### Corollary D.3 .

Under the setting of Theorem D.2, for $c_{t}=\sqrt{t+1}$ ($c_{-1}=c_{0}$) we have

$$ $\displaystyle\operatorname{\mathbb{E}}[f(\overline{x}^{T})-f(x^{*})]$ $\displaystyle\leq\frac{2\lambda_{1}[f(x^{0})-f(x^{*})]}{T}+\frac{D^{2}_{z}}{ \alpha\sqrt{T}}+\frac{2\sigma^{2}}{\sqrt{T}}.$ $$

###### Proof.

It follows from the well known inequality

$$ $\displaystyle\sum_{t=0}^{T-1}\frac{1}{\sqrt{t+1}}\leq 2\sqrt{T}.$ $$

∎

Now the proof of Theorem 3.5 follows immediately from Corollary D.3 setting $\lambda_{t}=\lambda=\frac{\beta}{1-\beta}$ (from Proposition 2.1). Moreover we have

$$ $\displaystyle\|z^{t}-x^{*}\|^{2}$ $\displaystyle\overset{(\ref{eq:z-def})}{=\joinrel=}\|(1+\lambda)x^{t}-\lambda x ^{t-1}-x^{*}\|^{2}$ $\displaystyle\leq(1+\lambda)\|x^{t}-x^{*}\|^{2}+\lambda\|x^{t-1}-x^{*}\|^{2}$ $\displaystyle\leq(1+2\lambda)\max_{t}\|x^{t}-x^{*}\|^{2},$ (26) $$

thus $D^{2}_{z}\leq(1+2\lambda)\max_{t}\|x^{t}-x^{*}\|^{2}=\frac{1+\beta}{1-\beta}D^
{2}$, where $D^{2}=\max_{t\in[T]}\|x^{t}-x^{*}\|$.

###### Theorem D.2 .

###### Proof.

###### Corollary D.3 .

###### Proof.

#### D.3 Proof of Theorem 3.6

Here we state and prove the corresponding theorem for AdaSPS in the IMA setting.

###### Theorem D.4 .

Assume that each $f_{i}$ is convex and $L_{i}$-smooth. Then, the iterates of IMA with

$$ $\displaystyle\eta_{t}=\min\left\{\frac{f_{S_{t}}(x^{t})-\ell_{S_{t}}^{*}}{c\| \nabla f_{S_{t}}(x^{t})\|^{2}}\frac{1}{\sqrt{\sum_{s=0}^{t}f_{S_{s}}(x^{s})- \ell_{S_{s}}^{*}}},\eta_{t-1}\right\}\text{ and }\lambda_{t}\geq 0,$ $$

where $\eta_{-1}=+\infty$, $c>0$ and $(\lambda_{t})$ a decreasing sequence, converge as

$$ $\displaystyle\operatorname{\mathbb{E}}[f(\overline{x}^{T})-f(x^{*})]\leq\frac{ \tau^{2}}{T}+\frac{\tau\sigma}{\sqrt{T}},$ $$

where $\overline{x}^{T}=\frac{1}{T}\sum_{t=0}^{T-1}x^{t}$, $\tau=\left(\lambda_{1}\sqrt{f(x^{0})-f(x^{*})}+\frac{cLD^{2}_{z}}{2}+\frac{1}{
2c}\right)$ and $D^{2}_{z}=\max_{t}\|z^{t}-x^{*}\|$.

###### Proof.

Using Equation 21 from the proof of Theorem 3.5, we have

$$ $\displaystyle 2\sum_{t=0}^{T-1}\operatorname{\mathbb{E}}[f(x^{t})-f(x^{*})]-2 \lambda_{1}[f(x^{0})-f(x^{*})]\leq\operatorname{\mathbb{E}}\left[\frac{D^{2}_{ z}}{\eta_{T-1}}+\sum_{t=0}^{T-1}\operatorname{\mathbb{E}}\eta_{t}\|\nabla f_{S _{t}}(x^{t})\|^{2}\right].$ (27) $$

By Lemma C.7 we have

$$ $\displaystyle\frac{D^{2}_{z}}{\eta_{T-1}}\leq cLD^{2}_{z}\sqrt{\sum_{s=0}^{T-1 }f_{S_{s}}(x^{s})-\ell_{S_{s}}^{*}},$ (28) $$

and by Lemma C.5 we have

$$ $\displaystyle\sum_{t=0}^{T-1}\eta_{t}\|\nabla f_{S_{t}}(x^{t})\|^{2}\leq\sum_{ t=0}^{T-1}\frac{f_{S_{t}}(x^{t})-\ell_{S_{t}}^{*}}{2c\sqrt{\sum_{s=0}^{t}f_{S_ {s}}(x^{s})-\ell_{S_{s}}^{*}}}\leq\frac{1}{c}\sqrt{\sum_{s=0}^{T-1}f_{S_{s}}(x ^{s})-\ell_{S_{s}}^{*}}.$ (29) $$

Now combining Equations 27, 28 and 29 and using Jensen’s inequality we get

$$ $\displaystyle 2\sum_{t=0}^{T-1}\operatorname{\mathbb{E}}[f(x^{t})-f(x^{*})]-2 \lambda_{1}[f(x^{0})-f(x^{*})]$ $\displaystyle\leq\operatorname{\mathbb{E}}\left[cLD^{2}_{z}\sqrt{\sum_{s=0}^{T -1}f_{S_{s}}(x^{s})-\ell_{S_{s}}^{*}}+\frac{1}{c}\sqrt{\sum_{s=0}^{T-1}f_{S_{s }}(x^{s})-\ell_{S_{s}}^{*}}\right]$ $\displaystyle=\left(cLD^{2}_{z}+\frac{1}{c}\right)\operatorname{\mathbb{E}} \left[\sqrt{\sum_{s=0}^{T-1}f_{S_{s}}(x^{s})-\ell_{S_{s}}^{*}}\right]$ $\displaystyle=\left(cLD^{2}_{z}+\frac{1}{c}\right)\operatorname{\mathbb{E}} \left[\sqrt{\sum_{s=0}^{T-1}(f_{S_{s}}(x^{s})-f_{S_{s}}(x^{*})+f_{S_{s}}(x^{*} )-\ell_{S_{s}}^{*})}\right]$ $\displaystyle\leq\left(cLD^{2}_{z}+\frac{1}{c}\right)\sqrt{\sum_{s=0}^{T-1} \operatorname{\mathbb{E}}[f(x^{s})-f(x^{*})]+\sigma^{2}}.$ (30) $$

Rearranging eq. 30 we have

$$ $\displaystyle 2\sum_{t=0}^{T-1}\operatorname{\mathbb{E}}[f(x^{t})-f(x^{*})] \leq 2\lambda_{1}[f(x^{0})-f(x^{*})]+\left(cLD^{2}_{z}+\frac{1}{c}\right)\sqrt {\sum_{s=0}^{T-1}\operatorname{\mathbb{E}}[f(x^{s})-f(x^{*})]+\sigma^{2}}.$ (31) $$

Now let us choose $c_{q}=\sqrt{f(x^{0})-f(x^{*})}$. Then

$$ $\displaystyle 2\lambda_{1}[f(x^{0})-f(x^{*})]$ $\displaystyle=2\lambda_{1}c_{q}\sqrt{f(x^{0})-f(x^{*})}$ $\displaystyle\leq 2\lambda_{1}c_{q}\sqrt{\sum_{s=0}^{T-1}\operatorname{\mathbb {E}}[f(x^{s})-f(x^{*})]+\sigma^{2}},$ (32) $$

Combining eq. 31 and eq. 32 we get

$$ $\displaystyle\sum_{t=0}^{T-1}\operatorname{\mathbb{E}}[f(x^{t})-f(x^{*})]\leq \left(\lambda_{1}c_{q}+\frac{cLD^{2}_{z}}{2}+\frac{1}{2c}\right)\sqrt{\sum_{s= 0}^{T-1}\operatorname{\mathbb{E}}[f(x^{s})-f(x^{*})]+\sigma^{2}}.$ (33) $$

Squaring both sides of eq. 33 we have

$$ $\displaystyle\left(\sum_{t=0}^{T-1}\operatorname{\mathbb{E}}[f(x^{t})-f(x^{*}) ]\right)^{2}\leq\tau^{2}\left(\sum_{t=0}^{T-1}\operatorname{\mathbb{E}}[f(x^{t })-f(x^{*})]+T\sigma^{2}\right),$ (34) $$

where $\tau=\left(\lambda_{1}c_{q}+\frac{cLD^{2}_{z}}{2}+\frac{1}{2c}\right)$. Now we use Lemma C.6 in eq. 33 to get

$$ $\displaystyle\sum_{t=0}^{T-1}\operatorname{\mathbb{E}}[f(x^{t})-f(x^{*})]\leq \tau^{2}+\tau\sigma\sqrt{T}.$ (35) $$

Finally by Jensen’s inequality and eq. 35 we have

$$ $\displaystyle\operatorname{\mathbb{E}}[f(\overline{x}^{T})-f(x^{*})]\leq\frac{ \tau^{2}}{T}+\frac{\tau\sigma}{\sqrt{T}},$ $$

as wanted.
∎

Now the proof of Theorem 3.6 follows immediately from Theorem D.4 setting $\lambda=\frac{\beta}{1-\beta}$ (from Proposition 2.1). Moreover, using eq. 26 we have $D^{2}_{z}\leq\frac{1+\beta}{1-\beta}D^{2}$.

###### Theorem D.4 .

###### Proof.

### Appendix E Additional Experiments

#### E.1 Extra Deterministic Experiment

Here we have included an extra experiment in the deterministic setting, this one for logistic regression. We have performed a grid search to find the best $\beta$ for MomPSmax since no optimal choices for $\beta$ are known for the general convex case. We use the same $\beta$ for HB and ALR-MAG for direct comparison. For the step-size of HB we used $\gamma=2(1-\beta)/L$ as recommended in Ghadimi et al. []. The results are presented in Figure 7. In this problem we observe that our step-size (MomPSmax) is the fastest, having similar performace with ALR-MAG and GD Polyak. For the deterministic setting we have used a 12 core AMD Ryzen 5 5600H CPU to run the experiments.

Figure: (a) Relative Error Plot
Refer to caption: x10.png

#### E.2 Stochastic Heavy Ball with Constant Step-size

As we saw in Corollary 3.4, when $\gamma_{b}\leq\frac{1}{2L_{\max}}$ then we have new theoretical guarantees for SHB with the constant step-size $\gamma=\frac{1-\beta}{2L_{\max}}$. The most recent analysis of constant SHB is from Liu et al. [] with step-size $\gamma=\frac{(1-\beta)^{2}}{L}\min\left\{\frac{1}{4-\beta+\beta^{2}},\frac{1}{
2\sqrt{2\beta+2\beta^{2}}}\right\}$. As mentioned in the main text our step-size is larger when $\beta\geq\sqrt{5}-2\approx 0.236$. In Figure 8 we provide a numerical comparison of these results for logistic regression with synthetic data. We observe that SHB with the Liu et al. [] step-size has faster convergence for the first iterations but it reaches a plateau much earlier. Moreover, note that as $\beta\to 1$ both step-sizes have similar performance.

Figure: (a) $\beta=0.1$
Refer to caption: x12.png

#### E.3 Other Choices

In this paper, we have provided convergence guarantees for SHB update rule $x^{t+1}=x^{t}-\gamma_{t}\nabla f_{S_{t}}(x^{t})+\beta(x^{t}-x^{t-1})$ with the MomSPSmax step-size given by $\gamma_{t}=(1-\beta)\min\left\{\frac{f_{S_{t}}(x^{t})-\ell_{S_{t}}^{*}}{c\|
\nabla f_{S_{t}}(x^{t})\|^{2}},\gamma_{b}\right\}$. However, as discussed in the main paper a more natural step-size would be to choose the well known SPSmax on the SHB rule. We call this new update rule SPSmax with naive momentum. Here we numerically compare these two updates. As we see in Figure 9, when $\beta$ is “small” and close to 0 0 then both SPSmax with naive momentum and MomSPSmax are close in performance. Of course, when $\beta=0$ they are both equal to standard SPSmax. However, as $\beta$ gets larger the performance of SPSmax with naive momentum gets worse and less stable and at some point it diverges. For $\beta\geq 0.7$ there is no SPSmax with naive momentum, since it cannot be numerically computed.

Figure: (a) $\beta=0.0$
Refer to caption: x21.png

Furthermore, recall that our decreasing variants MomDecSPS and MomAdaSPS are respectively given by

$$ $\displaystyle\gamma_{t}=\min\left\{\frac{(1-\beta)[f_{S_{t}}(x^{t})-\ell_{S_{t }}^{*}]}{c\sqrt{t+1}\|\nabla f_{S_{t}}(x^{t})\|^{2}},\frac{\gamma_{t-1}\sqrt{t }}{\sqrt{t+1}}\right\},$ (MomDecSPS) $\displaystyle\gamma_{t}=\min\left\{\frac{(1-\beta)[f_{S_{t}}(x^{t})-\ell_{S_{t }}^{*}]}{c\|\nabla f_{S_{t}}(x^{t})\|^{2}\sqrt{\sum_{s=0}^{t}f_{S_{s}}(x^{s})- \ell_{S_{s}}^{*}}},\gamma_{t-1}\right\}.$ (MomAdaSPS) $$

The above step-sizes are theoretically inspired by IMA, as explained in Section C.3, but again a natural question is what happens if we take the “correcting factor” $1-\beta$ outside the minimum. We call these step-sizes Alternative MomDecSPS and Alternative MomAdaSPS and are given by

$$ $\displaystyle\gamma_{t}=(1-\beta)\min\left\{\frac{f_{S_{t}}(x^{t})-\ell_{S_{t} }^{*}}{c\sqrt{t+1}\|\nabla f_{S_{t}}(x^{t})\|^{2}},\frac{\gamma_{t-1}\sqrt{t}} {\sqrt{t+1}}\right\},$ (Alt MomDecSPS) $\displaystyle\gamma_{t}=(1-\beta)\min\left\{\frac{f_{S_{t}}(x^{t})-\ell_{S_{t} }^{*}}{c\|\nabla f_{S_{t}}(x^{t})\|^{2}\sqrt{\sum_{s=0}^{t}f_{S_{s}}(x^{s})- \ell_{S_{s}}^{*}}},\gamma_{t-1}\right\}.$ (Alt MomAdaSPS) $$

In Figure 10 and Figure 11 we see that when $\beta=0$ we have similiar performances, however when $\beta>0$ then our proposed step-sizes are significantly better. A possible reason is that for both Alt MomDecSPS and Alt MomadaSPS it holds $\gamma_{t}\leq(1-\beta)\gamma_{t-1}$ for all $t$. Thus inductively we get $\gamma_{t}\leq(1-\beta)^{t}\gamma_{0}$. So when $\beta>0$, $(1-\beta)^{t}\to 0$ as $t\to\infty$ which means that the step-size is too small and it barely updates $x^{t}$.

Figure: (a) $\beta=0.0$
Refer to caption: x29.png

Figure: (a) $\beta=0.0$
Refer to caption: x33.png

#### E.4 Is γ b subscript 𝛾 𝑏 \gamma_{b} italic_γ start_POSTSUBSCRIPT italic_b end_POSTSUBSCRIPT needed?

Recall the definition of MomPSmax, $\gamma_{t}=(1-\beta)\min\left\{\frac{f(x^{t})-f(x^{*})}{\|\nabla f(x^{t})\|^{2
}},\gamma_{b}\right\}$. The more natural choice would be without the upper bound $\gamma_{b}$. Moreover, as in Section E.3, we can have the naive version of both of the above step-sizes, i.e. with no correcting factor $1-\beta$. 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 on synthetic data. The method is SHB with $\beta=0.97$ for the least squares and $\beta=0.3$ for the logistic regression.

$$ $\displaystyle\gamma_{t}=(1-\beta)\frac{f(x^{t})-f(x^{*})}{\|\nabla f(x^{t})\|^ {2}}$ (MomPS) $\displaystyle\gamma_{t}=\frac{f(x^{t})-f(x^{*})}{\|\nabla f(x^{t})\|^{2}}$ (Alt MomPS) $\displaystyle\gamma_{t}=(1-\beta)\min\left\{\frac{f(x^{t})-f(x^{*})}{\|\nabla f (x^{t})\|^{2}},\gamma_{b}\right\}$ (MomPS max ) $\displaystyle\gamma_{t}=\min\left\{\frac{f(x^{t})-f(x^{*})}{\|\nabla f(x^{t}) \|^{2}},\gamma_{b}\right\}.$ (Alt MomPS max ) $$

We see that if $\gamma_{b}$ is chosen large enough, all these step-sizes have similar performance. However, we highlight the fact that convergence guarantees are known only for MomPSmax.

Figure: (a) $\gamma_{b}=0.1$
Refer to caption: x37.png

Figure: (a) $\gamma_{b}=1$
Refer to caption: x41.png

#### E.5 Plots for C 1 subscript 𝐶 1 C_{1} italic_C start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT and C 2 subscript 𝐶 2 C_{2} italic_C start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT

Below it a plot of the constants $C_{1}$ and $C_{2}$ as functions of $\beta$ with $\gamma_{b}=2$ and $\alpha=1$ (recall $\gamma_{b}\geq\alpha$) from Theorem 3.1. Both functions are hyperbolas and the vertical line is the constant $\beta=\frac{\alpha}{2\gamma_{b}-\alpha}$. We see that in the interval $\beta\in\left[0,\frac{\alpha}{2\gamma_{b}-\alpha}\right)$ both functions are increasing.

Figure: Figure 14: Plots for the constants $C_{1}$ and $C_{2}$ as functions of $\beta$. Here we have chosen $\gamma_{b}=2$ and $\alpha=1$.
Refer to caption: extracted/5646932/figures/SHB+SPS/consts.png

#### E.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 2. We also include some extra experiments in Figures 15, 16, 17 and 18.

**Table 2: Logistic regression experiment**
| \CT@drsc@<br>Hyper-parameter | Value |
| --- | --- |
| Architecture | Logistic Regression |
| GPUs | 1x NVIDIA GeForce RTX 3050 |
| Batch-size | See caption of each plot |
| Epochs | 100 |
| Trials | 5 |
| Weight Decay | 0.0 |
| \CT@drsc@ |  |

Figure: Figure 15: LIBSVM dataset: vehicle, Batch size: 16
Refer to caption: x45.png

Figure: Figure 16: LIBSVM dataset: letter, Batch size: 256
Refer to caption: x46.png

Figure: Figure 17: LIBSVM dataset: vehicle, Batch size: 85
Refer to caption: x47.png

Figure: Figure 18: LIBSVM dataset: glass, Batch size: 32
Refer to caption: x48.png

#### E.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.

**Table 3: CIFAR-10 experiment**
| \CT@drsc@<br>Hyper-parameter | Value |
| --- | --- |
| Architecture | ResNet 10/34 [He et al., ], Densenet 121 [Huang et al., ] |
| GPUs | 1x Nvidia RTX 6000 Ada Generation |
| Batch-size | 256 |
| Epochs | 100 for Densenet 121, 200 for ResNet 10, ResNet 34 |
| Trials | 5 |
| Weight Decay | 0.0 |
| \CT@drsc@ |  |

Figure: Figure 19: Resnet 18 on CIFAR-10
Refer to caption: x49.png

Figure: Figure 20: DenseNet 121 on CIFAR-10
Refer to caption: x50.png

## Peer Reviews

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

## Suggestions
1. Please, add explicit assumptions in the text about interpolation and smoothness etc., that you use in the theorems in some section like "Preliminaries". Sometimes it is hard to follow, whether you mean interpolation or bounded variance and it is hard to find them in the text.
1. Please, increase font size in the figures.
soundness: 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.
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.
weaknesses: The Section 3 is written in rather overwhelming way. After each theorem goes a paragraph, mentioning some special cases and how the provided result matches the results from other existing works, which is rather hard to comprehend for a new reader. You could change these paragraphs to small tables with small supporting captions. This should take about the same size as the paragraph in text, although it would be much more evident. Additionally, after every theorem follows a lot of corollaries, that again consider some special cases. It is easy to get lost in so many special cases. I think, some of these corollaries can also be moved to appendix, leaving some mention about them in the main part. Since the main result is introduction of Polyak step size and adaptive Polyak step size to SHB, it should be the central point of main part of the paper. 
Overall, despite the seemingly good results, considering everything above and some suggestions in Question part, I think the main text needs polishing and not ready yet for publication.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: **Questions**

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

**Minor issues**

1. Line 55

   has efficiently analyzed => has been efficiently analyzed.

2. Line 400, 410

   Please be consistent with step size and step-size.

3. Line 1098

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

**References**

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

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

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

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

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

[6] Liu, Yanli, Yuan Gao, and Wotao Yin. An improved analysis of stochastic gradient descent with momentum. *Advances in Neural Information Processing Systems* 33, 2020.
soundness: 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.
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.
weaknesses: Although the paper is a solid work, my foremost concern is its limited technical novelty. 

1. Limited technical contributions

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

2. Strong assumptions

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

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

---

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

The assumption of bounded iterates in Section 3.2 constitutes a strong condition, which may not hold even for quadratics without additional constraints. In the absence of this assumption, is it possible to derive meaningful theoretical results for the MomDecSPS and MomAdaSPS methods?
soundness: 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.
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.
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.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: None.
soundness: 3
strengths: 1.The writing and exposition are clear.

2.The proposed methods are practical to implement and have theoretical improvements compared with existing works.
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.
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)

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

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

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

Decision: Accept (Poster)