On the Convergence Proof of AMSGrad and a New Version

Introduction and our contributions

One of the most popular algorithms for training deep neural networks is stochastic gradient descent (SGD) and its variants. Among the various variants of SGD, the algorithm with the adaptive moment estimation Adam is widely used in practice. However, Reddi et al. have recently shown that the convergence proof of Adam is problematic and proposed a variant of Adam called AMSGrad to solve this issue.

Our contribution. In this paper, we point out a flaw in the convergence proof of AMSGrad. We then fix this flaw by providing a new convergence proof for AMSGrad in the case of special parameters. In addition, in the case of general parameters, we propose a new and slightly modified version of AMSGrad.

To provide more details, let us recall AMSGrad in Algorithm 1, in which the mathematical notation can be fully found in Section 2.

Let $x_{t}$ and $v_{t}$ be the sequences obtained from Algorithm 1, $\alpha_{t}=\frac{\alpha}{\sqrt{t}}$, $\beta_{1}=\beta_{1,1}$, $\beta_{1,t}\leq\beta_{1}$ for all $t\in[T]$ and $\frac{\beta_{1}}{\sqrt{\beta_{2}}}\leq 1$. Assume that $\mathcal{F}$ has bounded diameter $D_{\infty}$ and $\lVert{\nabla f_{t}(x)}\rVert_{\infty}\leq G_{\infty}$ for all $t\in[T]$ and $x\in\mathcal{F}$. For $x_{t}$ generated using AMSGrad (Algorithm 1), we have the following bound on the regret:

In their proof for Theorem A, Reddi et al. resolved an issue on the so-called telescopic sum in the convergence proof of Adam ([2, Theorem 10.5]). Specifically, Reddi et al. adjusted $\hat{v}_{t}$ such that all components in the vector

are always positive. However, there is another issue (showed in Section 3) in the convergence proof of Adam that AMSGrad unfortunately neglects. The issue affects both the correctness of Reddi et al.’s proof and the upper bound for the regret in Theorem A. To deal with the issue in a general way, we propose to modify Algorithm 1 such that all components in the vector

are always positive. The differences with (1.0.1) are highlighted in the boxes for clarity.

Paper roadmap. We begin with preliminaries in Section 2. We show where the proof of Theorem A becomes invalid in Section 3. After that, we suggest two ways to resolve the issue in Sections 4 and 5.

Subsequent works. The first version of this paper publicly appeared on arXiv on 7 April 2019. On 19 April 2019, Reddi et al. revised their original proofshttps://arxiv.org/abs/1904.09237. The revised proof does not suffer from the issue pointed out in Section 3 of this paper, although yielding a constant factor missing in the original claims.

Preliminaries

where $x^{*}=\text{argmin}_{x\in\mathcal{F}}\sum_{t=1}^{T}f_{t}(x)$.

where $(-)^{{\sf T}}$ denotes the transpose of $(-)$.

Issue in the convergence proof of AMSGrad

Before showing the issue in the convergence proof of AMSGrad, let us recall and prove the following inequality, which also appears in .

Algorithm 1 achieves the following guarantee, for all $T\geq 1$:

and $\prod_{\mathcal{F},\sqrt{\hat{V}_{t}}}(x^{*})=x^{*}$ for all $x^{*}\in\mathcal{F}$. For all $1\leq t\leq T$, put $g_{t}=\nabla_{x}f_{t}(x_{t})$. Using Lemma 2.8 with $u_{1}=x_{t+1}$ and $u_{2}=x^{*}$, we have

Moreover, by Lemma 2.2, we have $f_{t}(x^{*})-f_{t}(x_{t})\geq g_{t}^{\text{T}}(x^{*}-x_{t})$, where $g_{t}^{{\sf T}}$ denotes the transpose of vector $g_{t}$. This means that

On the other hand, for all $t\geq 2$, we have

where the inequality is from the fact that $ab\leq a^{2}/2+b^{2}/2$ for any $a,b$. Hence,

Since $\beta_{1,t}\leq\beta_{1}(1\leq t\leq T)$, we obtain

where the last inequality is from the assumption that $\beta_{1,t}\leq\beta_{1}<1(1\leq t\leq T)$. Therefore,

and we obtain the desired bound for $R(T)$. ∎

Issue in the convergence proof of AMSGrad. We denote the terms on the right hand-side of the upper bound for $R(T)$ in Lemma 3.1 as

