Adam$^+$: A Stochastic Method with Adaptive Variance Reduction

1 INTRODUCTION . Adaptive gradient methods ( Duchi et al. , 2011 ; McMahan & Streeter , 2010 ; Tieleman & Hinton , 2012 ; Kingma & Ba , 2014 ; Reddi et al. , 2019 ) are one of the most important variants of Stochastic Gradient Descent ( SGD ) in modern machine learning applications . Contrary to SGD , adaptive gradient methods typically require little parameter tuning still retaining the computational efficiency of SGD . One of the most used adaptive methods is Adam ( Kingma & Ba , 2014 ) , which is considered by practitioners as the de-facto default optimizer for deep learning frameworks . Adam computes the update for every dimension of the model parameter through a moment estimation , i.e. , the estimates of the first and second moments of the gradients . The estimates for first and second moments are updated using exponential moving averages with two different control parameters . These moving averages are the key difference between Adam and previous adaptive gradient methods , such as Adagrad ( Duchi et al. , 2011 ) . Although Adam exhibits great empirical performance , there still remain many mysteries about its convergence . First , it has been shown that Adam may not converge for some objective functions ( Reddi et al. , 2019 ; Chen et al. , 2018b ) . Second , it is unclear what is the benefit that the moving average brings from theoretical point of view , especially its effect on the convergence rate . Third , it has been empirically observed that adaptive gradient methods can have worse generalization performance than its non-adaptive counterpart ( e.g. , SGD ) on various deep learning tasks due to the coordinate-wise learning rates ( Wilson et al. , 2017 ) . The above issues motivate us to design a new algorithm which achieves the best of both worlds , i.e. , provable convergence with benefits from the moving average and enjoying good generalization performance in deep learning . Specifically , we focus on the following optimization problem : min w∈Rd F ( w ) , where we only have access to stochastic gradients of F . Note that F could possibly be nonconvex in w. Due to the non-convexity , our goal is to design a stochastic first-order algorithm to find the -stationary point , i.e. , finding w such that E [ ‖∇F ( w ) ‖ ] ≤ , with low iteration complexity . Our key contribution is the design and analysis of a new stochastic method named Adam+ . Adam+ retains some of the key components of Adam but it also has several noticeable differences : ( i ) it does not maintain the moving average of second moment estimate but instead computes the moving average of first moment estimate at extrapolated points ; ( ii ) its adaptive step size is formed not by dividing the square root of coordinate-wise second moment estimate but instead by dividing the root of the norm of first moment estimate . These features allow us to establish the adaptive convergence of Adam+ . Different from existing adaptive methods where the adaptive convergence depends on the growth rate of stochastic gradients ( Duchi et al. , 2011 ; McMahan & Streeter , 2010 ; Kingma & Ba , 2014 ; Luo et al. , 2019 ; Reddi et al. , 2019 ; Chen et al. , 2018a ; b ; Ward et al. , 2019 ; Li & Orabona , 2019 ; Chen et al. , 2019 ) , our adaptive convergence is due to the adaptive variance reduction property of our first order moment estimate . In existing literature , the variance reduction is usually achieved by large mini-batch ( Goyal et al. , 2017 ) or recursive variance reduction ( Fang et al. , 2018 ; Zhou et al. , 2018 ; Pham et al. , 2020 ; Cutkosky & Orabona , 2019 ) . In contrast , we do not necessarily require large minibatch or computing stochastic gradients at two points per-iteration to achieve the variance reduction . In addition , we also establish a fast rate that matches the state-of-theart complexity under the same conditions of a variant of Adam+ . Table 1 provides an overview of our results and a summary of existing results . We refer readers to Section F for a comprehensive survey of other related work . We further corroborate our theoretical results with an extensive empirical study on various deep learning tasks . Our contributions are summarized below . • We propose a new algorithm with adaptive step size , namely Adam+ , for general nonconvex optimization . We show that it enjoys a new type of data-dependent adaptive convergence that depends on the variance reduction property of first moment estimate . Notably , this data-dependent complexity does not require the presence of sparsity in stochastic gradients to guarantee fast convergence as in previous works ( Duchi et al. , 2011 ; Kingma & Ba , 2014 ; Reddi et al. , 2019 ; Chen et al. , 2019 ; 2018a ) . To the best of our knowledge , this is the first work establishing such new type of data-dependent complexity . • We show that a general variant of our algorithm can achieve O ( −3.5 ) worst-case complexity , which matches the state-of-the-art complexity guarantee under the Hessian Lipschitz assumption ( Cutkosky & Mehta , 2020 ) . • We demonstrate the effectiveness of our algorithms on image classification , language modeling , and automatic speech recognition . Our empirical results show that our proposed algorithm consistently outperforms Adam on all tasks , and it achieves comparable performance with the best-tuned SGD and momentum SGD . Algorithm 1 Adam+ : Good default settings for the tested machine learning problems are α = 0.1 , a = 1 , β = 0.1 , 0 = 10−8 . 1 : Require : α , a ≥ 1 : stepsize parameters 2 : Require : β ∈ ( 0 , 1 ) : Exponential decay rates for the moment estimate 3 : Require : gt ( w ) : unbiased stochastic gradient with parameters w at iteration t 4 : Require : w0 : Initial parameter vector 5 : z0 = g0 ( w0 ) 6 : for t = 0 , . . . , T do 7 : Set ηt = αβ a max ( ‖zt‖1/2 , 0 ) 8 : wt+1 = wt − ηtzt 9 : ŵt+1 = ( 1− 1/β ) wt + 1/β ·wt+1 10 : zt+1 = ( 1− β ) zt + βgt+1 ( ŵt+1 ) 11 : end for 2 ALGORITHM AND THEORETICAL ANALYSIS . In this section , we introduce our algorithm Adam+ ( presented in Algorithm 1 ) and establish its convergence guarantees . Adam+ resembles Adam in several aspects but also has noticeable differences . Similar to Adam , Adam+ also maintains an exponential moving average of first moment ( i.e. , stochastic gradient ) , which is denoted by zt , and uses it for updating the solution in line 8 . However , the difference is that the stochastic gradient is evaluated on an extrapolated data point ŵt+1 , which is an extrapolation of two previous updates wt and wt+1 . Similar to Adam , Adam+ also uses an adaptive step size that is proportional to 1/‖zt‖1/2 . Nonetheless , the difference lies at its adaptive step size is directly computed from the square root of the norm of first moment estimate zt . In contrast , Adam uses an adaptive step size that is proportional to 1/ √ vt , where vt is an exponential moving average of second moment estimate . These two key components of Adam+ , i.e. , extrapolation and adaptive step size from the root norm of the first moment estimate , make it enjoy two noticeable benefits : variance reduction of first moment estimate and adaptive convergence . We shall explain these two benefits later . Before moving to the theoretical analysis , we would like to make some remarks . First , it is worth mentioning that the moving average estimate with extrapolation is inspired by the literature of stochastic compositional optimization ( Wang et al. , 2017 ) . Wang et al . ( 2017 ) showed that the extrapolation helps balancing the noise in the gradients , reducing the bias in the estimates and giving a faster convergence rate . Here , our focus and analysis techniques are quite different . In fact , Wang et al . ( 2017 ) focuses on the compositional optimization while we consider a general nonconvex optimization setting . Moreover , the analysis in ( Wang et al. , 2017 ) mainly deals with the error of the gradient estimator caused by the compositional nature of the problem , while our analysis focuses on carefully designing adaptive normalization to obtain an adaptive and fast convergence rates . A similar extrapolation scheme has been also employed in the algorithm NIGT by Cutkosky & Mehta ( 2020 ) . In later sections , we will also provide a more general variant of Adam+ which subsumes NIGT as a special case . Another important remark is that the update of Adam+ is very different from the famous Nesterov ’ s momentum method . In Nesterov ’ s momentum method , the update of wt+1 uses the stochastic gradient at an extrapolated point ŵt+1 = wt+1 +γ ( wt+1−wt ) with a momentum parameter γ ∈ ( 0 , 1 ) . In contrast , in Adam+ the update of wt+1 is using the moving average estimate at an extrapolated point ŵt+1 = wt+1 + ( 1/β − 1 ) ( wt+1 −wt ) . Finally , Adam+ does not employ coordinate-wise learning rates as in Adam , and hence it is expected to have better generalization performance according to Wilson et al . ( 2017 ) . 2.1 ADAPTIVE VARIANCE REDUCTION AND ADAPTIVE CONVERGENCE . In this subsection , we analyze Adam+ by showing its variance reduction property and adaptive convergence . To this end , we make the following assumptions . Assumption 1 . There exists positive constants L , ∆ , LH , σ and an initial solution w0 such that ( i ) F is L-smooth , i.e. , ‖∇F ( x ) −∇F ( y ) ‖ ≤ L ‖x− y‖ , ∀x , y ∈ Rd . ( ii ) For ∀x ∈ Rd , we have access to a first-order stochastic oracle at time t gt ( x ) such that E [ gt ( x ) ] = ∇F ( x ) , E ‖gt ( x ) −∇F ( x ) ‖2 ≤ σ2 . ( iii ) ∇F is a LH -smooth mapping , i.e. , ‖∇2F ( x ) −∇2F ( y ) ‖ ≤ LH‖x− y‖ , ∀x , y ∈ Rd . ( iv ) F ( w0 ) − F∗ ≤ ∆ < ∞ , where F∗ = infw∈Rd F ( w ) . Remark : Assumption 1 ( i ) and ( ii ) , ( iv ) are standard assumptions made in literature of stochastic non-convex optimization ( Ghadimi & Lan , 2013 ) . Assumption ( iii ) is the assumption that deviates from typical analysis of stochastic methods . We leverage this assumption to explore the benefit of moving average , extrapolation and adaptive normalization . It is also used in some previous works for establishing fast rate of stochastic first-order methods for nonconvex optimization ( Fang et al. , 2019 ; Cutkosky & Mehta , 2020 ) and this assumption is essential to get fast rate due to the hardness result in ( Arjevani et al. , 2019 ) . It is also the key assumption for finding a local minimum in previous works ( Carmon et al. , 2018 ; Agarwal et al. , 2017 ; Jin et al. , 2017 ) . We might also assume that the stochastic gradient estimator in Algorithm 1 satisfies the following variance property . Assumption 2 . Assume that E [ ‖g0 ( w0 ) − ∇F ( w0 ) ‖2 ] ≤ σ20 and E [ ‖gt ( wt ) − ∇F ( wt ) ‖2 ] ≤ σ2m , t ≥ 1 . Remark : When g0 ( resp . gt ) is implemented by a mini-batch stochastic gradient with mini-batch size S , then σ20 ( resp . σ 2 m ) can be set as σ 2/S by Assumption 1 ( ii ) . We differentiate the initial variance and intermediate variance because they contribute differently to the convergence . We first introduce a lemma to characterize the variance of the moving average gradient estimator zt . Lemma 1 . Suppose Assumption 1 and Assumption 2 hold and a ≥ 1 . Then , there exists a sequence of random variables δt satisfying ‖zt −∇F ( wt ) ‖ ≤ δt for ∀t ≥ 0 , E [ δ2t+1 ] ≤ ( 1− β 2 ) E [ δ2t ] + 2β2σ2m + E [ CL2H‖wt+1 −wt‖4 β3 ] ≤ ( 1− β 2 ) E [ δ2t ] + 2β2σ2m + E [ CL2Hα 4β4a−3‖zt‖2 ] , where C = 1944 . Remark : Note that δt is an upper bound of ‖zt − ∇F ( wt ) ‖ , the above lemma can be used to illustrate the variance reduction effect for the gradient estimator zt . To this end , we can bound ‖zt‖2 ≤ 2δ2t + 2‖∇F ( wt ) ‖2 , then the term CL2α4β4a−3δ2t can be canceled with −β/4δ2t with small enough α . Hence , we have Eδ2t+1 ≤ ( 1 − β/4 ) E [ δ2t ] + 2β2σ2m + cE [ ‖∇F ( wt ) ‖2 ] with a small constant c. As the algorithm converges with E [ ‖∇F ( wt ) ‖2 ] and β decreases to zero , the variance of zt will also decrease . Indeed , the above recursion of zt ’ s variance resembles that of the recursive variance reduced gradient estimators ( e.g. , SPIDER ( Fang et al. , 2018 ) , STORM ( Cutkosky & Orabona , 2019 ) ) . The benefit of using Adam+ is that we do not need to compute stochastic gradient twice at each iteration . We can now state our convergence rates for Algorithm 1 . Theorem 1 . Suppose Assumption 1 and Assumption 2 hold . Suppose ‖∇F ( w ) ‖ ≤ G for any w ∈ Rd . By choosing the parameters such that α4 ≤ 1 36CL2H , α ≤ 14L , a = 1 and 0 = β a , we have 1 T T∑ t=1 E ‖∇F ( wt ) ‖2 ≤ GE [ ∑T t=1 ‖zt‖ ] T + ∆ αT + 18σ20 βT + 30βσ2m . ( 1 ) In addition , suppose the initial batch size is T0 and the intermediate batch size is m , and choose β = T−b with 0 ≤ b ≤ 1 , we have 1 T T∑ t=1 E ‖∇F ( wt ) ‖2 ≤ E [ G ∑T t=1 ‖zt‖ ] T + ∆ αT + 18σ2 T 1−bT0 + 30σ2 mT b . ( 2 ) Theorem 2 . Suppose Assumption 1 and Assumption 2 hold . By choosing parameters such that 640α3L 3/2 H ≤ 1/120 , a = 1 , 0 = 0 , β = 1/T s with s = 2/3 then it takes T = O ( −4.5 ) number of iterations to ensure that 1 T T∑ t=1 E [ ‖∇F ( wt ) ‖3/2 ] ≤ 3/2 , 1 T E [ T∑ t=1 δ 3/2 t ] ≤ 3/2 . Remarks : • From Theorem 1 , we can observe that the convergence rate of Adam+ crucially depends on the growth rate of E [ ∑T t=1 ‖zt‖ ] , which gives a data-dependent adaptive complexity . If E [ ∑T t=1 ‖zt‖ ] ≤ Tα with α < 1 , then the algorithm converges . Smaller α implies faster convergence . Our goal is to ensure that 1T ∑T t=1 E‖∇F ( wt ) ‖2 ≤ 2 . Choosing b = 1−α , m = O ( 1 ) and T0 = T 1−α = O ( −2 ) , and we end up with T = O ( − 2 1−α ) complexity . • Theorem 2 shows that in the ergodic sense , the Algorithm Adam+ always converges , and the variance gets smaller when the number of iteration gets larger . Theorem 2 rules out the case that the magnitude of zt converges to a constant and the bound ( 2 ) in Theorem 1 becomes vacuous . • To compare with Adam-style algorithms ( e.g. , Adam , AdaGrad ) , these algorithms ’ convergence depend on the growth rate of stochastic gradient , i.e. , ∑d i=1 ‖g1 : T , i‖/T , where g1 : T , i = [ g1 , i , . . . , gT , i ] denotes the i-th coordinate of all historical stochastic gradients . Hence , the data determines the growth rate of stochastic gradient . If the stochastic gradients are not sparse , then its growth rate may not be slow and these Adam-style algorithms may suffer from slow convergence . In contrast , for Adam+ the convergence can be accelerated by the variance reduction property . Note that we have E [ ∑T t=1 ‖zt‖ ] /T ≤ E [ ∑T t=1 ( δt + ‖∇F ( wt ) ‖ ) ] /T . Hence , Adam+ ’ s convergence depends on the variance reduction property of zt . 2.2 A GENERAL VARIANT OF ADAM+ : FAST CONVERGENCE WITH LARGE MINI-BATCH Next , we introduce a more general variant of Adam+ by making a simple change . In particular , we keep all steps the same as in Algorithm 1 except the adaptive step size is now set as ηt = αβa max ( ‖zt‖p , 0 ) , where p ∈ [ 1/2 , 1 ) is parameter . We refer to this general variant of Adam + as power normalized Adam+ ( Nadam+ ) . This generalization allows us to compare with some existing methods and to establish fast convergence rate . First , we notice that when setting p = 1 and a = 5/4 and β = 1/T 4/7 , Nadam+ is almost the same as the stochastic method NIGT ( Cutkosky & Mehta , 2020 ) with only some minor differences . However , we observed that normalizing by ‖zt‖ leads to slow convergence in practice , so we are instead interested in p < 1 . Below , we will show that NAdam+ with p < 1 can achieve a fast rate of 1/ 3.5 , which is the same as NIGT . Theorem 3 . Under the same assumption as in Theorem 1 , further assume σ20 = σ2/T0 and σ2m = σ2/m . By using the step size ηt = αβ 4/3 max ( ‖zt‖2/3 , 0 ) in Algorithm 1 with CL2α4 ≤ 1/14 , 0 = 2β4/3 , in order to have E [ ‖∇F ( wτ ) ‖ ] ≤ for a randomly selected solution wτ from { w1 , . . . , wT } , it suffice to set β = O ( 1/2 ) , T = O ( −2 ) , the initial batch size T0 = 1/β = O ( −1/2 ) , the intermediate batch size as m = 1/β3 = O ( −3/2 ) , which ends up with the total complexity O ( −3.5 ) . Remark : Note that the above theorem establishes the fast convergence rate for Nadam+ with p = 2/3 . Indeed , we can also establish a fast rate of Adam+ ( where p = 1/2 ) in the order ofO ( 1/ 3.625 ) with details provided in the Appendix E .

Summary: This paper proposes the Adam+ algorithm that maintains an exponential moving average of the first moment and normalizes it by its $p$-th moment for some $p \in (1/2, 1)$. When $p = 2/3$, with appropriate hyperparameters, Adam+ achieves the state-of-the-art complexity $O(1/\epsilon^{3.5})$ to obtain an approximate first-order stationary point for smooth objectives with smooth Hessians. The proof technique is similar to that for SCGD by Wang et al. (2017), which construct a Lyapunov function for algorithms of this kind.