The Case for Full-Matrix Adaptive Regularization

1 INTRODUCTION . Stochastic gradient descent is the workhorse behind the recent deep learning revolution . This simple and age-old algorithm has been supplemented with a variety of enhancements to improve its practical performance , and sometimes its theoretical guarantees . Amongst the acceleration methods there are three main categories : momentum , adaptive regularization , and variance reduction . Momentum ( in its various incarnations , like heavy-ball or Nesterov acceleration ) is the oldest enhancement . It has a well-developed theory , and is known to improve practical convergence in a variety of tasks , small and large . It is also easy to implement . Variance reduction is the most recent advancement ; in theory and practice , it is mostly applicable to convex optimization , and is thus less influential in deep learning . This brings us to adaptive regularization : the most sophisticated , hard to implement , and debated acceleration method . While state-of-the-art optimizers such as Adam and AdaGrad ( Kingma & Ba , 2014 ; Duchi et al. , 2011 ) do use adaptive regularization , they do so in a very limited form : with diagonal matrices , often marketed as per-coordinate adaptive learning-rate methods . Despite solid theoretical guarantees , the practical value of diagonal adaptive regularization as compared to “ vanilla ” SGD has been the subject of much debate ( Wilson et al. , 2017 ) . However , the efficacy of full-matrix adaptive regularization has been relatively unexplored . This is due to the prohibitive computational cost associated with full-matrix operations : full AdaGrad requires taking the inverse square root of a large matrix . In this paper , we present GGT , a practical solution to the computational problems plaguing fullmatrix adaptive regularization , making this technique scalable for modern deep models . At the heart of our method is a simple , GPU-friendly way to apply the inverse square root of the low-rank second-moment matrix of recent gradients ; see Figure 1 . GGT ’ s running time is comparable to state-of-the-art optimizers . We proceed to show that full-matrix preconditioning allows for much better exploitation of anisotropic curvature in loss landscapes . First , we show synthetic experiments which demonstate clear benefits of GGT over baselines , especially when the problem is ill-conditioned . Then , we implement GGT at scale , and show that the benefits translate to faster training on standard deep learning benchmarks . Our improvement is most salient in complicated landscapes like RNN training . Our algorithm comes with theoretical guarantees . We give the first proof of convergence to firstorder critical points for an algorithm with adaptive regularization in a stochastic non-convex setting , featuring a rate which is dependent on an adaptive ratio . We show examples where our bound is stronger than that for SGD , providing some theoretical basis for our empirical findings . 1.1 RELATED WORK . Since the introduction of AdaGrad ( Duchi et al. , 2011 ) , diagonal adaptive regularization has been a mainstay in the machine learning practitioner ’ s toolbox . A quick perusal of the literature shows that these methods have continued to thrive in the deep learning era , and appear in all major frameworks ( Abadi et al. , 2016 ; Paszke et al. , 2017 ; Chen et al. , 2015 ) . By citation count ( or GitHub search hits ) , Adam ( Kingma & Ba , 2014 ) is by far the most popular adaptive optimizer for training a variety of modern deep models . For this reason , this paper ’ s exposition is targeted towards a full-matrix dropin replacement for Adam ; however , our techniques extend straightforwardly to a plethora of variants , like RMSprop ( Tieleman & Hinton , 2012 ) , Adadelta ( Zeiler , 2012 ) , Nadam ( Dozat , 2016 ) , etc . Full-matrix adaptive regularization has existed alongside the more commonly used diagonal-matrix manifestation since their common inception in ( Duchi et al. , 2011 ) ; however , a major obstacle to the scalability of these methods is the need for the storage and inversion of square matrices in the model dimension . This becomes prohibitively expensive in dimension greater than 104 , while state-of-theart models regularly exceed 107 parameters . Matrix sketching has been employed to approximate the AdaGrad preconditioner ( Krummenacher et al. , 2016 ; Mehta et al. , 2016 ) ; however , the sketched estimate for the matrix inverse can be sensitive to noise . In the former , the authors report a 5-10× overhead over AdaGrad , even with < 105 model parameters ; we could not find a usable GPU implementation for their requisite rank-1 QR update . ( Gupta et al. , 2018 ) propose a way to do AdaGrad with Kronecker products of full-matrix preconditioners , a more limited setting which requires knowledge of the model ’ s structure . Finally , as we argue in Section 3.1 , there is intrinsic value of “ forgetting ” past curvature using an exponential window . With this , a low-rank preconditioning matrix naturally arises , allowing us to bypass the computational need for sketching in the model dimension or architecture-dependent restriction of the preconditioner . Our algorithm bears a superficial resemblance to L-BFGS ( Liu & Nocedal , 1989 ) , a version of BFGS ( Broyden , 1970 ; Fletcher , 1970 ; Goldfarb , 1970 ; Shanno , 1970 ) which uses a sliding window of gradient history . Although some are viable for large-scale implementation , these quasi-Newton methods , along with ( subsampled , online , cubic-regularized ) Newton methods ( Erdogdu & Montanari , 2015 ; Agarwal et al. , 2017b ; Luo et al. , 2016 ; Hazan et al. , 2007 ; Agarwal et al. , 2017a ; Carmon et al. , 2017 ) exhibit very different dynamics than the standard optimizers in deep learning , and thus have not seen widespread adoption . We find recent deep learning applications of secondorder methods ( e.g . ( Martens & Grosse , 2015 ; Martens et al. , 2018 ) ) to be intriguing , though outside the scope of this paper . Recently , the role of adaptive regularization has been a hotly contested topic . In ( Wilson et al. , 2017 ) , the authors suggest that properly-tuned SGD exhibits superior generalization to adaptive methods . In turn , ( Keskar & Socher , 2017 ) propose switching the optimizer from Adam to SGD at the end of training , to reap the advantages of each . Influentially , Adam ’ s convergence has been the object of recent scrutiny ( Reddi et al. , 2018 ) . However , Adam continues to enjoy successful convergence in practice ; the problematic construction involves pathological outlier gradients . We do not use the analyses of Adam or AMSGrad . Several parallel works ( Li & Orabona , 2018 ; Zou & Shen , 2018 ; Ward et al. , 2018 ; Chen et al. , 2018a ; b ; Zhou et al. , 2018 ) have studied the convergence of adaptive methods for non-convex optimization , matching the asymptotic iteration complexity of SGD . Apart from our algorithmic contribution , our work is ( to our knowledge ) the first attempt to characterize the advantage of adaptivity in terms of the dimension and geometry of the optimization problem . 2 THE GGT ALGORITHM . Our main algorithmic contribution is GGT , an efficient first-order algorithm for full-matrix adaptive preconditioning . In brief , GGT uses the preconditioner from full-matrix AdaGrad , with gradient history attenuated exponentially as in Adam , and truncated to a window parameter r. The name GGT acts as a convenient mnemonic for the gradient second-moment matrix GG > maintained by full-matrix AdaGrad , even though we never compute this matrix . The mathematical specification of GGT is given in Algorithm 1 , in the usual model of stochastic optimization ( see Section 4 ) , with gradients ∇̃f ( x ) . Notice that the coordinate-wise scaling of Adam is recovered by zeroing out the off-diagonal entries of GG > . Algorithm 1 GGT adaptive optimizer 1 : Input : initializer x1 , window size r , learning rate schedule { ηt } , β2 ≤ 1 , ε > 0 . 2 : for t = 1 , . . . , T do 3 : Receive stochastic gradient ∇̃f ( xt ) . 4 : Let Gt = [ gt gt−1 . . . gt−r+1 ] , where gt−k : = βk2 ∇̃f ( xt−k ) , or 0 if k ≥ t. 5 : Update xt+1 ← xt − ηt · [ ( GtG > t ) 1/2 + εI ] −1 ∇̃f ( xt ) . 6 : end for GGT provides the power of full-matrix adaptive regularization at a cost not much larger than SGD . This crucially exploits the fact only a small window of historical gradients are used for preconditioning . The intuition for using a small window , as opposed to the entire history , is clear ( and time-tested , by the ubiquity of Adam ) : the curvature of the loss surface changes , rendering previous gradient information obsolete . We expand on the benefits of forgetting gradients in section 3.1 . The fact that the preconditioning matrix is based on a small window of gradients implies that it has low rank . GGT exploits this fact by computing the inverse square root of the empirical covariance matrix indirectly , as outlined in Figure 1 . In effect , instead of inverting a full matrix in the dimension of parameters , using the special matrix structure GGT inverts a matrix of dimension window-size . The remainder of this section will discuss efficient implementation and some heuristics . GGT has provable guarantees even for non-convex optimization : it is guaranteed to converge to a first-order critical point . Its rate of convergence is never significantly slower than that of SGD , and in some favorable geometric conditions , can be significantly faster . These theoretical bounds are made precise in section 4 . 2.1 FAST LOW-RANK PRECONDITIONING . The window parameter r should be roughly the number of copies of the model that fit in RAM ; in our large-scale experiments , we use r = 200 . A pessimistic but principled choice is r = Θ ( 1/ ( 1−β2 ) ) , which truncates on the time scale of the exponential attenuation . Our key observation , highlighted in Figure 1 , is that the inversion of the large low-rank matrix GG > can be performed by diagonalizing the small matrix G > G , along with some extremely GPU-friendly matrix-vector operations . The basic intuition is contained in Figure 1 , but it remains to include the εI term . We derive the full update here . Let G ∈ Rd×r , v ∈ Rd be arbitrary , with r ≤ d. Write the singular value decomposition G = UΣV > , with U ∈ Rd×d , Σ ∈ Rd×r , V ∈ Rr×r . Let Σd ∈ Rd×d : = [ Σ 0 ] , and let Σr ∈ Rr×r be its top left block . Let U = : [ Ur Ud−r ] , so that the columns of Ur ∈ Rd×r are an orthonormal basis for the column space of G , and Ud−r ∈ Rd× ( d−r ) its orthogonal component , noting that UrU > r + Ud−rU > d−r = Id . Then , we have [ ( GG > ) 1/2 + εI ] −1 v = [ ( UΣ2dU > ) 1/2 + εUU > ] −1 v = U ( Σd + εI ) −1U > v = [ Ur ( Σr + εIr ) −1U > r + Ud−r ( εId−r ) −1U > d−r ] v = Ur ( Σr + εIr ) −1U > r v + 1 ε ( Id −UrU > r ) v = 1 ε v + Ur [ ( Σr + εIr ) −1 − 1 ε Ir ] U > r v. The first term is none other than an SGD update step . The rest can be computed by taking the eigendecomposition G > G = VΣ2rV > , giving Ur = GV √ Σr † . We prefer this to taking the direct SVD of G , which is > 10 times slower on GPU . Using a cyclic buffer to store and update Gt , the algorithm takes O ( dr2 + r3 ) ( sequential ) time per iteration , and O ( dr ) memory in total . Iterating over the model parameters to update Gt incurs the same overhead cost as usual adaptive optimizers . The r × d matrix multiplication and r × r SVD operations benefit from decades of extensive hardware-level optimizations . In the experiments in Section 3 , we observed a ∼ 1.3× ( CNN ) and ∼ 2× ( RNN ) running-time overhead over SGD ; we note that this ratio could be even smaller in reinforcement learning ( where the environment causes the time bottleneck ) , or universally with a more optimized implementation .

Summary: adaptive versions of sgd are commonly used in machine learning. adagrad, adadelta are both popular adaptive variations of sgd. These algorithms can be seen as preconditioned versions of gradient descent where the preconditioner applied is a matrix of second-order moments of the gradients. However, because this matrix turns out to be a pxp matrix where p is the number of parameters in the model, maintaining and performing linear algebra with this pxp matrix is computationally intensive. In this paper, the authors show how to maintain and update this pxp matrix by storing only smaller matrices of size pxr and rxr, and performing 1. an SVD of a small matrix of size rxr 2. matrix-vector multiplication between a pxr matrix and rx1 vector. Given that rxr is a small constant sized matrix and that matrix-vector multiplication can be efficiently computed on GPUs, this matrix adapted SGD can be made scalable. The authors also discuss how to adapt the proposed algorithm with Adam style updates that incorporate momentum. Experiments are shown on various architectures (CNN, RNN) and comparisons are made against SGD, ADAM.