Mean-field Behaviour of Neural Tangent Kernel for Deep Neural Networks
Authors: Soufiane Hayou
Arnaud Doucet
Judith Rousseau
Venue:  (2020)

Abstract: 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 ( Schoenhols 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 this 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.

Under review as a conference paper at ICLR 2020
MEAN-FIELD BEHAVIOUR OF NEURAL TANGENT KERNEL FOR DEEP NEURAL NETWORKS
Anonymous authors Paper under double-bling review
ABSTRACT
Recent work by Jacot et al. (2018) has showed that training a neural network of any kind with gradient descent parameter space is equivalent kernel gradient descent in function space with respect to the Neural Tangent Kernel (NTK). Lee al. (2019) built on this result show that the output of neural network trained using full batch gradient descent can be approximated by linear model for wide networks. In parallel, a recent line of studies (Schoenholz et al. (0117). Hayou ct al. (2019)) suggested that 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.
INTRODUCTION
Deep neural networks have achieved state-of- sfate-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 not convex, Gradient Descent (GD) methods are often used successfully 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) (ZZu et al. 2018)).
The training dynamics of wide neural networks with GD is directly linked to kernel methods. Indeed, Jacot et al. (218)) showed that training a neural network with full batch GD in parameter space is equivalent to a functional GD i.e. GD functional space with respect to kernel called Neural Tangent Kernel (NTK). Du et al. (2019) used similar approach to prove that full batch GD converges to global minima for shallow neural networks and Karakida et al.) 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 tensor formalism
that express most of the computations in neural networks. Lee et al.) (2019) studied 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 al. (2019) gives 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. (211)), Lee et al. (2018), Sccoenholz et al. (2017). Yang & Schoenholz (2017). These works analyze the forward/backwwr some quantities through the network at the initial step as 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-paramette 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 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 initialization known as the Edge
propagation of

---

Under review as a conference paper at ICLR 2020
of Chaos (Yang & Schoenholz 201 leads to better training dynamies and that 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 proof' are detailed 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. (0188) 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 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 two layers convolutional neural network with ReLU activation and provided spectral decomposition of the kernel, while in Arora al. (2019). the authors propose an algorithm to compute the NTK for convolutional neural networks. However, for finite width neural networks, Arooa 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 amean-fiele approximation), and we study the behaviour of the NTK as the network depth goes to infinity assuming the width is infinite.
EDGE OF CILAOS AND ACTIVATION FUNCTION
2019) and ScheeShhhheozzeeoo et al. (017)) 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 special initialization can lead 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 value 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 CIFARIO.
Recent works by HHyou et al.
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 the NTK as the depth L goes infinity. Our main results
are:
1. With an Initialization the ordered/chaotic phase, the NTK converges exponentially to a constant kernel with respect to the depth I., making the training impossible for DÑNs (Lemma I and Proposition
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) ?
3. The Edge of Chaos initialization leads to 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 an initialization the ordered/chaotic phase (Propositon) 2
2

---

