Title: Convergence of Adam Under Relaxed Assumptions
Venue: NeurIPS 2023
Accepted: yes
Keywords: Non-convex optimization

Adam

Convergence

Variance reduction
Abstract: In this paper, we provide a rigorous proof of convergence of the Adaptive Moment Estimate (Adam) algorithm for a wide class of optimization objectives. Despite the popularity and efficiency of the Adam algorithm in training deep neural networks, its theoretical properties are not yet fully understood, and existing convergence proofs require unrealistically strong assumptions, such as globally bounded gradients, to show the convergence to stationary points. In this paper, we show that Adam provably converges to $\epsilon$-stationary points with $\mathcal{O}(\epsilon^{-4})$ gradient complexity under far more realistic conditions. The key to our analysis is a new proof of boundedness of gradients along the optimization trajectory of Adam, under a generalized smoothness assumption according to which the local smoothness (i.e., Hessian norm when it exists) is bounded by a sub-quadratic function of the gradient norm. Moreover, we propose a variance-reduced version of Adam with an accelerated gradient complexity of $\mathcal{O}(\epsilon^{-3})$.

Convergence of Adam Under Relaxed Assumptions

Abstract

In this paper, we provide a rigorous proof of convergence of the Adaptive Moment Estimate (Adam) algorithm for a wide class of optimization objectives. Despite the popularity and efficiency of the Adam algorithm in training deep neural networks, its theoretical properties are not yet fully understood, and existing convergence proofs require unrealistically strong assumptions, such as globally bounded gradients, to show the convergence to stationary points. In this paper, we show that Adam provably converges to ϵ -stationary points with O ( ϵ -4 ) gradient complexity under far more realistic conditions. The key to our analysis is a new proof of boundedness of gradients along the optimization trajectory of Adam, under a generalized smoothness assumption according to which the local smoothness (i.e., Hessian norm when it exists) is bounded by a sub-quadratic function of the gradient norm. Moreover, we propose a variance-reduced version of Adam with an accelerated gradient complexity of O ( ϵ -3 ) .

1 Introduction

In this paper, we study the non-convex unconstrained stochastic optimization problem

The Adaptive Moment Estimation (Adam) algorithm [23] has become one of the most popular optimizers for solving (1) when f is the loss for training deep neural networks. Owing to its efficiency and robustness to hyper-parameters, it is widely applied or even sometimes the default choice in many machine learning application domains such as natural language processing [44; 4; 13], generative adversarial networks [35; 21; 55], computer vision [14], and reinforcement learning [28; 33; 40]. It is also well known that Adam significantly outperforms stochastic gradient descent (SGD) for certain models like transformer [50; 24; 1].

Despite its success in practice, theoretical analyses of Adam are still limited. The original proof of convergence in [23] was later shown by [37] to contain gaps. The authors in [37] also showed that for a range of momentum parameters chosen independently with the problem instance , Adam does not necessarily converge even for convex objectives. However, in deep learning practice, the hyper-parameters are in fact problem-dependent as they are usually tuned after given the problem and weight initialization. Recently, there have been many works proving the convergence of Adam for non-convex functions with various assumptions and problem-dependent hyper-parameter choices. However, these results leave significant room for improvement. For example, [12; 19] prove the convergence to stationary points assuming the gradients are bounded by a constant, either explicitly or implicitly. On the other hand, [51; 45] consider weak assumptions, but their convergence results are still limited. See Section 2 for more detailed discussions of related works.

To address the above-mentioned gap between theory and practice, we provide a new convergence analysis of Adam without assuming bounded gradients , or equivalently, Lipschitzness of the objective function. In addition, we also relax the standard global smoothness assumption, i.e., the Lipschitzness of the gradient function, as it is far from being satisfied in deep neural network training. Instead, we

consider a more general, relaxed, and non-uniform smoothness condition according to which the local smoothness (i.e., Hessian norm when it exists) around x is bounded by a sub-quadratic function of the gradient norm ∥∇ f ( x ) ∥ (see Definition 3.2 and Assumption 2 for the details). This generalizes the ( L 0 , L 1 ) smoothness condition proposed by [49] based on language model experiments. Even though our assumptions are much weaker and more realistic, we can still obtain the same O ( ϵ -4 ) gradient complexity for convergence to an ϵ -stationary point.

The key to our analysis is a new technique to obtain a high probability, constant upper bound on the gradients along the optimization trajectory of Adam, without assuming Lipschitzness of the objective function. In other words, it essentially turns the bounded gradient assumption into a result that can be directly proven. Bounded gradients imply bounded stepsize at each step, with which the analysis of Adam essentially reduces to the simpler analysis of AdaBound [31]. Furthermore, once the gradient boundedness is achieved, the analysis under the generalized non-uniform smoothness assumption is not much harder than that under the standard smoothness condition. We will introduce the technique in more details in Section 5. We note that the idea of bounding gradient norm along the trajectory of the optimization algorithm can be use in other problems as well. For more details, we refer the reader to our concurrent work [26] in which we present a set of new techiniques and methods for bounding gradient norm for other optimization algorithms under a generalized smoothness condition.

Another contribution of this paper is to show that the gradient complexity of Adam can be further improved with variance reduction methods. To this end, we propose a variance-reduced version of Adam by modifying its momentum update rule, inspired by the idea of the STORM algorithm [9]. Under additional generalized smoothness assumption of the component function f ( · , ξ ) for each ξ , we show that this provably accelerates the convergence with a gradient complexity of O ( ϵ -3 ) . This rate improves upon the existing result of [47] where the authors obtain an asymptotic convergence of their approach to variance reduction for Adam in the non-convex setting, under the bounded gradient assumption.

1.1 Contributions

In light of the above background, we summarize our main contributions as follows.

We develop a new analysis to show that Adam converges to stationary points under relaxed assumptions. In particular, we do not assume bounded gradients or Lipschitzness of the objective function. Furthermore, we also consider a generalized non-uniform smoothness condition where the local smoothness or Hessian norm is bounded by a sub-quadratic function of the gradient norm. Under these more realistic assumptions, we obtain a dimension free gradient complexity of O ( ϵ -4 ) if the gradient noise is centered and bounded.

We generalize our analysis to the setting where the gradient noise is centered and has sub-Gaussian norm, and show the convergence of Adam with a gradient complexity of O ( ϵ -4 log 3 . 25 (1 /ϵ )) .

We propose a variance-reduced version of Adam (VRAdam) with provable convergence guarantees. In particular, we obtain the accelerated O ( ϵ -3 ) gradient complexity.

2 Related work

In this section, we discuss the relevant literature related to convergence of Adam and the generalized smoothness condition, and defer additional related work on variants of Adam and variance reduction methods to Appendix A.

Convergence of Adam. Adam was first proposed by Kingma and Ba [23] with a theoretical convergence guarantee for convex functions. However, Reddi et al. [37] found a gap in the proof of this convergence analysis, and also constructed counter-examples for a range of hyper-parameters on which Adam does not converge. That being said, the counter-examples depend on the hyperparameters of Adam, i.e., they are constructed after picking the hyper-parameters. Therefore, it does not rule out the possibility of obtaining convergence guarantees for problem-dependent hyperparameters, as also pointed out by [42; 51].

Many recent works have developed convergence analyses of Adam with various assumptions and hyper-parameter choices. Zhou et al. [54] show Adam with certain hyper-parameters can work on the counter-examples of [37]. De et al. [10] prove convergence for general non-convex functions

assuming gradients are bounded and the signs of stochastic gradients are the same along the trajectory. The analysis in [12] also relies on the bounded gradient assumption. Guo et al. [19] assume the adaptive stepsize is upper and lower bounded by two constants, which is not necessarily satisfied unless assuming bounded gradients or considering the AdaBound variant [31]. [51; 45] consider very weak assumptions. However, they show either 1) 'convergence' only to some neighborhood of stationary points with a constant radius, unless assuming the strong growth condition; or 2) convergence to stationary points but with a slower rate.

Generalized smoothness condition. Generalizing the standard smoothness condition in a variety of settings has been a focus of many recent papers. Recently, [49] proposed a generalized smoothness condition called ( L 0 , L 1 ) smoothness, which assumes the local smoothness or Hessian norm is bounded by an affine function of the gradient norm. The assumption was well-validated by extensive experiments conducted on language models. Various analyses of different algorithms under this condition were later developed [48; 34; 52; 17; 38; 8]. One recent closely-related work is [45] which studies converges of Adam under the ( L 0 , L 1 ) smoothness condition. However, their results are still limited, as we have mentioned above. In this paper, we consider an even more general smoothness condition where the local smoothness is bounded by a sub-quadratic function of the gradient norm, and prove the convergence of Adam under this condition. In our concurrent work [26], we further analyze various other algorithms in both convex and non-convex settings under similar generalized smoothness conditions following the same key idea of bounding gradients along the trajectory.

3 Preliminaries

Notation. Let ∥·∥ denote the Euclidean norm of a vector or spectral norm of a matrix. For any given vector x , we use ( x ) i to denote its i -th coordinate and x 2 , √ x , | x | to denote its coordinate-wise square, square root, and absolute value respectively. For any two vectors x and y , we use x ⊙ y and x/y to denote their coordinate-wise product and quotient respectively. We also write x ⪯ y or x ⪰ y to denote the coordinate-wise inequality between x and y , which means ( x ) i ≤ ( y ) i or ( x ) i ≥ ( y ) i for each coordinate index i . For two symmetric real matrices A and B , we say A ⪯ B or A ⪰ B if B -A or A -B is positive semi-definite (PSD). Given two real numbers a, b ∈ R , we denote a ∧ b := min { a, b } for simplicity. Finally, we use O ( · ) , Θ( · ) , and Ω( · ) for the standard big-O, big-Theta, and big-Omega notation.

3.1 Description of the Adam algorithm

Algorithm 1 ADAM

The formal definition of Adam proposed in [23] is shown in Algorithm 1, where Lines 5-9 describe the update rule of iterates { x t } 1 ≤ t ≤ T . Lines 5-6 are the updates for the first and second order momentum, m t and v t , respectively. In Lines 7-8, they are re-scaled to ˆ m t and ˆ v t in order to correct the initialization bias due to setting m 0 = v 0 = 0 . Then the iterate is updated by x t +1 = x t -h t ⊙ ˆ m t where h t = η/ ( √ ˆ v t + λ ) is the adaptive stepsize vector for some parameters η and λ .

3.2 Assumptions

In what follows below, we will state our main assumptions for analysis of Adam.

3.2.1 Function class

We start with a standard assumption in optimization on the objective function f whose domain lies in a Euclidean space with dimension d .

Assumption 1. The objective function f is differentiable and closed within its open domain dom( f ) ⊆ R d and is bounded from below, i.e., f ∗ := inf x f ( x ) > -∞ .

Remark 3.1 . A function f is said to be closed if its sub-level set { x ∈ dom( f ) | f ( x ) ≤ a } is closed for each a ∈ R . In addition, a continuous function f over an open domain is closed if and only f ( x ) tends to infinity whenever x approaches to the boundary of dom( f ) , which is an important condition to ensure the iterates of Adam with a small enough stepsize η stay within the domain with high probability. Note that this condition is mild since any continuous function defined over the entire space R d is closed.

Besides Assumption 1, the only additional assumption we make regarding f is that its local smoothness is bounded by a sub-quadratic function of the gradient norm. More formally, we consider the following ( ρ, L 0 , L ρ ) smoothness condition with 0 ≤ ρ < 2 .

Definition 3.2. A differentiable real-valued function f is ( ρ, L 0 , L ρ ) smooth for some constants ρ, L 0 , L ρ ≥ 0 if the following inequality holds almost everywhere in dom( f )