The issue in the proof of the convergence theorem of AMSGrad [3, Theorem 4] becomes on examining the term (3.1.3). Indeed, in [3, page 18], Reddi et al. usedConcretely, on page 18 of , it is stated that “The […] inequality use the fact that $\beta_{1,t}\leq\beta_{1}$.” the property that $\beta_{1,t}\leq\beta_{1}$, and hence

to replace all $\beta_{1,t}$ by $\beta_{1}$ as

However, the first inequality (in red) is not guaranteed because the quantity

in (3.1.3) may be both negative and positive as shown in Example 3.2. This is also a neglected issue in the convergence proofs in Kingma and Ba [2, Theorem 10.5], Luo et al. [5, Theorem 4], Bock et al. [6, Theorem 4.4], and Chen and Gu [7, Theorem 4.2].

We use the function in the Synthetic Experiment of Reddi et al. [3, Page 6]

with the constraint set $\mathcal{F}=$. The optimal solution is $x^{*}=-1$. By the proof of [3, Theorem 1], the initial point $x_{1}=1$. By Algorithm 1, $m_{0}=0$, $v_{0}=0$, and $\hat{v}_{0}=0$. We choose $\beta_{1}=0.9$, $\beta_{1,t}=\beta_{1}\lambda^{t-1}$, where $\lambda=0.001$, $\beta_{2}=0.999$, and $\alpha_{t}=\alpha/\sqrt{t}$, where $\alpha=0.001$. Under this setting, we have $f_{1}(x_{1})=1010x_{1}$, $f_{2}(x_{2})=-10x_{2}$, $f_{3}(x_{3})=-10x_{3}$ and hence

Since $x_{1}-\alpha_{1}m_{1}/\sqrt{\hat{v}_{1}}>0$, we have

Since $x_{2}-\alpha_{2}m_{2}/\sqrt{\hat{v}_{2}}>0$, we obtain

Outline of our solution. Let us rewrite (3.1.3) as

Omitting the term $\sum_{i=1}^{d}\frac{\sqrt{\hat{v}_{T,i}}}{2\alpha_{T}(1-\beta_{1,T})}(x_{T+1,i}-x^{*}_{,i})^{2}$, we obtain

in which the differences with Reddi et al. are highlighted in the boxes, namely, $\boxed{\beta_{1,t}}$ and $\boxed{\beta_{1,t-1}}$ instead of $\beta_{1}$.

We suggest two ways to overcome these differences depending on the setting of $\beta_{1,t}(1\leq t\leq T)$:

In Section 4: If either $\beta_{1,t}\overset{\Delta}{=}\beta_{1}\lambda^{t-1}$ or $\beta_{1,t}\overset{\Delta}{=}1/t$, $(1\leq t\leq T)$, where $0\leq\beta_{1}<1$ and $0<\lambda<1$, then we give a new convergence theorem for AMSGrad in Section 4.

In Section 5: If the setting for $\beta_{1,t}(1\leq t\leq T)$ is general, as in the statement of Theorem A, then we suggest a new (slightly modified) version for AMSGrad in Section 5.

New convergence theorem for AMSGrad

When either $\beta_{1,t}\overset{\Delta}{=}\beta_{1}\lambda^{t-1}$ or $\beta_{1,t}\overset{\Delta}{=}1/t$, $(1\leq t\leq T)$, where $0\leq\beta_{1}<1$ and $0<\lambda<1$, Theorem A can be fixed as follows, in which the upper bounds of the regret $R(T)$ are changed.

Let $x_{t}$ and $v_{t}$ be the sequences obtained from Algorithm 1, $\alpha_{t}=\frac{\alpha}{\sqrt{t}}$, either $\beta_{1,t}=\beta_{1}\lambda^{t-1}$, where $\lambda\in(0,1)$, or $\beta_{1,t}=\frac{\beta_{1}}{t}$ for all $t\in[T]$ and $\gamma=\frac{\beta_{1}}{\sqrt{\beta_{2}}}\leq 1$. Assume that $\mathcal{F}$ has bounded diameter $D_{\infty}$ and $\lVert{\nabla f_{t}(x)}\rVert_{\infty}\leq G_{\infty}$ for all $t\in[T]$ and $x\in\mathcal{F}$. For $x_{t}$ generated using AMSGrad (Algorithm 1), we have the following bound on the regret. Then, there is some $1\leq t_{0}\leq T$ such that AMSGrad achieves the following guarantee for all $T\geq 1$:

