A High Probability Analysis of Adaptive SGD with Momentum

Introduction

Despite the incredible popularity of stochastic gradient methods in practical machine learning applications, our theoretical understanding of these methods is still not complete. In particular, adaptive learning rates methods like AdaGrad (Duchi et al., 2011) have been mainly studied in the convex domain, with few analyses in the non-convex domain (Li & Orabona, 2019; Ward et al., 2019). However, even in these latter analyses, the assumptions used are very strong and/or the results limited.

In particular, there are two main problems with the previous analyses of Stochastic Gradient Descent (SGD) and its variants in the nonconvex setting. First, the classic analysis of convergence for SGD in the nonconvex setting uses an analysis in expectation. However, expectation bounds do not rule out extremely bad outcomes. As pointed out by Harvey et al. (2019a), it is a misconception that for the algorithms who have expectation bounds it is enough to pick the best of several independent runs to have a high probability guarantee: It can actually be a computational inefficient procedure. Moreover, in practical applications like deep learning, it is often the case that only one run of the algorithm is used since that the training process may take long time. Hence, it is essential to get high probability bounds which guarantee the performance of the algorithm on single runs.

Another very common assumption used in most of the previous papers is the one of bounded stochastic gradients. This is a rather strong assumption and it is false even in the deterministic optimization of a convex quadratic function, e.g., $f(x)=x^{2}$.

In this work, we overcome both these problems. We prove high probability convergence rates only assuming that the noises on the gradients are well-behaved, i.e., subgaussian. In this way, we allow for unbounded gradients and unbounded noise. The weak assumptions, the nonconvex analysis, and the adaptive learning rates make our results particularly challenging to obtain. Indeed, high probability bounds for bounded stochastic gradients are almost trivial to obtain but of limited applicability. Overall, we believe this paper is the first one to prove such guarantees.

In this short paper, we present a high probability analysis of SGD with momentum and adaptive learning rates, with weak assumptions on function and stochastic gradients. So, first in Theorem 1 we prove high probability bounds for the gradients of classic momentum SGD step size $O(\frac{1}{\sqrt{t}})$ in the nonconvex setting. Then, in Theorem 2 we prove for the first time high probability convergence rates for the gradients of AdaGrad with momentum in the nonconvex setting. In particular, we also show that the high probability bounds are adaptive to the level of noise.

Related Work

Stochastic momentum methods. Sutskever et al. (2013) discussed the importance of classic momentum methods in deep learning, which is nowadays widely used in the training of neural networks. On the convergence of stochastic momentum methods, Yang et al. (2016) studied a unified momentum method and provided expectation bounds in the rate of $O(\frac{1}{\sqrt{T}})$ in both convex and nonconvex setting. However, the results hold only for Lipschitz functions. Gadat et al. (2018) provided an in-depth description of the stochastic heavy-ball method. Moreover, for the non-convex functions, they showed some almost sure convergence results. Loizou & Richtárik (2017) provided a general analysis for the momentum variants of several classes of stochastic optimization algorithms and proved the linear rate convergence for quadratic and smooth functions. To the best of our knowledge, there are no high probability bounds for nonconvex stochastic momentum methods without using strong assumptions.

Nonconvex convergence of adaptive methods. In recent years, a variety of adaptive SGD algorithms have been developed to automatically tune the step size by using the past stochastic gradients. The first adaptive algorithm was AdaGrad Duchi et al. (2011), designed to adapt to sparse gradients. Li & Orabona (2019) and Ward et al. (2019) showed the convergence of variants of AdaGrad with a rate of $O(\ln T/\sqrt{T})$ in the non-convex case. Moreover, Li & Orabona (2019) showed that AdaGrad with non-coordinate-wise learning rates is adaptive to the level of noise. Zou et al. (2019) studied AdaGrad with a unified momentum and Chen et al. (2019) considered a large family of Adam-like algorithms (Kingma & Ba, 2015) including AdaGrad with momentum. Yet, all of these works prove on bounds in expectation and most of them use the very strong assumption of bounded stochastic gradients.

High probability bounds. The results on high probability bounds are relatively rare compared to those in expectation, which are easier to obtain. Kakade & Tewari (2009) used Freeman’s inequality to prove high probability bounds for an algorithm solving the SVM objective function. For classic SGD, Harvey et al. (2019b) and Harvey et al. (2019a) used a generalized Freedman’s inequality to prove bounds in non-smooth and strongly convex case, while Jain et al. (2019) proved the optimal bound for the last iterate of SGD with high probability. As far as we know, there are currently no high probability bounds for adaptive methods in the nonconvex setting.