Remark 3.3 . When ρ = 1 , Definition 3.2 reduces to the ( L 0 , L 1 ) smoothness condition in [49]. When ρ = 0 or L ρ = 0 , it reduces to the standard smoothness condition.

Assumption 2. The objective function f is ( ρ, L 0 , L ρ ) smooth with 0 ≤ ρ < 2 .

The standard smooth function class is very restrictive as it only contains functions that are upper and lower bounded by quadratic functions. The ( L 0 , L 1 ) smooth function class is more general since it also contains, e.g., univariate polynomials and exponential functions. Assumption 2 is even more general and contains univariate rational functions, double exponential functions, etc. See Appendix D.1 for the formal propositions and proofs. We also refer the reader to our concurrent work [26] for more detailed discussions of examples of ( ρ, L 0 , L ρ ) smooth functions for different ρ s.

It turns out that bounded Hessian norm at a point x implies local Lipschitzness of the gradient in the neighborhood around x . In particular, we have the following lemma.

Lemma 3.4. Under Assumptions 1 and 2, for any a > 0 and two points x ∈ dom( f ) , y ∈ R d such that ∥ y -x ∥ ≤ a L 0 + L ρ ( ∥∇ f ( x ) ∥ + a ) ρ , we have y ∈ dom( f ) and

Remark 3.5 . Lemma 3.4 can be actually used as the definition of ( ρ, L 0 , L ρ ) smooth functions in place of Assumption 2. Besides the local gradient Lipschitz condition, it also suggests that, as long as the update at each step is small enough, the iterates will not go outside of the domain.

̸

For the special case of ρ = 1 , choosing a = max {∥∇ f ( x ) ∥ , L 0 /L 1 } , one can verify that the required locality size in Lemma 3.4 satisfies a L 0 + L 1 ( ∥∇ f ( x ) ∥ + a ) ≥ 1 3 L 1 . In this case, Lemma 3.4 states that ∥ x -y ∥ ≤ 1 / (3 L 1 ) implies ∥∇ f ( y ) -∇ f ( x ) ∥ ≤ 2( L 0 + L 1 ∥∇ f ( x ) ∥ ) ∥ y -x ∥ . Therefore, it reduces to the local gradient Lipschitz condition for ( L 0 , L 1 ) smooth functions in [49; 48] up to numerical constant factors. For ρ = 1 , the proof is more involved because Grönwall's inequality used in [49; 48] no longer applies. Therefore we defer the detailed proof of Lemma 3.4 to Appendix D.2.

3.2.2 Stochastic gradient

We consider one of the following two assumptions on the stochastic gradient ∇ f ( x t , ξ t ) in our analysis of Adam.

Assumption 3. The gradient noise is centered and almost surely bounded. In particular, for some σ ≥ 0 and all t ≥ 1 ,

where E t -1 [ · ] := E [ · | ξ 1 , . . . , ξ t -1 ] is the conditional expectation given ξ 1 , . . . , ξ t -1 .

Assumption 4. The gradient noise is centered with sub-Gaussian norm. In particular, for some R ≥ 0 and all t ≥ 1 ,

where E t -1 [ · ] := E [ · | ξ 1 , . . . , ξ t -1 ] and P t -1 [ · ] := P [ · | ξ 1 , . . . , ξ t -1 ] are the conditional expectation and probability given ξ 1 , . . . , ξ t -1 .

Assumption 4 is strictly weaker than Assumption 3 since an almost surely bounded random variable clearly has sub-Gaussian norm, but it results in a slightly worse convergece rate up to poly-log factors (see Theorems 4.1 and 4.2). Both of them are stronger than the most standard bounded variance assumption E [ ∥∇ f ( x t , ξ t ) -∇ f ( x t ) ∥ 2 ] ≤ σ 2 for some σ ≥ 0 , although Assumption 3 is actually a common assumption in existing analyses under the ( L 0 , L 1 ) smoothness condition (see e.g. [49; 48]). The extension to the bounded variance assumption is challenging and a very interesting future work as it is also the assumption considered in the lower bound [3]. We suspect that such an extension would be straightforward if we consider a mini-batch version of Algorithm 1 with a batch size of S = Ω( ϵ -2 ) , since this results in a very small variance of O ( ϵ 2 ) and thus essentially reduces the analysis to the deterministic setting. However, for practical Adam with an O (1) batch size, the extension is challenging and we leave it as a future work.

4 Results

In the section, we provide our convergence results for Adam under Assumptions 1, 2, and 3 or 4. To keep the statements of the theorems concise, we first define several problem-dependent constants. First, we let ∆ 1 := f ( x 1 ) -f ∗ < ∞ be the initial sub-optimality gap. Next, given a large enough constant G > 0 , we define

where L can be viewed as the effective smoothness constant along the trajectory if one can show ∥∇ f ( x t ) ∥ ≤ G and ∥ x t +1 -x t ∥ ≤ r at each step (see Section 5 for more detailed discussions). We will also use c 1 , c 2 to denote some small enough numerical constants and C 1 , C 2 to denote some large enough ones. The formal convergence results under Assumptions 1, 2, and 3 are presented in the following theorem, whose proof is deferred in Appendix E.

Theorem 4.1. Suppose Assumptions 1, 2, and 3 hold. Denote ι := log(1 /δ ) for any 0 < δ < 1 . Let G be a constant satisfying G ≥ max { 2 λ, 2 σ, √ C 1 ∆ 1 L 0 , ( C 1 ∆ 1 L ρ ) 1 2 -ρ } . Choose

Let T = max { 1 β 2 , C 2 ∆ 1 G ηϵ 2 } . Then with probability at least 1 -δ , we have ∥∇ f ( x t ) ∥ ≤ G for every 1 ≤ t ≤ T , and 1 ∑ T ∥∇ f ( x t ) ∥ 2 ≤ ϵ 2 .

Note that G , the upper bound of gradients along the trajectory, is a constant that depends on λ, σ, L 0 , L ρ , and the initial sub-optimality gap ∆ 1 , but not on ϵ . There is no requirement on the second order momentum parameter β sq , although many existing works like [12; 51; 45] need certain restrictions on it. We choose very small β and η , both of which are O ( ϵ 2 ) . Therefore, from the choice of T , it is clear that we obtain a gradient complexity of O ( ϵ -4 ) , where we only consider the leading term. We are not clear whether the dependence on ϵ is optimal or not, as the Ω( ϵ -4 ) lower bound in [3] assumes the weaker bounded variance assumption than our Assumpion 3. However, it matches the state-of-the-art complexity among existing analyses of Adam.

One limitation of the dependence of our complexity on λ is O ( λ -2 ) , which might be large since λ is usually small in practice, e.g., the default choice is λ = 10 -8 in the PyTorch implementation. There are some existing analyses on Adam [12; 51; 45] whose rates do not depend explicitly on λ or only depend on log(1 /λ ) . However, all of them depend on poly( d ) , whereas our rate is dimension free. The dimension d is also very large, especially when training transformers, for which Adam is widely used. We believe that independence on d is better than that on λ , because d is fixed given the architecture of the neural network but λ is a hyper-parameter which we have the freedom to tune. In fact, based on our preliminary experimental results on CIFAR-10 shown in Figure 1, the performance of Adam is not very sensitive to the choice of λ . Although the default choice of λ is 10 -8 , increasing it up to 0 . 01 only makes minor differences.

As discussed in Section 3.2.2, we can generalize the bounded gradient noise condition in Assumption 3 to the weaker sub-Gaussian noise condition in Assumption 4. The following theorem formally shows the convergence result under Assumptions 1, 2, and 4, whose proof is deferred in Appendix E.6.

0

100

200

300

400

Epoch

0

20

40

60

80

Test Error (%)

λ=1e-4

λ=1e-3

λ=0.01

λ=0.1

λ=1

(a) CNN

0

100

200

300

400

Epoch

0

20

40

60

80

Test Error (%)

λ=1e-4

λ=1e-3

λ=0.01

λ=0.1

λ=1

(b) ResNet-Small

0

100

200

300

400

Epoch

0

20

40

60

80

Test Error (%)

λ=1e-4

λ=1e-3

λ=0.01

λ=0.1

λ=1

(c) ResNet110

Theorem 4.2. Suppose Assumptions 1, 2, and 4 hold. Denote ι := log(2 /δ ) and σ := R √ 2 log(4 T/δ ) for any 0 < δ < 1 . Let G be a constant satisfying G ≥ max { 2 λ, 2 σ, √ C 1 ∆ 1 L 0 , ( C 1 ∆ 1 L ρ ) 1 2 -ρ } . Choose

Let T = max { 1 β 2 , C 2 ∆ 1 G ηϵ 2 } . Then with probability at least 1 -δ , we have ∥∇ f ( x t ) ∥ ≤ G for every 1 ≤ t ≤ T , and 1 T ∑ T t =1 ∥∇ f ( x t ) ∥ 2 ≤ ϵ 2 .

Note that the main difference of Theorem 4.2 from Theorem 4.1 is that σ is now O ( √ log T ) instead of a constant. With some standard calculations, one can show that the gradient complexity in Theorem 4.2 is bounded by O ( ϵ -4 log p (1 /ϵ )) , where p = max { 3 , 9+2 ρ 4 } < 3 . 25 .

5 Analysis

5.1 Bounding the gradients along the optimization trajectory

We want to bound the gradients along the optimization trajectory mainly for two reasons. First, as discussed in Section 2, many existing analyses of Adam rely on the assumption of bounded gradients, because unbounded gradient norm leads to unbounded second order momentum ˆ v t which implies very small stepsize, and slow convergence. On the other hand, once the gradients are bounded, it is straightforward to control ˆ v t as well as the stepsize, and therefore the analysis essentially reduces to the easier one for AdaBound. Second, informally speaking 1 , under Assumption 2, bounded gradients also imply bounded Hessians, which essentially reduces the ( ρ, L 0 , L ρ ) smoothness to the standard smoothness. See Section 5.2 for more formal discussions.

In this paper, instead of imposing the strong assumption of globally bounded gradients, we develop a new analysis to show that with high probability, the gradients are always bounded along the trajectory of Adam until convergence. The essential idea can be informally illustrated by the following 'circular" reasoning that we will make precise later. On the one hand, if ∥∇ f ( x t ) ∥ ≤ G for every t ≥ 1 , it is not hard to show the gradient converges to zero based on our discussions above. On the other hand, we know that a converging sequence must be upper bounded. Therefore there exists some G ′ such that ∥∇ f ( x t ) ∥ ≤ G ′ for every t ≥ 1 . In other words, the bounded gradient condition implies the convergence result and the convergence result also implies the boundedness condition, forming a circular argument.

This circular argument is of course flawed. However, we can break the circularity of reasoning and rigorously prove both the bounded gradient condition and the convergence result using a contradiction

argument. Before introducing the contradiction argument, we first need to provide the following useful lemma, which is the reverse direction of a generalized Polyak-Lojasiewicz (PL) inequality, whose proof is deferred in Appendix D.3.

Lemma5.1. Under Assumptions 1 and 2, we have ∥∇ f ( x ) ∥ 2 ≤ 3(3 L 0 +4 L ρ ∥∇ f ( x ) ∥ ρ )( f ( x ) -f ∗ ) .

Define the function ζ ( u ) := u 2 3(3 L 0 +4 L ρ u ρ ) over u ≥ 0 . It is easy to verify that if ρ < 2 , ζ is increasing and its range is [0 , ∞ ) . Therefore, ζ is invertible and ζ -1 is also increasing. Then, for any constant G > 0 , denoting F = ζ ( G ) , Lemma 5.1 suggests that if f ( x ) -f ∗ ≤ F , we have

