Convergence of adaptive algorithms for constrained weakly convex optimization
Year: 2021

Abstract:
We analyze the adaptive first order algorithm AMSGrad, for solving a constrained stochastic optimization problem with a weakly convex objective. We prove the Õ(t−1/2) rate of convergence for the squared norm of the gradient of Moreau envelope, which is the standard stationarity measure for this class of problems. It matches the known rates that adaptive algorithms enjoy for the specific case of unconstrained smooth nonconvex stochastic optimization. Our analysis works with mini-batch size of 1, constant first and second order moment parameters, and possibly unbounded optimization domains. Finally, we illustrate the applications and extensions of our results to specific problems and algorithms.

heading: 1 Introduction
text: Adaptive first order methods have become a mainstay of neural network training in recent years. Most of these methods build on the AdaGrad framework [14], which is a modification of online gradient descent by incorporating the sum of the squared gradients in the step size rule. Based on the practical shortcomings of AdaGrad for training neural networks, RMSprop [32] and Adam [21] proposed to use exponential moving averages for gradients and squared gradients (also known as moment estimations) with parameters β1 and β2, respectively. These methods have seen a huge practical success.
The recent work [28] identified a technical issue that affects Adam and RMSprop and proposed a new Adam-variant called AMSGrad that does not suffer from the same problem. Theoretical properties of AMSGrad, AdaGrad and their variants for nonconvex optimization problems are studied in a number of recent papers [3, 4, 11, 24, 34, 37]. These works focus on unconstrained smooth stochastic optimization, where the standard analysis framework of the stochastic gradient descent (SGD) [18] can be used. Convergence of adaptive methods for the more general setting of constrained and/or nonsmooth stochastic nonconvex optimization has remained open, while these settings have broad practical applications [7, 13, 20, 25, 27, 33].
In this work, we take a step towards this direction and establish the convergence of AMSGrad for solving the problem
min x∈X {f(x) = Eξ∼P [f(x; ξ)]} , (1)
where f : Rd → R is ρ-weakly convex, X ⊆ Rd is closed convex, and ξ is a r.v. following a fixed unknown distribution P. This template captures the setting of previous analyses when f is L-smooth, as this implies L-weak convexity, and X = Rd. However, there exist many applications when X 6= Rd [20, 25, 27, 33] or when f is not L-smooth [7, Section 2.1],[13, 15]. Constrained stochastic minimization with nonconvexity presents challenges not met in the convex setting [5, 19]. In particular, until the recent work [7], even for SGD, increasing mini-batch sizes were required for convergence in constrained nonconvex optimization. To study the behavior of AMSGrad for solving (1), we build on the analysis framework of [7].
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
Contributions. We show that AMSGrad achieves O(log(T )/ √ T ) rate for near-stationarity, see (5), for solving (1). Key specifications for this result are the following:
• We can use a mini-batch size of 1. • We can use constant moment parameters β1, β2 which are used in practice [1, 4, 21, 28]. • We do not assume boundedness of the domain X .
Next, we particularize our results for constrained optimization with L-smooth objectives and for a variant of RMSprop. We also extend our analysis for the scalar version of AdaGrad with first order moment estimation. For easy reference, we compare our results with state-of-the-art in Table 1. Finally, in a numerical experiment for robust phase retrieval, we observe that AMSGrad is more robust to variation of initial step sizes, compared to SGD and SGD with momentum.
Algorithm 1 AMSGrad [28] Input: x1 ∈ X , αt = α√t , for t ≥ 1, α > 0, β1 < 1, β2 < 1, m0 = v0 = 0, v̂0 = δ1, 1 ≥ δ > 0. for t = 1, 2 . . . T do
Sample ξt ∼ P iid and set gt such that Eξt [gt] ∈ ∂f(xt) (see Assumption 1) mt = β1mt−1 + (1− β1)gt vt = β2vt−1 + (1− β2)g2t v̂t = max(v̂t−1, vt) xt+1 = P v̂ 1/2 t X (xt − αtv̂ −1/2 t mt)
end for Output: xt∗ , where t∗ is selected uniformly at random from {1, . . . , T}.

heading: 1.1 Examples of weakly convex problems
text: The class of problems we consider in this paper include constrained problems with L-smooth objectives which are, for example, studied in [5] in the context of adversarial attacks and adaptive methods. Other important examples with weak convexity are composite objectives h(c(x)), where h is a convex Lipschitz continuous function and c is a smooth map with Lipschitz continuous Jacobian. Concrete examples of weakly convex problems are listed in [7, Section 2.1], which include robust phase retrieval, sparse dictionary learning, Conditional Value-at-Risk, and many others.

