summaries: 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. The paper studies the limiting behavior of neural tangent kernels when the depth grows to infinity. They show that the obtained limit kernels are trivial (a constant) unless one uses 'edge of caos' initialization, in which case they are close to the identity. The authors compare the convergence for different activations, showing a slower convergence (hence better propagation) for some piecewise smooth activations compared to ReLU. For residual networks, the 'edge of caos' behavior is claimed to be in place regardless of the initialization.
abstractText: 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.
references.authors: S. Arora S.S. Du W. Hu Z. Li R. Salakhutdinov R. Wang
references.title: On exact computation with an infinitely wide neural net
references.authors: A. Bietti J. Mairal
references.title: On the inductive bias of neural tangent kernels
references.venue: arXiv Preprint arXiv:1905.12173,
references.authors: L. Chizat F. Bach
references.title: A note on lazy training in supervised differentiable programming
references.venue: arXiv preprint arXiv:1812.07956,
references.authors: D.A. Clevert T. Unterthiner S. Hochreiter
references.title: Fast and accurate deep network learning by exponential linear units (elus)
references.venue: International Conference on Learning Representations,
references.authors: S.S. Du J.D. Lee H. Li L. Wang X. Zhai
references.title: Gradient descent finds global minima of deep neural networks
references.venue: arXiv preprint arXiv:1811.03804,
references.authors: S.S. Du J.D. Lee Y. Tian B. Poczos A Singh
references.title: Gradient descent learns one-hidden-layer CNN: Don’t be afraid of spurious local minima
references.authors: S.S. Du X. Zhai B. Poczos A. Singh
references.title: Gradient descent provably optimizes over-parameterized neural networks
references.authors: B. Ghorbani S. Mei T. Misiakiewicz A. Montanari
references.title: Linearized two-layers neural networks in high dimension
references.venue: arXiv preprint arXiv:1904.12191,
references.authors: X. Glorot Y. Bengio
references.title: Understanding the difficulty of training deep feedforward neural networks
references.venue: International Conference on Artificial Intelligence and Statistics,
references.authors: S. Hayou A. Doucet J. Rousseau
references.title: On the impact of the activation function on deep neural networks training
references.authors: K. He X. Zhang S. Ren J. Sun
references.title: Delving deep into rectifiers: Surpassing human-level performance on imagenet classification
references.authors: W. Hu C. Junchi Li L. Li J Liu
references.title: On the diffusion approximation of nonconvex stochastic gradient descent
references.venue: arXiv preprint arXiv:1705.07562,
references.authors: J. Huang H.T Yau
references.title: Dynamics of deep neural networks and neural tangent hierarchy
references.venue: arXiv preprint arXiv:1909.08156,
references.authors: A. Jacot F. Gabriel C. Hongler
references.title: Neural tangent kernel: Convergence and generalization in neural networks
references.venue: 32nd Conference on Neural Information Processing Systems,
references.authors: R. Karakida S. Akaho S. Amari
references.title: Universal statistics of Fisher information in deep neural networks: Mean field approach
references.venue: arXiv preprint arXiv:1806.01316,
references.authors: M. Kubo R. Banno H. Manabe M. Minoji
references.title: Implicit regularization in over-parameterized neural networks
references.venue: arXiv preprint arXiv:1903.01997,
references.authors: J. Lee Y. Bahri R. Novak S.S. Schoenholz J. Pennington J. Sohl-Dickstein
references.title: Deep neural networks as Gaussian processes
references.venue: 6th International Conference on Learning Representations,
references.authors: J. Lee L. Xiao S. Schoenholz Y. Bahri J. Sohl-Dickstein J. Pennington
references.title: Wide neural networks of any depth evolve as linear models under gradient descent
references.authors: D. Lei Z. Sun Y. Xiao W.Y. Wang
references.title: Implicit regularization of stochastic gradient descent in natural language processing: Observations and implications
references.venue: arXiv preprint arXiv:1811.00659,
references.authors: Q. Li C. Tai W E
references.title: Stochastic modified equations and adaptive stochastic gradient algorithms
references.venue: arXiv preprint arXiv:1511.06251,
references.authors: Q. Nguyen M. Hein
references.title: Optimization landscape and expressivity of deep CNNs
references.authors: S.S. Schoenholz J. Gilmer S. Ganguli J. Sohl-Dickstein
references.title: Deep information propagation
references.venue: 5th International Conference on Learning Representations,
references.authors: G. Yang
references.title: Scaling limits of wide neural networks with weight sharing: Gaussian process behavior, gradient independence, and neural tangent kernel derivation
references.authors: G. Yang S. Schoenholz
references.title: Mean field residual networks: On the edge of chaos
references.venue: Advances in Neural Information Processing Systems,
references.authors: C. Zhang S. Bengio M. Hardt B. Recht O. Vinyals
references.title: Understanding deep learning requires rethinking generalization
references.venue: arXiv preprint arXiv:1611.03530,
references.authors: D. Zou Y. Cao D. Zhou Q. Gu
references.title: Stochastic gradient descent optimizes over-parameterized deep ReLU networks
references.venue: arXiv preprint arXiv:1811.08888,
references.authors: Hayou
references.title: For two inputs x, x′ ∈ R, define Σ(x, x′) = E[yl(x)yl(x′)] and let c(x, x′) be the corresponding correlation. Let f be the correlation function defined implicitly by c = f(c). In the limit of infinitely wide networks, we have the following results (Hayou et al
references.authors: Jacot
references.title: 2018), we have that K(x, x′) = Kl−1(x, x′)Σ̇l(x, x′) + Σ(x, x′). where the definition of Σ̇(x, x′) is given in Proposition 1 below. Definition 1
references.venue: Let φ : R→ R
references.authors: d ||d
references.title: Depending on the value of σw, the behaviour of K (x, x′) changes
references.venue: Hayou et al. (2019),
references.authors: Li
references.title: We can in particular bound the difference between equation 16 and the continuous time SDE approximation (see also Hu et al
references.authors: nL Jacot
references.title: 2018) proved that, with GD, the kernel K
sections.heading: 1 INTRODUCTION
sections.text: 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.
sections.heading: 2 MOTIVATION AND RELATED WORK
sections.text: 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
sections.heading: 3 NEURAL NETWORKS AND NEURAL TANGENT KERNEL
sections.heading: 3.1 SETUP AND NOTATIONS
sections.text: Consider a neural network model consisting of L layers (yl)1≤l≤L, with yl : Rnl−1 → Rnl , n0 = d and let θ = (θl)1≤l≤L be the flattened vector of weights and bias indexed by the layer’s index and p be the dimension of θ. Recall that θl has dimension nl + 1. The output f of the neural network is given by some transformation s : RnL → Ro of the last layer yL(x); o being the dimension of the output (e.g. number of classes for a classification problem). For any input x ∈ Rd, we thus have f(x, θ) = s(yL(x)) ∈ Ro. As we train the model, θ changes with time t and we denote by θt the value of θ at time t and ft(x) = f(x, θt) = (fj(x, θt), j ≤ o). Let D = (xi, yi)1≤i≤N be the data set and let X = (xi)1≤i≤N , Y = (yj)1≤j≤N be the matrices of input and output respectively, with dimension d×N and o×N . For any function g : Rd×o → Rk, k ≥ 1, we denote by g(X ,Y) the matrix (g(xi, yi))1≤i≤N of dimension k ×N . Jacot et al. (2018) studied the behaviour of the output of the neural network as a function of the training time t when the network is trained using a gradient descent algorithm. Lee et al. (2019) built on this result to linearize the training dynamics. We recall hereafter some of these results. For a given θ, the empirical loss is given by L(θ) = 1N ∑N i=1 `(f(xi, θ), yi). The full batch GD algorithm is given by θ̂t+1 = θ̂t − η∇θL(θ̂t) (1)
  where η > 0 is the learning rate. Let T > 0 be the training time and Ns = T/η be the number of steps of the discrete GD equation 1. The continuous time system equivalent to equation 1 with step ∆t = η is given by
  dθt = −∇θL(θt)dt (2) This differs from the result by Lee et al. (2019) since we use a discretization step of ∆t = η. It is well known that this discretization scheme leads to an error of order O(η) (see Appendix). As in Lee et al. (2019), Equation (2) can be re-written as
  dθt = − 1
  N ∇θf(X , θt)T∇z`(f(X , θt),Y)dt
  where∇θf(X , θt) is a matrix of dimension oN × p and∇z`(f(X , θt),Y) is the flattened vector of dimension oN constructed from the concatenation of the vectors ∇z`(z, yi)|z=f(xi,θt), i ≤ N . As a result, the output function ft(x) satisfies the following ordinary differential equation
  dft(x) = ∇θf(x, θt)dθt = − 1
  N ∇θf(x, θt)∇θf(X , θt)T∇z`(ft(X ),Y)dt ∈ Ro (3)
  The Neural Tangent Kernel (NTK) KLθ is defined as the o× o dimensional kernel satisfying: for all x, x′ ∈ Rd,
  KLθt(x, x ′) = ∇θf(x, θt)∇θf(x′, θt)T ∈ Ro×o
  = L∑ l=1 ∇θlf(x, θt)∇θlf(x′, θt)T . (4)
  We also define KLθt(X ,X ) as the oN × oN matrix defined blockwise by
  KLθt(X ,X ) =  KLθt(x1, x1) K L θt (x1, x2) · · · KLθt(x1, xN ) KLθt(x2, x1) · · · · · · K L θt (x2, xN ) · · · · · · · · · · · ·
  KLθt(xN , x1) K L θt (xN , x2) · · · KLθt(xN , xN )
  
  By applying equation 3 to the vector X , one obtains
  dft(X ) = − 1
  N KLθt(X ,X )∇z`(ft(X ),Y)dt, (5)
  meaning that for all j ≤ N dft(xj) = − 1NK L θt (xj ,X )∇z`(ft(X ),Y)dt.
  Infinite width dynamics : In the case of a fully connected feedforward neural network (FFNN) of depth L and widths n1, n2, ..., nL, Jacot et al. (2018) proved that, with GD, the kernel KLθt converges to a kernel KL which depends only on L (number of layers) for all t < T when n1, n2, ..., nL → ∞, where T is an upper bound on the training time, under the technical assumption ∫ T 0 ||∇z`(ft(X ,Y))||2dt <∞ almost surely with respect to the initialization weights. The infinite width limit of the training dynamics is given by
  dft(X ) = − 1
  N KL(X ,X )∇z`(ft(X ),Y)dt, (6)
  We note hereafter K̂L = KL(X ,X ). As an example, with the quadratic loss `(z, y) = 12 ||z − y|| 2, equation 6 is equivalent to
  dft(X ) = − 1
  N K̂L(ft(X )− Y)dt, (7)
  which is a simple linear model that has a closed-form solution given by
  ft(X ) = e− 1 N K̂ Ltf0(X ) + (I − e− 1 N K̂ Lt)Y. (8) For general input x ∈ Rd, we then have
  ft(x) = f0(x) +K L(x,X )KL(X ,X )−1(I − e− 1N K̂ Lt)(Y − f0(X )) (9) Note that in order for ft(x) to be defined, K̂L must be invertible. Indeed, it turns out that training with dynamics 6 is only possible if the NTK is invertible. We shed light on this behaviour in the following Lemma. Lemma 1 (Trainability of the Neural Network and Invertibility of the NTK). Assume f0(X ) 6= Y . Then with dynamics defined by equation 8, ||ft(X )− Y|| converges to 0 as t→∞ if and only if K̂L is non-singular. Moreover, if K̂L is singular, there exists a constant C > 0 such that for all t > 0,
  ||ft(X )− Y|| ≥ C
  Lemma 1 shows that an invertible NTK is crucial for trainability. Since KLθt = K L is constant w.r.t to training time, it is completely determined at the initialization step. It is therefore intuitive to study the impact of the initialization on the NTK, particularly as the number of layers L grows (Deep Neural Networks), which is our focus in this paper. Another interesting aspect is the impact of the NTK on the generalization error of the neural network model. To see this, if the NTK is constant for example (i.e. there exists a constant δ such as KL(x, x′) = δ for all x 6= x′, this example is useful in the next section), then the second part of ft(x) in equation 9 is constant w.r.t x. Therefore, the generalization function ft(x) of the model 9 is entirely given by its value at time zero f0(x), which means that the generalisation error Ex,y[||ft(x)− y||] remains of order O(1). In the next section, we show that the initialization and the activation function have major impact on the invertibility and ’expressivity’ of NTK. More precisely, we show that :
  1. Under some constraints, the NTK KL (or a scaled version of the NTK) converges to a limiting NTK K∞ as L goes to infinity (otherwise it diverges)
  2. A special initialization known as the Edge of Chaos (EOC) leads to an invertible K∞ which makes it useful for training DNNs
  3. The EOC initialization gives a sub-exponential rate for this convergence (w.r.t L), which means for the same depth L, the EOC gives ’richer’ limiting NTK, and therefore leading to better generalization properties
  4. The smoothness of the activation can further slow this convergence, leading to ’richer’ limiting NTK (the convergence to the limiting trivial kernel is slower)
  5. Adding Residual connections leads to sub-exponential convergence rate for the NTK (w.r.t to L) and we no longer need the Edge of Chaos
