SUPER-ADAM: Faster and Universal Framework of Adaptive Gradients

1 Introduction . In the paper , we consider solving the following stochastic optimization problem : min x∈X f ( x ) : = Eξ∼D [ f ( x ; ξ ) ] , ( 1 ) where f ( x ) denotes a smooth and possibly nonconvex loss function , and ξ is a random example variable following an unknown data distribution D. Here X = Rd or X ⊂ Rd is a compact and convex set . The problem ( 1 ) frequently appears in many machine learning applications such as the expectation loss minimization . Recently , Stochastic Gradient Descent ( SGD ) [ 14 ] is commonly used to solve the problem ( 1 ) such as Deep Neural Networks ( DNNs ) training [ 18 , 20 ] , due to only requiring a mini-batch samples or even one sample at each iteration . Adaptive gradient methods are one of the most important variants of SGD , which use adaptive learning rates and possibly incorporate momentum techniques , so they generally require less parameter tuning and enjoy faster convergence rate than SGD . Meanwhile , compared to SGD , adaptive gradient methods escape saddle points faster [ 31 ] . Thus , recently adaptive gradient methods have been widely developed and studied . For example , the first adaptive gradient method i.e. , Adagrad has been proposed in [ 12 ] , which significantly outperforms the vanilla SGD under the sparse gradient setting . Subsequently , some variants of Adagrad e.g. , SC-Adagra [ 28 ] and SAdagrad [ 9 ] have been proposed for ( strongly ) convex optimization . Unfortunately , Adagrad has been found that it does not be well competent to the dense gradient setting and the nonconvex setting . To address this drawback , some other efficient variants of 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Adagrad , e.g. , Adadelta [ 37 ] , Adam [ 22 ] , have been presented by using exponential moving average instead of the arithmetic average . Adam [ 22 ] recently has been shown great successes in current machine learning problems , e.g. , it is a default method of choice for training DNNs [ 17 ] and contrastive learning [ 7 ] . Unfortunately , Reddi et al . [ 29 ] still showed that Adam is frequently divergent in some settings where the gradient information quickly disappear . To deal with this issue , some variants of Adam algorithm , e.g. , AMSGrad [ 29 ] , YOGI [ 36 ] and generalized Adam [ 8 ] have been proposed . Specifically , AMSGrad [ 29 ] applies an extra ‘ long term memory ’ variable to preserve the past gradient information in order to handle the convergence issue of Adam . YOGI [ 36 ] introduces an adaptive denominator constant , and studies effect of the mini-batch size in its convergence . Subsequently , Chen et al . [ 8 ] studied the convergence of a class of Adam-type algorithms for nonconvex optimization . Zhou et al . [ 39 ] analyzed the convergence of a class of adaptive gradient algorithms for nonconvex optimization , and the result shows the advantage of adaptive gradient methods over SGD in sparse stochastic gradient setting . Meanwhile , Liu et al . [ 24 ] studied the variances of these adaptive algorithms . More recently , Guo et al . [ 19 ] presented a novel convergence analysis for a family of Adam-style methods ( including Adam , AMSGrad , Adabound , etc . ) with an increasing or large momentum parameter for the first-order moment . Although the above these adaptive gradient methods show some good empirical performances , their generalization performance is worse than SGD ( with momentum ) on many deep learning tasks due to using the coordinate-wise learning rates [ 35 ] . Thus , recently some adaptive gradient methods have been proposed to improve the generalization performance of Adam . For example , AdamW [ 26 ] and Padam [ 6 ] improve the generalization performance of Adam by decoupling weight decay regularization and introducing a partial adaptive parameter , respectively . Luo et al . [ 27 ] proposed a new variant of Adam ( i.e. , Adabound ) by employing dynamic bounds on learning rates to improve the generalization performance . Subsequently , AdaBelief [ 40 ] has been presented to obtain a good generalization by adopting the stepsize according to the ‘ belief ’ in the current gradient direction . In addition , the norm version of AdaGrad ( i.e. , AdaGrad-Norm ) [ 34 ] has been proposed to obtain a good generalization performance . So far , the above adaptive gradient methods still suffer from a high gradient complexity of O ( −4 ) for finding -stationary point in the worst case without considering sparsity of gradient . More recently , some faster variance-reduced adaptive gradient methods such as STORM [ 11 ] , Adaptive Normalized SGD [ 10 ] , Adam+ [ 25 ] have been proposed . For example , STORM applies the momentum-based variance reduced technique to obtain a lower gradient complexity of Õ ( −3 ) . To the best of our knowledge , all these existing adaptive gradient methods only use some specific adaptive learning rates with focusing on either pure theoretical or empirical aspects . Thus , it is desired to design a universal framework for the adaptive gradient methods on both theoretical analysis and practical algorithms to solve the generic problems . To fill this gap , in the paper , we propose a faster and universal framework of adaptive gradients , i.e. , SUPER-ADAM algorithm , by introducing a universal adaptive matrix . Moreover , we provide a novel convergence analysis framework for the adaptive gradient methods under the nonconvex setting based on the mirror descent algorithm [ 5 , 15 ] . In summary , our main contributions are threefold : 1 ) We propose a faster and universal framework of adaptive gradients ( i.e. , SUPER-ADAM ) by introducing a universal adaptive matrix that includes most existing adaptive gradients . Moreover , our framework can flexibly integrate the momentum and variance-reduced techniques . 2 ) We provide a novel convergence analysis framework for the adaptive gradient methods in the nonconvex setting under the milder conditions ( Please see Table 1 ) . 3 ) We apply a momentum-based variance reduced gradient estimator [ 11 , 32 ] to our algorithm ( SUPER-ADAM ( τ = 1 ) ) , which makes our algorithm reach a faster convergence rate than the classic adaptive methods . Specifically , under smoothness of each component function f ( x ; ξ ) , we prove that the SUPER-ADAM ( τ = 1 ) achieves the best known gradient complexity of Õ ( −3 ) for finding an -stationary point of the problem ( 1 ) , which matches the lower bound for stochastic smooth nonconvex optimization [ 1 ] . Under smoothness of the function f ( x ) , we prove that the SUPER-ADAM ( τ = 0 ) achieves a gradient complexity of Õ ( −4 ) . 2 Preliminaries . 2.1 Notations . ‖ · ‖ denotes the ` 2 norm for vectors and spectral norm for matrices , respectively . Id denotes a d-dimensional identity matrix . diag ( a ) ∈ Rd denotes a diagonal matrix with diagonal entries a = ( a1 , · · · , ad ) . For vectors u and v , up ( p > 0 ) denotes element-wise power operation , u/v denotes element-wise division and max ( u , v ) denotes element-wise maximum . 〈u , v〉 denotes the inner product of two vectors u and v. For two sequences { an } and { bn } , we write an = O ( bn ) if there exists a positive constant C such that an ≤ Cbn , and Õ ( · ) hides logarithmic factors . A 0 ( 0 ) denotes a positive ( semi ) definite matrix . δmin ( A ) and δmax ( A ) denote the smallest and largest eigenvalues of the matrix A , respectively . 2.2 Adaptive Gradient Algorithms . In the subsection , we review some existing typical adaptive gradient methods . Recently , many adaptive algorithms have been proposed to solve the problem ( 1 ) , and achieve good performances . For example , Adagrad [ 12 ] is the first adaptive gradient method with adaptive learning rate for each individual dimension , which adopts the following update form : xt+1 = xt − ηtgt/ √ vt , ( 2 ) where gt = ∇f ( xt ; ξt ) and vt = 1t ∑t j=1 g 2 j , and ηt = η√ t with η > 0 is the step size . In fact , ηt only is the basic learning rate that is the same for all coordinates of variable xt , while ηt√vt , i is the effective learning rate for the i-th coordinate of xt , which changes across the coordinates . Adam [ 22 ] is one of the most popular exponential moving average variant of Adagrad , which combines the exponential moving average technique with momentum acceleration . Its update form is : mt = α1mt−1 + ( 1− α1 ) ∇f ( xt ; ξt ) , vt = α2vt−1 + ( 1− α2 ) ( ∇f ( xt ; ξt ) ) 2 m̂t = mt/ ( 1− αt1 ) , v̂t = vt/ ( 1− αt2 ) , xt+1 = xt − ηtm̂t/ ( √ v̂t + ε ) , ∀ t ≥ 1 ( 3 ) where α1 , α2 ∈ ( 0 , 1 ) and ε > 0 , and ηt = η√t with η > 0 . However , Reddi et al . [ 29 ] found a divergence issue of the Adam algorithm , and proposed a modified version of Adam ( i.e. , Amsgrad ) , which adopts a new step instead of the debiasing step in ( 3 ) to ensure the decay of the effective learning rate , defined as v̂t = max ( v̂t−1 , vt ) , xt+1 = xt − ηtmt/ √ v̂t . ( 4 ) Algorithm 1 SUPER-ADAM Algorithm 1 : Input : Total iteration T , and tuning parameters { µt , αt } Tt=1 , γ > 0 ; 2 : Initialize : x1 ∈ X , sample one point ξ1 and compute g1 = ∇f ( x1 ; ξ1 ) ; 3 : for t = 1 , 2 , . . . , T do 4 : Generate an adaptive matrix Ht ∈ Rd×d ; // Given two examples to update Ht : 5 : Case 1 : given β ∈ ( 0 , 1 ) , λ > 0 and v0 = 0 , 6 : vt = βvt−1 + ( 1− β ) ∇f ( xt ; ξt ) 2 , Ht = diag ( √ vt + λ ) ; 7 : Case 2 : given β ∈ ( 0 , 1 ) , λ > 0 and b0 = 0 , 8 : bt = βbt−1 + ( 1− β ) ‖∇f ( xt ; ξt ) ‖ , Ht = ( bt + λ ) Id ; 9 : Update x̃t+1 = arg minx∈X { 〈gt , x〉+ 12γ ( x− xt ) THt ( x− xt ) } ; 10 : Update xt+1 = ( 1− µt ) xt + µtx̃t+1 ; 11 : Sample one point ξt+1 , and compute gt+1 = αt+1∇f ( xt+1 ; ξt+1 ) + ( 1 − αt+1 ) [ gt + τ ( ∇f ( xt+1 ; ξt+1 ) −∇f ( xt ; ξt+1 ) ) ] , where τ ∈ { 0 , 1 } ; 12 : end for 13 : Output : ( for theoretical ) xζ chosen uniformly random from { xt } Tt=1 ; ( for practical ) xT . Due to using the coordinate-wise learning rates , these adaptive gradient methods frequently have worse generalization performance than SGD ( with momentum ) [ 35 ] . To improve the generalization performance of Adam , AdamW [ 26 ] uses a decoupled weight decay regularization , defined as mt = α1mt−1 + ( 1− α1 ) ∇f ( xt ; ξt ) , vt = α2vt−1 + ( 1− α2 ) ( ∇f ( xt ; ξt ) ) 2 m̂t = mt/ ( 1− αt1 ) , v̂t = vt/ ( 1− αt2 ) , xt+1 = xt − ηt ( αm̂t/ ( √ v̂t + ε ) + λxt ) , ( 5 ) where α1 , α2 ∈ ( 0 , 1 ) , α > 0 , λ > 0 and ε > 0 . More recently , to further improve generalization performance , AdaBelief [ 40 ] adopts a stepsize according to ‘ belief ’ in the current gradient direction , mt = α1mt−1 + ( 1− α1 ) ∇f ( xt ; ξt ) , vt = α2vt−1 + ( 1− α2 ) ( ∇f ( xt ; ξt ) −mt ) 2 + ε m̂t = mt/ ( 1− αt1 ) , v̂t = vt/ ( 1− αt2 ) , xt+1 = xt − ηtm̂t/ ( √ v̂t + ε ) , ∀ t ≥ 1 ( 6 ) where α1 , α2 ∈ ( 0 , 1 ) , and ηt = η√t with η > 0 , and ε > 0 . At the same time , to improve generalization performance , recently some effective adaptive gradient methods [ 34 , 23 , 11 ] have been proposed with adopting the global adaptive learning rates instead of coordinate-wise counterparts . For example , AdaGrad-Norm [ 34 ] applies a global adaptive learning rate to the following update form , for all t ≥ 1 xt = xt−1 − η∇f ( xt−1 ; ξt−1 ) /bt , b2t = b2t−1 + ‖∇f ( xt−1 ; ξt−1 ) ‖2 , b0 > 0 , ( 7 ) where η > 0 . The adaptive-SGD [ 23 ] adopts a global adaptive learning rate , defined as for all t ≥ 1 ηt = k ( ω + ∑t−1 i=1 ‖∇f ( xi ; ξi ) ‖2 ) 1/2+ε , xt+1 = xt − ηt∇f ( xt ; ξt ) , ( 8 ) where k > 0 , ω > 0 , and ε ≥ 0 . Subsequently , STORM [ 11 ] not only uses a global adaptive learning rate but also adopts the variance-reduced technique in gradient estimator to accelerate algorithm , defined as for all t ≥ 1 ηt = k ( ω + ∑t i=1 ‖∇f ( xi ; ξi ) ‖2 ) 1/3 , xt+1 = xt − ηtgt , ( 9 ) gt+1 = ∇f ( xt+1 ; ξt+1 ) + ( 1− cη2t ) ( gt −∇f ( xt ; ξt+1 ) ) , where k > 0 , ω > 0 and c > 0 .

Summary: This manuscript introduces a unified framework (i.e., SUPER-ADAM) of adaptive gradient methods for nonconvex optimization. It is proved that a universal convergence measure converges with a fast rate, which recovers the $O(\epsilon^{-3})$ complexity by STORM. Another result is that it also recovers $O(\epsilon^{-4})$ result by SGD, when individual function does not have a Lipschitz-continuous gradient. Numerical experiments on image classification and language modeling shows that SUPER-ADAM has faster convergence rate when training deep neural networks.