heading: 1.2 Related work
text: Adaptive algorithms based on AdaGrad [14] and Adam [1, 21, 28] are classically analyzed in online optimization framework with convex objective functions. Recent works studied the behavior of these methods for nonconvex optimization [2–4, 11, 24, 34, 36, 37]. The common characteristic of these results is that they are based on the well established proof templates of SGD [18] that only works in the simplest case of unconstrained smooth stochastic minimization. Moreover, as mentioned in [1], unconstrained optimization makes it easier to use a constant β1 parameter in Adam-type methods. In particular, many results for constrained optimization require a fast diminishing schedule for β1 parameter, while a constant parameter is used in practice [5, 21, 28].
The specific case of (1) with L-smooth f is studied by [5], where the authors proposed a zeroth order variant of AMSGrad. This result applies for the specific case of β1 = 0 which corresponds to a variant of RMSprop [28, 32]. More importantly, since its analysis follows the one of [19], increasing mini-batch sizes of the order √ t are required [5, Theorem 2].
As also mentioned in [5, 7], analysis of SGD for constrained problems introduces specific difficulties that are not observed in the convex case. Due to this, classical works analyzing SGD for nonconvex constrained optimization used large mini-batches to ensure convergence [19]. Showing convergence for SGD for constrained optimization with a single sample had been an open question until [7] gave a positive answer in the framework of weakly convex stochastic optimization, which includes constrained smooth stochastic optimization as a special case.
Weakly convex optimization is well studied with SGD based methods [7, 10, 15]. Recent work by [26], considers momentum SGD for solving (1). However, this algorithm (i) does not use momentum with β2 and (ii) uses non-adaptive, scalar, fixed step size: in the notation of Algorithm 1, v̂t = 1, αt = α/ √ T . These make the algorithm less practical, while simpler for analysis. Moreover the original work of Davis and Drusvyatskiy [7] analyzed standard SGD with no momentum or adaptive step sizes. As incorporating these mechanisms introduce more error terms, we need tighter estimations (for example see (9)) and techniques to handle time-dependent adaptive and diagonal step sizes.
The difficulty of handling the combination of nonconvexity, adaptive step sizes, momentum and constrained sets is mentioned in [5, Section 4.3]. In particular, in terms of our analysis, (i) adaptive step size introduces coupling between the step sizes, iterates, and the proximal point defined in (4); (ii) time-dependent diagonal step size requires an analysis framework based on variable metrics. Both of these issues were not the case in earlier works for weakly convex optimization [7, 26], and are among the key aspects of our analysis. For details, please see Lemma 1, Lemma 2 and discussions therein.
Another promising direction of research concerns nonsmooth nonconvex problems under more general assumptions. For instance, tameness and Hadamard semi-differentiability are used in [8] and [35], respectively, where convergence guarantees are established for SGD-based methods. Because of the generality of the problem class in these works, the algorithms studied there are simpler than the Adam-type algorithms considered in this paper, and the stationarity measures are less standard [35].

