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