In other words, if ρ < 2 , the gradient is bounded within any sub-level set, even though the sub-level set could be unbounded. Then, let τ be the first time the sub-optimality gap is strictly greater than F , truncated at T +1 , or formally,

Then at least when t < τ , we have f ( x t ) -f ∗ ≤ F and thus ∥∇ f ( x t ) ∥ ≤ G . Based on our discussions above, it is not hard to analyze the updates before time τ , and one can contruct some Lyapunov function to obtain an upper bound on f ( x τ ) -f ∗ . On the other hand, if τ ≤ T , we immediately obtain a lower bound on f ( x τ ) , that is f ( x τ ) -f ∗ > F , by the definition of τ in (3). If the lower bound is greater than the upper bound, it leads to a contradiction, which shows τ = T +1 , i.e., the sub-optimality gap and the gradient norm are always bounded by F and G respectively before the algorithm terminates. We will illustrate the technique in more details in the simple deterministic setting in Section 5.3, but first, in Section 5.2, we introduce several prerequisite lemmas on the ( ρ, L 0 , L ρ ) smoothness.

5.2 Local smoothness

In Section 5.1, we informally mentioned that ( ρ, L 0 , L ρ ) smoothness essentially reduces to the standard smoothness if the gradient is bounded. In this section, we will make the statement more precise. First, note that Lemma 3.4 implies the following useful corollary.

Corollary 5.2. Under Assumptions 1 and 2, for any G > 0 and two points x ∈ dom( f ) , y ∈ R d such that ∥∇ f ( x ) ∥ ≤ G and ∥ y -x ∥ ≤ r := min { 1 5 L ρ G ρ -1 , 1 5( L ρ -1 0 L ρ ) 1 /ρ } , denoting L := 3 L 0 +4 L ρ G ρ , we have y ∈ dom( f ) and

The proof of Corollary 5.2 is deferred in Appendix D.4. Although the inequalities in Corollary 5.2 look very similar to the standard global smoothness condition with constant L , it is still a local condition as it requires ∥ x -y ∥ ≤ r . Fortunately, at least before τ , such a requirement is easy to satisfy for small enough η , according to the following lemma whose proof is deferred in Appendix E.5.

Lemma 5.3. Under Assumption 3, if t < τ and choosing G ≥ σ , we have ∥ x t +1 -x t ∥ ≤ ηD where D := 2 G/λ .

Then as long as η ≤ r/D , we have ∥ x t +1 -x t ∥ ≤ r which satisfies the requirement in Corollary 5.2. Then we can apply the inequalities in it in the same way as the standard smoothness condition. In other words, most classical inequalities derived for standard smooth functions also apply to ( ρ, L 0 , L ρ ) smooth functions.

5.3 Warm-up: analysis in the deterministic setting

In this section, we consider the simpler deterministic setting where the stochastic gradient ∇ f ( x t , ξ t ) in Algorithm 1 or (18) is replaced with the exact gradient ∇ f ( x t ) . As discussed in Section 5.1, the key in our contradiction argument is to obtain both upper and lower bounds on f ( x τ ) -f ∗ . In the following derivations, we focus on illustrating the main idea of our analysis technique and ignore minor proof details. In addition, all of them are under Assumptions 1, 2, and 3.

In order to obtain the upper bound, we need the following two lemmas. First, denoting ϵ t := ˆ m t -∇ f ( x t ) , we can obtain the following informal descent lemma for deterministic Adam.

Lemma 5.4 (Descent lemma, informal) . For any t < τ , choosing G ≥ λ and a small enough η ,

where ' ⪅ ' omits less important terms.

Compared with the standard descent lemma for gradient descent, there is an additional term of ∥ ϵ t ∥ 2 in Lemma 5.4. In the next lemma, we bound this term recursively.

Lemma 5.5 (Informal) . Choosing β = Θ( ηG ρ +1 / 2 ) , if t < τ , we have

The proof sketches of the above two lemmas are deferred in Appendix B. Now we combine them to get the upper bound on f ( x τ ) -f ∗ . Define the function Φ t := f ( x t ) -f ∗ + 2 η λβ ∥ ϵ t ∥ 2 . Note that for any t < τ , (4) + 2 η λβ × (5) gives

The above inequality shows Φ t is non-increasing and thus a Lyapunov function. Therefore, we have

where in the last inequality we use Φ 1 = f ( x 1 ) -f ∗ = ∆ 1 since ϵ 1 = ˆ m 1 -∇ f ( x 1 ) = 0 in the deterministic setting.

As discussed in Section 5.1, if τ ≤ T , we have F < f ( x τ ) -f ∗ ≤ ∆ 1 . Note that we are able to choose a large enough constant G so that F = G 2 3(3 L 0 +4 L ρ G ρ ) is greater than ∆ 1 , which leads to a contradiction and shows τ = T +1 . Therefore, (6) holds for all 1 ≤ t ≤ T . Taking a summation over 1 ≤ t ≤ T and re-arranging terms, we get

if choosing T ≥ 8 G ∆ 1 ηϵ 2 , i.e., it shows convergence with a gradient complexity of O ( ϵ -2 ) since both G and η are constants independent of ϵ in the deterministic setting.

5.4 Extension to the stochastic setting

In this part, we briefly introduce how to extend the analysis to the more challenging stochastic setting. It becomes harder to obtain an upper bound on f ( x τ ) -f ∗ because Φ t is no longer non-increasing due to the existence of noise. In addition, τ defined in (3) is now a random variable. Note that all the derivations, such as Lemmas 5.4 and 5.5, are conditioned on the random event t < τ . Therefore, one can not simply take a total expectation of them to show E [Φ t ] is non-increasing.

Fortunately, τ is in fact a stopping time with nice properties. If the noise is almost surely bounded as in Assumption 3, by a more careful analysis, we can obtain a high probability upper bound on f ( x τ ) -f ∗ using concentration inequalities. Then we can still obtain a contradiction and convergence under this high probability event. If the noise has sub-Gaussian norm as in Assumption 4, one can change the definition of τ to

for appropriately chosen F and σ . Then at least when t < τ , the noise is bounded by σ . Hence we can get the same upper bound on f ( x τ ) -f ∗ as if Assumption 3 still holds. However, when t ≤ T , the lower bound f ( x τ ) -f ∗ > F does not necessarily holds, which requires some more careful analyses. The details of the proofs are involved and we defer them in Appendix E.

6 Variance-reduced Adam

In this section, we propose a variance-reduced version of Adam (VRAdam). This new algorithm is depicted in Algorithm 2. Its main difference from the original Adam is that in the momentum update

rule (Line 6), an additional term of (1 -β ) ( ∇ f ( x t , ξ t ) -∇ f ( x t -1 , ξ t )) is added, inspired by the STORM algorithm [9]. This term corrects the bias of m t so that it is an unbiased estimate of ∇ f ( x t ) in the sense of total expectation, i.e., E [ m t ] = ∇ f ( x t ) . We will also show that it reduces the variance and accelerates the convergence.

Aside from the adaptive stepsize, one major difference between Algorithm 2 and STORM is that our hyper-parameters η and β are fixed constants whereas theirs are decreasing as a function of t . Choosing constant hyper-parameters requires a more accurate estimate at the initialization. That is why we use a mega-batch S 1 to evaluate the gradient at the initial point to initialize m 1 and v 1 (Lines 2-3). In practice, one can also do a full-batch gradient evaluation at initialization. Note that there is no initialization bias for the momentum, so we do not re-scale m t and only re-scale v t . We also want to point out that although the initial mega-batch gradient evaluation makes the algorithm a bit harder to implement, constant hyper-parameters are usually easier to tune and more common in training deep neural networks. It should be not hard to extend our analysis to time-decreasing η and β and we leave it as an interesting future work.

Algorithm 2 VARIANCE-REDUCED ADAM (VRADAM)

1: Input: β, β sq , η, λ, T, S 1 , x init

2: Draw a batch of samples S 1 with size S 1 and use them to evaluate the gradient ∇ f ( x init , S 1 ) .

3: Initialize m 1 = ∇ f ( x init , S 1 ) , v 1 = β sq m 2 1 , and x 2 = x init -ηm 1 | m 1 | + λ .

4: for t = 2 , · · · , T do

5: Draw a new sample ξ t and perform the following updates:

6: m t = (1 -β ) m t -1 + β ∇ f ( x t , ξ t )+(1 -β ) ( ∇ f ( x t , ξ t ) -∇ f ( x t -1 , ξ t ))

8: ˆ v t = v t 1 -(1 -β sq ) t

7: v t = (1 -β sq ) v t -1 + β sq ( ∇ f ( x t , ξ t )) 2

9:

x

t

+1

η

ˆ

v

10: end for

In addition to Assumption 1, we need to impose the following assumptions which can be viewed as stronger versions of Assumptions 2 and 3, respectively.

Assumption 5. The objective function f and the component function f ( · , ξ ) for each fixed ξ are ( ρ, L 0 , L ρ ) smooth with 0 ≤ ρ < 2 .

Assumption 6. The random variables { ξ t } 1 ≤ t ≤ T are sampled i.i.d. from some distribution P such that for any x ∈ dom( f ) ,

Remark 6.1 . Assumption 6 is stronger than Assumption 3. Assumption 3 applies only to the iterates generated by the algorithm, while Assumption 6 is a pointwise assumption over all x ∈ dom( f ) and further assumes an i.i.d. nature of the random variables { ξ t } 1 ≤ t ≤ T . Also note that, similar to Adam, it is straightforward to generalize the assumption to noise with sub-Gaussian norm as in Assumption 4.

6.1 Analysis

In this part, we briefly discuss challenges in the analysis of VRAdam. The detailed analysis is deferred in Appendix F. Note that Corollary 5.2 requires bounded update ∥ x t +1 -x t ∥ ≤ r at each step. For Adam, it is easy to satisfy for a small enough η according to Lemma 5.3. However, for VRAdam, obtaining a good enough almost sure bound on the update is challenging even though the gradient noise is bounded. To bypass this difficulty, we directly impose a bound on ∥∇ f ( x t ) -m t ∥ by changing the definition of the stopping time τ , similar to how we deal with the sub-Gaussian noise condition for Adam. In particular, we define

Then by definition, both ∥∇ f ( x t ) ∥ and ∥∇ f ( x t ) -m t ∥ are bounded by G before time τ , which directly implies bounded update ∥ x t +1 -x t ∥ . Of course, the new definition brings new challenges to lower bounding f ( x τ ) -f ∗ , which requires more careful analyses specific to the VRAdam algorithm. Please see Appendix F for the details.

t

+

=

x

t

-

√

⊙

m

t

λ

6.2 Convergence guarantees for VRAdam

In the section, we provide our main results for convergence of VRAdam under Assumptions 1, 5, and 6. We consider the same definitions of problem-dependent constants ∆ 1 , r, L as those in Section 4 to make the statements of theorems concise. Let c be a small enough numerical constant and C be a large enough numerical constant. The formal convergence result is shown in the following theorem.

Theorem 6.2. Suppose Assumptions 1, 5, and 6 hold. For any 0 < δ < 1 , let G > 0 be a constant satisfying G ≥ max { 2 λ, 2 σ, √ C ∆ 1 L 0 /δ, ( C ∆ 1 L ρ /δ ) 1 2 -ρ } . Choose 0 ≤ β sq ≤ 1 and β = a 2 η 2 , where a = 40 L √ Gλ -3 / 2 . Choose

Then with probability at least 1 -δ , we have ∥∇ f ( x t ) ∥ ≤ G for every 1 ≤ t ≤ T , and 1 T ∑ T t =1 ∥∇ f ( x t ) ∥ 2 ≤ ϵ 2 .

