Title: Closing the gap between the upper bound and lower bound of Adam's iteration complexity
Venue: NeurIPS 2023
Accepted: yes
Keywords: Adam

Convergence

Upper Bound

Lower Bound
Abstract: Recently,  Arjevani et al. [1]  establish a lower bound of iteration complexity for the first-order optimization under an $L$-smooth condition and a bounded noise variance assumption. However, a thorough review of existing literature on Adam's convergence reveals a noticeable gap: none of them meet the above lower bound. In this paper, we close the gap by deriving a new convergence guarantee of Adam, with only an $L$-smooth condition and a bounded noise variance assumption. Our results remain valid across a broad spectrum of hyperparameters. Especially with properly chosen hyperparameters, we derive an upper bound of the iteration complexity of Adam and show that it meets the lower bound for first-order optimizers. To the best of our knowledge, this is the first to establish such a tight upper bound for Adam's convergence. Our proof utilizes novel techniques to handle the entanglement between momentum and adaptive learning rate and to convert the first-order term in the Descent Lemma to the gradient norm, which may be of independent interest.

Closing the Gap Between the Upper Bound and the Lower Bound of Adam's Iteration Complexity

Abstract

Recently, Arjevani et al. [1] establish a lower bound of iteration complexity for the first-order optimization under an L -smooth condition and a bounded noise variance assumption. However, a thorough review of existing literature on Adam's convergence reveals a noticeable gap: none of them meet the above lower bound. In this paper, we close the gap by deriving a new convergence guarantee of Adam, with only an L -smooth condition and a bounded noise variance assumption. Our results remain valid across a broad spectrum of hyperparameters. Especially with properly chosen hyperparameters, we derive an upper bound of iteration complexity of Adam and show that it meets the lower bound for first-order optimizers. To the best of our knowledge, this is the first to establish such a tight upper bound for Adam's convergence. Our proof utilizes novel techniques to handle the entanglement between momentum and adaptive learning rate and to convert the first-order term in the Descent Lemma to the gradient norm, which may be of independent interest.

1 Introduction

First-order optimizers, also known as gradient-based methods, make use of gradient (first-order derivative) information to find the minimum of a function. They have become a cornerstone of many machine learning algorithms due to the efficiency as only gradient informaiton is required, and the flexibility as gradients can be easily computed for any function represented as directed acyclic computational graph via auto-differentiation [2, 25].

Therefore, it is fundamental to theoretically understand the properties of these first-order methods. Recently, Arjevani et al. [1] establish a lower bound on the iteration complexity of stochastic firstorder methods. Formally, for a well-studied setting where the objective is L -smooth and a stochastic oracle can query the gradient unbiasly with bounded variance (see Assumption 1 and 2), any stochastic first-order algorithm requires at least ε -4 queries (in the worst case) to find an ε -stationary point, i.e., a point with gradient norm at most ε . Arjevani et al. [1] further show that the above lower bound is tight as it matches the existing upper bound of iteration complexity of SGD [15].

On the other hand, among first-order optimizers, Adam [20] becomes dominant in training stateof-the-art machine learning models [3, 18, 4, 11]. Compared to vanilla stochastic gradient descent (SGD), Adam consists of two more key components: (i) momentum to accumulate historical gradient

information and (ii) adaptive learning rate to rectify coordinate-wise step sizes. The psedo-code of Adam is given as Algorithm 1. While the sophisticated design of Adam enables its empirical superiority, it brings great challenges for the theoretical analysis. After examining a series of theoretical works on the upper bound of iteration complexity of Adam [33, 9, 10, 36, 16, 27, 34], we find that none of them match the lower bound for first-order optimizers: they not only consume more queries than the lower bound to reach ε -stationary iterations but also requires additional assumptions (see Section 3 for a detailed discussion).

This theoretical mismatch becomes even more unnatural given the great empirical advantage of Adam over SGD, which incites us to think:

Is the gap between the upper and lower bounds for Adam a result of the inherent complexity induced by Adam's design, or could it be attributed to the proof techniques not being sharp enough?

This paper answers the above question, validating the latter hypothesis, by establishing a new upper bound on iteration complexity of Adam for a wide range of hyperparameters that cover typical choices. Specifically, our contribution can be summarized as follows:

We examine existing works that analyze the iteration complexity of Adam, and find that none of them meets the lower bound of first-order optimization algorithms;

We derive a new convergence guarantee of Adam with only assuming L -smooth condition and bounded variance assumption (Theorem 1), which holds for a wide range of hyperparameters covering typical choices;

With chosen hyperparameters, we further tighten Theorem 1 and show that the upper bound on the iteration complexity of Adam meets the lower bound, closing the gap (Theorem 2). Our upper bound is tighter than existing results by a logarithmic factor, in spite of weaker assumption.

To the best of our knowledge, this work provides the first upper bound on the iteration complexity of Adam without additional assumptions other than L -smooth condition and bounded variance assumption. It is also the first upper bound matching the lower bound of first-order optimizers.

Organization of this paper. The rest of the paper is organized as follows: in Section 2, we first present the notations and settup of analysis in this paper ; in Section 3, we revisit the existing works on the iteration complexity of Adam; in Section 4, we present a convergence analysis of Adam with general hyperparameters (Theorem 1); in Section 5, we tighten Theorem 1 with a chosen hyperparameter, and derive an upper bound of Adam's iteration complexity which meets the lower bound; in Section 6, we discuss the limitation of our results.

2 Preliminary

The Adam algorithm is restated in Agorithm 1 for convenient reference. Note that compared to the orignal version of Adam in Kingma and Ba [20], the bias-correction terms are omitted to simplify the analysis, and our analysis can be immediately extended to the original version of Adam because the effect of bias-correction term decays exponentially. Also, in the original version of Adam, the adaptive learning rate is η √ ν t + λ 1 d instead of η √ ν t . However, our setting is more challenging and our result can be easily extend to the original version of Adam, since the λ term makes the adaptive learning rate upper bounded and eases the analysis.

Algorithm 1 Adam

Input: Stochastic oracle O , learning rate η > 0 , initial point w 1 ∈ R d , initial conditioner ν 0 ∈ R + , initial momentum m 0 , momentum parameter β 1 , conditioner parameter β 2 , number of epoch T

1: Sample r ∼ Unif { 1 , · · · , T }

2: For t = 1 → T :

3: Generate a random z t , and query stochastic oracle g t = O f ( w t , z t )

4: Calculate ν t = β 2 ν t -1 +(1 -β 2 ) g glyph[circledot] 2 t

5: Calculate m t = β 1 m t -1 +(1 -β 1 ) g t

6: Update w t +1 = w t -η 1 √ ν t glyph[circledot] m t

7: EndFor

Output: w r