provided $\beta_{1,t}=\beta_{1}\lambda^{t-1}$, and

provided $\beta_{1,t}=\frac{\beta_{1}}{t}$.

To prove Theorem 4.1, we need the following Lemmas 4.2, 4.3, and 4.4.

From the definition of $\hat{v}_{t}$ in AMSGrad’s algorithm, it is implied that $\hat{v}_{t}=\max\{v_{1},...,v_{t}\}$. Therefore, there is some $1\leq s\leq t$ such that $\hat{v}_{t}=v_{s}$. Hence,

where the last inequality is by Lemma 2.4. ∎

If either $\beta_{1,t}=\beta_{1}\lambda^{t-1}$ or $\beta_{1,t}=\beta_{1}/t$, then there exists some $t_{0}$ such that for every $t>t_{0}$,

Since $\hat{v}_{t,i}\geq\hat{v}_{t-1,i}$, it is sufficient to prove that there exists some $t_{0}$ such that for every $t>t_{0}$,

When $\beta_{1,t}=\beta_{1}/t$, from (4.3.1) we have

When $\beta_{1,t}=\beta_{1}\lambda^{t-1}$, (4.3.1) have the following form

Since $\beta_{1}$ and $\lambda$ are smaller than $1$, it is easy to see that when $t$ is sufficiently large, meaning that $t>t_{0}$ for some $t_{0}$, the left-hand side of (4.3.2) is $1-O(1/t^{2})$ and the left-hand side of (4.3.3) is larger than $1-\beta_{1}\lambda^{t-2}=1-O(\lambda^{t-2})$. Therefore, (4.3.2) and (4.3.3) hold when $t$ is sufficiently large. ∎

For the parameter settings and conditions assumed in Theorem 4.1, we have

The proof is almost identical to that of [3, Lemma 2]. Since for all $t\geq 1$, $\hat{v}_{t,i}\geq v_{t,i}$, we have

where the second inequality is by Lemma 2.3, the third inequality is from the properties of $\beta_{1,k}\leq 1$ and $\beta_{1,k}\leq\beta_{1}$ for all $1\leq k\leq T$, and the fourth inequality is obtained by applying Lemma 2.4 to $\sum_{k=1}^{t}\beta_{1}^{t-k}$. Therefore,

where the second inequality is by Lemma 2.7. Therefore

It is sufficient to consider $\sum_{t=1}^{T}\frac{1}{\sqrt{t}}\sum_{k=1}^{t}\gamma^{t-k}|{g_{k,i}}|$. Firstly, $\sum_{t=1}^{T}\frac{1}{\sqrt{t}}\sum_{k=1}^{t}\gamma^{t-k}|{g_{k,i}}|$ can be expanded as

Changing the role of $|g_{1,i}|$ as the common factor, we obtain

where the last inequality is by Lemma 2.4, we obtain

where the first inequality is by Lemma 2.3 and the last inequality is by Lemma 2.5, we obtain

To prove Theorem 4.1, by Lemma 3.1, we need to bound the terms (3.1.3), (3.1.4), and (3.1.5). First, we consider (3.1.4). We have

where the equality is by the assumption that $\alpha_{t}=\alpha/\sqrt{t}$ and the last inequality is by Lemma 4.4. Next, we consider (3.1.5). The bound for (3.1.5) depends on either $\beta_{1,t}=\beta_{1}\lambda^{t-1}(0<\lambda<1)$ or $\beta_{1,t}=\frac{\beta_{1}}{t}$. Recall that by assumption, $\lVert{x_{m}-x_{n}}\rVert_{\infty}\leq D_{\infty}$ for any $m,n\in\{1,...,T\}$, $\alpha_{t}=\alpha/\sqrt{t}$. If $\beta_{1,t}=\beta_{1}\lambda^{t-1}(0<\lambda<1)$, then,

where the first inequality is from Lemma 4.2 and the assumption that $\beta_{1}\leq 1$, the last inequality is by Lemma 2.4. If $\beta_{1,t}=\frac{\beta_{1}}{t}$, then,

where the first inequality is from Lemma 4.2 and the assumption that $\beta_{1}\leq 1$, and the last inequality is by Lemma 2.6.