Note that the choice of G , the upper bound of gradients along the trajectory of VRAdam, is very similar to that in Theorem 4.1 for Adam. The only difference is that now it also depends on the failure probability δ . Similar to Theorem 4.1, there is no requirement on β sq and we choose a very small β = O ( ϵ 2 ) . However, the variance reduction technique allows us to take a larger stepsize η = O ( ϵ ) (compared with O ( ϵ 2 ) for Adam) and obtain an accelerated gradient complexity of O ( ϵ -3 ) , where we only consider the leading term. We are not sure whether it is optimal as the Ω( ϵ -3 ) lower bound in [3] assumes the weaker bounded variance condition. However, our result significantly improves upon [47], which considers a variance-reduced version of Adam by combining Adam and SVRG [22] and only obtains asymptotic convergence in the non-convex setting. Similar to Adam, our gradient complexity for VRAdam is dimension free but its dependence on λ is O ( λ -2 ) . Another limitation is that, the dependence on the failure probability δ is polynomial, worse than the poly-log dependence in Theorem 4.1 for Adam.

7 Conclusion and future works

In this paper, we proved the convergence of Adam and its variance-reduced version under less restrictive assumptions compared to those in the existing literature. We considered a generalized nonuniform smoothness condition, according to which the Hessian norm is bounded by a sub-quadratic function of the gradient norm almost everywhere. Instead of assuming the Lipschitzness of the objective function as in existing analyses of Adam, we use a new contradiction argument to prove that gradients are bounded by a constant along the optimization trajectory. There are several interesting future directions that one could pursue following this work.

Relaxation of the bounded noise assumption. Our analysis relies on the assumption of bounded noise or noise with sub-Gaussian norm. However, the existing lower bounds in [3] consider the weaker bounded variance assumption. Hence, it is not clear whether the O ( ϵ -4 ) complexity we obtain for Adam is tight in this setting. It will be interesting to see whether one can relax the assumption to the bounded variance setting. One may gain some insights from recent papers such as [16; 46] that analyze AdaGrad under weak noise conditions. An alternative way to show the tightness of the O ( ϵ -4 ) complexity is to prove a lower bound under the bounded noise assumption.

Potential applications of our technique. Another interesting future direction is to see if the techniques developed in this work for bounding gradients (including those in the the concurrent work [26]) can be generalized to improve the convergence results for other optimization problems and algorithms. We believe it is possible so long as the function class is well behaved and the algorithm is efficient enough so that f ( x τ ) -f ∗ can be well bounded for some appropriately defined stopping time τ .

Understanding why Adam is better than SGD. We want to note that our results can not explain why Adam is better than SGD for training transformers, because [26] shows that non-adaptive SGD converges with the same O ( ϵ -4 ) gradient complexity under even weaker conditions. It would be interesting and impactful if one can find a reasonable setting (function class, gradient oracle, etc) under which Adam or other adaptive methods provably outperform SGD.

Acknowledgments

This work was supported, in part, by the MIT-IBM Watson AI Lab and ONR Grants N00014-20-12394 and N00014-23-1-2299. We also acknowledge support from DOE under grant DE-SC0022199, and NSF through awards DMS-2031883 and DMS-1953181.

References

Kwangjun Ahn, Xiang Cheng, Minhak Song, Chulhee Yun, Ali Jadbabaie, and Suvrit Sra. Linear attention is (maybe) all you need (to understand transformer optimization). arXiv preprint arXiv:2310.01082 , 2023.

Zeyuan Allen-Zhu and Elad Hazan. Variance reduction for faster non-convex optimization. In International conference on machine learning , pages 699-707. PMLR, 2016.

Yossi Arjevani, Yair Carmon, John C Duchi, Dylan J Foster, Nathan Srebro, and Blake Woodworth. Lower bounds for non-convex stochastic optimization. Mathematical Programming , 199 (1-2):165-214, 2023.

Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, T. J. Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeff Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. ArXiv , abs/2005.14165, 2020.

Congliang Chen, Li Shen, Fangyu Zou, and Wei Liu. Towards practical adam: Non-convexity, convergence theory, and mini-batch acceleration. The Journal of Machine Learning Research , 23(1):10411-10457, 2022.

Xiangyi Chen, Sijia Liu, Ruoyu Sun, and Mingyi Hong. On the convergence of a class of adam-type algorithms for non-convex optimization. arXiv preprint arXiv:1808.02941 , 2018.

Ziyi Chen, Yi Zhou, Yingbin Liang, and Zhaosong Lu. Generalized-smooth nonconvex optimization is as efficient as smooth nonconvex optimization. arXiv preprint arXiv:2303.02854 , 2023.

Michael Crawshaw, Mingrui Liu, Francesco Orabona, Wei Zhang, and Zhenxun Zhuang. Robustness to unbounded smoothness of generalized signsgd. Advances in Neural Information Processing Systems , 35:9955-9968, 2022.

Ashok Cutkosky and Francesco Orabona. Momentum-based variance reduction in non-convex sgd. ArXiv , abs/1905.10018, 2019.

Soham De, Anirbit Mukherjee, and Enayat Ullah. Convergence guarantees for rmsprop and adam in non-convex optimization and an empirical comparison to nesterov acceleration. arXiv: Learning , 2018.

Aaron Defazio, Francis Bach, and Simon Lacoste-Julien. Saga: A fast incremental gradient method with support for non-strongly convex composite objectives. Advances in neural information processing systems , 27, 2014.

Alexandre D'efossez, Léon Bottou, Francis R. Bach, and Nicolas Usunier. A simple convergence proof of adam and adagrad. arXiv: Machine Learning , 2020.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. ArXiv , abs/1810.04805, 2019.

Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. ArXiv , abs/2010.11929, 2020.

Cong Fang, Chris Junchi Li, Zhouchen Lin, and Tong Zhang. Spider: Near-optimal nonconvex optimization via stochastic path-integrated differential estimator. Advances in neural information processing systems , 31, 2018.

Matthew Faw, Isidoros Tziotis, Constantine Caramanis, Aryan Mokhtari, Sanjay Shakkottai, and Rachel 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.

Matthew Faw, Litu Rout, Constantine Caramanis, and Sanjay Shakkottai. Beyond uniform smoothness: A stopped analysis of adaptive sgd. arXiv preprint arXiv:2302.06570 , 2023.

Sébastien Gadat and Ioana Gavra. Asymptotic study of stochastic adaptive algorithms in non-convex landscape. The Journal of Machine Learning Research , 23(1):10357-10410, 2022.

Zhishuai Guo, Yi Xu, Wotao Yin, Rong Jin, and Tianbao Yang. A novel convergence analysis for algorithms of the adam family. ArXiv , abs/2112.03459, 2021.

Hideaki Iiduka. Theoretical analysis of adam using hyperparameters close to one without lipschitz smoothness. Numerical Algorithms , pages 1-39, 2023.

Phillip Isola, Jun-Yan Zhu, Tinghui Zhou, and Alexei A. Efros. Image-to-image translation with conditional adversarial networks. 2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , pages 5967-5976, 2016.

Rie Johnson and Tong Zhang. Accelerating stochastic gradient descent using predictive variance reduction. In NIPS , 2013.

Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. CoRR , abs/1412.6980, 2014.

Frederik Kunstner, Jacques Chen, Jonathan Wilder Lavington, and Mark Schmidt. Noise is not the main factor behind the gap between sgd and adam on transformers, but sign descent might be. arXiv preprint arXiv:2304.13960 , 2023.

Lihua Lei, Cheng Ju, Jianbo Chen, and Michael I Jordan. Non-convex finite-sum optimization via scsg methods. Advances in Neural Information Processing Systems , 30, 2017.

Haochuan Li, Jian Qian, Yi Tian, Alexander Rakhlin, and Ali Jadbabaie. Convex and non-convex optimization under generalized smoothness. arXiv preprint arXiv:2306.01264 , 2023.

Zhize Li, Hongyan Bao, Xiangliang Zhang, and Peter Richtárik. Page: A simple and optimal probabilistic gradient estimator for nonconvex optimization. In International conference on machine learning , pages 6286-6295. PMLR, 2021.

Timothy P. Lillicrap, Jonathan J. Hunt, Alexander Pritzel, Nicolas Manfred Otto Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. CoRR , abs/1509.02971, 2015.

Deyi Liu, Lam M Nguyen, and Quoc Tran-Dinh. An optimal hybrid variance-reduced algorithm for stochastic composite nonconvex optimization. arXiv preprint arXiv:2008.09055 , 2020.

Zijian Liu, Perry Dong, Srikanth Jagabathula, and Zhengyuan Zhou. Near-optimal highprobability convergence for non-convex stochastic optimization with variance reduction. arXiv preprint arXiv:2302.06032 , 2023.

Liangchen Luo, Yuanhao Xiong, Yan Liu, and Xu Sun. Adaptive gradient methods with dynamic bound of learning rate. ArXiv , abs/1902.09843, 2019.

Julien Mairal. Optimization with first-order surrogate functions. In International Conference on Machine Learning , pages 783-791. PMLR, 2013.

Volodymyr Mnih, Adrià Puigdomènech Badia, Mehdi Mirza, Alex Graves, Timothy P. Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. ArXiv , abs/1602.01783, 2016.

Jiang Qian, Yuren Wu, Bojin Zhuang, Shaojun Wang, and Jing Xiao. Understanding gradient clipping in incremental gradient methods. In International Conference on Artificial Intelligence and Statistics , 2021.

Alec Radford, Luke Metz, and Soumith Chintala. Unsupervised representation learning with deep convolutional generative adversarial networks. CoRR , abs/1511.06434, 2015.

Sashank J. Reddi, Ahmed Hefny, Suvrit Sra, Barnabas Poczos, and Alex Smola. Stochastic variance reduction for nonconvex optimization. In Maria Florina Balcan and Kilian Q. Weinberger, editors, Proceedings of The 33rd International Conference on Machine Learning , volume 48 of Proceedings of Machine Learning Research , pages 314-323, New York, New York, USA, 20-22 Jun 2016. PMLR. URL https://proceedings.mlr.press/v48/reddi16.html .

Sashank J. Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence of adam and beyond. ArXiv , abs/1904.09237, 2018.

Amirhossein Reisizadeh, Haochuan Li, Subhro Das, and Ali Jadbabaie. Variance-reduced clipping for non-convex optimization. arXiv preprint arXiv:2303.00883 , 2023.

Nicolas Roux, Mark Schmidt, and Francis Bach. A stochastic gradient method with an exponential convergence _rate for finite training sets. In F. Pereira, C.J. Burges, L. Bottou, and K.Q. Weinberger, editors, Advances in Neural Information Processing Systems , volume 25. Curran Associates, Inc., 2012. URL https://proceedings.neurips.cc/paper_files/paper/ 2012/file/905056c1ac1dad141560467e0a99e1cf-Paper.pdf .

John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. ArXiv , abs/1707.06347, 2017.

Shai Shalev-Shwartz and Tong Zhang. Stochastic dual coordinate ascent methods for regularized loss minimization. Journal of Machine Learning Research , 14(1), 2013.

Naichen Shi, Dawei Li, Mingyi Hong, and Ruoyu Sun. Rmsprop converges with proper hyper-parameter. In International Conference on Learning Representations , 2021.

Quoc Tran-Dinh, Nhan H Pham, Dzung T Phan, and Lam M Nguyen. Hybrid stochastic gradient descent algorithms for stochastic nonconvex optimization. arXiv preprint arXiv:1905.05920 , 2019.

Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. ArXiv , abs/1706.03762, 2017.

Bohan Wang, Yushun Zhang, Huishuai Zhang, Qi Meng, Zhirui Ma, Tie-Yan Liu, and Wei Chen. Provable adaptivity in adam. ArXiv , abs/2208.09900, 2022.

Bohan Wang, Huishuai Zhang, Zhiming Ma, and Wei 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.

