The Curse of Depth in Kernel Regime
Abstract
Recent work by Jacot et al.
(2018) has shown that training a neural network of any kind with gradient descent is strongly related to kernel gradient descent in function space with respect to the Neural Tangent Kernel (NTK).
Empirical results in (Lee et al., 2019) demonstrated high performance of a linearized version of training using the so-called NTK regime.
In this paper, we show that the large depth limit of this regime is unexpectedly trivial and we fully characterize the convergence rate to this trivial regime
Introduction
The Neural Tangent Kernel (Jacot et al., 2018), a.k.a the NTK, has been the main focus of a growing number of works aiming to understand the inductive bias of Deep Neural Networks (DNNs).
To cite a few, Bietti and Mairal (2019) ; Karakida et al.
; Yang (2019); Arora et al.
Bietti and
Bach (2021).
In the so-called NTK regime (inﬁnite width), the whole training procedure is reduced to a linear model given by the ﬁrst order Taylor expansion of the output function near its initialization value.
It was shown in Lee et al.
(2019), that such simple models could surprisingly achieve high performance.
However, most experiments with NTK regime were conducted on shallow networks and have not sufﬁciently covered DNNs.
NTK regime (Inﬁnite width) for DNNs.
The inﬁnite width limit of the NTK for different architectures was studied by Yang (2019), who introduced a tensor framework that allows the derivation of recursive formulas for the NTK.
Information propagation.
In parallel, information propagation in inﬁnite width DNNs has been studied in several works (Hayou et al., 2019; Lee et al., 2018;
Schoenholz et al., 2017; Yang and
Schoenholz, 2017a; Poole et al., 2016) where the authors identify a set of hyper-parameters known as the Edge of Chaos ( EOC ) and activation functions ensuring a deep propagation of the information carried by the input.
This ensures that the network output does not ‘forget’ the input information as the depth grows.
In this paper, we show that this has a direct impact on the NTK.
Contributions.
There have been few attempts to understand the large depth limit of the NTK regime
Xiao et al.
; Huang et al.
; Bietti and Bach (2021); however, none of these works have characterized the limiting NTK and more importantly the exact convergence rate of the NTK to this limiting kernel.
The closest work is Huang et al.
(2020) where the authors considered a scaled version of ResNet with ReLU and proved an upper bound on the convergence rate of order O ( polylog L L ) .
In this paper, we improve this result in many ways: we prove that the convergence to the limiting NTK happens with a rate of Θ(log ( L ) L − 1 ) for different architectures and activation functions.
Note that for NTK regime, the lower bound is more important as it ensures a sub-exponential convergence rate of the NTK to its trivial limiting kernel(e.g.
constant).
We also show that the large depth behaviour of the NTK is initialization-sensitive; in particular, we prove that for FFNNs, we obtain an exponential convergence rate if the initialization is not on the EOC (this is a generalization of Xiao et al.
(2020)), which is not the case with ResNet.
Neural Networks and Neural Tangent Kernel 2.1 Setup and notations
Consider a neural network model consisting of L layers of widths ( n l ) 1 ≤ l ≤ L , n 0 =
d is the input dimension, and let θ =
( θ l ) 1 ≤ l ≤ L be the ﬂattened vector of weights and bias indexed by the layer’s index, and p be the dimension of θ .
The output f of the neural network is given by some mapping s : R n L → R o of the last layer y L ( x ) ; o being the dimension of the output (e.g. number of classes for a classiﬁcation problem).
For any input x ∈ R d , we thus have f ( x, θ )
=
s ( y L ( x ))
∈ R o .
We denote by θ t the value of θ at training time (step) t and f t ( x ) =
f
( x, θ t ) .
Let D =
( x i , z i ) 1 ≤ i ≤ N be the dataset, and let X = ( x i ) 1 ≤ i ≤ N , Z = ( z j ) 1 ≤ j ≤ N be the sequences of inputs and outputs respectively.
We assume that there is no collinearity in the input dataset X , i.e. for all x, x (cid:48) ∈ X and α ∈ R , x (cid:48) (cid:54) =
αx .
We also assume that X ⊂ E where E ⊂ R d is a compact set.
The NTK is deﬁned as the o × o dimensional kernel satisfying for all x, x
(cid:48) ∈ R d
[Equation]
The NTK regime (Inﬁnite width).
For a fully connected feedforward neural network, Jacot et al. (2018) proved that K Lθ t converges pointwise to a kernel K L (depends only on L ) for all t < T when min { n 1 , n 2 , ..., n L } → ∞ , where T is a constant.
In this limit, for the quadratic loss, f t is given by a simple linear model
[Equation]
where ˆ K L = K L ( X , X ) and γ ( x, X ) =
K L ( x, X )( ˆ K L ) − 1 .
Hereafter, we refer to f ∞ by the "NTK regime".
For the cross-entropy loss, Lee et al.
(2019) introduced some approximations to obtain the NTK regime.
These approximations are implemented in Novak et al.
Scale invariance.
f ∞ is scale invariant in the sense that it does not change if we scale the kernel by some depth dependent scalar since for any a L > 0 ,
[Equation]
Thus, studying the NTK regime with kernel K L is equivalent to studying the NTK regime with any scaled kernel
K L /a L .
In Theorems 1 and 2, we study scaled kernels to mitigate an exploding kernel effect in the limit of large depth, as the NTK regime solution remains unchanged.
Generalization in the NTK regime.
As observed in Du et al.
(2019), the convergence rate (w.r.t time) of f t to f ∞ (inﬁnite training time) is given by the smallest eigenvalue of ˆ K L .
If ˆ K L becomes singular in the large depth limit, then the performance of NTK regime decreases and might even be trivial.
Notice also that we can write f t ( x ) − f 0 ( x ) =
(cid:80) Ni
=1 a i K L ( x i , x ) for some a 1 , ..., a N ∈ R , i.e. the ‘residual’ f t − f 0 belongs to the
Reproducing Kernel Hilbert space of the K L .
Asymptotic Neural Tangent Kernel regime
In this section, we characterize the behaviour of K L as L goes to ∞ .
We prove that K L converges to a kernel K ∞ (which is trivial) with an initialization-and-architecture-dependent convergence rate.
3.1 Deep NTK of a FeedForward Neural Network (FFNN)
Consider an FFNN of depth L , widths ( n l ) 1 ≤ l ≤ L , weights w l and bias b l .
For some input x ∈ R d , the forward propagation using the NTK parameterization (similar to Jacot et al.
(2018)) is given by
[Equation]
We initialize the model randomly with w lij , b li iid ∼ N (0 , 1) , where N ( µ, σ 2 ) denotes the normal distribution of mean µ and variance σ 2 .
In the limit of inﬁnite width, the neurons ( y li ( . ))
i,l converge to Gaussian processes (Neal, 1995; Lee et al., 2018; Matthews et al., 2018;
Hayou et al., 2019; Schoenholz et al., 2017).
Hereafter, we denote by q l ( x, x (cid:48) )
the covariance between y l 1 ( x ) and y l 1 ( x (cid:48) )
( y l 1 can be replaced by any y li since ( y li )
i are i.i.d.
See appendix 4 for a comprehensive review of the signal propagation theory).
We deﬁne the correlation c l ( x, x (cid:48) ) .
For the ﬁrst layer, we have q 1 ( x, x (cid:48) )
=
σ 2 b + σ 2 w d x · x (cid:48) .
For (cid:15) ∈ (0 , 1) , we deﬁne the set B (cid:15) by: FFNN : B (cid:15) = { ( x, x (cid:48) ) ∈ R d
[Equation]
( x, x (cid:48) ) ≤ 1 − (cid:15) } , and we assume that there exists (cid:15) > 0 , such that for all x (cid:54) =
x (cid:48) ∈ X , ( x, x (cid:48) )
∈ B (cid:15) .
Edge of Chaos (EOC).
Given an input x , we denote by q l ( x ) the variance of y l ( x ) .
The asymptotic behaviour of q l ( x ) was studied in Lee et al.
(2018), Schoenholz et al.
(2017), and Hayou et al.
Under general regularity conditions, q l ( x ) converges to a point q ( σ b , σ w ) > 0 independent of x .
The asymptotic behaviour of the correlation c l ( x, x (cid:48) ) between y l ( x ) and y l ( x (cid:48) ) for any two inputs x and x (cid:48) is driven by the choice of ( σ b , σ w ) ; Schoenholz et al.
(2017) showed that if σ 2 w E [ φ (cid:48) ( (cid:112) q ( σ b , σ w ) Z ) 2 ] < 1 , where Z ∼ N (0 , 1) , then c l ( x, x (cid:48) ) converges to 1 exponentially quickly; this is called the ordered phase.
If σ 2 w E [ φ (cid:48) ( (cid:112) q ( σ b , σ w ) Z ) 2 ] >then c l ( x, x (cid:48) ) converges to c < 1 , which is then referred to as the chaotic phase.
The authors deﬁne the EOC as the set of parameters ( σ b , σ w ) that satisfy σ 2 w E [ φ (cid:48) ( (cid:112) q ( σ b , σ w ) Z ) 2 ] = 1 .
The EOC was studied in (Hayou et al., 2019) where the authors showed that the correlation converges to 1 with a polynomial rate (see Section 4 in the appendix).
The following proposition establishes that any initialization on the Ordered or Chaotic phase, leads to a trivial limiting NTK as L becomes large.
Proposition 1 (NTK with Ordered/Chaotic Initialization) .
Let ( σ b , σ w ) be either in the ordered or in the chaotic phase.
Then, there exist λ > 0 such that for all (cid:15) ∈ (0 , 1) , there exists γ > 0 such that
[Equation]
The proof of Proposition 1 relies on the asymptotic analysis of the second moment of the gradient.
We refer the reader to Section 6 in the appendix for more details.
Proposition 1 show that K L becomes trivial exponentially quickly w.r.t deph.
In this case, the NTK regime yields trivial performance, i.e. no better than that of a random classiﬁer.
Empirically, we ﬁnd that with depth L = 30 , the NTK training fails when the network is initialized on the Ordered phase ??????????
(Section ?? ).
In the next theorem, we show that the NTK explodes in the limit of large depth when the network is initialized on the EOC.
Leveraging our remark on the scale invariance property of the NTK (see Eq.
, we show that a scaled version of the kernel converges with a polynomial rate to the degenerate kernel, meaning that the inﬁnite-depth NTK regime is also trivial in this case, although the convergence is much slower.
The notation g ( x ) =
Θ( m ( x )) means there exist two constants A, B >
such that A m ( x ) ≤ g ( x ) ≤ B m ( x ) .
Theorem 1 (NTK on the EOC) .
Let ( σ b , σ w ) be on the EOC and ˜ K L = K L /L .
We have that
[Equation]
Moreover, there exists λ ∈ (0 , 1) such that for all (cid:15) ∈ (
[Equation]
[Equation] where, • ˜ K ∞ ( x, x (cid:48) )
=
σ 2 w (cid:107) x
(cid:107) (cid:107) x (cid:48) (cid:107) d (1 − (1 − λ )
1 x (cid:54) =
x (cid:48) ) with λ = 1 / 4 , for φ = ReLU
• ˜ K ∞ ( x, x (cid:48) )
=
q (1 − (1 − λ ) 1 x (cid:54) =
x (cid:48) ) where q > 0 is a constant and λ = 1 / 3 , for φ =
Tanh
We refer the reader to Section 1 in the appendix for the proof details.
Theorem 1 shows that the EOC initialization yields a polynomial convergence rate (w.r.t L ) of ˜ K L to ˜ K ∞ .
This is important knowing that ˜ K ∞ is trivial and brings hardly any information on x .
Indeed, the convergence rate of ˜ K L to ˜ K ∞ is Θ(log( L ) L − 1 ) .
This means that as L grows, the kernel ˜ K with EOC is still much further from the trivial kernel ˜ K ∞ compared to the Ordered/Chaotic initialization.
Thus, the EOC alleviates the curse of depth for NTK regime .
However, as shown in table 1, NTK regime fails for very deep networks ( L = 300 ).
3.2 Residual Neural Networks (ResNet)
Another important feature of DNNs, which is known to be highly inﬂuential, is their architecture.
For residual networks, the next theorem shows that for any σ w > 0 , the NTK of a ResNet explodes (exponentially) as L grows.
However, a normalized version ¯ K L = K L /α L of the NTK of a ResNet will always have a polynomial convergence rate to a limiting trivial kernel. [Table]
Theorem 2 (NTK for ResNet) .
Consider a ResNet satisfying
[Equation]
where F is a dense layer (Eq.
(3) ) with ReLU activation.
Let K Lres be the corresponding NTK, and ¯ K Lres =
K Lres /α L (Normalized NTK) with α L = L (1 + σ 2 w 2 )
L − 1 .
Then, we have
[Equation]
Moreover, there exists a constant λ ∈ (0 , 1) such that for all (cid:15) ∈ (0 , 1)
[Equation]
where ¯ K
[Equation]
The proof techniques used in Theorem 2 are similar to those used in the proof of theorem 1.
Details are provided in the appendix.
Theorem 2 shows that the NTK of a ReLU ResNet explodes exponentially w.r.t L .
However, the normalized kernel ¯ K Lres =
K Lres ( x, x (cid:48) ) /α
L converges to a limiting kernel ¯ K ∞ res at the exact polynomial rate Θ(log( L ) L − 1 ) for all σ w > 0 .
This suggests that ResNet act by default as an FFNN that is initialized on the EOC.
However, ¯ K Lres converges to a trivial kernel, which means that, even with ResNet, the performance of the NTK regime will decrease as we increase the depth, although it happens with a polynomial rate.
Table 1 shows the performance of the NTK regime versus standard SGD training on CIFAR10.
While the NTK regime fails with L = 300 for both
Ordered/EOC initializations
, it yields non-trivial performance when initialized on the EOC with L = 30 , which is not the case with an Ordered phase.
With ResNet, the performance is similar for both initializations which conﬁrms the results of theorem 2.
However, for all initializations schemes, NTK regime fails for L = 300 while standard SGD training succeeds.
We now leverage the previous results to obtain the asymptotic behaviour of the spectrum of the kernels studied in Theorems 1, 2 and Proposition 1, on the unit sphere S d − 1 = { x ∈ R
d : (cid:107) x (cid:107) 2 = 1 } .
On the sphere S d − 1 , all of these kernels (namely K L for FFNN on the Ordered/Chaotic phase, ˜ K L for FFNN/CNN on the EOC, and ¯ K Lres for ResNets) are dot-product kernels, i.e. for any of these kernels, denoted by κ L , there exists a function g L such that κ L ( x, x (cid:48) )
=
g L ( x · x (cid:48) ) for all x, x (cid:48)
∈ S d − 1 (we refer the reader to appendix 4 for more details).
This type of kernels is known to be diagonalizable on the sphere S d − 1 and its eigenfunctions are the so-called Spherical Harmonics of S d − 1 .
Many concurrent results have observed this fact Geifman et al.
(2020); Cao et al.
Bietti and Bach (2021).
In the next proposition, we leverage the results of Section 3 to study the aforementioned kernels from a spectral perspective.
Proposition 2 (Spectral decomposition on S d − 1 ) .
Let κ L be either, K L for an FFNN with L layers initialized on the Ordered phase (Proposition 1),
˜ K L for an FFNN with L layers initialized on the EOC (Theorem 1), or ¯ K Lres for a ResNet with L Fully Connected layers (Theorem 2).
Then, for all L ≥ 1 , there exists ( µ Lk )
k ≥ such that for all x, x (cid:48)
∈ S d −κ L ( x, x (cid:48) )
= (cid:88) k ≥µ
Lk
N ( d,k ) (cid:88) j =
Y k,j ( x ) Y k,j ( x (cid:48) ) .
[Figure]
Moreover, we have that 0 < µ ∞ 0 =
lim L →∞ µ L 0 < ∞ , and for all k ≥ 1 , lim L →∞ µ Lk = 0 .
The proof of Proposition 2 is based on a result from spectral theory analysis.
The limiting eigenvalues are obtained by a simple application of the dominated convergence theorem.
Proposition 2 shows that in the limit of large L , the kernel κ L becomes close to the trivial kernel κ ∞ ( x, x (cid:48) )
(cid:55)→ µ ∞
Y 0 , 0 ( x ) Y 0 , 0 ( x (cid:48) ) , where Y 0 , 0 is the constant function in the spherical harmonics class.
Therefore, in the limit of inﬁnite depth, the RKHS of the kernel κ L is reduced to the space of constant functions, conﬁrming that
the
NTK regime is trivial in this limit (recall that f ∞ − f 0 is in the RKHS of κ L ).
Fig 1 illustrates this deterioration of the spectrum as the depth increases.
Notice that with EOC, the deterioration happens with a much slower rate, which is expected from theorems 1 and 2.
Conclusion and Limitations
In this paper, we have shown that the inﬁnite depth limit of the NTK regime is trivial and cannot explain the performance of DNNs.
However, this convergence is initialization dependent.
These ﬁndings add to a recent line of research which shows that the inﬁnite width approximation of the NTK does not fully capture the training dynamics of DNNs (Chizat and Bach, 2018; Ghorbani et al., 2019; Huang and Yau, 2020; Hanin and Nica, 2019).
Appendix Setup and notations
In the appendix, we prove more general results than those presented in the main text (we have chosen to present partial results for the sake of simplicity).
The reader will ﬁnd versions of Propositions 1, 2, and Theorems 1, 2, that include Convolutional Neural Networks.
0.1 Neural Tangent Kernel
Consider a neural network model consisting of L layers ( y l ) 1 ≤ l ≤ L , with y l : R n l − 1 → R n l , n 0 = d and let θ =
( θ l ) 1 ≤ l ≤ L be the ﬂattened vector of weights and bias indexed by the layer’s index and p be the dimension of θ .
Recall that θ l has dimension n l + 1 .
The output f of the neural network is given by some transformation s : R n L → R o of the last layer y L ( x ) ; o being the dimension of the output (e.g. number of classes for a classiﬁcation problem).
For any input x ∈ R d , we thus have f ( x, θ )
=
s ( y L ( x ))
∈ R o .
As we train the model, θ changes with time t
and we denote by θ t the value of θ at time t and f t ( x ) =
f
( x, θ t )
=
( f j ( x, θ t ) , j ≤ o ) .
Let D =
( x i , z i ) 1 ≤ i ≤ N be the data set and let X = ( x i ) 1 ≤ i ≤ N , Z = ( z j ) 1 ≤ j ≤ N be the matrices of input and output respectively, with dimension d × N and o × N .
For any function g : R d × o
→ R k , k ≥ 1 , we denote by g ( X , Z ) the matrix ( g ( x i , z i ))
1 ≤ i ≤ N of dimension k × N .
Jacot et al.
(2018) studied the behaviour of the output of the neural network as a function of the training time t when the network is trained using a gradient descent algorithm.
Lee et al.
(2019) built on this result to linearize the training dynamics.
We recall hereafter some of these results.
For a given θ , the empirical loss is given by L ( θ ) = 1 N (cid:80) [Equation] [Equation] f
[Equation] i [Equation] [Equation]
The full batch GD algorithm is given by
[Equation]
where η > 0 is the learning rate.
Let T > 0 be the training time and N s = T/
η be the number of steps of the discrete GD (1).
The continuous time system equivalent to (1) with step ∆ t = η is given by
[Equation]
This differs from the result by Lee et al.
(2019) since we use a discretization step of ∆ t = η .
It is well known that this discretization scheme leads to an error of order O ( η ) (see Appendix).
Equation (2) can be re-written as
[Equation]
where ∇ θ f ( X , θ t ) is a matrix of dimension oN × p and ∇ z (cid:48) (cid:96) ( f ( X , θ t ) , Z ) is the ﬂattened vector of dimension oN constructed from the concatenation of the vectors ∇ z (cid:48)
(cid:96) ( z (cid:48) , z i ) | z (cid:48) =
f ( x i ,θ t ) , i ≤ N .
As a result, the output function f t ( x ) =
f
( x, θ t ) ∈ R o satisﬁes the following ODE
[Equation]
The Neural Tangent Kernel (NTK) K Lθ is deﬁned as the o × o dimensional kernel satisfying: for all x, x (cid:48) ∈ R d ,
[Equation]
We also deﬁne K Lθ t ( X , X ) as the oN × oN matrix deﬁned blockwise by
[Equation]
By applying (3) to the vector X , one obtains
[Equation]
meaning that for all j ≤ N
[Equation]
Inﬁnite width dynamics.
In the case of an FFNN, Jacot et al.
(2018) proved that, with GD, the kernel K Lθ t converges to a kernel K L which depends only on L (number of layers) for all t < T when n 1 , n 2 , ..., n L → ∞ , where T is an upper bound on the training time, under the technical assumption (cid:82) T 0 ||∇ z (cid:96) ( f t ( X , Z ))
||
2 dt < ∞
a.s. with respect to the initialization weights.
The inﬁnite width limit of the training dynamics is given by
[Equation]
We note hereafter ˆ K L = K L ( X , X ) .
As an example, with the quadratic loss (cid:96) (
[Equation]
6) is equivalent to
[Equation]
which is a simple linear model that has a closed-form solution given by
[Equation]
For general input x ∈ R d , we have
[Equation]
where γ ( x ) =
K L ( x, X ) K L ( X , X ) − 1 [-.0.
-].
0.2 Architectures
Let φ be the activation function.
We consider the following architectures (FFNN and CNN)
• FeedForward Fully-Connected Neural Network (FFNN) Consider an FFNN of depth L , widths ( n l ) 1 ≤ l ≤ L , weights w l and bias b l .
For some input x ∈ R d , the forward propagation using the NTK parameterization is given by
[Equation]
• Convolutional Neural Network (CNN/ConvNet) Consider a 1D convolutional neural network of depth L , denoting by [ m : n ] the set of integers { m, m +1 , ..., n } for n ≤ m , the forward propagation is given by
[Equation]
where i ∈ [1 : n l ] is the channel number, α ∈ [0 : M − 1] is the neuron location in the channel, n l is the number of channels in the l th layer, and M is the number of neurons in each channel, ker l
= [ − k : k ] is a ﬁlter with size 2 k +
and v l =
n l − 1 (2 k + 1) .
Here, w l ∈ R n l × n l − 1 × (2 k +1) .
We assume periodic boundary conditions, which results in having y li,α = y li,α + M = y li,α − M and similarly for l = 0 , x i,α + M
= x i,α = x
i,α − M 0 .
For the sake of simpliﬁcation, we consider only the case of 1D CNN, the generalization to a m D CNN for m ∈ N is straightforward.
Proof techniques
The techniques used in the proofs range from simple algebraic manipulation to tricky inequalities.
Lemmas 1, 2, 3, 4.
The proofs of these lemmas are simple and follow the same inductive argument as in the proof of the original NTK result in Jacot et al.
Note that these results can also be obtained by simple application of the Master Theorem in Yang (2020) using the framework of Tensor Programs.
Proposition 1, Theorems 1, 2.
The proof of these results follow two steps; Firstly, estimating the asymptotic behaviour of the NTK in the limit of large depth; secondly, controling these behaviour using upper/lower bounds.
We analyse the asymptotic behaviour of the NTK of FFNN using existing results on signal propagation in deep FFNN.
However, for CNNs, the dynamics are a bit trickier since they involve convolution operators; We use some results from the theory of Circulant Matrices for this purpose.
It is relatively easy to control the dynamics of the NTK in the Ordered/Chaotic phase, however, the dynamics become a bit complicated on the Edge of Chaos and technical lemmas which we call Appendix Lemmas are introduced for this purpose.
Proposition 2.
The spectral decomposition of zonal kernels on the sphere is a classical result in spectral theory which was recently applied to Neural Tangent Kernel Geifman et al.
(2020); Cao et al.
In order to prove the convergence of the eigenvalues, we use Dominated Convergence Theorem, leveraging the asymptotic results in Proposition 1 and Theorems 1, 2.
The inﬁnite width limit 2.1 Forward propagation
FeedForward Neural Network.
For some input x ∈ R d , the propagation of this input through the network is given by
[Equation]
Where φ : R → R is the activation function.
When we take the limit n l − 1 → ∞ recursively over l , this implies, using Central Limit Theorem, that y li ( x ) is a Gaussian variable for any input x .
This gives an error of order
O (1 / √ n l − 1 ) (standard Monte Carlo error).
More generally, an approximation of the random process y li ( . ) by a Gaussian process was ﬁrst proposed by Neal (1995) in the single layer case and has been extended to the multiple layer case by Lee et al.
(2018) and Matthews et al.
The limiting Gaussian process kernels follow a recursive formula given by, for any inputs x, x (cid:48) ∈ R d
[Equation]
where Ψ φ is a function that only depends on φ .
This provides a simple recursive formula for the computation of the kernel κ l ; see, e.g., Lee et al.
(2018) for more details.
Convolutional Neural Networks.
The inﬁnite width approximation with 1D CNN yields a recursion for the kernel.
However, the inﬁnite width here means inﬁnite number of channels, with a Monte Carlo error of O (1 / √ n l − 1 ) .
The kernel in this case depends on the choice of the neurons in the channel and is given by
[Equation]
so that
[Equation]
The convolutional kernel κ lα,α (cid:48) has the ‘self-averaging’ property; i.e. it is an average over the kernels corre- sponding to different combination of neurons in the previous layer.
However, it is easy to simplify the analysis in this case by studying the average kernel per channel deﬁned by ˆ κ l = 1 N 2 (cid:80) α,α (cid:48)
κ lα,α (cid:48) .
Indeed, by summing terms in the previous equation and using the fact that we use circular padding, we obtain
[Equation]
This expression is similar in nature to that of FFNN.
We will use this observation in the proofs.
Note that our analysis only requires the approximation that, in the inﬁnite width limit, for any two inputs x, x (cid:48) , the variables y li ( x ) and y li ( x (cid:48) ) are Gaussian with covariance κ l ( x, x (cid:48) ) for FFNN, and y li,α ( x ) and y li,α (cid:48)
( x (cid:48) ) are
Gaussian with covariance κ lα,α (cid:48) ( x, x (cid:48) ) for CNN.
We do not need the much stronger approximation that the process y li ( x ) ( y li,α ( x ) for CNN) is a Gaussian process.
Residual Neural Networks.
The inﬁnite width limit approximation for ResNet yields similar results with an additional residual terms.
It is straighforward to see that, in the case of a ResNet with FFNN-type layers, we have that
[Equation]
whereas for ResNet with CNN-type layers, we have that
[Equation]
2.2 Gradient Independence
In the mean-ﬁeld literature of DNNs, an omnipresent approximation in prior literature is that of the gradient independence which is similar in nature to the practice of feedback alignment (Lillicrap et al., 2016).
This approximation states that, for wide neural networks, the weights used for forward propagation are independent from those used for back-propagation.
When used for the computation of Neural Tangent Kernel, this approximation was proven to give the exact computation for standard architectures such as FFNN, CNN and ResNets Yang (2020)
(Theorem D.1).
This result has been extensively used in the literature as an approximation before being proved to yields exact computation for the NTK, and theoretical results derived under this approximation were veriﬁed empirically; see references below.
Gradient Covariance back-propagation.
Analytical formulas for gradient covariance back-propagation were derived using this result, in (Hayou et al., 2019; Schoenholz et al., 2017; Yang and Schoenholz, 2017b; Lee et al., 2018; Poole et al., 2016; Xiao et al., 2018; Yang, 2019).
Empirical results showed an excellent match for
FFNN in Schoenholz et al.
(2017), for Resnets in Yang (2019) and for CNN in Xiao et al.
Neural Tangent Kernel.
The Gradient Independence approximation was implicitly used in Jacot et al.
(2018) to derive the inﬁnite width Neural Tangent Kernel (See Jacot et al.
(2018), Appendix A.1).
Authors have found that this inﬁnite width NTK computed with the Gradient Independence approximation yields excellent match with empirical (exact) NTK.
We use this result in our proofs and we refer to it simply by the Gradient Independence.
Discussion on Assumption
Assumption 1.
We assume that for all x, x (cid:48) ∈ X , q 1 α,α (cid:48) ( x, x (cid:48) ) is independent of α, α (cid:48) .
Assumption 1 implies that, there exists some function e : ( x, x (cid:48) )
(cid:55)→ e ( x, x (cid:48) ) such that for all α, α (cid:48) , x, x (cid:48) (cid:88) j (cid:88)
[Table]
[Equation]
This system has N 2 M 2 equations and N × 2 n 0 × M variables.
Therefore, in the case n 0 >> 1 , the set of solutions S is large.
By using Assumption 1, we restrict our analysis to this case.
Hereafter, for all CNN analysis, for some function G and set E , taking the supremum sup ( x,x (cid:48) ) ∈
E G ( x, x (cid:48) ) should be interpreted as sup ( x,x (cid:48) )
∈
E ∩X 2 G ( x, x (cid:48) ) .
Another justiﬁcation to assumption 1 can be attributed a self-averaging property of the dynamics of the correlation inside a CNN.
We refer the reader to the proof of Appendix lemma 3 for more details.
Warmup: Results from the Mean-Field theory of DNNs 4.1 Notation
For FFNN layers, let q l ( x ) :
= q l ( x, x ) be the variance of y l 1 ( x ) (the choice of the index 1 is not important since, in the inﬁnite width limit, the random variables ( y li ( x )) i ∈ [1: N l ] are iid).
Let q l ( x, x (cid:48) ) , resp.
c l 1 ( x, x (cid:48) ) be the covariance, resp.
the correlation between y l 1 ( x ) and y l 1 ( x (cid:48) ) .
For Gradient back-propagation, let ˜ q l ( x, x (cid:48) ) be the Gradient covariance deﬁned by ˜ q l ( x, x (cid:48) )
=
E (cid:104) ∂
L ∂y l 1 ( x ) ∂
L ∂y l 1 ( x (cid:48) )
(cid:105) where L is some loss function.
Similarly, let ˜ q l ( x ) be the Gradient variance at point x .
We also deﬁne ˙
q l ( x, x (cid:48) )
=
σ 2 w E [ φ (cid:48)
( y l − 1 1 ( x )) φ (cid:48)
( y l − 1 1 ( x (cid:48) ))] .
For CNN layers, we use similar notation across channels.
Let q lα ( x ) be the variance of y l 1 ,α, ( x ) (the choice of the index 1 is not important here either since, in the limit of inﬁnite number of channels, the random variables ( y li,α ( x ))
i ∈ [1: N l ] are iid).
Let q lα,α (cid:48)
( x, x (cid:48) )
the covariance between y l 1 ,α ( x ) and y l 1 ,α (cid:48)
( x (cid:48) ) , and c lα,α (cid:48) ( x, x (cid:48) )
the corresponding correlation.
We also deﬁne the pseudo-covariance ˆ q lα,α (cid:48)
( x, x (cid:48) )
=
σ 2 b + σ 2 w E [ φ ( y l − 1 1 ,α ( x )) φ ( y l − 1 1 ,α (cid:48)
( x (cid:48) ))] and ˙ q lα,α (cid:48) ( x, x (cid:48) )
=
σ 2 w E [ φ ( y l − 1 1 ,α ( x ))
φ ( y l − 1 1 ,α (cid:48) ( x (cid:48) ))] .
The Gradient covariance is deﬁned by ˜ q lα,
[Equation]
4.1.1 Covariance propagation
Covariance propagation for FFNN.
In Section 2.1, we derived the covariance kernel propagation in an
FFNN.
For two inputs x, x (cid:48) ∈ R d , we have
[Equation]
this can be written as
[Equation]
with c l − 1 :=
c l − 1 ( x, x (cid:48) ) .
With ReLU, and since ReLU is positively homogeneous (i.e. φ ( λx ) = λφ ( x ) for λ ≥ 0 ), we have that
[Equation]
where f is the ReLU correlation function given by Hayou et al.
[Equation]
Covariance propagation for CNN.
The only difference with FFNN is that the independence is across channels and not neurons.
Simple calculus yields
[Equation]
Observe that
[Equation]
With ReLU [Equation] we have
[Equation]
Covariance propagation for ResNet with ReLU.
In the case of ResNet, only an added residual term shows up in the recursive formula.
For a ResNet with FFNN layers
[Equation]
the recursion reads
[Equation]
with CNN layers, we have instead
[Equation]
4.1.2 Gradient Covariance back-propagation
Gradient back-propagation for FFNN.
The gradient back-propagation is given by
[Equation]
where L is some loss function.
Using the Gradient Independence 2.2, we have as in Schoenholz et al.
[Equation]
where χ ( q l ( x ))
=
σ 2 w E [ φ ( (cid:112) q l ( x ) Z ) 2 ] .
Gradient Covariance back-propagation for CNN.
We have that
[Equation]
Moreover,
[Equation]
Using the Gradient Independence 2.2, and taking the average over the number of channels we have that
[Equation]
We can get similar recursion to that of the FFNN case by summing over α and using the periodic boundary condition, this yields
[Equation]
4.1.3 Edge of Chaos (EOC)
Let x ∈ R
d be an input.
The convergence of q l ( x ) as l increases has been studied by Schoenholz et al.
(2017) and Hayou et al.
In particular, under weak regularity conditions, it is proven that q l ( x ) converges to a point q ( σ b , σ w ) > 0 independent of x as l → ∞ .
The asymptotic behaviour of the correlations c l ( x, x (cid:48) ) between y l ( x ) and y l ( x (cid:48) ) for any two inputs x and x (cid:48) is also driven by ( σ b , σ w ) : the dynamics of c l is controlled by a function f i.e. c l +1 = f
( c l ) called the correlation function.
The authors deﬁne the EOC as the set of parameters ( σ b , σ w ) such that σ 2 w E [ φ (cid:48) ( (cid:112) q ( σ b , σ w ) Z ) 2 ] = 1 where Z ∼ N (0 , 1) .
Similarly the Ordered, resp.
Chaotic, phase is deﬁned by σ 2 w E [ φ (cid:48) ( (cid:112) q ( σ b , σ w ) Z ) 2 ] < 1 , resp.
σ 2 w E [ φ (cid:48) ( (cid:112) q ( σ b , σ w ) Z ) 2 ] > 1 .
On the
Ordered phase, the gradient will vanish as it backpropagates through the network, and the correlation c l ( x, x (cid:48) ) converges exponentially to 1 .
Hence the output function becomes constant (hence the name ’Ordered phase’).
On the Chaotic phase, the gradient explodes and the correlation converges exponentially to some limiting value c < 1 which results in the output function being discontinuous everywhere (hence the ’Chaotic’ phase name).
On the EOC, the second moment of the gradient remains constant throughout the backpropagation and the correlation converges to 1 at a sub-exponential rate, which allows deeper information propagation.
Hereafter, f will always refer to the correlation function .
We initialize the model with w lij , b li iid ∼ N (0 , 1) , where N ( µ, σ 2 ) denotes the normal distribution of mean µ and variance σ 2 .
In the remainder of this appendix, we assume that the following conditions are satisﬁed
• The input data is a subset of a compact set E of R d , and no two inputs are co-linear.
• All calculations are done in the limit of inﬁnitely wide networks.
4.2 Some results from the information propagation theory Results for FFNN with Tanh activation.
Fact 1.
For any choice of σ b , σ w ∈ R + , there exist q, λ > 0 such that for all l ≥ 1 , sup x ∈ R −
≤ e −
Fact 2.
On the Ordered phase, there exists γ > 0 such that sup x,x (cid:48) ∈ R d | c l ( x, x (cid:48) )
− 1 | ≤ e − γl .
(Equation (8) in Schoenholz et al.
Fact 3.
Let ( σ b , σ w ) ∈ EOC.
Using the same notation as in fact 4, we have that sup ( x,x (cid:48) )
∈ B (cid:15) |
1 − c l ( x, x (cid:48) ) | =
O ( l − 1 ) .
(Proposition 3 in Hayou et al.
Fact 4.
Let B (cid:15) = {
( x, x (cid:48) ) ∈ R
d : c 1 ( x, x (cid:48) ) < 1 − (cid:15) } .
On the chaotic phase, there exist c < 1 such that for all (cid:15) ∈ (0 , 1) , there exists γ > 0 such that sup ( x,x (cid:48) )
∈ B (cid:15)
|
c l ( x, x (cid:48) ) − c | ≤ e − γl .
(Equations (8) and (9) in Schoenholz et al.
Fact (Correlation function) .
The correlation function f is deﬁned by f ( x ) =
[Equation]
q where q is given in Fact 1 and Z 1 , Z 2 are iid standard Gaussian variables.
Fact 6.
f has a derivative of any order j ≥ 1 given by
[Equation]
As a result, we have that f ( j ) (1) =
σ 2 w q j − 1 E [ φ ( j ) ( Z 1 ) 2 ] > 0 for all j ≥ 1 .
The proof of the previous fact is straightforward following the same integration by parts technique as in the proof of Lemma 1 in Hayou et al.
The result follows by induction.
Fact 7.
Let ( σ b , σ w ) ∈ EOC.
We have that f (cid:48) (1) =
1 (by deﬁnition of EOC).
As a result, the Taylor expansion of f near 1 is given by
[Equation]
where α, ζ > 0 .
Proof.
The proof is straightforward using fact 6, and integral-derivative interchanging.
Results for FFNN with ReLU activation.
Fact 8.
The ordered phase for ReLU is given by Ord = { ( σ b , σ w ) ∈ ( R + ) 2 : σ w < √ 2 } .
Moreover, for any ( σ b , σ w ) ∈ Ord , there exist λ such that for all l ≥ 1 , sup x ∈ R d | q l ( x, x ) − q | ≤
[Equation]
where q =
[Equation]
The proof is straightforward using equation (12).
Fact 9.
For any ( σ b , σ w ) in the Ordered phase, there exist λ such that for all l ≥ 1 , sup ( x,x (cid:48) )
∈ R d | c l ( x, x (cid:48) ) − 1 | ≤ e − λl .
The proof of this claim follows from standard Banach Fixed point theorem in the same fashion as for Tanh in Schoenholz et al.
Fact 10.
The Chaotic phase for ReLU is given by Ch
= { ( σ b , σ w ) ∈
( R + ) 2 : σ w > √ 2 } .
Moreover, for any ( σ b , σ w ) ∈
Ch , for all l ≥ 1 , x ∈ R d , q l ( x, x ) (cid:38) ( σ 2 w / 2) l .
The variance explodes exponentially on the Chaotic phase, which means the output of the Neural Network can grow arbitrarily in this setting.
Hereafter, when no activation function is mentioned, and when we choose " ( σ b , σ w ) on the Ordered/Chaotic phase", it should be interpreted as " ( σ b , σ w ) on the Ordered phase" for ReLU and " ( σ b , σ w ) on the Ordered/Chaotic phase" for Tanh.
Fact 11.
For ReLU FFNN on the EOC, we have that q l ( x, x ) =
σ 2 w d || x
|| 2 for all l ≥ 1 .
The proof is straightforward using equation 12 and that ( σ b , σ w ) =
(0 , √ 2) on the EOC.
Fact 12.
The EOC of ReLU is given by the singleton { ( σ b , σ w ) =
(0 , √ 2) } .
In this case, the correlation function of an FFNN with ReLU is given by
[Equation]
(Proof of Proposition 1 in Hayou et al.
Fact 13.
Let ( σ b , σ w ) ∈ EOC.
Using the same notation as in fact 4, we have that
[Equation]
(Follows straightforwardly from Proposition 1 in Hayou et al.
Fact 14.
We have that
[Equation]
16) with s = 2 √ 2 3 π and b = √ 2 30 π .
This result was proven in Hayou et al.
Results for CNN with Tanh activation function.
Fact 15.
For any choice of σ b , σ w ∈ R + , there exist q, λ > 0 such that for all l ≥ 1 , sup α,α (cid:48) sup x ∈ R d | q lα,
(cid:48) ( x, x ) − q | ≤ e λl .
(Equation Xiao
and variance convergence
The behaviour of the correlation c lα,α (cid:48) ( x, x (cid:48) ) was studied in Xiao et al.
(2018) only in the case x (cid:48)
=
x .
We give a comprehensive analysis of the asymptotic behaviour of c lα,α (cid:48)
( x, x (cid:48) ) in the next section.
General results on the correlation function.
Fact 16.
Let f be either the correlation function of Tanh or ReLU.
We have that
• f (1) =
(Lemma 2 in Hayou et al.
• On the ordered phase 0 < f (cid:48) (1) < 1 (By deﬁnition).
• On the Chaotic phase f (cid:48) (1) > 1 (By deﬁnition).
• On the EOC, f (cid:48) (1) =
1 (By deﬁnition).
• On the Ordered phase and the EOC, 1 is the unique ﬁxed point of f (Hayou et al.
• On the Chaotic phase, f has two ﬁxed points, 1 which is unstable, and c < 1 which is a stable ﬁxed point Schoenholz et al.
Fact 17.
Let (cid:15) ∈ (0 , 1) .
On the Ordered/Chaotic phase, with either ReLU or Tanh, there exists α ∈ (0 , 1) , γ > 0 such that
[Equation]
Proof.
This result follows from a simple ﬁrst order expansion inequality.
For Tanh on the Ordered phase, we have that
[Equation]
where ζ l = sup
[Equation]
1) | f (cid:48)(cid:48) ( t ) | → | f [Equation] (1) | .
We conclude for Ordered phase with Tanh using fact 2.
The same argument can be used for Chaotic phase with Tanh using fact 4; in this case, α = f (cid:48)
( c ) where c is the unique stable ﬁxed point of the correlation function f .
On the Ordered phase with ReLU, let ˜ f be the correlation function.
It is easy to see that ˜ f (cid:48)
( c ) =
σ 2 w 2 f (cid:48) ( c ) where f is given in fact 12.
f (cid:48) ( x ) = 1 − √ 2 π (1 − x ) 1 / 2 + O ((1 − x ) 3 / 2 ) .
Therefore, there exists l 0 , ζ > 0 such that for l > l 0 ,
[Equation]
We conclude using fact 9.
Asymptotic behaviour of the correlation in FFNN.
Appendix Lemma 1 (Asymptotic behaviour of c l for ReLU) .
Let ( σ b , σ w ) ∈ EOC and (cid:15) ∈ (0 , 1) .
{+There exist universal constants κ, κ (cid:48) , κ (cid:48)(cid:48) > 
 (that do not depend on any parameter) such that+}
[Equation]
and,
[Equation]
Proof.
Let ( x, x (cid:48) )
∈ B (cid:15) and s = 2 √ 2 3 π .
From the preliminary results, we have that lim
l →∞ sup x,x (cid:48) ∈
R d 1 − c l ( x, x (cid:48) )
=
(fact 13).
Using fact 14, we have uniformly over B (cid:15) ,
[Equation]
where s, b > 0 , this yields
[Equation]
Thus, letting b (cid:48) = b 2 + 5 16 s 3 , as l goes to inﬁnity
[Equation]
and by summing and equivalence of positive divergent series
[Equation]
Moreover, since γ − 1 / [Equation] 1 = γ
[Equation]
) , using the same argument multiple times and inverting the formula yields
[Equation]
where κ = 9 π 2 2 . Note that, by Appendix Lemma 5 (section ?? ), the O bound can be chosen in a way that it does not depend on ( x, x (cid:48) ) , it depends only on (cid:15) ; this concludes the proof for the ﬁrst part of the result.
Using fact 12, we have that
[Equation]
Thus, it follows that
[Equation]
for some universal constant κ (cid:48)(cid:48) uniformly over the set B (cid:15) , which concludes the proof.
We prove a similar result for an FFNN with Tanh activation.
Appendix Lemma 2 (Asymptotic behaviour of c l for Tanh) .
Let ( σ b , σ w ) ∈ EOC and (cid:15) ∈ (0 , 1) .
We have
[Equation]
where κ =
[Equation]
and ζ = f 3 (1) > 0 .
Moreover, we have that
[Equation]
Proof.
Let ( x, x (cid:48) )
∈ B (cid:15) and λ l := 1 − c l ( x, x (cid:48) ) .
Using a Taylor expansion of f near 1 (fact 7), there exist α, ζ > 0 such that
[Equation]
Here also, we use the same technique as in the previous lemma.
We have that
[Equation]
By summing (divergent series), we have that λ − 1 l ∼ [Equation]
Therefore,
[Equation]
By summing a second time, we obtain
[Equation]
Using the same technique once again, we obtain
[Equation]
This yields
[Equation]
In a similar fashion to the previous proof, we can force the upper bound in O to be independent of x using Appendix Lemma 5.
This way, the bound depends only on (cid:15) .
This concludes the ﬁrst part of the proof. For the second part, observe that f (cid:48) ( x ) = 1 + ( x − 1) f (cid:48)(cid:48) (1) + O (( x − 1) 2 ) , hence
[Equation]
which concludes the proof.
4.3 Large depth behaviour of the correlation in CNNs
For CNNs, the inﬁnite width will always mean the limit of inﬁnite number of channels.
Recall that, by deﬁnition, ˆ
[Equation]
Unlike FFNN, neurons in the same channel are correlated since they share the same ﬁlters.
Let x, x (cid:48) be two inputs and α, α (cid:48) two nodes in the same channel i .
Using Central Limit Theorem in the limit of large n l (number of channels), we have
[Equation]
Let c lα,α (cid:48)
( x, x (cid:48) ) be the corresponding correlation.
Since q lα,α ( x, x ) converges exponentially to q which depends neither on x nor on α , the mean-ﬁeld correlation as in Schoenholz et al.
; Hayou et al.
(2019) is given by
[Equation]
where
[Equation]
and Z 1 , Z 2 are independent standard normal variables.
The dynamics of c lα,α (cid:48) become similar to those of c l in an FFNN under assumption 1.
We show this in the proof of Appendix Lemma 3.
In Xiao et al.
, authors studied only the limiting behaviour of correlations c lα,α (cid:48)
( x, x ) (same input x ), however, they do not study c lα,α (cid:48)
( x, x (cid:48) )
when x (cid:54) = x (cid:48) .
We do this in the following Lemma, which will prove also useful for the main results of the paper.
Appendix Lemma 3 (Asymptotic behaviour of the correlation in CNN with Tanh) .
We consider a CNN with Tanh activation function.
Let ( σ b , σ w ) ∈ ( R + ) 2 and (cid:15) ∈ (0 , 1) .
Let B (cid:15) =
{ ( x, x (cid:48) ) ∈ R d : sup α,α (cid:48) c
1 α,α (cid:48)
( x, x (cid:48) ) < 1 − (cid:15) } .
The following statements hold
If ( σ b , σ w ) are on the Ordered phase, then there exists β > 0 such that
[Equation]
If ( σ b , σ w ) are on the Chaotic phase, then for all (cid:15) > 0 there exists β > 0 and c ∈ (0 , 1) such that sup
[Equation]
Under Assumption 1, if ( σ b , σ w ) ∈ EOC , then we have
[Equation]
where κ =
[Equation]
[Equation]
> 0 , and f is the correlation function given in Fact 5.
We prove statements 1 and 2 for general inputs, i.e. without using Assumption 1.
The third statement requires Assumption 1.
Proof.
Let ( x, x (cid:48) )
∈ R d .
Without using assumption 1, we have that
[Equation]
Writing this in matrix form yields
[Equation]
where C l = (( c lα,α + β ( x, x (cid:48) ))
α ∈ [0: N − 1] ) β ∈ [0: N − 1] is a vector in R N 2 , U is a convolution matrix and f is applied element-wise.
As an example, for k = 1 , U is given by
U =      ...

 
For general k , U is a Circulant symmetric matrix with eigenvalues λ 1 > λ 2 ≥ λ 3 ... ≥ λ N 2 .
The largest eigenvalue of U is given by λ 1 = 2 k +
and its equivalent eigenspace is generated by the vector e 1 = 1 N (1 , 1 , ..., 1) ∈ R N 2 .
This yields
[Equation]
where β = log( λ 1 λ 2 ) .
This provides another justiﬁcation to Assumption 1; as l grows, and assuming that C l → e 1 (which we show in the remainder of this proof), C l exhibits a self-averaging property since C l ≈ [Equation]
[Equation]
− 1 . This system concentrates around the average value of the entries of C l as l grows.
Since the variances converge to a constant q as l goes to inﬁnity (fact 15), this approximation implies that the entries of C l become almost equal as l goes to inﬁnity, thus making assumption 1 almost satisﬁed in deep layers.
Let us now prove the statements.
Let ( σ b , σ w ) be in the Ordered phase, ( x, x (cid:48) )
∈ R d and c lm = min
α,α (cid:48)
c lα,α (cid:48) ( x, x (cid:48) ) .
Using the fact that f is non-decreasing, we have that c lα,α (cid:48)
( x
[Equation]
x
[Equation]
Taking the minimum again over α, α (cid:48) , we have c lm ≥ c l − 1 m , therefore c lm is non-decreasing and converges to the unique ﬁxed point of f which is c = 1 .
This proves that sup α,α (cid:48) | c lα,α (cid:48) ( x, x (cid:48) ) − 1 | → 0 .
Moreover, the convergence rate is exponential using the fact that (fact 16) 0 < f (cid:48) (1) < 1 .
To see this, observe that
[Equation]
Knowing that sup ζ ∈ [
[Equation] f [Equation] ( ζ ) → f (cid:48) (1) < 1 , we conclude.
Moreover, the convergence is uniform in ( x, x (cid:48) ) since the convergence rate depends only on f (cid:48) (1) .
2. Let (cid:15) ∈ (0 , 1) .
In the chaotic phase, the only difference is the limit c
=
c 1 <
and the Supremum is taken over B (cid:15) to avoid points where c 1 ( x, x (cid:48) )
= 1 .
In the Chaotic phase (fact 16), f has two ﬁxed points, 1 is an unstable ﬁxed point and c 1 ∈ (0 , 1) which is the unique stable ﬁxed point.
We conclude by following the same argument.
Let (cid:15) ∈ (0 , 1) and ( σ b , σ w ) ∈ EOC.
Using the same argument of monotony as in the previous cases and that f has 1 as unique ﬁxed point, we have that lim l →∞ sup x
[Equation]
α [Equation] −
[Equation]
x
[Equation]
) | = 0 .
From fact 7, the Taylor expansion of f near 1 is given by
[Equation]
where α = f (cid:48)(cid:48) (1) and ζ = f (3) (1) .
Using fact 6, we know that f ( k ) (1) =
σ 2 w q k − 1 E
[ φ ( k ) ( √ qZ ) 2 ] .
Therefore, we have α > 0 , and ζ < 0 .
Under assumption 1, it is straightforward that for all α, α (cid:48) , and l ≥c lα,α (cid:48)
( x, x (cid:48) )
=
c l ( x, x (cid:48) )
i.e. c lα,α (cid:48) are equal for all α, α (cid:48) .
The dynamics of c l ( x, x (cid:48) ) are exactly the dynamics of the correlation in an FFNN.
We conclude using Appendix Lemma 2.
It is straightforward that the previous Appendix Lemma extend to ReLU activation, with slightly different dynamics.
In this case, we use Appendix Lemma 1 to conclude for the third statement.
Appendix Lemma 4 (Asymptotic behaviour of the correlation in CNN with ReLU-like activation functions) .
We consider a CNN with ReLU activation.
Let ( σ b , σ w ) ∈ ( R + ) 2 .
Let ( σ b , σ w ) ∈
( R + ) 2 and (cid:15) ∈ (0 , 1) .
The following statements hold
If ( σ b , σ w ) are on the Ordered phase, then there exists β > 0 such that
[Equation]
If ( σ b , σ w ) are on the Chaotic phase, then there exists β > 0 and c ∈ (0 , 1) such that
[Equation]
Under Assumption 1, if ( σ b , σ w ) ∈ EOC , then
[Equation]
where κ, κ (cid:48) > 0 are universal constants.
Proof.
The proof is similar to the case of Tanh in Appendix Lemma 3.
The only difference is that we use Appendix Lemma 1 to conclude for the third statement.
A technical tool for the derivation of uniform bounds
Results in Theorem 1 and 2 and Proposition 1 involve a supremum over the set B (cid:15) .
To obtain such results, we need a ’uniform’ Taylor analysis of the correlation c l ( x, x (cid:48) )
(see the next section) where uniformity is over ( x, x (cid:48) )
∈ B (cid:15) .
It turns out that such result is trivial when the correlation follows a dynamical system that is controlled by a non-decreasing function.
We clarify this in the next lemma.
Appendix Lemma 5 (Uniform Bounds) .
Let A ⊂ R be a compact set and g a non-decreasing function on A .
Deﬁne the sequence ζ l by ζ l = g ( ζ l − 1 ) and ζ 0 ∈ A .
Assume that there exist α l , β l that do not depend on ζ l , with β l
= o ( α l ) , such that for all ζ 0 ∈ A , ζ l = α l + O ζ 0 ( β l ) where O ζ 0 means that the O bound depends on ζ 0 .
Then, we have that
[Equation]
i.e. we can choose the bound O to be independent of ζ 0 .
Proof.
Let ζ 0 ,m = min A and ζ 0 ,M = max A .
Let ( ζ m,l ) and ( ζ M,l ) be the corresponding sequences.
Since g is non-decreasing, we have that for all ζ 0 ∈ A , ζ m,l ≤ ζ l ≤ ζ M,l .
Moreover, by assumption, there exists
M 1 , M 2 > 0 such that [Equation]
[Equation]
and
[Equation]
ζ
[Equation]
[Equation]
therefore,
[Equation]
which concludes the proof.
Note that Appendix Lemma 5 can be easily extended to Taylor expansions with ‘ o ’ instead of ‘ O ’.
We will use this result in the proofs, by refereeing to Appendix Lemma 5.
Proofs of Section 3: Large Depth Behaviour of Neural Tangent Kernel 6.1 Proofs of the results of Section 3.
In this section, we provide proofs for the results of Section 3.1 in the paper.
Recall that Lemma 1 in the paper is a generalization of Theorem 1 in Jacot et al.
and is reminded here.
The proof is simple and follows similar induction techniques as in Jacot et al.
Lemma 1 (Generalization of Th.
1 in Jacot et al.
Consider an FFNN of the form (3) .
Then, as n 1 , n 2 , ..., n L − 1 → ∞ , we have for all x, x (cid:48) ∈ R d , i, i (cid:48) ≤ n L , K Lii [Equation] ( x, x (cid:48) )
=
δ ii (cid:48)
K L ( x, x (cid:48) ) , where K L ( x, x (cid:48) ) is given by the recursive formula
[Equation]
where q l ( x, x (cid:48) )
=
σ 2 b + σ 2 w E [ φ ( y l − 1 1 ( x )) φ ( y l − 1 1 ( x (cid:48) ))] and ˙ q l ( x, x (cid:48) )
=
σ 2 w E [ φ (cid:48)
( y l − 1 1 ( x )) φ (cid:48)
( y l − 1 1 ( x (cid:48) ))] .
Proof.
The proof for general σ w is similar to when σ w = 1 (Jacot et al.
(2018)) which is a proof by induction.
For l ≥ 2 and i ∈ [1 : n l ]
[Equation]
Therefore,
[Equation]
[Table]
[Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation] [Table] [Equation] [Table]
[Equation]
[Table]
[Equation]
[Table] [Equation] [Table] Using the induction hypothesis, namely that as n 0 , n 1 , ..., n l
[Table]
1 → ∞ , for all j, j
[Table]
≤ n l and all x, x
[Table]
we then obtain for all n l , as n 0 , n 1 , ..., n l − 1 → ∞ σ 2 w n l
[Table]
and letting n l go to inﬁnity, the law of large numbers, implies that
[Equation]
Moreover, we have that
[Equation]
which ends the proof.
We now provide the recursive formula satisﬁed by the NTK of a CNN, namely Lemma 2 of the paper.
Lemma 2 (Inﬁnite width dynamics of the NTK of a CNN) .
Consider a CNN of the form ( ?? )
, then we have that for all x, x (cid:48) ∈ R d , i, i (cid:48) ≤ n
and α, α [Equation] ∈ [0 : M − 1]
[Equation]
For l ≥ 2 , as n 1 , n 2 , ..., n l − 1 → ∞ recursively, we have for all i, i (cid:48) ≤ n l , α, α (cid:48) ∈ [0 : M − 1] ,
K l ( i,α ) , ( i (cid:48) ,α (cid:48) )
( x, x (cid:48) )
=
δ ii (cid:48)
K lα,α (cid:48) ( x, x (cid:48) ) , where K lα,α (cid:48) is given by the recursive formula
[Equation]
where Ψ l − 1 α,α (cid:48)
= ˙
q lα,α (cid:48)
K l − 1 α,α (cid:48)
+ ˆ q lα,α (cid:48) , and ˆ q lα,α , ˙
q lα,α (cid:48) are deﬁned in Lemma 1, with y l − 1 1 ,α ( x ) , y l − 1 1 ,α (cid:48)
( x (cid:48) ) in place of y l − 1 1 ( x ) , y l − 1 1 ( x (cid:48) ) .
Proof.
Let x, x (cid:48) be two inputs.
We have that
[Equation]
therefore
[Equation]
Assume the result is true for l − 1 , let us prove it for l .
Let θ 1: l − 1 be model weights and bias in the layers 1 to l − 1 .
Let ∂ θ 1: l − 1 y li,α [Equation] x ) = ∂y li,α
[Equation]
x
[Equation]
[Equation]
We have that
[Equation]
this yields
[Equation]
as n 1 , n 2 , ..., n l − 2 → ∞ and using the induction hypothesis, we have
[Equation]
note that K l − [Equation]
j,α +
[Equation]
x, x [Equation]
K l [Equation] [Equation]
x [Equation] x [Equation]
) for all j since the variables are iid across the channel index j .
Now letting n l − 1 → ∞ , we have that
[Equation]
We conclude using the fact that
[Equation]
To alleviate notations, we use hereafter the notation K L for both the NTK of FFNN and CNN.
For FFNN, it represents the recursive kernel K L given by lemma 1, whereas for CNN, it represents the recursive kernel K Lα,α (cid:48) for any α, α (cid:48) , which means all results that follow are true for any α, α (cid:48) .
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 for all (cid:15) ∈ (0 , 1) , there exists γ > 0 such that
[Equation]
We will use the next lemma in the proof of proposition 1.
Appendix Lemma 6.
Let ( a l ) be a sequence of non-negative real numbers such that ∀ l ≥ 0 , a l +1 ≤ αa l + ke − βl , where α ∈ (0 , 1) and k, β > 0 .
Then there exists γ > 0 such that ∀ l ≥ 0 , a l ≤ e − γl .
Proof.
Using the inequality on a l , we can easily see that
[Equation]
where we divided the sum into two parts separated by index l/ 2 and upper-bounded each part.
The existence of γ is straightforward.
Now we prove Proposition[-
-]
Proof.
We prove the result for FFNN ﬁrst.
Let x, x (cid:48) be two inputs.
From lemma 1, we have that
[Equation]
where q 1 ( x, x (cid:48) )
=
σ 2 b + σ 2 w d x T x (cid:48) and q l ( x, x (cid:48) )
=
σ 2 b + σ 2 w E f ∼N
[Equation]
0 ,q l − 1 ) [ φ ( f ( x ))
φ ( f ( x (cid:48) ))] and ˙ q l ( x, x (cid:48) )
=
σ 2 w E f ∼N
[Equation]
0 ,q l − 1 ) [ φ
[Equation]
( f
[Equation]
x ))
φ
[Equation]
( f ( x (cid:48) ))] .
From facts 1, 2, 4, 9, 17, in the ordered/chaotic phase, there exist k, β, η, l 0 > 0 and α ∈ (0 , 1) such that for all l ≥ l 0 we have
[Equation]
and
[Equation]
Therefore, there exists M > 0 such that for any l ≥ l 0 and x, x (cid:48) ∈ R d
[Equation]
Letting r l = sup
( x,x (cid:48) ) ∈ B (cid:15) | K l ( x, x (cid:48) ) − k 1 − α | , we have
[Equation]
We conclude using Appendix Lemma 6.
Under Assumption 1, the proof is similar for CNN, using Appendix Lemmas 3 and 4.
Now, we show that the Initialization on the EOC improves the convergence rate of the NTK wrt L .
We ﬁrst prove two preliminary lemmas that will be useful for the proof of the next proposition.
Hereafter, the notation g ( x ) =
Θ( m ( x )) means there exist two constants A, B > 0 such that Am ( x ) ≤ g ( x ) ≤ Bm ( x ) .
Appendix Lemma 7.
Let A, B, Λ ⊂ R + be three compact sets, and ( a l ) , ( b l ) , ( λ l ) be three sequences of non-negative real numbers such that for all ( a 0 , b 0 , λ 0 )
∈
A × B × Λ
[Equation]
where α ∈ N ∗ independent of a 0 , b 0 , λ 0 , q ( b 0 ) ≥ 0 is a limit that depends on b 0 , and β ∈ (0 , 1) .
Assume the ‘ O ’ and ‘ o ’ depend only on A, B, Λ ⊂ R .
Then, we have
[Equation]
Proof.
Let A, B, Λ ⊂ R be three compact sets and ( a 0 , b 0 , λ 0 ) ∈
A × B × Λ .
It is easy to see that there exists a constant G > 0 independent of a 0 , b 0 , λ 0 such that | a l | ≤ G × l + | a 0 | for all l ≥ 0 .
Letting r l = a l l , we have that for l ≥[-
-]
[Equation]
where O bound depends only on A, B, Λ .
Letting x l = r l − q 1+ α , there exists M > 0 that depends only on
A, B, Λ , and l 0 > 0 that depends only on α such that for all l ≥ l[-
-]
[Equation]
Let us deal with the right hand inequality ﬁrst.
By induction, we have that
[Equation]
By taking the logarithm of the ﬁrst term in the right hand side and using the fact that (cid:80) lk = l 0 1 k = log( l ) + O (1) , we have
[Equation]
where the bound Θ does not depend on l 0 .
For the second part, observe that
[Equation]
and
[Equation]
Since α ≥ 1 ( α ∈ N ∗ ), then the serie with term k α − β is divergent and we have that
[Equation]
Therefore, it follows that
[Equation]
This proves that
[Equation]
where the ‘ o ’ bound depends only on A, B, Λ .
Using the same approach for the left-hand inequality, we prove that
[Equation]
This concludes the proof.
The next lemma is a different version of the previous lemma which will be useful for other applications.
Appendix Lemma 8.
Let A, B, Λ ⊂ R + be three compact sets, and ( a l ) , ( b l ) , ( λ l ) be three sequences of non-negative real numbers such that for all ( a 0 , b 0 , λ 0 )
∈
A × B × Λ
[Equation]
b l [Equation] l − 2 ) , where α ∈ N ∗ , κ (cid:54) =
both do not depend on a 0 , b 0 , Λ 0 , q ( b o ) ∈ R + is a limit that depends on b 0 .
Assume the ‘ O ’ and ‘ o ’ depend only on A, B, Λ ⊂ R .
Then, we have
[Equation]
Proof.
Let A, B, Λ ⊂ R be three compact sets and ( a 0 , b 0 , λ 0 ) ∈
A × B × Λ .
Similar to the proof of Appendix Lemma 7, there exists a constant G > 0 independent of a 0 , b 0 , λ 0 such that | a l | ≤ G × l + | a 0 | for all l ≥ 0 , therefore ( a l /l ) is bounded.
Let r l = a l l .
We have
[Equation]
Let x l = r l − q 1+ α .
It is clear that λ l = 1 − α/l + O ( l − 3 / 2 ) .
Therefore, using appendix lemma 7 with β = 1 / 2 , we have r l → q 1+ α uniformly over a 0 , b 0 , λ 0 .
Thus, assuming κ > 0 (for κ < 0 , the analysis is the same), there exists κ 1 , κ 2 , M, l 0 > 0 that depend only on A, B, Λ such that for all l ≥ l[-
-]
[Equation]
It follows that
[Equation]
and
[Equation]
Recall that we have
[Equation]
and
[Equation]
so that
[Equation]
Therefore, we obtain
[Equation]
where C 1 > 0 is a constant.
Similarly, there exists a constant C 2 > 0 such that
[Equation]
Moreover, having that ( l − α − 1)!
l ! ∼ l − 1 − α yields
[Equation]
where C (cid:48) and ‘ o ’ depend only on A, B, Λ .
Using the same analysis, we get x l ≥ C [Equation] log
[Equation]
l
[Equation] + [Equation] log( l )) where C (cid:48)(cid:48) and ‘ o ’ depend only on A, B, Λ , which concludes the proof.
Theorem 1 (Neural Tangent Kernel on the Edge of Chaos) .
Let φ be ReLU or Tanh, ( σ b , σ w ) ∈ EOC and ˜
K L = K L /L .
We have that
[Equation]
Moreover, there exists a constant λ ∈ (0 , 1) such that for all (cid:15) ∈ (0 , 1)
[Equation]
where
• if φ is ReLU-like, then ˜ K ∞ ( x, x (cid:48) )
=
σ 2 w (cid:107) x
(cid:107) (cid:107) x (cid:48) (cid:107) d
[Equation]
[Equation]
• if φ is Tanh, then ˜ K ∞ ( x, x (cid:48) )
=
q (1 − (1 − λ ) 1 x (cid:54) =
x (cid:48) ) where q > 0 is a constant.
Proof.
We start by proving the results for an FFNN architecture, the proof machinery remains the same for CNN under assumption:cnn.
Case 1: FFNN.
Let (cid:15) ∈ (0 , 1) , E ⊂ R d , ( σ b , σ w ) ∈ EOC, and x, x (cid:48) ∈ R d . Recall that c l ( x, x (cid:48) )
=
q l ( x,x (cid:48) ) √ q l ( x,x ) q l ( x (cid:48) ,x (cid:48) ) .
Let γ l := 1 − c l ( x, x (cid:48) ) and f be the correlation function deﬁned by the recursive equation c l {++
+} = f
( c l ) (See appendix 4).
By deﬁnition, we have that ˙
q l ( x, x ) = f (cid:48)
( c l − 1 ( x, x (cid:48) )) .
Let us ﬁrst prove the result for ReLU.
• φ = ReLU :
From fact 11 in the appendix, we know that, when choosing the hyper-parameters ( σ w , σ b ) on the EOC for ReLU, the variance q l ( x, x ) is constant w.r.t l and is given by q l ( x, x ) =
q 1 ( x, x ) = σw d || x || 2 .
Moreover, from fact [-
-]16, we have that ˙
q l ( x, x ) = 1 .
Therefore
[Equation]
which concludes the proof for K L ( x, x ) .
Note that the results is ’exact’ for ReLU, which means the upper bound O ( L − 1 ) is valid but not optimal in this case.
However, we will see that this bound is optimal for Tanh.
From Appendix Lemma 1, we have that
[Equation]
and
[Equation]
Using Appendix Lemma 8 with a l = K l +1 ( x, x
[Equation]
[Equation] b l = q l +1 ( x, x (cid:48) ) [Equation] λ l = f (cid:48) ( c l ( x, x (cid:48) ))
[Equation]
we conclude that [Equation] Using the compactness of B (cid:15)
[Equation]
we conclude that
[Equation]
• φ =
Tanh : The proof in the case of Tanh is slightly different from that of ReLU. We use different technical lemmas to conclude
From Appendix Lemma 2, we have that sup
[Equation]
x,x (cid:48)
[Equation]
= O
[Equation]
l − 3 ) where κ = 2 f (cid:48)(cid:48) (1) > 0 and ζ = f 3 (1)> 0 .
Moreover, we have that sup
[Equation]
x,x [Equation]
[Equation]
We conclude in the same way as in the case of ReLU using Appendix Lemma 8.
The only difference is that, in this case, the limit of the sequence b l = q l +1 ( x, x (cid:48) ) is the limiting variance q (from facts 3, 1) does not depend on ( x, x (cid:48) ) .
Case 2: CNN.
Under Assumption 1, the NTK of a CNN is the same as that of an FFNN.
Therefore, the results on the NTK of FFNN are all valid to the NTK of CNN K lα,α (cid:48) for any α, α (cid:48) .
6.2 Proofs of the results of Section 3.2 on ResNets
In this section, we provide proofs for lemmas 3 and 4 together with Theorem ??
and proposition ??
on ResNets.
Lemma 3 in the paper gives the recursive formula for the mean-ﬁeld NTK of a ResNet with Fully Connected blocks.
Lemma 3 (NTK of a ResNet with Fully Connected layers in the inﬁnite width limit) .
Let x, x (cid:48) be two inputs and K res, 1 be the exact NTK for the Residual Network with 1 layer.
Then, we have
• For the ﬁrst layer (without residual connections), we have for all x, x (cid:48) ∈ R
[Equation]
where x · x (cid:48) is the inner product in R d .
• For l ≥ 2 , as n 1 , n 2 , ..., n L − 1 → ∞ , we have for all i,
[Equation]
where K lres ( x, x (cid:48) ) is given by the recursive formula have for all x, x (cid:48) ∈ R d and l ≥ 2 , as n 1 , n 2 , ..., n l → ∞ recursively, we have
[Equation]
Proof.
The ﬁrst result is the same as in the FFNN case since we assume there is no residual connections between the ﬁrst layer and the input.
We prove the second result by induction.
• Let x, x (cid:48) ∈ R d .
We have
[Equation]
The proof is similar to the FeedForward network NTK
For l ≥ and i ∈ : n l ]
[Equation]
Therefore, we obtain
[Equation]
where
[Equation]
Using the induction hypothesis, as n 0 , n 1 , ..., n l − 1 → ∞ , we have that
[Equation]
where I (cid:48)
=
σ 2 w n l w
[Equation]
+1 ii ( φ
[Equation]
( y li ( x ))
+ φ (cid:48)
( y li ( x (cid:48) )
[Equation]
) K lres ( x, x (cid:48) ) .
As n l → ∞ , we have that I (cid:48)
→ 0 .
Using the law of large numbers, as n l → ∞
[Equation]
Moreover, we have that
[Equation]
j [Equation] Now we proof the recursive formula for ResNets with Convolutional layers.
Lemma 4 (NTK of a ResNet with Convolutional layers in the inﬁnite width limit) .
Let K res, 1 be the exact
NTK for the ResNet with 1 layer.
Then
• For the ﬁrst layer (without
residual
connections), we have for all x, x
(cid:48) ∈
R d
[Table]
• For l ≥ 2 , as n 1 , n 2 , ..., n l − 1 → ∞ recursively, we have for all i,
[Table]
, where K res
[Table]
α
[Table]
is given by the recursive formula for all x, x (cid:48) ∈ R d , using the same notations as in lemma 2,
[Table]
[Equation]
[Table]
where Ψ
[Equation]
Proof.
Let x, x (cid:48) be two inputs.
We have that
[Equation]
Assume the result is true for l − [Equation] let us prove it for l .
Let θ 1: l − 1 be model weights and bias in the layers 1 to l − 1 .
Let ∂ θ 1: l − 1 y li,α ( x ) =
∂y li,α ( x ) ∂θ 1: l − 1 .
We have that
[Equation]
this yields
[Equation]
where
[Equation]
As n 1 , n 2 , ..., n l − 2 → ∞ and using the induction hypothesis
[Equation]
we have
[Equation]
Note that K l − 1 ( j,α + β ) , ( j,α (cid:48) + β ) ( x, x
[Equation]
) = K [Equation] [Equation] α + β ) , (
[Equation]
α (cid:48)
+ β ) ( x, x (cid:48) )
for all j since the variables are iid across the channel index j .
Now letting n l − 1 → ∞ , we have that
[Equation]
where f (cid:48) ( c l − 1 α + β,α
[Equation]
β
[Equation]
x
[Equation]
x [Equation]
We conclude using the fact that
[Equation]
Before moving to the main theorem on ResNets, We ﬁrst prove a Lemma on the asymptotic behaviour of c l for
ResNet.
Appendix Lemma 9 (Asymptotic expansion of c l for ResNet) .
Let (cid:15) ∈ (0 , 1) and σ w > 0 .
We have for FFNN
[Equation]
where κ σ w , κ (cid:48) σ w > 0 are two constants that depend on σ w .
Moreover, we have that
[Equation]
where f is the ReLU correlation function given in fact 12 and κ
[Equation]
σ W > 0 is a constant that depends on σ w .
Moreover, this result holds also for CNNs where the supremum should be replaced by sup [Equation] x
[Equation]
x [Equation] B
[Equation]
sup α,α [Equation]
Proof.
We ﬁrst prove the result for ResNet with fully connected layers, then we generalize it to convolutional layers.
Let (cid:15) ∈ (0 , 1) .
• Let x (cid:54) =
x (cid:48) ∈ R d , and c l :=
c l ( x, x (cid:48) ) .
It is straightforward that the variance terms follow the recursive form
[Equation]
Leveraging this observation, we have that
[Equation]
where f is the ReLU correlation function given in fact 12 and α
= σ 2 w 2 .
Recall that
[Equation]
As in the proof of Appendix Lemma 1, let γ l = 1 − c l , therefore, using Taylor expansion of f neargiven in fact 14 yields
[Equation]
This form is exactly the same as in the proof of Appendix Lemma 1 with s (cid:48)
= αs 1+ α and b (cid:48)
=
αb [-
 +-]1+ α .
Thus, following the same analysis we conclude.
For the second result, observe that the derivation is the same as in Appendix Lemma 1. • Under Assumption 1, results of FFNN hold for CNN.
The next theorem shows that no matter what the choice of σ w > 0 , the normalized NTK of a ResNet will always have a subexponential convergence rate to a limiting ¯ K ∞ res .
Theorem 2 (NTK for ResNet) .
Consider a ResNet satisfying
[Equation]
where F is either a convolutional or dense layer with ReLU activation.
Let K Lres be the corresponding NTK and ¯ K Lres =
K Lres /α
L (Normalized NTK) with α L =
L (1 + σ 2 w 2 )
L − 1 .
If the layers are convolutional assume Assumption 1 holds.
Then, we have
[Equation]
Moreover, there exists a constant λ ∈ (0 , 1) such that for all (cid:15) ∈ (0 , 1)
[Equation]
where ¯ K
[Equation]
Proof.
We start by proving the result for a ResNet architecture with fully-connected layers.
As in the previous proofs, the result easily extends to a ResNet with convolutional layers under assumption:cnn.
Case 1: ResNet with fully-connected layers.
Let (cid:15) ∈ (0 , 1) , E ⊂ R d , and x, x (cid:48) ∈ R d .
We ﬁrst prove the result for the diagonal terms K Lres ( x, x ) , we deal afterwards with off-diagonal terms K Lres ( x, x (cid:48) ) .
• Diagonal terms: from fact 12, we have that ˙
q l ( x, x ) =
σ 2 w 2 f (1) =
σ 2 w 2 .
Moreover, it is easy to see that the variance terms for a ResNet follow the recursive formula q l ( x, x ) = q l −
1 ( x, x ) + σ [Equation] / 2 × q
l
−
( x, x ) , hence
[Equation]
Recall that the recursive formula of NTK of a ResNet with fully-connected layers is given by (Appendix Lemma 3)
[Equation]
Hence, for the diagonal terms we obtain
[Equation]
Letting ˆ K l res =
K l res / (cid:16) 1 + σ 2 w l (cid:17) l − 1 yields
[Equation]
Therefore, ¯ K lres [Equation] x, [Equation]
[Equation]
x
[Equation]
2 , the conclusion is straightforward since E is compact and ˆ K 1 res ( x, x ) is continuous which implies that it is uniformly bounded on E .
• Off-diagonal terms: the argument is similar to that of Theorem with few key
differences
From Appendix Lemma 9 we have that
[Equation]
where κ σ w , κ (cid:48) σ w > 0 .
Moreover, we have that
[Equation]
Let α = σ 2 w 2 .
We also have ˙ q l +1 ( x, x (cid:48) )
=
αf (cid:48) ( c l ( x, x (cid:48) ))
where f is the ReLU correlation function given in fact 12.
It follows that for all ( x, x (cid:48) )
∈ B (cid:15)
[Equation]
for some constant ζ (cid:54) =
that does not depend on x, x (cid:48) .
The bound O does not depend on x, x (cid:48) either.
Now let a l = K l +1 res ( x,x (cid:48) )
(1+ α ) l .
Using the recursive formula of the NTK, we obtain a l = λ l a l − 1 + b l where λ l = 1 − 3 l − 1 + ζ log( l ) l 2 + O
l − 3 ) , b l = σ 2 w d (cid:112) (cid:107) x
cid:107)(cid:107) x (cid:48) (cid:107) f
( c l ( x, x (cid:48) ))
=
q ( x, x (cid:48) )
+ O ( l − 2 ) with q ( x, x (cid:48) )
=
σ 2 w d (cid:112) (cid:107) x (cid:107)(cid:107) x (cid:48) (cid:107) and where we used the fact that c l ( x, x (cid:48) )
=
1 + O ( l − 2 ) (Appendix
Lemma 1) and the formula for ResNet variance terms given by equation (18).
Observe that all bounds O are independent from the inputs ( x, x (cid:48) ) .
Therefore, using Appendix Lemma 8, we have
[Equation]
which can also be written as
[Equation]
We conclude by observing that K Lres ( x, x (cid:48) ) /
( L − 1)(1 + α )
L − 1 = K Lres ( x
[Equation]
x (cid:48) ) /L
(1 + α ) L − 1 +
O ( L − 1 ) where O can be chosen to depend only on (cid:15) .
Case 2: ResNet with Convolutional layers.
Under Assumption 1, the dynamics of the correlation and
NTK are exactly the same for FFNN, hence all results on FFNN apply to CNN.
6.2.1 Beyond ResNet: Scaled ResNet
The term α L in the residual NTK might cause numerical stability issues for NTK training, and the triviality of the limiting kernel yields a trivial NTK regime solution (recall that f t − f 0 belongs to the RKHS of the NTK; see section 2).
It turns out that we can improve the performance of NTK training of ResNets with a simple scaling of the ResNet blocks.
Theorem 3 (Scaled ResNet) .
Consider a ResNet satisfying
[Equation]
where F is either a convolutional or dense layer ( (3) and ( ?? ) ) with ReLU activation.
Then the results of
Theorem 2 apply with α L =
L 1+ σ 2 w / 2 and the convergence rate Θ(log
( L ) − 1 ) .
Theorem 3 shows that scaling the residual blocks by 1 / √ l has two important effects on the NTK: ﬁrst, it stabilizes the NTK which only grows as L 1+ σ 2 w 2 instead of L (1 + σ 2 w 2 )
L − 1 ; second, it drastically slows down the convergence rate to the limiting (trivial) ¯ K ∞ res .
Both properties are highly desirable for NTK training.
The second property in particular means that with the scaling, we can ‘NTK train’ deeper ResNets compared to the non-scaled ResNet.
A more aggressive scaling was studied in Huang et al.
(2020), where authors scale the blocks with 1 /L instead of our scaling 1 / √ l , and show that it also stabilizes the NTK of ResNet.
This is the main topic of the next chapter.
We particularly show that a suitable scaling ensures that the limiting NTK is universal, i.e. we can approximate any continuous function on some compact set K with a function from the Reproducing Kernel Hilbert Space of the limiting NTK.
This is a desirable property since the second term in the solution of the NTK regime lives in the RKHS of the NTK.
Now let us prove the Scaled Resnet result.
Before that, we prove the following Lemma
Appendix Lemma 10.
Consider a Residual Neural Network with the following forward propagation equations
[Equation]
where F is either a convolutional or dense layer (equations 3 and ?? ) with ReLU activation.
Then there exists ζ, ∇ > 0 such that for all (cid:15) ∈ (0 , 1) [Equation] x [Equation]
x (cid:48) )
∈ B
[Equation]
where the bound ‘ o ’ depends only on (cid:15) .
For CNN, under Assumption 1, the result holds and the supremum is taken also over α, α (cid:48) , i.e.
[Equation]
Proof.
We ﬁrst start with the dense layer case.
Let (cid:15) ∈ (0 , 1) and ( x, x (cid:48) )
∈ B (cid:15) be two inputs and denote by c l :=
c l ( x, x (cid:48) ) .
Following the same machinery as in the proof of Appendix Lemma 9, we have that
[Equation]
where α l = σ 2 w 2 l .
Using fact 12, it is straightforward that f (cid:48) ≥ 0 , hence f is non-decreasing.
Therefore, c l ≥ c l − 1 and c l converges to a ﬁxed point c .
Let us prove that c = 1 .
By contradiction, suppose c < 1 so that f ( c ) − c >
( f has a unique ﬁxed point which is 1 ).
This yields
[Equation]
by summing, this leads to c l − c ∼ ( f ( c ) − c ) log( l ) which is absurd since f ( c ) (cid:54) =
c ( f has only 1 as a ﬁxed point).
We conclude that c = 1 .
Using the non-decreasing nature of f , it is easy to conclude that the convergence is uniform over B (cid:15) .
Now let us ﬁnd the asymptotic expansion of 1 − c l .
Recall the Taylor expansion of f near 1 given in fact[-
-]
[Equation]
where s = 2 √ 2 3 π and b = √ 2 30 π .
Letting γ l = 1 − c l , we obtain
[Equation]
which yields
[Equation]
therefore, we have that
[Equation]
and 1 − c l ∼ ζ log( l ) 2 where ζ = 16 /s 2 σ 4 w . we can further expand the asymptotic approximation to have
[Equation]
where ∇ > 0 .
the ‘ o ’ holds uniformly for ( x, x (cid:48) )
∈ B (cid:15) as in the proof of Appendix Lemma 1.
This result holds for a ResNet with CNN layers under Assumption 1 since the dynamics are the same in this case.
Proof of Theorem
Proof.
We use the same techniques as in the non scaled case.
Let us prove the result for fully connected layers, the proof for convolutional layers follows the same analysis.
Let (cid:15) ∈ (0 , 1) and x, x (cid:48)
∈ B (cid:15) be two inputs.
We ﬁrst prove the result for the diagonal term K Lres ( x, x ) then K Lres ( x, x (cid:48) ) .
• We have that ˙
q l ( x, x ) = σ 2 w 2 l f ( ) =
σ 2 w 2 l .
Moreover, we have q l ( x, x ) =
q l − 1 ( x, x ) + σ 2 w / 2 l × q l − 1 ( x, x ) =
[ (cid:81) lk
=1 (1 + σ 2 w / 2 k )]
σ 2 w d (cid:107) x (cid:107) 2 .
Recall that
[Equation]
letting k (cid:48) l =
[Equation]
[Equation]
) ,we have that
[Equation]
using the fact that
(cid:81) lk =
1 + σ w / 2 k ) = Θ( l
σ
2 w 2 ) , we conclude for K lres ( x, x )
Recall that
[Equation]
Let c l :=
c l ( x, x (cid:48) ) .
From Appendix Lemma 10 we have that
[Equation]
ζ = 16 s 2 σ 4 w and ∇ > 0 .
Using the Taylor expansion of f (cid:48) as in Appendix Lemma 1, it follows that
[Equation]
where ζ (cid:48) = ∇√ 2 πζ .
We obtain
[Equation]
[Table]
where ζ (cid:48)(cid:48)
=
σ 2 w 2 ζ (cid:48) .
Letting a l =
K
[Equation]
1 res (
[Equation]
x (cid:48) ) (cid:81)
[Equation]
=1 (
[Equation]
2 w /
[Equation]
) , we obtain a
[Equation]
λ
[Equation]
+ b l where λ l = 1 − l − 1 − 3 l − 1 log( l ) − 1 + O ( l − 1 log( l ) − 2 ) , b l = (cid:112) q 1 ( x, x ) (cid:112) q 1 ( x (cid:48) , x (cid:48) )
f ( c l ( x, x (cid:48) ))
=
q ( x, x (cid:48) )
+ O (log( l ) − 2 ) with q = (cid:112) q 1 ( x, x ) (cid:112) q 1 ( x (cid:48) , x (cid:48) ) and where we used the fact that c l = 1 + O (log( l ) − 2 ) (Appendix Lemma 10).
Now we proceed in the same way as in the proof of Appendix Lemma 8.
Let x l = a l l − q , then there exists M 1 , M 2 > 0 such that
[Equation]
therefore, there exists l 0 independent of ( x, x [Equation]
) such that for all l ≥ lx [Equation] ≤
[Equation] [Equation]
and x l ≥ x [Equation] [Equation] (cid:89) k
=
[Equation] [Equation]
M
[Equation]
l (cid:88) k
[Equation]
ll (cid:89)
j = k + (1 − 1 j ) k − 1 log
[Equation]
after simpliﬁcation, we have that
[Equation]
where we have used the asymptotic approximation of the Logarithmic Intergal function Li ( x ) = (cid:82) t 1 log( t ) ∼ x →∞ x log( x ) we conclude that α L
=
L × (cid:81) lk
=1 (1 + σ 2 w / 2 k ) ∼ L 1+ σ 2 w 2 and the convergence rate of the NTK is now Θ(log( L ) − 1 ) which is better than Θ( L − 1 ) .
The convergence is uniform over the set B (cid:15) .
In the limit of large L , the matrix NTK of the scaled resnet has the following form
[Equation]
where U is the matrix of ones, and M L has all elements but the diagonal equal to 1 and the diagonal terms are O
( L − 1 log( L ))
→ 0 .
Therefore, M L is inversible for large L which makes ˆ K lres also inversible.
Moreover, observe that the convergence rate for scaled resnet is log( L ) − 1 which means that for the same depth L , the NTK remains far more expressive for scaled resnet compared to standard resnet, this is particularly important for the generalization.
[Table]
Does Scaled ResNet outperforms ResNet with SGD?
We train standard ResNet with depths 32, 50, and 104 on CIFAR100 with SGD.
We use a decaying learning rate schedule; we start with 0.1 and divide by 10 after n e / 2 epochs, where n e is the total number of epochs; we scale again, by 10, after n e / 4 epochs.
We use a batch size of 128, and we train the model with 160 epochs.
Theorem 3 shows that the NTK of Scaled ResNet is more stable compared to the NTK of standard ResNet.
Although this result is limited to NTK training, we investigate the impact of scaling on SGD training.
Table 2 displays test accuracy for standard ResNet and scaled ResNet after 10 and 160 epochs; Scaled ResNet outperforms ResNet and converges faster.
However, it is not clear whether this is linked to the NTK, or caused by something else.
We leave this for future work.
6.3 Spectral decomposition of the limiting NTK 6.3.1 Review on Spherical Harmonics
We start by giving a brief review of the theory of Spherical Harmonics MacRobert (1967).
Let S d − 1 be the unit sphere in R d deﬁned by S d − 1 = { x ∈ R
d : (cid:107) x (cid:107) 2 = 1 } .
For some k ≥ 1 , there exists a set
( Y k,j ) 1 ≤ j ≤ N ( d,k ) of Spherical Harmonics of degree k with N ( d, k ) =
2 k + d − 2 k (cid:0) k
+
[Equation]
3 d − 2 (cid:1) .
The set of functions ( Y k,j ) k ≥ 1 ,j ∈ [1: N ( d,k )] form an orthonormal basis with respect to the uniform measure on the unit sphere S d − 1 .
For some function g , the Hecke-Funk formula is given by
[Equation]
where ν d − 1 is the uniform measure on the unit sphere S d − 1 , Ω d is the volume of the unit sphere S d − 1 , and P dk is the multi-dimensional Legendre polynomials given explicitly by Rodrigues’ formula
[Equation]
( P [Equation] [Equation] ≥ 0 form an orthogonal basis of L 2 ([ − 1 , 1] , (1 − t 2 ) d − 32 dt ) , i.e.
[Equation]
where δ ij is the Kronecker symbol.
Moreover, we have
[Equation]
Using the Heck-Funk formula, we can easily conclude that any dot product kernel on the unit sphere S d − 1 , i.e. and kernel of the form κ ( x, x (cid:48) )
=
g ( (cid:104) x, x (cid:48) (cid:105) ) can be decomposed on the Spherical Harmonics basis.
Indeed, for any x, x (cid:48) ∈ S
d − 1 , the decomposition on the spherical harmonics basis yields
[Equation]
Using the Hecke-Funk formula yields
[Equation]
we conclude that
[Equation]
where µ k =
[Equation]
We use these result in the proof of the next theorem.
Proposition 2 (Spectral decomposition) .
Let κ L be either, the NTK ( K L ) for an FFNN with L layers initialized on the Ordered phase, The Average NTK ( AK L ) for an FFNN with L layers initialized on the EOC, or the
Normalized NTK ( ¯ K Lres ) for a ResNet with L layers (Fully Connected).
Then, for all L ≥ 1 , there exists ( µ Lk )
k ≥ such that for all x, x (cid:48)
∈ S d −
[Equation]
( Y k,j ) k ≥ 0 ,j ∈ [1: N ( d,k )] are spherical harmonics of S d − 1 , and N ( d, k ) is the number of harmonics of order k .
Moreover, we have that 0 < µ ∞ 0 =
lim L →∞ µ L 0 < ∞ , and for all k ≥ 1 , lim L →∞ µ Lk = 0 .
Proof.
From the recursive formulas of the NTK for FFNN, CNN and ResNet architectures, it is straightforward that on the unit sphere S d − 1 , the kernel κ L is zonal in the sense that it depends only on the scalar product, more precisely, for all L ≥ 1 , there exists a function g L such that for all x, x (cid:48) ∈ S d − [Equation] x
[Equation]
( (cid:104) x [Equation] x (cid:48) (cid:105) ) using the previous results on Spherical Harmonics, we have that for all x, x (cid:48) ∈ S d −κ L ( x, x (cid:48) )
= (cid:88) k ≥µ
Lk
N ( d,k ) (cid:88) j =
Y k,j ( x ) Y k,j ( x (cid:48) )
where µ Lk = Ω d − 1 Ω d (cid:82) 1 − 1 g L ( t ) P dk ( t )(1 − t 2 ) ( d − 3) / 2 dt .
For k = 0 , we have that for all L ≥ 1 , µ L 0 =
Ω
d − 1 Ω d (cid:82) 1 − 1 g L ( t )(1 − t 2 ) ( d − 3) / 2 dt .
By a simple dominated convergence argument, we have that lim L →∞ µ L 0 = qλ
Ω d − 1 Ω d (cid:82) 1 − 1 (1 − t 2 ) ( d − 3) / 2 dt > 0 , where q, λ are given in Theorems 1, 2 and Proposition 1 (where we take q = 1 for the Ordered/Chaotic phase initialization in Proposition 1).
Using the same argument, we have that for k ≥ 1 , lim L
[Equation]
=
[Equation]
=
[Equation]
d
[Equation]
P
[Equation]
L 2 ([ − 1 , 1] , (
[Equation]
2 ) d − 32 dt ) = 0 .
Further experimental results