An Improved Analysis of Stochastic Gradient Descent with Momentum

Introduction

Stochastic gradient methods have been a widespread practice in machine learning. They aim to minimize the following empirical risk:

In general, stochastic gradient methods can be written as

The idea behind SGDM originates from Polyak’s heavy-ball method for deterministic optimization. For strongly convex and smooth objectives, heavy-ball method enjoys an accelerated linear convergence rate over gradient descent . However, the theoretical understanding of its stochastic counterpart is far from being complete.

In the case of fixed stepsize and momentum weight, most of the current results only apply to restrictive settings. In and , the behavior of SGDM on least square regression is analyzed and linear convergence is established. analyzes the local convergence rate of SGDM for strongly convex and smooth functions, where the initial point $x^{0}$ is assumed to be close enough to the minimizer $x^{*}$. provides global convergence of SGDM, but only for objectives with uniformly bounded gradients, thus excluding many machine learning models such as Ridge regression. Very recently, presents a convergence bound of $\mathcal{O}(\frac{1}{k\alpha}+\frac{\alpha}{1-\beta})$ for general smooth nonconvex objectivesHere $k$ is the number of iterations. Note that in , a different but equivalent formulation of SGDM is analyzed; their stepsize $\gamma$ is effectively $\frac{\alpha}{1-\beta}$ in our setting.. When $\beta=0$, this recovers the classical convergence bound of $\mathcal{O}(\frac{1}{k\alpha}+\alpha)$ of SGD . However, the size of stationary distribution $\mathcal{O}(\frac{\alpha}{1-\beta})$ is $\frac{1}{1-\beta}$ times larger than that of SGD. This factor is not negligible, especially when large $\beta$ values such as $0.99$ and $0.995$ is applied . Therefore, their result does not explain the competitiveness of SGDM compared to SGD. Concurrent to this work, shows that SGDM converges as fast as SGD under convexity and strong convexity, and that it is asymptotically faster than SGD for overparameterized models. Remarkably, their analysis considers a different stepsize and momentum weight schedule from this work, and applies to arbitrary sampling without assuming the bounded variance of the gradient noise.

In deep learning, SGDM is often applied with various parameter tuning rules to achieve efficient training. One of the most widely adopted rules is called “constant and drop", where a constant stepsize is applied for a long period and is dropped by some constant factor to allow for refined training, while the momentum weight is either kept unchanged (usually $0.9$) or gradually increasing. We call this strategy Multistage SGDM and summarize it in Algorithm 1. Practically, (multistage) SGDM was successfully applied to training large-scale neural networks , and it was found that appropriate parameter tuning leads to superior performance . Since then, (multistage) SGDM has become increasingly popular .

At each stage, Multistage SGDM (Algorithm 1) requires three parameters: stepsize, momentum weight, and stage length. In and , doubling argument based rules are analyzed for SGD on strongly convex objectives, where the stage length is doubled whenever the stepsize is halved. Recently, certain stepsize schedules are shown to yield faster convergence for SGD on nonconvex objectives satisfying growth conditions , and a nearly optimal stepsize schedule is provided for SGD on least square regression . These results consider only the momentum-free case. Another recent work focuses on the asymptotic convergence of SGDM (i.e., without convergence rate) , which requires the momentum weights to approach either or $1$, and therefore contradicts the common practice in neural network training. In summary, the convergence rate of Multistage SGDM (Algorithm 1) has not been established except for the momentum-free case, and the role of parameters in different stages is unclear.

In this work, we provide new convergence analysis for SGDM and Multistage SGDM that resolve the aforementioned issues. A comparison of our results with prior work can be found in Table 1.

We show that for both strongly convex and nonconvex objectives, SGDM (2) enjoys the same convergence bound as SGD. This helps explain the empirical observations that SGDM is at least as fast as SGD . Our analysis relies on a new observation that, the update direction $m^{k}$ of SGDM (2) has a controllable deviation from the current full gradient $\nabla f(x^{k})$, and enjoys a smaller variance. Inspired by this, we construct a new Lyapunov function that properly handles this deviation and exploits an auxiliary sequence to take advantage of the reduced variance.

Compared to aforementioned previous work, our analysis applies to not only least squares, does not assume uniformly bounded gradient, and improves the convergence bound.