Ruiqi Wang and Diego Klabjan. Divergence results and convergence of a variance reduced version of adam. ArXiv , abs/2210.05607, 2022.

Bohang Zhang, Jikai Jin, Cong Fang, and Liwei Wang. Improved analysis of clipping algorithms for non-convex optimization. Advances in Neural Information Processing Systems , 33:1551115521, 2020.

J. Zhang, Tianxing He, Suvrit Sra, and Ali Jadbabaie. Why gradient clipping accelerates training: A theoretical justification for adaptivity. arXiv: Optimization and Control , 2019.

Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiv Kumar, and Suvrit Sra. Why are adaptive methods good for attention models? Advances in Neural Information Processing Systems , 33:15383-15393, 2020.

Yushun Zhang, Congliang Chen, Naichen Shi, Ruoyu Sun, and Zhimin Luo. Adam can converge without any modification on update rules. ArXiv , abs/2208.09632, 2022.

Shen-Yi Zhao, Yin-Peng Xie, and Wu-Jun Li. On the convergence and improvement of stochastic normalized gradient descent. Science China Information Sciences , 64, 2021.

Dongruo Zhou, Jinghui Chen, Yuan Cao, Yiqi Tang, Ziyan Yang, and Quanquan Gu. On the convergence of adaptive gradient methods for nonconvex optimization. arXiv preprint arXiv:1808.05671 , 2018.

Zhiming Zhou, Qingru Zhang, Guansong Lu, Hongwei Wang, Weinan Zhang, and Yong Yu. Adashift: Decorrelation and convergence of adaptive learning rate methods. ArXiv , abs/1810.00143, 2018.

Jun-Yan Zhu, Taesung Park, Phillip Isola, and Alexei A. Efros. Unpaired image-to-image translation using cycle-consistent adversarial networks. 2017 IEEE International Conference on Computer Vision (ICCV) , pages 2242-2251, 2017.

Fangyu Zou, Li Shen, Zequn Jie, Weizhong Zhang, and Wei Liu. A sufficient condition for convergences of adam and rmsprop. 2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , pages 11119-11127, 2018.

A Additional related work

In this section, we discuss additional related work on variants of Adam and variance reduction methods.

Variants of Adam. After Reddi et al. [37] pointed out the non-convergence issue with Adam, various variants of Adam that can be proved to converge were proposed [56; 18; 6; 5; 31; 54]. For example, AMSGrad [37] and AdaFom [6] modify the second order momentum so that it is non-decreasing. AdaBound [31] explicitly imposes upper and lower bounds on the second order momentum so that the stepsize is also bounded. AdaShift [54] uses a new estimate of the second order momentum to correct the bias. There are also some works [53; 18; 20] that provide convergence guarantees of these variants. One closely related work to ours is [47], which considers a variancereduced version of Adam by combining Adam and SVRG [22]. However, they assume bounded gradients and can only get an asymptotic convergence in the non-convex setting.

Variance reduction methods. The technique of variance reduction was introduced to accelerate convex optimization in the finite-sum setting [39; 22; 41; 32; 11]. Later, many works studied variance-reduced methods in the non-convex setting and obtained improved convergence rates for standard smooth functions. For example, SVRG and SCSG improve the O ( ϵ -4 ) gradient complexity of stochastic gradient descent (SGD) to O ( ϵ -10 / 3 ) [2; 36; 25]. Many new variance reduction methods [15; 43; 29; 27; 9; 30] were later proposed to further improve the complexity to O ( ϵ -3 ) , which is optimal and matches the lower bound in [3]. Recently, [38; 7] obtained the O ( ϵ -3 ) complexity for the more general ( L 0 , L 1 ) smooth functions. Our variance-reduced Adam is motivated by the STORM algorithm proposed by [9], where an additional term is added in the momentum update to correct the bias and reduce the variance.

B Proof sketches of informal lemmas in Section 5.3

In this section, we provide the proof sketches of the informal lemmas in Section 5.3. We focus on illustrating the ideas rather than rigorous proof details. Please see Appendix E for more rigorous and detailed proofs of Adam in the stochastic setting.

Proof Sketch of Lemma 5.4. By the definition of τ , for all t < τ , we have f ( x t ) -f ∗ ≤ F which implies ∥∇ f ( x t ) ∥ ≤ G . Then from the update rule (18) in Proposition E.1 provided later in Appendix E, it is easy to verify ˆ v t ⪯ G 2 since ˆ v t is a convex combination of { ( ∇ f ( x s )) 2 } s ≤ t . Let h t := η/ ( √ ˆ v t + λ ) be the stepsize vector and denote H t := diag ( h t ) . We know

As discussed in Section 5.2, when η is small enough, we can apply Corollary 5.2 to obtain

where in the first (approximate) inequality we ignore the second order term 1 2 L ∥ x t +1 -x t ∥ 2 ∝ η 2 in Corollary 5.2 for small enough η ; the equality applies the update rule x t +1 -x t = -H t ˆ m t = -H t ( ∇ f ( x t ) + ϵ t ) ; in the second inequality we use 2 a ⊤ Ab ≤ ∥ a ∥ 2 A + ∥ b ∥ 2 A for any PSD matrix A and vectors a and b ; and the last inequality is due to (7).

Proof Sketch of Lemma 5.5. By the update rule (18) in Proposition E.1, we have

For small enough η , we can apply Corollary 5.2 to get

where the second inequality is due to L = O ( G ρ ) and ∥ x t +1 -x t ∥ = O ( η ) ∥ ˆ m t ∥ ; and the last inequality uses ˆ m t = ∇ f ( x t ) + ϵ t and Young's inequality ∥ a + b ∥ 2 ≤ 2 ∥ a ∥ 2 +2 ∥ b ∥ 2 . Therefore,

where the first inequality uses (8) and Young's inequality ∥ a + b ∥ 2 ≤ (1 + u ) ∥ a ∥ 2 +(1+1 /u ) ∥ b ∥ 2 for any u > 0 ; the second inequality uses (1 -α t +1 )(1 + α t +1 / 2) ≤ 1 -α t +1 / 2 and (9); and in the last inequality we use β ≤ α t +1 and choose β = Θ( ηG ρ +1 / 2 ) which implies O ( η 2 G 2 ρ /α t +1 ) ≤ λβ 16 G ≤ β/ 4 .

C Probabilistic lemmas

In this section, we state several well-known and useful probabilistic lemmas without proof.

Lemma C.1 (Azuma-Hoeffding inequality) . Let { Z t } t ≥ 1 be a martingale with respect to a filtration {F t } t ≥ 0 . Assume that | Z t -Z t -1 | ≤ c t almost surely for all t ≥ 0 . Then for any fixed T , with probability at least 1 -δ ,

Lemma C.2 (Optional Stopping Theorem) . Let { Z t } t ≥ 1 be a martingale with respect to a filtration {F t } t ≥ 0 . Let τ be a bounded stopping time with respect to the same filtration. Then we have E [ Z τ ] = E [ Z 0 ] .

D Proofs related to ( ρ, L 0 , L ρ ) smoothness

In this section, we provide proofs related to ( ρ, L 0 , L ρ ) smoothness. In what follows, we first provide a formal proposition in Appendix D.1 showing that univariate rational functions and double exponential functions are ( ρ, L 0 , L ρ ) smooth with ρ < 2 , as we claimed in Section 3.2.1, and then provide the proofs of Lemma 3.4, Lemma 5.1, and Corollary 5.2 in Appendix D.2, D.3 and D.4 respectively.

D.1 Examples

Proposition D.1. Any univariate rational function P ( x ) /Q ( x ) , where P, Q are two polynomials, and any double exponential function a ( b x ) , where a, b > 1 , are ( ρ, L 0 , L ρ ) smooth with 1 < ρ < 2 . However, they are not necessarily ( L 0 , L 1 ) smooth.

Proof of Proposition D.1. We prove the proposition in the following four parts:

Univariate rational functions are ( ρ, L 0 , L ρ ) smooth with 1 < ρ < 2 . Let f ( x ) = P ( x ) /Q ( x ) where P and Q are two polynomials. Then the partial fractional decomposition of f ( x ) is given by

̸

where w ( x ) is a polynomial, A ir , B ir , C ir , a i , b i , c i are all real constants satisfying b 2 i -4 c i < 0 for each 1 ≤ i ≤ n which implies x 2 + b i x + c i > 0 for all x ∈ R . Assume A ij i = 0 without loss of

generality. Then we know f has only finite singular points { a i } 1 ≤ i ≤ m and has continuous first and second order derivatives at all other points. To simplify notation, denote

Then we have f ( x ) = w ( x ) + ∑ m i =1 ∑ j i r =1 p ir ( x ) + ∑ n i =1 ∑ k i r =1 q ir ( x ) . For any 3 / 2 < ρ < 2 , we know that ρ > r +2 r +1 for any r ≥ 1 . Then we can show that

where the first equality is because one can easily verify that the first and second order derivatives of p ij i dominate those of all other terms when x goes to a i , and the second equality is because ∣ ∣ p ′ ij i ( x ) ∣ ∣ ρ = O ( ( x -a i ) -ρ ( j i +1) ) , ∣ ∣ p ′′ ij i ( x ) ∣ ∣ = O ( ( x -a i ) -( j i +2) ) , and ρ ( j i +1) > j i +2 (here we assume j i ≥ 1 since otherwise there is no need to prove (10) for i ). Note that (10) implies that, for any L ρ > 0 , there exists δ i > 0 such that

Similarly, one can show lim x →∞ | f ′ ( x ) | ρ | f ′′ ( x ) | = ∞ , which implies there exists M > 0 such that

We know B is a compact set and therefore the continuous function f ′′ is bounded within B , i.e., there exists some constant L 0 > 0 such that

Combining (11), (12), and (13), we have shown

which completes the proof of the first part.

Rational functions are not necessarily ( L 0 , L 1 ) smooth. Consider the ration function f ( x ) = 1 /x . Then we know that f ′ ( x ) = -1 /x 2 and f ′′ ( x ) = 2 /x 3 . Note that for any 0 < x ≤ min { ( L 0 +1) -1 / 3 , ( L 1 +1) -1 } , we have

which shows f is not ( L 0 , L 1 ) smooth for any L 0 , L 1 ≥ 0 .

Double exponential functions are ( ρ, L 0 , L ρ ) smooth with 1 < ρ < 2 . Let f ( x ) = a ( b x ) , where a, b > 1 , be a double exponential function. Then we know that

For any ρ > 1 , we have

where the first equality is a direct calculation; the second equality uses change of variable y = b x ; and the last equality is because exponential function grows faster than linear function. Then we complete the proof following a similar argument to that in Part 1.

Double exponential functions are not necessarily ( L 0 , L 1 ) smooth. Consider the double exponential function f ( x ) = e ( e x ) . Then we have

For any x ≥ max { log( L 0 +1) , log( L 1 +1) } , we can show that

which shows f is not ( L 0 , L 1 ) smooth for any L 0 , L 1 ≥ 0 .

Define

D.2 Proof of Lemma 3.4

Before proving Lemma 3.4, we need the following lemma that generalizes (a special case of) Grönwall's inequality.

Lemma D.2. Let u : [ a, b ] → [0 , ∞ ) and ℓ : [0 , ∞ ) → (0 , ∞ ) be two continuous functions. Suppose u ′ ( t ) ≤ ℓ ( u ( t )) for all t ∈ ( a, b ) . Denote function ϕ ( w ) := ∫ 1 ℓ ( w ) dw . We have for all t ∈ [ a, b ] ,

Proof of Lemma D.2. First, by definition, we know that ϕ is increasing since ϕ ′ = 1 ℓ > 0 . Let function v be the solution of the following differential equation