heading: 1.3 Notation
text: We adopt the convention of using the standard operations ab, a2, a/b, a1/2, 1/a, max{a, b} as element-wise, given two vectors a, b ∈ Rd. To denote ith element of the vector at ∈ Rd, we use the notation at,i. All-ones vector is denoted as 1. Given a vector a ∈ Rd, we define the matrix diag(a) ∈ Rd×d as the diagonal matrix with elements of a in the diagonal. For any set X , indicator function IX is given by IX (x) = 0 if x ∈ X and IX (x) = +∞ otherwise. Given the elements vi > 0, i = 1, . . . , d, we define a weighted norm ‖x‖2v := 〈x,diag(v)x〉. The weighted projection operator onto X is defined as P vX (x) = argminy∈X ‖y − x‖2v. A standard property of this operator is nonexpansiveness: ∀x, y ∈ Rd, ‖P vX (y)− P vX (x)‖v ≤ ‖y − x‖v . Due to nonconvexity, we cannot use standard definition of subgradients to form a global underestimator. Regular subdifferential, denoted as ∂f , for nonconvex functions [30, Ch. 8] is defined as the set of vectors q ∈ Rd such that, ∀x, y ∈ Rd, q ∈ ∂f(x) if
f(y) ≥ f(x) + 〈y − x, q〉+ o(‖y − x‖), as y → x. (2) When f is convex, this reduces to standard definition of a subdifferential and when f is differentiable, this set coincides with {∇f(x)}. We say that f is ρ-weakly convex w.r.t. ‖ · ‖, if f(x) + ρ2‖x‖ 2 is convex. An equivalent representation for weakly convex functions is that, ∀x, y ∈ Rd, where q ∈ ∂f(x) [7, Lemma 2.1],
f(y) ≥ f(x) + 〈y − x, q〉 − ρ 2 ‖y − x‖2. (3)
Moreover, we say f is L-smooth, if it holds that, ∀x, y ∈ Rd, ‖∇f(x) − ∇f(y)‖ ≤ L‖x − y‖. Given random iterates x1, . . . , xt, we denote the filtration generated by these realizations as Ft = σ(x1, . . . , xt), and the corresponding conditional expectation as Et[·] = E[·|Ft]. By the law of total expectation, E [Et[·]] = E[·]. We also sometimes use Eξ to denote expectation w.r.t. randomness of ξ.
We now present the assumptions of our analysis.
Assumption 1. • f : Rd → R is ρ-weakly convex with respect to norm ‖ · ‖. • The set X ⊆ Rd is convex and closed. • We can obtain iid samples ξt ∼ P, and gt such that Eξt [gt] ∈ ∂f(xt) and ‖gt‖∞ ≤ G,∀t. • f is lower bounded: f? ≤ f(x),∀x ∈ X .
A few remarks are in order for Assumption 1. First, we do not require boundedness of the domain X . Second, weak convexity assumption is weaker than smoothness assumption on f and the assumption of bounded gradients is standard for adaptive algorithms [3, 4, 11]. In principle, it is possible to relax the bounded stochastic subgradient assumption to the weaker requirement E‖gt‖2 ≤ G as in [37, Remark 6. (ii)] with a slightly worse and complicated convergence rate. For simplicity, we use Assumption 1. The third item in Assumption 1 concerning sampling a stochastic oracle is standard in the weakly convex optimization literature [7, Assumption A2], [26, Assumption A1]. Remark 1. We note that when f is ρ-weakly convex w.r.t. ‖ · ‖, then it is ρ√
δ -weakly convex w.r.t.
‖ · ‖ v̂ 1/2 t , ∀t, since v̂t,i ≥ δ > 0 (see Algorithm 1). We denote ρ̂ = ρ√δ .
It is easy to verify this remark by noticing that x 7→ f(x) + ρ2‖x‖ 2 is convex and ρ̂2‖x‖ 2
v̂ 1/2 t
≥ ρ2‖x‖ 2.

heading: 2 Algorithm
text: We analyze the algorithm AMSGrad (see Algorithm 1) proposed in [28]. On top of Adam [21], it includes a step to ensure monotonicity of the exponential moving average of squared gradients. It is standard in stochastic nonconvex optimization to output a randomly selected iterate [7, 18, 19], which we also adopt. We next define the composite objective
ϕ(x) = f(x) + IX (x).
For nonsmooth problems, the standard stationarity measures such as the norm of (sub)gradients are no longer applicable, see [7, 26] and [13, Section 4]. This motivates the following definitions that, as we show below, relate to a relaxed form of stationarity. Based on ϕ and a parameter ρ̄ > 0, we define the proximal point of xt and the Moreau envelope, respectively as
x̂t = prox v̂ 1/2 t
ϕ/ρ̄ (xt) = argmin y
{ ϕ(y) + ρ̄
2 ‖y − xt‖2v̂1/2t
} , (4)
ϕt1/ρ̄(xt) = miny
{ ϕ(y) + ρ̄
2 ‖y − xt‖2v̂1/2t
} .
We compare the definitions with that of [7]. Due to the use of variable metric v̂t in adaptive methods, we have a time dependent Moreau envelope, where the corresponding vector v̂t is used for defining the norm. Important considerations for these quantities are the uniqueness of x̂t and the smoothness of ϕt1/ρ̄. As we shall see now, choice of ρ̄ is critical for ensuring these. In light of Remark 1, selecting ρ̄ > ρ̂ = ρ√
δ , and by using the arguments in [7, for example, Lemma 2.2], [29, Theorem 31.5], it
follows x̂t is unique and ϕt1/ρ̄ is smooth with the gradient∇ϕ t 1/ρ̄(xt) = ρ̄v̂ 1/2 t (xt − x̂t).
Near stationarity. Near-stationarity conditions follow from the optimality condition of x̂t: 0 ∈ ∂ϕ(x̂t) + ρ̄v̂
1/2 t (x̂t − xt), where we also use v̂t,i ≤ G2: ‖xt − x̂t‖2 v̂ 1/2 t = 1ρ̄2 ‖∇ϕ t 1/ρ̄(xt)‖ 2 v̂ −1/2 t dist2(0, ∂ϕ(x̂t)) ≤ G‖∇ϕt1/ρ̄(xt)‖ 2 v̂ −1/2 t
ϕ(x̂t) ≤ ϕ(xt). (5)
Consistent with previous literature for weakly convex optimization [7, 26], we state the convergence guarantees in terms of the norm of the gradient of Moreau envelope. Given (5), this means that the iterate xt is near stationary: it is close to its proximal point x̂t and x̂t is approximately stationary.

