M EAN - FIELD B EHAVIOUR OF N EURAL T ANGENT K ERNEL FOR D EEP N EURAL N ETWORKS
Anonymous authors
Paper under double-blind review
A BSTRACT
Recent work by Jacot et al.
(2018) has showed that training a neural network of any kind with gradient descent in parameter space is equivalent to kernel gradient descent in function space with respect to the Neural Tangent Kernel (NTK).
Lee et al.
(2019) built on this result to show that the output of a neural network trained using full batch gradient descent can be approximated by a linear model for wide networks.
In parallel, a recent line of studies (Schoenholz et al.
(2017), Hayou et al.
(2019)) suggested that a special initialization known as the Edge of Chaos leads to good performance.
In this paper, we bridge the gap between these two concepts and show the impact of the initialization and the activation function on the NTK as the network depth becomes large.
We provide experiments illustrating our theoretical results.
I NTRODUCTION
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.
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 inﬁnite width networks.
The inﬁnite 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 efﬁcient algorithm to compute exactly the NTK for convolutional architectures (Convolutional NTK or CNTK).
In all of these papers, authors studied only the effect of inﬁnite width on the NTK.
The aim of this paper is to tackle the inﬁnite 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.
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.
We illustrate these theoretical results through simulations.
All the proofs are detailed in the Supplementary Material which also includes additional theoretical and experimental results.
M OTIVATION AND R ELATED WORK N EURAL T ANGENT K ERNEL
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 inﬁnite 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 inﬁnite width limit, Lee et al.
(2018) showed that a ﬁrst 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 ﬁnite 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 ﬁll this gap, Huang & Yau (2019) studied the dynamics of the NTK as a function of the training time for ﬁnite width neural networks and showed that the NTK dynamics follow an inﬁnite hierarchy of ordinary differential equations baptised Neural Tangent Hierarchy (NTH).
In this paper, we consider the limit of inﬁnite width neural networks (mean-ﬁeld approximation), and we study the behaviour of the NTK as the network depth goes to inﬁnity assuming the width is inﬁnite.
E DGE OF C HAOS AND A CTIVATION F UNCTION
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 inﬁnity.
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.
O UR 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 inﬁnity.
Our main results are :
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 inﬁnite 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)
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
N EURAL N ETWORKS AND N EURAL T ANGENT K ERNEL 3. S ETUP AND NOTATIONS
Consider a neural network model consisting of L layers ( y l ) 1 ≤ l ≤ L , with y l : R n l − 1 → R n l , n 0 = d and let θ =
( θ l ) 1 ≤ l ≤ L be the ﬂattened vector of weights and bias indexed by the layer’s index and p be the dimension of θ .
Recall that θ l has dimension n l + 1 .
The output f of the neural network is given by some transformation s : R n L → R o of the last layer y L ( x ) ; o being the dimension of the output (e.g. number of classes for a classiﬁcation problem).
For any input x ∈ R d , we thus have f ( x, θ )
=
s ( y L ( x ))
∈ R o .
As we train the model, θ changes with time t
and we denote by θ t the value of θ at time t and f t ( x ) =
f
( x, θ t )
=
( f j ( x, θ t ) , j ≤ o ) .
Let D =
( x i , y i ) 1 ≤ i ≤ N be the data set and let X =
( x i ) 1 ≤ i ≤ N , Y = ( y j ) 1 ≤ j ≤ N be the matrices of input and output respectively, with dimension d × N and o × N .
For any function g : R d × o
→ R k , k ≥ 1 , we denote by g ( X , Y ) the matrix ( g ( x i , y i ))
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 ( θ ) = 1 N (cid:80)
Ni =1 (cid:96) ( f ( x i , θ ) , y i ) .
The full batch GD algorithm is given by
[Equation]
where η > 0 is the learning rate.
Let T > 0 be the training time and N s = 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
[Equation]
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
[Equation]
where ∇ θ f ( X , θ t ) is a matrix of dimension oN × p and ∇ z (cid:96) ( f ( X , θ t ) , Y ) is the ﬂattened vector of dimension oN constructed from the concatenation of the vectors ∇ z (cid:96)
( z, y i ) | z
= f
( x i ,θ t ) , i ≤ N .
As a result, the output function f t ( x ) satisﬁes the following ordinary differential equation
[Equation]
The Neural Tangent Kernel (NTK) K Lθ is deﬁned as the o × o dimensional kernel satisfying: for all x, x (cid:48) ∈ R d ,
[Equation]
We also deﬁne K Lθ t ( X , X ) as the oN × oN matrix deﬁned blockwise by
[Equation]
 By applying equation 3 to the vector X , one obtains