It is straightforward to verify that the solution to (14) satisfies

Then it suffices to show ϕ ( u ( t )) ≤ ϕ ( v ( t )) ∀ t ∈ [ a, b ] . Note that

where the inequality is because u ′ ( t ) ≤ ℓ ( u ( t )) by the assumption of this lemma and v ′ ( t ) = ℓ ( v ( t )) by (14). Since ϕ ( u ( a )) -ϕ ( v ( a )) = 0 , we know for all t ∈ [ a, b ] , ϕ ( u ( t )) ≤ ϕ ( v ( t )) .

With Lemma D.2, one can bound the gradient norm within a small enough neighborhood of a given point as in the following lemma.

Lemma D.3. Suppose f is ( ρ, L 0 , L ρ ) smooth for some ρ, ρ, L 0 , L ρ ≥ 0 . For any a > 0 and points x, y ∈ dom( f ) satisfying ∥ y -x ∥ ≤ a L 0 + L ρ ( ∥∇ f ( x ) ∥ + a ) ρ , we have

Proof of Lemma D.3. Denote functions z ( t ) := (1 -t ) x + ty and u ( t ) := ∥∇ f ( z ( t )) ∥ for 0 ≤ t ≤ 1 . Note that for any 0 ≤ t ≤ s ≤ 1 , by triangle inequality,

We know that u ( t ) = ∥∇ f ( z ( t )) ∥ is differentiable since f is second order differentiable 2 . Then we have

Let ϕ ( w ) := ∫ w 0 1 ( L 0 + L ρ v ρ ) ∥ y -x ∥ dv . By Lemma D.2, we know that

Denote ψ ( w ) := ∫ w 0 1 ( L 0 + L ρ v ρ ) dv = ϕ ( w ) · ∥ y -x ∥ . We have

=

ψ

(

∥∇

f

(

x

)

∥

+

a

)

Since ψ is increasing, we have ∥∇ f ( y ) ∥ ≤ ∥∇ f ( x ) ∥ + a .

With Lemma D.3, we are ready to prove Lemma 3.4.

Proof of Lemma 3.4. Denote z ( t ) := (1 -t ) x + ty for some y ∈ R d satisfying ∥ y -x ∥ ≤ a L 0 + L ρ ( ∥∇ f ( x ) ∥ + a ) ρ . We first show y ∈ dom( f ) by contradiction. Suppose y / ∈ dom( f ) , let us define t b := inf { 0 ≤ t ≤ 1 | z ( t ) / ∈ X} and z b := z ( t b ) . Then we know z b is a boundary point of X . Since f is a closed function with an open domain, we have

On the other hand, by the definition of t b, we know z ( t ) ∈ X for every 0 ≤ t < t b. Then by Lemma D.3, for all 0 ≤ t < t b, we have ∥∇ f ( z ( t )) ∥ ≤ ∥∇ f ( x ) ∥ + a . Therefore for all 0 ≤ t < t b

which contradicts with (15). Therefore we have shown y ∈ dom( f ) . We have

where the last inequality is due to Lemma D.3.

D.3 Proof of Lemma 5.1

Proof of Lemma 5.1. Denote G := ∥∇ f ( x ) ∥ and L := 3 L 0 + 4 L ρ G ρ . Let y := x -1 2 L ∇ f ( x ) . Then we have

where the inequality can be easily verified considering both cases of G ≤ ( L 0 /L ρ ) 1 /ρ and G ≥ ( L 0 /L ρ ) 1 /ρ . Then based on Corollary 5.2, we have y ∈ dom( f ) and

which completes the proof.

D.4 Proof of Corollary 5.2

Proof of Corollary 5.2. First, Lemma 3.4 states that for any a > 0 ,

If ∥∇ f ( x ) ∥ ≤ G , we choose a = max { G, ( L 0 /L ρ ) 1 /ρ } . Then it is straightforward to verify that

Therefore we have shown for any x, y satisfying ∥ y -x ∥ ≤ r ,

Next, let z ( t ) := (1 -t ) x + ty for 0 ≤ t ≤ 1 . We know

where the inequality is due to (16).

E Convergence analysis of Adam

In this section, we provide detailed convergence analysis of Adam. We will focus on proving Theorem 4.1 under the bounded noise assumption (Assumption 3) in most parts of this section except Appendix E.6 where we will show how to generalize the results to noise with sub-Gaussian norm (Assumption 4) and provide the proof of Theorem 4.2.

For completeness, we repeat some important technical definitions here. First, we define

as the deviation of the re-scaled momentum from the actual gradient. Given a large enough constant G defined in Theorem 4.1, denoting F = G 2 3(3 L 0 +4 L ρ G ρ ) , we formally define the stopping time τ as

i.e., τ is the first time when the sub-optimality gap is strictly greater than F , truncated at T +1 to make sure it is bounded in order to apply Lemma C.2. Based on Lemma 5.1 and the discussions below it, we know that if t < τ , we have both f ( x t ) -f ∗ ≤ F and ∥∇ f ( x t ) ∥ ≤ G . It is clear to see that τ is a stopping time 3 with respect to { ξ t } t ≥ 1 because the event { τ ≥ t } is a function of { ξ s } s<t and independent of { ξ s } s ≥ t . Next, let

be the stepsize vector and H t := diag ( h t ) be the diagonal stepsize matrix. Then the update rule can be written as

Finally, as in Corollary 5.2 and Lemma 5.3, we define the following constants.

E.1 Equivalent update rule of Adam

The bias correction steps in Lines 7-8 make Algorithm 1 a bit complicated. In the following proposition, we provide an equivalent yet simpler update rule of Adam.

Proposition E.1. Denote α t = β 1 -(1 -β ) t and α sq t = β sq 1 -(1 -β sq ) t . Then the update rule in Algorithm 1 is equivalent to

where initially we set ˆ m 1 = ∇ f ( x 1 , ξ 1 ) and ˆ v 1 = ( ∇ f ( x 1 , ξ 1 )) 2 . Note that since 1 -α 1 = 1 -α sq 1 = 0 , there is no need to define ˆ m 0 and ˆ v 0 .

Proof of Proposition E.1. Denote Z t = 1 -(1 -β ) t . Then we know α t = β/Z t and m t = Z t ˆ m t . By the momentum update rule in Algorithm 1, we have

Note that Z t satisfies the following property

Next, we verify the initial condition. By Algorithm 1, since we set m 0 = 0 , we have m 1 = β ∇ f ( x 1 , ξ 1 ) . Therefore we have ˆ m 1 = m 1 /Z 1 = ∇ f ( x 1 , ξ 1 ) since Z 1 = β . Then the proof is completed by applying the same analysis on v t and ˆ v t .

E.2 Useful lemmas for Adam

In this section, we list several useful lemmas for the convergence analysis. Their proofs are all deferred in Appendix E.5.

First note that when t < τ , all the quantities in the algorithm are well bounded. In particular, we have the following lemma.

Next, we provide a useful lemma regarding the time-dependent re-scaled momentum parameters in (18).

In the next lemma, we provide an almost sure bound on ϵ t in order to apply Azuma-Hoeffding inequality (Lemma C.1).

Finally, the following lemma hides messy calculations and will be useful in the contradiction argument.

Lemma E.5. Denote

Under the parameter choices in either Theorem 4.1 or Theorem 4.2, we have I 1 ≤ I 2 and I 1 /T ≤ ϵ 2 .

Then we have

E.3 Proof of Theorem 4.1

Before proving the main theorems, several important lemmas are needed. First, we provide a descent lemma for Adam.

Lemma E.6. If t < τ , choosing G ≥ σ + λ and η ≤ min { r D , λ 6 L } , we have

Proof of Lemma E.6. By Lemma E.2, we have if t < τ ,

Since we choose η ≤ r D , by Lemma 5.3, we have ∥ x t +1 -x t ∥ ≤ r if t < τ . Then we can apply Corollary 5.2 to show that for any t < τ ,

where the second inequality uses (17) and (19); the third inequality is due to Young's inequality a ⊤ Ab ≤ 1 3 ∥ a ∥ 2 A + 3 4 ∥ b ∥ 2 A and ∥ a + b ∥ 2 A ≤ 2 ∥ a ∥ 2 A +2 ∥ b ∥ A for any PSD matrix A ; the second last inequality uses η ≤ λ 6 L ; and the last inequality is due to (19).

The following lemma bounds the sum of the error term ∥ ϵ t ∥ 2 before the stopping time τ . Since its proof is complicated, we defer it in Appendix E.4.

Lemma E.7. If G ≥ 2 σ and η ≤ min { r D , λ 3 / 2 β 6 L √ G , σβ DL } , with probability 1 -δ

Combining Lemma E.6 and Lemma E.7, we obtain the following useful lemma, which simultaneously bounds f ( x t ) -f ∗ and ∑ τ -1 t =1 ∥∇ f ( x t ) ∥ 2 .

Lemma E.8. If G ≥ 2 max { λ, σ } and η ≤ min { r D , λ 3 / 2 β 6 L √ G , σβ DL } , then with probability at least 1 -δ ,

Proof of Lemma E.8. By telescoping, Lemma E.6 implies

Lemma E.7 could be written as

With Lemma E.8, we are ready to complete the contradiction argument and the convergence analysis. Below we provide the proof of Theorem 4.1.

Proof of Theorem 4.1. According to Lemma E.8, there exists some event E with P ( E ) ≥ 1 -δ , such that conditioned on E , we have

By the definition of τ , if τ ≤ T , we have

Based on Lemma E.5, we have I 1 ≤ I 2 , which leads to a contradiction. Therefore, we must have τ = T +1 conditioned on E . Then, Lemma E.8 also implies that under E ,

where the last inequality is due to Lemma E.5.

E.4 Proof of Lemma E.7

In order to prove Lemma E.7, we need the following several lemmas.

Lemma E.9. Denote γ t -1 = (1 -α t )( ϵ t -1 + ∇ f ( x t -1 ) - ∇ f ( x t )) . If G ≥ 2 σ and η ≤ min { r D , λ 3 / 2 β 6 L √ G } , we have for every 2 ≤ t ≤ τ ,

Proof of Lemma E.9. According to the update rule (18), we have

Since we choose η ≤ r D , by Lemma 5.3, we have ∥ x t -x t -1 ∥ ≤ r if t ≤ τ . Therefore by Corollary 5.2, for any 2 ≤ t ≤ τ ,

Therefore

where the first inequality uses Young's inequality ∥ a + b ∥ 2 ≤ (1 + u ) ∥ a ∥ 2 +(1+1 /u ) ∥ b ∥ 2 for any u > 0 ; the second inequality is due to

the third inequality uses (24) and Young's inequality; and in the last inequality we choose η ≤ λ 3 / 2 β 6 L √ G , which implies 2 η 2 L 2 λ 2 β ≤ λβ 16 G ≤ β 2 ≤ α t 2 . Then by (23), we have

Lemma E.10. Denote γ t -1 = (1 -α t )( ϵ t -1 + ∇ f ( x t -1 ) - ∇ f ( x t )) . If G ≥ 2 σ and η ≤ min { r D , σβ DL } , with probability 1 -δ ,

Proof of Lemma E.10. First note that

Since τ is a stopping time, we know that 1 τ ≥ t is a function of { ξ s } s<t . Also, by definition, we know γ t -1 is a function of { ξ s } s<t . Then, denoting

we know that E t -1 [ X t ] = 0 , which implies { X t } t ≤ T is a martingale difference sequence. Also, by Assumption 3 and Lemma E.4, we can show that for all 2 ≤ t ≤ T ,

Then by the Azuma-Hoeffding inequality (Lemma C.1), we have with probability at least 1 -δ ,

where in the last inequality we use Lemma E.3.