sections.heading: 4 IMPACT OF THE INITIALIZATION AND THE ACTIVATION FUNCTION ON THE NEURAL TANGENT KERNEL
sections.text: 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 nL is straightforward.
  Consider a FFNN of depth L, widths (nl)1≤l≤L, weights wl and bias bl. For some input x ∈ Rd, the forward propagation is given by
  y1i (x) = d∑ j=1 w1ijxj + b 1 i , y l i(x) = nl−1∑ j=1 wlijφ(y l−1 j (x)) + b l i, for l ≥ 2, (10)
  where φ is the activation function.
  We initialize the model with wlij iid∼ N (0, σ 2 w nl−1 ) and bli iid∼ N (0, σ2b ), where N (µ, σ2) denotes the normal distribution of mean µ and variance σ2. For some x, we denote by ql(x) the variance of yl(x). The convergence of ql(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 ql(x) converges to a point q(σb, σw) > 0 independent of x as l→∞. Also the asymptotic behaviour of the correlations between yl(x) and yl(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 σ2wE[φ′( √ q(σb, σw)Z)
  2] = 1 where Z ∼ N (0, 1). Similarly the Ordered, resp. Chaotic, phase is defined by σ2wE[φ′( √ q(σb, σw)Z)
  2] < 1, resp. σ2wE[φ′( √ 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 (Ai)1≤i≤n of R and infinitely differentiable functions g1, g2, ..., gn such that φ(2) = ∑n i=1 1Aigi, 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′∈Rd
  |KL(x, x′)− λ| ≤ e−γL →L→∞ 0
  As a result, as L goes to infinity,KL converges to a constant kernelK∞(x, x′) = λ for all x, x′ ∈ Rd. The training is then impossible. Indeed, we have KL(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
  KLθt(X ,X ) = ∇θf(X , θt)∇θf(X , θt) T = L∑ l=1 ∇θlf(X , θt)∇θlf(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 ∇θlf(X , θt)∇θlf(X , θt)T are of the same order. Hence, it is more convenient to study the average NTK (ANTK hereafter) given by KL/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 ∈ Rd, K L(x,x) L = σ2w||x|| 2 d + K0(x,x)
  L . Moreover, there exist A, λ ∈ (0, 1) such that
  sup x 6=x′∈Rd ∣∣KL(x, x′) L −λσ 2 w d ||x||||x′|| ∣∣ ≤ A L , K∞(x, x ′) = σ2w||x||‖x′‖ d (1−(1−λ)1x 6=x′)
  2. If φ is in S , then, there exists q > 0 such that K L(x,x) L = q + K0(x,x) L → q. Moreover, there
  exist B,C, λ ∈ (0, 1) such that
  B log(L)
  L ≤ sup x 6=x′∈Rd ∣∣KL(x, x′) L −qλ ∣∣ ≤ C log(L) L , K∞(x, x ′) = q(1−(1−λ)1x 6=x′)
  Since 0 < λ < 1, on the EOC there exists a matrix J invertible such thatKL(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 andO(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)
  yli(x) = y l−1 i (x) + nl−1∑ j=1 wlijφ(y l−1 j (x)) + b l i, for l ≥ 2. (11)
  with initialization parameters σb = 0 and σw > 0 and φ is the ReLU activation function. Let KLres be the corresponding NTK. Then for all x ∈ Rd, K L res(x,x) αL×2L = ||x||2 d +O(γL) and there exists λ ∈ (0, 1) such that
  sup x 6=x′∈Rd ∣∣KLres(x, x′) αL × 2L − ||x|| × ||x ′|| d λ ∣∣ = O(L−1),
  where αl and γl are given by
  • if σw < √ 2, then αL = 1 and γL = ( 1+σ2w/2 2 ) L • if σw = √ 2, then αL = L and γL = L−1
  • if σw > √ 2, then αL = ( 1+σ2w/2 2 ) L and γL = ( 1+σ2w/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 KLres(x, x ′)/αL2 L where x 6= x′ converges to a limiting kernel similar to K∞ with a rateO(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.
sections.heading: 5 EXPERIMENTS
sections.text: 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.
sections.heading: 5.1 CONVERGENCE RATE OF KL AS L GOES TO INFINITY
sections.text: Propositions 1, 2 and 3 give theoretical convergence rates for quantities of the form ∣∣KL α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.
sections.heading: 5.2 IMPACT OF THE INITIALIZATION AND SMOOTHNESS OF THE ACTIVATION ON THE OVERALL PERFORMANCE
sections.text: 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
  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. (2015) 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 classifier with 10 classes.
sections.heading: 6 CONCLUSION AND LIMITATIONS
sections.text: 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.
sections.heading: A APPENDIX: PROOFS OF SECTION 3: NEURAL NETWORKS AND NEURAL TANGENT KERNEL
sections.text: Lemma 1 (Trainability of the Neural Network and Invertibility of the NTK). Assume f0(X ) 6= Y . Then with dynamics (8), ||ft(X )− Y|| converges to 0 as t→∞ if and only if K̂L is non-singular. Moreover, if KL is singular, there exists a constant C > 0 such that for all t > 0,
  ||ft(X )− Y|| ≥ C
  Proof. Assume f0(X ) 6= Y . Let K̂L = QTDQ 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̂ Lt = QT e− 1 NDtQ = QTDiag(e− di N t)1≤i≤oNQ where (di)1≤i≤oN are the eigenvalues. We also have ||ft(X ) − Y|| = ||e− 1 N K̂
  Lt(f0(X ) − Y)||. Therefore, the equivalence holds true.
  Moreover, assume K̂L in singular. Let Zt = Q(ft(X ) − Y)QT . We have that Zt = e− 1 NDtZ0. Since D has at least one zero diagonal value, then there exists j ∈ {1, 2, ..., oN} such that for all t, (Zt)j = (Z0)j , and we have
  ||ft(X )− Y|| = ||Zt|| ≥ |(Zt)j| = |(Z0)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/η]
  ||θtk − θ̂k|| ≤ ηC ′
  Proof. For t ∈ [0, T ], we define the stepwise constant system θ̃t = θ̂bt/ηc. Let t ∈ [0, T ], we have
  θ̃t = θ0 − η bt/ηc−1∑ k=0 ∇θL(θ̂k)
  = θ0 − ∫ t
  0 ∇θL(θ̃s)ds+ η ∫ t/η b∗ct/η−1 ∇θL(θ̂b∗cs)ds
  Therefore, ||θt − θ̃t|| ≤ ∫ t
  0 ||∇θL(θs)−∇θL(θ̃s)||ds+ η ∫ t/η b∗ct/η−1 ||∇θL(θ̂b∗cs)||ds
  ≤ C ∫ t
  0
  ||θs − θ̃s||ds+ η(t/η − b∗ct/η)||∇θL(θ̂b∗ct/η)||+ η||∇θL(θ̂b∗ct/η−1)||
  Moreover, for any k ∈ [0, b∗cT/η], we have
  ||θ̂k − θ0|| ≤ (1 + ηC)||θ̂k−1 − θ0|| ≤ (1 + ηC)T/η||θ̂1 − θ0|| ≤ eCT ||θ̂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
  η(t/η − b∗ct/η)||∇θL(θ̂b∗ct/η)||+ η||∇θL(θ̂b∗ct/η−1)|| ≤ η(2||∇θL(θ̂0)||+ C||θ̂b∗ct/η − θ0||+ C||θ̂b∗ct/η−1 − θ0||) ≤ ηC̃
  Now we have
  ||θt − θ̃t|| ≤ C ∫ t
  0
  ||θs − θ̃s||ds+ ηC̃,
  so we can conclude using Gronwall’s lemma.
sections.heading: B PROOFS OF SECTION 4: IMPACT OF THE INITIALIZATION AND THE ACTIVATION FUNCTION ON THE NEURAL TANGENT KERNEL
sections.text: We first 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 (nl)1≤l≤L, weights wl and bias bl. For some input x ∈ Rd, the forward propagation is given by
  y1i (x) = d∑ j=1 w1ijxj + b 1 i , y l i(x) = nl−1∑ j=1 wlijφ(y l−1 j (x)) + b l i, for l ≥ 2, (12)
  where φ is the activation function.
  We initialize the model with wlij iid∼ N (0, σ 2 w nl−1 ) and bli iid∼ N (0, σ2b ), where N (µ, σ2) denotes the normal distribution of mean µ and variance σ2. For some x, we denote by ql(x) the variance of yl(x). In general, ql(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 σ2wE[φ′( √ q(σb, σw)Z)
  2] = 1 where Z ∼ N (0, 1). Similarly the Ordered, resp. Chaotic, phase is defined by σ2wE[φ′( √ q(σb, σw)Z)
  2] < 1, resp. σ2wE[φ′( √ q(σb, σw)Z)
  2] > 1 (see Hayou et al. (2019) for more details). For two inputs x, x′ ∈ Rd, define Σl(x, x′) = E[yl(x)yl(x′)] and let cl(x, x′) be the corresponding correlation. Let f be the correlation function defined implicitly by cl+1 = f(cl). In the limit of infinitely wide networks, we have the following results (Hayou et al. (2019)) :
  • Σl(x, x′) = σ2b + σ2wEz∼N (0,Σl−1)[φ(z(x))φ(z(x′))]. • There exist q, λ > 0 such that, for all supx∈Rd |Σl(x, x)− q| ≤ e−λl. • On the Ordered phase, there exists γ > 0 such that supx,x′∈Rd |cl(x, x′)− 1| ≤ e−γl.
  • On the chaotic phase, there exist γ > 0 and c < 1 such that supx 6=x′∈Rd |cl(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) = σ2b + σ 2 wE[φ(
  √ qZ1)φ( √ qZ(x))]
  q
  where Z(x) = xZ1 + √
  1− x2Z2 and Z1, Z2 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) = σ2wq j−1E[φ(j)(Z1)φ(j)(Z(x))]
  • From Jacot et al. (2018), we have that
  Kl(x, x′) = Kl−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 (Ai)1≤i≤n of R and infinitely differentiable functions g1, g2, ..., gn such that φ(2) = ∑n i=1 1Aigi, 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 (al) be a sequence of non-negative real numbers such that ∀l ≥ 0, al+1 ≤ αal+ke−βl, where α ∈ (0, 1) and k, β > 0. Then there exists γ > 0 such that ∀l ≥ 0, l ≤ e−γl.
  Proof. Using the inequality on al, we can easily see that
  al ≤ a0αl + k l−1∑ j=0 αje−β(l−j)
  ≤ a0α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′∈Rd
  |KL(x, x′)− λ| ≤ e−γL →L→∞ 0
  Proof. From Jacot et al. (2018), we have that
  Kl(x, x′) = Kl−1(x, x′)Σ̇l(x, x′) + Σl(x, x′)
  where Σ1(x, x′) = σ2b + σ2w d x Tx′ and Σl(x, x′) = σ2b + σ 2 wEf∼N (0,Σl−1)[φ(f(x))φ(f(x′))] and Σ̇l(x, x′) = Ef∼N (0,Σl−1)[φ′(f(x))φ′(f(x′))]. In the ordered/chaotic phase, Hayou et al. (2019) showed that there exist k, γ, l0 > 0 and α ∈ (0, 1) such that for all l ≥ l0 we have
  sup x,x′∈Rd
  |Σl(x, x′)− k| ≤ e−γl
  and sup
  x,x′∈Rd Σ̇l(x, x′) ≤ α.
  Therefore we have for any l ≥ l0 and x, x′ ∈ Rd
  Kl(x, x′) ≤ αKl−1(x, x′) + k + e−γl.
  Letting rl = Kl(x, x′)− k1−α , we have
  rl ≤ αrl−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.
  Lemma 3. Let (al), (bl), (λl) be three sequences of real numbers such that al = al−1λl + bl
  λl = 1− α
  l +O(l−1−β)
  bl = q +O(l −1)
  where α ∈ N∗, β, q ∈ R+ and α > β − 1. Then,
  al l = q 1 + α +O(l−min(1,β))
  Proof. It is easy to see that there exists a constant G > 0 |al| ≤ G× l + |a0| for all l ≥ 0, therefore (al/l) is bounded. Now let ζ = min(1, β) and rl = all . We have
  rl = rl−1(1− 1 l )(1− α l +O(l−1−β)) + q l +O(l−2)
  = rl−1(1− 1 + α
  l ) +
  q l +O(l−1−ζ).
  Letting xl = ∣∣rl − q1+α ∣∣, there is exist a constant M > 0 such that
  xl ≤ xl−1(1− 1 + α
  l ) +
  M
  l1+ζ .
  Hence, we have
  xl ≤ x0 l∏
  k=1
  (1− 1 + α k ) +M l∑ k=1 l∏ j=k+1 (1− 1 + α j ) 1 k1+ζ .
  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 k1+ζ ∼k→∞ kα−ζ
  so that, l∑
  k=1
  k! (k − α− 1)! 1 k1+ζ ∼ l∑ k=1 kα−ζ
  ∼ ∫ l
  1
  tα−ζdt
  ∼ 1 α− ζ + 1 lα−ζ+1
  therefore, l∑
  k=1 l∏ j=k+1 (1− 1 + α j ) 1 k1+ζ = (l − α− 1)! l! l∑ k=1
  k! (k − α− 1)! 1 k1+ζ
  ∼ 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 (al), (bl), (λl) be three sequences of real numbers such that
  al = al−1λl + bl
  λl = 1− α
  l + κ
  log(l)
  l2 +O(l−1−β)
  bl = q +O(l −1)
  where α ∈ N∗, β, q ∈ R+ and α > β − 1, β ≥ 1. Then, there exists A,B > 0 such that
  A log(l) l ≤ |al l − q 1 + α | ≤ B log(l) l
  Proof. It is easy to see that there exists a constant G > 0 |al| ≤ G× l + |a0| for all l ≥ 0, therefore (al/l) is bounded. Let rl = all . We have
  rl = rl−1(1− 1 l )(1− α l + κ log(l) l2 +O(l−1−β)) + q l +O(l−2)
  = rl−1(1− 1 + α
  l ) + rl−1κ
  log(l) l2 + q l +O(l−2)
  Let xl = rl − q1+α . It is clear that λl = 1 − α/l + O(l −3/2). Therefore, using Lemma 3 with β = 1/2, we have rl → q1+α . Thus, there exists κ1, κ2,M, l0 > 0 such that for all l ≥ l0
  xl−1(1− 1 + α
  l ) + κ1
  log(l) l2 − M l2 ≤ xl ≤ xl−1(1− 1 + α l ) + κ2 log(l) l2 + M l2
  Similarly to the proof of Lemma 3, it follows that
  xl ≤ xl0 l∏
  k=l0
  (1− 1 + α k ) + l∑ k=l0 l∏ j=k+1 (1− 1 + α j ) κ2 log(k) +M k2
  and
  xl ≥ x0 l∏
  k=0
  (1− 1 + α k ) + l∑ k=l0 l∏ j=k+1 (1− 1 + α j ) κ1 log(k)−M k2
  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 k2 ∼k→∞ log(k)kα−1
  Therefore, we obtain l∑
  k=1
  k! (k − α− 1)! κ1 log(k)−M k2 ∼ l∑ k=1 log(k)kα−1
  ∼ ∫ l
  1
  log(t)tα−1dt
  ∼ C1lα log(α)
  where C1 > 0 is a constant. Similarly, there exists a constant C2 > 0 such that l∑
  k=1
  k! (k − α− 1)! κ2 log(k) +M k2 ∼ C2lα 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 ∈ Rd, K L(x,x) L = σ2w||x|| 2 d + K0(x,x)
  L . Moreover, there exist A, λ ∈ (0, 1) such that
  sup x 6=x′∈Rd ∣∣KL(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 + K0(x,x) L → q. Moreover, there
  exist B,C, λ ∈ (0, 1) such that B log(L)
  L ≤ sup x6=x′∈Rd ∣∣KL(x, x′) L − qλ ∣∣ ≤ C log(L) L
  Proof. We use some results from Hayou et al. (2019) in this proof.
  Let x, x′ ∈ Rd and clx,x′ = Σ(x,x′)√
  Σ(x,x)Σ(x′,x′) . Let γl := 1 − clx,x′ and f be the correlation
  function defined by the recursive equation cl+1 = f(xl). From the preliminary results, we know that Σl(x, x) = σ 2 w
  d ||x|| 2 and that Kl(x, x′) = Kl−1(x, x′)Σ̇l(x, x′) + Σl(x, x′).
  This concludes the proof for KL(x, x). We denote s = 2 √
  2 3π . From Hayou et al. (2019), we
  have on the EOC γl+1 = γl − sγ3/2l +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/2l+1 − γ −1/2 l = s 2 +O(γl) = s 2 +O(l −2), we have γ−1/2l = s 2 l+O(1). Therefore, clx,x′ = 1− 9π 2 2l2 +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 ′(clx,x′) = 1− 3
  l +O(l−2)
  Moreover, qlx,x′ = q +O(l −2) where q is the limiting variance of yl. Using Lemma 3, we conclude that K l(x,x′) l = 1 4 σ2w d ||x||||x
  ′|| + O(l−1). Since cx,x′ is bounded, this result is uniform in x, x′. Therefore, we can take the supremum over x, x′ ∈ Rd.
  1.2 We prove the result when φ(2)(x) = 1x<0g1(x) + 1x≥0g2(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−x2)(k−2)/2 gk(x) where gk ∈ C ∞.
  We have f ′′(x) = σ2wqE[φ′′( √ qZ1)φ ′′( √ qU2(x))]
  = σ2wqE[φ′′( √ qZ1)1U2(x)<0g1( √ qU2(x))] + σ 2 wqE[φ′′( √ qZ1)1U2(x)>0g2( √ qU2(x))].
  Let G(x) = E[φ′′(√qZ1)1U2(x)<0g1( √ qU2(x))] then
  G′(x) = E[φ′′( √ qZ1)(Z1 − x√ 1− x2 Z2)δU2(x)=0 1√ 1− x2 g1( √ qU2(x))]
  + E[φ′′( √ qZ1)1U2(x)<0 √ q(Z1 − x√ 1− x2 Z2)g ′ 1( √ qU2(x))].
  It is easy to see that G′(x) = 1√ 1−x2G1(x) where G1 ∈ C 1. A similar analysis can be applied to the second term of f ′′. We conclude that there exists g3 ∈ C∞ such that f (3)(x) = 1√
  1−x2 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 − t2 to obtain a Taylor expansion near 1. Simple algebra shows that the function t→ f(1− t2) has a Taylor expansion near 0:
  f(1− t2) = 1− t2f ′(1) + t 4
  2 f ′′(1) +
  t6 6 f (3)(1) +O(t8).
  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− cl, there exist α, β > 0 such that
  λl+1 = λl − αλ2l − βλ3l +O(λ4l )
  therefore,
  λ−1l+1 = λ −1 l (1− αλl − βλ 2 +O(λ3l )) −1
  = λ−1l (1 + αλl + βλ 2 l +O(λ 3 l )) = λ−1l + α+ βλl +O(λ 2 l ).
  By summing (divergent series), we have that λ−1l ∼ l βq . Therefore,
  λ−1l+1 − λ −1 l − α = βα −1l−1 +O(l−2)
  By summing a second time, we obtain
  λ−1l = αl + βα −1 log(l) +O(1)
  so that λl = α−1l−1 − α−1β log(l)l2 +O(l −2).
  Using the fact that f ′(x) = 1 + (x − 1)f ′′(1) + O((x − 1)2), we have f ′(clx,x′) = 1− 2l +κ log(l) l2 +O(l
  −2). We can now conclude using Lemma 4. Using again the agrument of the boundedness of c1x,x′ , we can take the supremum.
  Proposition 3. Consider the following network architecture (FFNN with residual connections)
  yli(x) = y l−1 i (x) + nl−1∑ j=1 wlijφ(y l−1 j (x)) + b l i, for l ≥ 2. (13)
  with initialization parameters σb = 0 and σw > 0. Let KLres be the corresponding NTK. For all x ∈ Rd, K L res(x,x) L×2L = σ 2 w ||x||2 d +O(L −1) and there exists λ ∈ (0, 1) such that
  sup x 6=x′∈Rd ∣∣KLres(x, x′) L× 2L − σ2w ||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 KLres satisfies the following recursive equation
  Klres(x, x ′) = Kl−1res (x, x ′)(Σ̇l(x, x′) + 1) + Σl(x, x′)
  To see this, with FeedForward Neural Network, we have ( Jacot et al. (2018))
  ∂yl+1i ∂θ1:l (x)( ∂yl+1i ∂θ1:l (x′))t = nl∑ j,j′ wl+1ij w l+1 ij′ φ ′(ylj) ∂yli ∂θ1:l (x)( ∂yli ∂θ1:l (x))t + IW
  whereas for Residual Networks, we have an extra term
  ∂yl+1i ∂θ1:l (x)( ∂yl+1i ∂θ1:l (x′))t = ∂yli ∂θ1:l (x)( ∂yli ∂θ1:l (x′))t + nl∑ j,j′ wl+1ij w l+1 ij′ φ ′(ylj) ∂yli ∂θ1:l (x)( ∂yli ∂θ1:l (x))t + I ′W
  where IW , I ′W are quantities that converge to 0 as nl grows (self-averaging). The additional term results in an added term Kl(x, x′) in the formula of Kl+1(x, x′).
  we already now that Σ̇l(x, x) = 1. Moreover, we have Σl(x, x) = Σl−1(x, x) + σ2w/2Σ l−1(x, x) = (1 + σ2w/2) l−1 σ2w 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 αL2L 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 c1(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 = (Xi)1≤i≤n, Y = (Yi)1≤i≤N be two centered Gaussian vectors in Rn. Let g ∈ D2(Rn,R). Then we have
  E[g(X)]−E[g(Y )] = 1 2 ∫ 1 0 ∑ 1≤i,j,≤n (E[XiXj ]−E[YiYj ])E[ ∂g ∂xi∂xj ( √ 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− t2Y )] we have that
  G′(t) = n∑ i=1 E[(Xi − t√ 1− t2 Yi) ∂g ∂xi (tX + √ 1− t2Y )]
  Moreover, it is easy to see that for any random vector Z in Rn we have E[Zig(Z)] =∑n j=1 cov(Xi, Xj)E[ ∂g ∂xj (Z)], this yields
  G′(t) = t n∑ i,j=1 (E[XiXj ]− E[YiYj ])E[ ∂g ∂xi∂xj (tX + √ 1− t2Y )]
  We conclude by integrating G′(t) between 0 and 1.
  Now let D = {(xi, zi) : 1 ≤ i ≤ N} be the datapoints. Using the same notations as in the previous chapter, let yL(xi) denotes on the neurons of lth layer (the neurons are iid). Assume c1xi,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 ∈ D2 such that for all i, j ∈ [|1, N |], ∂g∂xi∂xj ≥ 0, there exist β > 0, ζN > 0 such that
  E[g(yLeoc(X))] ≤ E[g(yLord(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 σ20 + 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[yLeoc(Xi)yLeoc(Xj)] ≤ E[yLord(Xi)yLord(Xj)]− κL−β where β = 1 for a smooth activation functions in S and β = 2 for ReLU-like activation functions (see Hayou et al. (2019)). Let λN = infi,j infu∈[0,1] E[ ∂g∂xi∂xj ( √ 1− uX + √ uY )]. Using Lemma 5, we have
  E[yLeoc(Xi)yLeoc(Xj)] ≤ E[yLeoc(Xi)yLord(Xj)]− 1
  2 κN2λNL
  −β
  As a simple application, using the function g(X) = ∏N i=1 1xi≤ti , we have the following Lemma 7. Let t1, t2, ..., tN ∈ R. Then there exist β > 0, ζN > 0 such that
  P(yLeoc(x1) ≤ t1, ..., yLeoc(xN ) ≤ tN ) ≤ P(yLord(x1) ≤ t1, ..., yLord(xN ) ≤ tN )− ζN Lβ
  as a Corollary, we have a generalized form of Slepian’s Lemma for yLeoc(X) and y L ord(X). Corollary 1 (Max Range and Min Range).
  P(max i yLeoc(xi) ≥ t) ≥ P(max i yLord(xi) ≥ t) + ζN Lβ
  and P(min
  i yLeoc(xi) ≤ t) ≥ P(max i yLord(xi) ≤ t) + ζN Lβ
  where β = 1 for activation functions of type S and β = 2 for ReLU-like activation functions.
sections.heading: D TRAINING WITH SGD INSTEAD OF GD
sections.text: 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
  θ̂t+1 = θ̂t − η∇θL(S)(θ̂t), (15)
  where L(S) = 1S ∑S i=1 `(f(x̃i, θ), ỹi) where (x̃i, ỹi)1≤i≤S is a randomly selected batch of size S. Then for all θ
  ∇θL(S)(θ)−∇θL(θ) = ∑ i Zi(S) S (∇θ`(fθ(xi), yi)− E0(θ))− N∑ i=1 (∇θ`(fθ(xi), yi)− E0(θ)) N
  where Zi(S) = 1 if observation i belongs to the batch (x̃j , ỹj), j ≤ S and equals 0 otherwise and E0(θ) = E0∇θ`(f(X1, θ), Y1). We have
  tr [ Cov ( N∑ i=1 (∇θ`(fθ(xi), yi)− E0(θ)) N )] = p∑ l=1 Var (∂`(fθ(X1), Y1)/∂θl) N .
  So that if S = o(N) and if
  tr (Cov (∇θ`(f(X1, θ), Y1))) = o(S)
  where Cov(·) denotes the covariance matrix under P0. Then
  ∇θL(S)(θ)−∇θL(θ) = ZS(θ)√
  S + oP0(S
  −1/2)
  where ZS(θ) converges in distribution (as S goes to infinity) to a Gaussian random vector with covariance matrix Σ(θ) = Cov (∇θ`(f(X1, θ), Y1)) and we have, neglecting the term oP0(S−1/2),
  θ̂t+1 = θ̂t − η∇θL(θ̂t) + η√ S Z(θt). (16)
  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))
  dθt = −∇θL(θt)dt+ √ η
  S Σ(θt)
  1 2 dWt. (17)
  SGD updates can therefore be seen as a discretization of the previous SDE with time step ∆ = η, and where Σ(θt) 1 2 is the square-root matrix of Σ(θt) = Cov (∇θ`(f(X1, θt), Y1)) and (Wt)t≥0 a standard Brownian motion.
  Since the dynamics of θt are described by an SDE, the dynamics of ft 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 ft(X ) is the solution of the following SDE
  dft(X ) = [− 1
  N KLθt(X ,X )∇z`(ft(X ), Y ) +
  1
  2
  η S Γt(X )]dt+
  √ η
  S ∇θf(X , θt)Σ(θt)
  1 2 dWt (18)
  where Γt(X ) is the concatenated vector of (Γt(x) = (Tr(Σ(θt) 1 2∇2fi(x, θt)Σ(θt) 1 2 ))1≤i≤o)x∈X and ∇2fi(x, θ) is the Hessian of fi (ith 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 ft. We denote by ft,i the ith coordinate of ft, i.e., for an input x, ft(x) = (ft,i(x))1≤i≤k. Let i ∈ 1, ..., k, we have
  dft,i(x) = ∇θfi(x, θt)dθt + 1
  2
  η S Tr(Σ(θt) 1 2∇2fi(x, θt)Σ(θt) 1 2 )dt
  = [−∇θft,i(x)∇θft(X )∇z`(ft(X ), Y ) + 1
  2
  η S Tr(Σ(θt) 1 2∇2fi(x, θt)Σ(θt) 1 2 )]dt
  +
  √ η
  S ∇θfi(x, θt)Σ(θt)
  1 2 dWt
  where ∇2fi(x, θt) is the hessian of fi with respect to θ. Aggregating these equations with respect to i yields
  dft(x) = [− 1
  N ∇θft(x)∇θft(X )∇z`(ft(X ), Y ) +
  1
  2
  η S Γt(x)]dt+
  √ η
  S ∇θf(x, θt)Σ(θt)
  1 2 dWt
  where Γt(x) = (Tr(Σ(θt) 1 2∇2fi(x, θt)Σ(θt) 1 2 ))1≤i≤k.
  Therefore, the dynamics of the vector ft(X ) is given by
  dft(X ) = [− 1
  N Kθt(X ,X )∇z`(ft(X ), Y ) +
  1
  2
  η S Γt(X )]dt+
  √ η
  S ∇θf(X , θt)Σ(θt)
  1 2 dWt
  where Γt(X ) is the concatenated vector of (Γt(x))x∈X .
  With the quadratic loss `(z, y) = 12 ||z − y|| 2, the SDE equation 18 is equivalent to
  dft(X ) = [− 1
  N KLθt(X ,X )(ft(X )− Y) +
  1
  2
  η S Γt(X )]dt+
  √ η
  S ∇θf(X , θt)Σ(θt)
  1 2 dWt. (19)
  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 ft for wide FeedForward neural networks : In the case of a fully connected feedforward neural network (FFNN hereafter) of depth L and widths n1, n2, ..., nL, Jacot et al. (2018) proved that, with GD, the kernel KLθt converges to a kernel K
  L that depends only on L (number of layers) for all t < T when n1, n2, ..., nl →∞, where T is an upper bound on the training time, under the technical assumption ∫ T 0 ||∇z`(ft(X ,Y))||2dt < ∞ almost surely with respect to the initialization. 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 ft(X ) for wide networks is given by
  dft(X ) = − 1
  N K̂L(ft(X )−Mt)dt+
  √ η
  S ∇θf(X , θt)Σ(θt)
  1 2 dWt,
  where K̂L = KL(X ,X ) and Mt = Y − ηN2S (K̂ L)−1Γt(X ). This is an Ornstein–Uhlenbeck process whose closed-form expression is given by
  ft(X ) = e− t N K̂ L f0(X ) + (I − e− t N K̂ L )Y +At(X ) (20)
  where At(X ) = − η2S ∫ t 0 e− t−s N K̂ L Γs(X )ds + √ η S ∫ t 0 e− t−s N K̂ L∇θf(X , θs)Σ(θs) 1 2 dWs; see supplementary material for the proof. So for any (test) input x ∈ Rd, we have
  ft(x) = f0(x) +K L(x,X )(K̂L)−1(I − e− t−s N K̂ L )(Y − f0(X )) + Zt(x) +Rt(x), (21)
  where Rt(x) = √ η S ∫ t 0 [KL(x,X )(K̂L)−1(e− tN K̂L − I)∇θf(X , θs) + ∇θf(x, θs)]Σ(θs) 1 2 dWs
  and Zt(x) = η2S [∫ t 0 Γs(x)ds+ ∫ t 0 K(x,X )(K̂L)−1(I − e− (t−s) N K̂ L )Γs(X )ds ] .
  Proof. Using the approximation of the NTK by KL as n1, n2, ..., nL →∞, the dynamics of ft(X ) for wide networks are given by
  dft(X ) = − 1
  N K̂L(ft(X )−Mt)dt+
  √ η
  S ∇θf(X , θt)Σ(θt)
  1 2 dWt,
  To solve it, we use the change of variable At = e t N K̂ L ft(X ). Using Ito’s lemma, we have
  dAt = 1
  N K̂LAtdt+ e
  t N K̂ L
  dft(X )
  = 1
  N K̂Le
  t N K̂ L
  Mtdt+
  √ η
  S e
  t N K̂ L ∇θf(X , θt)Σ(θt) 1 2 dWt
  By integrating, we conclude that
  ft(X ) = e− t N K̂ L f0(X ) + 1
  N ∫ t 0 K̂Le −(t−s) N K̂ L Msds+ √ η S ∫ t 0 e− t−s N K̂ L ∇θf(X , θs)Σ(θs) 1 2 dWs
  we conclude for ft(X ) using the fact that 1N ∫ t 0 K̂Le −(t−s) N K̂ L Msds = (I − e− t N K̂ L
  )Y − η
  2S ∫ t 0 e −(t−s) N K̂ L Γsds.
  Recall that for any input x ∈ Rd,
  dft(x) = [− 1
  N KL(x,X )(ft(X )− Y ) +
  1
  2
  η S Γt(x)]dt+
  √ η
  S ∇θf(x, θt)Σ(θt)
  1 2 dWt
  To prove the expression of ft(x) for general x ∈ Rd, we substitute ft(X ) by its value in the SDE of ft(x) and integrate.
sections.heading: E ADDITIONAL EXPERIMENTS : IMPACT OF THE PARAMETRIZATION
sections.text: Figure 4 shows 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.