[Equation]
meaning that for all j ≤ N df t ( x j ) = − 1 N K Lθ t ( x j , X ) ∇ z (cid:96) ( f t ( X ) , Y ) dt .
Inﬁnite width dynamics :
In the case of a fully connected feedforward neural network (FFNN) of depth L and widths n 1 , n 2 , ..., n L , Jacot et al.
(2018) proved that, with GD, the kernel K Lθ t converges to a kernel K L which depends only on L (number of layers) for all t < T when n 1 , n 2 , ..., n L → ∞ , where T is an upper bound on the training time, under the technical assumption (cid:82) T 0 ||∇ z (cid:96) ( f t ( X , Y ))
||
2 dt < ∞ almost surely with respect to the initialization weights.
The inﬁnite width limit of the training dynamics is given by
[Equation]
We note hereafter ˆ K L = K L ( X , X ) .
As an example, with the quadratic loss (cid:96) ( z, y )
=
12 || z − y ||
2 , equation 6 is equivalent to
[Equation]
(7) which is a simple linear model that has a closed-form solution given by
[Equation]
For general input x ∈ R d , we then have
[Equation]
Note that in order for f t ( x ) to be deﬁned, ˆ 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 f 0 ( X ) (cid:54) =
Y .
Then with dynamics deﬁned by equation 8, || f t ( 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 ,
[Equation]
Lemma 1 shows that an invertible NTK is crucial for trainability.
Since K Lθ 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 K L ( x, x (cid:48) )
=
δ for all x (cid:54) =
x (cid:48) , this example is useful in the next section), then the second part of f t ( x ) in equation 9 is constant w.r.t x .
Therefore, the generalization function f t ( x ) of the model 9 is entirely given by its value at time zero f 0 ( x ) , which means that the generalisation error E x,y [ || f t ( 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 :
Under some constraints, the NTK K L (or a scaled version of the NTK) converges to a limiting NTK K ∞ as L goes to inﬁnity (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
I MPACT OF THE I NITIALIZATION AND THE A CTIVATION FUNCTION ON THE N EURAL T ANGENT K ERNEL
In this section we study the impact of the initialization and the activation function on the limiting NTK for Fully-connected Feed-forward Neural Networks (FFNN).
We prove that only an initialization on the Edge of Chaos (EOC) leads to an invertible NTK for deep neural networks.
All other initializations will lead to a trivial non-invertible NTK.
We also show that the smoothness of the activation function plays a major role in the behaviour of NTK.
To simplify notations, we restrict ourselvs to the case s ( x ) = x and o = 1 , since generalization to any function s and any n L is straightforward.
Consider a FFNN of depth L , widths ( n l ) 1 ≤ l ≤ L , weights w l and bias b l .
For some input x ∈ R d , the forward propagation is given by
[Equation]
where φ is the activation function.
We initialize the model with w lij iid ∼ N (0 , σ 2 w n l − 1 ) and b li iid ∼ N (0 , σ 2 b ) , where N ( µ, σ 2 ) denotes the normal distribution of mean µ and variance σ 2 .
For some x , we denote by q l ( x ) the variance of y l ( x ) .
The convergence of q l ( x ) as l increases is studied in Lee et al.
(2018), Schoenholz et al.
(2017) and
Hayou et al.
In particular, under weak regularity conditions they prove that q l ( x ) converges to a point q ( σ b , σ w ) > 0 independent of x as l → ∞ .
Also the asymptotic behaviour of the correlations between y l ( x ) and y l ( x (cid:48) ) for any two inputs x and x (cid:48) is driven by ( σ b , σ w ) ; the authors deﬁne the
EOC as the set of parameters ( σ b , σ w )
such that σ 2 w E [ φ (cid:48) ( (cid:112) q ( σ b , σ w ) Z ) 2 ] = 1 where Z ∼ N (0 , 1) .
Similarly the Ordered, resp.
Chaotic, phase is deﬁned by σ 2 w E [ φ (cid:48) ( (cid:112) q ( σ b , σ w ) Z ) 2 ] < 1 , resp.
σ 2 w E [ φ (cid:48) ( (cid:112) q ( σ b , σ w ) Z ) 2 ] > 1 ; more details are recalled in Section 2 of the supplementary material.
It turns out that the EOC plays also a crucial role on the NTK.
Let us ﬁrst deﬁne two classes of activation functions.
Deﬁnition 1.
Let φ : R → R be a measurable function.
Then
φ is said to be ReLU-like if there exist λ, β ∈ R such that φ ( x ) =
λx for x > 0 and φ ( x ) = βx for x ≤ 0 .
φ is said to be in S if φ (0) = 0 , φ is twice differentiable, and there exist n ≥ 1 , a partition ( A i ) 1 ≤ i ≤ n of R and inﬁnitely differentiable functions g 1 , g 2 , ..., g n such that φ (2)
=
(cid:80) ni
=1 1 A i g i , where φ (2) is the second derivative of φ .
The class of ReLU-like activations includes ReLU and Leaky-ReLU, whereas the S class includes, among others, Tanh, ELU and SiLU (Swish).
The following proposition establishes that any initialization on the Ordered or Chaotic phase, leads to a trivial limiting NTK as the number of layers L becomes large.
Proposition 1 (Limiting Neural Tangent Kernel with Ordered/Chaotic Initialization) .
Let ( σ b , σ w ) be either in the ordered or in the chaotic phase.
Then, there exist λ, γ > 0 such that
[Equation]
As a result, as L goes to inﬁnity, K L converges to a constant kernel K ∞ ( x, x (cid:48) )
=
λ for all x, x (cid:48) ∈ R d .
The training is then impossible.
Indeed, we have K L ( X , X )
≈ λU where U is the matrix with all elements equal to one, i.e. ˆ K L is at best degenerate and asymptotically (in L ) non invertible, rendering the training impossible by Lemma 1.
We illustrate empirically this result in Section 5.
Recall that the (matrix) NTK for input data X is given by
[Equation]
As shown in Schoenholz et al.
(2017) and Hayou et al.
(2019), an initialization on the EOC preserves the norm of the gradient as it back-propagates through the network.
This means that the terms ∇ θ l f ( X , θ t ) ∇ θ l f ( X , θ t ) T are of the same order.
Hence, it is more convenient to study the average
NTK (ANTK hereafter) given by K L /L .
Note that the invertibility of the NTK is equivalent to that of the ANTK.
The next proposition shows that on the EOC, the ANTK converges to an invertible kernel as L → ∞ at a sub-exponential rate.
Moreover, by choosing an activation function in the class S , we can slow the convergence of ANTK with respect to L , which means that, for the same depth L , a smooth activation function from the class S leads to ’richer’ NTK which is crucial for the generalization error of deep models as discussed in Section 3.
This conﬁrms the ﬁndings in (Hayou et al., 2019).
Proposition 2 (Neural Tangent Kernel on the Edge of Chaos) .
Let φ be a non-linear activation function and ( σ b , σ w ) ∈ EOC .
If φ is ReLU-like, then for all x ∈ R d , K L ( x,x ) L
[Equation]
Moreover, there exist
A, λ ∈ (0 , 1) such that
[Equation]
If φ is in S , then, there exists q > 0 such that K L ( x,x ) L
[Equation]
x
[Equation]
→ q .
Moreover, there exist B, C, λ ∈ (0 , 1) such that
[Equation]
Since 0 < λ < 1 , on the EOC there exists a matrix J invertible such that K L ( X , X ) =
L × J (1+ o (1)) as L → ∞ .
Hence, although the NTK grows linearly with L , it remains asymptotically invertible.
This makes the training possible for deep neural networks when initialized on the EOC, contrariwise to an initialization on the Ordered/Chaotic phase, see Proposition 1).
However the limiting kernels
K ∞ carry (almost) no information on x, x (cid:48) and have therefore little expressive power.
Interestingly the convergence rate of the ANTK to K ∞ is slow in L ( O ( L − 1 ) for ReLU-like activation functions and O (log( L ) L − 1 ) for activation functions of type S ).
This means that as L grows, the NTK remains expressive compared to the Ordered/Chaotic phase case (exponential convergence rate).
This is particularly important for the generalization part (see equation 9).
The log( L ) gain obtained when using smooth activation functions of type S means we can train deeper neural networks with this kind of activation functions compared to the ReLU-like activation functions and could explain why
ELU and Tanh tend to perform better than ReLU and Leaky-ReLU (see Section 5).
Another important feature of deep neural network which is known to be highly inﬂuential is their architecture.
The next proposition shows that adding residual connections to a ReLU network leads to a polynomial rate for wide range of initialization parameters.
Proposition 3 (Residual connections) .
Consider the following network architecture (FFNN with residual connections)
[Equation]
with initialization parameters σ b = 0 and σ w > 0 and φ is the ReLU activation function.
Let K Lres be the corresponding NTK.
Then for all x ∈ R d , K Lres ( x,x ) α L × 2 L = || x
|| 2 d [Equation] O ( γ L ) and there exists λ ∈ (0 , 1) such that
[Equation]
where α l and γ l are given by
• if σ w < √ 2 , then α L = 1 and γ L = ( 1+ σ 2 w / 2)
L • if σ w = √ 2 , then α L = L and γ L =
L −
[Figure]
• if σ w > √ 2 , then α L =
( 1+ σ 2 w / 2) L and γ L = ( 1+ σ 2 w / 2) − L
Proposition 3 shows that the NTK of a ReLU FFNN with residual connections explodes exponentially with respect to L .
However, the normalised kernel K Lres ( x, x (cid:48) ) /α
L 2 L where x (cid:54) =
x (cid:48) converges to a limiting kernel similar to K ∞ with a rate O
( L − 1 ) for all σ w > 0 .
We say that residual networks ’live’ on the Edge of Chaos, i.e. no matter what the choice of σ w is, the convergence rate of the NTK w.r.t
L is polynomial and there is no Ordered/Chaotic phase in this case.
This could potentially explain why residual networks perform better than FFNN (RELU) in many tasks when the initialization is not on the EOC.
We illustrate this result in section 5.
E XPERIMENTS
In this section, we illustrate empirically the theoretical results obtained in the previous sections.
We ﬁrst illustrate the results of Propositions 1, 2 and 3.
Then, we conﬁrm the impact of the EOC and
Activation function on the overall performance of the model (FFNN), on MNIST and CIFAR datasets.
5. C ONVERGENCE RATE OF K L AS L GOES TO INFINITY
Propositions 1, 2 and 3 give theoretical convergence rates for quantities of the form (cid:12)(cid:12)
K L α L − K ∞ | .
We illustrate these results in Figure 1.
Figure 1a shows a convergence rate approximately equal to
O ( L − 1 ) for ReLU and ELU.
Recall that for ELU the exact rate is O (log( L ) L − 1 ) but one cannot observe experimentally the logarithmic factor.
However, ELU performs indeed better than ReLU (see
Table 1) which might be explained by this log( L ) factor.
Figure 1b demonstrates that this convergence occurs at an exponential convergence rate in the Ordered phase for both ReLU and ELU, and Figure 1c the convergence rate in the case of FFNN with residual connections.
As predicted by Proposition 3, the convergence rate O ( L − 1 ) is independent of the parameter σ w .
5. I MPACT OF THE INITIALIZATION AND SMOOTHNESS OF THE ACTIVATION ON THE OVERALL PERFORMANCE
We train FFNN of width 300 and depths L ∈ { 200 , 300 } and width ∈ { 200 , 300 } with SGD and categorical cross-entropy loss.
Training with full batch GD is practically impossible for DNNs, so we use SGD instead (see Section D in the Appendix for more details about how the results extend to
SGD) with a batchsize of 64 and a learning rate 10 − 3 for L = 100 and 10 − 4 for L ∈ 200 , 300 (this learning rate was found by a grid search of exponential step size 10).
For each activation function, we use an initialization on the EOC when it exists, we add the symbol (EOC) after the activation when this is satisﬁed.
We use ( σ b , σ w ) =
(0 , √ 2) for ReLU, ( σ b , σ w ) =
(0 . 2 , 1 . 227) for ELU and ( σ b , σ w ) =
(0 . 2 , 1 . 302) for Tanh.
These values are all on the EOC (see Hayou et al.
for more details).
Table 1 displays the test accuracy for different activation functions on MNIST and CIFAR10 after 10 and 100 training epochs for depth 300 and width 300.
Functions in class S (ELU and Tanh) perform much better than ReLU-like activation functions (ReLU, Leaky-Relu − α with α ∈ { 0 . 01 , 0 . 02 , 0 . 03 } ).
Even with Parametric ReLU (PReLU) where the parameter of the leaky-ReLU is also learned by backpropagation, we obtain only a small improvement over ReLU.
For
[Figure]
[Table]
activation functions that do not have an EOC, such as Softplus and Sigmoid, we use He initialization for MNIST and Glorot initialization for CIFAR10 (see He
et al.
and Glorot & Bengio (2010)).
For Softplus and Sigmoid, the training algorithm is stuck at a low test accuracy ∼ 10% which is the test accuracy of a uniform random classiﬁer with 10 classes.
C ONCLUSION AND LIMITATIONS
That the training dynamics of deep neural networks is equivalent to a Functional Gradient Descent with respect to the Neural Tangent Kernel.
In the inﬁnite width limit, the NTK has a closed-form expression.
This approximation sheds light on how the NTK impacts the training dynamics: it controls the training rate and the generalization function.
Using this approximation for wide neural networks (Mean-ﬁeld approximation), we show that for an initialization in the Ordered/Chaotic phase, NTK converges exponentially fast to a non-invertible kernel as the number of layers goes to inﬁnity, making training impossible.
An initialization on the EOC leads to an invertible ANTK (and NTK) even for an inﬁnite number of layers: the convergence rate is O
( L − 1 ) for ReLU-like activation functions and O (log( L ) L − 1 ) for a class of smooth activation functions.
However, recent ﬁndings showed that the inﬁnite width approximation of the NTK does not fully capture the dynamics of the training of DNNs.
A recent line of work showed that the NTK for ﬁnite width neural networks changes with time and might even be random (Chizat & Bach (2018),
Ghorbani et al.
(2019), Huang & Yau (2019), Arora et al.
Therefore, we believe that the NTK is a useful tool to partially understand wide deep neural networks (have insights on hyper-parameters choices for example) and not a tool to train neural networks.
A A PPENDIX : P ROOFS OF S ECTION 3: N EURAL N ETWORKS AND N EURAL T ANGENT K ERNEL
Lemma 1 (Trainability of the Neural Network and Invertibility of the NTK) .
Assume f 0 ( X ) (cid:54) =
Y .
Then with dynamics (8), || f t ( 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 ,
[Equation]
Proof.
Assume f 0 ( X ) (cid:54) =
Y .
Let ˆ K L = Q T DQ be the spectral decomposition of the empirical NTK; i.e. Q is an orthogonal matrix and D is a diagonal matrix.
We have that e − 1 N ˆ K L t =
Q T e − 1 N Dt Q = Q T Diag ( e − diN t ) 1 ≤ i ≤ oN Q where ( d i ) 1 ≤ i ≤ oN are the eigenvalues.
We also have || f t ( X ) − Y|| =
||
e − 1 N ˆ K L t ( f 0 ( X ) − Y )
|| .
Therefore, the equivalence holds true.
Moreover, assume ˆ K L in singular.
Let Z t =
Q ( f t ( X ) − Y ) Q T .
We have that Z t =
e − 1 N Dt Z 0 .
Since D has at least one zero diagonal value, then there exists j ∈ { 1 , 2 , ..., oN } such that for all t , ( Z t )
j = ( Z 0 )
j , and we have
[Equation]
Lemma 2 (Discretization Error for Full-Batch Gradient Descent) .
Assume ∇ θ L is C -lipschitz, then there exists C (cid:48) > 0 that depends only on C and T such that
[Equation]
Proof.
For t ∈ [
[Equation]
T ]
[Equation]
we deﬁne the stepwise constant system ˜ θ t
= ˆ θ
[Equation]
t/η (cid:99) .
Let t ∈
[0 , T ]
[Equation]
we have
[Equation]
Therefore,
[Equation]
) ||
Moreover, for any k ∈ [
[Equation]
(cid:98)∗(cid:99) T/η ] , we have
[Equation]
where we have used log(1 + ηC ) ≤ ηC .
Using this result, there exists a constant ˜ C depending on T and C such that
[Equation]
Now we have
[Equation]
so we can conclude using Gronwall’s lemma.
B P ROOFS OF S ECTION 4: I MPACT OF THE I NITIALIZATION AND THE A CTIVATION FUNCTION ON THE N EURAL T ANGENT K ERNEL
We ﬁrst recall the results obtained in Lee et al.
(2018), Schoenholz et al.
(2017) and Hayou et al.
(2019) where the impact of the EOC (Edge of Chaos) on the initialization is studied.
We also present some results that we will be used below.
Consider a FFNN of depth L , widths ( n l ) 1 ≤ l ≤ L , weights w l and bias b l .
For some input x ∈ R d , the forward propagation is given by
[Equation]
where φ is the activation function.
We initialize the model with w
[Equation]
iid ∼
[Equation] [Equation]
w n l − 1 ) and b li iid ∼ N (0 , σ 2 b ) , where N ( µ, σ 2 ) denotes the normal distribution of mean µ and variance σ 2 .
For some x , we denote by q l ( x ) the variance of y l ( x ) .
In general, q l ( x ) converges to a point q ( σ b , σ w ) > 0 independent of x as l → ∞ .
The EOC is deﬁned by the set of parameters ( σ b , σ w ) such that σ 2 w E [ φ (cid:48)
( (cid:112) q ( σ b , σ w ) Z ) 2 ] = 1 where Z ∼ N (0 , 1) .
Similarly the Ordered, resp.
Chaotic, phase is deﬁned by σ 2 w E [ φ (cid:48) ( (cid:112) q ( σ b , σ w ) Z ) 2 ] < 1 , resp.
σ 2 w E [ φ (cid:48) ( (cid:112) q ( σ b , σ w ) Z ) 2 ] > 1 (see Hayou et al.
for more details).
For two inputs x, x (cid:48) ∈ R d , deﬁne Σ l ( x, x (cid:48) )
=
E [ y l ( x ) y l ( x (cid:48) )] and let c l ( x, x (cid:48) ) be the corresponding correlation.
Let f be the correlation function deﬁned implicitly by c l +
= f
( c l ) .
In the limit of inﬁnitely wide networks, we have the following results (Hayou et al.
• Σ l ( x, x (cid:48) )
=
σ 2 b + σ 2 w E z ∼N (0 , Σ l − 1 ) [ φ ( z ( x ))
φ
( z ( x (cid:48) ))] .
• There exist q, λ > 0 such that, for all sup x ∈ R d | Σ l ( x, x ) − q
|
≤ e − λl .
• On the Ordered phase, there exists γ > 0 such that sup x,x (cid:48) ∈ R d | c l ( x, x (cid:48) )
− 1 | ≤ e − γl .
• On the chaotic phase, there exist γ > 0 and c < 1 such that sup x (cid:54) =
x (cid:48) ∈ R d | c l ( x, x (cid:48) )
− c | ≤ e − γl .
• For ReLU network on the EOC, we have that Σ l ( x, x ) =
σ 2 w d || x
|| 2 for all l ≥ 1 .
Moreover, we have
[Equation]
• In general, we have
[Equation]
q where Z ( x ) = xZ 1 + √ 1 − x 2 Z 2 and Z 1 , Z 2 are iid standard Gaussian variables.
• On the EOC, we have f (cid:48) (1) =
• If φ is k -times differentiable, then f is k -times differentiable and for all 1 ≤ j ≤ k , we have f ( j ) ( x ) =
σ 2 w q j −
[Equation]
[ φ ( j ) ( Z 1 ) φ ( j ) ( Z ( x )
[Equation]
• From Jacot et al.
(2018), we have that
[Equation]
where the deﬁnition of ˙Σ l ( x, x (cid:48) ) is given in Proposition 1 below.
Deﬁnition 1.
Let φ : R → R be a measurable function.
Then
φ is said to be ReLU-like if there exist λ, β ∈ R such that φ ( x ) =
λx for x > 0 and φ ( x ) = βx for x ≤ 0 .
φ is said to be in S if φ (0) = 0 , φ is twice differentiable, and there exist n ≥ 1 , a partition ( A i ) 1 ≤ i ≤ n of R and inﬁnitely differentiable functions g 1 , g 2 , ..., g n such that φ (2)
=
(cid:80) ni
=1 1 A i g i , where φ (2) is the second derivative of φ .
The following two lemmas will be useful to prove the results of Section 3 in the main paper.
Lemma 2.
Let ( a l ) be a sequence of non-negative real numbers such that ∀ l ≥ 0 , a l +1 ≤ αa l + ke − βl , where α ∈ (0 , 1) and k, β > 0 .
Then there exists γ > 0 such that ∀ l ≥ 0 , l ≤ e − γl .
Proof.
Using the inequality on a l , we can easily see that
[Equation]
where we divided the sum into two parts separated by index l/ 2 and upper-bounded each part.
The existence of γ is straightforward.
Proposition 1 (Limiting Neural Tangent Kernel with Ordered/Chaotic Initialization) .
Let ( σ b , σ w ) be in the ordered or chaotic phase.
Then, there exist λ, γ > 0 such that
[Equation]
Proof.
From Jacot et al.
(2018), we have that
[Equation]
where Σ 1 ( x, x (cid:48) )
=
σ 2 b + σ 2 w d x T x (cid:48) and Σ l ( x, x (cid:48) )
=
σ 2 b + σ 2 w E f ∼N (0 , Σ l − 1 ) [ φ ( f ( x )) φ ( f ( x (cid:48) ))] and ˙Σ l ( x, x (cid:48) )
=
E f ∼N (0 , Σ l − 1 ) [ φ (cid:48)
( f ( x )) φ (cid:48)
( f ( x (cid:48) ))] .
In the ordered/chaotic phase, Hayou et al.
(2019) showed that there exist k, γ, l 0 > 0 and α ∈ (0 , 1) such that for all l ≥ l 0 we have
[Equation]
and
[Equation]
Therefore we have for any l ≥ l 0 and x, x (cid:
) ∈ R d
[Equation]
Letting r l = K l ( x, x (cid:48) )
− k 1 − α , we have
[Equation]
We can now conclude using Lemma 2.
Now, we show that the Initialization on the EOC leads to an invertible NTK even if the number of layers L goes to inﬁnity.
We ﬁrst prove two preliminary lemmas that will be useful for the proof of the next proposition.
Lemma 3.
Let ( a l ) , ( b l ) , ( λ l ) be three sequences of real numbers such that
[Equation]
where α ∈ N ∗ , β, q ∈ R + and α > β − 1 .
Then,
[Equation]
Proof.
It is easy to see that there exists a constant G > 0 | a l | ≤ G × l + | a 0 | for all l ≥ 0 , therefore ( a l /l ) is bounded.
Now let ζ = min(1 , β ) and r l = a l l .
We have
[Equation]
Letting x l = (cid:12)(cid:12)
r l − q 1+ α (cid:12) (cid:12) , there is exist a constant M > 0 such that
[Equation]
Hence, we have
[Equation]
By taking the logarithm of the ﬁrst term in the right hand side and using the fact that (cid:80) lk =1 1 k ∼ log( l ) , we have
[Equation]
For the second part, observe that
[Equation]
−
[Equation]
− 1)! l !
[Equation]
and
[Equation]
so that, l (cid:88)
[Equation]
therefore, l (cid:88)
[Equation]
We can now conclude using the fact that α > β − 1 .
We now introduce a different form of the previous Lemma that will be useful for other applications.
Lemma 4.
Let ( a l ) , ( b l ) , ( λ l ) be three sequences of real numbers such that
[Equation]
where α ∈ N ∗ , β, q ∈ R + and α > β − 1 , β ≥ 1 .
Then, there exists A, B > 0 such that
[Equation]
Proof.
It is easy to see that there exists a constant G > 0 | a l | ≤ G × l + | a 0 | for all l ≥ 0 , therefore ( a l /l ) is bounded.
Let r l = a l l .
We have
[Equation]
Let x l = r l − q 1+ α .
It is clear that λ l = 1 − α/l + O ( l − 3 / 2 ) .
Therefore, using Lemma 3 with β = 1 / 2 , we have r l → q 1+ α .
Thus, there exists κ 1 , κ 2 , M, l 0 > 0 such that for all l ≥ l
[Equation]
Similarly to the proof of Lemma 3, it follows that
[Equation]
k = ll (cid:89) j = k + (1 − 1 + α j ) κ 2 log( k ) + M kand x l ≥ xl (cid:89)
k =(1 − 1 + α k )
+ l (cid:88) k
= ll (cid:89)
j = k + (1 − 1 + α j ) κ 1 log( k ) − M k
Recall that we have l (cid:89) k
= (1 − 1 + α k ) ∼ l − 1 − α and l (cid:89)
j = k + (1 − 1 + α j ) =
( l − α − 1)!
l ! k !
( k − α − 1)!
so that k !
( k − α − 1)!
κ 1 log( k ) − M k∼ k →∞ log( k ) k α −
Therefore, we obtain l (cid:88) k
= k !
( k − α − 1)!
κ 1 log( k ) − M k∼ l (cid:88) k
= log( k ) k
α −∼ (cid:90) l 1 log( t ) t α − 1 dt ∼ C 1 l α log( α ) where C 1 > 0 is a constant.
Similarly, there exists a constant C 2 > 0 such that
[Equation]
We conclude using the fact that ( l − α − 1)!
l ! ∼ l − 1 − α .
Proposition 2 (Neural Tangent Kernel on the Edge of Chaos) .
Let φ be a non-linear activation function and ( σ b , σ w ) ∈ EOC .
If φ is ReLU-like, then for all x ∈ R d , K L ( x,x ) L
[Equation]
Moreover, there exist A, λ ∈ (0 , 1) such that
[Equation]
If φ is in S , then, there exist q > 0 such that K L ( x,x ) L
[Equation]
→ q .
Moreover, there exist B, C, λ ∈ (0 , 1) such that
[Equation]
L
[Equation]
Proof.
We use some results from Hayou et al.
in this proof.
Let x, x (cid:48) ∈ R d and c lx,x (cid:48)
=
Σ ( x,x (cid:48) ) √ Σ ( x,x )Σ ( x (cid:48) ,x (cid:48) ) .
Let γ l := 1 − c lx,x (cid:48) and f be the correlation function deﬁned by the recursive equation c l +
= f
( x l ) .
From the preliminary results, we know that Σ l ( x, x ) =
σ 2 w d || x
||
and that K l ( x, x (cid:48) )
=
K l − 1 ( x, x (cid:48) )
˙Σ l ( x, x (cid:48) )
+ Σ l ( x, x (cid:48) ) .
This concludes the proof for K L ( x, x ) .
We denote s = 2 √ 2 3 π .
From Hayou et al.
(2019), we have on the EOC γ l +1 =
γ l − sγ 3 / 2 l + O ( γ
[Equation] [Equation]
so that
[Equation]
Thus, as l goes to inﬁnity γ − 1 /
[Equation] [Equation]
γ
[Equation]
∼ s and by summing and equivalence of positive divergent series γ − 1 / 2 l ∼ s 2 l.
Moreover, since γ − 1 / 2 l +1 − γ − 1 / 2 l = s 2 + O ( γ l ) = s 2 + O ( l − 2 ) , we have γ − 1 / 2 l = s 2 l + O (1) .
Therefore, c lx,x (cid:48)
= 1 − 9 π 2 2 l 2 + O ( l − 3 ) .
we also have
[Equation]
Thus, it follows that
[Equation]
+ O ( l − 2 )
Moreover, q lx,x (cid:48)
=
q + O ( l − 2 ) where q is the limiting variance of y l .
Using Lemma 3, we conclude that K l ( x,x (cid:48) )
l
[Equation]
|
[Equation]
Since c x,x (cid:48) is bounded, this result is uniform in x, x (cid:48) .
Therefore, we can take the supremum over x, x (cid:48) ∈ R d .
We prove the result when φ (2) ( x ) = 1 x< 0 g 1 ( x ) + 1 x ≥
g 2 ( x ) .
The generalization to the whole class is straightforward.
Let f be the correlation function.
We ﬁrst show that for all k ≥ 3 f ( k ) ( x ) = 1 (1 − x 2 ) ( k − 2) /
2 g k ( x ) where g k ∈ C ∞ .
We have
[Equation]
Let G ( x ) =
E [ φ (cid:48)(cid:48) ( √ qZ 1 )1 U 2 ( x ) < 0 g 1 ( √ qU 2 ( x ))] then
[Equation]
It is easy to see that G (cid:48)
( x ) = 1 √ 1 − x 2 G 1 ( x ) where G 1 ∈ C 1 .
A similar analysis can be applied to the second term of f (cid:48)(cid:48) .
We conclude that there exists g 3 ∈ C ∞ such that f (3) ( x ) = 1 √ 1 − x 2 g ( x ) .
We obtain the result by induction.
Since f ( k ) are potentially not deﬁned at 1, we use the change of variable x = 1 − t 2 to obtain a Taylor expansion near 1.
Simple algebra shows that the function t → f (1 − t 2 ) has a Taylor expansion near 0:
[Equation]
Therefore,
[Equation]
Letting λ l := 1 − c l , there exist α, β > 0 such that
[Equation]
therefore,
[Equation]
By summing (divergent series), we have that λ − 1 l ∼ lβ q .
Therefore,
[Equation]
By summing a second time, we obtain
[Equation]
O
[Equation]
1) so that λ l = α − 1 l − 1 − α − 1 β log( l ) l+
O
( l − 2 ) .
Using the fact that f (cid:48) ( x ) = 1 + ( x − 1) f (cid:48)(cid:48) (1) + O (( x − 1) 2 ) , we have f (cid:48)
( c lx,x (cid:48) )
= 1 − 2 l + κ log( l ) l 2 + O ( l − 2 ) .
We can now conclude using Lemma 4.
Using again the agrument of the boundedness of c 1 x,x (cid:48) , we can take the supremum.
Proposition 3.
Consider the following network architecture (FFNN with residual connections)
[Equation]
with initialization parameters σ b = 0 and σ w > 0 .
Let K Lres be the corresponding NTK.
For all x ∈ R d , K Lres ( x,x ) L × 2 L
[Equation]
|
[Equation]
|
[Equation]
d + O ( L − 1 ) and there exists λ ∈ (0 , 1) such that
[Equation]
Proof.
We only give a sketch of the proof.
A more rigorous proof can be easily done but it unecessary.
As in the feedforward without residual connections case, it is easy to see that K Lres satisﬁes the following recursive equation
[Equation]
To see this, with FeedForward Neural Network, we have ( Jacot et al. [Equation] [Equation] I
[Equation] whereas for Residual Networks, we have an extra term ∂y l {++
 i [Equation]+}
where I W , I (cid:48)
W are quantities that converge to 0 as n l grows (self-averaging).
The additional term results in an added term K l ( x, x (cid:48) ) in the formula of K l +1 ( x, x (cid:48) ) .
we already now that ˙Σ l ( x, x ) = 1 .
Moreover, we have Σ l ( x, x ) =
Σ l − 1 ( x, x ) + σ 2 w / 2Σ l − 1 ( x, x )
=
(1 + σ 2 w / 2) l − 1 σ 2 w d ||
d
|| .
Depending on the value of σ w , the behaviour of K l ( x, x (cid:48) ) changes.
However, from Hayou et al.
(2019), we have that ˙Σ l ( x, x (cid:48) )
= 1 − βl − 1 + O ( l − 2 ) .
So by scaling with α L 2 L and using Lemma 3, we conclude on the convergence rate of O ( l − 1 ) .
We can take the supremum as the result of the boundedness of c 1 ( x, x (cid:48) ) .
C I MPACT OF THE I NITIALIZATION THE OUTPUT FUNCTION
In this section, we show how an initialization on the EOC impacts on the output function of the neural network.
More precisely, we show that it leads to a larger range compared to an initialization in the Ordered/chaotic phase.
Let us start by a simple Lemma that compares the expectations of some smooth mapping with respect to two different Gaussian vectors.
Lemma 5.
Let X =
( X i ) 1 ≤ i ≤ n , Y =
( Y i ) 1 ≤ i ≤ N be two centered Gaussian vectors in R n .
Let g ∈ D
( R n , R ) .
Then we have
[Equation]
The result of Lemma 5 is valid when the second derivatives of g exist only in the distribution sense (e.g. Dirac mass).
Proof.
We deﬁne the function G on R by
[Equation]
we have that
[Equation]
Moreover, it is easy to see that for any random vector Z in R n we have E
[ Z i g ( Z )]
= (cid:80) nj =1 cov [Equation] X i , X j ) E [ [Equation]
j ( Z )] , this yields
[Equation]
We conclude by integrating G (cid:48) ( t ) between 0 and 1 .
Now let D = { ( x i , z i ) : 1 ≤ i ≤ N } be the datapoints.
Using the same notations as in the previous chapter, let y L ( x i ) denotes on the neurons of l th layer (the neurons are iid).
Assume c
1 x i ,x j ≥ 0 for all i, j ∈ [1 , N ] (this is almost always the case, but in general, we can re-scale the input data to satisfy this assumption).
We have the following result
Lemma 6.
Let φ be a non ReLU-like activation function and ( σ b , σ w ) / ∈ EOC .
Then there exists ( σ b,EOC , σ w,EOC ) ∈ EOC such that for any function g ∈ D 2 such that for all i, j ∈ [ | 1 , N | ] , ∂g ∂x i ∂x j ≥ 0 , there exist β > 0 , ζ N > 0 such that
[Equation]
Proof.
Let ( σ b , σ w ) / ∈ EOC and q be the corresponding limiting variance.
From the previous chapter, it is easy to see that there exists σ 0 > 0 such that σ 20 + E [ φ ( √ qZ ) 2 ] E [ φ (cid:48) ( √ qZ ) 2 ] = q .
Let ( σ b,EOC , σ w,EOC ) =
( σ 0 , 1 / (cid:112) E [ φ (cid:48) ( √ qZ ) 2 ]) ∈ EOC .
There exists a constant κ > 0 (independent of N ) such that for all i, j ∈ [ | 1 , N | ] , E [ y Leoc ( X i )
y Leoc
X j )
] ≤ E [ y Lord ( X i )
y Lord ( X j )] − κL − β where β
= 1 for a smooth activation functions in S and β = 2 for ReLU-like activation functions (see Hayou et al.
Let λ N = inf
i,j inf u ∈ [0 , 1] E [ ∂g ∂x i ∂x j ( √ 1 − uX + √ uY )] .
Using Lemma 5, we have
[Equation]
As a simple application, using the function g ( X ) = (cid:81)
Ni =
1 x i ≤ t i , we have the following
Lemma 7.
Let t 1 , t 2 , ..., t N ∈ R .
Then there exist β > 0 , ζ N > 0 such that
[Equation]
β as a Corollary, we have a generalized form of Slepian’s Lemma for y Leoc ( X ) and y Lord ( X ) .
Corollary 1 (Max Range and Min Range) .
[Equation]
and
[Equation]
where β
= 1 for activation functions of type S and β = 2 for ReLU-like activation functions.
D T RAINING WITH SGD INSTEAD OF GD
In this section, we extend the results of the NTK to the case of SGD.
We use an approximation of the
SGD dynamics by a diffusion process.
We assume implicitly the existence of the triplet (Ω , P , F ) where Ω is the probability space, P is a probability measure on Ω , and F is the natural ﬁltration of the Brownian motion.
Under boundedness conditions, when using SGD, the gradient update can be seen as a GD with a Gaussian noise (Hu et al., 2018; Li et al., 2017).
More precisely, let S
= o ( N ) be the batchsize.
The SGD update is given by
[Equation]
[Equation] where L ( S ) = 1 S (cid:80) Si =1 (cid:96) ( f (˜ x i , θ ) , ˜ y i ) where (˜ x i , ˜ y i ) 1 ≤ i ≤ S is a randomly selected batch of size S .
Then for all θ ∇
[Equation]
N where Z i ( S ) = 1 if observation i belongs to the batch (˜ x j , ˜ y j ) , j ≤ S and equals 0 otherwise and
E 0 ( θ ) =
E 0 ∇ θ (cid:96) ( f ( X 1 , θ ) , Y 1 ) .
We have
[Equation]
N
[Equation]
N . So that if S = o ( N ) and if
[Equation]
where Cov ( · ) denotes the covariance matrix under P 0 .
Then
[Equation]
where Z S ( θ ) converges in distribution (as S goes to inﬁnity) to a Gaussian random vector with covariance matrix Σ( θ )
=
Cov ( ∇ θ (cid:96) ( f ( X 1 , θ ) , Y 1 )) and we have, neglecting the term o P {+
+} ( S − 1 / 2 ) ,
[Equation]
We can in particular bound the difference between equation 16 and the continuous time SDE approximation (see also Hu et al.
(2018) and Li et al.
[Equation]
SGD updates can therefore be seen as a discretization of the previous SDE with time step ∆ =
η , and where Σ( θ t ) 12 is the square-root matrix of Σ( θ t ) =
Cov ( ∇ θ (cid:96) ( f ( X 1 , θ t ) , Y 1 )) and ( W t ) t ≥ 0 a standard Brownian motion.
Since the dynamics of θ t are described by an SDE, the dynamics of f t can also be described by an
SDE which can be obtained from Itô’s lemma.
Proposition 4.
Under the dynamics of the SDE equation 17, the vector f t ( X ) is the solution of the following SDE
[Equation]
where Γ t ( X ) is the concatenated vector of (Γ t ( x )
=
(Tr(Σ( θ t[-) 12 ∇ 
-]
[Equation]
12 ∇ 2 f i ( x, θ t )Σ( θ t ) 12 )) 1 ≤ i ≤ o ) x ∈X and ∇ {+
+} f i ( x, θ ) is the Hessian of f i ( i th component of f ) with respect to θ .
Proof.
Since θ t is a diffusion process, we can use Itô’s lemma to deduce how the randomness propagates to f t .
We denote by f t,i the i th coordinate of f t , i.e., for an input x , f t ( x ) =
( f t,i ( x ))
1 ≤ i ≤ k .
Let i ∈ 1 , ..., k , we have
[Equation]
where ∇ 2 f i ( x, θ t ) is the hessian of f i with respect to θ .
Aggregating these equations with respect to i yields
[Equation]
where Γ t ( x ) =
(Tr(Σ( θ t ) 12 ∇
f
i ( x, θ t )Σ( θ t ) 12 )) 1 ≤ i ≤ k .
Therefore, the dynamics of the vector f t ( X ) is given by
[Equation]
where Γ t ( X ) is the concatenated vector of (Γ t ( x )) x ∈X .
With the quadratic loss (cid:96) ( z, y ) =
12 || z − y ||
2 , the SDE equation 18 is equivalent to
[Equation]
This is an Ornstein-Uhlenbeck process (mean-reverting process) with time dependent parameters.
The additional term Γ t is due to the randomness of the mini-batch, it can be seen as a regularization term and could partly explain why SGD gives better generalization errors compared to GD (Kubo et al.
(2019), Lei et al.
Dynamics of f t for wide FeedForward neural networks :
In the case of a fully connected feedforward neural network (FFNN hereafter) of depth L and widths n 1 , n 2 , ..., n L , Jacot et al.
(2018) proved that, with GD, the kernel K Lθ t converges to a kernel K L that depends only on L (number of layers) for all t < T when n 1 , n 2 , ..., n l → ∞ , where T is an upper bound on the training time, under the technical assumption (cid:82) T 0 ||∇ z (cid:96) ( f t ( X , Y ))
||
2 dt < ∞ almost surely with respect to the initialization.
For SGD, we assume that the convergence result of the NTK holds true as well, this is illustrated empirically in ﬁgure 3 but we leave the theoretical proof for future work.
With this approximation, the dynamics of f t ( X ) for wide networks is given by
[Equation]
where ˆ K L = K L ( X , X ) and M t = Y − ηN 2 S
( ˆ K L ) − 1 Γ t ( X ) .
This is an Ornstein–Uhlenbeck process whose closed-form expression is given by
[Equation]
where A t ( X ) =
{+− η 2 S (cid:82) t 
 e − t − sN ˆ K L Γ s ( X ) ds + (cid:112)+}
ηS (cid:82) {+t 
+}
e − t − sN ˆ K L ∇ θ f
( X , θ s )Σ( θ s ) 12 dW s ; see supplementary material for the proof.
So for any (test) input x ∈ R d , we have
[Equation]
[Equation]
where R t ( x ) = (cid:112) ηS (cid:82) t
[ K L ( x, X )( ˆ K L ) − 1 ( e − tN ˆ K L − I ) ∇ θ f
( X , θ s )
+ ∇ θ f
( x, θ s )]
Σ( θ s ) 12 dW s and Z t ( x ) =
η 2 S (cid:104)(cid:82) t
Γ s ( x ) ds + (cid:82) t 0 K ( x, X )( ˆ K L ) − 1 ( I − e − ( t − s ) N ˆ K L )
Γ s ( X ) ds (cid:105) .
[Figure]
Proof.
Using the approximation of the NTK by K L as n 1 , n 2 , ..., n L → ∞ , the dynamics of f t ( X ) for wide networks are given by
[Equation]
To solve it, we use the change of variable A t =
e tN ˆ K L f t ( X ) .
Using Ito’s lemma, we have
[Equation]
By integrating, we conclude that
[Equation]
we conclude for f t ( X ) using the fact that 1 N (cid:82) t 0 ˆ K L e − ( t − s [Equation] N [Equation]
η 2 S
[Equation]
t 0 e − ( t − s ) N ˆ K L Γ s ds .
Recall that for any input x ∈ R d ,
[Equation]
To prove the expression of f t ( x ) for general x ∈ R d , we substitute f t ( X ) by its value in the SDE of f t ( x ) and integrate.