heading: 3 Convergence

heading: 3.1 Main result
text: We start with our main theorem that shows that the norm of the gradient of Moreau envelope converges to 0 at the claimed rate, resulting in near-stationarity of xt∗ as in (5), see Remark 1 for definition of ρ̂. Theorem 1. Let Assumption 1 hold. Let β1 < 1, β2 < 1, γ = β21 β2 < 1, ρ̄ = 2ρ̂. Then, for iterate xt∗ generated by Algorithm 1, it follows that
E‖∇ϕt ∗
1/ρ̄(xt∗)‖ 2 v̂ −1/2 t∗ ≤ 2dG α √ δT (1− β1)
[ C1 + (1 + log T )C2 + C3 ] ,
where C1 = 8ραG+ 1dG ( ϕ11/ρ̄(x1)− f ? ) , C2 = 2ρα 2√
(1−β2)(1−γ)
( 10G√ δ ) , C3 = 8Gρ ∑d i=1 Ev̂ 1/2 T+1,i.
The bound in Theorem 1 has complicated constants as it is usual for adaptive algorithms in nonconvex case [3, 4]. These constants are slightly simplified and the proof of Theorem 1 in Appendix A includes the non-simplified version. Next, we explain and interpret the bound in terms of dependence to key parameters.

heading: 3.2 Discussion on Theorem 1
text: In the context of near-stationarity (5), Theorem 1 states that to have xt∗ in Algorithm 1 such that ‖∇ϕt∗1/ρ̄(xt∗)‖v̂−1/2t ≤ , we require Õ(
−4) iterations. This matches the known complexities for adaptive methods in unconstrained smooth stochastic optimization [1, 3, 4, 11, 24, 34, 36, 37], and SGD-type methods in weakly convex optimization [7, 26].
Our first remark is about the metric of the norm used for the gradient of Moreau envelope in Theorem 1. We then discuss the dependence of our bound w.r.t. important quantities. Remark 2. By (5), one has ‖∇ϕt∗1/ρ̄(xt∗)‖ 2
v̂ −1/2 t∗
= ρ̄2‖xt∗−x̂t∗‖2 v̂ 1/2 t∗ . We note that ‖xt∗−x̂t∗‖2 v̂ 1/2 t∗ ≥
√ δ‖xt∗−x̂t∗‖2 as v̂t,i ≥ δ. It also holds that v̂t,i ≤ G2. Therefore, one can convert our guarantees to ‖xt∗ − x̂t∗‖2 or ‖∇ϕt ∗ (xt∗)‖ by multiplying the right hand side by appropriate quantities depending on δ or G. We leave the result with the metric, as δ and G are the worst case bounds.
Knowledge of ρ. To run the algorithm, one does not need to know the weak convexity parameter ρ. The parameters ρ̄ and ρ̂ are merely for analysis purposes [8, 26], and the convergence rate holds for any choice of step size αt, independent of ρ.
Dependence w.r.t. β1. Comparing with the previous work, the scaling of our bound in terms of β1 is (1− β1)−1 matching the dependence for the unconstrained setting [1, 11]. Dependence w.r.t. d. Standard dependence in the convergence rates of Adam-type algorithms for unconstrained case is d/ √ T [1, 11].1
Even though in Theorem 1, the constant C3 has worst case dependence d2, this is merely due to assumptions. The main reason is that we do not assume boundedness of the sequence xt, instead we prove the necessary result for the analysis in Lemma 1. However, this result gives a bound for ‖xt − x̂t‖, which is naturally dimension dependent. We used this bound in (10), where we need to bound ‖xt − x̂t‖∞. In particular, if we had assumed a bound for ‖xt − x̂t‖∞, then in (10) we could have used it instead of Lemma 1 to have standard d/ √ T in C3. Boundedness assumption also would remove a factor of 1√ δ
in the bound, as those terms appear in the steps where we avoid boundedness assumption. However, for generality, we do not assume boundedness.
Dependence w.r.t. δ. Our bound has a polynomial dependence of 1/δ similar to [1, 3, 4]. In [11], a more refined technique from [34] is used to have a logarithmic dependence of 1/δ. This technique, used on the case of smooth unconstrained problems in these works, did not seem to apply to our setting.
1We note that in [3] better dependence is obtained by using step sizes in the order of 1√ d , which we do not consider.