Problem Set-Up

Assumptions.

In this paper we focus on the optimization problem

where $f$ is bounded from below and we denote its infimum by $f^{\star}$. We do not assume the function to be convex. We consider stochastic optimization algorithms that have access to a noisy estimate of the gradient of $f$. This covers the ubiquitous SGD (Robbins & Monro, 1951), as well modern variants as AdaGrad. We are interested in studying the convergence of the gradients to zero, because without additional assumptions it is the only thing we can study in the nonconvex setting.

We make the following assumption on the objective function $f(\boldsymbol{x})$:

It is easy to see that this assumption is necessary to have the convergence of the gradients to zero. Indeed, without smoothness the norm of the gradients does not go to zero even in the convex case, e.g., consider the function $f(x)=|x-1|$.

We will also make the following assumption on the variance of the noise.

The condition (B2) has been used by Nemirovski et al. (2009) and Harvey et al. (2019a) to prove high probability convergence guarantees. Intuitively, it implies that the tails of the noise distribution are dominated by tails of a Gaussian distribution. Note that, by Jensen’s inequality, this condition implies a bounded variance on the stochastic gradients.

A General Analysis for Algorithms with Momentum

In this section, we will consider a generic stochastic optimization algorithm with Polyak’s momentum (Polyak, 1964; Qian, 1999; Sutskever et al., 2013), also known as the Heavy-ball algorithm or classic momentum, see Algorithm 1.

First, we want to point out that there two forms of Heavyball algorithms are possible. The first one is in Algorithm 1, while the second one is

This second is used in many practical implementation, see, for example, PyTorch (Paszke et al., 2019). It would seem that there is no reason to prefer one over the other. However, here we argue that the classic form of momentum is the right one if we want to use adaptive learning rates. To see why, let’s unroll the updates in both cases. Using the update in Algorithm 1, we have

In words, in the first case the update is composed by a sum of weighted gradients, each one multiplied by a learning rate we decided in the past. On the other hand, in the update (2) the update is composed by a sum of weighted gradients, each one multiplied by the current learning rate. From the analysis point of view, the second update destroys the independence between the past and the future, introducing a dependency that breaks our analysis, unless we introduce very strict conditions on the gradients. On the other hand, the update in Algorithm 1 allows us to carry out the analysis because each learning rate was chosen only with the knowledge of the past. Note that this is known problem in adaptive algorithms: the lack of independence between past and present is exactly the reason why Adam fails to converge on simple 1d convex problems, see for example the discussion in Savarese et al. (2019).

It is interesting to note that usually people argue that these two types updates for momentum are usually considered equivalent. This seems indeed true only if the learning rates are not adaptive.

Assumptions on learning rates.

$\boldsymbol{\eta}_{t}$ is non-increasing, i.e., $\boldsymbol{\eta}_{t+1}\leq\boldsymbol{\eta}_{t}$, $\forall t$.

$\boldsymbol{\eta}_{t}$ is independent with $\xi_{t}$.

The first assumption is very common (e.g., Duchi et al., 2011; Reddi et al., 2018; Li & Orabona, 2019; Chen et al., 2019; Zhou et al., 2018). Indeed, AdaGrad has the non-increasing step sizes by the definition. Also, Reddi et al. (2018) have claimed that the main issue of the divergences of Adam and RMSProp lies in the positive definiteness of $1/\boldsymbol{\eta}_{t}-1/\boldsymbol{\eta}_{t-1}$.

The need of the second assumption is technical and shared by similar analysis (Li & Orabona, 2019; Savarese et al., 2019). Indeed, Li & Orabona (2019) showed that delayed step sizes can avoid the possible deviation brought by the step sizes that include the current noise.

High probability guarantee.

Adaptive learning rates and in general learning rates that are decided using previous gradients become stochastic variables. This makes the high probability analysis more complex. Hence, we use a new concentration inequality for martingales in which the variance is treated as a random variable, rather than a deterministic quantity. We use this concentration in the proof of Lemma 2. Our proof, in the Appendix, merges ideas from the related results in Beygelzimer et al. (2011, Theorem 1) and Lan et al. (2012, Lemma 2). A similar result has also been shown by Jin et al. (2019, Lemma 6).