For the more popular SGDM in the multistage setting (Algorithm 1), we establish its convergence and demonstrate that the multistage strategy are faster at initial stages. Specifically, we allow larger stepsizes in the first few stages to boost initial performance, and smaller stepsizes in the final stages decrease the size of stationary distribution. Theoretically, we properly redefine the aforementioned auxiliary sequence and Lyapunov function to incorporate the stagewise parameters.

To the best of our knowledge, this is the first convergence guarantee for SGDM in the multistage setting.

2 Other related work

Nesterov’s momentum achieves optimal convergence rate in deterministic optimization , and has also been combined with SGD for neural network training . Recently, its multistage version has been analyzed for convex or strongly convex objectives . Other forms of momentum for stochastic optimization include PID Control-based methods , Accelerated SGD , and Quasi-Hyperbolic Momentum . In this work, we restrict ourselves to heavy-ball momentum, which is arguably the most popular form of momentum in current deep learning practice.

Notation and Preliminaries

The following assumption is effective throughout, which is standard in stochastic optimization.

Smoothness: The objective $f(x)$ in (1) is $L-$smooth.

Independent samples: the random samples $\{\zeta_{k}\}_{k=1}^{\infty}$ are independent.

Unless otherwise noted, all the proof in the paper are deferred to the appendix.

Key Ingredients of Convergence Theory

In this section, we present some key insights for the analysis of stochastic momentum methods. For simplicity, we first focus on the case of fixed stepsize and momentum weight, and make proper generalizations for Multistage SGDM in App. C.

In this section, we make the following observation on the role of momentum:

With a momentum weight $\beta\in[0,1)$, the update vector $m^{k}$ enjoys a reduced “variance" of $(1-\beta)\sigma^{2}$, while having a controllable deviation from the full gradient $g^{k}$ in expectation.

First, without loss of generality, we can take $m^{0}=0$, and express $m^{k}$ as

$m^{k}$ is a moving average of the past stochastic gradients, with smaller weights for older onesNote the sum of weights $(1-\beta)\sum_{i=1}^{k}\beta^{k-i}=1-\beta^{k}\rightarrow 1$ as $k\rightarrow\infty$..

we have the following result regarding the “variance" of $m^{k}$, which is measured between $m^{k}$ and its deterministic version $(1-\beta)\sum_{i=1}^{k}\beta^{k-i}g^{i}$.

Under Assumption 1, the update vector $m^{k}$ in SGDM (2) satisfies

Lemma 1 follows directly from the property of the moving average.

On the other hand, $(1-\beta)\sum_{i=1}^{k}\beta^{k-i}g^{i}$ is a moving average of all past gradients, which is in contrast to SGD. It seems unclear how far is $(1-\beta)\sum_{i=1}^{k}\beta^{k-i}g^{i}$ from the ideal descent direction $g^{k}$, which could be unbounded unless stronger assumptions are imposed. Previous analysis such as and make the blanket assumption of bounded $\nabla f$ to circumvent this difficulty.

In this work, we provide a different perspective to resolve this issue.

From Lemma 2, we know the deviation of $\frac{1}{1-\beta^{k}}(1-\beta)\sum_{i=1}^{k}\beta^{k-i}g^{i}$ from $g^{k}$ is controllable sum of past successive iterate differences, in the sense that the coefficients $a_{k,i}$ decays linearly for older ones. This inspires the construction of a new Lyapunov function to handle the deviation brought by the momentum, as we shall see next.

2 A new Lyapunov function

Let us construct the following Lyapunov function for SGDM:

In the Lyapunov function (5), $\{c_{i}\}_{i=1}^{\infty}$ are positive constants to be specified later corresponding to the deviation described in Lemma 2. Since the coefficients in (4) converges linearly to as $k\rightarrow\infty$, we can choose $\{c_{i}\}_{i=1}^{\infty}$ in a diminishing fashion, such that this deviation can be controlled, and $L^{k}$ defined in (5) is indeed a Lyapunov function under strongly convex and nonconvex settings (see Propositions 1 and 2).

In (5), $z^{k}$ is an auxiliary sequence defined as

This auxiliary sequence first appeared in the analysis of deterministic heavy ball methods in , and later applied in the analysis of SGDM . It enjoys the following property.

Since the coefficients of the deviation in Lemma 2 converges linearly to as $k\rightarrow\infty$, we can choose $\{c_{i}\}_{i=1}^{\infty}$ in a diminishing fashion, such that this deviation can be controlled. Remarkably, we shall see in Sec. 4 that with $c_{1}=\mathcal{O}\left(\frac{L}{1-\beta}\right)$, $L^{k}$ defined in (5) is indeed a Lyapunov function under strongly convex and nonconvex settings, and that SGDM converges as fast as SGD.

