Gram-Gauss-Newton Method: Learning Overparameterized Neural Networks for Regression Problems

1 INTRODUCTION . First-order methods such as Stochastic Gradient Descent ( SGD ) are currently the standard choice for training deep neural networks . The merit of first-order methods is obvious : they only calculate the gradient and therefore are computationally efficient . In addition to better computational efficiency , SGD has even more advantages among the first-order methods . At each iteration , SGD computes the gradient only on a mini-batch instead of all training data . Such randomness introduced by sampling the mini-batch can lead to better generalization ( Hardt et al. , 2015 ; Keskar et al. , 2016 ; Masters & Luschi , 2018 ; Mou et al. , 2017 ; Zhu et al. , 2018 ) and better convergence ( Ge et al. , 2015 ; Jin et al. , 2017a ; b ) , which is crucial when the function class is highly overparameterized deep neural networks . Recently there is a huge body of works trying to develop more efficient first-order methods beyond SGD ( Duchi et al. , 2011 ; Kingma & Ba , 2014 ; Luo et al. , 2019 ; Liu et al. , 2019 ) . Second-order methods , despite their better convergence rate , are rarely used to train deep neural networks . At each iteration , the algorithm has to compute second order information , for example , the Hessian or its approximation , which is typically an m by m matrix where m is the number of parameters of the neural network . Moreover , the algorithm needs to compute the inverse of this matrix . The computational cost is prohibitive and usually it is not even possible to store such a matrix . Recently , there is a series of works ( Du et al. , 2018b ; a ; Zou et al. , 2018 ; Allen-Zhu et al. , 2018a ; Oymak & Soltanolkotabi , 2018 ; Arora et al. , 2019b ; a ; Cao & Gu , 2019 ; Zou & Gu , 2019 ) considering the optimization of neural networks based on the idea of neural tangent kernel ( NTK ) ( Jacot et al. , 2018 ) . Roughly speaking , the idea of NTK is to linearly approximate the output of a network w.r.t . the parameters at a local region , thus resulting in a kernel feature map , which is the gradient of the output w.r.t . the parameters . Jacot et al . ( 2018 ) shows that when the width of the network tends to infinity , the NTK tends to be unchanged during gradient flow since the network only needs o ( 1 ) change ( goes to zero as the width tends to infinity ) of parameter to fit the data , so the change of kernel is also o ( 1 ) . As for finite-width networks , however , the NTK changes and previous works ( Allen-Zhu et al. , 2018b ; Arora et al. , 2019a ) show that for sufficiently wide neural networks the optimization dynamic of GD/SGD is equivalent to that of using GD/SGD to solve an NTK kernel regression problem where the kernel is slowly evolving ( see Lemma 1 in Section 2 for a precise description ) . A natural question then arises : Can we gain acceleration by directly solving kernel regression w.r.t . the NTK at each step ? In this paper , we give a positive answer to this question and reveal the connection between NTK regression and Gauss-Newton method . We propose a novel optimization method – the Gram-GaussNewton ( GGN ) method . Instead of doing gradient descent , GGN solves the kernel regression w.r.t . the NTK at each step of the optimization . Following this idea , we theoretically prove that for overparameterized networks , GGN enjoys quadratic convergence compared to the linear rate of gradient descent . Besides theoretical fast convergence rate , GGN is also very efficient in practice . In fact , GGN is implicitly a reformulation of the Gauss-Newton method ( see Section 3.1 for details ) which is a classic second-order algorithm often used for solving nonlinear regression problems with square loss . In the Gauss-Newton method , one uses J > J as an approximation of the Hessian ( see Section 2 for a formal description ) where J is the Jacobian matrix . However , the original Gauss-Newton method faces challenges when used for training deep neural networks . Most seriously , the size of the approximate Hessian J > J is m by m , where m is the number of parameters of the neural network . Moreover , for overparameterized neural networks , J > J is not invertible , which may make the algorithm intractable for training commonly-used neural networks . GGN bypasses the difficulty stated above as follows . Instead of using J > J as approximate Hessian and applying Newton-type method , each step of GGN only involves the Gram matrix JJ > whose size is n by n where n is the number of data . Furthermore , as already mentioned , to get better generalization performance , it is crucial to use mini-batch to introduce sampling noise when calculating derivatives . Therefore , like SGD , we also use mini-batch in GGN . In this case , the size of the Gram matrix further reduces to b by b , where b is the batch size . Though conventional wisdom may suggest that applying mini-batch scheme to second-order methods will introduce biased estimation of the accelerated gradient direction , we give the first convergence result for mini-batch second-order method on overparameterized networks . Regarding computational complexity , we show that at each iteration , the overhead of GGN is small compared to SGD : the extra computation of GGN is mainly the matrix product JJ > and the inverse of this matrix whose size is small for a mini-batch . Detailed analyses can be found in Section 3.3 . We next conduct experiments on two regression tasks to study the effectiveness of the GGN algorithm . We demonstrate that in these two real applications , using a practical neural network ( e.g. , ResNet-32 ) with standard width , our proposed GGN algorithm can converge faster and achieve better performance than several baseline algorithms . 1.1 RELATED WORKS . Despite the prevalence of first-order methods for training deep neural networks , there have been continuing efforts in developing practical second-order methods ( Becker et al. , 1988 ; Pascanu & Bengio , 2013 ) . We summarize some of these works below . The main approach for these methods is to develop delicate approximations of the second-order information matrix so that the update direction can be computed as efficiently as possible . For example , Botev et al . ( 2017 ) proposed a recursive block-diagonal approximation of the Hessian . The blocks are Kronecker factored and can be efficiently computed and inverted . Grosse and Martens in a series of works developed the K-FAC method ( Martens & Grosse , 2015 ; Grosse & Martens , 2016 ) . The key idea is a Kronecker-factored approximation of the Fisher information matrix , which is used as the second-order matrix in natural gradient methods . These works received considerable attention and have been further improved ( Wu et al. , 2017 ; George et al. , 2018 ; Martens et al. , 2018 ) . Bernacchia et al . ( 2018 ) derived an exact expression of the natural gradient update , but only works for linear networks . Different from all these works , our GGN algorithm does not try to approximate the second-order matrix whose size is inevitably huge . Instead , we present an easy-to-compute solution of the updating direction , reducing the computational cost significantly . One exceptional concurrent work Ren & Goldfarb ( 2019 ) also aims to use the exact Gauss-Newton update . They focus on reducing the complexity of inverting approximate Hessian by Sherman-Morrison-Woodbury Formula and require subtle implementation tricks to use backpropagation . In contrast , GGN has simpler update rule and better guarantee for neural networks . In a concurrent and independent work , Zhang et al . ( 2019a ) showed that natural gradient method and K-FAC have a linear convergence rate for sufficiently wide networks in full-batch setting . In contrast , our method enjoys a higher-order ( quadratic ) convergence rate guarantee for overparameterized networks , and we focus on developing a practical and theoretically sound optimization method . We also reveal the relation between our method and NTK kernel regression , so using results based on NTK ( Arora et al. , 2019b ) , one can easily give generalization guarantee of our method . Another independent work ( Achiam et al. , 2019 ) proposed a preconditioned Q-learning algorithm which has similar form of our update rule . Unlike the methods considered in Zhang et al . ( 2019a ) ; Achiam et al . ( 2019 ) which contain the learning rate that needed to be tuned , our derivation of GGN does not introduce a learning rate term ( or understood as suggesting that the learning rate can be fixed to be 1 to get good performance which is verified in Figure 2 ( c ) ) . 2 NEURAL TANGENT KERNEL AND THE CLASSIC GAUSS-NEWTON METHOD FOR NONLINEAR LEAST SQUARES REGRESSION . Nonlinear least squares regression problem is a general machine learning problem . Given data pairs { xi , yi } ni=1 and a class of nonlinear functions f , e.g . neural networks , parameterized by w , the nonlinear least squares regression aims to solve the optimization problem min w∈Rm L ( w ) = 1 2 n∑ i=1 ( f ( w , xi ) − yi ) 2 . ( 1 ) In the seminal work ( Jacot et al. , 2018 ) , the authors consider the case when f is a neural network with infinite width . They showed that optimization on this problem using gradient flow involves a special kernel which is called neural tangent kernel ( NTK ) . The follow-up works further extended the relation between optimization and NTK which can be concluded in the following lemma : Lemma 1 ( Lemma 3.1 in Arora et al . ( 2019a ) , see also Dou & Liang ( 2019 ) ; Mei et al . ( 2019 ) ) . Consider optimizing problem ( 1 ) by gradient descent with infinitesimally small learning rate : dwtdt = −∇L ( wt ) . where wt is the parameters at time t. Let ft = ( f ( wt , xi ) ) ni=1 ∈ Rn be the network outputs on all xi ’ s at time t , and y = ( yi ) ni=1 be the desired outputs . Then ft follows the following evolution : dft dt = −Gt · ( ft − y ) , ( 2 ) where Gt is an n × n positive semidefinite matrix , i.e . the Gram matrix w.r.t . the NTK at time t , whose ( i , j ) -th entry is 〈∇wf ( wt , xi ) , ∇wf ( wt , xj ) 〉 . The key idea of Jacot et al . ( 2018 ) and its extensions ( Du et al. , 2018b ; a ; Zou et al. , 2018 ; Allen-Zhu et al. , 2018a ; Oymak & Soltanolkotabi , 2018 ; Lee et al. , 2019 ; Yang , 2019 ; Arora et al. , 2019b ; a ; Cao & Gu , 2019 ; Zou & Gu , 2019 ) is that when the network is sufficiently wide , the Gram matrix at initialization G0 is close to a fixed positive definite matrix defined by the infinite-width kernel and Gt is close to G0 during training for all t. Under this situation , Gt remains invertible , and the above dynamics is then identical to the dynamics of solving kernel regression with gradient flow w.r.t . the current kernel at time t. In fact , Arora et al . ( 2019a ) rigorously proves that a fully-trained sufficiently wide ReLU neural network is equivalent to the kernel regression predictor . As pointed out in Chizat & Bach ( 2018 ) , the idea of NTK can be summarized as a linear approximation using first order Taylor expansion . We give an example of this idea on the NTK at initialization : f ( w , xi ) − f ( w0 , xi ) ≈ ∇wf ( w0 , xi ) · ( w −w0 ) , ( 3 ) where∇wf ( w0 , x ) can then be viewed as an explicit expression of feature map at x , w −w0 is the parameter in reproducing kernel Hilbert space ( RKHS ) induced by NTK and f ( w , xi ) − f ( w0 , xi ) the target value . The idea of linear approximation is also used in the classic Gauss-Newton method ( Golub , 1965 ) to obtain an acceleration algorithm for solving nonlinear least squares problem ( 1 ) . Concretely , at iteration t , Gauss-Newton method takes the following first-order approximation : f ( w , xi ) − f ( wt , xi ) ≈ ∇wf ( wt , xi ) · ( w −wt ) , ( 4 ) where wt stands for the parameter at iteration t. We note that this is also the linear expansion for deriving NTK at time t. According to Eq . ( 1 ) and ( 4 ) , to update the parameter , one can instead solve the following problem . wt+1 = argmin w 1 2 ‖ft + Jt ( w −wt ) − y‖22 , ( 5 ) where ft , y have the same meaning as in Lemma 1 , and Jt = ( ∇wf ( wt , x1 ) , · · · , ∇wf ( wt , xn ) ) > ∈ Rn×m is the Jacobian matrix . A necessary and sufficient condition for w to be the solution of Eq . ( 5 ) is ( J > t Jt ) · ( w −wt ) = −J > t ( ft − y ) . ( 6 ) Below we will denote Ht : = J > t Jt ∈ Rm×m . For under-parameterized model ( i.e. , the number of parameters m is less than the number of data n ) , Ht is invertible , and the update rule is wt+1 = wt −H−1t J > t ( ft − y ) . ( 7 ) This can also be viewed as an approximate Newton ’ s method using Ht = J > t Jt to approximate the Hessian matrix . In fact , the exact Hessian matrix is ∇2w 1 2 n∑ i=1 ( f ( wt , xi ) − yi ) 2 = J > t Jt + n∑ i=1 ( f ( wt , xi ) − yi ) ∇2wf ( wt , xi ) . ( 8 ) In the case when f is only mildly nonlinear w.r.t . w at data point xi ’ s , ∇2wf ( wt , xi ) ≈ 0 , and Ht is close to the real Hessian . In this situation , the behavior of the Gauss-Newton method is similar to that of Newton ’ s method , and thus can achieve a superlinear convergence rate ( Golub , 1965 ) .

Summary: The authors propose a scalable second order method for optimization using a quadratic loss. The method is inspired by the Neural Tangent kernel approach, which also allows them to provide global convergence rates for GD and batch SGD. The algorithm has a computational complexity that is linear in the number of parameters and requires to solve a system of the size of the minibatch.  They also show experimentally the advantage of using their proposed methods over SGD.