Adaptive Gradient Methods Can Be Provably Faster than SGD with Random Shuffling

1 INTRODUCTION . We consider the finite sum minimization problem in stochastic optimization : min x2Rd f ( x ) = 1 n nX i=1 fi ( x ) , ( 1 ) where f is the objective function and its component functions fi : Rd ! R are smooth and possibly non-convex . This formulation has been used extensively in training neural networks today . Stochastic gradient descend ( SGD ) and its variants have shown to be quite effective for solving this problem , whereas recent works demonstrate another prominent line of gradient-based algorithms by introducing adaptive step sizes to automatically adjust the learning rate ( Duchi et al. , 2011 ; Tieleman & Hinton , 2012 ; Kingma & Ba , 2014 ) . Despite the superior performance of adaptive gradient methods in many tasks ( Devlin et al. , 2019 ; Vaswani et al. , 2017 ) , their theoretical convergence remains the same or even worse for non-convex objectives , compared to SGD . In general non-convex settings , it is often impractical to discuss optimal solutions . Therefore , the attention of analysis turns to stationary points instead . Many works have been proposed to study first-order ( Chen et al. , 2019 ; Zhou et al. , 2018 ; Zaheer et al. , 2018 ; Ward et al. , 2018 ; Zhou et al. , 2018 ) and second-order ( Allen-Zhu , 2018 ; Staib et al. , 2019 ) stationary points . Table 1 summarized some previous best-known results for finding first-order stationary points . One might notice that the best dependence on the total iteration number T of adaptive gradient methods matches that of vanilla SGD . In addition , with the introduction of incremental sampling techniques , an even better convergence of SGD can be obtained ( Haochen & Sra , 2019 ; Nguyen et al. , 2020 ) . This gap between theory and practice of adaptive gradient methods has been an open problem that we aim to solve in this paper . Motivated by the analysis of sampling techniques , we rigorously prove that adaptive gradient methods exhibit a faster non-asymptotic convergence rate that matches the best result on SGD . In particular , we make the following contributions : • Our main contribution ( Theorem 1,2,3 ) is to prove that two variants of AdaGrad can find Õ ( T 1/2 ) -approximate first-order stationary points under the strong growth condition assumption ( Schmidt & Roux , 2013 ) . This improves previous best convergence results of adaptive gradient methods and shuffling SGD by factors of O ( T 1/4 ) and O ( T 1/6 ) , Our analysis points out two key components that lead to better convergence results of adaptive gradient methods : the epoch-wise analysis of random shuffling can incorporate the benefit of full gradients ; the adaptive learning rates along with the strong growth condition provide better improvement of objective value in consecutive epochs . Finite Sum Minimization vs . Expectation Minimization The comparison in Table 1 shows the convergence rates in the non-convex setting with respect to first-order stationary points . The results in the first two categories apply to the general expectation minimization problem with f ( x ) = Ezf ( x , z ) . Whereas the convergences for expectation minimization naturally transform to finite sum minimization , the statements remain asymptotic , meaning Ekrf ( x ) k ⇠ O ( T↵ ) , where the expectation is taken to compensate the stochastic gradients . Many efforts have been made to reduce variance in finite sum minimization ( Johnson & Zhang , 2013 ; Reddi et al. , 2016 ; Haochen & Sra , 2019 ) . In particular , non-asymptotic results can be gained using random shuffling , under which the dependency on sample size n seems to be unavoidable ( Haochen & Sra , 2019 ) . 2 PRELIMINARIES . A typical setting of machine learning using gradient methods is the finite sum minimization in equation ( 1 ) . In this problem , the number of samples n is usually very large , rendering the evaluation of full gradients expensive . Therefore , a mini-batch gradient is introduced to approximate the full gradient . Mini-batch gradient descent is often carried out in epochs , where each epoch includes several iterations of parameter updates . This epoch-wise implementation can easily incorporate shuffling techniques , which have proven to be effective for SGD both in theory and practice . We aim to analyze the convergence rate of adaptive gradient methods under this framework , where the objective can be non-convex . Throughout this paper , we restrict the discussions of convergence to achieving ✏-approximate first-order stationary point defined as x satisfying krf ( x ) k ✏ . We leave for future work analysis related to saddle points and second-order stationary points . We want to show that adaptive gradient methods can find x such that krf ( x ) k = Õ ( T 1/2 ) in T epochs . Notations . v2 denotes the matrix vv > and kvk is the l2-norm of vector v ; diag ( V ) , kV k , min ( V ) and max ( V ) are the diagonal matrix , the spectral norm , the largest and smallest non-zero eigenvalues of the matirx V , respectively . For alphabets with subscripts , vi : j denotes the collection of { vi , vi+1 , ... , vj } and v : denotes the entire set of v· ; similar notations are used for alphabets with double subscripts . Let [ n ] = { 1 , ... , n } , O ( · ) , Õ ( · ) be the standard asymptotic notations . Denote ei as the unit vector with its i-th component being 1 and e the all-one vector whose dimension depend on the context . As a clarification , we use T to denote the number of epochs ( instead of the number of iterations in Table 1 ) starting from this section . AdaGrad-type methods . As opposed to SGD , adaptive gradient methods assign a coordinate-wise adaptive learning rate to the stochastic gradient . We formulate the generic AdaGrad-type optimizers , including their full and diagonal versions , as follows . At the i-th iteration of epoch t , the parameter is updated by : xt , i+1 = xt , i ⌘tV 1/2t , i gt , i , ( full version ) xt , i+1 = xt , i ⌘tdiag ( Vt , i ) 1/2gt , i , ( diagonal version ) where gt , i is the mini-batch gradient of the objective at xt , i , the matrix Vt , i contains second-moment calculated using all the past stochastic gradients and ⌘t is the step size of epoch t. The initial parameter of epoch t + 1 is taken to be the parameter updated by epoch t , i.e . xt+1,1 = xt , m+1 , where we have m iterations in each epoch . The full version is impractical for high-dimensional x . Thus the diagonal version is often preferred in literature . As an example , the second-moment matrix in AdaGrad is taken to be Vt , i = ( P t 1 s=1 P m j=1 g 2 s , j + P i j=1 g 2 t , j ) /t . SGD can also be written into this general form where we set Vt , i to be the identity matrix . Sampling Strategy . Random shuffling , also known as sampling without replacement , is an oftenused technique to accelerate the convergence of SGD . The idea is to sample a random permutation of function indices [ n ] for each epoch and slide through this permutation to get the mini-batch gradients for the iterations in this epoch . Some implementations shuffle the set [ n ] uniformly independently for each epoch while others shuffle the set once during initialization and use the same permutation for all epochs . Generally speaking , suppose we have a permutation = ( 1 , ... , n ) at epoch t , we define the set Bt , i = { j : ( i 1 ) nm < j i n m } where m is the number of iterations in one epoch . Then the mini-batch gradient is taken to be gt , i = m/n · P j2Bt , i rfj ( xt , i ) . This sampling method of mini-batches benefits the theoretical analysis of SGD by providing a bounded error between the full gradient and the aggregation of mini-batch gradients in one epoch ( Haochen & Sra , 2019 ; Nguyen et al. , 2020 ) . A naive observation that backups this point can be made by assuming xt,1 = ... = xt , m , since [ mi=1Bt , i = [ n ] , we would have P m i=1 gt , i = rf ( xt,1 ) . Then full gradient can be used to obtain convergence better than plain SGD . Random shuffling for AdaGrad . Unlike SGD , in adaptive methods , it is hard to approximate the full gradient with the aggregation of mini-batch gradient updates in one epoch due to the presence of the second moments . As we will show in experiments , the simple shuffling variant that only changes the sampling method of mini batches in AdaGrad does not lead to better convergence . The major difficulty hampering the analysis of this variant is that the second-moment matrix uses all the gradient information in history without distinguishment . Thus to be able to leverage the benefit of the full gradient , we propose to study a slight modification of AdaGrad . Formally , we shuffle the set [ n ] once at initialization and obtain the mini-batch gradients in a random shuffling manner . We update the parameters by the same rules of AdaGrad-type methods described above where the second-moment matrix is taken to be : Vt , i = mX j=i+1 g 2 t 1 , j + iX j=1 g 2 t , j . ( AdaGrad-window ) The difference between AdaGrad-window and AdaGrad is that the former only use the latest m mini-batch gradients instead of an epoch-wise average of all the mini-batch gradients in history . The step size is ⌘t = ⌘/ p t where ⌘ is a constant for both methods . The updates of AdaGrad-window is also very similar to the GGT method ( Agarwal et al. , 2019 ) without momentum . However , GGT uses the full matrix inversion , where our analysis applies to both full and diagonal versions . 3 MAIN RESULTS . We will show that AdaGrad-window has the convergence rate of Õ ( T 1/2 ) for non-convex problems under some mild assumptions . This is a significant improvement compared with previous best convergence results of adaptive gradient methods and random shuffling SGD , which are of order O ( T 1/4 ) and Õ ( T 1/3 ) respectively . The key towards our convergence rate improvement is twofold : the epoch-wise analysis of random shuffling enables us to leverage the benefit of full gradients ; the adaptive learning rates and the strong growth condition endow a better improvement of objective value in consecutive epochs . In order to achieve this better convergence , we first state the assumptions and important concepts used in the proof . Apart from the general assumptions ( A1 ) and ( A2 ) used in previous analysis ( Fang et al. , 2018 ; Zaheer et al. , 2018 ; Ward et al. , 2018 ) , we pose another assumption described below in ( A3 ) to characterize the consistency between individual gradients and the full gradient . Assumptions . We assume the following for AdaGrad-window : ( A1 ) The objective function is lower bounded and component-wise L-smooth , i.e . 9f⇤ 2 R s.t . f ( x ) f⇤ > 1 , 8x and krfi ( x ) rfi ( y ) k L kx yk , 8x , y , i . ( A2 ) The mini-batch gradients in the algorithm are uniformly upper bounded , i.e . 9G 2 R s.t . kgt , ik G , 8t , i . ( A3 ) The objective function satisfies the strong growth condition with constant r2 , i.e . 8x , 1 n P n i=1 krfi ( x ) k2 r2krf ( x ) k2 . The strong growth condition assumption is essentially enforcing the norms of individual gradients to be at the same scale as the norm of the full gradient . This condition was originally used to derive a faster convergence for SGD in the context of convex finite sum minimization ( Schmidt & Roux , 2013 ) . It was further explored to analyze SGD-type methods after showing its closed relationship with interpolation ( Ma et al. , 2018 ; Vaswani et al. , 2019 ; Gower et al. , 2019 ) . Built upon these previous analysis , we will give a more in-depth discussion for this assumption in section 6 . Under these assumptions , the following theorems show our convergence result for the full and diagonal versions of AdaGrad-window . Theorem 1 ( The convergence rate of full AdaGrad-window ) . For any T > 4 , set ⌘ = m 5/4 , denote C1 = m5/4 p 11/3 + 8 · r2 ( f ( x1,1 ) f⇤ +G ) / p 2 and C2 = 5m5/4 p 11/3 + 8 · r2L/ p 2 as constants independent of T . We have : min 1tT krf ( xt,1 ) k 1p T ( C1 + C2 lnT ) . ( 2 ) Theorem 2 ( The convergence rate of diagonal AdaGrad-window ) . For any T > 4 , set ⌘ = m 5/4 , denote C 01 = m5/4 p 11/3 + 8 · r2 ⇣ f ( x1,1 ) f⇤ +G p d ⌘ / p 2 and C 02 = 5m5/4 p 11/3 + 8 · r2d3/2L/ p 2 as constants independent of T . We have : min 1tT krf ( xt,1 ) k 1p T ( C 01 + C 0 2 lnT ) . ( 3 ) The interpretation of these two theorems is that we are able to find an approximate first-order stationary point such that krf ( x ) k = Õ ( T 1/2 ) within T epochs using both versions . We notice that the convergence rate of AdaGrad-window matches that of GD when m = 1 , which denotes that our results are relatively tight with respect to T . The complete proof is included in the appendix . We will give the intuition and key lemmas explaining how to utilize random shuffling and second moments to obtain these results in the next section . In addition , we also prove for another variant of AdaGrad , namely , AdaGrad-truncation with secondmoment matrix defined as V1 , i = m · I and Vt , i = mk P m j=1 gt 1 , jk2 · I when t > 1 . This second-moment matrix is very similar to the norm version of AdaGrad ( Ward et al. , 2018 ) whereas we use the aggregation of mini-batch gradients in the previous epoch as the coefficient . AdaGradtruncation is beneficial since the formulation leads to a fast and simple implementation without needing to discuss the full and diagonal versions . Due to the space limitation , we list the result below and defer the discussions to the appendix . Theorem 3 ( The convergence rate of AdaGrad-truncation ) . For any T > 4 , set ⌘ =p 3/ ( 10m1/2Lr ) · p f ( x1,1 f⇤ +G2 ) / p L+ 2 , denote C = 80m ( Lr+2 ) p f ( x1,1 ) f⇤ +G2 as constants independent of T . We have : min 1tT krf ( xt,1 ) k Cp T . ( 4 )

Summary: This paper shows that adaptive learning rates are beneficial for finding critical points of finite-sum optimization problems. In particular, with appropriate learning rates, a variant of adagrad can find a epsilon critical point in \tilde O(1/epsilon^2) iterations. This improves upon previous results of either O(1/\epsilon^3) or O(1/\epsion^4) in various situations. The key new assumption is a “consistency condition” that bounds how big individual example gradients can be when the overall gradient is small.