Notations. For a, b ∈ Z ≥ 0 and a ≤ b , denote [ a, b ] = { a, a +1 , · · · , b -1 , b } . For any two vectors w , v ∈ R d , denote w glyph[circledot] v as the Hadamard product (i.e., coordinate-wise multiplication) between w and v . When analyzing Adam, we denote the true gradient at iteration t as G t = ∇ f ( w t ) , and the sigma algebra before iteration t as F t = σ ( g 1 , · · · , g t -1 ) . We denote conditional expectation as E |F t [ ∗ ] = E [ ∗|F t ] . We also use asymptotic notations o , O , Ω , and Θ , where h 2 ( x ) = o x → x 0 ( h 1 ( x )) means that lim x → x 0 h 2 ( x ) h 1 ( x ) = 0 (when the context is clear, we abbreviate x → x 0 and only use o ( h 1 ( x )) ); h 2 ( x ) = O ( h 1 ( x )) means that there exists constant γ independent of x such that h 2 ( x ) ≤ γh 1 ( x ) ; h 2 ( x ) = Ω( h 1 ( x )) means that h 1 ( x ) = O ( h 2 ( x )) ; and h 2 ( x ) = Θ( h 1 ( x )) means that h 2 ( x ) = O ( h 1 ( x )) and h 2 ( x ) = Ω( h 1 ( x )) .

Objective function. In this paper, we consider solving the following optimization problem: min w ∈ R d f ( w ) . We make the following assumption on the objective function f .

Assumption 1 (On objective function) . We assume f to be non-negative. We further assume that f satisfies L -smooth condition, i.e., f is differentiable, and the gradient of f is L -Lipschitz.

We denote the set of all objective functions satisfying Assumption 1 as F ( L ) .

Stochastic oracle. As f is differentiable, we can utilize the gradient of f (i.e., ∇ f ) to solve the above optimization problem. However, the ∇ f is usually expensive to compute. Instead, we query a stochastic estimation of ∇ f through a stochastic oracle O . Specifically, the stochastic oracle O consists of a distribution P over a measurable space Z and a mapping O f : R d ×Z → R d . We make the following asssumption on O .

Assumption 2 (On stochastic oracle) . We assume that O is unbiased, i.e., ∀ w ∈ R d , E z ∼P O f ( w , z ) = ∇ f ( w ) . We further assume O has bounded variance, i.e., ∀ w ∈ R d , E z ∼P [ ‖ O f ( w , z ) -∇ f ( w ) ‖ 2 ] ≤ σ 2 .

We denote the set of all stochastic oracles satisfying Assumption 2 with variance bound σ 2 as O ( σ 2 ) .

Algorithm. Adam belongs to first-order optimization algorithms, which is defined as follows:

Definition 1 (First-order optimization algorithm) . An algorithm A is called a first-order optimization algorithm, if it takes an input w 1 and hyperparameter θ , and produces a sequence of parameters as follows: first sample a random seed r from some distribution P r * , set w A ( θ ) 1 = w 1 and then update the parameters as

where z 1 , z 2 , · · · , z t are sampled i.i.d. from P .

Iteration complexity. Denote the set of all first-order optimization algorithms as A first . We next introduce iteration complexity to measure the convergence rate of optimization algorithms.

Definition 2 (Iteration complexity) . The iteration complexity of first-order optimization algorithm A is defined as

Furthermore, the iteration complexity of the family of first-order optimization algorithms A first is

It should be noticed that the iteration complexity of the family of first-order optimization algorithms is a lower bound of the iteration complexity of a specific first-order optimization algorithm, i.e., ∀ A ∈ A first , C ε ( A , ∆ , L, σ 2 ) ≥ C ε (∆ , L, σ 2 ) .

3 Related works: none of existing upper bounds match the lower bound

In this section, we examine existing works that study the iteration complexity of Adam, and defer a discussion of other related works to Appendix A. Specifically, we find that none of them match the lower bound for first-order algorithms provided in [1] (restated as follows).

Note that in the above bound, we omit the dependence of the lower bound over ∆ , L , and σ 2 , which is a standard practice in existing works (see Cutkosky and Mehta [8], Xie et al. [32], Faw et al. [13] as examples) because the dependence over the accuracy ε can be used to derive how much additional iterations is required for a smaller target accuracy and is thus of more interest. In this paper, when we say "match the lower bound", we always mean that the upper bound has the same order of ε as the lower bound.

Generally speaking, existing works on the iteration complexity of Adam can be divided into two categories: they either (i) assume that gradient is universally bounded or (ii) make stronger assumptions on smoothness. Below we respectively explain how these two categories of works do not match the lower bound in [1].

The first line of works, including Zaheer et al. [33], De et al. [9], Défossez et al. [10], Zou et al. [36], Guo et al. [16], assume that the gradient norm of f is universally bounded, i.e., ‖∇ f ( w ) ‖ ≤ G , ∀ w ∈ R d . In other words, what they consider is another iteration complexity defined as follows:

This line of works do not match the lower bound due to the following two reasons: First of all, the upper bound they derive is O ( log 1 /ε ε 4 ) , which has an additional log 1 /ε factor more than the lower bound; secondly, the bound they derive is for C ε ( A , ∆ , L, σ 2 , G ) . Note that F ( L ) ∩{ f : ‖∇ f ‖ ≤ G } is a proper subset of F ( L ) for any G , where a simple example in F ( L ) but without bounded gradient is the quadratic function f ( x ) = ‖ x ‖ 2 . Therefore, we have that

and thus the upper bound on C ε ( A , ∆ , L, σ 2 , G ) does not apply to C ε ( A , ∆ , L, σ 2 ) . Moreover, their upper bound of C ε ( A , ∆ , L, σ 2 , G ) tends to ∞ as G →∞ , which indicates that if following their analysis, the upper bound of C ε ( A , ∆ , L, σ 2 ) would be infinity based on Eq. (1).

The second line of works [27, 34, 30] additionally assume a mean-squared smoothness property besides Assumption 1 and 2, i.e., E z ∼P ‖ O f ( w , z ) -O f ( v , z ) ‖ 2 ≤ L ‖ w -v ‖ 2 . Denote ˜ O ( σ 2 , L ) glyph[defines] { O : E z ∼P ‖ O f ( w , z ) -O f ( v , z ) ‖ 2 ≤ L ‖ w -v ‖ 2 , ∀ w , v ∈ R d }∩ O ( σ 2 ) . The iteration complexity that they consider is defined as follows:

The rate derived in [27, 34, 30] is O ( log 1 /ε ε 6 ) , which is derived by minimizing the upper bounds in [27, 34, 30] with respect to the hyperparameter of adaptive learning rate β 2 . According to Arjevani et al. [1], the lower bound of iteration complexity of ˜ C ε ( A , ∆ , L, σ 2 ) is Ω( 1 ε 3 ) and smaller than the original lower bound Ω( 1 ε 4 ) , resulting in an even larger gap between the upper and lower bounds.

Recently, there is a concurrent work [21] which does not require bounded gradient assumption and mean-squared smoothness property but poses a stronger assumption on the stochastic oracle: the set of stochastic oracles they consider is ˜ ˜ O = { O : ∀ w ∈ R d , E z ∼P O f ( w , z ) = ∇ f ( w ) , P ( ‖ O f ( w , z ) -∇ f ( w ) ‖ 2 ≤ σ 2 ) = 1 } . ˜ ˜ O is a proper subset of O because a simple example is that O f ( w , z ) = ∇ f ( w ) + z where z is a standard gaussian variable. Therefore, their result does not provide a valid upper bound of C ε ( A , ∆ , L, σ 2 ) .