Then we are ready to prove Lemma E.7.

Proof of Lemma E.7. By Lemma E.9, we have for every 2 ≤ t ≤ τ ,

Taking a summation over t from 2 to τ , we have

where the first inequality uses Lemma E.10; and the second inequality uses Lemma E.3 and ∥ ϵ 1 ∥ 2 = ∥∇ f ( x 1 , ξ 1 ) -∇ f ( x 1 ) ∥ 2 ≤ σ 2 . Then we complete the proof by multiplying both sides by 2 /β .

E.5 Omitted proofs for Adam

In this section, we provide all the omitted proofs for Adam including those of Lemma 5.3 and all the lemmas in Appendix E.2.

Proof of Lemma 5.3. According to Lemma E.2, if t < τ ,

Proof of Lemma E.2. By definition of τ , we have ∥∇ f ( x t ) ∥ ≤ G if t < τ . Then Assumption 3 directly implies ∥∇ f ( x t , ξ t ) ∥ ≤ G + σ . ∥ ˆ m t ∥ can be bounded by a standard induction argument as follows. First note that ∥ ˆ m 1 ∥ = ∥∇ f ( x 1 , ξ 1 ) ∥ ≤ G + σ . Supposing ∥ ˆ m k -1 ∥ ≤ G + σ for some k < τ , then we have

Then we can show ˆ v t ⪯ ( G + σ ) 2 in a similar way noting that ( ∇ f ( x t , ξ t )) 2 ⪯ ∥∇ f ( x t , ξ t ) ∥ 2 ≤ ( G + σ ) 2 . Given the bound on ˆ v t , it is straight forward to bound the stepsize h t .

Proof of Lemma E.3. First, when t ≥ 1 /β , we have (1 -β ) t ≤ 1 /e . Therefore,

Next, note that when t < 1 /β , we have (1 -β ) t ≤ 1 -1 2 βt . Then we have

Therefore we have ∑ T t =2 α 2 t ≤ 3(1 + β 2 T ) .

Proof of Lemma E.4. We prove ∥ ϵ t ∥ ≤ 2 σ for all t ≤ τ by induction. First, note that for t = 1 , we have

Now suppose ∥ ϵ t -1 ∥ ≤ 2 σ for some 2 ≤ t ≤ τ . According to the update rule (18), we have

which implies

Since we choose η ≤ r D , by Lemma 5.3, we have ∥ x t -x t -1 ∥ ≤ ηD ≤ r if t ≤ τ . Therefore by Corollary 5.2, we have for any 2 ≤ t ≤ τ ,

where the last inequality uses the choice of η and β ≤ α t . Therefore we have ∥ ϵ t ∥ ≤ 2 σ which completes the induction. Then it is straight forward to show

Proof of Lemma E.5. We first list all the related parameter choices below for convenience.

We will show I 1 /I 2 ≤ 1 first. Note that if denoting W = 3 L λG 2 , we have

Below are some facts that can be easily verified given the parameter choices.

(a) By the choice of G , we have G 2 ≥ 6∆ 1 (3 L 0 + 4 L ρ G ρ ) = 6∆ 1 L for large enough C 1 , which implies W ≤ 1 2∆ 1 λ .

(b) By the choice of T , we have ηβT ≤ η β + C 2 ∆ 1 Gβ ϵ 2 .

(d) By the choice of η , we have η/β ≤ c 2 σλ LG √ ι , which implies Wσ 2 √ ι · η β ≤ 3 c 2 σ 3 G 3 ≤ 1 200 for small enough c 2 .

(e) By the choice of β and (a), we have Wσ 2 ∆ 1 G √ ιβ ϵ 2 ≤ σ 2 G √ ιβ 2 λϵ 2 ≤ 1 100 C 2 for small enough c 1 . Therefore,

where the first inequality is due to Facts (a-c); the second inequality uses σ ≤ G , ι ≥ 1 , and √ a + b ≤ √ a + √ b for a, b ≥ 0 ; and the last inequality is due to Facts (d-e).

Next, we will show I 1 /T ≤ ϵ 2 . We have

where in the first inequality we use T ≥ C 2 ∆ 1 G ηϵ 2 and √ a + b ≤ √ a + √ b for a, b ≥ 0 ; the second inequality uses T ≥ 1 β 2 ; the second equality uses the parameter choice of β ; and in the last inequality we choose a large enough C 2 and small enough c 1 .

E.6 Proof of Theorem 4.2

Proof of Theorem 4.2. We define stopping time τ as follows

Then it is straightforward to verify that τ 1 , τ 2 , τ are all stopping times.

Since we want to show P ( τ ≤ T ) is small, noting that { τ ≤ T } = { τ = τ 1 ≤ T } ∪ { τ = τ 2 ≤ T } , it suffices to bound both P ( τ = τ 1 ≤ T ) and P ( τ = τ 2 ≤ T ) .

First, we know that

where the fourth inequality uses Assumption 4; and the last inequality uses σ = R √ 2 log(4 T/δ ) .

Next, if τ = τ 1 ≤ T , by definition, we have f ( x τ ) -f ∗ > F , or equivalently,

On the other hand, since for any t < τ , under the new definition of τ , we still have

Then we know that Lemma E.8 still holds because all of its requirements are still satisfied, i.e., there exists some event E with P ( E ) ≤ δ/ 2 , such that under its complement E c ,

By Lemma E.5, we know I 1 ≤ I 2 , which suggests that E c ∩ { τ = τ 1 ≤ T } = ∅ , i.e., { τ = τ 1 ≤ T } ⊂ E . Then we can show

Therefore,

and under the event E c ∩ { τ = T +1 } , we have τ = T +1 and

where the last inequality is due to Lemma E.5.

F Convergence analysis of VRAdam

In this section, we provide detailed convergence analysis of VRAdam and prove Theorem 6.2. To do that, we first provide some technical definitions 4 . Denote

as the deviation of the momentum from the actual gradient. From the update rule in Algorithm 2, we can write

where we define

Let G be the constant defined in Theorem 6.2 and denote F := G 2 3(3 L 0 +4 L ρ G ρ ) . We define the following stopping times as discussed in Section 6.1.

It is straight forward to verify that τ 1 , τ 2 , τ are all stopping times. Then if t < τ , we have

Then we can also bound the update ∥ x t +1 -x t ∥ ≤ ηD where D = 2 G/λ if t < τ (see Lemma F.3 for the details). Finally, we consider the same definition of r and L as those for Adam. Specifically,

F.1 Useful lemmas

We first list several useful lemmas in this section without proofs. Their proofs are deferred later in Appendix F.3.

To start with, we provide a lemma on the local smoothness of each component function f ( · , ξ ) when the gradient of the objective function f is bounded.

Lemma F.1. For any constant G ≥ σ and two points x ∈ dom( f ) , y ∈ R d such that ∥∇ f ( x ) ∥ ≤ G and ∥ y -x ∥ ≤ r/ 2 , we have y ∈ dom( f ) and

where r and L are defined in (27) .

With the new definition of stopping time τ in (26), all the quantities in Algorithm 2 are well bounded before τ . In particular, the following lemma holds.

Next, we provide the following lemma which bounds the update at each step before τ .

Lemma F.3. if t < τ , ∥ x t +1 -x t ∥ ≤ ηD where D = 2 G/λ .

The following lemma bounds ∥ W t ∥ when t ≤ τ .

Lemma F.4. If t ≤ τ , G ≥ 2 σ , and η ≤ r 2 D ,

Finally, we present some inequalities regarding the parameter choices, which will simplify the calculations later.

Lemma F.5. Under the parameter choices in Theorem 6.2, we have

F.2 Proof of Theorem 6.2

Before proving the theorem, we will need to present several important lemmas. First, note that the descent lemma still holds for VRAdam.

Lemma F.6. If t < τ , choosing G ≥ σ + λ and η ≤ min { r 2 D , λ 6 L } , we have

The proof is essentially the same as that of Lemma E.6.

Proof of Lemma F.6.

Lemma F.7. Choose G ≥ max { 2 σ, 2 λ } , S 1 ≥ 1 2 β 2 T , and η ≤ min { r 2 D , λ 3 / 2 40 L √ β G } . We have

Proof of Lemma F.7. By Lemma F.4, we have

where in the second inequality we choose η ≤ λ 3 / 2 40 L √ β G . Therefore, noting that λβ 16 G ≤ β/ 2 , by (25), we have

Taking a summation over 2 ≤ t ≤ τ and re-arranging the terms, we get

Taking expectations on both sides, noting that

by the Optional Stopping Theorem (Lemma C.2), we have

where in the second inequality we choose S 1 ≥ 1 2 β 2 T which implies E [ ∥ ϵ 1 ∥ 2 ] ≤ σ 2 /S 1 ≤ 2 σ 2 β 2 T .

Lemma F.8. Under the parameter choices in Theorem 6.2, we have

Proof of Lemma F.8. First note that according to Lemma F.5, it is straight forward to verify that the parameter choices in Theorem 6.2 satisfy the requirements in Lemma F.6 and Lemma F.7. Then by Lemma F.6, if t < τ ,

Taking a summation over 1 ≤ t < τ , re-arranging terms, multiplying both sides by 8 G η , and taking an expection, we get

where the last inequality is due to Lemma F.5. Then (28) + (29) gives

which completes the proof.

With all the above lemmas, we are ready to prove the theorem.

Proof of Theorem 6.2. First note that according to Lemma F.5, it is straight forward to verify that the parameter choices in Theorem 6.2 satisfy the requirements in all the lemmas for VRAdam.

Then, first note that if τ = τ 1 ≤ T , we know f ( x τ ) -f ∗ > F by the definition of τ . Therefore,

where the second inequality uses Markov's inequality; the third inequality is by Lemma F.8; and the last inequality is due to Lemma F.5.

Similarly, if τ 2 = τ ≤ T , we know ∥ ϵ τ ∥ > G . We have

where the second inequality uses Markov's inequliaty; the third inequality is by Lemma F.8; and the last inequality is due to Lemma F.5. where the last inequality is due to Lemma F.5. Therefore,

where the last inequality is due to P ( τ = T +1) = 1 -P ( τ ≤ T ) ≥ 1 -δ/ 2 ≥ 1 / 2 . Then we can get

Let F := { 1 T ∑ T t =1 ∥∇ f ( x t ) ∥ 2 > ϵ 2 } be the event of not converging to stationary points. By Markov's inequality, we have

Therefore,

i.e., with probability at least 1 -δ , we have both τ = T +1 and 1 T ∑ T t =1 ∥∇ f ( x t ) ∥ 2 ≤ ϵ 2 .

F.3 Proofs of lemmas in Appendix F.1

Proof of Lemma F.1. This lemma is a direct corollary of Corollary 5.2. Note that by Assumption 6, we have ∥∇ f ( x, ξ ) ∥ ≤ G + σ ≤ 2 G . Hence, when computing the locality size and smoothness constant for the component function f ( · , ξ ) , we need to replace the constant G in Corollary 5.2 with 2 G , that is why we get a smaller locality size of r/ 2 and a larger smoothness constant of 4 L .

Proof of Lemma F.2. The bound on ∥ m t ∥ is by the definition of τ in (26). All other quantities for VRAdam are defined in the same way as those in Adam (Algorithm 1), so they have the same upper bounds as in Lemma E.2.

where

Then we can bound

where the second inequality uses Lemma F.1; and the last inequality is due to ∥ x t -x t -1 ∥ ≤ η ∥ m t -1 ∥ /λ ≤ η ( ∥∇ f ( x t -1 ) ∥ + ∥ ϵ t -1 ∥ ) /λ . Then, we have

Proof of Lemma F.5. These inequalities can be obtained by direct calculations.