heading: 3.3 Analysis
text: In this section, we will flesh out the main ideas of our proof with four lemmas. We will continue with a proof sketch to show how the pieces come together.

heading: 3.3.1 Preliminary results
text: We start with a result showing that under Assumption 1, the quantity ‖xt−x̂t‖ from (5) stays bounded. We remark that third term on RHS in (9) arises as a spurious term due to time-dependent diagonal step sizes, which was not the case in previous works on weakly convex optimization with scalar step sizes [26]. Our result in the next lemma is the main tool for us to avoid assuming boundedness of X . The proof of this lemma given in Appendix A combines the definition of x̂t with weak convexity. Lemma 1. Let Assumption 1 hold. Let ρ̄ > ρ̂, and v̂t ≥ δ > 0 (see Algorithm 1). It follows that
‖xt − x̂t‖2 ≤ D̂2 := 4dG2
δ(ρ̄− ρ̂)2 .
A key aspect in the analysis of adaptive algorithms is the dependence of v̂t and gt that couples x̂t and gt (see (4)), preventing taking expectation of 〈xt − x̂t, gt〉 that we use to obtain the stationarity measure in the proof. As this was not the case in prior works [7, 26], we need a more refined analysis. Lemma 2. Let Assumption 1 hold. Let qt = Et[gt] ∈ ∂f(xt), then it follows that
αtEt〈xt− x̂t, gt〉 ≥ αt(ρ̄− ρ̂)Et‖xt− x̂t‖2v̂1/2t − (αt−1−αt)
√ dD̂G− ρ̄− ρ̂
4ρ̄ Et‖x̂t− x̂t−1‖2v̂1/2t−1
− αt−1 2 Et‖mt−1‖2v̂−1/2t−1 − ( 1 2 + ρ̄ ρ̄− ρ̂ ) α2t−1√ δ Et‖gt‖2.
Interpreting Lemma 2. We review the terms in this bound to gain some intuition. The first term in the RHS is the stationarity measure (see (5)), second term will sum to a constant, fourth and fifth terms will sum to log(T ) by Lemma 4. Handling the third term in RHS is not as obvious, but we can show that we can cancel it using the contribution from another part of the analysis that we detail in the full proof (see (9)).
One critical issue for Adam-type algorithms is to obtain results with constant β1 parameter. A recent work [1] studied this problem for constrained convex problems. The following lemma also plays an important role in our analysis. Lemma 3. [1, Lemma 1] Let mt = β1mt−1 + (1− β1)gt. Then for any vectors At−1, At,
〈At, gt〉 = 1
1− β1 (〈At,mt〉 − 〈At−1,mt−1〉) + 〈At−1,mt−1〉+ β1 1− β1 〈At−1 −At,mt−1〉.
This lemma derives a decomposition for handling constant β1 parameter in the beginning of the analysis. As explained in Section 3.1 of the abovementioned paper, using a decomposition for mt later in the analysis results in a requirement of decreasing β1, especially for constrained problems, which we would like to avoid.
Next lemma is a standard estimation used for the analysis of Adam-based methods, dating back to [21]. For easy reference we point out to [1] where this bound is included as a separate lemma with tighter estimations than previous works, due to using a constant β1. It bounds the sum of the norms of first moment vectors multiplied by the step size sequence.
Lemma 4. Let β1 < 1, β2 < 1, γ = β21 β2 < 1, then it holds that
T∑ t=1 α2t ‖mt‖2v̂−1/2t ≤ (1− β1)α 2√ (1− β2)(1− γ) dG(1 + log T ).