Main result.

We can now present our general lemma, that allows to analyze SGD with momentum with adaptive learning rates. We will then instantiate it for particular examples.

Assume (A, B1, B2, C1, C2). Then, for any $\delta\in(0,1)$, with probability at least $1-\delta$, the iterates of Algorithm 1 satisfy

Lemma 2 accomplishes the task of upper bounding the inner product $\sum_{t=1}^{T}\langle\boldsymbol{\eta}_{t}\nabla f(\boldsymbol{x}_{t}),\boldsymbol{m}_{t}\rangle$. Then, it is easy to lower bound the l.h.s by $\sum_{t=1}^{T}\langle\boldsymbol{\eta}_{T},\nabla f(\boldsymbol{x}_{t})^{2}\rangle$ using the assumption (C1), followed by the upper bound of $\sum_{t=1}^{T}\|\boldsymbol{\eta}_{t}\boldsymbol{g}_{t}\|^{2}$ based on the setting of $\boldsymbol{\eta}_{t}$.

1 SGD with Momentum with 1t1𝑡\frac{1}{\sqrt{t}} Learning Rates

As a warm-up, we now use Lemma 2 to prove a high probability convergence guarantee for the simple case of deterministic learning rates of $\eta_{t,i}=\frac{c}{\sqrt{t}}$.

Let $T$ the number of iterations of Algorithm 1. Assume (A, B1, B2). Set step size $\boldsymbol{\eta}_{t}$ as $\eta_{t,i}=\frac{c}{\sqrt{t}},i=1,\cdots,d$, where $c\leq\frac{1-\mu^{T}}{4M(3-2\mu)}$. Then, for any $\delta\in(0,1)$, with probability at least $1-\delta$, the iterates of Algorithm 1 satisfy

2 AdaGrad with Momentum

Now, we are going to prove the convergence rate of a variant AdaGrad in which we use momentum and learning rates that do not contain the current gradient. That is, the step sizes are defined as $\boldsymbol{\eta}_{t}=(\eta_{t,j})_{j=1,\dots,d}$

where $\alpha,\beta>0$. Removing the current gradient from the learning rate was proposed in Li & Orabona (2019); Savarese et al. (2019). Following the naming style in (Savarese et al., 2019), we denote this variant by Delayed AdaGrad.

Obviously, (3) satisfies (C1) and (C2). Hence, we are able to employ Lemma 2 to analyze this variant. Moreover, for Delayed AdaGrad, we upper bound $\sum_{t=1}^{T}\|\boldsymbol{\eta}_{t}\boldsymbol{g}_{t}\|^{2}$ with the following lemma, whose proof is in the Appendix.

Assume (A, B1, B2). Let $\boldsymbol{\eta}_{t}$ set as in (3), where $\alpha,\beta>0$. Then, for any $\delta\in(0,1)$, with probability at least $1-\delta$, we have

We now present the convergence guarantee for Delayed AdaGrad with momentum.

Assume (A, B1, B2). Let $\boldsymbol{\eta}_{t}$ set as in (3), where $\alpha,\beta>0$ and $4\alpha\leq\frac{\sqrt{\beta}(1-\mu)^{2}}{2M(1+\mu)}$. Then, for any $\delta\in(0,1)$, with probability at least $1-\delta$, the iterates of Algorithm 1 satisfy

where $C(T)=O\left(\frac{1}{\alpha}+\frac{d\left(\alpha+\sigma^{2}\left(\alpha\ln\frac{T}{\delta}+\frac{\ln\frac{1}{\delta}}{1-\mu}\right)\right)}{1-\mu}\right)$.

Observe that when $\sigma=0$, the convergence rate recovers the rate of Gradient Descent if $O(\frac{1}{T})$ with a constant learning rate. On the other hand, in the noisy case, it matches the rate of SGD $O(\frac{\sigma}{\sqrt{T}})$ with the optimal worst-case learning rate of $O(\frac{1}{\sigma\sqrt{t}})$. In other words, with a unique learning rate, we recover two different optimal convergence rates that requires two different learning rates and the knowledge of $\sigma$. This adaptivity of Delayed AdaGrad was already proved in Li & Orabona (2019), but only in expectation and without a momentum term.

Dependency on μ𝜇\mu.

