First-Order Preconditioning via Hypergradient Descent

1 INTRODUCTION . High-dimensional nonlinear optimization problems often present a number of difficulties , such as strongly-correlated parameters and variable scaling along different directions in parameter space ( Martens , 2016 ) . Despite this , deep neural networks and other large-scale machine learning models applied to such problems typically rely on simple variations of gradient descent to train , which is known to be highly sensitive to these difficulties . While this approach often works well in practice , addressing the underlying issues directly could provide stronger theoretical guarantees , accelerate training , and improve generalization . Adaptive learning rate methods such as Adam ( Kingma and Ba , 2014 ) , Adagrad ( Duchi et al. , 2011 ) , and RMSProp ( Tieleman and Hinton , 2012 ) provide some degree of higher-order approximation to re-scale updates based on per-parameter behavior . Newton-based methods make use of the curvature of the loss surface to both re-scale and rotate the gradient in order to improve convergence . Natural gradient methods ( Amari , 1998 ) do the same in order to enforce smoothness in the evolution of the model ’ s conditional distribution . In each of the latter two cases , the focus is on computing a specific linear transformation of the gradient that improves the conditioning of the problem . This transformation is typically known as a preconditioning , or curvature , matrix . In the case of quasiNewton methods , the preconditioning matrix takes the form of the inverse Hessian , while for natural gradient methods it ’ s the inverse Fisher information matrix . Computing these transformations is typically intractable for high-dimensional problems , and while a number of approximate methods exist for both ( e.g. , ( Byrd et al. , 1996 ; Martens and Grosse , 2015 ; Grosse and Martens , 2016 ) ) , they are often still too expensive for the performance gain they provide . These approaches also suffer from rigid inductive biases regarding the nature of the problems to which they are applied in that they seek to compute or approximate specific transformations . However , in large , non-convex problems , the optimal gradient transformation may be less obvious , or may even change over the course of training . In this paper , we address these issues through a method we term first-order preconditioning ( FOP ) . Unlike previous approaches , FOP doesn ’ t attempt to compute a specific preconditioner , such as the inverse Hessian , but rather uses first-order hypergradient descent ( Maclaurin et al. , 2015 ) to learn an Algorithm 1 Learned First-Order Preconditioning ( FOP ) 1 : Require : model parameters θ , objective function J , FOP matrix M , learning rate , hypergra- dient learning rate ρ 2 : for t = 1,2 , ... do 3 : Draw data x ( t ) , y ( t ) ∼ D 4 : Perform forward pass : ŷ ( t ) = fθ ( t ) ( x ( t ) ) 5 : Compute loss J ( y ( t ) , ŷ ( t ) ) 6 : Update inference parameters θ : θ ( t+1 ) ← θ ( t ) − M ( t ) M ( t ) > ∇θ ( t ) J 7 : Update preconditioning matrices : M ( t+1 ) ←M ( t ) + ρ ( ∇θ ( t ) J [ ∇θ ( t−1 ) J ] > +∇θ ( t−1 ) J [ ∇θ ( t ) J ] > ) M ( t ) 8 : Cache∇θ ( t ) J adaptable transformation online directly from the task objective function . Our method adds minimal computational and memory cost to standard deep network training and results in improved convergence speed and generalization compared to standard approaches . FOP can be flexibly applied to any gradient-based optimization problem , and we show that when used in conjunction with standard optimizers , it improves their performance . To our knowledge , this is also the first successful application of any hypergradient method to the ImageNet dataset ( Deng et al. , 2009 ) . 2 FIRST ORDER PRECONDITIONING . 2.1 THE BASIC APPROACH . Consider a parameter vector θ and a loss function J . A traditional gradient update with a preconditioning matrix P can be written as θ ( t+1 ) = θ ( t ) − P ( t ) ∇θ ( t ) J ( t ) . ( 1 ) Our goal is to learn P . However , while we place no other constraints on our preconditioner , in order to ensure that it is positive semi-definite , and therefore does not reverse the direction of the gradient , we reparameterize P as the product of a matrix M with itself , changing Equation 1 to : θ ( t+1 ) = θ ( t ) − M ( t ) M ( t ) > ∇θ ( t ) J ( t ) . ( 2 ) Under reasonable assumptions , gradient descent is guaranteed to converge with the use of even a random symmetric , positive-definite preconditioner , as we show in Appendix A.3 . Because θ ( t ) is a function ofM ( t−1 ) , we can then backpropagate from the loss at iteration t to the previous iteration ’ s preconditioner via a simple application of the chain rule : ∂J ( t ) ∂M ( t−1 ) = ∂J ( t ) ∂θ ( t ) ∂θ ( t ) ∂M ( t−1 ) . ( 3 ) By applying the chain and product rules to Equation 2 , the gradient with respect to the preconditioner is then simply ∇M ( t−1 ) J ( t ) = − ( ∇θ ( t ) J ( t ) [ ∇θ ( t−1 ) J ( t−1 ) ] > +∇θ ( t−1 ) J ( t−1 ) [ ∇θ ( t ) J ( t ) ] > ) M ( t ) . ( 4 ) We provide a more detailed derivation in Appendix A.1 . Note that ideally we would compute ∇M ( t ) J ( t+1 ) to update M ( t ) , but as we don ’ t have access to J ( t+1 ) yet , we follow the example of Almeida et al . ( 1998 ) and assume that J is smooth enough that this does not have an adverse effect on training . The basic approach is summarized in Algorithm 1 . We use supervised learning as an example , but the same method applies to any gradient-based optimization . The preconditioned gradient can then be passed to any standard optimizer to produce an update for M . For example , we describe the procedure for using FOP with momentum ( Polyak , 1964 ) in Section 2.3 . For multilayer networks , in order to make the simplest modification to normal backpropagation , we learn a separate M for each layer , not a global curvature matrix . To get an intuition for the behavior of FOP compared to standard algorithms , we observed its trajectories on a set of low-dimensional optimization problems ( Figure 1 ) . Interestingly , while FOP converged in fewer iterations than SGD and Adam ( Kingma and Ba , 2014 ) , it took more jagged paths along the objective function surface , suggesting that while it takes more aggressive steps , it is perhaps also able to change direction more rapidly . 2.2 SPATIAL PRECONDITIONING FOR CONVOLUTIONAL NETWORKS . In order to further reduce the computational cost of FOP in convolutional networks ( CNNs ) , we implemented layer-wise spatial preconditioners , sharing the matrices across both input and output channels ( results shown in Section 4 ) . More concretely , if a convolutional layer has spatial kernels with shape k × k , we can learn a preconditioner that is k2 × k2 . To implement this , when θ is a 4- tensor of kernels of shape k×k×I×O , where I andO are the input and output channels , respectively , we can reshape it to a matrix of size k2 × IO , left-multiply it by the learned curvature matrix , and then reshape it back to its original dimensions . When k is small , as is typically the case in deep CNNs , this preconditioner is small as well , resulting in both computational and memory efficiency . For large fully-connected networks , we adopt a low-rank preconditioner to save computation . In Appendix A.4 , we offer further details and show that FOP is still effective even when the rank is very low . 2.3 FOP FOR MOMENTUM . FOP can be implemented alongside any standard optimizer , such as gradient descent with momentum ( Polyak , 1964 ) . Given a parameter vector θ and a loss function J , a basic momentum update with FOP matrix M is typically expressed as two steps : v ( t+1 ) = αv ( t ) +M ( t ) M ( t ) > ∇θ ( t ) J ( 5 ) θ ( t+1 ) = θ ( t ) − v ( t+1 ) , ( 6 ) where v is the velocity term and α is the momentum parameter . Combining Equations 5 and 6 allows us to write the full update as θ ( t+1 ) = θ ( t ) − αv ( t ) − M ( t ) M ( t ) > ∇θ ( t ) J . ( 7 ) If , as in Maclaurin et al . ( 2015 ) , we were meta-learning M or only updating M after a certain number of iterations , we would then have to backpropagate through v to calculate the gradient for M . As we are updating M online , however , we only need to calculate ∇M ( t ) J . Therefore , the update is the same as for standard gradient descent . The experiments in Section 4 were performed using this modification of momentum . 3 RELATED WORK . Almeida et al . ( 1998 ) introduced the idea of using gradients from the objective function to learn optimization parameters such as the learning rate or a curvature matrix . However , their preconditioning matrix was strictly diagonal , amounting to an approximate Newton algorithm ( Martens , 2016 ) , and they only tested their framework on simple optimization problems with either gradient descent or SGD . This is related to common deep learning optimizers such as Adagrad ( Duchi et al. , 2011 ) and Adam ( Kingma and Ba , 2014 ) , which also learn adaptive per-parameter learning rates , but unlike FOP , do not induce a rotation on the gradient . As we note in Section 2.3 , an advantage of FOP is that it may also be used in conjunction with any such optimizer . More recently , Maclaurin et al . ( 2015 ) applied the backpropagation-across-iterations approach in a neural network context , terming the process hypergradient descent . Their method backpropagates through multiple iterations of the training process of a relatively shallow network to meta-learn a learning rate . However , this method can incur significant memory and computational cost for large models and long training times . Baydin et al . ( 2017 ) instead proposed an online framework directly inherited from Almeida et al . ( 1998 ) that used hypergradient-based optimization in which the learning rate is updated after each iteration . Our method extends this idea to not only learn existing optimizer parameters ( e.g. , learning rate , momentum ) , but to introduce novel ones in the form of a non-diagonal , layer-specific preconditioning matrix for the gradient . It ’ s also important to discuss the relationship of FOP to other , non-hypergradient preconditioning methods for deep networks . These mostly can be sorted into one of two categories , quasi-Newton algorithms and natural gradient approaches . Quasi-Newton methods seek to learn an approximation of the inverse Hessian . L-BFGS , for example , does this through tracking the differences between gradients across iterations ( Byrd et al. , 1996 ) . This is significantly different from FOP , although the outer product of gradients used in the update for FOP can also be an approximation of the Hessian or Gauss-Newton matrices ( Bottou et al. , 2016 ) . Natural gradient methods , such as K-FAC ( Martens and Grosse , 2015 ) and KFC ( Grosse and Martens , 2016 ) , which approximate the inverse Fisher information matrix , bear a much stronger resemblance to FOP . However , there are notable differences . First , unlike FOP , these methods perform extra computation to ensure the invertibility of their curvature matrices . Second , the learning process for the preconditioner in these methods is completely different , as they do not backpropagate across iterations .

Summary: This paper studies hypergradient descent for precondition matrices. The goal is to learn an adaptable preconditioning for the task while training. Specifically, they take the gradient of the loss wrt the precondition matrix and update the precondition matrix to decrease the loss. They reparametrize the precondition matrix to ensure it is positive-definite and provide low-rank approximations and they provide cheap approximations for CNNs.