Now, let us turn to the Multistage SGDM (Algorithm 1), which has been very successful in neural network training. However, its convergence still remains unclear except for the momentum-free case. To establish convergence, we require the parameters of Multistage SGDM to satisfy

where $\alpha_{i},\beta_{i},$ and $T_{i}$ are the stepsize, momentum weight, and stage length of $i$th stage, respectively, and $A_{1},A_{2}$ are properly chosen constants. In principle, one applies larger stepsizes $\alpha_{i}$ at the initial stages, which will accelerate initial convergence, and smaller stepsizes for the final stages, which will shrink the size of final stationary distribution. As a result, (7) stipulates that less iterations are required for stages with large stepsizes and more iterations for stages with small stepsizes. Finally, (7) requires the momentum weights to be monotonically increasing, which is consistent with what’s done in practice . often, using constant momentum weight also works.

Under the parameter choices in (7), let us define the auxiliary sequence $z^{k}$ by

This $\{z^{k}\}_{k=1}^{\infty}$ sequence reduces to (6) when a constant stepsize and momentum weight are applied. Furthermore, the observations made in Lemmas 1, 2, and 3 can also be generalized (see Lemmas 4, 5, 6, and 7 in App. C). In Sec. 5. we shall see that with (7) and appropriately chosen $\{c_{i}\}_{i=1}^{\infty}$, $L^{k}$ in (5) also defines a Lyapunov function in the multistage setting, which in turn leads to the convergence of Multistage SGDM.

Convergence of SGDM

In this section, we proceed to establish the convergence of SGDM described in (2). First, by following the idea presented in Sec. 3, we can show that $L^{k}$ defined in (5) is a Lyapunov function.

Let Assumption 1 hold. In (2), let $\alpha\leq\frac{1-\beta}{2\sqrt{2}L\sqrt{\beta+\beta^{2}}}$. Let $\{c_{i}\}_{i=1}^{\infty}$ in (5) be defined by

By telescoping (9), we obtain the stationary convergence of SGDM under nonconvex settings.

Let Assumption 1 hold. In (2), let $\alpha\leq\alpha\leq\min\{\frac{1-\beta}{L(4-\beta+\beta^{2})},\frac{1-\beta}{2\sqrt{2}L\sqrt{\beta+\beta^{2}}}\}$. Then,

Now let us turn to the strongly convex setting, for which we have

Let Assumption 1 hold. Assume in addition that $f$ is $\mu-$strongly convex. In (2), let $\alpha\leq\min\{\frac{1-\beta}{5L},\frac{1-\beta}{L\left(3-\beta+2\beta^{2}+\frac{48\sqrt{\beta}}{25}\frac{2L+18\mu}{L}\right)}\}$. Then, there exists positive constants $c_{i}$ for (5) such that for all $k\geq k_{0}\coloneqq\lfloor\frac{\log 0.5}{\log\beta}\rfloor$, we have

The choices of $\{c_{i}\}_{i=1}^{\infty}$ is similar to those of Proposition 1 and can be found in App. B.4. With Proposition 2, we immediately have

Let Assumption 1 hold and assume in addition that $f$ is $\mu-$strongly convex. Under the same settings as in Proposition 2, for all $k\geq k_{0}=\lfloor\frac{\log 0.5}{\log\beta}\rfloor$ we have

Let Assumption 1 hold and assume in addition that $f$ is $\mu-$strongly convex. Under the same settings as in Proposition 2, for all $k\geq k_{0}=\lfloor\frac{\log 0.5}{\log\beta}\rfloor$ we have

Under nonconvex settings, the classical convergence bound of SGD is $\mathcal{O}\left(\frac{f(x^{1})-f^{*}}{k\alpha}+L\alpha\sigma^{2}\right)$ with $\alpha=\mathcal{O}(\frac{1}{L})$ (see, e.g., Theorem 4.8 of ). Therefore, Theorem 1 tells us that with $\alpha=\mathcal{O}(\frac{1-\beta}{L})$, SGDM achieves the same convergence bound as SGD.

In contrast, the radius of the stationary distribution for SGDM in and is $\mathcal{O}(\frac{\alpha\sigma^{2}}{1-\beta})$, and the latter one also assumes that $\nabla f$ is uniformly bounded.