4 Convergence analysis of Adam with only Assumptions 1 and 2

As discussed in Section 3, existing works on analyzing Adam require additional assumptions besides Assumption 1 and 2. In this section, we provide the first convergence analysis of Adam with only Assumption 1 and 2, which naturally gives an upper bound on the iteration complexity C ε ( A , ∆ , L, σ 2 ) . In fact, our analysis even holds when the stochastic oracle satisfies the following more general assumption.

Assumption 3 (Coordinate-wise affine noise variance) . We assume that O is unbiased, i.e., ∀ w ∈ R d , E z ∼P O f ( w , z ) = ∇ f ( w ) . We further assume O has coordinate-wise affine variance, i.e., ∀ w ∈ R d and ∀ i ∈ [ d ] , E z ∼P [ | ( O f ( w , z )) i | 2 ] ≤ σ 2 0 + σ 2 1 ∂ i f ( w ) 2 .

One can easily observe that Assumption 3 is more general than Assumption 2 since Assumption 2 immediately indicates Assumption 3 with σ 0 = σ and σ 1 = 1 . We consider Assumption 3 not only because it is more general but also because it allows the noise to grow with the norm of the true gradient, which is usually the case in machine learning practice [14, 19].

Our analysis under Assumption 1 and Assumption 3 is then given as follows.

Theorem 1. Let A be by Adam (Algorithm 1) and θ = ( η, β 1 , β 2 ) are the hyperparameters of A . Let Assumption 1 and 2 hold. Then, if 0 ≤ β 1 ≤ √ β 2 -8 σ 2 1 (1 -β 2 ) β -2 2 and β 2 < 1 , we have

where ν 0 ,i is the i -th coordinate of ν 0 ,

A proof sketch is given in Section 4.2 and the full proof is deferred to Appendix.

The right-hand side in Eq. (2) looks messy at the first glance. We next explain Theorem 1 in detail and make the upper bound's dependence over hyperparameters crystally clear.

4.1 Discussion on Theorem 1

Required assumptions and conditions. As mentioned previously, Theorem 1 only requires Assumption 1 and 2, which aligns with the setting of the lower bound (Proposition 1). To our best knowledge, this is the first analysis of Adam without additional assumptions.

glyph[negationslash]

As for the range of β 1 and β 2 , one can immediately see that the condition β 1 ≤ √ β 2 -8 σ 2 1 (1 -β 2 ) β -2 2 degenerates to β 1 ≤ √ β 2 in the bounded gradient case (i.e., σ 1 = 0 ), the weakest condition required in existing literature [36]. When σ 1 = 0 , such a condition is stronger than β 1 ≤ √ β 2 . We point out that this is not due to technical limitations but instead agrees with existing counterexamples for Adam: Reddi et al. [26], Zhang et al. [34] show that when σ 1 = 0 , there exists a counterexample glyph[negationslash]

satisfying Assumption 1 and Assumption 3 and a pair of ( β 1 , β 2 ) with β 1 < √ β 2 and Adam with ( β 1 , β 2 ) diverges over such a counterexample.

Dependence over β 2 , η , and T . Here we consider the influence of β 2 , η , and T while fixing β 1 constant (we will discuss the effect of β 1 in Section 6). With logarithmic factors ignored and coefficients hidden, C 1 , C 2 and the right-hand-side of Eq. (2) can be rewritten with asymptotic notations as

where ˜ O denotes O with logarithmic terms ignored. Consequently, the dependence of Eq. (2) over β 2 , η and T becomes

Here we consider two cases: (i). β 2 and η are independent over T , and (ii). β 2 and η are dependent over T . For case (i), based on the above equation, one can easily observe that the averaged gradient norm 1 T E ∑ T t =1 ‖∇ f ( w t ) ‖ will converge to the threshold O ( η 2 √ 1 -β 2 + 4 √ 1 -β 2 + η 4 √ 1 -β 2 ) with rate O (1 / √ T ) . This aligns with the observation in [27, 34] that Adam will not converge to the stationary point with constant β 2 .

For case (ii), in order to ensure convergence, i.e., min t ∈ [ T ] E ‖ G t ‖ 1 → 0 as T → ∞ , a sufficient condition is that the right-hand-side of the above equation is o ( T ) . Specifically, by choosing η = Θ( T -a ) and 1 -β 2 = Θ( T -b ) , we obtain that

By simple calculation, we obtain that the right-hand side of the above inequality is o (1) as T →∞ if and only if b > 0 , 1 > a > 0 and b -a < 1 . Moreover, the minimum of the right-hand side of the above inequality is ˜ O (1 /T 1 4 ) , which is achieved at a = 1 2 and b = 1 . Such a minimum implies an upper bound of the iteration complexity which at most differs from the lower bound by logarithmic factors as solving ˜ O (1 /T 1 4 ) = ε gives T = ˜ O ( 1 ε 4 ) . In Theorem 2, we will further remove the logarithmic factor by giving a refined proof when a = 1 2 and b = 1 and close the gap between the upper and lower bounds.