Finally, we will bound (3.1.3). From the inequality (3) and replacing $\alpha_{t}$ with $\frac{\alpha}{\sqrt{t}}(1\leq t\leq T)$, we obtain

By Lemma 4.3, there is some $t_{0}(1\leq t_{0}\leq T)$ such that $\frac{\sqrt{t\hat{v}_{t,i}}}{1-\beta_{1,t}}\geq\frac{\sqrt{(t-1)\hat{v}_{t-1,i}}}{1-\beta_{1,t-1}}$ for all $t>t_{0}$. Therefore,

where the second inequality is obtained by omitting the term $\frac{1}{2\alpha}\sum_{i=1}^{d}\sum_{t=2}^{t_{0}}(x_{t,i}-x^{*}_{,i})^{2}\frac{\sqrt{(t-1)\hat{v}_{t-1,i}}}{1-\beta_{1,t-1}}$, and the last inequality is by Lemma 4.2 and the assumption that $\beta_{1,t}\leq\beta_{1}(1\leq t\leq T)$. Summing up, if $\beta_{1,t}=\beta_{1}\lambda^{t-1}$, then, from (4), (4), and (4), we obtain

If $\beta_{1,t}=\frac{\beta_{1}}{t}$, then, from from (4), (4), and (4), we obtain

The following corollary shows that, when either $\beta_{1,t}=\beta_{1}\lambda^{t-1}$ or $\beta_{1,t}=1/t$, $(1\leq t\leq T)$, where $0\leq\beta_{1}<1$ and $0<\lambda<1$, the average regret of AMSGrad converges.

With the same assumption as in Theorem 4.1, AMSGrad achieves the following guarantee:

The result is obtained by using Theorem 4 and the following fact:

where the inequality is from the assumption that $\lVert{g_{t}}\rVert_{\infty}\leq G_{\infty}$ for all $t\in[T]$. ∎

New version of AMSGrad optimizer: AdamX

With this Algorithm 2, the regret is bounded as follows.

Let $x_{t}$ and $v_{t}$ be the sequences obtained from Algorithm 2, $\alpha_{t}=\frac{\alpha}{\sqrt{t}}$, $\beta_{1}=\beta_{1,1}$, $\beta_{1,t}\leq\beta_{1}$ for all $t\in[T]$ and $\frac{\beta_{1}}{\sqrt{\beta_{2}}}\leq 1$. Assume that $\mathcal{F}$ has bounded diameter $D_{\infty}$ and $\lVert{\nabla f_{t}(x)}\rVert_{\infty}\leq G_{\infty}$ for all $t\in[T]$ and $x\in\mathcal{F}$. For $x_{t}$ generated using the AdamX (Algorithm 2), we have the following bound on the regret:

To prove Theorem 5.1, we need the following Lemmas 5.2, 5.3, and 5.4.

We will prove (5.2.1) by induction on $t$. Recall that by the update rule on $\hat{v}_{t}$, we have $\hat{v}_{1}\overset{\Delta}{=}v_{1}$ and $\hat{v}_{t}\overset{\Delta}{=}\max\{\frac{(1-\beta_{1,t})^{2}}{(1-\beta_{1,t-1})^{2}}\hat{v}_{t-1},v_{t}\}$ if $t\geq 2$. Therefore,

and the (5.2.1) holds for all $1\leq j\leq t-1$. Since

For all $t\geq 1$, we have $\sqrt{\hat{v}_{t}}\leq\frac{G_{\infty}}{1-\beta_{1}}$, where $\hat{v}_{t}$ is in Algorithm 2.

Therefore, there is some $1\leq s\leq t$ such that $\hat{v}_{t}=\frac{(1-\beta_{1,t})^{2}}{(1-\beta_{1,s})^{2}}v_{s}$. Hence,

For the parameter settings and conditions assumed in Theorem 5.1, we have

by Lemma 5.2, we have $\hat{v}_{t,i}\geq v_{t,i}$, and hence the proof is the same as that of Lemma 4.4. ∎

Similarly to the proof of Theorem 4.1, we need to bound (3.1.3), (3.1.4), and (3.1.5). By using Lemma 5.4, we obtain the same bound for (3.1.4) as in the proof of Theorem 4.1, that is,

where the last inequality is by Lemma 5.4. Now we bound (3.1.5). By the assumption that $\lVert{x_{m}-x_{n}}\rVert_{\infty}\leq D_{\infty}$ for any $m,n\in\{1,...,T\}$, $\alpha_{t}=\alpha/\sqrt{t}$, and $\beta_{1,t}=\beta_{1}\lambda^{t-1}\leq\beta_{1}\leq 1$, we obtain