In Theorem 2 and Corollary 1, the convergence bounds hold for $k\geq k_{0}=\lfloor\frac{\log 0.5}{\log\beta}\rfloor$, where $k_{0}$ is a mild constantFor example, we have $k_{0}=6$ for the popular choice $\beta=0.9$..

The convergence bound of SGD under strong convexity is $\mathcal{O}\left((1-\alpha\mu)^{k}+\frac{L}{\mu}\alpha\sigma^{2}\right)$ (see, e.g, Theorem 4.6 of ), our result for SGDM in Corollary 1 recovers this when $\beta=0$.

Convergence of Multistage SGDM

In this section, we switch to the Multistage SGDM (Algorithm 1).

Let us first show that when the (7) is applied, we can define the constants $c_{i}$ properly so that (5) still produces a Lyapunov function.

Let Assumption 1 hold. In Algorithm 1, let the parameters satisfy (7) with $A_{1}=\frac{1}{24\sqrt{2}L}$. In addition, let

Then, we have $c_{i}>0$ for any $i\geq 1$. Furthermore, with $z^{k}$ defined in (8), for any $k\geq 1$, we have

where $\alpha(k),\beta(k)$ are the stepsize and momentum weight applied at $k$th iteration, respectively.

Let Assumption 1 hold. Under the same settings as in Proposition 3, let $\beta_{1}\geq\frac{1}{2}$ and let $A_{2}$ be large enough such that

On the left hand side of (11), we have the average of the averaged squared gradient norm of $n$ stages.

On the right hand side of (11), the first term dominates at initial stages, we can apply large $\alpha_{i}$ for these stages to accelerate initial convergence, and use smaller $\alpha_{i}$ for later stages so that the size of stationary distribution is small. In contrast, (static) SGDM need to use a small stepsize $\alpha$ to make the size of stationary distribution with small.

It is unclear whether the overall iteration complexity of Multistage SGDM is better than SGDM or not. Numerically, we do observe that Multistage SGDM is faster. We leave the possible improved analysis of Multistage SGDM to future work.

Experiments

In this section, we verify our theoretical claims by numerical experiments. For each combination of algorithm and training task, training is performed with $3$ random seeds $1,2,3$. Unless otherwise stated, we report the average of losses of the past $m$ batches, where $m$ is the number of batches for the whole dataset. Additional implementation details can be found in App. E.

Setup. The MNIST dataset consists of $n=60000$ labeled examples of $28\times 28$ gray-scale images of handwritten digits in $K=10$ classes $0,1,\dots,9$. For all algorithms, we use batch size $s=64$ (and hence number batches per epoch is $m=1874$), number of epochs $T=50$. The regularization parameter is $\lambda=5\times 10^{-4}$.

The effect of $\alpha$ in (static) SGDM. By Theorem 2 we know that, with a fixed $\beta$, a larger $\alpha$ leads to faster loss decrease to the stationary distribution. However, the size of the stationary distribution is also larger. This is well illustrated in Figure 1. For example, $\alpha=1.0$ and $\alpha=0.5$ make losses decrease more rapidly than $\alpha=0.1$. During later iterations, $\alpha=0.1$ leads to a lower final loss.

Multistage SGDM. We take $3$ stages for Multistage SGDM. The parameters are chosen according to (7): $T_{1}=3,T_{2}=6,T_{3}=21$, $\alpha_{i}=A_{2}/T_{i}$, $\beta_{i}=A_{1}/(c_{2}+\alpha_{i})$, where $A_{2}=2.0$ and $A_{1}=1.0$.Here, $A_{1}$ is not set by its theoretical value $\frac{1}{24L}$, since the dataset is very large and the gradient Lipschitz constant $L$ cannot be computed easily. We compare Multistage SGDM with SGDM with $(\alpha,\beta)=(0.66,0.9)$ and $(\alpha,\beta)=(0.095,0.9)$, where $0.66$, $0.095$ are the stepsizes of the first and last stage of Multistage SGDM, respectively. The training losses of initial and later iterations are shown in Figure 2.

We can see that SGDM with $(\alpha,\beta)=(0.66,0.9)$ converges faster initially, but has a higher final loss; while SGDM with $(\alpha,\beta)=(0.095,0.9)$ behaves the other way. Multistage SGDM takes the advantage of both, as predicted by Theorem 3. The performances of SGDM and Vanilla SGD with the same stepsize are similar.

2 Image classification