where the first inequality uses the update rule in Algorithm 2 and h t ⪯ η/λ by Lemma F.2; the second inequality is again due to Lemma F.2.

Proof of Lemma F.4. By the definition of W t , it is easy to verify that

Peer reviews:
comments: SUMMARY:
The paper removes the Lipschit gradient assumption for the adaptive SGD (ASGD), making the ASGD broader to wide applications. Under the weak assumption, the authors still proved the optimal convergence rate. Moreover, they propose a variance-reduced version  with an accelerated complexity. The results are interesting and novel.

STRENGTHS:
1. They further relaxed the ($L_0,L_1$)  Lipschit gradient assumption, which is more relastic.

2. They  develop a new analysis to show that ASGD can be convergent under the weak Lipschit gradient assumption but still with the same rate as existing works on ASGD. The results and techniques are novel.

3. A variance-reduced version of  ASGD is proposed with provable acceleration.

4. The authors show that the rates of ASGD and its variance-reduced version is dimension-free under certain cases.

WEAKNESSES:
1. The proof of variance-reduced version is interesting but without any numerical demonstrations. If the authors proposed an accelerated algorithm, it is convening to present some numerics because the theory is not the focus now.

2. In the contribution part, the authors claim that they do not assume bounded gradients. This is because you assume the a.s. bounded noise. which has been studied by previous work [Li and Orabona, 2019].

QUESTIONS:
There are very recent works on using bounded variance noises [Matthew Faw et al. COLT 2022, Bohang Wang et al. COLT 2023]. Can the results in this paper be extended to that case?

LIMITATIONS:
YES
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-06-22

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a new proof strategy for the convergence of Adam in the\nnon-convex setting. The new analysis relaxes the typical assumptions\nin the following ways: 1) it assumes relaxed smoothness, where the norm of the Hessian\ngrows sub-quadratically with the gradient norm; 2) it does not require\nbounded gradients. With a deterministic gradient oracle, the obtained rate is $O(1/\\epsilon^2)$.\nFor the stochastic setting, the author shows a convergence rate of $O(1/\\epsilon^4)$\nwith high probability.\n\nFurthermore, the author proposes a variance-reduced version of Adam and proves\na rate of $O(1/\\epsilon^3)$. The above rates in stochastic setting is dimension-dependent\nand can be dimension-free if the Hessian norm is sub-affine on the gradient norm.\n\nThe main limitation is that although the bounded gradient assumption is removed,\nthe bounded noise is still needed.\n\nSTRENGTHS:\n1. The paper is well-written and easy to follow. Specifically, the technical\nchallenges and the intuition behind the new analysis are discussed comprehensively.\n2. Removing the bounded gradient assumption is generally challenging. Although\nthere are a few works for AdaGrad that address this, existing work targeting\nthe removal of this assumption for Adam only converges to a neighborhood of\nstationary points.\n3. The paper considers the smoothness condition where the norm of the\nHessian is bounded by a sub-quadratic function of the gradient norm, which\nis more general than $(L_0, L_1)$-smoothness.\n\nWEAKNESSES:\n1. Although the work relaxes the bounded stochastic gradient assumption\n   commonly required in the analysis of Adam, assuming almost surely bounded\n   noise is still somewhat strong and impractical. Since the proof technique,\n   the contradiction argument, heavily relies on this assumption, it is\n   unclear whether it can be used in more realistic settings.\n2. It is difficult to assess the usefulness of the proposed variance reduction method\nsince the convergence result does not improve upon previous methods and there\nis no experimental evidence.\n\n\nMinor:\n\nThe result presented for the deterministic setting is not formal. It would be\nbetter to have a formal theorem, at least in the appendix.\n\nQUESTIONS:\n1. In Theorem 4.1, the same value of $\\beta_{sp}$ and $\\beta$ are selected,\nwhich differ from conventional choice of $(1 - \\beta)^2 < 1 - \\beta_{sp}$.\nCould the author please explain the reason behind this?\n2. Why do we have a dependence on $d$ in the convergence result when $\\rho \\geq 1$\nand suddenly become dimension-free when $\\rho < 1$? Also, some results for Adam\nassume the infinity norm of the gradient is upper bounded by a constant [1], which\nnatural leads to dimension dependence. It seems that by assuming bounded 2-norm, the\nresults become dimension-free [2, 3]. Could the author provide further elaborate on this matter?\n3. It seems that if we have access to problem-dependent parameters and the initial conditions (e.g., the gradient norm at the first iteration), SGD can also converge with the optimal rate [4]. What is the benefit of adaptive methods (if they also require tuning) compared to SGD in the more generalized smoothness setting?\n\nReferences:\n1. Défossez, A., Bottou, L., Bach, F., & Usunier, N. (2020). A simple convergence proof of adam and adagrad. arXiv preprint arXiv:2003.02395.\n2. Chen, X., Liu, S., Sun, R., & Hong, M. (2018). On the convergence of a class of adam-type algorithms for non-convex optimization. arXiv preprint arXiv:1808.02941.\n3. Guo, Z., Xu, Y., Yin, W., Jin, R., & Yang, T. (2021). A novel convergence analysis for algorithms of the adam family and beyond. arXiv preprint arXiv:2104.14840.\n4. Li, H., Qian, J., Tian, Y., Rakhlin, A., & Jadbabaie, A. (2023). Convex and Non-Convex Optimization under Generalized Smoothness. arXiv preprint arXiv:2306.01264.\n\nLIMITATIONS:\nThe limitations are addressed in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the convergence of Adam over non-convex objectives. To begin with, this paper proposes a new non-uniform smoothness condition called $(\\rho, L_0,L_1)$ smoothness condition, which generalizes $(L_0,L_1)$ smoothness condition proposed in [Zhang et al. 2019]. The authors then prove the high-probability convergence rate $\\mathcal{O}(1/\\sqrt{T})$ of Adam under $(\\rho, L_0,L_1)$ smoothness condition and affine noise variance assumption. The authors then propose Variance Reduced Adam (VRAdam) by combining the gradient estimation of STORM into Adam, and derive the convergence rate $\\mathcal{O}(1/\\sqrt[3]{T^2})$ of VRAdam.\n\nSTRENGTHS:\n1. This paper provides the first $O(1/\\sqrt{T})$ convergence rate of Adam without the bounded gradient assumption.\n\n2. The constructed stopping time is interesting and can be of independent interest.\n\nWEAKNESSES:\n1. Although this paper provides the first $O(1/\\sqrt{T})$ convergence rate of Adam without the bounded gradient assumption, the analysis is somewhat restricted because it requires a non-zero $\\lambda$ and the convergence rate has a polynomial dependence over $1/\\lambda$. However, in the previous analyses of Adam including [Défossez et al., 2020; Zhang et al., 2022, Wang et al., 2022], it is allowed that $\\lambda=0$ or the rate has a logarithmic dependence over $1/\\lambda$. I treat this as a weakness of this paper because, in practice, $\\lambda$ is set to very close to $0$ (for example, $10^{-8}$ as the default value in PyTorch) and may result in very loose bound.\n\n2.  As mentioned in this paper, the assumed noise condition is still stronger than the assumptions used for the analysis of other optimizers, for example, affine variance assumption.\n\n**Minor Issue**:\n\nThe statement with respect to [Zhang et al., 2022; Wang et al., 2022] in line 28 is not inproper. In fact, if you pick $\\beta_{sq}$ according to $\\epsilon$ in [Zhang et al., 2022; Wang et al., 2022] just as this paper does, you can derive the convergence to stationary points (but with a slower rate). Therefore, a proper statement should be \"but the convergence rate of [Zhang et al., 2022; Wang et al., 2022] is slower\".\n\n**References**\n\nDefossez et al., A Simple Convergence Proof of Adam and Adagrad, 2019\n\nZhang et al., Adam Can Converge Without Any Modification On Update Rules, 2022\n\nWang et al., Provable adaptivity in Adam, 2022\n\nQUESTIONS:\nI will consider increasing my score according to the responses to the questions.\n\n1. Although the studied smoothness assumption is more general, what is the motivation to study it? Specifically, according to [Zhang et al. 2019], $(L_0,L_1)$ smoothness condition seems to be enough to capture the training process of neural networks.\n\n2. The convergence rate has a dependence on the initial gradient norm $\\Vert \\nabla f(x_1) \\Vert $. Does this mean that Adam is sensitive to the initialization? I expect some explanations here.\n\nLIMITATIONS:\nNone", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

comments: SUMMARY:
This paper studies the convergence of the Adam algorithm. Under a more general local smoothness assumption, the convergence of Adam to stationary points is proved without assuming boundedness of the loss gradient. Here the key technique is to show that the loss gradient along the trajectory is indeed bounded, using a proof by contradiction. A variance-reduced variant of Adam is proposed to achieve an accelerated gradient complexity.

STRENGTHS:
The convergence of Adam is definitely an important question, and this paper seems to be a significant contribution by relaxing the bounded gradient assumption. The argument of bounding the gradients along the optimization trajectory looks neat and easy to understand. Overall the result of the paper is solid and novel, and the writing is also easy to follow.

WEAKNESSES:
It seems that the proof doesn't suggest the benefit of the momentum. It would be helpful if the authors can comment on this and provide some insights on the theoretical understanding of the momentum term.

It would also be helpful if the authors can summarize the recent results on Adam convergence in the form of a table.

QUESTIONS:
- Under the current local smoothness condition, how do we compare the performance of Adam and other algorithms like SGD? Does the current result suggest any advantage of Adam?

LIMITATIONS:
The authors have adequately addressed the limitations.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-10

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper provides convergence results for Adaptive Moment Estimate (Adam) and its variance-reduced variant under generalized smooth and bouned-noise assumptions.\n\nSTRENGTHS:\nThis paper mainly studies the convergence of Adaptive Moment Estimate (Adam) under a generalized smooth assumption. The authors drop the commonly used globally bounded gradients assumption and provide a new analysis framework, which is based on a contradiction argument, to show that gradients are bounded along the optimization trajectory. Based on this, the authors could deduce a high probability convergence bound of $\\mathcal{O}(\\epsilon^{-4})$ for Adam under the generalized smooth assumption. The paper also provides a variance-reduced version of Adam and improves the gradient complexity to $\\mathcal{O}(\\epsilon^{-3})$. \n\nIn general, the assumptions are standard in the analysis of adaptive methods including unbiased gradient estimate and bounded noise. The convergence results make sense and are highly valuable since the convergence result of Adam under unbounded gradients still leaves empty, to my best knowledge. As Adam is a widely used optimizer in deep learning field, the theoretical results (like parameter setting in Theorem 4.1 and Theorem 4.3) in this paper could help practitioners to better understand and use the algorithm. The central idea of analysis is inspiring.\n\nWEAKNESSES:\nFirst, the convergence result requires a dedicated tuned step-size and momentum parameters, specifically requiring prior knowledge such as generalized smooth parameter $(L_0, L_1)$ and the noise level $\\sigma$. However, in most realistic situations, these parameters could be unknown or hard to obtain. Thus, the result is more valuable in the theoretical aspect but may not be so useful in the experimental aspect.\n\nSecond, the parameter setting in Theorem 4.1 contradicts the real-one default setting on deep learning packages (such as Pytorch or Tensorflow) where $\\beta_1 = 0.9 $ is smaller than $\\beta_2 = 0.999$. It may be more persuasive to do some experiments showing the convergence under this new parameter setting.\n\nThird, the high probability convergence result in Theorem 6.2 does not provide an optimal rate to the probability margin $\\delta$, leaving space for further improvement. In addition, since VRAdam is a new algorithm, it better requires some experiments to show its convergence under the parameter setting in Theorem 6.2 and its faster convergence rate than Adam.\n\nQUESTIONS:\nsee the weakness part\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-25"}