Dependence over λ . Our analysis allows λ = 0 in the adaptive learning rate η 1 √ ν t + λ 1 d . In contrast, some existing works [16, 21] require non-zero λ and their iteration complexity has polynomial dependence over 1 λ , which is less desired as λ can be as small as 10 -8 in practice (e.g., in PyTorch's default setting). Furthermore, compared to their setting, our setting is more challenging as non-zero λ immediately provides an upper bound of the adaptive learning rate.

4.2 Proof Sketch of Theorem 1

In this section, we demonstrate the proof idea of Theorem 1. Generally speaking, our proof is inspired by (i). the construction of the Lyapunov function for SGDM [22] and (ii) the construction of auxiliary function and the conversion from regret bound to gradient bound for AdaGrad [31], but the adaptation of these techniques to Adam is highly non-trivial, as SGDM does not hold an adaptive learning rate, and the adaptive learning rate of AdaGrad is monotonously decreasing. Below we sketch the proof by identifying three key challenges in the proof and provide our solutions respectively.

Challenge I: Disentangle the stochasticity in stochastic gradient and adaptive learning rate. For simplicity, let us first consider the case where β 1 = 0 , i.e., where the momentum m t degenerates to

the stochastic gradient g t . According to the standard descent lemma, we have that

The first challenge arises from bounding the "First Order" term above. To facilitate the understanding of the difficulty, we compare the "First Order" term of Adam to the corresponding "First Order" term of SGD, i.e., -η E 〈 G t , g t 〉 . By directly applying E |F t g t = G t , we obtain that the "First-Order" term of SGD equals to -η E ‖ G t ‖ 2 . However, as for Adam, we do not even know what E |F t 1 √ ν t glyph[circledot] g t is given that the stochasticity in g t and ν t entangles. A common practice is to use a surrogate adaptive learning rate ˜ ν t measurable with respect to F t , to approximate the real adaptive learning rate ν t . This leads to the following equation:

One can immediately see that "First Order Main" terms equals to E [ 〈 G t , -η 1 √ ˜ ν t glyph[circledot] G t 〉 ] < 0 , but now we need to handle the "Error" term. In existing literature, such a term is mostly bypassed by applying the bounded gradient assumption [10, 36], which, however, we do not assume.

Solution to Challenge I. Inspired by recent advance in the analysis of AdaGrad [31], we consider the auxiliary function ξ t = E [ η 〈 G t , -1 √ ˜ ν t +1 glyph[circledot] G t 〉 ] , where we choose ˜ ν t = β 2 ν t -1 +(1 -β 2 ) σ 2 0 1 d . In the following lemma, we show that the error term can be controlled using ξ t , parallel to (Lemma 4. [31]).

Lemma 1 (Informal version of Lemma 7 with β 1 = 0 ) . Let all conditions in Theorem 1 hold. Then,

In the right-hand-side of inequality (4), one can easily observe that the first term can be controlled by "First Order Main" term, and the third term is as small as the "Second Order" term. However, the second term seems annoying - in the analysis of AdaGrad [31], there is no 1 / √ β 2 factor, making the corresponding term a telescoping, but this is no longer true due to the existence of the 1 / √ β 2 factor. We resolve this difficulty by looking at the sum of 1 √ β 2 ξ t -1 -ξ t over t from 1 to T , which gives O ((1 -β 2 ) ∑ T -1 t =1 ξ t ) . By further noticing that ˜ ν t +1 ≥ β 2 ˜ ν t , we have

The right-hand-side term can thus be controlled by the "First Order Main" term when β 2 is close to 1 .

Remark 1. Compared to the analysis of AdaGrad in [31], our proof technique has two-fold novelties. First, our auxiliary function has an additional (1 -β 2 ) σ 2 0 1 d term, which is necessary for the analysis of Adam as it makes ˜ ν t lower bounded from 0 (AdaGrad does not need this, as ν t -1 of AdaGrad itself is lower bounded). Secondly, as discussed above, the "AdaGrad version" of second term in the right-hand-side of inequality (4) is a telescoping, the sum of which can be bounded straightforwardly.

glyph[negationslash]

Challenge II: Handle the mismatch between stochastic gradient and momentum. In the analysis above, we assume β 1 = 0 . Additional challenges arise when we move to the case where β 1 = 0 . Specifically, following the same routine, the "First Order Main" term now becomes E [ 〈 G t , -η 1 √ ˜ ν t glyph[circledot] m t ] . It is hard to even estimate whether such a term is negative or not, given that m t and ˜ ν t still has entangled stochasticity, and the conditional expectation of m t also differs from G t , both due to the existence of historical gradient.

Solution to Challenge II. Inspired by the state-of-art analysis of SGDM [22], which leverage the potential function f ( v t ) with v t = w t -β w t -1 1 -β , we propose to use the potential function f ( u t ) with

We again focus on the "First Order" term, which can be written as

Here approximate equation ( ∗ ) is due to Assumption 1 and that w t is close to u t , and approximate equation ( ◦ ) is due to Lemma 1 and ˜ ν t = β 2 ν t -1 + (1 -β 2 ) σ 2 0 ≈ β 2 ν t -1 (of course, these are informal statements. Please refer to Appendix C for the detailed proof). With the above methodology, we arrive at the following lemma.

Lemma 2 (Informal Version of Lemma 8) . Let all conditions in Theorem 1 holds. Then,

Summing the above lemma over t from 1 to T , we obtain

We then encounter the second challenge.

Challenge III: Convert Eq. (6) to a bound of gradient norm. Although we have derived a regret bound, i.e., a bound of ∑ T t =1 E [ ‖ 1 4 √ ˜ ν t glyph[circledot] G t ‖ 2 ] , we need to convert it into a bound of E [ ‖ G t ‖ 2 ] . In existing works [36, 10, 16] which assumes bounded gradient, such a conversion is straightforward because (their version of) ˜ ν t is upper bounded. However, we do not assume bounded gradient and ˜ ν t can be aribitrarily large, making E [ ‖ 1 4 √ ˜ ν t glyph[circledot] G t ‖ 2 ] arbitrarily small than E [ ‖ G t ‖ 2 ] .

Solution to Challenge III. As this part involves coordinate-wise analysis, we define g t,i , G t,i , ν t,i , and ˜ ν 1 t,i respectively as the l -th coordinate of g t , G t , ν t , and ˜ ν 1 t . To begin with, note that due to Cauchy's inequality and Hölder's inequality,

Therefore, we only need to derive an upper bound of ∑ T t =1 E [ ‖ 4 ˜ ν t ‖ 2 ] , which is achieved by the following divide-and-conque methodology. Firstly, when | G t,i | ≥ σ 0 σ 1 , we can show 2 E |F t | g t,i | 2 ≥ 2 | G t,i | 2 ≥ E |F t | g t,i | 2 . Then, through a direct calculation, we obtain that

and thus

Secondly, when | G t,i | < σ 0 σ 1 , define { ¯ ν t,i } ∞ t =0 as ¯ ν 0 ,l = ν 0 ,l , ¯ ν t,i = ¯ ν t -1 ,i + | g t,i | 2 1 | G t,i | < σ 0 σ 1 . One can easily observe that ¯ ν t,i ≤ ν t,i , and thus

Putting the above two estimations together, we derive that

The above methodology can be summarized as the following lemma.

Lemma 3. Let all conditions in Theorem 1 hold. Then,

Based on Lemma 3, we can derive the estimation of ∑ T t =1 E [ ‖ 4 √ ˜ ν t ‖ 2 ] since ˜ ν t is close to ν t . The proof is then completed by combining the estimation of ∑ T t =1 E [ ‖ 4 √ ˜ ν t ‖ 2 ] (Eq. (6)) and Eq. (7).

5 Gap-closing upper bound on the iteration complexity of Adam

In this section, based on a refined proof of Stage II of Theorem 1 (see Appendix C) under the specific case η = Θ(1 / √ T ) and β 2 = 1 -Θ(1 /T ) , we show that the logarithmic factor in Theorem 1 can be removed and the lower bound can be achieved. Specifically, we have the following theorem.

Theorem 2. Let Assumption 1 and Assumption 2 hold. Then, select the hyperparameters of Adam as η = a √ T , β 2 = 1 -b T and β 1 = c √ β 2 , where a, b > 0 and 0 ≤ c < 1 are independent of T . Then, let w τ be the output of Adam in Algorithm 1, and we have

where

As a result, let A be Adam in Algorithm 1, we have C ε ( A , ∆ , L, σ 2 ) = O ( 1 ε 4 ) .

The proof of Theorem 2 is based on a refined solution of Challenge II in the proof of Theorem 1 under the specific hyperparameter settings, and we defer the concrete proof to Appendix D. Below we discuss on Theorem 2, comparing it with practice, with Theorem 1 and existing convergence rate of Adam, and with the convergence rate of AdaGrad.

Alignment with the practical hyperparameter choice. The hyperparameter setting in Theorem 2 indicates that to achieve the lower bound of iteration complexity, we need to select small η and close-to1 β 2 , with less requirement over β 1 . This agrees with the hyperparameter setting in deep learning libaries, for example, η = 10 -3 , β 2 = 0 . 999 , and β 1 = 0 . 9 in PyTorch.

Comparison with Theorem 1 and existing works. To our best knowledge, Theorem 2 is the first to derive the iteration complexity O ( 1 ε 4 ) . Previously, the state-of-art iteration complexity is O ( log 1 /ε ε 4 ) [10] where they additionally assume bounded gradient. Theorem 2 is also tighter than Theorem 1 (while Theorem 1 holds for more general hyperparameter settings). As discussed in Section 4.1, if applying the hyperparameter setting in Theorem 2 (i.e., η = a √ T , β 2 = 1 -b T and β 1 = c √ β 2 ) to Theorem 1, we will obtain that E ‖∇ f ( w τ ) ‖ ≤ O (poly(log T ) / 4 √ T ) and C ε ( A , ∆ , L, σ 2 ) = O ( log 1 /ε ε 4 ) , worse than the upper bound in Theorem 2 and the lower bound in Proposition 1 by a logarithmic factor.

Comparison with AdaGrad. AdaGrad [12] is another popular adaptive optimizer. Under Assumptions 1 and 2, the state-of-art iteration complexity of AdaGrad is O ( log 1 /ε ε 4 ) [13], which is worse than Adam by a logarithmic factor. Here we show that such a gap may be not due to the limitation of analysis, and can be explained by analogizing AdaGrad to Adam without momentum as SGD with diminishing learning rate to SGD with constant learning rate. To start with, the update rule of AdaGrad is given as

We first show that in Algorithm 1, if we allow the hyperparameters to be dynamical, i.e.,

then Adam is equivalent to AdaGrad by setting η t = η √ t , β 1 ,t = 0 , and β 2 ,t = 1 -1 t . Specifically, by setting µ t = t ν t in Eq. (9), we have Eq. (9) is equivalent to with Eq. (8) (by replacing ν t by µ t in Eq. (8)). Comparing the above hyperparameter setting with that in Theorem 2, we see that the above hyperparameter setting can be obtained by changing T to t and setting c = 0 in Theorem 2. This is similar to the relationship between SGD with diminishing learning rate Θ(1 / √ t ) and SGD with diminishing learning rate Θ(1 / √ T ) . Recall that the iteration complexity of SGD with diminishing learning rate Θ(1 / √ t ) also has an additional logarithmic factor than SGD with constant learning rate, which may explain the gap between AdaGrad and Adam.

6 Limitations

Despite that our work provides the first result closing the upper bound and lower bound of the iteration complexity of Adam, there are several limitations listed as follows:

Dependence over the dimension d . The bounds in Theorem 1 and Theorem 2 is monotonously increasing with respect to d . This is undesired since the upper bound of iteration complexity of SGD is invariant with respect to d . Nevertheless, removing such an dependence over d is technically hard since we need to deal with every coordinate separately due to coodinate-wise learning rate, while the descent lemma does not hold for a single coordinate but combines all coordinates together. To our best knowledge, all existing works on the convergene of Adam also suffers from the same problem. We leave removing the dependence over d as an important future work.

No better result with momentum. It can be observed that in Theorem 1 and Theorem 2, the tightest bound is achieved when β 1 = 0 (i.e., no momentum is applied). This contradicts with the common wisdom that momentum helps to accelerate. Although the benefit of momentum is not very clear even for simple optimizer SGD with momentum, we view this as a limitation of our work and defer proving the benefit of momentum in Adam as a future work. Also, our result does not imply that setting β 1 is not as critical as setting β 2 . The primary objective of this paper is to characterize the dependence on ε , and the importance of setting β 1 might be justified in other ways or characterizations. To help readers gain a deeper understanding of this issue, we include experiments to illustrate the dependence of performance on β 1 in Appendix E.

Acknowledgments and Disclosure of Funding

This work was founded by the CAS Project for Young Scientists in Basic Research under Grant No. YSBR-034 and the Innovation Funding of ICT, CAS under Grant No.E000000.

References

Y. Arjevani, Y. Carmon, J. C. Duchi, D. J. Foster, N. Srebro, and B. Woodworth. Lower bounds for non-convex stochastic optimization. Mathematical Programming , pages 1-50, 2022.

J. Bradbury, R. Frostig, P. Hawkins, M. J. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, J. VanderPlas, S. Wanderman-Milne, and Q. Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http://github.com/google/jax .

A. Brock, J. Donahue, and K. Simonyan. Large scale gan training for high fidelity natural image synthesis. In International Conference on Learning Representations , 2018.

T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al. Language models are few-shot learners. Advances in neural information processing systems , 33:1877-1901, 2020.

X. Chen, S. Liu, R. Sun, and M. Hong. On the convergence of a class of Adam-type algorithms for non-convex optimization. arXiv preprint arXiv:1808.02941 , 2018.

X. Chen, C. Liang, D. Huang, E. Real, K. Wang, Y. Liu, H. Pham, X. Dong, T. Luong, C.-J. Hsieh, et al. Symbolic discovery of optimization algorithms. arXiv preprint arXiv:2302.06675 , 2023.

M. Crawshaw, M. Liu, F. Orabona, W. Zhang, and Z. Zhuang. Robustness to unbounded smoothness of generalized signSGD. arXiv preprint arXiv:2208.11195 , 2022.

A. Cutkosky and H. Mehta. Momentum improves normalized SGD. In International conference on machine learning , pages 2260-2268. PMLR, 2020.

S. De, A. Mukherjee, and E. Ullah. Convergence guarantees for RMSProp and ADAM in non-convex optimization and an empirical comparison to Nesterov acceleration. arXiv preprint arXiv:1807.06766 , 2018.

A. Défossez, L. Bottou, F. Bach, and N. Usunier. A simple convergence proof of Adam and Adagrad. Transactions on Machine Learning Research , 2022.

A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations , 2020.

J. Duchi, E. Hazan, and Y. Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research , 12(7), 2011.

M. Faw, I. Tziotis, C. Caramanis, A. Mokhtari, S. Shakkottai, and R. Ward. The power of adaptivity in SGD: Self-tuning step sizes with unbounded gradients and affine variance. In Conference on Learning Theory , pages 313-355. PMLR, 2022.

W. A. Fuller. Measurement error models . John Wiley & Sons, 2009.

S. Ghadimi and G. Lan. Stochastic first-and zeroth-order methods for nonconvex stochastic programming. SIAM Journal on Optimization , 23(4):2341-2368, 2013.

Z. Guo, Y. Xu, W. Yin, R. Jin, and T. Yang. A novel convergence analysis for algorithms of the Adam family. arXiv preprint arXiv:2112.03459 , 2021.

K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 770-778, 2016.

J. D. M.-W. C. Kenton and L. K. Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of NAACL-HLT , pages 4171-4186, 2019.

F. Khani and P. Liang. Feature noise induces loss discrepancy across groups. In International Conference on Machine Learning , pages 5209-5219. PMLR, 2020.

D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

H. Li, A. Jadbabaie, and A. Rakhlin. Convergence of Adam under relaxed assumptions. arXiv preprint arXiv:2304.13972 , 2023.

Y. Liu, Y . Gao, and W. Yin. An improved analysis of stochastic gradient descent with momentum. Advances in Neural Information Processing Systems , 33:18261-18271, 2020.

L. Luo, Y. Xiong, Y. Liu, and X. Sun. Adaptive gradient methods with dynamic bound of learning rate. arXiv preprint arXiv:1902.09843 , 2019.

S. Merity, C. Xiong, J. Bradbury, and R. Socher. Pointer sentinel mixture models. arXiv preprint arXiv:1609.07843 , 2016.

A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. volume 32, 2019.

S. J. Reddi, S. Kale, and S. Kumar. On the convergence of Adam and beyond. arXiv preprint arXiv:1904.09237 , 2019.

N. Shi, D. Li, M. Hong, and R. Sun. RMSprop converges with proper hyper-parameter. In International Conference on Learning Representations , 2021.

K. Simonyan and A. Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv:1409.1556 , 2014.

A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

B. Wang, Y. Zhang, H. Zhang, Q. Meng, Z.-M. Ma, T.-Y. Liu, and W. Chen. Provable adaptivity in Adam. arXiv preprint arXiv:2208.09900 , 2022.

B. Wang, H. Zhang, Z. Ma, and W. Chen. Convergence of adagrad for non-convex objectives: Simple proofs and relaxed assumptions. In The Thirty Sixth Annual Conference on Learning Theory , pages 161-190. PMLR, 2023.

X. Xie, P. Zhou, H. Li, Z. Lin, and S. Yan. Adan: Adaptive nesterov momentum algorithm for faster optimizing deep models. arXiv preprint arXiv:2208.06677 , 2022.

M. Zaheer, S. Reddi, D. Sachan, S. Kale, and S. Kumar. Adaptive methods for nonconvex optimization. Advances in neural information processing systems , 31, 2018.

Y. Zhang, C. Chen, N. Shi, R. Sun, and Z.-Q. Luo. Adam can converge without any modification on update rules. arXiv preprint arXiv:2208.09632 , 2022.

Z. Zhou, Q. Zhang, G. Lu, H. Wang, W. Zhang, and Y. Yu. Adashift: Decorrelation and convergence of adaptive learning rate methods. arXiv preprint arXiv:1810.00143 , 2018.

F. Zou, L. Shen, Z. Jie, W. Zhang, and W. Liu. A sufficient condition for convergences of Adam and RMSProp. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 11127-11135, 2019.

A Other Related works

Section 3 has provided a detailed discussion over existing convergence analysis of Adam. In this section, we briefly review other related works. Adam is proposed with a convergence analysis in online optimization [20]. The proof, however, is latter shown to be flawed in Reddi et al. [26] as it requires the adaptive learning rate of Adam to be non-increasing. This motivates a line of works modifying Adam to ensure convergence. The modifications include enforcing the adaptive learning rate to be non-increasing [26, 5], imposing upper bound and lower bound of the adaptive learning rate [23], and using different approach to estimate second-order momentum [35, 7]. Recently, Chen et al. [6] discover a new optimizer Lion through Symbolic Discovery, which uses sign operation to replace the adaptive learning rate in Adam, achieving comparable performance of Adam with less memory costs.

B Auxilliary Lemmas

The following two lemmas are useful when bounding the second-order term.

Lemma 4. Assume we have 0 < β 2 < 1 and a sequence of real numbers ( a n ) ∞ n =1 . Let b 0 > 0 and b n = β 2 b n -1 +(1 -β 2 ) a 2 n . Then, we have

Lemma 5. Assume we have 0 < β 2 1 < β 2 < 1 and a sequence of real numbers ( a n ) ∞ n =1 . Let b 0 > 0 , b n = β 2 b n -1 +(1 -β 2 ) a 2 n , c 0 = 0 , and c n = β 1 c n -1 +(1 -β 1 ) a n . Then, we have

Proof. To begin with,

Applying Cauchy's inequality, we obtain

Summing the above inequality over n from 1 to T then leads to

The proof is completed.

The following lemma bound the update norm of Adam.

Proof. We have that

Here the second inequality is due to Cauchy's inequality. The proof is completed.

C Proof of Theorem 1

This section collects the proof of Theorem 1. As a part of the proof, we first provide formal descriptions of Lemma 1, Lemma 2, and Lemma 3, and their corresponding proofs. We then proceed to prove Theorem 1 leveraging these lemmas.

C.1 Formal description of Lemma 1, Lemma 2, and Lemma 3 and their proof

Lemma 7 (Formal version of Lemma 1) . Let all conditions in Theorem 1 hold. Then, we have

Proof. To start with,

As for I.1.1, we have

where inequality ( ∗ ) uses Lemma 6, inequality ( ◦ ) is due to Holder's inequality, and inequality ( · ) is due to Assumption 3. Applying mean-value inequality respectively to ∑ d i =1 η (1 -β 1 ) √ 1 -β 2 (√ 1 -β 1 √ β 2 ) 3 E |F t | G t,i | √ ˜ ν t,i σ 0 √ E |F t g 2 t,i ( √ ν t,i + √ ˜ ν t,i ) 2 and ∑ d η (1 -β 1 ) √ 1 -β 2 |F t | G t,i | √ |F t g 2

(√

β

√

1

β

i

=1

1

-

˜

ν

t,i

2

the above inequality can be bounded by

Here the inequality is due to ˜ ν t,i = (1 -β 2 ) σ 2 0 + β 2 ν t -1 ,i ≥ (1 -β 2 ) σ 2 0 . Meanwhile, we have

)

(

ν

t,i t,i

+

√

˜

ν

t,i

)

2

E

σ

1

|

G

t,i

|

E

√

, we obtain that the right-hand-side of

3

√

Applying the above inequality back to Eq. (10), we obtain that

Furthermore, due to Assumption 1, we have (we define G 0 glyph[defines] G 1 )

which further leads to

Applying the above inequality back to Eq. (11) leads to that

As for I.1.2, we have

With Inequalities (12) and (13), we conclude that

Lemma 8 (Formal version of Lemma 2) . Let all conditions in Theorem 1 holds. Then,

Proof. According to the definition of u t , we have

where Eq. ( ∗ ) is due to m t = β 1 m t -1 +(1 -β 1 ) g t .

Applying the above equation to the "First Order" term, we find that it can be decomposed as

Here we apply Lemma 7 to bound I.1. We proceed by bounding I.2 and I.3 respectively. As for I.2, we have

As for I.3, we directly apply Assumption 1 and obtain

All in all, we summarize that the "First Order" term can be bounded by

Furthermore, the "Second Order" term can be directly bounded by

Applying the estimations of the first-order term and the second-order term to the descent lemma then gives

The proof is completed.

Lemma 9 (Lemma 3, restated) . Let all conditions in Theorem 1 hold. Then,

Proof of Lemma 3. To begin with, we have that

On the other hand, we have that

As a conclusion,

On the other hand, as stated in Section 4.2, we define { ¯ ν t,i } ∞ t =0 as ¯ ν 0 ,i = ν 0 ,i , ¯ ν t,i = β 2 ¯ ν t -1 ,i + (1 -β 2 ) | g t,i | 2 1 | G t,i | < σ 2 0 σ 2 1 . One can easily observe that ¯ ν t,i ≤ ν t,i , and thus

All in all, summing the above two inequalities together, we obtain that

Since ∀ t ,

combining with √ β 2 ˜ ν 1 ,i = √ β 2 ( β 2 ¯ ν 0 ,i +(1 -β 2 ) σ 2 0 ) and E √ ˜ ν t +1 ,i = E √ β 2 ν t,i +(1 -β 2 ) σ 2 0 ≥ E √ β 2 ¯ ν t,i +(1 -β 2 ) σ 2 0 , we obtain

Leveraging Eq. (16), we then obtain that

where in the last inequality we use the concavity of h ( x ) = ln x . Solving the above inequality with respect to ∑ T +1 t =1 ∑ d i =1 E √ ˜ ν t,i then gives

The proof is then completed by applying the definition of C 2 .

C.2 Proof of Theorem 1

Proof of Theorem 1. Summing the inequality in Lemma 8 over t from 1 to T and collecting the terms, we obtain

where we define

to simplify the notations, inequality ( ∗ ) is due to that ˜ ν t +1 ,i ≥ β 2 ˜ ν t +1 ,i and β 1 ≤ √ β 2 -8 σ 2 1 (1 -β 2 ) β -2 2 , and inequality ( ◦ ) is due to Lemma 5. Simple rearrangement of the above inequality then gives

Then, according to Cauchy's inequality, we have

Meanwhile, by Lemma 3, we have

Combining the above inequality and Eq. (17) gives

The proof is then completed.

D Proof of Theorem 2

Proof. To start with, we have that

where the last inequality is due to that

Furthermore, we have

Conclusively, we obtain

Secondly, as β 2 → 1 as T →∞ , β 1 ≤ √ β 2 -8 σ 2 1 (1 -β 2 ) β -2 2 holds for large enough T , and thus Theorem 1 holds. Applying the value of β 1 , β 2 , and η to Eq. (16), we obtain that

Summing Eq. (19) with respect to i then gives

where the second inequality is due to Eq. (20), the second-to-last inequality is due to Eq. (18), and the last inequality is due to Jensen's inequality. Solving the above ineqaulity with respect to

Therefore, by Cauchy's inequality, we have

Since

The proof is completed.

E Experiments

As mentioned in Section 6, one of the limitations of our theory is that it can not provide better results when momentum is present. To complement such a limitation, we initialize a empirical study of the effect of momentum in Adam as follows.

Experimental setting. We use Adam training on Cifar 10 with ResNet18 [28] and VGG13 [17] and wikitext2 with two-layer Transformer [29] for 50 epoch and record its training loss at 50 epoch

as a measure for the optimization speed. Smaller loss indicates better optimization. The batch size is set 1024 for Cifar10 dataset and 100 for WikiText2 Dataset.

The results are given in Table 1. Our discoveries are:

Momentum can benefit the optimization when the β is not too large.

For all datasets, larger β 1 (Setting β 1 close to 1) will worse the optimization.

Connection with Theorem 1. One can easily observe that in Theorem 1 both C 1 and C 2 polynomially depend on 1 1 -β 1 and thus so does E ∑ T t =1 ‖∇ f ( w t ) ‖ = ˜ O ( 1 1 -β 1 ) . Therefore, Theorem 1 is aligned with the experimental results in the sense that both our theory and the experimental results indicate that the Adam will become worse when the β 1 is close to 1. However, Theorem 1 cannot explain the benefit of using momentum (by setting β 1 larger than 0). This may be due to that our Theorem 1 is a worst-case analysis. We conjecture that theoretically deriving the benefit of momentum requires restricting the underlying objective function to a more specific range, which we leave as a future work.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a convergence analysis of Adam under only the smoothness and bounded variance conditions.\n\nSTRENGTHS:\n- The strength of the paper is to present a convergence analysis of Adam under only the smoothness and bounded variance conditions, in contrast to the existing analyses (Section 3). We are interested in the analysis of Adam under only the two conditions, since the conditions are more natural and realistic than the boundedness of the gradient norm of the objective function and the Lipschitz continuity of the stochastic gradient.  \n-  The abstract indicates that \"Especially with properly chosen hyperparameters, we derive an upper bound of iteration complexity of Adam and show that it meets the lower bound for first-order optimizers.\" Proposition 1 shows that the iteration complexity $\\mathcal{C}$ of a first-order optimizer is $\\Omega (1/\\epsilon^4)$, that is, there exist $c_1, c_2 > 0$ such that $c_1/\\epsilon^4 \\leq \\mathcal{C} \\leq c_2/\\epsilon^4$. Theorem 2 implies that Adam satisfies that $\\mathcal{C} = O(1/\\epsilon^4)$. Hence, the paper concludes the claim of the abstract.\n\nWEAKNESSES:\nI understand the motivation of the paper. Meanwhile, many theoretical and practical results of Adam have been presented. Hence, unfortunately, I do not find that the results in the paper are surprised as compared with the existing ones. Please see Questions.\n\nQUESTIONS:\nI have the following concerns. If the authors could address the concerns, then I think that the paper has the contribution and novelty.\n- Theorem 1/Theorem 2 implies the upper bounds of the gradient norm of Adam using a constant step-size $\\eta$/a diminishing step-size $\\eta = a/\\sqrt{T}$. The theorems lead to the finding that, e.g., (Lines 159--162) \"Theorem 1 holds for general choices of hyperparameters since the only condition posed on hyperparameters is $\\beta_1 < \\beta_2$. Such condition covers a wide range of hyperparameters, e.g., the default setting $\\beta_1 = 0.9$ and $\\beta_2 = 0.999$ in PyTorch [19].\" I understand that using $\\beta_1$ and $\\beta_2$ close to 1 is useful to implement Adam. However, I cannot understand that Theorem 1 implies the claims in (Lines 159--162). Does Theorem 1 imply that $\\beta_1, \\beta_2 \\approx 1 \\Rightarrow$ the right-hand side of (2) is small? Can the authors show that the right-hand side of (2) is a decreasing function with respect to $\\beta_1$ and $\\beta_2$?\n- The above discussion holds for Theorem 2 and Lines 250--253. Does Theorem 2 imply that $a, b \\approx 0 \\Rightarrow$ the upper bound of the gradient norm is small? Can the authors show that it is decreasing when $a$ and $b$ are small?\n- The above two concerns mean that there is a gap between theory (Theorems 1 and 2) and practice ($\\beta_1 = 0.9$ and $\\beta_2 = 0.999$). Can the authors bridge the gap?\n- Can the authors provide lower bounds of the gradient norm of Adam, such as $\\mathbb{E}\\sum_{i=1}^T \\Vert \\nabla f(w_t) \\Vert \\geq$ $C \\times$  (the right-hand side of (2))?\n- Paper [25] and its references are sufficient for both theory and practice.  Can the authors show more novelty and contributions from Theorems 1 and 2?\n- Related to the above concern, I believe that practical results (e.g., numerical results and new setting of hyperparameters) are provided based on the novelty and contributions of the paper. Can the authors provide such practical results? For example, it would be nicer to compare numerically Adam using $\\beta_1 = 0.9$ and $\\beta_2 = 0.999$ with Adam using the new setting of $\\beta_1$ and $\\beta_2$ based on Theorems 1 and 2.\n\nLIMITATIONS:\nThere is no potential negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-24"}