For the task of training ResNet18 on the CIFAR-10 dataset, we compare Multistage SGDM, a baseline SGDM, and YellowFin , an automatic momentum tuner based on heuristics from optimizing strongly convex quadratics. The initial learning rate of YellowFin is set to $0.1$,We have experimented with initial learning rates $0.001$ (default), $0.01$, $0.1$ and $0.5$, each repeated $3$ times; we found that the choice $0.1$ is the best in terms of the final training loss. and other parameters are set as their default values. All algorithms are run for $T=50$ epochs and the batch size is fixed as $s=128$.

For Multistage SGDM, the parameters choices are governed by (7): the stage lengths are $T_{1}=5$, $T_{2}=10$, and $T_{3}=35$. Take $A_{1}=1.0$, $A_{2}=2.0$, set the per-stage stepsizes and momentum weights as $\alpha_{i}=A_{2}/T_{i}$ and $\beta_{i}=A_{1}/(A_{1}+\alpha_{i})$, for stages $i=1,2,3$. For the baseline SGDM, the stepsize schedule of Multistage SGDM is applied, but with a fixed momentum $\beta=0.9$.

In Figure 3, we present training losses and end-of-epoch validation accuracy of the tested algorithms. We can see that Multistage SGDM performs the best. Baseline SGDM is slightly worse, possibly because of its fixed momentum weight.

Summary and Future Directions

In this work, we provide new theoretical insights into the convergence behavior of SGDM and Multistage SGDM. For SGDM, we show that it is as fast as plain SGD in both nonconvex and strongly convex settings. For the widely adopted multistage SGDM, we establish its convergence and show the advantage of stagewise training.

There are still open problems to be addressed. For example, (a) Is it possible to show that SGDM converges faster than SGD for special objectives such as quadratic ones? (b) Are there more efficient parameter choices than (7) that guarantee even faster convergence?

References

Appendix A Proof of Preliminary Lemmas

It is straightforward to see that the same conclusion holds for $i<j$.

Finally, we know from the item 4 in Assumption 1 that

A.2 Proof of Lemma 2

where we have applied Cauchy-Schwarz in the first inequality.

By applying triangle inequality and the smoothness of $f$ (item 1 in Assumption 1), we further have

Therefore, by defining $a^{\prime}_{k,j}=\frac{1-\beta}{1-\beta^{k}}L^{2}\sum_{i=1}^{j}\beta^{k-i}(k-i)$, we get

Furthermore, $a^{\prime}_{k,j}$ can be calculated as

Combining this with (12), we finally arrive at

A.3 Proof of Lemma 3

Let us consider the cases of $k=1$ and $k\geq 2$ separately.

Appendix B Main Theory for SGDM

In order to prove Proposition 1, let us first show an auxiliary result.

Take Assumption 1. Then, for $z^{k}$ defined in (6), we have

where we have applied Lemma 3 in the second step.

where $\rho_{0}>0$ can be any positive constant (to be determined later).

Combining (16) and the last inequality, we arrive at

Since $z^{k}=x^{k}$ when $k=1$ and $z^{k}=\frac{1}{1-\beta}x^{k}-\frac{\beta}{1-\beta}x^{k-1}$ when $k\geq 2$, it can be verified that $z^{k}-x^{k}=-\frac{\beta}{1-\beta}\alpha m^{k-1}$. Consequently,

On the other hand, from Lemma 1 we know that

Finally, using $1-\beta^{k-1}<1$ and $\rho_{0}=\frac{1-\beta}{2L\alpha}$ gives

B.2 Proof of Proposition 1

where $\rho_{0}=\frac{1-\beta}{2L\alpha}$.

In the rest of the proof, let us show that the sum of the last three terms in (22) is non-positive.

Therefore, in order to make the sum of the last three terms of (22) to be non-positive, we need to have

Since $1-\beta^{k}<1$, it suffices to enforce the following for all $i\geq 1$:

And in order for $c_{i}>0$ for all $i\geq 1$, we can determine $c_{1}$ by

we have $\sum_{i=1}^{\infty}i\beta^{i}=\frac{\beta}{(1-\beta)^{2}}$ and

Notice that $\alpha\leq\frac{1-\beta}{4L\sqrt{\beta+\beta^{2}}}$ ensures $c_{1}>0$.

With the choices of $c_{i}$ in (23) and (24), the sum of the last three terms of (22) is non-positive. Therefore,

B.3 Proof of Theorem 1