heading: 3.3.2 Proof sketch of Theorem 1
text: The proof sketch of the theorem is a careful combination of the preliminary results mentioned in the previous section. The sketch includes the necessary bounds, but omits the tedious estimations required in some steps. The full proof with the details is given in Appendix A.
Proof sketch of Theorem 1. We sum the result of Lemma 3 and use A1 = A0. with m0 = 0. We note that we have At = ρ̄αt(xt − x̂t), for t ≥ 1.
T∑ t=1 〈At, gt〉 = β1 1− β1 〈AT ,mT 〉+ T∑ t=1 〈At,mt〉+ β1 1− β1 T−1∑ t=1 〈At −At+1,mt〉. (6)
After plugging in the value of At, (6) becomes
T∑ t=1 ρ̄αt〈xt − x̂t, gt〉 ≤ β1ρ̄αT 1− β1 〈xT − x̂T ,mT 〉+ T∑ t=1 ρ̄αt〈xt − x̂t,mt〉
+ β1ρ̄
1− β1 T−1∑ t=1 〈αt(xt − x̂t)− αt+1(xt+1 − x̂t+1),mt〉. (7)
LHS of this bound is suitable for applying Lemma 2 to obtain the stationarity measure. We have to estimate the three terms on the RHS. It is easy to bound the first term using Cauchy-Schwarz inequality and Lemma 1. Other two terms require longer estimations which we sketch below.
• Bound for β1ρ̄1−β1 ∑T−1 t=1 〈αt(xt − x̂t)− αt+1(xt+1 − x̂t+1),mt〉 in (7).
Decomposing this term gives
〈αt(xt − x̂t)− αt+1(xt+1 − x̂t+1),mt〉 = (αt − αt+1)〈xt+1 − x̂t+1,mt〉+ αt〈xt − xt+1,mt〉 + αt〈x̂t+1 − x̂t,mt〉.
For the first term, we use that αt ≥ αt+1 and Cauchy-Schwarz inequality
T−1∑ t=1 (αt − αt+1)〈xt+1 − x̂t+1,mt〉 ≤ α1D̂ √ dG.
For the second term we deduce by Cauchy-Schwarz inequality and nonexpansiveness of the projection
αt〈xt − xt+1,mt〉 ≤ α2t ‖mt‖2v̂−1/2t .
For the third term, we use Young’s inequality to obtain the bound
T−1∑ t=1 β1ρ̄ 1− β1 〈αt(xt−x̂t)−αt+1(xt+1−x̂t+1),mt〉 ≤ β1ρ̄ 1− β1 α1D̂ √ dG+ T∑ t=1 β1ρ̄α 2 t 1− β1 ‖mt‖2v̂−1/2t
+ T∑ t=1 ρ̄− ρ̂ 4 ‖x̂t+1 − x̂t‖2v̂1/2t + ρ̄2 (ρ̄− ρ̂) β21 (1− β1)2 T∑ t=1 α2t ‖mt‖2v̂−1/2t , (8)
• Bound for ∑T t=1 ρ̄αt〈xt − x̂t,mt〉 in (7).
We proceed similar to [7], with a tighter estimation (resulting in the negative term on RHS) to obtain
ϕt+11/ρ̄ (xt+1) ≤ ϕ t 1/ρ̄(xt) + ρ̄αt〈x̂t− xt,mt〉+
ρ̄ 2 ‖x̂t− xt+1‖2v̂1/2t+1−v̂1/2t − ρ̄− ρ̂ 2 ‖x̂t− x̂t+1‖2v̂1/2t+1
+ ρ̄
2 α2t ‖mt‖2v̂−1/2t . (9)
Then we manipulate the fourth term on RHS with standard ‖a− b‖2 ≤ 2‖a‖2 + 2‖b‖2, and Lemma 1,
ρ̄ 2 ‖x̂t − xt+1‖2v̂1/2t+1−v̂1/2t ≤ ρ̄‖x̂t − xt‖2v̂1/2t+1−v̂1/2t + Gρ̄√ δ ‖xt − xt+1‖2v̂1/2t
≤ ρ̄D̂2 d∑ i=1 (v̂ 1/2 t+1,i − v̂ 1/2 t,i ) + Gρ̄√ δ α2t ‖mt‖2v̂−1/2t . (10)
We use this estimation in (9) and sum to get
ρ̄αt T∑ t=1 〈xt − x̂t,mt〉 ≤ ϕ11/ρ̄(x1)− ϕ T+1 1/ρ̄ (xT+1) + T∑ t=1 ( 1 2 + G√ δ ) ρ̄α2t ‖mt‖2v̂−1/2t
+ ρ̄D̂2 d∑ i=1 v̂ 1/2 T+1,i − T∑ t=1 ρ̄− ρ̂ 2 ‖x̂t − x̂t+1‖2v̂1/2t+1 . (11)
We collect (8) and (11) in (7). Finally, we have to obtain the stationarity criterion on the LHS of (7) by taking conditional expectation. This is not immediate due to coupling of x̂t, v̂t, and gt. We use Lemma 2 to handle this issue and the negative term in (11) is utilized to cancel the third term in the RHS of the result of Lemma 2. Then, we use (5), plug in Lemma 4 and ρ̄ = 2ρ̂ to conclude.

heading: 4 Applications & Extensions

