Mean-field Behaviour of Neural Tangent Kernel for Deep Neural Networks

1 INTRODUCTION . Deep neural networks have achieved state-of-the-art results on numerous tasks ; see , e.g. , Nguyen & Hein ( 2018 ) , Du et al . ( 2018b ) , Zhang et al . ( 2017 ) . Although the loss function is not convex , Gradient Descent ( GD ) methods are often used successfully to learn these models . It has been actually recently shown that for certain overparameterized deep ReLU networks , GD converges to global minima ( ( Du et al. , 2018a ) ) . Similar results have been obtained for Stochastic Gradient Descent ( SGD ) ( ( Zou et al. , 2018 ) ) . The training dynamics of wide neural networks with GD is directly linked to kernel methods . Indeed , Jacot et al . ( 2018 ) showed that training a neural network with full batch GD in parameter space is equivalent to a functional GD i.e . a GD in a functional space with respect to a kernel called Neural Tangent Kernel ( NTK ) . Du et al . ( 2019 ) used a similar approach to prove that full batch GD converges to global minima for shallow neural networks and Karakida et al . ( 2018 ) linked the Fisher Information Matrix to the NTK and studied its spectral distribution for infinite width networks . The infinite width limit for different architectures was studied by Yang ( 2019 ) who introduced a tensor formalism that can express most of the computations in neural networks . Lee et al . ( 2019 ) studied a linear approximation of the full batch GD dynamics based on the NTK and gave an method to approximate the NTK for different architectures . Finally , Arora et al . ( 2019 ) gives an efficient algorithm to compute exactly the NTK for convolutional architectures ( Convolutional NTK or CNTK ) . In all of these papers , authors studied only the effect of infinite width on the NTK . The aim of this paper is to tackle the infinite depth limit . In parallel , the impact of the initialization and activation function on the performance of wide deep neural networks has been studied in Hayou et al . ( 2019 ) , Lee et al . ( 2018 ) , Schoenholz et al . ( 2017 ) , Yang & Schoenholz ( 2017 ) . These works analyze the forward/backward propagation of some quantities through the network at the initial step as a function of the initial parameters and the activation function . They propose a set of parameters and activation functions so as to ensure a deep propagation of the information at initialization . While experimental results in these papers suggest that such selection also leads to overall better training procedures ( i.e . beyond the initialization step ) , it remains unexplained why this is the case . In this paper , we link the initialization hyper-parameters and the activation function to the behaviour of the NTK which controls the training of DNNs , this could potentially explain the good performance . We provide a comprehensive study of the impact of the initialization and the activation function on the NTK and therefore on the resulting training dynamics for wide and deep networks . In particular , we show that an initialization known as the Edge of Chaos ( Yang & Schoenholz , 2017 ) leads to better training dynamics and that a class of smooth activation functions discussed in ( Hayou et al. , 2019 ) also improves the training dynamics compared to ReLU-like activation functions ( see also Clevert et al . ( 2016 ) ) . We illustrate these theoretical results through simulations . All the proofs are detailed in the Supplementary Material which also includes additional theoretical and experimental results . 2 MOTIVATION AND RELATED WORK . NEURAL TANGENT KERNEL Recent work by Jacot et al . ( 2018 ) has shown that the training dynamics of neural networks are captured by the Neural Tangent Kernel ( NTK ) . In the infinite width limit ( wide neural networks ) , the NTK converges to a kernel that remains unchanged as the training time grows . While this is only true in the infinite width limit , Lee et al . ( 2018 ) showed that a first order linear approximation of the training dynamics ( approximation of the NTK by its value at the initialization step ) leads to comparable performances for different architectures . More recently , Bietti & Mairal ( 2019 ) studied the RKHS of the NTK for a two layers convolutional neural network with ReLU activation and provided a spectral decomposition of the kernel , while in Arora et al . ( 2019 ) , the authors propose an algorithm to compute the NTK for convolutional neural networks . However , for finite width neural networks , Arora et al . ( 2019 ) observed a gap between the performances of the linear model derived from the NTK and the deep neural network , which is mainly due to the fact that the NTK changes with time . To fill this gap , Huang & Yau ( 2019 ) studied the dynamics of the NTK as a function of the training time for finite width neural networks and showed that the NTK dynamics follow an infinite hierarchy of ordinary differential equations baptised Neural Tangent Hierarchy ( NTH ) . In this paper , we consider the limit of infinite width neural networks ( mean-field approximation ) , and we study the behaviour of the NTK as the network depth goes to infinity assuming the width is infinite . EDGE OF CHAOS AND ACTIVATION FUNCTION Recent works by Hayou et al . ( 2019 ) and Schoenholz et al . ( 2017 ) have shown that weight initialization plays a crucial role in the training speed of deep neural networks ( DNNs ) . In Schoenholz et al . ( 2017 ) , the authors demonstrate that only a special initialization can lead to good performance . This initialization is known as the ’ Edge of Chaos ’ since it represents a transition between two phases : an ordered phase and a chaotic phase . When the DNN is initialized on the ordered phase , the output function of the DNN is constant almost everywhere , because the correlation of the outputs of two different inputs converges to 1 as the number of layers becomes large . On the other hand , when the DNN is initialized on the Chaotic phase , the output function is discontinuous almost everywhere as the depth goes to infinity . In this case , the correlation between the outputs of two different inputs converges to a value c such that |c| < 1 , therefore , very close inputs may lead to very different outputs . In Hayou et al . ( 2019 ) , authors give a comprehensive analysis of the Edge of Chaos , and further show that a certain class of smooth activation functions outperform the ReLU-like activation functions in term of test accuracy on MNIST and CIFAR10 . OUR CONTRIBUTIONS In this paper , we bridge the gap between the two previous concepts of Neural Tangent Kernel and Edge of Chaos Initialization for DNNs . More precisely , we study the impact of the Edge of Chaos initialization and the activation function on the NTK as the depth L goes to infinity . Our main results are : 1 . With an Initialization on the ordered/chaotic phase , the NTK converges exponentially to a constant kernel with respect to the depth L , making the training impossible for DNNs ( Lemma 1 and Proposition 1 ) 2 . The Edge of Chaos initialization leads to an invertible NTK even in the infinite depth limit , making the model trainable even for very large depths ( Proposition 2 ) 3 . The Edge of Chaos initialization leads to a sub-exponential convergence rate of the NTK to the limiting NTK ( w.r.t to L ) , which means that the ’ information ’ carried by the NTK propagates deeper compared to an initialization on the ordered/chaotic phase ( Propositon 2 ) 4 . Using a certain class S of smooth activation functions can further slow this convergence , making this class of activation functions more suitable for DNNs 5 . When adding Residual connections , we no longer need the initialisation on the Edge of Chaos , and the convergence of the NTK to the limiting NTK is always at a polynomial rate 3 NEURAL NETWORKS AND NEURAL TANGENT KERNEL . 3.1 SETUP AND NOTATIONS . Consider a neural network model consisting of L layers ( yl ) 1≤l≤L , with yl : Rnl−1 → Rnl , n0 = d and let θ = ( θl ) 1≤l≤L be the flattened vector of weights and bias indexed by the layer ’ s index and p be the dimension of θ . Recall that θl has dimension nl + 1 . The output f of the neural network is given by some transformation s : RnL → Ro of the last layer yL ( x ) ; o being the dimension of the output ( e.g . number of classes for a classification problem ) . For any input x ∈ Rd , we thus have f ( x , θ ) = s ( yL ( x ) ) ∈ Ro . As we train the model , θ changes with time t and we denote by θt the value of θ at time t and ft ( x ) = f ( x , θt ) = ( fj ( x , θt ) , j ≤ o ) . Let D = ( xi , yi ) 1≤i≤N be the data set and let X = ( xi ) 1≤i≤N , Y = ( yj ) 1≤j≤N be the matrices of input and output respectively , with dimension d×N and o×N . For any function g : Rd×o → Rk , k ≥ 1 , we denote by g ( X , Y ) the matrix ( g ( xi , yi ) ) 1≤i≤N of dimension k ×N . Jacot et al . ( 2018 ) studied the behaviour of the output of the neural network as a function of the training time t when the network is trained using a gradient descent algorithm . Lee et al . ( 2019 ) built on this result to linearize the training dynamics . We recall hereafter some of these results . For a given θ , the empirical loss is given by L ( θ ) = 1N ∑N i=1 ` ( f ( xi , θ ) , yi ) . The full batch GD algorithm is given by θ̂t+1 = θ̂t − η∇θL ( θ̂t ) ( 1 ) where η > 0 is the learning rate . Let T > 0 be the training time and Ns = T/η be the number of steps of the discrete GD equation 1 . The continuous time system equivalent to equation 1 with step ∆t = η is given by dθt = −∇θL ( θt ) dt ( 2 ) This differs from the result by Lee et al . ( 2019 ) since we use a discretization step of ∆t = η . It is well known that this discretization scheme leads to an error of order O ( η ) ( see Appendix ) . As in Lee et al . ( 2019 ) , Equation ( 2 ) can be re-written as dθt = − 1 N ∇θf ( X , θt ) T∇z ` ( f ( X , θt ) , Y ) dt where∇θf ( X , θt ) is a matrix of dimension oN × p and∇z ` ( f ( X , θt ) , Y ) is the flattened vector of dimension oN constructed from the concatenation of the vectors ∇z ` ( z , yi ) |z=f ( xi , θt ) , i ≤ N . As a result , the output function ft ( x ) satisfies the following ordinary differential equation dft ( x ) = ∇θf ( x , θt ) dθt = − 1 N ∇θf ( x , θt ) ∇θf ( X , θt ) T∇z ` ( ft ( X ) , Y ) dt ∈ Ro ( 3 ) The Neural Tangent Kernel ( NTK ) KLθ is defined as the o× o dimensional kernel satisfying : for all x , x′ ∈ Rd , KLθt ( x , x ′ ) = ∇θf ( x , θt ) ∇θf ( x′ , θt ) T ∈ Ro×o = L∑ l=1 ∇θlf ( x , θt ) ∇θlf ( x′ , θt ) T . ( 4 ) We also define KLθt ( X , X ) as the oN × oN matrix defined blockwise by KLθt ( X , X ) = KLθt ( x1 , x1 ) K L θt ( x1 , x2 ) · · · KLθt ( x1 , xN ) KLθt ( x2 , x1 ) · · · · · · K L θt ( x2 , xN ) · · · · · · · · · · · · KLθt ( xN , x1 ) K L θt ( xN , x2 ) · · · KLθt ( xN , xN ) By applying equation 3 to the vector X , one obtains dft ( X ) = − 1 N KLθt ( X , X ) ∇z ` ( ft ( X ) , Y ) dt , ( 5 ) meaning that for all j ≤ N dft ( xj ) = − 1NK L θt ( xj , X ) ∇z ` ( ft ( X ) , Y ) dt . Infinite width dynamics : In the case of a fully connected feedforward neural network ( FFNN ) of depth L and widths n1 , n2 , ... , nL , Jacot et al . ( 2018 ) proved that , with GD , the kernel KLθt converges to a kernel KL which depends only on L ( number of layers ) for all t < T when n1 , n2 , ... , nL → ∞ , where T is an upper bound on the training time , under the technical assumption ∫ T 0 ||∇z ` ( ft ( X , Y ) ) ||2dt < ∞ almost surely with respect to the initialization weights . The infinite width limit of the training dynamics is given by dft ( X ) = − 1 N KL ( X , X ) ∇z ` ( ft ( X ) , Y ) dt , ( 6 ) We note hereafter K̂L = KL ( X , X ) . As an example , with the quadratic loss ` ( z , y ) = 12 ||z − y|| 2 , equation 6 is equivalent to dft ( X ) = − 1 N K̂L ( ft ( X ) − Y ) dt , ( 7 ) which is a simple linear model that has a closed-form solution given by ft ( X ) = e− 1 N K̂ Ltf0 ( X ) + ( I − e− 1 N K̂ Lt ) Y . ( 8 ) For general input x ∈ Rd , we then have ft ( x ) = f0 ( x ) +K L ( x , X ) KL ( X , X ) −1 ( I − e− 1N K̂ Lt ) ( Y − f0 ( X ) ) ( 9 ) Note that in order for ft ( x ) to be defined , K̂L must be invertible . Indeed , it turns out that training with dynamics 6 is only possible if the NTK is invertible . We shed light on this behaviour in the following Lemma . Lemma 1 ( Trainability of the Neural Network and Invertibility of the NTK ) . Assume f0 ( X ) 6= Y . Then with dynamics defined by equation 8 , ||ft ( X ) − Y|| converges to 0 as t→∞ if and only if K̂L is non-singular . Moreover , if K̂L is singular , there exists a constant C > 0 such that for all t > 0 , ||ft ( X ) − Y|| ≥ C Lemma 1 shows that an invertible NTK is crucial for trainability . Since KLθt = K L is constant w.r.t to training time , it is completely determined at the initialization step . It is therefore intuitive to study the impact of the initialization on the NTK , particularly as the number of layers L grows ( Deep Neural Networks ) , which is our focus in this paper . Another interesting aspect is the impact of the NTK on the generalization error of the neural network model . To see this , if the NTK is constant for example ( i.e . there exists a constant δ such as KL ( x , x′ ) = δ for all x 6= x′ , this example is useful in the next section ) , then the second part of ft ( x ) in equation 9 is constant w.r.t x . Therefore , the generalization function ft ( x ) of the model 9 is entirely given by its value at time zero f0 ( x ) , which means that the generalisation error Ex , y [ ||ft ( x ) − y|| ] remains of order O ( 1 ) . In the next section , we show that the initialization and the activation function have major impact on the invertibility and ’ expressivity ’ of NTK . More precisely , we show that : 1 . Under some constraints , the NTK KL ( or a scaled version of the NTK ) converges to a limiting NTK K∞ as L goes to infinity ( otherwise it diverges ) 2 . A special initialization known as the Edge of Chaos ( EOC ) leads to an invertible K∞ which makes it useful for training DNNs 3 . The EOC initialization gives a sub-exponential rate for this convergence ( w.r.t L ) , which means for the same depth L , the EOC gives ’ richer ’ limiting NTK , and therefore leading to better generalization properties 4 . The smoothness of the activation can further slow this convergence , leading to ’ richer ’ limiting NTK ( the convergence to the limiting trivial kernel is slower ) 5 . Adding Residual connections leads to sub-exponential convergence rate for the NTK ( w.r.t to L ) and we no longer need the Edge of Chaos

Summary: This paper discussed the property of the NTK with the increasing depth L with the help of the Edge of Chaos initialization. The authors show that if deep neural networks are not properly initialized, the NTK can have a large condition number, which leads to the poor performance of training and generalization. Moreover, the authors also introduce the conditions that make the neural network trainable by decreasing the convergence rate to a nearly constant kernel w.r.t the depth L by using the specific Edge of Chaos initialization as well as different activations and use residual connections. Experiment results show that the theoretical results are well aligned with the practice.