Title: MEAN-FIELD BEHAVIOUR OF NEURAL TANGENT KERNEL FOR DEEP NEURAL NETWORKS

Abstract: 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.

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 (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 flattened 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 classification 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 : 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.

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 .

For a given θ, the empirical loss is given by L

(θ) = 1 N N i=1 (f (x i , θ), y i ).

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 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

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, y i ) |z=f (xi,θt) , i ≤ N . As a result, the output function f t (x) satisfies the following ordinary differential equation

EQUATION

The Neural Tangent Kernel (NTK) K L θ is defined as the o × o dimensional kernel satisfying: for all x, x ∈ R d ,

EQUATION

We also define K L θt (X , X ) as the oN × oN matrix defined blockwise by

K L θt (X , X ) =     K L θt (x 1 , x 1 ) K L θt (x 1 , x 2 ) • • • K L θt (x 1 , x N ) K L θt (x 2 , x 1 ) • • • • • • K L θt (x 2 , x N ) • • • • • • • • • • • • K L θt (x N , x 1 ) K L θt (x N , x 2 ) • • • K L θt (x N , x N )    

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 (f t (X ), Y)dt.

Infinite 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

T 0 ||∇ z (f t (X , Y))|| 2 dt <

∞ almost surely with respect to the initialization weights. The infinite width limit of the training dynamics is given by

EQUATION

We note hereafter KL = K L (X , X ). As an example, with the quadratic loss (z, y) = 1 2 ||z -y|| 2 , equation 6 is equivalent to

EQUATION

which is a simple linear model that has a closed-form solution given by

EQUATION

Note that in order for f t (x) to be defined, KL 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 ) = Y.

Then with dynamics defined by equation 8, ||f t (X ) -Y|| converges to 0 as t → ∞ if and only if KL is non-singular. Moreover, if KL is singular, there exists a constant C > 0 such that for all t > 0, ||f t (X ) -Y|| ≥ C

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 ) = δ for all x = x , 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 :

1. Under some constraints, the NTK K L (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

4. IMPACT OF THE INITIALIZATION AND THE ACTIVATION FUNCTION ON THE NEURAL TANGENT KERNEL:

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 l ij iid ∼ N (0, σ 2 w n l-1 ) and b l i 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. (2019) . 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 ) for any two inputs x and x is driven by (σ b , σ w ); the authors define the EOC as the set of parameters