Finally, we will bound (3.1.3). By the inequality (3) and replacing $\alpha_{t}=\frac{\alpha}{\sqrt{t}}(1\leq t\leq T)$, we obtain

Moreover, by the update rule of Algorithm 2, we have

Therefore, $\hat{v}_{t,i}\geq\frac{(1-\beta_{1,t})^{2}}{(1-\beta_{1,t-1})^{2}}\hat{v}_{t-1,i}$, and hence

Now by the positivity of the essential formula $\frac{\sqrt{t\hat{v}_{t,i}}}{1-\beta_{1,t}}-\frac{\sqrt{(t-1)\hat{v}_{t-1,i}}}{1-\beta_{1,t-1}}$, we obtain

where the last inequality is by Lemma 5.3. Hence we obtain the desired upper bound for $R(T)$. ∎

With the same assumption as in Theorem 5.1, and for all $0\leq\beta_{1,t}<1$ satisfying

By Theorem 5.1, it is sufficient to consider the term

on the right hand side of the upper bound for $R(T)$ in Theorem 5.1. Because $\frac{dD_{\infty}^{2}G_{\infty}}{2\alpha(1-\beta_{1})^{2}}$ is bounded and does not depend on $T$, the statement follows. ∎

When either $\beta_{1,t}=\beta_{1}\lambda^{t-1}$ for some $\lambda\in(0,1)$, or $\beta_{1,t}=\frac{1}{t}$ in Theorem 5.1, we obtain the following guarantee that the average regret of AdamX converges.

With the same assumption as in Theorem 5.1, and either $\beta_{1,t}=\beta_{1}\lambda^{t-1}$ for some $\lambda\in(0,1)$, or $\beta_{1,t}=\frac{1}{t}$, AdamX achieves the following guarantee:

By Corollary 5.5, it is sufficient to consider the term

When $\beta_{1,t}=\beta_{1}\lambda^{t-1}$ for some $\lambda\in(0,1)$, we have

where the first inequality is from the property that $\beta_{1}\leq 1$, and the last inequality is from Lemma 2.4. When $\beta_{1,t}=\frac{1}{t}$, we obtain

where the last inequality is from Lemma 2.6. Now, by combining (5) and (5) with Corollary 5.5, we obtain the desired result. ∎

Experiments

While we consider our main contributions as the theoretical analyses on AMSGrad and AdamX in the previous sections, we provide experimental results in this section for AMSGrad and AdamX. Concretely, we use the PyTorch code for AMSGradhttps://pytorch.org/docs/stable/_modules/torch/optim/adam.html via setting the boolean flag amsgrad = True. The code for AdamX is based on that of AMSGrad, with corresponding modifications as in Algorithm 2. The parameters for AMSGrad and AdamX are identical in our experiments, namely $(\beta_{1},\beta_{2})=(0.9,0.999)$, the term added to the denominator to improve numerical stability is $\epsilon=10^{-8}$, and and additionally we set $\beta_{1,t}=\beta_{1}\lambda^{t-1}$ with $\lambda=0.001$ to make use of Corollary 5.6 on the convergence of AdamX.

The learning rate is scheduled for both optimizers AMSGrad and AdamX as follows: $10^{-3}$, $10^{-4}$, $10^{-5}$, $10^{-6}$, $10^{-6}/2$ if the epoch is correspondingly in the ranges $$,$$, $$,$$, $$. We use CIFARhttps://www.cs.toronto.edu/ kriz/cifar.html-10 (containing 50000 training images and 10000 test images of size$32\times 32$) as the dataset and the residual networks ResNet18 and PreActResNet18 for training with batch size is 128. The testing result is given in Figure 1 where one can see that AMSGrad and AdamX behaves similarly, which supports our theoretical results on the convergence of both AMSGrad (Section 4) and AdamX (Section 5).

Conclusion

We have shown that the convergence proof of AMSGrad is problematic, and presented various fixes for it, which include a new and slightly modified version called AdamX. Along the lines, we also observe that the issue has been neglected in various works such as in [2, Theorem 10.5], [5, Theorem 4], [6, Theorem 4.4], [7, Theorem 4.2]. Our work helps ensure the theoretical foundation of those optimizers.

References