heading: 4.1 Applications
text: RMSprop. The counterexamples presented in [28] show that RMSprop, similar to Adam might diverge in simple problems. Setting β1 = 0 in AMSGrad [28] results in an algorithm similar to RMSprop, with the difference of having v̂t as the output of the max step. Therefore, our analysis also applies to this version of RMSprop with similar guarantees.
Corollary 1. Let β1 = 0. Then, for a variant of RMSprop [28], Theorem 1 applies.
It is easy to see that β1 = 0 gives a better bound in Theorem 1. This is in fact common for the bounds of Adam-type algorithms even in the convex case [28]. Setting nonzero momentum parameters β1, β2 do not predict improvement, however, in practice they are routinely observed to improve performance.
SGD with momentum. When v̂t = 1,∀t, AMSGrad reduces to an algorithm similar to SGD with momentum. Lack of diagonal step sizes in this case simplifies the analysis as weighted projections are not used in the algorithm. This specific case is studied in the recent work [26], with a slightly different way to set mt. Our analysis can be seen as an alternative derivation of convergence for a method similar to [26].
Constrained smooth optimization. A special case of (1) is when f isL-smooth. In this case, standard convergence measure is the gradient mapping [19], which is used in [5]: Gλ(x) = v̂1/4t λ−1(x − P v̂ 1/2 t
X (x − λv̂ −1/2 t ∇f(x))). It is instructive to observe that when X = Rd, then ‖Gλ(x)‖ = ‖∇f(x)‖ v̂ −1/2 t ≥ G−1‖∇f(x)‖which is the stationarity measure for smooth unconstrained problems. When X 6= Rd, gradient mapping is used as the stationarity measure [7, 19, 26]. As illustrated in [7], for the specific case of constrained smooth minimization, norm of the Moreau envelope is of the same order as the norm of the gradient mapping, therefore, the results can be converted to guarantees on gradient mapping norms. Using similar ideas as in [12, Theorem 3.5], [7], one can show that ‖G1/ρ̄(xt)‖ ≤ Cg,m‖∇ϕt1/ρ̄(xt)‖v̂−1/2t , for a constant Cg,m (see Appendix B.1).