In the rest the proof, we will bound $R_{1}$ and $R_{2}$ appropriately.

First, let us show that $R_{1}\geq\frac{\alpha}{2}$ when $\rho_{0}=\frac{1-\beta}{2L\alpha}$ as in (26) and $\alpha\leq\min\{\frac{1-\beta}{L(4-\beta+\beta^{2})},\frac{1-\beta}{2L\sqrt{\beta+\beta^{2}}}\}$.

Since $\alpha\leq\frac{1-\beta}{2\sqrt{2}L\sqrt{\beta+\beta^{2}}}$, we have

Therefore, in order to ensure $R_{1}\geq\frac{\alpha}{2}$ where $R_{1}$ is defined in (28), it suffices to have

Applying $\rho_{0}=\frac{1-\beta}{2L\alpha}$ yields

where we have applied $\alpha\leq\frac{1-\beta}{L(4-\beta+2\beta^{2})}$ in the last step.

Now let us turn to $R_{2}$. By (32) we know that

Since $\rho_{0}=\frac{1-\beta}{2L\alpha}$, we have

By applying $\alpha\leq\min\{\frac{1-\beta}{L(4-\beta+\beta^{2})},\frac{1-\beta}{2\sqrt{2}L\sqrt{\beta+\beta^{2}}}\}\leq\frac{1-\beta}{3.75L}<\frac{1-\beta}{4L}$, we further have

By putting (34) and (35) into (30), we finally obtain

B.4 Proof of Proposition 2

In order to prove Proposition 2, we will set

Take $\rho_{0}=\frac{1-\beta}{2L\alpha}$ in (22), we have

Let us first derive a lower bound of the first term on the right hand side of (36).

where $\rho$, $\rho_{1}>0$ are to be determined later.

On the other hand, we have from (18) that

Putting these two inequalities into (38) and rearranging gives

Taking $\rho=\frac{1}{1-\beta}$ and $\rho_{1}=\frac{1}{\beta}$ gives

Since $\alpha\leq\frac{1-\beta}{5L}$, we have that

Therefore, by $\alpha\leq\frac{1-\beta}{L\left(3-\beta+2\beta^{2}+\frac{48\sqrt{\beta}}{25}\frac{2L+18\mu}{L}\right)}$ we have

By combining (44) with (41), we further obtain

In the rest of the proof, we will show that if the constants $c_{i}$ are chosen such that

Then, we have $c_{i}>0$ for all $i\geq 1$ and

And therefore, we will have the desired result:

First of all. by $k\geq\frac{\log 0.5}{\log\beta}$, we know that $\beta^{k}\leq\frac{1}{2}$, and (47) gives

Therefore, in order for (51) to hold, it suffices to set

On the other hand, (50) is also equivalent to

Therefore, in order to have $c_{i}>0$ for all $i\geq 1$, we can set

Since $\beta\leq\sqrt{\beta}\leq 1+B_{1}=1-\alpha\mu\frac{1}{1+\frac{8\mu}{L}}$ and

for any $q\in(0,1)$, (52) is equivalent to

Since $\alpha\leq\frac{1-\beta}{L}$, we further have

Since $B_{1}=-\frac{\alpha\mu}{1+\frac{8\mu}{L}}$ and $\alpha\leq\frac{1-\beta}{5L}$, it can be verified that $\frac{\beta}{1+B_{1}}\leq\sqrt{\beta}$ for all $\beta\in[0,1)$ and $\mu\leq L$. Therefore,

As a result, in order to have (53), it suffices to set

Since $\alpha\leq\frac{1-\beta}{5L}$, we have

which is exactly our choice of $c_{1}$ as in (49).

B.5 Proof of Theorem 2

From Proposition 2 we know that for all $k\geq k_{0}=\lfloor\frac{\log 0.5}{\log\beta}\rfloor$,

where we have applied $\alpha\leq\frac{1-\beta}{5L}$ in the last step. Therefore,

By $c_{i}\geq 0$ for all $i\geq 1$ and (42), we conclude that

B.6 Proof of Corollary 1

In fact, by (6) we can express $x^{k}$ as a convex combination of $\{z^{i}\}_{i=1}^{k}$:

The desired result follows directly from the convexity of $f$ and Theorem 2.

Appendix C Generalizations of Lemmas 1, 2, and 3 for Multistage SGDM

In order to establish the convergence of Multistage SGDM(Algorithm 1), we need to generalize the Lemmas 1 and 2 , which play a key role in the convergence of SGDM in (2).