Observe that the convergence upper bound increases over $\mu\in(0,1)$ and the optimal upper bound is achieved when taking the momentum parameter $\mu=0$. In words, the algorithms without momentums have the best theoretical results. This is a known caveat for this kind of analysis and a similar behavior w.r.t. $\mu$ is present, e.g., in Zou et al. (2018, Theorem 1) for algorithms with Polyak’s momentum.

Conclusion and Future Work.

In this work, we present a high probability analysis of adaptive SGD with Polyak’s momentum in the nonconvex setting. Without using the common assumption of bounded gradients nor bounded noise, we give the high probability bound for SGD with Polyak’s momentum with step size $O(\frac{1}{\sqrt{t}})$ and for Delayed AdaGrad with momentum. In particular, to the best of our knowledge, this is the first high probability convergence guarantee for adaptive methods.

In the future, we plan to extend our results to more adaptive methods, such as Adam and AMSGrad (Tieleman & Hinton, 2012). Moreover, we will explore other forms of momentum such as exponential moving average and Nesterov’s momentum (Nesterov, 1983).

Acknowledgements

This material is based upon work supported by the National Science Foundation under grant no. 1925930 “Collaborative Research: TRIPODS Institute for Optimization and Learning”.

References

Appendix A Appendix

By Jensen’s inequality, it follows that for any $c\in$,

Also it can be verified that $\exp(x)\leq x+\exp(9x^{2}/16)$ for all $x$, hence for $|\kappa|\in[0,4/3]$ we get

where in the second inequality, we used (4). Besides, $kx\leq 3k^{2}/8+2x^{2}/3$ holds for any $k$ and $x$. Hence for $|\kappa|\geq 4/3$, we get

where in the second inequality we used (4). Combining (5) and (6), we get $\forall\kappa$,

By Markov’s inequality, $P\left(Y_{T}\geq\frac{1}{\delta}\right)\leq\delta$, and $Y_{T}=\exp\left(\lambda\sum_{t=1}^{T}Z_{t}-\frac{3}{4}\lambda^{2}\sum_{t=1}^{T}\sigma_{t}^{2}\right),$ we have

To prove Lemma 2, we first need the following technical Lemma.

We prove these equalities by induction. When $T=1$, they obviously hold. Now, for $k<T$, assume that $\sum_{t=1}^{k}a_{t}\sum_{i=1}^{t}b_{i}=\sum_{t=1}^{k}b_{t}\sum_{i=t}^{k}a_{i}$. Then, we have

Hence, by induction, the equality is proved.

Similarly, for second equality assume that for $k<T$ we have $\sum_{t=1}^{k}a_{t}\sum_{i=0}^{t-1}b_{i}=\sum_{t=0}^{k-1}b_{t}\sum_{i=t+1}^{k}a_{i}$. Then, we have

By the smoothness of $f$ and the definition of $\boldsymbol{x}_{t+1}$, we have

We now upper bound $-\langle\nabla f(\boldsymbol{x}_{t}),\boldsymbol{m}_{t}\rangle$.

where the second inequality is due to the smoothness of $f$. Hence, iterating the inequality we have

Thus, denoting by $\boldsymbol{\epsilon}_{t}=\boldsymbol{g}_{t}-\nabla f(\boldsymbol{x}_{t})$ and summing (8) over $t$ from $1$ to $T$, we obtain

We then upper bound $S_{T}$. Denote by $L_{t}:=-\frac{1-\mu^{T-t+1}}{1-\mu}\langle\nabla f(\boldsymbol{x}_{t}),\boldsymbol{\eta}_{t}\boldsymbol{\epsilon}_{t}\rangle$, and $N_{t}:=\frac{(1-\mu^{T-t+1})^{2}}{(1-\mu)^{2}}\|\boldsymbol{\eta}_{t}\nabla f(\boldsymbol{x}_{t})\|^{2}\sigma^{2}$. Using the assumptions on the noise, for any $1\leq t\leq T$, we have

Finally, we upper bound $\sum_{t=1}^{T}\|\boldsymbol{m}_{t}\|^{2}$. From the convexity of $\|\cdot\|^{2}$, we have

Summing over $t$ from $1$ to $T$, we have

where in the first equality we used $\boldsymbol{m}_{0}=0$. Reordering the terms, we have that

