Towards Practical Second Order Optimization for Deep Learning

1 Introduction . Second order methods are among the most powerful algorithms in mathematical optimization . Algorithms in this family often use a preconditioning matrix to transform the gradient before applying each step . Classically , the preconditioner is the matrix of second-order derivatives ( i.e. , the Hessian ) in the context of exact deterministic optimization ( e.g. , Fletcher , 2013 ; Lewis & Overton , 2013 ; Nocedal , 1980 ) . While second-order methods often have significantly better convergence properties than first-order methods , the size of typical problems prohibits their use in practice , as they require quadratic storage and cubic computation time for each gradient update . Approximate algorithms such as quasi-Newton methods are aimed at significantly reducing these requirements ; nonetheless , they still impose non-trivial memory costs equivalent to storing several copies of the model ( and often quadratic computation , as in the popular two-loop recursion ( Nocedal , 1980 ) ) , which severely limits their use at the immense scale of present-day deep learning . Arguably , one of the greatest challenges of modern optimization is to bridge this gap between theoretical and practical optimization towards making second-order methods feasible to implement and deploy at immense scale . Besides the compelling scientific and mathematical developments it may stimulate , this challenge has also a clear real-world significance : recent practice of training deep learning models suggests that the utility of common first-order methods is quickly reaching a plateau , in large part because their time-per-step is already negligible ( compared to other parts of the computation ) and can not be optimized further ; thus , the only way to obtain faster training performance is by drastically reducing the number of update steps . To this end , utilizing second-order methods seem a very natural and promising approach . In this paper we attempt to narrow the gap between theory and practice of second-order methods , focusing on second-order adaptivemethods for stochastic optimization . These methods can be thought of as full-matrix analogues of common adaptive algorithms such as AdaGrad ( Duchi et al. , 2011 ; McMahan & Streeter , 2010 ) and Adam ( Kingma & Ba , 2014 ) : they precondition each gradient with a second moment matrix , akin to a covariance matrix , that accumulates the outer products of the stochastic gradients . Full-matrix versions are potentially more powerful than first-order methods as they can exploit statistical correlations between ( gradients of ) different parameters ; geometrically , they can scale and rotate gradients whereas first order methods only scale gradients . However they suffer from similar prohibitive runtime and memory costs as Hessian-based methods . Recent developments in the space of second-order methods , on which we focus on in this paper , include the K-FAC ( Heskes , 2000 ; Martens & Grosse , 2015 ) and Shampoo ( Gupta et al. , 2018 ) algorithms that exploit the structure of deep networks ( and more generally , models described by a collection of tensors ) for mitigating the space and runtime costs of full-matrix second-order algorithms . These methods approximate each preconditioning matrix using a factored representation that stems from the network structure . However , in very large applications , such algorithms are still impractical due to a number of numerical and infrastructural pitfalls and are difficult to parallelize . Contributions . We provide solutions to practical concerns and challenges that arise in implementing and using second-order methods at large scale . Our focus will be on the Shampoo algorithm , but most of the challenges we address are relevant to the implementation of many other second-order methods . These include : • We design and implement an pipelined version of the optimization algorithm , critically exploiting the heterogeneity and computing power of CPU-Accelerator coupled architectures ; • We extend Shampoo in a number of ways so as to make it applicable to a larger range of deep architectures ; in particular , the extensions allow Shampoo to be used for training very large layers such as embedding layers ubiquitous in language and translation models ; • We replace expensive spectral decompositions ( e.g. , SVD ) used formanipulating preconditioners with an efficient and numerically-stable iterative method for computing roots of PSD matrices ; • We describe practical challenges and limitations we faced in our design , which we argue could be useful for the design considerations of next-generation accelerator hardware architectures . Our distributed implementation demonstrates significant improvements in performance , both in terms of number of steps , and often in actual wall-clock time , on some extremely large deep learning tasks : • Machine translation : we train Transformer models ( Vaswani et al. , 2017 ) on the WMT ’ 14 English to French translation task ( Bojar et al. , 2014 ) in half as many steps compared to state-of-the-art ( well tuned Adam ) , resulting with up to 45 % reduction in wall-time . • Language modeling : we trained BERT ( Devlin et al. , 2018 ) in 16 % fewer steps and achieve higher masked-LM accuracy compared to state-of-the-art optimizer ( You et al. , 2019 ) at 32K batch size ; overall wall-time decreased by 4 % from 3.8 to 3.65 hours . ( For this task , our system has not yet been tuned for performance ; we discuss several possible optimizations below . ) • Click-Through Rate ( CTR ) prediction : we trained the DLRM model ( Naumov et al. , 2019 ) on the terabyte Criteo dataset ( Criteo Labs , 2015 ) at 64K batch size in half as many steps as the current state-of-the-art optimizer , with a wall-time reduction of 37.5 % . We achieve a new state-of-the-art performance of 80.56 % AUC ( ≈ 0.3 % improvement ) on this task . ( An improvement of 0.1 % is considered significant ; see Rong et al. , 2020 ; Wang et al. , 2017 . ) • Image classification : we achieve MLPerf target accuracy of 75.9 % ( Mattson et al. , 2019 ) at 32K batch size on the standard ResNet-50 ImageNet benchmark in 10 % fewer steps than previous state-of-the-art . Here we do not see wall-time gains , mainly because the problem is too small ( only few thousand steps for convergence which does not allow for amortization of costs ) . However , we expect that one would be able to better exploit parallelism via improved software and hardware support . We note that one of our main points in this work was to demonstrate wall-time speedups with secondorder methods implemented on a real-world distributed setup being used to train state-of-the-art deep models . In our view , this is important for influencing future hardware accelerator design and runtime software . Indeed , first-order methods have received huge investments in tuning , implementation , platform support and tailored accelerator hardware over the last decade ; we believe there are numerous opportunities to improve the per-step time performance of preconditioned methods as well . For example , our results provide a concrete justification for incorporating 64bit accumulation units in hardware for distributed training , adding larger on-chip memory , better model parallelism and tighter coupling between accelerators and CPUs , which would make second order methods feasible across more domains and models . Related work . Classic techniques for addressing the high storage and computation costs of secondorder methods mostly belong to the quasi-Newton or the trust-region families of algorithms ( Conn et al. , 2000 ; Nocedal & Wright , 2006 ) . Traditionally , these methods need nearly-accurate gradients in order to construct useful quadratic approximations and implement reliable line searches , rendering them as suitable for training with very large batch sizes , and resulting in expensive iterations that make the overall algorithm slow compared with stochastic first-order methods ( see , e.g. , Bollapragada et al. , 2018 for a recent account ) . Hence , our focus in this paper is on adaptive second-order methods which are directly applicable in a stochastic setting . That said , our effort could be relevant to quasi-Newton and trust-region methods as well : e.g. , each iteration of typical trust-region methods amounts to solving a certain generalized eigenvalue problem , which presents numerical difficulties of similar nature to those encountered in matrix root/inverse computations , being addressed here . Various approximations to the preconditioning matrix have been proposed in the recent literature ( e.g. , Gonen & Shalev-Shwartz , 2015 ; Erdogdu & Montanari , 2015 ; Agarwal et al. , 2016 ; Xu et al. , 2016 ; Pilanci & Wainwright , 2017 ) . However , so far the only prevalent and pragmatic approximation is the diagonal approximation . Some recent approaches for approximating a full-matrix preconditioner are K-FAC ( Martens & Grosse , 2015 ) , Shampoo ( Gupta et al. , 2018 ) and GGT ( Agarwal et al. , 2018 ) . K-FAC uses a factored approximation of the Fisher-information matrix as a preconditioner . While our focus in this paper is on Shampoo , we believe that many of the techniques presented here could also be applied to make K-FAC practical in large scale ( see Appendix C ) . GGT uses a clever trick to compute a low-rank approximation to the AdaGrad preconditioner . However , GGT maintains several hundred copies of the gradient in memory , which is too expensive even for mid-sized models . Ba et al . ( 2017 ) took a first important step at experimenting with distributed K-FAC for training deep models , using a single machine with 8 GPUs to simulate a distributed environment for training . In contrast , a main thrust of our work is to demonstrate wall-time speedups with second-order methods on a real-world distributed setup used for training state-of-the-art deep models , that call for design considerations crucially different than in ( Ba et al. , 2017 ) . More recently , Osawa et al . ( 2019 ) scaled up K-FAC for training convolutional networks , but fell short of reaching the accuracy of first order methods , despite making changes to data augmentation and model architecture . 2 Preliminaries . Adaptive preconditioning methods . First order methods iteratively update the parameters solely based on gradient information : wt+1 = wt − ηt ḡt where wt and ḡt are ( column ) vectors in Rd . Here ḡt denotes a linear combination of the current and past gradients g1 , . . . , gt , where different algorithms use different combinations . Preconditioned methods take the form wt+1 = wt − Pt ḡt where Pt is an d × d matrix . Whereas in Newton-type methods this matrix is related to the Hessian matrix of second-order derivatives , adaptive preconditioning is based on gradient-gradient correlations . The parameters of a deep network are structured as a set of tensors of order two ( i.e. , a matrix ) , three , or four . For simplicity of presentation we focus on the matrix case—however our design , analysis , and implementation hold for tensors of arbitrary order . We denote the space of parameters by the matrix W ∈ Rm×n and an estimate of its gradient by G. Full matrix Adagrad flattens W , G to vectors of dimension mn , it thus requires m2n2 space to store the preconditioner and m3n3 time to perform the update . m and n are in the 1000 ’ s in state-of-the-art models , thus rendering full-matrix preconditioning impractical . For this reason , both AdaGrad and Adam constrain the preconditioning matrices to be diagonal . Shampoo bridges the gap between full matrix preconditioning and the diagonal version by approximating the matrices . The Shampoo algorithm . We describe Shampoo in the context of the Online Convex Optimization ( OCO ) framework , which generalizes stochastic optimization ( see , e.g. , Shalev-Shwartz , 2012 ; Hazan , 2016 ) . In OCO , learning progresses in rounds where on round t the learner receives an input Xt and then uses the parameters Wt to form a prediction denoted ŷt . After making the prediction , the true outcome yt is revealed . The discrepancy between the true and predicted outcomes is assessed by a loss function ` which takes values in R+ . The learner then uses the discrepancy to update the matrix to Wt+1 and prepare for the next round . For instance , the input on round t can be an example xt ∈ Rn for which the learner predicts ŷ = f ( Wt , xt ) where f : Rm → R and the loss is a function ` : R ×R→ R+ such as ` ( ŷ , y ) = ( y − ŷ ) 2 or ` ( ŷ , y ) = log ( 1 + exp ( −yŷ ) ) . Stochastic gradient methods use the gradient Gt = ∇W ` ( f ( W , xt ) , yt ) , thus Gt ∈ Rm×n if the parameters are shaped as a matrix W ∈ Rm×n . For matrix-shaped parameters , Shampoo tracks two statistics over the course of its run , Lt and Rt , which are defined as follows : Lt = ²Im + ∑t s=1 GsG T s ; Rt = ²In + ∑t s=1 G T sGs . Note that Lt ∈ Rm×m and Rt ∈ Rn×n . These are used to precondition the gradient and update W : Wt+1 = Wt − η L−1/4t GtR −1/4 t . The primary complexity of Shampoo arises from the computation of L−1/4t and R −1/4 t , which was naively implemented using spectral decompositions ( i.e. , SVD ) .

Summary: This work addresses practical challenges in applying full matrix pre-conditioner methods (such as Shampoo) on problems involving large datasets and architectures trained using a distributed setup.  In particular, this work presents a practical extension for the Shampoo algorithm by (1) using only a left or right preconditioner for large layers (2) computing inverse pth roots via coupled Newton iteration algorithms (3) distributing preconditioner computation across CPU cores in a CPU-GPU/TPU cluster and (4) delaying preconditioner computation to occur only once per several steps.  The proposed modifications lead to an implementation of Shampoo that consistently decreases the number of training steps and in certain cases provides a direct wall time improvement over Adagrad/Adam.