(σ b , σ w ) such that σ 2 w E[φ ( q(σ b , σ w )Z) 2 ] = 1 where Z ∼ N (0, 1). Similarly the Ordered, resp. Chaotic, phase is defined by σ 2 w E[φ ( q(σ b , σ w )Z) 2 ] < 1, resp. σ 2 w E[φ ( 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 first define two classes of activation functions. Definition 1. Let φ : R → R be a measurable function. Then 1. φ is said to be ReLU-like if there exist λ, β ∈ R such that φ(x) = λx for x > 0 and φ(x) = βx for x ≤ 0.

2. φ 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 infinitely differentiable functions g 1 , g 2 , ..., g n such that

φ (2) = n i=1 1 Ai 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

sup x,x ∈R d |K L (x, x ) -λ| ≤ e -γL → L→∞ 0

As a result, as L goes to infinity,

K L converges to a constant kernel K ∞ (x, x ) = λ for all x, x ∈ 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. KL 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

K L θt (X , X ) = ∇ θ f (X , θ t )∇ θ f (X , θ t ) T = L l=1 ∇ θ l f (X , θ t )∇ θ l f (X , θ t ) T

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 confirms the findings 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.

1. If φ is ReLU-like, then for all x ∈ R d , K L (x,x) L = σ 2 w ||x|| 2 d + K 0 (x,x) L

. Moreover, there exist A, λ ∈ (0, 1) such that

sup x =x ∈R d K L (x, x ) L -λ σ 2 w d ||x||||x || ≤ A L , K ∞ (x, x ) = σ 2 w ||x|| x d (1-(1-λ)1 x =x )

2. If φ is in S, then, there exists q > 0 such that K L (x,x)

L = q + K 0 (x,x) L → q. Moreover, there exist B, C, λ ∈ (0, 1) such that B log(L) L ≤ sup x =x ∈R d K L (x, x ) L -qλ ≤ C log(L) L , K ∞ (x, x ) = q(1-(1-λ)1 x =x )

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 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 influential 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 L res be the corresponding NTK. Then for all x ∈ R d ,

K L res (x,x) α L ×2 L = ||x|| 2 d + O(γ L )

and there exists λ ∈ (0, 1) such that

sup x =x ∈R d K L res (x, x ) α L × 2 L - ||x|| × ||x || d λ = O(L -1 ),

where α l and γ l are given by 

• if σ w < √ 2, then α L = 1 and γ L = ( 1+σ 2 w /2 2 ) L • if σ w = √ 2, then α L = L and γ L = L -1

• if σ w > √ 2, then α L = ( 1+σ 2 w /2 2

) L and γ L = (

1+σ 2 w /2 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 L res (x, x )/α L 2 L where x = x 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.

5. EXPERIMENTS:

In this section, we illustrate empirically the theoretical results obtained in the previous sections. We first illustrate the results of Propositions 1, 2 and 3. Then, we confirm the impact of the EOC and Activation function on the overall performance of the model (FFNN), on MNIST and CIFAR10 datasets.

5.1. CONVERGENCE RATE OF K L AS L GOES TO INFINITY:

Propositions 1, 2 and 3 give theoretical convergence rates for quantities of the form 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.2. IMPACT 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 satisfied. 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. (2019) 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 

6. CONCLUSION 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 infinite 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-field 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 infinity, making training impossible. An initialization on the EOC leads to an invertible ANTK (and NTK) even for an infinite 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 findings showed that the infinite 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 finite width neural networks changes with time and might even be random (Chizat & Bach (2018) , Ghorbani et al. (2019) , Huang & Yau (2019) , Arora et al. (2019) ). 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.

We provide in Section A and Section B the proof of the theoretical results presented in the main document. Section C provides additional theoretical results while Section ?? presents additional experimental results.

A APPENDIX: PROOFS OF SECTION 3: NEURAL NETWORKS AND NEURAL TANGENT KERNEL

Lemma 1 (Trainability of the Neural Network and Invertibility of the NTK). Assume f 0 (X ) = Y.

Then with dynamics (8), ||f t (X ) -Y|| converges to 0 as t → ∞ if and only if KL is non-singular. Moreover, if K L is singular, there exists a constant C > 0 such that for all t > 0,

||f t (X ) -Y|| ≥ C Proof. Assume f 0 (X ) = Y.

Let KL = 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 KL t = Q T e -1 N Dt Q = Q T Diag(e -d i N t ) 1≤i≤oN Q where (d i ) 1≤i≤oN are the eigenvalues. We also have ||f t (X ) -Y|| = ||e -1 N KL t (f 0 (X ) -Y)||. Therefore, the equivalence holds true. Moreover, assume KL 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

||f t (X ) -Y|| = ||Z t || ≥ |(Z t )j| = |(Z 0 )j|

Lemma 2 (Discretization Error for Full-Batch Gradient Descent). Assume ∇ θ L is C-lipschitz, then there exists C > 0 that depends only on C and T such that

sup k∈[0,T /η] ||θ t k -θk || ≤ ηC

Proof. For t ∈ [0, T ], we define the stepwise constant system θt = θ t/η . Let t ∈ [0, T ], we have

θt = θ 0 -η t/η -1 k=0 ∇ θ L( θk ) = θ 0 - t 0 ∇ θ L( θs )ds + η t/η * t/η-1 ∇ θ L( θ * s )ds Therefore, ||θ t -θt || ≤ t 0 ||∇ θ L(θ s ) -∇ θ L( θs )||ds + η t/η * t/η-1 ||∇ θ L( θ * s )||ds ≤ C t 0 ||θ s -θs ||ds + η(t/η - * t/η)||∇ θ L( θ * t/η )|| + η||∇ θ L( θ * t/η-1 )||

Moreover, for any k ∈ [0, * T /η], we have

|| θk -θ 0 || ≤ (1 + ηC)|| θk-1 -θ 0 || ≤ (1 + ηC) T /η || θ1 -θ 0 || ≤ e CT || θ1 -θ 0 ||

where we have used log(1 + ηC) ≤ ηC. Using this result, there exists a constant C depending on T and C such that (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.

η(t/η - * t/η)||∇ θ L( θ * t/η )|| + η||∇ θ L( θ * t/η-1 )|| ≤ η(2||∇ θ L( θ0 )|| + C|| θ * t/η -θ 0 || + C|| θ * t/η-1 -θ 0 ||) ≤ η C Now

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 l ij iid ∼ N (0, σ 2 w n l-1 ) and b l i 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 defined by the set of parameters (σ b , σ w ) such that σ 2 w E[φ ( q(σ b , σ w )Z) 2 ] = 1 where Z ∼ N (0, 1). Similarly the Ordered, resp. Chaotic, phase is defined by Hayou et al. (2019) for more details). For two inputs x, x ∈ R d , define Σ l (x, x ) = E[y l (x)y l (x )] and let c l (x, x ) be the corresponding correlation. Let f be the correlation function defined implicitly by c l+1 = f (c l ). In the limit of infinitely wide networks, we have the following results (Hayou et al. (2019) ) :

σ 2 w E[φ ( q(σ b , σ w )Z) 2 ] < 1, resp. σ 2 w E[φ ( q(σ b , σ w )Z) 2 ] > 1 (see

• Σ l (x, x ) = σ 2 b + σ 2 w E z∼N (0,Σ l-1 ) [φ(z(x))φ(z(x ))].

• 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 ∈R d |c l (x, x ) -1| ≤ e -γl .

• On the chaotic phase, there exist γ > 0 and c < 1 such that sup

x =x ∈R d |c l (x, x ) -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 f (x) = x→1- x + 2 √ 2 3π (1 -x) 3/2 + O((1 -x) 5/2 )

• In general, we have

f (x) = σ 2 b + σ 2 w E[φ( √ qZ 1 )φ( √ qZ(x))] 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 (1) = 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-1 E[φ (j) (Z 1 )φ (j) (Z(x))]

• From Jacot et al. ( 2018), we have that

K l (x, x ) = K l-1 (x, x ) Σl (x, x ) + Σ l (x, x ).

where the definition of Σl (x, x ) is given in Proposition 1 below.

Definition 1. Let φ : R → R be a measurable function. Then 1. φ is said to be ReLU-like if there exist λ, β ∈ R such that φ(x) = λx for x > 0 and φ(x) = βx for x ≤ 0.

2. φ 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 infinitely differentiable functions g 1 , g 2 , ..., g n such that

φ (2) = n i=1 1 Ai 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

a l ≤ a 0 α l + k l-1 j=0 α j e -β(l-j) ≤ a 0 α l + k l 2 e -βl/2 + k l 2 α l/2

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

sup x,x ∈R d |K L (x, x ) -λ| ≤ e -γL → L→∞ 0

Proof. From Jacot et al. ( 2018), we have that

K l (x, x ) = K l-1 (x, x ) Σl (x, x ) + Σ l (x, x )

where

Σ 1 (x, x ) = σ 2 b + σ 2 w d x T x and Σ l (x, x ) = σ 2 b + σ 2 w E f ∼N (0,Σ l-1 ) [φ(f (x))φ(f (x ))] and Σl (x, x ) = E f ∼N (0,Σ l-1 ) [φ (f (x))φ (f (x ))].

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

sup x,x ∈R d |Σ l (x, x ) -k| ≤ e -γl

and sup

x,x ∈R d Σl (x, x ) ≤ α.

Therefore we have for any l ≥ l 0 and x, x ∈ R d

K l (x, x ) ≤ αK l-1 (x, x ) + k + e -γl .

Letting r l = K l (x, x ) -k 1-α , we have r l ≤ αr l-1 .

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 infinity. We first prove two preliminary lemmas that will be useful for the proof of the next proposition.

Under review as a conference paper at ICLR 2020 Lemma 3. Let (a l ), (b l ), (λ l ) be three sequences of real numbers such that

a l = a l-1 λ l + b l λ l = 1 - α l + O(l -1-β ) b l = q + O(l -1 )

where α ∈ N * , β, q ∈ R + and α > β -1.

Then,

a l l = q 1 + α + O(l -min(1,β) )

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

r l = r l-1 (1 - 1 l )(1 - α l + O(l -1-β )) + q l + O(l -2 ) = r l-1 (1 - 1 + α l ) + q l + O(l -1-ζ ).

Letting x l = r l -q 1+α , there is exist a constant M > 0 such that

x l ≤ x l-1 (1 - 1 + α l ) + M l 1+ζ .

Hence, we have

x l ≤ x 0 l k=1 (1 - 1 + α k ) + M l k=1 l j=k+1 (1 - 1 + α j ) 1 k 1+ζ .

By taking the logarithm of the first term in the right hand side and using the fact that l k=1

1 k ∼ log(l), we have l k=1 (1 - 1 + α k ) ∼ l -1-α

For the second part, observe that l j=k+1

(1 -

1 + α j ) = (l -α -1)! l! k! (k -α -1)! and k! (k -α -1)! 1 k 1+ζ ∼ k→∞ k α-ζ so that, l k=1 k! (k -α -1)! 1 k 1+ζ ∼ l k=1 k α-ζ ∼ l 1 t α-ζ dt ∼ 1 α -ζ + 1 l α-ζ+1 therefore, l k=1 l j=k+1 (1 - 1 + α j ) 1 k 1+ζ = (l -α -1)! l! l k=1 k! (k -α -1)! 1 k 1+ζ ∼ 1 α -ζ + 1 l -ζ

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

a l = a l-1 λ l + b l λ l = 1 - α l + κ log(l) l 2 + O(l -1-β ) b l = q + O(l -1 )

where α ∈ N * , β, q ∈ R + and α > β -1, β ≥ 1. Then, there exists A, B > 0 such that

A log(l) l ≤ | a l l - q 1 + α | ≤ B log(l) l

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

r l = r l-1 (1 - 1 l )(1 - α l + κ log(l) l 2 + O(l -1-β )) + q l + O(l -2 ) = r l-1 (1 - 1 + α l ) + r l-1 κ log(l) l 2 + q l + O(l -2 ) 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 0

x l-1 (1 - 1 + α l ) + κ 1 log(l) l 2 - M l 2 ≤ x l ≤ x l-1 (1 - 1 + α l ) + κ 2 log(l) l 2 + M l 2

Similarly to the proof of Lemma 3, it follows that

x l ≤ x l0 l k=l0 (1 - 1 + α k ) + l k=l0 l j=k+1 (1 - 1 + α j ) κ 2 log(k) + M k 2 and x l ≥ x 0 l k=0 (1 - 1 + α k ) + l k=l0 l j=k+1 (1 - 1 + α j ) κ 1 log(k) -M k 2

Recall that we have

l k=1 (1 - 1 + α k ) ∼ l -1-α and l j=k+1 (1 - 1 + α j ) = (l -α -1)! l! k! (k -α -1)! so that k! (k -α -1)! κ 1 log(k) -M k 2 ∼ k→∞ log(k)k α-1

Therefore, we obtain

l k=1 k! (k -α -1)! κ 1 log(k) -M k 2 ∼ l k=1 log(k)k α-1 ∼ 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

l k=1 k! (k -α -1)! κ 2 log(k) + M k 2 ∼ C 2 l α log(α)

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.

1. If φ is ReLU-like, then for all x ∈ R d , K L (x,x) L = σ 2 w ||x|| 2 d + K 0 (x,x) L

. Moreover, there exist A, λ ∈ (0, 1) such that

sup x =x ∈R d K L (x, x ) L -λ σ 2 w d ||x||||x || ≤ A L 2. If φ is in S, then, there exist q > 0 such that K L (x,x) L = q + K 0 (x,x) L → q. Moreover, there exist B, C, λ ∈ (0, 1) such that B log(L) L ≤ sup x =x ∈R d K L (x, x ) L -qλ ≤ C log(L) L

Proof. We use some results from Hayou et al. (2019) in this proof.

Let x, x ∈ R d and c l x,x = Σ ( x,x ) √ Σ ( x,x)Σ ( x ,x )

. Let γ l := 1 -c l x,x and f be the correlation function defined by the recursive equation c l+1 = f (x l ). From the preliminary results, we know that Σ l (x, x) = σ 2 w d ||x|| 2 and that K l (x, x ) = K l-1 (x, x ) Σl (x, x ) + Σ l (x, x ). 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(γ 5/2 l ) so that γ -1/2 l+1 = γ -1/2 l (1 -sγ 1/2 l + O(γ 3/2 l )) -1/2 = γ -1/2 l (1 + s 2 γ 1/2 l + O(γ 3/2 l )) = γ -1/2 l + s 2 + O(γ l ).

Thus, as l goes to infinity γ

-1/2 l+1 -γ -1/2 l

∼ s 2 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 l x,x = 1 -9π 2 2l 2 + O(l -3

). we also have

f (x) = 1 π arcsin(x) + 1 2 = 1 - √ 2 π (1 -x) 1/2 + O((1 -x) 5/2 ).

Thus, it follows that

f (c l x,x ) = 1 - 3 l + O(l -2 )

Moreover, q l x,x = q + O(l -2 ) where q is the limiting variance of y l . Using Lemma 3, we conclude that K l (x,x )

l = 1 4 σ 2 w d ||x||||x || + O(l -1

). Since c x,x is bounded, this result is uniform in x, x . Therefore, we can take the supremum over x, x ∈ R d . 1. 2. We prove the result when φ (2) (x) = 1 x<0 g 1 (x) + 1 x≥0 g 2 (x). The generalization to the whole class is straightforward. Let f be the correlation function. We first 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 f (x) = σ 2 w qE[φ ( √ qZ 1 )φ ( √ qU 2 (x))] = σ 2 w qE[φ ( √ qZ 1 )1 U2(x)<0 g 1 ( √ qU 2 (x))] + σ 2 w qE[φ ( √ qZ 1 )1 U2(x)>0 g 2 ( √ qU 2 (x))]. Let G(x) = E[φ ( √ qZ 1 )1 U2(x)<0 g 1 ( √ qU 2 (x))] then G (x) = E[φ ( √ qZ 1 )(Z 1 - x √ 1 -x 2 Z 2 )δ U2(x)=0 1 √ 1 -x 2 g 1 ( √ qU 2 (x))] + E[φ ( √ qZ 1 )1 U2(x)<0 √ q(Z 1 - x √ 1 -x 2 Z 2 )g 1 ( √ qU 2 (x))].

It is easy to see that

G (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 . 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 defined 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:

f (1 -t 2 ) = 1 -t 2 f (1) + t 4 2 f (1) + t 6 6 f (3) (1) + O(t 8 ).

Therefore,

f (x) = 1 + (x -1)f (1) + (x -1) 2 2 f (1) + (1 -x) 3 6 f (3) (1) + O((x -1) 4 ).

Letting λ l := 1 -c l , there exist α, β > 0 such that

λ l+1 = λ l -αλ 2 l -βλ 3 l + O(λ 4 l ) therefore, λ -1 l+1 = λ -1 l (1 -αλ l -βλ 2 + O(λ 3 l )) -1 = λ -1 l (1 + αλ l + βλ 2 l + O(λ 3 l )) = λ -1 l + α + βλ l + O(λ 2 l ).

By summing (divergent series), we have that λ -1 l ∼ l βq . Therefore,

λ -1 l+1 -λ -1 l -α = βα -1 l -1 + O(l -2

) By summing a second time, we obtain

λ -1 l = αl + βα -1 log(l) + O(1) so that λ l = α -1 l -1 -α -1 β log(l) l 2 + O(l -2 ). Using the fact that f (x) = 1 + (x -1)f (1) + O((x -1) 2 ), we have f (c l x,x ) = 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 , 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 L res be the corresponding NTK. For all x ∈ R d ,

K L res (x,x) L×2 L = σ 2 w ||x|| 2 d + O(L -1

) and there exists λ ∈ (0, 1) such that

sup x =x ∈R d K L res (x, x ) L × 2 L -σ 2 w ||x|| × ||x || d λ = O(L -1 )

.

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 L res satisfies the following recursive equation

K l res (x, x ) = K l-1 res (x, x )( Σl (x, x ) + 1) + Σ l (x, x )

To see this, with FeedForward Neural Network, we have ( Jacot et al. (2018) )

∂y l+1 i ∂θ 1:l (x)( ∂y l+1 i ∂θ 1:l (x )) t = n l j,j w l+1 ij w l+1 ij φ (y l j ) ∂y l i ∂θ 1:l (x)( ∂y l i ∂θ 1:l (x)) t + I W

whereas for Residual Networks, we have an extra term

∂y l+1 i ∂θ 1:l (x)( ∂y l+1 i ∂θ 1:l (x )) t = ∂y l i ∂θ 1:l (x)( ∂y l i ∂θ 1:l (x )) t + n l j,j w l+1 ij w l+1 ij φ (y l j ) ∂y l i ∂θ 1:l (x)( ∂y l i ∂θ 1:l (x)) t + I W

where I W , I 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 ) in the formula of K l+1 (x, x ).

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 ) changes. However, from Hayou et al. (2019) , we have that Σl (x, x ) = 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 ).

C IMPACT OF THE INITIALIZATION 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 2 (R n , R). Then we have E[g(X)] -E[g(Y )] = 1 2 1 0 1≤i,j,≤n (E[X i X j ] -E[Y i Y j ])E[ ∂g ∂x i ∂x j ( √ 1 -uX + √ uY )]du (14)

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 define the function G on R by

G(t) = E[g(tX + 1 -t 2 Y )]

we have that

G (t) = n i=1 E[(X i - t √ 1 -t 2 Y i ) ∂g ∂x i (tX + 1 -t 2 Y )]

Moreover, it is easy to see that for any random vector Z in R n we have

E[Z i g(Z)] = n j=1 cov(X i , X j )E[ ∂g ∂xj (Z)], this yields G (t) = t n i,j=1 (E[X i X j ] -E[Y i Y j ])E[ ∂g ∂x i ∂x j (tX + 1 -t 2 Y )]

We conclude by integrating G (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 xi,xj ≥ 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 ∂xi∂xj ≥ 0, there exist β > 0, ζ N > 0 such that E[g(y L eoc (X))] ≤ E[g(y L ord (X))] - ζ N L β

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 σ

2 0 + E[φ( √ qZ) 2 ] E[φ ( √ qZ) 2 ] = q. Let (σ b,EOC , σ w,EOC ) = (σ 0 , 1/ E[φ ( √ qZ) 2 ]) ∈ EOC. There exists a constant κ > 0 (independent of N ) such that for all i, j ∈ [|1, N |], E[y L eoc (X i )y L eoc (X j )] ≤ E[y L ord (X i )y L ord (X j )] -κL -β

where β = 1 for a smooth activation functions in S and β = 2 for ReLU-like activation functions (see Hayou et al. (2019) 

). Let λ N = inf i,j inf u∈[0,1] E[ ∂g ∂xi∂xj ( √ 1 -uX + √ uY )]. Using Lemma 5, we have E[y L eoc (X i )y L eoc (X j )] ≤ E[y L eoc (X i )y L ord (X j )] - 1 2 κN 2 λ N L -β

As a simple application, using the function g(X) = N i=1 1 xi≤ti , we have the following Lemma 7. Let t 1 , t 2 , ..., t N ∈ R. Then there exist β > 0, ζ N > 0 such that

P(y L eoc (x 1 ) ≤ t 1 , ..., y L eoc (x N ) ≤ t N ) ≤ P(y L ord (x 1 ) ≤ t 1 , ..., y L ord (x N ) ≤ t N ) - ζ N L β

as a Corollary, we have a generalized form of Slepian's Lemma for y L eoc (X) and y L ord (X). Corollary 1 (Max Range and Min Range).

P(max:

i y L eoc (x i ) ≥ t) ≥ P(max i y L ord (x i ) ≥ t) + ζ N L β and P(min i y L eoc (x i ) ≤ t) ≥ P(max i y L ord (x i ) ≤ t) + ζ N L β

where β = 1 for activation functions of type S and β = 2 for ReLU-like activation functions.

D TRAINING 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 filtration 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

where L (S) = 1 S S i=1 (f (x i , θ), ỹi ) where (x i , ỹi ) 1≤i≤S is a randomly selected batch of size S. Then for all θ

∇ θ L (S) (θ) -∇ θ L(θ) = i Z i (S) S (∇ θ (f θ (x i ), y i ) -E 0 (θ)) - N i=1 (∇ θ (f θ (x i ), y i ) -E 0 (θ)) N

where Z i (S) = 1 if observation i belongs to the batch (x j , ỹj ), j ≤ S and equals 0 otherwise and

E 0 (θ) = E 0 ∇ θ (f (X 1 , θ), Y 1 ). We have tr Cov N i=1 (∇ θ (f θ (x i ), y i ) -E 0 (θ)) N = p l=1 Var (∂ (f θ (X 1 ), Y 1 )/∂θ l ) N . So that if S = o(N ) and if tr (Cov (∇ θ (f (X 1 , θ), Y 1 ))) = o(S)

where Cov(•) denotes the covariance matrix under P 0 . Then

∇ θ L (S) (θ) -∇ θ L(θ) = Z S (θ) √ S + o P0 (S -1/2 )

where Z S (θ) converges in distribution (as S goes to infinity) to a Gaussian random vector with covariance matrix Σ(θ) = Cov (∇ θ (f (X 1 , θ), Y 1 )) and we have, neglecting the term o P0 (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. (2017) )

EQUATION

SGD updates can therefore be seen as a discretization of the previous SDE with time step ∆ = η, and where Σ(θ t )

EQUATION

)) 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

df t (X ) = [- 1 N K L θt (X , X )∇ z (f t (X ), Y ) + 1 2 η S Γ t (X )]dt + η S ∇ θ f (X , θ t )Σ(θ t ) 1 2 dW t (18)

where Γ t (X ) is the concatenated vector of (Γ t (x) = (Tr(Σ(θ t )

1 2 ∇ 2 f i (x, θ t )Σ(θ t ) 1 2 )) 1≤i≤o ) x∈X and ∇ 2 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

2 )]dt