heading: 4.2 Extension: Scalar AdaGrad with momentum
text: An alternative adaptive algorithm is AdaGrad [14] and its variants with first order moment estimation are referred to as AdamNC [28] or AdaFOM [4]. In unconstrained smooth stochastic optimization, it has been observed that the same proof structure applies to AMSGrad and AdaGrad-based methods simultaneously [4, 11]. However, in our setting, the analysis we developed for AMSGrad does not directly apply to AdaGrad-based methods.
The main reason is that vt in the case of AdaGrad does not admit a lower bound separated from 0, unlike AMSGrad where 0 < δ ≤ v̂t. The uniform lower bound is necessary for converting regular weak convexity assumption w.r.t. norm ‖ · ‖ to the one w.r.t. the weighted norm ‖ · ‖
v 1/2 t in the sense of Remark 1. Naively assuming the existence of ρ̂ in Remark 1 is not consistent, since vt is not separated from zero due to vt ≥ δ√t in AdaGrad, and hence, the norm ‖ · ‖v1/2t is not well-defined.
In this section, we provide partial results on this direction. In particular, we show that the scalar version of AdaGrad, that is used in [22–24, 34], along with its variant with first order moment estimation also has the same convergence rate. In the framework of Algorithm 1, scalar (non-diagonal) version of these methods iterate as, for iid sampled ξt ∼ P, Eξt [gt] ∈ ∂f(xt), mt = β1mt−1 + (1− β1)gt vt = 1 t (δ + 1 d ∑t j=1 ‖gj‖2)
xt+1 = PX (xt − αt√vtmt), (12)
where PX denotes standard Euclidean projection. The factor of d−1 in front of gradient norms is to normalize the step size, as `2-norm is dimension dependent. This factor only affects the dimension dependence of the bound. In this case, one does not need the time-dependent definitions for Moreau envelope and x̂t: x̂t = prox1/ρ̄(xt) and ϕ1/ρ̄(x) = miny∈X f(y) + ρ̄ 2‖y − x‖
2, due to lack of weighted projection. The proof then is similar to [7] with AdaGrad step sizes. The difficulties arising due to adaptive step sizes and existence of β1, are handled using Lemmas 1, 3 and 4. Theorem 2. Let Assumption 1 hold. Then, for the method sketched in (12), with β1 < 1, αt = α√t
E‖∇ϕ1/2ρ(xt∗)‖2 ≤ 2G √ ρ
α √ T
[ C1 + ( 1 + log ( TG2
δ + 1
)) C2 ] ,
where C1 = ϕ1/2ρ(x1)− f? + 12αdG 2
√ δ(1−β1)
, C2 = 8ρα 2d
(1−β1)2 .
The terms in the bound are simplified as Theorem 1 and their non-simplified variants are in the proof in Appendix A. Our analysis takes care of the coupling between the AdaGrad step sizes and iterates of the algorithm. However, unlike our results with AMSGrad, in this case the analysis does not cover the diagonal case, where the step sizes are set by using elementwise squared gradients g2t . We leave it as an open question to derive similar results for AdaGrad-based methods with diagonal step sizes.

heading: 5 Numerical experiment
text: This section illustrates the potential advantages of adaptive algorithms, in particular AMSGrad, for solving a prototypical weakly convex problem, compared to SGD and SGD with momentum (also referred in [26] as stochastic heavy ball: SHB). As popular in the literature of weakly convex stochastic optimization methods [8, 10, 26], we will compare the algorithms in terms of their robustness to initial step sizes. We note that “robustness to tuning” of algorithms is also investigated in the context of deep learning in the literature and the advantage of adaptive algorithms such as Adam/AMSGrad is observed [6, 31]. In particular, we solve the robust phase retrieval problem [9, 16, 17]: minx∈Rd 1n ∑n i=1 |〈ai, x〉2 − bi|, where A = [a1, . . . , an]> ∈ Rn×d, n = 300, d = 50. Weak convexity of this problem is well-known [7, 9].
We briefly recall the setup from [26] that considered SGD with momentum for solving this problem. The data is generated as A = QD, with a standard normal distributed matrix Q ∈ Rn×d and D = linspace(1/κ, 1, d), where κ ≥ 1 controls the conditioning. We generate the solution x? as a standard normal random vector with unit norm. The observations are generated as b = Ax? + δη
2We make the “effective initial step sizes” of algorithms equal. In particular we pick αSGD0 = αMSGD0 = αAMS0
β2 √ maxi(g 2 1,i) , since the initial step size of AMSGrad is α0√ v̂21 .
where elements of η ∈ Rn have distribution N (0, 25) and δ = diag(δ1, . . . , δd) is such that |{i ∈ [n] : δi = 1}|/n = 0.2, which means that 20% of the observations are corrupted. With this setup, it is proven in [9, Lemma B.12] that only solutions of the problem are {x?,−x?}. Therefore, for the algorithms, we will use f(xk) − f(x?) ≤ ε as the stopping criterion. We run stochastic subgradient method (SGD) [7], momentum SGD (SHB) [26] and AMSGrad that we analyzed in this paper. For all algorithms, the step size is chosen as αk = α0/ √ k. We varied the initial step size2 between 0.01 and 10 for all algorithms, and we plotted the number of epochs that the algorithms take to reach to f(x)− f(x?) ≤ 0.1. In terms of other algorithmic parameters, we use both β = 0.1 and β = 0.01 for SHB, as in in [26] and β1 = β2 = 0.99 as popular, for AMSGrad.
We present the results in Figure 1 for varying values of κ = {1, 10, 100}, where each setup is run for 50 times, medians are drawn as lines and the region between 20th and 80th percentiles are shaded. It has been observed in [26] that SHB improves the robustness of SGD to initial step sizes. We observe in Figure 1 that AMSGrad shows a more robust behavior to initial step size compared to both algorithms. We note that our findings indicate the potential of AMSGrad and adaptive methods for weakly convex optimization. Moreover, our findings about robustness of adaptive algorithms to tuning is consistent with the findings from deep learning literature [6, 31].
Acknowledgments and Disclosure of Funding
Most of the work was done while Ahmet Alacaoglu and Yura Malitsky were at EPFL.
This project received funding from NSF Award 2023239; DOE ASCR under Subcontract 8F-30039 from Argonne National Laboratory; the Wallenberg Al, Autonomous Systems and Software Program (WASP) funded by the Knut and Alice Wallenberg Foundation, with the project number 305286; the European Research Council (ERC) under the European Union’s Horizon 2020 research and innovation programme (grant agreement no 725594 - timedata); the Swiss National Science Foundation (SNSF) under grant number 200021_178865/1; the Department of the Navy, Office of Naval Research (ONR) under a grant number N62909-17-1-2111; and the Hasler Foundation Program: Cyber Human Systems (project number 16066).

Review summaries:
The paper provides a convergence analysis for the AMSGrad algorithm for the more general class of functions that are weakly convex. The convergence guarantee is comparable to existing results for SGD, and the main advantage is that the AMSGrad algorithm does not require knowledge of the weak convexity parameter and it automatically adapts to it. The paper extends the convergence analysis to other algorithms, such as SGD with momentum and the AdaGrad algorithm with scalar step sizes. 

This paper analyzes AMSGrad when applied to a broad family of nonsmooth, nonconvex, stochastic, constrained optimization problems. Convergence rates based on the norm of the Moreau envelope's gradient (although in a rescaled norm) are presented. The rates nearly match the O(1/T^{1/4}) rate of the subgradient method, off by a logarithmic term. The primary innovation in this work is tackling technical barriers to handle first and second order moment parameters that complicate analysis.