Exact Stochastic Newton Method for Deep Learning: the feedforward networks case.

The inclusion of second-order information into Deep Learning optimization has drawn consistent interest as a way forward to improve upon gradient descent methods . Estimating the second-order update is computationally expensive , which drastically limits its usage scope and forces the use of various truncations and approximations . This work demonstrates that it is possible to solve the Newton direction in the stochastic case exactly . We consider feedforward networks as a base model , build a second-order Lagrangian which we call Sifrian , and provide a closed-form formula for the exact stochastic Newton direction under some monotonicity and regularization conditions . We propose a convexity correction to escape saddle points , and we reconsider the intrinsic stochasticity of the online learning process to improve upon the formulas . We finally compare the performance of the developed solution with well-established training methods and show its viability as a training method for Deep Learning . Optimization in Deep Learning is mainly dominated by first-order methods built around the central concept of backpropagation ( LeCun et al. , 1988 ) . Second-order methods have exceptional theoretical properties in deterministic optimization ( Boyd et al. , 2004 ; Nocedal & Wright , 2006 ) , but such properties do not translate well into the stochastic case ( LeCun et al. , 2012 ; Bottou et al. , 2018 ) . First-order methods such as the Stochastic Gradient Descent ( SGD ) ( Robbins & Monro , 1951 ) are relatively simple and have been adaptively improved for Deep Learning ( Duchi et al. , 2011 ; Tieleman & Hinton , 2012 ; Kingma & Ba , 2014 ; Reddi et al. , 2019 ; Yao et al. , 2020 ) . Inherently , secondorder methods might seems inadequate for Deep learning due to increased computational cost , poor clock-wall performance , and the non-convex nature of Deep Learning ( LeCun et al. , 2012 ) . Furthermore , Newton method might even reduce the generalization capabilities of training ( Wadia et al. , 2021 ; Amari et al. , 2020 ) . Despite these various limitations , substantial effort was deployed to include Hessian information into the optimization process ( e.g . Byrd et al. , 2011 ; Sohl-Dickstein et al. , 2014 ; Byrd et al. , 2016 ; Agarwal et al. , 2017 ; Berahas et al. , 2019 ; Anil et al. , 2020 ; Goldfarb et al. , 2020 ; Castera et al. , 2021 ) . Several approaches exist such as the Gauss-Newton method ( e.g Schraudolph , 2002 ; Botev et al. , 2017 ) , diagonal approximation of the Hessian ( e.g . Bordes et al. , 2009 ; Schaul et al. , 2013 ) , iterative low-rank updates such as BFGS ( Broyden , 1970 ; Fletcher , 1970 ; Goldfarb , 1970 ; Shanno , 1970 ) ( see also , Liu & Nocedal , 1989 ; Schraudolph et al. , 2007 ; Bollapragada et al. , 2018 ) , or Hessian-Free methods which combine the fast Hessian matrix multiplication ( Pearlmutter , 1994 ) and the conjugate gradient algorithm ( e.g . Martens , 2010 ; Martens & Sutskever , 2012 ; Dauphin et al. , 2014 ) . The use of the Fisher matrix to capture curvature information in the space of distributions , instead of the Hessian , is another approach which yields the natural gradient ( Amari , 1998 ) , but suffers from the same computational issues as the Newton method . Within this context K-FAC method ( Martens & Grosse , 2015 ; Ba et al. , 2016a ; George et al. , 2018 ) mitigates some of the computational issues of the natural gradient method . In essence , several drawbacks and flaws limits the adoption of second-order methods as a standard for neural networks training . In this paper , we propose a novel approach to characterize the Newton update which help us derive an exact closed-form solution for the stochastic Newton method . Our method requires a suitable regularization of the neural network and strict monotonicity of the activation functions . We start this paper by introducing useful notations while deriving the well-known backpropagation algorithm for a feedforward network . Then we introduce a second-order Lagrangian which we call Sifrian , that will serve to characterize the Newton direction . We derive four types of equations from the Sifrian , and we provide an exact closed-form solution for the Newton direction in the stochastic case . We further propose a saddle-free version of our method , and add a randomization process to enhance our solution . In the last part of this paper , we show the applicability of our method for Deep Learning through diverse classification tasks using feedforward architectures . 1 PRELIMINARIES . Deep Learning and neural networks training could be seen as an optimization problem of an expected loss function ` over a distribution D of labeled samples ( x , d ) . In general , weights and biases ( W , β ) are the sought-after parameters of the network . The labeled samples distribution D is often unknown , but a large number of samples ( database D ) allows the approximation of the expected loss with an empirical risk : min W , β ( E ( x , d ) ∼D [ ` ( W , β , x , d ) ] ) → 1 |D| ∑ p∈D ` ( W , β , x ( p ) , d ( p ) ) , ( x ( p ) , d ( p ) ) p∈D i.i.d∼ D. ( 1 ) The loss or cost function ` is typically a cross entropy function or an l2 norm of the mismatches between the network outputs and the labels : ` ( W , β , ( x ( p ) , d ( p ) ) ) = 12 ∥∥∥x ( p ) n − d ( p ) ∥∥∥2 2 . Where n indicates the output layer number . We present more in detail the architecture of a Feedforward Neural Network ( FNN ) hereafter . 1.1 NOTATIONS AND FEEDFORWARD NEURAL NETWORKS ( FNN ) . In this section , we recall the main details of FNNs , which will be used as a standard model for Deep Learning1 . The notation used throughout this paper is similar to notations presented in LeCun et al . ( 1988 ) . The main equation governing the FNN , a.k.a . the forward model is the following : x ( p ) k = F ( Wk x ( p ) k−1 + βk ) , k ∈ [ 1 .. n ] , p ∈ D. ( 2 ) where D is the database , p designs one single sample from the database ( e.g . one single image or audio recording ) , k is the layer number and n is the total number of layers in the network . The initial input for sample p is x ( p ) 0 ( e.g . the vectorized input image data ) . The state variable x ( p ) is transformed at each layer k through a multiplication by a weight matrix Wk and an addition of a bias vector βk . The activation function F , which is typically a sigmoid or a Rectified Linear Unit ( ReLU ) , is applied element wise on the resulting activation vector : a ( p ) k = Wk x ( p ) k−1 + βk , and serves to introduce non-linearity in the neural network . 1.2 THE LAGRANGIAN AND BACKPROPAGATION . The origin of backpropagation could be traced back to the early 1970s and could be derived by casting Deep Learning as a constrained optimization problem . This section is similar to LeCun et al . ( 1988 ) . The main novelty in this section is the addition of a regularization term dependent on the state variable R ( x ) to the cost function ` : ` R = ` + R ( x ) . The regularization has to verify the admissibility criteria which we define hereafter : Definition 1 : A regularization term is admissible if the second derivative of the augmented cost function w.r.t . to the state variable is separable and non-null , i.e . : ∀ ( p , q ) ∈ D , ∀ ( k , m ) ∈ [ 1 .. n ] ( ∂ ` R ∂x ( p ) k x ( q ) m ) p , q , k , m ∝ δp=q . ( 3 ) An example of such regularization is the following function : R ( x ) = 1 2 ∑ p∈D ∑ k=1 .. n−1 〈 x ( p ) k , Λ ( p ) k x ( p ) k 〉 ( 4 ) 1Convolutional Neural Networks ( CNN ) are a type of FNN . The concepts of this paper apply also to CNN . ( Λ ( p ) k ) p , k are symmetric positive matrices . Such a choice might seem atypical since the regularization often concerns the network parameters , mainly the weights or the biases . This ” admissible ” regularization was introduced for the sole purpose of guaranteeing a non-null partial derivative of the cost ` w.r.t the state variable x ( p ) k . This last property is essential to solve Equ . 16 . We will show later that such a regularization minimizes also the norm of the gradient ( Barrett & Dherin , 2020 ; Smith et al. , 2021 ) . In order to derive the backpropagation algorithm we introduce the following Lagrangian as per the notation of LeCun et al . ( 1988 ) : L ( x , W , β , b ) = ` R + ∑ p∈D ∑ k=1 .. n 〈 x ( p ) k − F ( Wk x ( p ) k−1 + βk ) , b ( p ) k 〉 . ( 5 ) The Lagrangian contains the original cost function , the admissible regularization term , and the product of the forward equation with the adjoint state vectors ( b ( p ) k ) p , k . These adjoint vectors are de- fined for each layer of the network and each sample of the database ; however , their values are not fixed yet and will be chosen in a way that simplifies the gradient computation . If the state variable x verifies the forward equations , then the Lagrangian simplifies to the cost function : L ( x ( W , β ) , W , β , b ) = ` R ( x ( W , β ) ) . ( 6 ) The total derivative of the previous Lagrangian w.r.t . the weights or biases is the gradient and could be expressed in terms of partial derivatives as follows2 : dL ( x ( W , β ) , W , β , b ) dWk = ∂L∂Wk + ∑ p∈D ∑ m=1 .. n ( dx ( p ) m dWk ) ( ∂L ∂x ( p ) m ) = d ` RdWk , dL ( x ( W , β ) , W , β , b ) dβk = ∂L∂βk + ∑ p∈D ∑ m=1 .. n ( dx ( p ) m dβk ) ( ∂L ∂x ( p ) m ) = d ` Rdβk . ( 7 ) The partial derivatives of the Lagrangian w.r.t to { Wk , βk } k=1 .. n are straightforward to compute . The Fréchet derivatives ( dx ( p ) m dWk , dx ( p ) m dβk ) p , k , m are non-trivial to evaluate and the core idea of back- propagation is the selection of the adjoint state variables ( b ( p ) k ) p , k which cancels the superfluous terms . Such a simplification is achievable if : ( ∂L ∂x ( p ) k ) p∈D , k∈ [ 1 .. n ] = 0 . ( 8 ) Computing the previous partial derivative yields the following backpropagation equation : ∂L ∂x ( p ) k = b ( p ) k − 1k=1 .. n−1W T k+1∇F ( a ( p ) k+1 ) b ( p ) k+1 + ∂ ` R ∂x ( p ) k = 0 . ( 9 ) 1E is the indicator function , and it is equal to one if the underlying condition E is true and null otherwise . The nabla operator ∇F is a diagonal square matrix with an element-wise derivative of its argument ( along the diagonal ) . The resolution of the backpropagation system could be split into a boundary condition and a backward propagation system . Further details could be found in LeCun et al . ( 1988 ) . The particular choice of the adjoint state vectors ( b ( p ) k ) p∈D , k=1 .. n yields a simple formula for the gradient of the cost function : Gk = ∂L ∂Wk = − ∑ p∇F ( a ( p ) k ) b ( p ) k x ( p ) T k−1 , gk = ∂L ∂βk = − ∑ p∇F ( a ( p ) k ) b ( p ) k . ( 10 ) The backpropagation algorithm is fundamental for Deep Learning . It is simple and provides a straightforward solution to an otherwise tedious problem to solve . Unfortunately , developing an efficient second-order backpropagation method remains elusive . The current state-of-the-art is based on the use of the R−operator ( Pearlmutter , 1994 ) to compute the product of the Hessian with a given vector without explicitly calculating or storing the Hessian ( Martens , 2010 ; Dauphin et al. , 2014 ; Agarwal et al. , 2017 ) . In the following section , we provide a different framework for backpropagation which allows the characterization of the second-order Newton direction . 2A denominator layout notation for derivation is used throughout this paper .

Summary: This paper claims that it is possible to compute the Newton method's update exactly for deep neural networks (multi-layer perceptrons). The motivation is that Newton's method, as a second-order optimizer that includes loss curvature information, should improve upon first-order optimizers, such as gradient descent, that use only first derivatives (gradient) of the loss. Second order methods tipically converge in a smaller number of iterations, but each update is expensive to compute and that's why they are not widely used in practice. It must be noted that Newton's method is almost never mentioned in the context of Neural Networks (NN), because it is only guaranteed to converge for convex loss functions (which is clearly not the case for NN). However, the authors claim they have some tricks to fix that.