Under the assumptions of Theorem 3, the variance of update vector $m^{k}$ in Algorithm 1 satisfies

where b_{k,i}=\big{(}1-\beta(i)\big{)}\prod_{j=i+1}^{k}\beta(j).

To begin with, let us express $m^{k}$ by the past stochastic gradients:

where we have applied $m^{0}=0$ and defined

It can be verified that the sum of weights is

As a result, by applying Assumption 1 we have

Note that by setting $k=T_{1}+\dots+T_{n_{k}}+r_{k}$, we have

Under the assumptions of Theorem 3, the update vector $m^{k-1}$ in Algorithm 1 satisfies

By setting $k-1=T_{1}+\dots+T_{n_{k-1}}+r_{k-1}$, we have

C.2 Generalization of Lemma 3 for Multistage SGDM

where $\alpha(k)$ is the stepsize applied at the $k$th step.

Recall that the auxiliary sequence $z^{k}$ is defined by

where $A_{1}\equiv\frac{\alpha_{i}\beta_{i}}{1-\beta_{i}}$ and $\alpha_{i},\beta_{i}$ are the stepsize and momentum weight at the $i$th stage, respectively. Therefore, we also have

where $\alpha(k),\beta(k)$ are the stepsize and momentum weight applied at the $k$th step. Using this, we obtain

C.3 Generalization of Lemma 2 for Multistage SGDM

In Multistage SGDM(Algorithm 1), assume that the momentum weights at $n$ stages satisfy $\beta_{1}\leq\beta_{2}\leq...\leq\beta_{n}$. Then, we have

where b_{k,i}=\big{(}1-\beta(i)\big{)}\prod_{j=i+1}^{k}\beta(j) and $\beta(i)$ is the momentum weight applied at the $i$th iteration, and

By By (55), (56) and (57), we can compute that

where we have used (57) in the first and third equality, and Cauchy-Schwarz in the first inequality. In the last inequality, we have applied the triangle inequality and the $L-$smoothness of $f$.

In the Proposition 5 below, we shall see that $d_{k,i}\leq a_{k,i}$ for all $i\leq k-1$, where $a_{k,i}$ is defined in (58). Therefore,

$d_{k,i}$ defined in (59) and $a_{k,i}$ defined in (58) satisfy

We aim to show that $d_{k,i}\leq a_{k,i}$ for all $i\leq k-1$. Or equivalently, $d_{k,j}\leq a_{k,j}$ for all $j\leq k-1$.

In order to show $d_{k,j}\leq a_{k,j}$, we just need to show that

Let $k=T_{1}+T_{2}+\dots+T_{n_{k}}+r_{k}$, where $0\leq n_{k}\leq n-1$. If $n_{k}<n-1$, then $0\leq r_{k}\leq T_{n_{k}+1}-1$. If $n_{k}=n-1$, then $0\leq r_{k}\leq T_{n_{k}+1}=T_{n}$.

Since $j\leq k-1$, we have $j=T_{1}+\dots+T_{n_{j}}+r_{j}$, where $0\leq n_{j}\leq n_{k}$.

Now, let us compute the left hand side of (60) explicitly.

where we have applied $r_{j}\leq T_{n_{j}+1}$ if $n_{j}<n_{k}$ and $r_{j}\leq r_{k}$ if $n_{j}=n_{k}$ in the last term. Since

By applying these equalities on (61), we have

On the right hand side, the first $n_{j}$ terms are non-positive since $\beta_{1}\leq\beta_{2}\leq...\leq\beta_{n}$. Therefore,

By applying $\beta_{n_{j}+1}^{r_{j}}\geq\beta^{T_{n_{j}}+1}_{n_{j}+1}$ and $k-T_{1}-\dots-T_{n_{j}}-1=k-(j-r_{j})-1\geq 0$ (since $j\leq k-1$), we arrive at

Now let us consider two cases: $r_{k}>0$ and $r_{k}=0$.

In this case, we apply $\beta_{1}\leq...\leq\beta_{n}$ to get

Since $r_{k}>0$, iteration $k$ is at the $(n_{k}+1)-$th stage, we have $\beta(k)=\beta_{n_{k}+1}$, and the above inequality is exactly what we want to show in (60).

In this case, we apply $\beta_{1}\leq...\leq\beta_{n}$ to get