comments: SUMMARY:
This paper analyzes the iteration complexity of Adam. It Is first pointed out that upper bounds in prior work do not match existing lower bound; the reason is that the lower bound is proved under smoothness and bounded noise variance, while prior upper bounds make more assumptions. This paper then proves a general upper bound (Theorem 1) which only requires smoothness and bounded noise variance and matches the lower bound up to a logarithmic factor. Later in Theorem 2, a refined analysis is given which further removes the logarithmic factor, and thus giving an upper bound that matches the lower bound exactly.

STRENGTHS:
This paper analyzes the iteration complexity of Adam, which is a very important problem given Adam's popularity. Moreover, it is pointed out that existing upper bounds do not match lower bound, and this paper closes this gap, which is a nice contribution. The proof techniques may be of independent interest, such as the peeling-off strategy to handle the dependency between the momentum and the adaptive learning rate.

WEAKNESSES:
N/A

QUESTIONS:
Can SGD match the lower bound? Alternatively, can we show that Adam converges faster than SGD, which is usually observed in practice?

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-05

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nAdam is one of the most popular stochastic optimization algorithms especially in deep learning, the existing convergence theories do not achieve a tight upper bound that meets the lower bound. Moreover, many of them require additional assumptions, such as bounded gradient. This paper shows that Adam can achieve the tight bound of $O ( \\epsilon^{-4} )$ without such additional assumptions. This result closes the gap between the practical success of Adam and the theoretical sub-optimality.\n\nSTRENGTHS:\n- The theoretical result is strong. Although there are many existing works that analyze the convergence of Adam, this paper is the first one that proves that Adam can converge with $\\mathcal{O} ( 1 / \\sqrt{T} )$, or equivalently $\\mathcal{O} ( \\epsilon^{-4} )$.\n- Technical improvements for deriving the result are interesting and clearly explained.\n- This paper is easy to read, and the presentation is also clear at least to the experts on the convergence theory of stochastic optimization.\n\nWEAKNESSES:\n**Relation to [1] is not clear**\n\nAlthough I am basically positive about the result of this paper, I am not sure about how it relates to the well-known result about the non-convergence behavior of Adam by [1]. They showed that there always exits a problem in which Adam fails to converge, but the authors' result seems that Adam can always converge with the optimal rate (i.e., $\\mathcal{O} ( \\epsilon^{-4} )$). In my current understanding, this is because Theorem 2 requires $1 - \\beta_2 = \\Theta ( 1 / T )$, which means that, in order to ensure the convergence of Adam, we need to choose the hyper-parameter $\\beta_2$ depending on the total number of parameter updates $T$. I think it would be better to clarify the relationship clearly in Section 7. Though the authors mention [1] in the section, the relation to their theoretical result is not clear to me. When it becomes clear in the rebuttal period, I will raise my score.\n\n**There are no experiments**\n\n- I think the condition of $1 - \\beta_2 = \\Theta ( 1 / T )$ is crucial to achieve the optimal convergence rate, so it would be better to demonstrate it experimentally (I think a toy experiment is enough).\n\n**Minor comments**\n\n- The notations of the output of Adam are inconsistent in Algorithm 1 ($\\boldsymbol{w}_r$) and Theorem 2 ($\\boldsymbol{w}_\\tau$), which is a little confusing.\n- Finishing a paper with Related work is not common in my opinion.\n- It would be better to add equation numbers to all the equations for the ease of communication between the reviewers and the authors.\n\n**Typos**\n\n- line 20: uderstand -> understand\n- line 198: $\\boldsymbol{G_{t-2} \\rightarrow G_{t-1}}$\n- line 248: pratice -> practice\n\n**References**\n\n[1] Reddi, Sashank J., Satyen Kale, and Sanjiv Kumar. \"On the Convergence of Adam and Beyond.\" International Conference on Learning Representations. 2018.\n\nQUESTIONS:\n**Relation to [1]**\n\nAs I mentioned in the weaknesses section, I would like to know the relation between the results of this paper and [1].\n\n**What happens when using diminishing step size?**\n\nIn the analysis on SGD, diminishing step size (e.g., $\\eta_t = \\Theta ( 1 / \\sqrt{t} )$ is often used, because it is close to practical situations. Is it easy to extend this result to the case of diminishing step size?\n\nLIMITATIONS:\nThe limitations are mentioned in Section 6, which is clear to me.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

