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 a new optimizer called Adam+, with two main distinctions from standard Adam template: 1) the first order moment estimate is computed using the gradient evaluated at an extrapolated iterate. 2) the step size is scaled with the square root of the norm of the first order moment, rather than the exponential moving average (EMA) in the previous work. Under Lipschitz continuous gradient, Hessian and bounded gradient assumptions, the complexity for finding a point with small gradient norm is $\epsilon^{-3.5}$, which is a SOTA complexity. The practical performance of the algorithm is evaluated in a range of different tasks and the performance is shown to be consistently promising and comparable to SGD/Momentum SGD.