Since $r_{k}=0$, we have $\beta(k)=\beta_{n_{k}}$ and by $j\leq k-1$ we deduce that $n_{j}\leq n_{k}-1$ (Otherwise $j=T_{1}+\dots+T_{n_{j}}+r_{j}=T_{1}+\dots+T_{n_{k}}+r_{j}\geq T_{1}+\dots+T_{n_{k}}=k$). Therefore, we have

which is exactly what we want to show in (60).

Appendix D Main Theory for Multistage SGDM

In this section, we prove the main convergence theory of Multistage SGDM.

Proposition 3 is a generalization of Propositions 4 and 1 to the multistage case. Therefore, its proof is similar to those of Propositions 4 and 1.

First of all, by the smoothness of $f$ we have

where we have applied Lemma 6 in the second step. Note that $\alpha(k)$ is the stepsize applied at the $k-$th iteration.

where $\rho_{0,k}>0$ can be any positive constant.

By (8) we know that $z^{k}-x^{k}=-A_{1}m^{k-1}$, which leads to

Plugging (65) and (67) into (64) gives us

In the rest of the proof, we will show that the sum of the last three terms in (68) is non-positive.

Therefore, in order to make the sum of the last three terms of (68) to be non-positive, we need to enforce that

Since $1-\prod_{i=1}^{k}\beta(i)<1$, $\beta_{1}\leq\beta(k)\leq\beta_{n}$, and $\alpha_{1}\leq\alpha(k)\leq\alpha_{n}$, we need to enforce the following for all $i\geq 1$:

Recall that $\frac{\alpha_{i}\beta_{i}}{1-\beta_{i}}\equiv A_{1}$ for all $n$ stages $i=1,2,...,n$. This gives us

Since $\beta_{1}\leq\beta(k)$, it suffices to enforce that

Note that the equalities in (70) does not depend on $k$. In order for $c_{i}>0$ for all $i\geq 1$, we can determine $c_{1}$ by

we have $\sum_{i=1}^{\infty}i\beta_{n}^{i}=\frac{\beta_{n}}{(1-\beta_{n})^{2}}$ and

Notice that $A_{1}=\frac{1}{24\sqrt{2}L}$ and $\frac{1-\beta_{1}}{\beta_{1}}\leq 12\frac{1-\beta_{n}}{\sqrt{\beta_{n}+\beta^{2}_{n}}}$ ensures

With the choices of $c_{i}$ in (70) and (71), the sum of the last three terms of (68) is non-positive. Therefore,

Taking $\rho_{0,k}=\frac{1-\beta(k)}{2L\alpha(k)}$ in (73) gives

Finally, by applying Lemma 4 and Lemma 5, we arrive at

D.2 Proof of Theorem 3

In the rest the proof, we will bound $R_{1,i}$ and $R_{2,i}$ appropriately.

First, let us show that $R_{1,i}\geq\frac{\alpha(i)}{2}$ under $\rho_{0,i}=\frac{1-\beta(i)}{2L\alpha(i)}$ as in (69) and $\alpha(i)=\frac{A_{1}(1-\beta(i))}{\beta(i)}=\frac{1-\beta(i)}{24\sqrt{2}L\beta(i)}$.

Therefore, in order for $R_{1,i}\geq\frac{\alpha(i)}{2}$, it suffices to have

By $\beta(i)\geq\beta_{1}\geq\frac{1}{2}$ we know that

Therefore, $\frac{L\alpha^{2}(i)}{2}\leq\frac{\alpha(i)}{4}$. Furthermore, $\rho_{0,i}=\frac{1-\beta(i)}{2L\alpha(i)}$ yields

where in the inequality above, we have applied

Now let us turn to $R_{2,i}$. By (76) and (72) we know that

Since $\rho_{0,i}=\frac{1-\beta(i)}{2L\alpha(i)}$ and $\frac{\alpha(i)\beta(i)}{1-\beta(i)}\equiv A_{1}$, we have

By applying Lemmas 4 and 5, we further have

By putting (79) and (80) into (77) with $k=T_{1}+T_{2}+\dots+T_{n}$, we obtain

Dividing both sides by $\frac{1}{2}nA_{2}\equiv\frac{1}{2}n\alpha_{l}T_{l}$ gives

Appendix E Details of computational infrastructure

All experiments were performed on a computing server with Intel(R) Core(TM) i9-9940X CPU @ 3.30GHz and NVidia GeForce RTX 2080 P8. The weights of the neural networks are initialized by the default, random initialization routines.