Small steps and giant leaps: Minimal Newton solvers for Deep Learning

1 INTRODUCTION . Stochastic Gradient Descent ( SGD ) and back-propagation ( LeCun et al. , 1998 ) are the algorithmic backbone of current deep network training . The success of deep learning demonstrates the power of this combination , which has been successfully applied on various tasks with large datasets and very deep networks ( He et al. , 2016 ) . Yet , while SGD has many advantages , speed of convergence ( in terms of number of iterations ) is not necessarily one of them . While individual SGD iterations are very quick to compute and lead to rapid progress at the beginning of the optimisation , it soon reaches a slower phase where further improvements are achieved slowly . This can be attributed to entering regions of the parameter space where the objective function is poorly scaled . In such cases , rapid progress would require vastly different step sizes for different directions in parameter space , which SGD can not deliver . Second-order methods , such as Newton ’ s method and its variants , eliminate this issue by rescaling the gradient according to the local curvature of the objective function . For a scalar loss in R , this rescaling takes the form H−1J where H is the Hessian matrix ( second-order derivatives ) or an approximation of the local curvature in the objective space , and J is the gradient of the objective . They can in fact achieve local scale-invariance ( Wright & Nocedal , 1999 , p. 27 ) , and make provably better progress in the regions where gradient descent stalls . While they are unmatched in other domains , there are several obstacles to their application to deep models . First , it is impractical to invert or even store the Hessian , since it grows quadratically with the number of parameters , and there are typically millions of them . Second , any Hessian estimate is necessarily noisy and ill-conditioned due to stochastic sampling , to which classic inversion methods such as conjugate-gradient are not robust . In this paper , we propose a new algorithm that can overcome these difficulties and make second order optimisation practical for deep learning . We show in particular how to avoid the storage of any estimate of the Hessian matrix or its inverse . Instead , we treat the computation of the Newton update , H−1J , as solving a linear system that itself can be solved via gradient descent . The cost of solving this system is amortized over time by interleaving its steps with the parameter update steps . Our proposed method adds little overhead , since a Hessian-vector product can be implemented for modern networks with just two steps of automatic differentiation . Interestingly , we show that our method is equivalent to momentum SGD ( also known as the heavy-ball method ) with a single additional term , accounting for curvature . For this reason we named our method CURVEBALL . Unlike other proposals , the total memory footprint is as small as that of momentum SGD . This paper is structured as follows . We introduce relevant technical background in sec . 2 , and present our method in sec . 3 . We evaluate our method and show experimental results in sec . 4 . Related work is discussed in sec . 5 . Finally we summarise our findings in sec . 6 . 2 BACKGROUND . In order to make the description of our method self-contained , we succinctly summarise a few standard concepts in optimisation . Our goal is to find the optimal parameters of a model ( e.g . a neural network ) φ : Rp → Ro , with p parameters w ∈ Rp and o outputs ( the notation does not show the dependency on the training data , which is subsumed in φ for compactness ) . The quality of the outputs is evaluated by a loss function L : Ro → R , so finding w is reduced to the optimisation problem : 1 w∗ = arg min w L ( φ ( w ) ) = arg min w f ( w ) . ( 1 ) Perhaps the simplest algorithm to find an optimum ( or at least a stationary point ) of eq . 1 is gradient descent ( GD ) . GD updates the parameters using the iteration w ← w − βJ ( w ) , where β > 0 is the learning rate and J ( w ) ∈ Rp is the gradient ( or Jacobian ) of the objective function f with respect to the parameters w. A useful variant is to augment GD with a momentum variable z ( Polyak , 1964 ) , which can be interpreted as a decaying average of past gradients:2 z ← ρz − βJ ( w ) ( 2 ) w ← w + z ( 3 ) with a momentum parameter ρ . Momentum GD , as given by eq . 2-3 , can be shown to have faster convergence than GD for convex functions , remaining stable under higher learning rates , and exhibits somewhat better resistance to poor scaling of the objective function ( Nesterov , 2013 ; Goh , 2017 ) . One important aspect is that these advantages cost almost no additional computation and only a modest additional memory , which explains why it is widely used in practice . In neural networks , GD is usually replaced by its stochastic version ( SGD ) , where at each iteration one computes the gradient not of the model f = L ( φ ( w ) ) , but of the model ft = Lt ( φt ( w ) ) assessed on a small batch of samples , drawn at random from the training set . 2.1 SECOND-ORDER OPTIMISATION . As mentioned in section 1 , the Newton method is similar to GD , but steers the gradient by the inverse Hessian matrix , computing H−1J as a descent direction . However , inverting the Hessian may be numerically unstable or the inverse may not even exist . To address this issue , the Hessian is usually regularized with a parameter λ > 0 , obtaining what is known as the Levenberg ( Moré , 1978 ) method : z = − ( H + λI ) −1J , ( 4 ) w ← w + z , ( 5 ) where H ∈ Rp×p , J ∈ Rp and I ∈ Rp×p is the identity matrix . Note that , unlike for momentum GD ( eq . 2 ) , the new step z is independent of the previous step . To avoid burdensome notation , we omit the w argument in H ( w ) and J ( w ) , but they must be recomputed at each iteration . Intuitively , the effect of eq . 4 is to rescale the step appropriately for different directions — directions with high curvature require small steps , while directions with low curvature require large steps to make progress . Note also that Levenberg ’ s regularization loses the scale-invariance of the original Newton method , meaning that rescaling the function f changes the scale of the gradient and hence the regularised descent direction chosen by the method . An alternative that alleviates this issue is LevenbergMarquardt , which replaces I in eq . 4 with diag ( H ) . For non-convex functions such as deep networks , these methods only converge to a local minimum when the Hessian is positive-semidefinite ( PSD ) . 1We omit the optional regulariser term for brevity , but this does not materially change our derivations . 2Some sources use an alternative form with β in eq . 3 instead ( equivalent by a re-parametrisation of z ) , 2.2 AUTOMATIC DIFFERENTIATION AND BACK-PROPAGATION . In order to introduce fast computations involving the Hessian , we must take a short digression into how Jacobians are computed . The Jacobian of L ( φ ( w ) ) ( eq . 1 ) is generally computed as J = JφJL where Jφ ∈ Rp×o and JL ∈ Ro×1 are the Jacobians of the model and loss , respectively . In practice , a Jacobian is never formed explicitly , but Jacobian-vector products Jv are implemented with the back-propagation algorithm . We define ←− AD ( v ) = Jv ( 6 ) as the reverse-mode automatic differentiation ( RMAD ) operation , commonly known as backpropagation . Note that , because the loss is a scalar function , the starting projection vector v typically used in gradient descent is a scalar and we set v = 1 . For intermediate computations , however , it is generally a ( vectorized ) tensor of gradients ( see eq . 9 ) . A perhaps lesser known alternative is forward-mode automatic differentiation ( FMAD ) , which computes a vector-Jacobian product , from the other direction : −→ AD ( v′ ) = v′J ( 7 ) This variant is less commonly-known in deep learning as RMAD is appropriate to compute the derivatives of a scalar-valued function , such as the learning objective , whereas FMAD is more appropriate for vector-valued functions of a scalar argument . However , we will show later that FMAD is relevant in calculations involving the Hessian . The only difference between RMAD and FMAD is the direction of associativity of the multiplication : FMAD propagates gradients in the forward direction , while RMAD ( or back-propagation ) does it in the backward direction . For example , for the composition of functions a ◦ b ◦ c , −→ ADa◦b◦c ( v ) = ( ( vJa ) Jb ) Jc ( 8 ) ←− ADa◦b◦c ( v ′ ) = Ja ( Jb ( Jcv ′ ) ) ( 9 ) Because of this , both operations have similar computational overhead , and can be implemented similarly . 2.3 FAST HESSIAN-VECTOR PRODUCTS . Since the Hessian of learning objectives involving deep networks is not necessarily positive semidefinite ( PSD ) , it is common to use a surrogate matrix with this property , which prevents second-order methods from being attracted to saddle-points ( this problem is discussed by Dauphin et al . ( 2014 ) ) . One of the most widely used is the Gauss-Newton approximation ( Martens & Grosse , 2015 ; Botev et al. , 2017 ; Wright & Nocedal , 1999 , p. 254 ) : Ĥ = JφHLJ T φ , ( 10 ) When HL is PSD , which is the case for all convex losses ( e.g . logistic loss , Lp distance ) , the resulting Ĥ is PSD by construction . For the method that we propose , and indeed for any method that implicitly inverts the Hessian ( or its approximation ) , only computing Hessian-vector products Ĥv is required . As such , eq . 10 takes a very convenient form : Ĥv = Jφ ( HL ( JTφ v ) ) ( 11 ) = ←− ADφ ( HL ( −→ ADφ ( v ) ) ) . ( 12 ) The cost of eq . 12 is thus equivalent to that of two back-propagation operations . This is similar to a classic result ( Pearlmutter , 1994 ; Schraudolph , 2002 ) , but written in terms of common automatic differentiation operations . The intermediate matrix-vector product HLu has negligible cost : for example , for the squared-distance loss , HL = 2I ⇒ HLu = 2u . Similarly , for the multinomial logistic loss we have HL = diag ( p ) − ppT ⇒ HLu = p u − p ( pTu ) , where p is the vector of predictions from a softmax layer and is the element-wise product . These products thus require only element-wise operations . Algorithm 1 : CURVEBALL ( proposed ) . 1 : z0 = 0 2 : for t = 0 , ... , T − 1 do 3 : ∆z = Ĥ ( wt ) zt + J ( wt ) 4 : zt+1 = ρzt − β∆z 5 : wt+1 = wt + zt+1 6 : end for Algorithm 2 : Simplified Hessian-free method . 1 : for t = 0 , ... , T − 1 do 2 : z0 = −J ( wt ) 3 : for r = 0 , ... , R− 1 ( or convergence ) do 4 : zr+1 = CG ( zr , Ĥ ( wt ) zr , J ( wt ) ) 5 : end for 6 : wt+1 = wt + zR 7 : end for

Summary: Authors propose choosing direction by using a single step of gradient descent "towards Newton step" from an original estimate, and then taking this direction instead of original gradient. This direction is reused as a starting estimate for the next iteration of the algorithm. This can be efficiently implemented since it only relies on Hessian-vector products which are accessible in all major frameworks.