comments: SUMMARY:
This paper gives a new analysis of the Adam algorithm intended to close the gap between the upper bound of Adam's iteration complexity and the existing lower bound for first-order nonconvex optimization. The authors show that existing analysis of Adam either uses the bounded gradient assumption, achieves a suboptimal iteration complexity, or relies on the mean-squared smoothness assumption. Afterwards, they give a novel analysis of the algorithm that meets the lower bound in dependence on the desired accuracy $\epsilon$.

STRENGTHS:
- The new analysis of Adam is conducted under the same assumptions as standard SGD, unlike most prior work.
- The new analysis of Adam achieves the optimal $\frac{1}{\epsilon^{4}}$ complexity with no additional log factors.
- The authors introduce several technical tools that can be helpful in the analysis of adaptive algorithms more generally, for example the stochastic surrogates used are new.

WEAKNESSES:
- My main problem with this paper is that the proof, as it is, is very complicated to check. It would be very helpful if the authors included a section in the deterministic case (no stochasticity) with their full proof in this simplified setting.
- The bounded variance condition is a bit restrictive (see [1]), can the convergence of Adam be derived under any of the more general conditions mentioned in [1]?
- (Minor typos) line 134 should be log 1/\epsilon not log \epsilon.  Line 198 $G_{t-1}$ not $G_{t-2}$. Please use \left and \right for braces in line 229.

[1] Ahmed Khaled & Peter Richtárik, Better Theory for SGD in the Nonconvex World, TMLR 2023

QUESTIONS:
1. Please address my concern in the weaknesses section.

LIMITATIONS:
N/A.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-10