Convergence of adaptive algorithms for constrained weakly convex optimization

1 Introduction . 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 } . 1.1 Examples of weakly convex problems . 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 . 1.2 Related work . 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 ] . 1.3 Notation . 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 can not 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 .

Summary: 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.