Under review as a conference paper at ICLR 2020
4. Using 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 polynomial rate
3 NEURAL NETWORKS AND NEURAL TANGENT KERNEL
3.1 SETUP AND NOTATIONS
Consider neural network model consisting of layers (y')ısis with y' :-- + R" no and let 0 (0') 10') be the flattened vector of weights and bias indexed by the layer's index and p be the dimension of Recall that 0' has dimension n4 1. The output of the neural network is given by some transformation 8 RT - IR of the last layer (a:); o being the dimension of the output (e.g. number of classes for a classification problem). For any input fzz,0) 0) = (x)) E IR", As we train the model, / changes with time and we denote by 6 the value of time and fx( = f(z 0,) (z, 0t). S o). Let set and let X' (#i)igigN. y = (Uj) be the matrices of input and output respectively, with dimension d x N and o x N. For any function dd - IR k 2 1, we denote by g(X )") the matrix (g( x. y:)) SISN of dimension
Rd we thus have
(z,, yil be the data
N.
Jacot et al. 2018) studied the behaviour of the output of the neural network as a function of the training time when the network is trained using gradient descent algorithm. [eee al. (2019) built on this result to linearize the training dynamics. We recall hereafter some of these results.
For given 0, the empirical loss is given by [()) = t/2 Ei ((f(s,, m)). The full batch GD algorithm is given by
8.t1 = =0 - (o
(1)
where 7) >> 0 the learning rate.
Let T > 0 be the training time and N, = TT/n 1 be the number of steps of the discrete GD equation The continuous time system equivalent equation with step Al is given by
d0;- = V (0,
(2)
This differs from the result by Lee et al. (2019) since we use discretization step of At 17. It is well known that this discretization scheme leads an error of order O(1)) (see Appendix). As Lee et al. (2019), Equation (2) can be re-written as
d0, N (
V.L(((X))))))
where Val 0₂) is a matrix of dimension oN x and V_((((X' 0,), is the flattened vector of dimension oN constructed from the concatenation of the vectors V_l(z, ()))==((....))> S N. As result, the output function fe(x) satisfies the following ordinary differential equation
dff(z) (af d.),0, Vof(s,,,)Voo(X) (,._V///((((.V/VVVV(..))).).)
V((((())) R
The Neural Tangent Kernel (NTK) KJ defined as the dimensional kernel satisfying: for all Rd
3r, x'
Ky, (2, (11,x')= Vef(z, 0,) Vof(( 0,) E ox
(4)
V flr ((,
We also define (X,X) as the oN x oN matrix defined blockwise by
KP (21, K/ (z2
(11,
Kg (N)
K o uN))
K% (X,
(a (IN,
K (xN,
3

---

Under review as a conference paper at ICLR 2020
By applying equation 3 to the vector X, one obtains
1 N (V,
df((
(5)
meaning that for all j S N df((Tz)) = ----- (,
Infinite width dynamics: : In the case of a fully connected feedforward neural network (FFNN) of depth L and widths n1, n,,,..,, Jacot et al. (2018) proved that, with GD, the kernel K3, converges to kernel KL which depends only on I. (number of layers) for all t T when n1, 122 .... T + 00, where T is an upper bound on the training time, under the technical as- sumption Ied [Vz((f((X)))) ())llatt oo almost surely with respect the initialization weights. The infinite width limit of the training dynamics is given by
1 K (X, (2, 2()
(x)
(6)
We note hereafter K = (x, As an example, with the quadratic loss f(2, v) 2/1l2= - y/1². equation iis equivalent to
df((x) = - ()(()) - V)ddt,
(7)
which is simple linear model that has a closed-form solution given by fx( - fo(x) + (I ---e-- - ))'
(8)
For general input E _dd we then have
fe(x) = fo(x) + (r,
- (o--)(()))
Note that in order for fz)x to be defined, R must be invertible. Indeed, it turns out that training with dynamics is only possible if the NTK is invertible. We shed light on this behaviour in the following Lemma.
Lemma (Trainabilitty of the Neural Network and Invertibility of the NTK). Assume fo(X) + y. Then with dynamics defined by equation & (x') - >ill converges to 0 as t - oo if and only if K_. is sin-singular. Moreover, if KL is singular, there exists constant C > 0 such that for all t > 0.
(x) - V 2 C
Lemmal [ shows that invertible NTK is crucial for trainability. Since KJ, = K- & 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 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 6 such as K<r(,, x') = 5 for all x o x', this example is useful in the next section), then the second part of fx(x:) in equation 9 is constant w.r.t æ. Therefore, the generalization function fz) of the model 9 is entirely given by its value at time zero fo(r), which means that the generalisation error E2,y - yil] remains of order (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 I., 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.nt
to T.) and we no longer need the Edge of Chaos
4

---

Under review as a conference paper at ICLR 2020
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) and 1, since generalization any function and any n. is straightforward.
Consider a FFNN of depth L, widths (ni)igiSL weights w and bias b'. For some input x E R4, the forward propagation is given by
22 Σ j=1
(+) - Σ UB + b,, v((x)= -
w (1) + b/. for 1 2 2,
(10)
where @ is the activation function.
We initialize the model with wly Gid NO, o2 and by itd NO of), where NN,, o²) denotes the normal distribution of mean fl and variance a². For some x. we denote by (x) the variance of y (x).
The convergence of q' (z:) 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 ('(x) converges to point ( b. we
independent of as co. Also the asymptotic behaviour of the correlations between y (x) and y' (x') for any two inputs r and x' is driven by (as, Jw); the authors define the EOC as the set of parameters (w) such that (o' q ow) = 1 where 2 ~ NN(0, 1). Similarly the Ordered, resp. Chaotic, phase is defined by Bld' (/
< 1, resp. Elo't (Q((o, we 2 > 1; more details are recalled Section 2 of the supplementary material. It turns out that the EOC plays also crucial role on the NTK. Let us first define two classes of activation functions.
Definition 1. Let Φ R R be measurable function. Then
1. o is said to be ReLU-like there exist B E such that <(x) >x for > and d(x) = fx for S O.
2. o is said to be in S if 6(0) 0, is twice differentiable, and there exist n > 1, partition (Ad)ssign of R and infinitely differentiable functions 91, 92...., In such that
(x(2) = Σm= A,96 where ((2) is the second derivative of db.
The class of ReLu- activations includes ReLU and Leaky-ReLU, whereas the S class includes, among others, Tanh, ELU and SiLU (Swish). The following proposition establishes that any initial- ization 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 (ab, aw) be either in the ordered or in the chaotic phase. Then, there exist X, > 0 such that
sup (K Al S 74 L
2,,² 2
As a result, as L goes to infinity, K converges to a constant kernel Km 22) À for all x x' E Rd. The training is then impossible. Indeed, we have (x,,') R XU where is the matrix with all elements equal to one, i.e. KT 7 is at best degenerate and asymptotically (in L) non invertible, rendering the training impossible by Lemma We illustrate empirically this result in Section
Recall that the (matrix) NTK for input data X' is given by
K1, Kb, (X,x)
Vo,
5

---

Under review as a conference paper at ICLR 2020
As shown in Schoenholz et al. (2017) and Hayou al. (2019), initialization the EOC preserves the norm of the gradient as it back propagates through the network. This means that the terms fo, o, f,), 8,)" are of the same order. Hence, it is more convenient to study the average NTK (ANTK hereafter) given by /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 invertible kernel as L + oo 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 (0r, ow) EOC.
1. If à is ReLU-like, then for all x E Rdd (" A, (0, 1) such that
of 1all
+ (a.a) L
Moreover, there exist
d
KL (ar
2 10
o |ل L' (x,
A
sup ERd
|a
11
S
(1-(1_))1+47))
L
d
2. If a is in S, then, there exists q > 0 such that K=(a,))
((.x)) 1 + q. Moreover, there +
L.
exist B, C, E (0, 1) such that
B Blog(L)
KJ (²(xxx) - Clog(L)
=((1-)1)))((11111-111 (1-A) A)
sup EF''
L
Since 0 < x < 1, on the EOC there exists a matrix J/ invertible such that K 5 (XX = Lxx((+o(()) as L co. Hence, although the NTK grows linearly with L, 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 [T) However the limiting kernels Koc. carry (almost) no information on 2, x' and have therefore little expressive power. Interestingly the convergence rate of the ANTK to Koc. is slow in (O(L-)) for ReL.Uliike activation functions and O(log( -)) 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 9p. 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 5p
Another important feature of deep neural network which is known to be highly inffuential is their architecture. The next proposition shows that adding residual connections to a ReLU network leads polynomial rate for wide range of initialization parameters.
Proposition 3 (Residual connections). Consider the following network architecture (FFNN with residual connections)
v((x)= = vi 1-1 (x)
Σ 1=1
w 001 (x)) + b/,, for 2 2.
(11)
with initialization parameters o; = and ow > and Φ the ReLU activation function. Let KTe be
the corresponding NTK. Then for all Rdd, k),(), ||=||2 +((z)) and there exists X (0, 1) xrx
such that
(x,
| x |a'll A O(L-1).
sup Edd
an x 24
d
where on and in are given by
ow 2, then ay 1 and 7L
1+a /2, yr.
ofw 22, then az and 71. I--1

---

Under review as a conference paper at ICLR 2020
PuLU withinn N0 1.3) (Dadons end wUt Gu (Onderdd grane
lest L)
(2)
(a) EOC
(b) Ordered phase
(c) FFNN with residual connections
Figure I: Convergence rates for different initializations and architectures. (a) Edge of Chaos. (b)
Ordered phase. (c) Adding residual connections.
ifow 22, then ar. = (2 yk and "r, = (2
Proposition shows that the NTK of a ReLU FENN with residual connections explodes exponentially with respect to L. However, the normalised kernel K (I, x xz,za)/a,,2) )/r where x o x' converges to limiting kernel similar K with a rate (T.-1) for all aw > O. We say that residual networks 'live" on the Edge of Chaos, i.e. no matter what the choice of ow is, the convergence rate of the NTK w.r.t 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 EXPERIMENTS
In this section, we illustrate empirically the theoretical results obtained in the previous sections. We first illustrate the results of Propositions | and Then, we confirm the impact of the EOC and Activation function on the overall performance of the model (FFNN), on MNIST and CIFARIO datasets.
5.1 CONVERGENCE RATE OF KL AS L GOES TO INFINITY
and give theoretical convergence rates for quantities of the form K K K We illustrate these results in Figure [ Figure Ia shows a convergence rate approximately equal to O(L-²) for ReLU and ELU. Recall that for ELU the exact rate is ([log(L)L-) but one cannot observe experimentally the logarithmic factor. However, ELU performs indeed better than ReLU (see demonstrraees that this convergence occurs at an exponential convergence rate in the Ordered phase for both ReI.U and ELU, and Figure ic the convergence rate in the case of FFNN with residual connections. As predicted by Proposition
Propositions
Table[] which might be explained by this log(L) factor. Figure Ib
3 the convergence rate O(L-1) ) is independent of the parameter Ow-
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 E {200, 300} with SGD and categorical cross-entropy loss. Training with full batch GD is practically impossible for DNNs, we use SGD instead (see Section ii the Appendix for more details about how the results extend to SGD) with a batchsize of 64 and a learning rate 10 for L = 100 and 14 for E 200,30 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 (w) and (a,, ) = (0.2, 1.302) for Tanh. These values are all on the EOC (see Hayou et al. (2019) for more details). Table I 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 RRluu- with a E {0.01, 0.02, 0.03)). Even with Parametric ReLU (PReLU) where the parameter of the leaky-ReI.U is also learned by backpropagation, we obtain only small improvement over ReLUU. For
(0, v2) for ReLU, (06, )w) = (0.2, 1. 1.227) for ELU
7

---

Under review as a conference paper at ICLR 2020
Elu
Text
ReLu
ReLU
Training Epoch
Training Epoch
(a) (width,depth) = (200,100)
(b) (width,depth) = (200,200)
Figure 2: Test accuracy for different Activation Functions and (width, depth) on MNIST
Table I:
Test accuracy for FENN with width 300 and depth 300 for different activation functions on MNIST and CIFARIO. We show test accuracy after 10 epochs and 100 epochs
MNIST
CIFARIO
Activation
Epoch 10
Epoch 100
Epoch 10
Epoch 100
ReLU (EOC)
46.53 + 12.01
82.11 4.51
20.38 1.85
35.88 + 0.6
LReLUo.01 (EOC)
48.10 3.31
84.71 + 3.39
22.62 + 1.15
29.44 + 4.14
LReLU0.02 (EOC)
49.09 = 3.58
84.3. 3.98
18.62
30.78 = 6.33
LReLUo.0s (EOC)
50.94 4.48
8.749
21.19
3 2.32
PReLU
51.94
1.54
3.55
36.13 3.83
ELU (EOC)
91.63 + 2.21
96.07 0.13
33.81 +
46.14 + 1.49
Tanh (EOC)
91.16 + 1.21
05.75
1.88
42.40 + 1.13
Softplus
10.11 + 0.09
0.18
0..13
11.09 0.36
Sigmoid
9.85 + 0.11
9.10
10.65 0.25
10.33 + 0.17
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 (0100). For Softplus and Sigmoid, the training algorithm is stuck low test accuracy 10% which the test accuracy of a uniform random classifier with 10 classes.
CONCLUSION AND LIMITATIONS
That the training dynamics of deep neural networks is equivalent to Functional Gradient Descent with respect the Neural Tangent Kernel. In the infinite width limit, the NTK has a closed-form expression. This approximation sheds light 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 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 (and NTK) even for an infinite number of layers: the convergence rate is ((L-1) for ReLU-like activation functions and Olog([)) for class of smooth activation functions.
the EOC leads to invertible ANTK
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) 8). Arora Ghorbani et al. (2019), Huang & Yau (2019), Arora et al. (2019). Therefore, we believe that the NTK a useful tool to partially understand wide deep neural networks (have insights on parrmmprrreeters choices for example) and not tool to train neural networks.
8

---

Under review as a conference paper at ICLR 2020
REFERENCES
S. Arora, S.S. Du W. Hu, Li, R. Salakhutdinov. and R. Wang. On exact computation with an infinitely wide neural net. arXiv preprint UrXiv:1904.11955 2019.
A. Bietti and Mairal. On the inductive bias of neural tangent kernels. arXiv Preprint arXiv: 1905.12173 2019.
L. Chizat and F. Bach. A note lazy training in supervised differentiable programming. arXiv preprint ar arXiv:1812.07995 2018.
D.A. Clevert T. Unterthiner, and S. Hochreiter. Fast and accurate deep network learning by exponen- tial linear units (elus). International Conference on Learning Representations, 2016.
S.S. Du, J.D. Lee, H. Li, L. Wang. and X. Zhai. Gradient descent finds global minima of deep neural networks. arXiv preprint arXiv: 1111188111..00 2018a.
S.S. Du, J.D. Lee, Y. Tian, B. Poczos, and A Singh. Gradient descent learns one-hidden-layer CNN: Don't be afraid of spurious local minima. ICML., 20.8b.
S.S. Du, X. Zhai, B. Poczos, and A. Singh. Gradient descent provably optimizes over-parameterizzed neural networks. ICLR, 2019.
B. Ghorbani, S. Mei, Misiakiewicz, and A. Montanari. Linearized two-layers neural networks in high dimension. arXiv preprint 1904::1999..2/111 2019.
X. Glorot and Y. Bengio. Understanding the difficulty of training deep feedforward neural networks. International Conference on Artificial Intelligence and Statistics, 2010.
S. Hayou, A. Doucet, and J. Rousseau. On the impact of the activation function on deep neural networks training. ICML, 2019.
K. He, Zhang, S. Ren, and Sun. Delving deep into rectifiers: Surpassing human-level perfor- mance imagenet classification. ICCV, 2015.
W. Hu, C. Junchi Li, L. Li, and Liu. On the diffusion approximation of nonconvex stochastic gradient descent. arXiv preprint arXiv: 2018.
J. Huang and H.T Yau. Dynamics of deep neural networks and neural tangent hierarchy. arXiv preprint 19099v9909.08
2019.
A. Jacot, Gabriel, and C. Hongler. Neural tangent kernel: Convergence and generalization neural networks. 32nd Conference on Neural Information Processing Systems, 2018
R. Karakida, S. Akaho, and S. Amari. Universal statistics of Fisher information in deep neural networks: Mean field approach. arXiv preprint arXiv:1806.013/6, 2018.
M. Kubo, R. Banno, H. Manabe, and M. Minoji. Implicit regularization in over-parameterized neural networks. arXiv preprint arXiv: 1903.019777 2019.
Lee, Y. Bahri, R. Novak, S.S. Schoenholz, J. Pennington, and J. Sohl-Dickstein. Deep neural networks as Gaussian processes. 6th International Conference on Learning Representations, 2018.
J. Lee, L. Xiao, S. Schoenholz, Y. Bahri, J. Sohl-Dickstein, and J. Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. arXiv preprint arXiv: 1902.0672,, 2019.
D. Lei, Z. Sun, Xiao, and W.Y. Wang. Implicit regularization of stochastic gradient descent in natural language processing: Observations and implications. arXiv preprint arXiv: 1.XXv:18//.00666 1.00659, 2018.
Q. Li, C. Tai, and E.E. Stochastic modified equations and adaptive stochastic gradient algorithms. arXiv preprint arXiv:1511.0625,, 2017.
Q Nguyen and M. Hein. Optimization landscape and expressivity of deep CNNs. ICMI., 2018

---

Under review as a conference paper at ICLR 2020
S.S. Schoenholz, J. Gilmer. S. Ganguli, and J. Sohl-Dickstein, Deep information propagation. 5th International Conference on Learning Representations, 2017.
G. Yang. Scaling limits of wide neural networks with weight sharing: Gaussian process behavior, gradient independence, and neural tangent kernel derivation, arXiv preprint 2019.
G. Yang and S. Schoenholz. Mean field residual networks: On the edge of chaos. Advances in Neural Information Processing Systems, 30:2869 (202869-2869, 2017.
C. Zhang, S. Bengio, M. Hardt, B. Recht, and O. Vinyals. Understanding deep learning requires rethinking generalization. arXiv preprint a 2017.
D. Zou, Y. Cao, D. Zhou, and Q Gu. Stochastic gradient descent optimizes over-parameterized deep ReLU networks. arXiv preprint 1rXiv:1811.088800 2018.
10

---

Under review as a conference paper at ICLR 2020
We provide in Section A and Section the proof of the theoretical results presented in the main document. Section
provides additional theoretical results while Section ?? presents additional experimental results.
C
A APPENDIX: PROOFS OF SECTION 3: NEURAL NETWORKS AND NEURAL TANGENT KERNEL
Lemma (Trainability of the Neural Network and Invertibility of the NTK). Assume fo(x) + V. Then with dynamics 18). [1f((X)) - Y converges to 0 as + oc if and only if K is non-singular
Moreover, if K= singular, there exists constant C 0 such that for all
(1) - V > C
Proof. Assume fo( * y. Let R³ Q" DD be the spectral decomposition of the empirical NTK; i.e. Q is an orthogonal matrix and D is a diagonal matrix.
We have that 2 tht e- D. Q 2-" Diagle #1) siSoNQ where (d:) are the eigenvalues. We also have Ilf - -11 - 1-lle- - k (fo(X) -)]l Therefore, the equivalence holds true.
Moreover, assume R in singular. Let Z, = Since D has least one zero diagonal value, then there exists j E {1,2, ..., oNN such that for 4 (Ze), (Zo)j. and we have
We have that Ze = - Zo-
[[f((z) - -01 ||
2 (Z)j ((oo)ll
Lemma 2 (Discretization Error for Full-Batch Gradient Descent). Assume Vo is C-lipschitz, then there exists C' > 0 that depends only on C and T such that
110r, 0zll S n/C
sup kE(o,T/n]
Proof. For [0, T), we define the stepwise constant system 0e = - Let E [0, 7), we have
Wt/nl
0. 0o 1
voL(Ox)
0o VoL(o.) + / ToC(0..jsdd)
2//
-1t//-1
Therefore,
110, -000 / [[VoL(0.) VoL(O,)] -o +
1 [[VeL(0
10]t/--1
C / 110, 0ollds + - - -
Moreover, for any A E [0, [+ [+|TT/ l//), we have
10 0nl
5 (1+ (1+nC)110==+ 0oll
$ (1 + T//0// - 0oll
< CT 1 000
11

---

Under review as a conference paper at ICLR 2020
where have used log(1 + nC) $ ncC Using this result, there exists constant C depending T and C' such that
m(t/n - L*]t/m) [=]]//mm
Boll + C110 1-11/111- 0oll)
S nn
Now we have
110, - s. C J. 110, 0,llds n1O
so we can conclude using Gronwall's lemma.
B PROOFS OF SECTION 4: IMPACT OF THE INITIALIZATION AND THE ACTIVATION FUNCTION ON THE NEURAL TANGENT KERNEL
We first recall the results obtained in Lee et al. (018)) Schoenholz et al. 2017) and Hayou et al. (2019) where the impact of the EOC (Edge of Chaos) on the initialization studied. We also present some results that we will be used below.
Consider FENN of depth L, widths (mi)ssgs,, weights w and bias V/ For some input x E Rd, the forward propagation is given by
Mz-1
v = Σ whyy + b
v((a))
Σ 1=1
w Wly - b.. for 2 2,
(12)
where ø is the activation function.
We initialize the model with wly did ~ NN my -4m) and by id NOO, of), where Nu, o²) denotes the normal distribution of mean a and variance a². For some x, we denote by q (x) the variance of (z)) In general, g' (a:) converges to point ((0,, aw) > 0 independent of as + co. The EOC is defined by the set of parameters (ar, ow)) such that BB (1fl"(vv(an, aw [2) Similarly the Ordered, resp. Chaotic, phase defined by B.d'd (v ( 22) < 1, resp. Elo ow 2) > (see Hayou et al. (2019) for more details). For two inputs x, x, e Rd, define Σ' E'<x,x')= = (x)y' (x')] and let c' c'(x, x') be the corresponding correlation. Let f be the correlation function defined implicitly by c'+1 f(c'). In the limit of infinitely wide networks,
where Z (o 1).
have the following results (Hayou et al. (2019)
Σ xx, = + - -o
There exist q, À > 0 such that, for all supzers 15' (x, - of e-M
On the Ordered phase, there exists 7 > 0 such that supa, le (xx x') 1| S e-
On the chaotic phase, there exist 0 and < such that suPmyyarRR (c, cl <
e-
For ReLU network on the EOC, we have that E' x)) = ||||| for all 2 1. Moreover,
we have
2
f(x)
x x
(1
x²) 3/2 O(11 - 3/23
z-1
3π
In general, have
of + o F 7L
f(x)
where 7(x) x71 + V1 x2 7. and 71, Zz are iid standard Gaussian variables.
On the EOC, we have f1( =
If o is k-times differentiable, then f is k-times differentiable and for all 1 S J s k, we have (x) -' (B((())) ((((()((2)) ()) (7()))
12

---

Under review as a conference paper at ICLR 2020
From Jacot al. (2018), have that
(x,x²)) (I, (x, (x, x') + [xx x').
where the definition of (xx,'') is given in Proposition |
below.
Definition 1. Let ò + E be measurable function. Then
1. o is said to be ReLU-like there exist X, B R such that 6(x) xx for >> and ((x) = fz for x S O.
2. & is said to be in S if x(0) 0, á is twice difserentiable, and there exist nt 2 1, partition (Ai)isisa of R and infinitely differentiable functions 91, 92, -22;-19 In such that d(2) Ei=1 LA,g where 2(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 (ai) be sequence of non-negative real numbers such that v1 2 0, 0,111 I aaa he where a E (0, 1) and k. 3 > Then there exists y > 0 such that VI 2 0, S em
Proof. Using the inequality a,, can easily that
aj agd Σ -81-)111))
5=0
agd 1 + - 81/2 kiza"
1/2
where we divided the sum into two parts separated by index 1/2 and upper-bounded each part. The existence of is straightforward.
Proposition 1 (Limiting Neural Tangent Kernel with Ordered/Chhotii Initialization). Let (ob, aw) be in the ordered or chaotic phase. Then, there exist X, > such that
sup |K" Al s s- +
2, ER
Proof. From Jacot et al. (2018), we have that
(' = K'
where Σd (,,,,'') of a "T=T' and (x) of = (o)) and
('xxx')) By=N(oos--1jd((((()))) (f(x'))) (f the ordered/chaotic phase, Hayou al. 2019 showed that there exist k, 7,40 and E (0, 1) such that for all 2 to we have
(x, k < -m
sup
a.d'
and
sup S1x S a
a.t'RR
Therefore we have for any 2 lo and
2, e Rdd
Kt (z, x') s ak a t
Letting = (I,
k I we have
r S 071-1
We can now conclude using Lemma
Now, show that the Initialization the EOC leads invertible NTK even if the number of layers L goes to infinity. We first prove two preliminary lemmas that will useful for the proof of the next proposition.
13

---

Under review as a conference paper at ICLR 2020
Lemma 3. Let (a:), (bi), (At) be three sequences of real numbers such that
= a1-1 + bj
A = 1 a
7 + 0(1-100)
bi q 0(1-1)
where a E N,, ,89. Rt and a > Then,
q + 0(1 m)m((.0) 1 1 +
Proof. is easy to see that there exists a constant G |ai E G x + |aol for all / 2 0, therefore
(ar/l) is bounded. Now let Ç min(1, 8) and ri 4. We have
n- + )1
+ 0 + 0 1 + 0(1-3)
m-1 1 ++a a
+ - + +0 -
Letting 21 (n T
there is exist constant M 0 such that
1+a M
3 S 0111
Hence, we have
21 S ro
II (1-
M
Σ
II
(1
1+ a
k
41-5'
By taking the logarithm of the first term in the right hand side and using the fact that Σ we have
t/2 ~ log(!),
II
+ a
~ 11-
k:
k=
For the second part, observe that
II f=b+1
(1
1 a
(1
(1-a- a
1)!
k! (k
j
1)! 1!
and
k!
1
~ k--5
(k 1)! ki+t
so that,
k!
k
(k: a 1)! k
k=
k=l
1a
1
14-
a +
therefore,
Σ 1+ 1
(1
a 1)!
1
1
km1j=k41 kl+c
!!
(k a 1)! kl+c
km
1
-
We can now conclude using the fact that a: >
B 3-1.
14

---

Under review as a conference paper at ICLR 2020
We now introduce different form of the previous Lemma that will be useful for other applications. Lemma 4. Let (ar), (br), (At) be three sequences of real numbers such that
a 01-1A1 + bj
A 1
α
log(l)
0(1
1
= q 0
where a E N 33,q RR and > 1-1,8 1, 2
Then, there exists A, B > O such that
A (1g(()
01 q
s B log(()
1
1+a
Proof. is easy to see that there exists constant G lai S G x + laol for all / 0, therefore (ar/l) is bounded. Let 4. We have
7i 1 70
log(1)
0(-1--9))
+ 2 0(-2)
[
12
-1
1 + a
+ 18
log(1)
+
4
+ 0(1
p2
1
Let 21 11 Ita
It is clear that \₁ 1 - a/l + 0 3/2). 0(1-3/2).
Therefore, using Lemma with
1/2, we have ri- Ita
q Thus, there exists K1, K2; M, such that for all 2
1 1+a + K-
log(()
M
1+a
log(1) + R2 12
+ M
T-
12
12
san
sn
20
Similarly to the proof of Lemma
3 it follows that
S ølo
II
1+a + II
(1
1 a K2 log(k) + M
k2
k=lo
km kmlg k++
and
21 2 00
+
Σ
II
1
(1
1 1+
K1 log(k) M
j
k=o
k=
j=k+1
Recall that we have
1 I[ k=1
1 a
~ 111-
and
II j=b+1
(1
1+a,
(1 (1-ce a - 1)!
k!
1!
(k
1)!
so that
kd!
Ki log M
~--000 log(k:)h:a
(h: 1)!
222
Therefore, obtain
k!
Ki log (k) M
log(k)ba*
(k
1)!
122
log(t)(a-1og
log(a)
15

---

Under review as a conference paper at ICLR 2020
where C1 > 0 is constant. Similarly, there exists constant C2 such that
k!
R2 log(k) +
M
C21 log(a)
(k a 1)!
h2
We conclude using the fact that (1-a-11 ly!
111a
Proposition 2 (Neural Tangent Kernel on the Edge of Chaos), Let @ be non-linear activation function and (os, aw) EOC.
1. If à is ReLU-like, then for all x E Bdd K² (x,)))
(z,a),
+
Moreover, there exist
d
L
A, E (0,1) 1) such that
(x,,,")
A L
sup xfrrr ER.
x
|| |||2||11|
S
L
d
K²
(a,0))
2. If a is in S. then, there exist > 0 such that exist B,C, C,C,EE E (0, such that
q +
- q. Moreover, there
L.
Lo)(L)
K x²)
9 S Colog())
sup 2a
L
ERd
L
L
Proof.
We use some results from Hayou et al. (2019) in this proof.
Let e, x' E IRd and x.,' (E Let = 1 2, and ʃ be the correlation
function defined by the recursive equation d = f(x²). From the preliminary results,
we know that (xx x) = - | and that K' K'(x,x')= K² (x, x' ()) (z, x') + >n' (z, x').
This concludes the proof for K (zz x). We denote 2v 3
From Hayou et al. 2019))
have on the EOC M+1 = - 3/2 + 0 so that
1/11 - 1/2 - sy 0(0i (/3))) 7 1/2
1/2
1/2 + 0
2
= 1/2 +, O(n).
Thus, 1 goes to infinity
1/2
1/2
8
1+1 %
2
and by summing and equivalence of positive divergent series
7 1/2 I
-1/2 - 7/1 Moreover, since 7/41
1/2 =
0000m = 0(12), we have 7, -1/2 +(1)
Therefore, e we also have
2 + 0(13))
1 arcsin(r)
1
f((x)
π
V2
(1 - 1/2 + 0(11 - 2))0/
π
Thus, it follows that
3 3- 0
=1 =
Moreover, o' + 0(122) where g is the limiting variance of y'
Using Lemma 3 conclude that K'(x,z)) 1 o² bounded, this result is uniform x, r'. Therefore, we can take the supremum over x, x' d
d |*||||||||| O(1-)). Since cn is
16

---

Under review as a conference paper at ICLR 2020
1. We prove the result when (x) = Ix<u91 (z) + 12209((x). The generalization to the whole class is straightforward. Let be the correlation function. We first show that for all (1-29)(h 2779k (2) where gk E c
& 2 3 ((8 (x)
We have
f" = BB (V7Z1)0" (VqU((()))))
= o (1221)) (V (adp:(/))))Hodd(dd((//a///oooa/d
Let G(x) (VQZ)) (172 (VTU((1)) then
= - E[0"((2))
I
T )02)00)))
1 V1 (VTU2(x))
-1-2²
x²
+ Elod \1st))))
x
-I-
It is easy to that Gx(x) Vi-1² G1 where G1 E C³. A similar analysis can be applied to the second term of f We conclude that there exists ga E C such that ((3 g We obtain the result by induction.
Since f(k) are potentially not defined at I, we use the change of variable x = to obtain Taylor expansion near 1. Simple algebra shows that the function f(1 (2) has a Taylor expansion near O:
6 111-1 10- --
(1)+
6
Therefore,
f(x) = 1+ (z (x- (x - 1)2 f (1) +
(1
x/3
f( (1) O((x 1)9).
6
Letting A := = 1- c', there exist 5 > 0 such that
A141 A a - BA 0
therefore,
X -1 8x² + -(
(1+ + 8X3 + O()?))
-- BA + 0
By summing (divergent series), we have that X 1. Therefore,
D.
XT11 A a 3a 1 0(1-2)
By summing a second time, we obtain
X al + 3a log(1) + 011)
so that A = a 11-11 - a1 ( logil) + 0(1-2).
Using the fact that f'(x) = 1 + (x (1) + 0((x - 1)²), we have ((().r)) = 2 + log(1) of the boundedness of c ", we can take the supremum.
12 + ( We can now conclude using Lemma Using again the agrument
17

---

Under review as a conference paper at ICLR 2020
Proposition 3. Consider the following network architecture (FFNN with residual connections)
R1-1
vl(x) vi (x)
+
w Hy (x)) + b/, for 2 2.
(13)
1=1
with initialization parameters Tb and Tn O. Let Kiee be the corresponding NTK. For all Rdd K (x,zt
|| (L- and there exists X (0, 1) such that
Lx24
KL
|||||| x ||''|| Al = 0
sup ERK'
o"
L
Proof. We only give a sketch the proof. A more rigorous proof can be easily done but it unecessary. As in the feedforward without residual connections case, is easy that KL following recursive equation
res satisfies the
(x, = Kiedd (x, x' (2' (x, + + E (xr, x²)
To this, with FeedForward Neural Network, have (Jacot al. (2018)
0y: 1+1
0y:t (o'))²
Oyk 00
Qy (x))² Iw
(z)(
Σ
w w
('()))
(z)
001
00,
0011
3.5"
whereas for Residual Networks, we have an extra term
0yt+1
Oy:t +1 (x'))² = Oul (1))
Oy! (z)))*
Oyl (x)
Out (x))² + T'
(x)(
+ Σ
1+1 1+1
w
o ((()))
001 0011
301 0014
0014
0011
3.5"
where Iw, W are quantities that converge to 0 ni grows (self-averaging). The additional term results in an added term K3 (xr,ar) in the formula of k1+1(x, x').
we already now that (x, 1. Moreover, we have Ex x) Σ (x, x) + + (x,x))
(1+ = (2) a |d|. Depending on the value of aw, the behaviour of K' (x(xx')) x') changes. However, from Hayou et al. 2019), have that (x (x,xx')
91 ((. So by scaling with az
and using Lemma 3] we conclude on the convergence rate of 0(1-1). We can take the supremum as the result of the boundedness of c! (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, 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:)ssigm Y (Yi)ssigN be two centered Gaussian vectors in R., Let g E D2 (RR", IR). Then we have
0g (VI - uX + /uY)ldu (14)
E[g(X)) E[g))))))((())----) 2 Istdisn
Σ
dr,Or,
The result Lemma 5 is valid when the second derivatives of exist only in the distribution sense (e.g. Dirac mass)
18

---

Under review as a conference paper at ICLR 2020
Proof. We define the function G on R by
G(t) Elg(tX V (2Y)
we have that
0g - 12 Oxi -I-
(t = Σ B [
Yi)
(tX + VI-
Moreover, it is easy to see that for any random vector Z in R we have E[Zig(Z)
Σj=1 coo(X, Xy) [ dtzy
Og (Z)), this yields
Og B dr,dr, G'(t) =t Σ [X
(tX V PYY)
We conclude by integrating (1) between and 1.
Now let D = [[x 2;) 1 I S ]] be the datapoints. Using the same notations as in the previous chapter, let (xx,) denotes on the neurons of 7 layer (the neurons are iid). Assume c 2 0 for all j E [] (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 non ReLU-like activation function and (ov. Tw) & EOC. Then there FOCC such that for any function g E DD2 such that for all i, E
exists (ab, EOC, FOC)
[11,N], 0n;0n;
8g 2 0, there exist 3 > 0, CN > 0 such that
SN L
Eig(veo (X)) < Big(vera
Proof. Let (ob, & EOC and be the corresponding limiting variance. From the previous chapter, it is easy to that there exists dg >> 0 such that o [/e'( q. Let (o,, ww,Eo)) (ou, // ( EOC. There exists constant x (independent of N) such that for all i [1I, NI. B[ ee (X,)] < E (x
KL where 3 1 for smooth activation functions in S and B = 2 for ReLU-like activation functions (see Hayou et al. (011 ). Let
AN = inf,, inf E drjdr, (vI uX + VuY)). Using Lemma we have
0g
1 2 K
Elytne (X.) (X,,) < Elyvoo(X,
As simple application, using the function g(X) -1 In we have the following Lemma 7. Let t1, 12, IN R. Then there exist 3 0, ÇN >> O such that
SN
(a S 4,, (xN) S (x) < Pl (x1) S bi,..., (aNa S (N) 7,3
as a Corollary, we have a generalized form of Slepian's 's Lemma for y (X) and (X). Corollary (Max Range and Min Range).
Pmmma you (I;) 2 t) 2 P(max vord 2 2) + SN
and
ÇN
F(min ( <() Pm (2x) <
where 3 = 1 for activation functions of type S and 2 for ReI.Ullkke activation functions.
19

---

Under review as a conference paper at ICLR 2020
D TRAINING WITH SGD INSTEAD OF GD
In this section, extend the results of the NTK to the case of SGD. We use approximation of the SGD dynamics by a diffusion process. We assume implicitly the existence of the triplet (Ω, P, F) where 52 is the probability space. IP is probability measure Ωt and F is the natural filtration of the Brownian motion. Under boundedness conditions, when using SGD, the gradient update can be seen GD with Gaussian noise (Hu al. 2018 be the batchsize. The SGD update is given by
Li al. 2017). More precisely, let of (N)
O = - (
(15)
where ((8)
2 ((f(Z;,)) 0),0;) where (āi, Isiss is a randomly selected batch of size
Then for all 0
Z((5) S (Vol(J((m))))) (o( (ol(fo((z)) Mb) - Eo(0))
V VoL(5) VoL()- Σ
N
where Z,(S) = observation belongs to the batch (ā,, Vi).J S and equals otherwise and
Eo( E. Vo)(f(X1.0), Y1). We have
N
vol(fo(x))).)) - Fo(o)) )
Var (XL) 11)/00,)
tr
Cov
N
N
So that if = o(N) and if
tr (Cov O), Vi)) = o(S)
where Cov(-) denotes the covariance matrix under Po. Then
Zs( VS
VoL(s) (0) VoC(O)=
+ on, (S-1/2)
where (0) converges in distribution (as S goes to infinity) to Gaussian random vector with covariance matrix (0) = Cov (V VVl(f(X,, Oi) Yi)) and we have, neglecting the term op, (8-1/2),
0.4 0. VoL(())
"
Z(0.).
(16)
VS
can in particular bound the difference between equation approximation (see also Hu et al. (2018) and
16 and the continuous time SDE
Li et al. (017))
d0,- = - (ol())d + <(..)-d..
(17)
S
SGD updates can therefore be seen as a discretization of the previous SDE with time step A = and where E(6,) is the square-root matrix of Σ(O,) Cov (V) and (W,) standard Brownian motion.
Since the dynamics of 6g are described by an SDE, the dynamics of fa can also be described by an SDE which can be obtained from Itô's lemma.
Proposition 4. Under the dynamics of the SDE equation the vector Sc(X) is the solution of the following SDE
1 V df((x) =1- (X,, + (2)))d + Voo((2 (18)
1
25
where (2(X)) is the concatenated vector of (T) =
(ssss))
and 22f(2 8) is the Hessian of f; (ith component of f) with respect to 8.
20

---

Under review as a conference paper at ICLR 2020
Proof. Since 8, is a diffusion process, use Itô's lemma to deduce how the randomness propa- gates to ft. We denote by ft, the ith coordinate of fi, i... for an input 2. f((Ø) = (fsss= Let i L,..,,, k, we have
1 dfl,(x)) VoSi(e, 0,)d0, +
(,)
2
17 V2
= -t-Vff.lt
25
+ V " Vofile, ( 0,) dw,
where V2f((2,01) is the hessian of with respect to 8. Aggregating these equations with respect to i yields
df((x) [-- ~ V
2
1 (x))dd + V
25
where Ir(xx) (1. (2,((x)))((()))) (.))),,,))()))))))))))))))) Therefore, the dynamics of the vector fx(x') is given by
1 1 (x)]dtt
(x) = -
Y) + 2
where (x()) is the concatenated vector of (Te(z))xEx.
With the quadratic loss f(2, y) = y/13. the SDE equation 18 is equivalent
1 dft( - 11 1 (2, (f+(()) )/) TT(()))d V Vof[(x
(19)
25
This is Chistein-Uhlenbeck process (mean-reverting process) with time dependent parameters. The additional term I' 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 GD (Kubo et al. (2019), Lei et al. (2018)).
Dynamics of fo for wide Foward neural networks:
In the case of a fully connected feedforward neural network (FFNN hereafter) of depth L and widths 11, 122 Lz Jacot et al. (2018) proved that, with GD, the kernel K converges to kernel K* that depends only on L (number of layers) for all t < T when n1, n22 .... my + co, where T is an upper bound on the training time, under the technical assumption fo (V (X))") 122d co almost surely with respect to the initialization. For SGD, we assume that the convergence result of the NTK holds true as well, this illustrated empirically in figure 3 but we leave the theoretical proof for future work. With this approximation, the dynamics of (x( for wide networks is given by
1
n V
df((x)= - R (J((2) -M))d
where K4= K- (x and M, Y whose closed-form expression is given by
nN
(1,(x), This is an Omsttein-Uhlenbeck process 28
f(()= = 1 &
R fo(2)) + (T - y + Ax(x)
(20)
where A.(X) - 23 - ( + / - Vol(X 0, (.,. see sup- plementary material for the proof. So for any (test) input x E Rdd we have
s((x)) fo( + K² (x,
' R4 )0 fo(x)) (z) Rx(x)
(21)
where R((x)) /
- 1)Vol(X 0,) +
and (z(x) = 2/ [z) + 2) السي k4 (x))ss
21

---

Under review as a conference paper at ICLR 2020
Dopth L:50
1/14
Depth L-100
1//
Witth
(a) 1=0
(b) 1000
(c) t=1000
Figure 3: Ratio K K for three randomly selected pairs from MNIST dataset as a function of width for three training times t 0, t 100 and 1000 (training time is measured by SGD updates)
Proof. Using the approximation of the NTK by KL n1, n2..... + the dynamics of fe() for wide networks are given by
dx' 1
(f((x) MM)ddt + T / (ofx
To solve it, we use the change of variable At =et f((X)) Using Ito's lemma, we have
dA, N RL Aidt + te df((2)
Kke R² M,dt
1 _ Vof((x0))(()))
By integrating, we conclude that
fc(x) =e x R¹ fo(x) N
1 [ kte - Rt. M.ds -
7 [ - (o
JS
we conclude for Se(X) using the fact that N Sor - M,ds = (1 - - 2S -ic- k T,,ds.
Recall that for any input E Rdd
df. (f(()) -(=))=1~ (1,,
Vo 1
(x)))),)))))))
To prove the expression of fz( for general 1dd, we substitute (x( by its value in the SDE of fe(x) and integrate.
E ADDITIONAL EXPERIMENTS: IMPACT OF THE PARAMETRIZATION
Figure 4 shows the impact of different parametrizations the training dynamics. It shows that both standard parametrization and NTK parametrization have similar behaviour on different initialization regimes.
22

---

Under review as a conference paper at ICLR 2020
1.0
of Chass
14
Ordered prase
pawpammmtrizzttonn
standard par.mmrrionnmmmniiattion
Training Accuracy
0.8
0.8
Training
3.4
84
022
nik pat.metrietiiiinonn
82
standard standerd_perameenn on
epach
epach
(a) EOC
(b) Ordered phase
Figure 4: Impact of NTK parametrization on the Training
23


Reviews:
[{"review_id": "S1lrpO86YH", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper studies the role of weight/bias variance of neural network’s trainability via analyzing large depth behaviour of Neural Tangent Kernels(NTK). \n\nRecently NTK has been a popular topic of study in theoretical deep learning as it describes exact gradient descent dynamics of infinitely wide networks. Original NTK paper (Jacot et al. (2018)) and other follow up papers often gloss over role of weight/bias scales whereas in the setting of signal propagation or NNGP covariance, Schoenholz et al. (2017), Lee et al. (2018), Novak et al. (2019), Hayou et al. (2019) have shown understanding initialization scale is quite important as networks becomes deeper. This paper brings those analysis for NTK and discovers few interesting results. \n\nThe good weight/bias initialization scale that propagates signal for very deep network is denoted edge of chaos (EOC). The authors show that 1) for fully connected feed forward networks, outside initialization edge of chaos (EOC) the NTK converges exponentially to constant kernel. This indicates non-trainability. 2) With EOC initialization the convergence is polynomial and the NTK remains invertible for very large depth implying trainability. 3) Certain class of activation function (denoted class S, including ELU/Tanh/Swish) it has even slower convergence(O(log(L)/L)) compare to ReLU (O(1/L)). 4) Residual FC networks NTK has polynomial convergence for all weight and bias variance. \n\nIn terms of novelty, the paper is combining existing techniques and objects to study large depth behaviour of NTK. However the contribution of the paper is interesting and worth the ICLR audience to know about, especially with the current surge of interest in NTK.\n\nOne main weakness is that the experiments are weak and have a very weak connection to the early part of the paper. Section 5.1 is direct convergence comparison, which provides fair evidence. In section 5.2, I’m not certain whether experiments displayed connects to asymptotic NTK analysis. First of all, in order to obtain NTK in the infinite width limit, one has to use `NTK parameterization’ where one scales out 1/sqrt(fan_in) in network definition and not in weight initialization. It seems (by mention of He/Glorot init, and choice of learning rate O(1e-3/1e-4)) the authors experimented with standard parameterization. In this case the connection to very deep behaviour of NTK is not straightforward to actual network training since the training dynamics will be different.  I suggest authors try experiments in NTK parameterization and see if results are similar.\n\n\nFew comments:\nIt should be emphasized that the infinite width is taken first before taking infinite depth limit. There are subtle effects when depth/width are taken to infinity at the same time. The analysis on asymptotic behaviour of NTK at infinite depth is only valid after width is taken to infinity. \n\nFirst sentence of the abstract is too strong in the sense that NTK has limitations. NTK certainly does not work for any kind of network, so I suggest authors to down tone the sentence. \n\nP4 paragraph after Lemma1 `'K^L(x, x’) = cte’ is a typo?\n\nIn section 5.2, the authors’ claim full batch GD is practically impossible. One could accumulate gradients over minibatch to simulate full batch GD. \n\nP13 typo in last paragraph, `  NTKl'\n\nEDITS POST AUTHOR RESPONSE:\nRegarding point 5) I was referring to solving the memory problem. In practice, if one is sampling without replacement, there is no randomness accumulating gradient of the full epoch. \n\nI also thank the authors for the clarification. My score still remains the same. ", "extra_scores": {"experience_assessment": "I have published in this field for several years.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "HylpGdsQFS", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nDetailed Comments:\n1. 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.\n2. 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.\n3. A typo in the first paragraph in Sec 4, ouerselvs to ourselves.\n4. 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.\n5. 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.\n6. There are several typos in the appendix. Please go through the appendix and fix them.\n7. 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. \n8. Better use the Stirling’s approximation rather than k to infinity in the calculation of k!/(k-\\alpha-1)!\n9. 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?\n10. I also feel the content in Appendix D is unnecessary for this paper.\n\nOverall, 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.\n\n[1] Arora, Sanjeev, et al. \"Fine-Grained Analysis of Optimization and Generalization for Overparameterized Two-Layer Neural Networks.\" International Conference on Machine Learning. 2019.\n[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\n", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "BJgUpGxAFS", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nThe detailed characterization of these limiting kernels for different activations and initializations is interesting. Yet, the work is quite incremental and of limited significance, in that such EOC initialization was known to be important for controlling propagation of both activations and gradients with such networks, so it is no surprise that the NTK has similar properties, given that it basically consists of the sum of similar gradient covariances.\n\nsome comments:\n- title: isn't the \"mean-field behavior\" already subsumed in the definition of NTK?\n- contribution 4.: \"more suitable\" seems a bit strong if it's just a log(L) factor?\n- after lemma 1: generalization may not provide a strong argument here unless further discussed: the constant kernel is obviously bad for learning anything, but the EOC limiting kernel is also pretty bad for predicting anything outside the training data\n- section 3: typo \"ourselves\"\n- prop 3: specify that phi is the relu\n- the proof of prop 3 should be more detailed. Also, it is not obvious that you are giving the correct formula for the NTK of resnets\n- Table 1: is it meaningful to compare performance with such low accuracies?", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper focuses on studying the impact of initialization and activation functions on the Neural Tangent Kernel (NTK) type analysis. The authors claim to make a connection between NTK and edge of chaos analysis. The reviewers had some concern about (1) impact of smooth activations "any NTK-based training method for DNNs should use a Smooth Activation Function from the class S and the network should be initialized on the EOC" (2) proofs of residual networks (3) and why mixing NTK with EOC is interesting. Some of these concerns were addressed in the response. I do share the reviewer concerns about (2). The authors need to give a clear proof. I think this combination of NTK and EOC could be interesting but needs to be better motivated. As a result I do not recommend publication.
Decision: Reject