Combining things together, and taking $\lambda=\frac{2(1-\mu)^{2}}{3\|\boldsymbol{\eta}_{1}\|(1-\mu^{T})^{2}\sigma^{2}}$, with probability at least $1-\delta$, we have

Rearranging the terms, we get the stated bound.

A.2 Details of Section 4.1

The proof of this Theorem 1 makes use of the following additional Lemma on the tail of sub-gaussian noise.

Assume B2, then for any $\delta\in(0,1)$, with probability at least $1-\delta$, we have

With the fact that $\|\boldsymbol{a}+\boldsymbol{b}\|^{2}\leq 2\|\boldsymbol{a}\|^{2}+2\|\boldsymbol{b}\|^{2}$, we have

By Lemma 5, Lemma 2 and the union bound, we have that with probability at least $1-\delta$,

Rearranging the terms and lower bounding $\sum_{t=1}^{T}\|\nabla f(\boldsymbol{x}_{t})\|^{2}$ by $T\cdot\min_{1\leq t\leq T}\|\nabla f(\boldsymbol{x}_{t})\|^{2}$, we have the stated bound. ∎

A.3 Details of Section 4.2

For the proof of Lemma 3, we first need the following technical Lemma.

Let $a_{i}\geq 0,\cdots,T$ and $f:[0,+\infty)\rightarrow[0,+\infty)$ non-increasing function. Then

Denote by $s_{t}=\sum_{i=0}^{t}a_{i}$. Then, we have

Summing over $i=1,\cdots,T$, we have the stated bound. ∎

First, we separate $\sum_{t=1}^{T}\|\boldsymbol{\eta}_{t}\boldsymbol{g}_{t}\|^{2}$ into two terms:

Using Lemma 5 on (9), for $\delta\in(0,1)$, with probability at least $1-\frac{\delta}{2}$, we have

We now upper bound $\sum_{t=1}^{T}\|\boldsymbol{\eta}_{t+1}\boldsymbol{g}_{t}\|^{2}$:

where in the first inequality we used Lemma 6 and in the second inequality we used Jensen’s inequality. Then using Lemma 5 on (11), with probability at least $1-\frac{\delta}{2}$, we have

Putting things together, we have the stated bound. ∎

Finally, to prove Theorem 2, we need the two following Lemmas.

Let $x\geq 0$, $A,C,D\geq 0$, $B>0$, and $x^{2}\leq(A+Bx)(C+D\ln(A+Bx))$. Then.

If $x\geq 0$, and $x\leq C\sqrt{A+Bx}$, then $x\leq\max\left(2BC^{2},C\sqrt{2A}\right)$.

By Lemma 2 and Lemma 3, for $\delta\in(0,1)$, with probability at least $1-\frac{2}{3}\delta$, we have

where $K$ denotes $2\alpha^{2}d\ln\left(\sqrt{\beta+\frac{2T\sigma^{2}}{d}\ln\frac{2Te}{\delta}}+\sqrt{\frac{2}{d}}\sqrt{\sum_{t=1}^{T}\|\nabla f(\boldsymbol{x}_{t})\|^{2}}\right)$ for conciseness.

where in the second inequality we used $4\alpha\leq\frac{\sqrt{\beta}(1-\mu)}{2M(3-\mu)}$. Also, we have

By Lemma 5, with probability at least $1-\delta$, we have

where $A=\sqrt{\beta+2T\sigma^{2}\ln\frac{3Te}{\delta}}$, $B=\sqrt{2}$, $C=\frac{4(f(\boldsymbol{x}_{1})-f^{\star})}{\alpha}+\frac{8M(3-\mu)d\alpha\sigma^{2}}{\beta(1-\mu)}\ln\frac{3Te}{\delta}+\frac{3d(1-\mu^{T})^{2}\sigma^{2}}{\beta(1-\mu)^{2}}\ln\frac{3}{\delta}$ and $D=\frac{4\alpha dM(3-\mu)}{1-\mu}$. Using Lemma 7, we have that

We use this upper bound in the logarithmic term of (14). Thus, we have (13) again, this time with

Solving (14) by Lemma 8 and lower bounding $\sum_{t=1}^{T}\|\nabla f(\boldsymbol{x}_{t})\|^{2}$ by $T\min_{1\leq t\leq T}\|\nabla f(\boldsymbol{x}_{t})\|^{2}$, we get the stated bound. ∎