+ η S ∇ θ f i (x, θ t )Σ(θ t ) 1 2 dW t

where ∇ 2 f i (x, θ t ) is the hessian of f i with respect to θ. Aggregating these equations with respect to i yields

df t (x) = [- 1 N ∇ θ f t (x)∇ θ f t (X )∇ z (f t (X ), Y ) + 1 2 η S Γ t (x)]dt + η S ∇ θ f (x, θ t )Σ(θ t ) 1 2 dW t

where Γ t (x) = (Tr(Σ(θ t )

1 2 ∇ 2 f i (x, θ t )Σ(θ t ) 1 2 )) 1≤i≤k .

Therefore, the dynamics of the vector f t (X ) is given by

df t (X ) = [- 1 N K θ t (X , X )∇ z (f t (X ), Y ) + 1 2 η S Γ t (X )]dt + η S ∇ θ f (X , θ t )Σ(θ t ) 1 2 dW t

where Γ t (X ) is the concatenated vector of (Γ t (x)) x∈X .

With the quadratic loss (z, y) = 1 2 ||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. (2018) ).

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 T 0 ||∇ z (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 figure 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

df t (X ) = - 1 N KL (f t (X ) -M t )dt + η S ∇ θ f (X , θ t )Σ(θ t ) 1 2 dW t ,

where KL = K L (X , X ) and M t = Y -ηN 2S ( KL ) -1 Γ t (X ). This is an Ornstein-Uhlenbeck process whose closed-form expression is given by f t (X ) = e 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

df t (X ) = - 1 N KL (f t (X ) -M t )dt + η S ∇ θ f (X , θ t )Σ(θ t ) 1 2 dW t ,

To solve it, we use the change of variable A t = e t N KL f t (X ). Using Ito's lemma, we have Recall that for any input x ∈ R d ,

df t (x) = [- 1 N K L (x, X )(f t (X ) -Y ) + 1 2 η S Γ t (x)]dt + η S ∇ θ f (x, θ t )Σ(θ t ) 1 2 dW t

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. 

E ADDITIONAL EXPERIMENTS : IMPACT OF THE PARAMETRIZATION:



:



Figures and Tables:

Figure 1: Convergence rates for different initializations and architectures. (a) Edge of Chaos. (b) Ordered phase. (c) Adding residual connections.

Figure 2: Test accuracy for different Activation Functions and (width, depth) on MNIST

f (X , θ s )Σ(θ s ) 1 2 dW s ; see supplementary material for the proof. So for any (test) input x ∈ R d , we havef t (x) = f 0 (x) + K L (x, X )( KL ) -1 (I -e -t-s N KL )(Y -f 0 (X )) + Z t (x) + R t (x),(21)whereR t (x) = η S t 0 [K L (x, X )( KL ) -1 (e -t N KL -I)∇ θ f (X , θ s ) + ∇ θ f (x, θ s )]Σ(θ s ) 1 2 dW s and Z t (x) = η 2S t 0 Γ s (x)ds + t 0 K(x, X )( KL ) -1 (I -e -(t-s) N KL )Γ s (X )ds .

Figure 3: Ratio K L θt /K L for three randomly selected pairs from MNIST dataset as a function of width for three training times t = 0, t = 100 and t = 1000 (training time is measured by SGD updates)

Figure4shows the impact of different parametrizations on the training dynamics. It shows that both standard parametrization and NTK parametrization have similar behaviour on different initialization regimes.

Figure 4: Impact of NTK parametrization on the Training

Test accuracy for a FFNN with width 300 and depth 300 for different activation functions on MNIST and CIFAR10. We show test accuracy after 10 epochs and 100 epochs LReLU 0.01 (EOC) 48.10 ± 3.31 84.71 ± 3.39 22.62 ± 1.15 29.44 ± 4.14 LReLU 0.02 (EOC) 49.09 ± 3.58 84.3. ± 3.98 18.62 ± 4.56 30.78 ± 6.33 LReLU 0.03 (EOC) 50.94 ± 4.48 85.49 ± 2.71 21.19 ± 6.53 34.54 ± 2.32

Table content: <table><tr><td/><td colspan="2">MNIST</td><td/><td>CIFAR10</td></tr><tr><td>Activation</td><td>Epoch 10</td><td>Epoch 100</td><td>Epoch 10</td><td>Epoch 100</td></tr><tr><td>ReLU (EOC)</td><td colspan="4">46.53 ± 12.01 82.11 ± 4.51 20.38 ± 1.85 35.88 ± 0.6</td></tr><tr><td>PReLU</td><td>51.94 ± 5.51</td><td colspan="3">87.49 ± 1.58 22.95 ± 3.57 36.13 ± 3.83</td></tr><tr><td>ELU (EOC)</td><td>91.63 ± 2.21</td><td colspan="3">96.07 ± 0.13 33.81 ± 1.55 46.14 ± 1.49</td></tr><tr><td>Tanh (EOC)</td><td>91.16 ± 1.21</td><td colspan="3">95.75 ± 0.27 32.37 ± 1.88 42.40 ± 1.13</td></tr><tr><td>Softplus</td><td>10.11 ± 0.09</td><td colspan="3">10.13 ± 0.18 11.13 ± 0.15 11.09 ± 0.36</td></tr><tr><td>Sigmoid</td><td>9.85 ± 0.11</td><td>9.87 ± 0.10</td><td colspan="2">10.65 ± 0.25 10.33 ± 0.17</td></tr></table>

PROOFS OF SECTION 4: IMPACT OF THE INITIALIZATION AND THE ACTIVATION FUNCTION ON THE NEURAL TANGENT KERNELWe first recall the results obtained inLee et al. (2018),Schoenholz et al. (2017) andHayou et al.

Table content: <table><tr><td>we have</td><td/></tr><tr><td>t</td><td/></tr><tr><td>||θ t -θt || ≤ C</td><td>||θ s -θs ||ds + η C,</td></tr><tr><td>0</td><td/></tr><tr><td>so we can conclude using Gronwall's lemma.</td><td/></tr></table>

Peer review (iclr2020):
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. Detailed Comments: 1. At the paragraph after Lemma 1, the authors claimed that f_t(x) is entirely given by f_0(x), which means a generalization error of order O(1). This is a little inaccurate I think, as NTK can only characterize the training process and does not directly indicate generalization. Also, as t to infinity, the coefficient of initialization f_0(x) can be arbitrary small and thus the f_t(x) is entirely decided by initialization is not accurate as well. To analyze generalization with NTK, we need a little more, see [1]. This is not some core issue, but I think the authors should make the description more accurate. 2. The authors should explain what richer limiting NTK means at the end of the Sec 3. It is unclear how the convergence rate of the NTK related to the richer limiting NTK. 3. A typo in the first paragraph in Sec 4, ouerselvs to ourselves. 4. In definition 1, the second-order derivative of phi is defined via the indicator function 1_{A_i}? It’s better to use another notation like mathbf{1} or mathbb{1} to make it more clear. 5. I think for clarity the authors should give a formal definition of ANTK, or if it is unnecessary, better directly use the original NTK. Also, it is better to show the invertibility of the NTK is equivalent to ANTK more formally before Proposition 2. This conclusion is not so straightforward at the first view. Meanwhile, I think the invertibility of ANTK can just indicate a bad condition number of NTK? I don’t think it directly related to the invertibility of NTK. 6. There are several typos in the appendix. Please go through the appendix and fix them. 7. In the proof of Lemma 3 and Lemma 4, I don’t get why |a_l| < l + |a_0|. It may depends on the property of q? For sufficiently large q, if a_0=0, a_1 = q + O(1) which is not necessarily smaller than l=1? But to get a_l/l bounded, it is not necessary to use this. Also, I think it’s better not omit the constant in the dynamic system of Lemma 3. 8. Better use the Stirling’s approximation rather than k to infinity in the calculation of k!/(k-alpha-1)! 9. The lemma number of Appendix C is in a chaos. And I don’t quite get the idea behinds Appendix C. Is it relevant to the contents in the main text? 10. I also feel the content in Appendix D is unnecessary for this paper. Overall, this paper is interesting and gives a unified perspective on the recently developed NTK and Edge of Chaos initialization. It also sheds light on the impact of different activation function on NTK by generalizing the results of [2]. I feel this paper can help the communities have more understanding on the properties of deep neural networks. However, this paper can be better if the authors can polish their paper and reorganize the appendix part. [1] Arora, Sanjeev, et al. *Fine-Grained Analysis of Optimization and Generalization for Overparameterized Two-Layer Neural Networks.* International Conference on Machine Learning. 2019. [2] Hayou, S., Doucet, A. & Rousseau, J.. (2019). On the Impact of the Activation function on Deep Neural Networks Training. Proceedings of the 36th International Conference on Machine Learning, in PMLR 97:2672-2680