D ISENTANGLING T RAINABILITY AND G ENERALIZATION IN D EEP L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
A fundamental goal in deep learning is the characterization of trainability and generalization of neural networks as a function of their architecture and hyperparameters.
In this paper, we discuss these challenging issues in the context of wide neural networks at large depths where we will see that the situation simpliﬁes considerably.
To do this, we leverage recent advances that have separately shown: (1) that in the wide network limit, random networks before training are
Gaussian Processes governed by a kernel known as the Neural Network Gaussian
Process (NNGP) kernel, (2) that at large depths the spectrum of the NNGP kernel simpliﬁes considerably and becomes “weakly data-dependent”, and (3) that gradient descent training of wide neural networks is described by a kernel called the Neural Tangent Kernel (NTK) that is related to the NNGP.
Here we show that in the large depth limit the spectrum of the NTK simpliﬁes in much the same way as that of the NNGP kernel.
By analyzing this spectrum, we arrive at a precise characterization of trainability and a necessary condition for generalization across a range of architectures including Fully Connected Networks (FCNs) and Convolutional Neural Networks (CNNs).
In particular, we ﬁnd that there are large regions of hyperparameter space where networks can only memorize the training set in the sense they reach perfect training accuracy but completely fail to generalize outside the training set, in contrast with several recent results.
By comparing CNNs with- and without-global average pooling, we show that CNNs without average pooling have very nearly identical learning dynamics to FCNs while CNNs with pooling contain a correction that alters its generalization performance.
We perform a thorough empirical investigation of these theoretical results and ﬁnding excellent agreement on real datasets.
I NTRODUCTION
Machine learning models based on deep neural networks have attained state-of-the-art performance across a dizzying array of tasks including vision (Cubuk et al., 2019),
speech recognition (Park et al., 2019), machine translation (Bahdanau et al., 2014), chemical property prediction Gilmer et al.
(2017), diagnosing medical conditions Raghu et al.
(2019), and playing games Silver et al.
Historically, the rampant success of deep learning models has lacked a sturdy theoretical foundation; architectures, hyperparameters, and learning algorithms are more often than not selected by brute force search Bergstra & Bengio (2012) and heuristics Glorot & Bengio (2010).
Recently, signiﬁcant theoretical progress has been made on several fronts that have shown promise in making neural network design more systematic.
In particular, in the inﬁnite width (or channel) limit, the distribution of functions induced by neural networks with random weights and biases has been precisely characterized before, during, and after training.
The study of inﬁnite networks dates back to seminal work by Neal (1994) who showed that the distribution of functions given by single hidden-layer networks with random weights and biases in the inﬁnite-width limit are Gaussian Processes (GPs).
Recently, there has been renewed interest in studying random, inﬁnite, networks starting with concurrent work on “conjugate kernels” (Daniely et al., 2016; Daniely, 2017) and “mean-ﬁeld theory” (Poole et al., 2016; Schoenholz et al., 2017).
The former set of papers argued that the empirical covariance matrix of pre-activations became deterministic in the inﬁnite-width limit and called this the conjugate kernel of the network while the latter papers studied the properties of these limiting kernels along with the kernel describing distribution of gradients.
In particular, it was shown that the spectrum of the conjugate kernel of wide fully-connected networks approached a well-deﬁned, data-independent, limit when the depth exceeds a certain scale, ξ .
Networks with tanh -nonlinearities (among other bounded activations) exhibit a phase transition between two limiting spectral distributions of the conjugate kernel as a function of their hyperparameters with ξ diverging at the transition.
It was additionally hypothesized that networks were un-trainable when the conjugate kernel was sufﬁciently close to its limit.
Since then this analysis has been pushed to a wide range for architectures such as convolutions (Xiao et al., 2018), recurrent networks (Chen et al., 2018; Gilboa et al., 2019), networks with residual connections (Yang & Schoenholz, 2017), networks with quantized activations (Blumenfeld et al., 2019), the spectrum of the ﬁsher (Karakida et al., 2018), a range of activation functions Hayou et al.
(2018), and batch normalization (Yang et al., 2019).
In each case, it was observed that the spectra of the kernels correlated strongly with whether or not the architectures were trainable.
While these papers studied the properties of the conjugate kernels, especially the spectrum in the large-depth limit, a branch of concurrent work made a stronger statement: that many networks converge to Gaussian Processes as their width becomes large (Lee et al., 2018; Matthews et al., 2018; Novak et al., 2019b; Yang, 2019).
In this case, the Conjugate Kernel was referred to as the Neural Network
Gaussian Process (NNGP) kernel.
Together this work offered a signiﬁcant advance to our understanding of wide neural networks; however, this theoretical progress was limited to networks at initialization or after Bayesian posterior estimation and provided no link to gradient descent.
Moreover, there was some preliminary evidence that suggested the situation might be more nuanced than the qualitative link between the NNGP spectrum and trainability might suggest.
For example, Philipp et al.
(2017) showed that deep fullyconnected tanh -networks could be trained after the kernel reached its large-depth, data-independent, limit but that these networks did not generalize to unseen data.
In the last year, signiﬁcant theoretical clarity has been reached regarding the relationship between the
GP prior and the distribution following gradient descent.
In particular, Jacot et al.
(2018) along with followup work (Lee et al., 2019; Chizat et al., 2019) showed that the distribution of functions induced by gradient descent for inﬁnite-width networks is a Gaussian Process with a particular compositional kernel known as the Neural Tangent Kernel (NTK).
In addition to characterizing the distribution over functions following gradient descent in the wide network limit, the learning dynamics can be solved analytically throughout optimization.
In this paper, we leverage these developments and revisit the relationship between architecture, hyperparameters, trainability, and generalization in the large-depth limit for a variety of neural networks.
In particular, we make the following contributions:
We compute the large-depth asymptotics of several quantities related to trainability, including the largest eigenvalue of the NTK, λ max , and the condition number κ = λ max /λ min , where λ min is the smallest eigenvalue; see Table 1. 2.
We introduce the residual predictor ∆ ( l ) , namely the difference between the ﬁnite depth and inﬁnite depth NTK predictions, which is related to the model’s ability to generalize: the network fails to generalize if ∆ ( l ) is too small.
We show that the ordered and chaotic phases identiﬁed in Poole et al.
(2016) lead to markedly different limiting spectra of the NTK, which further indicates that, as a function of depth, the optimal learning rates ought to decay exponentially, linearly and remain roughly a constant in the chaotic, order-to-chaos and ordered phases, respectively.
We examine the differences in the above quantities for fully-connected networks (FCNs) and convolutional networks (CNNs) with and without pooling and precisely characterize the effect of pooling to these quantities.
We provide substantial experimental evidence supporting these claims, includes experiments that densely vary the hyperparameters of FCNs and CNNs with and without pooling.
Together these results provide a complete, analytically tractable, and dataset-independent theory for learning in very deep and wide networks.
In addition to being interesting in its own right our theory provides a strong test of the NTK theory.
Finally, our results provides clarity regarding the observation that for linear networks the learning rate must be decreased linearly in the depth of the [Table] network Saxe et al.
Here, we note that this is true only for networks that are initialized critically , i.e. on the order-to-chaos phase boundary.
B ACKGROUND
We summarize recent developments in the study of wide random networks.
We will keep our discussion relatively informal; see (Lee et al., 2018; Matthews et al., 2018; Novak et al., 2019b) for a more rigorous version of these arguments.
To simplify this discussion and as a warmup for the main text, we will consider the case of FCNs.
Consider a fully-connected network of depth L where each layer has a width N ( l ) and an activation function φ : R → R .
In this work we will take φ = erf however, most of the results will hold for a wide range of non-linearities though speciﬁcs - such as the phase diagram - can vary substantially.
For simplicity, we will take the width of the hidden layers to inﬁnity sequentially: N (1) → ∞ , . . . , N ( L − 1) → ∞ .
The network is parameterized by weights and biases that we take to be randomly initialized with W ( l ) ij , b ( l ) i ∼ N (0 , 1) along with hyperparameters, σ w and σ b that set the scale of the weights and biases.
Letting the pre-activations in layer l due to an input x be given by z ( l ) i ( x ) , the network is then described by the recursion,
[Equation]
Notice that as N ( l ) → ∞ , the sum ends up being over a large number of random variables and we can invoke the central limit theorem to conclude that the { z ( l +1) i } i ∈ [ N ] are i.i.d. Gaussian with zero mean.
Given a dataset of m points, the distribution over pre-activations can therefore be described completely by the covariance matrix between neurons in different inputs K ( l ) ( x, x (cid:48) )
=
E [ z ( l ) i ( x ) z ( l ) i ( x (cid:48) )] .
Inspecting Equation 75, we see that K ( l +1) ( x, x (cid:48) ) can be computed in terms of K ( l ) ( x, x (cid:48) ) as
[Equation]
for T , an appropriately deﬁned operator from the space of positive semi-deﬁnite matrices to itself.
Equation 2 describes a dynamical system on positive semi-deﬁnite matrices K ( x, x (cid:48) ) .
It was shown in Poole et al.
(2016) that ﬁxed points, K ∗ ( x, x (cid:48) ) , of these dynamics exist such that lim l →∞ K ( l ) ( x, x (cid:48) )
=
K ∗ ( x, x (cid:48) ) with K ∗ ( x, x (cid:48) )
=
q ∗ [ δ x,x (cid:48)
+ c ∗ (1 − δ x,x (cid:48) )] independent of the inputs x and x (cid:48) .
The values of q ∗ and c ∗ are determined by the hyperparameters, σ w and σ b .
However Equation 2 admits multiple ﬁxed points (e.g. c ∗ = 0 , 1 ) and the stability of these ﬁxed points plays a signiﬁcant role in determining the properties of the network.
Generically, there are large regions of the ( σ w , σ b ) plane in which the ﬁxed-point structure is constant punctuated by curves, called phase transitions, where the structure changes.
The rate at which K ( x, x (cid:48) ) approaches or departs K ∗ ( x, x (cid:48) ) can be determined by expanding Equation 2 about its ﬁxed point, δ K ( x, x (cid:48) )
=
K ( x, x (cid:48) )
− K ∗ ( x, x (cid:48) ) to ﬁnd
[Equation]
3) which exhibits exponential convergence to - or divergence from - the ﬁxed-point as δ K ( l ) ( x, x (cid:48) ) ∼ χ ( x, x (cid:48) ) l where χ ( x, x (cid:48) )
= σ 2 w ˙ T ( K ∗ ( x, x (cid:48) )) .
Since K ∗ ( x, x (cid:48) ) does not depend on x or x (cid:48) it follows that χ ( x, x (cid:48) ) will take on a single value, χ c ∗ , whenever x (cid:54) = x (cid:48) . If χ c ∗ > 1 then the ﬁxed point is unstable and, as discussed above, there will be another ﬁxed point that becomes stable, if χ c ∗ <then the ﬁxed point is stable, and if χ c ∗ = 1 then the hyperparameters lie at a phase transition.
As was shown in Poole et al.
, there is always a ﬁxed-point at c ∗ = 1 whose stability is determined by χ 1 .
This deﬁnes the order-to-chaos transition. Note, that χ c ∗ can be used to deﬁne a depth-scale, ξ c ∗ = − 1 / log( χ c ∗ ) that describes the number of layers over which K ( l ) approaches K ∗ . This provides a precise characterization of the NNGP kernel at large depths.
As discussed above, recent work (Jacot et al., 2018; Lee et al., 2019; Chizat et al., 2019) has connected the prior described by the NNGP with the result of gradient descent training using a quantity called the NTK.
To construct the NTK, suppose we enumerate all the parameters in the fully-connected network described above by θ α .
The ﬁnite width NTK is deﬁned by ˆΘ( x, x (cid:48) )
=
J ( x ) J ( x (cid:48) ) T where J iα ( x ) =
∂ θ α z Li
( x ) is the Jacobian evaluated at a point x .
The main result in Jacot et al.
(2018) was to show that in the inﬁnite-width limit, the NTK converges to a deterministic kernel Θ and remains constant over the course of training.
As such, at a time t during gradient descent training with an MSE loss, the expected outputs of an inﬁnitely wide network, µ t ( x ) =
E [ z Li ( x )] , evolve as
[Equation]
for train and test points respectively; see Section 2 in Lee et al.
Here Θ test, train denotes the NTK between the test inputs X test and training inputs X train and Θ train, train is deﬁned similarly. Since ˆΘ converges to Θ , the gradient ﬂow dynamics of real network also converge to the dynamics described by Equation 4 and Equation 5 (Jacot et al., 2018; Lee et al., 2019; Chizat et al., 2019; Yang, 2019; Arora et al., 2019; Huang & Yau, 2019).
As the training time, t tends to inﬁnity we note that these equations reduce to µ ( X train ) =
Y train and µ ( X test )
=
Θ test, train Θ − 1 train, train Y train .
Consequently we call the linear operator
[Equation] the “mean predictor” or “predictor” for short.
In addition to showing that the NTK describes networks during gradient descent, Jacot et al.
(2018) showed that the NTK could be computed in closed form in terms of T , ˙ T , and the NNGP as,
[Equation]
7) where Θ ( l ) is the NTK for the pre-activations at layer- l .
M ETRICS FOR T RAINABILITY AND G ENERALIZATION AT L ARGE D EPTH
We begin by discussing the interplay between the conditioning of Θ train,train and the trainability of wide networks.
We can write Equation 4 in terms of the spectrum of Θ train,train letting Θ train,train =
U T DU as,
[Equation]
where λ i are the eigenvalues of Θ train,train and ˜ µ t ( X train ) =
Uµ t ( X train ) , ˜ Y train =
UY train are the mean prediction and the labels respectively written in the eigenbasis of Θ train,train .
If we order the eigenvalues such that λ 0 ≥ · · · ≥ λ M then it has been hypothesized in e.g. Lee et al.
(2019) that the maximum feasible learning rate scales as η ∼ 2 /λ 0 as we verify empirically in section 4.
{+Plugging this scaling for η into Equation 8 we see that the smallest eigenvalue will converge exponentially at a rate given by κ = λ M /λ 
+}
the conditioning number.
It follows that if the conditioning number of the NTK associated with a neural network diverges then it will become untrainable and so we use κ as a metric for trainability.
We will see that at large depths, the spectrum of Θ train,train typically features a single large eigenvalue, λ max , and then a gap that is large compared with the rest of the spectrum.
We therefore will often refer to a typical eigenvalue in the bulk as λ rest and approximate the condition number as κ = λ max /λ rest .
In the large-depth limit we will see that Θ ( l ) converges to Θ ∗ independent of the data distribution. In this case Θ ∗ test,train will be a rank-1 constant matrix.
As such, the mean prediction deﬁned by Equation 5 completely fails to generalize.
We deﬁne the ﬁnite depth correction to the inﬁnite depth predictor 2 , [Equation] By the triangle inequality, the generalization error is lower bounded by
[Equation]
(cid:107) P (Θ ∗ )
Y train − Y test (cid:107) 2 is a constant independent of the test inputs and Equation 10 is large if (cid:107) ∆ ( l ) Y Train (cid:107) 22 is too small.
Therefore, a necessary condition for the network to generalize is that there exists some ρ > 0 such that
[Equation]
As such, we use ∆ ( l ) as a metric for generalization in this paper.
Our goal is therefore to characterize the evolution of the two metrics κ ( l ) and ∆ ( l ) in l .
We follow the methodology outlined in Schoenholz et al.
; Xiao et al.
(2018) to explore the spectrum of the NTK as a function of depth.
We will use this to make precise predictions relating trainability and generalization to the hyperparameters ( σ w , σ b , l ) .
Our main results are summarized in Table 1 which describes the evolution of λ ( l )max (the largest eigenvalue of Θ ( l ) ), λ ( l ) rest (the remaining eigenvalues), κ ( l ) , and ∆ ( l ) in three different phases (ordered, chaotic, and the phase transition) and their dependence on m , the size of the training set, the choices of architectures: FCN, CNN-F (convolution with ﬂattening) and CNN-P (convolution with pooling), and size, d , of the window in the pooling layer (which we always take to be the penultimate layer). We give a brief derivation of these results in Section 4 followed by a more detailed discussion in the appendix. However, it is useful to ﬁrst give a qualitative overview of the phenomenology. In the ordered phase, λ ( l )max → m ( q ∗ − lχ l 1 ) and λ ( l ) rest → lχ l 1 .
At large depths since χ 1 < 1 it follows that κ ( l ) → mq ∗ / ( lχ l 1 ) and so the condition number diverges exponentially quickly.
Thus, in the ordered phase we expect networks not to be trainable (or, speciﬁcally, the rate at which they learn will grow exponentially in their depth).
The predictor scales as lχ l 1 which goes to zero at the same rate as the divergence of κ ( l ) ; thus, in the ordered phase networks fail to train and generalize simultaneously.
By contrast in the chaotic phase we see that there is no gap between λ ( l )max and λ ( l ) rest and networks become perfectly conditioned and are trainable everywhere. However, in this regime we see that the predictor scales as l ( χ c ∗ /χ 1 ) l . Since, by deﬁnition, in the chaotic phase χ c ∗ < 1 and χ 1 >it follows that ∆ ( l ) → 0 over a depth ξ = − 1 / log( χ c ∗ /χ 1 ) . In the chaotic phase networks fail to generalize at a ﬁnite depth but remain trainable indeﬁnitely.
Finally, notice that introducing pooling modestly augments the depth over which networks can generalize in the chaotic phase but reduces the depth in the ordered phase.
We will explore all of these predictions in detail in section 5.  L ARGE -D EPTH A SYMPTOTICS OF THE NNGP AND NTK We now give a brief derivation of the results in table 1. To simplify the notation we will discuss fully-connected networks and then extend the results to CNNs with pooling (CNN-P) and without pooling (CNN-F).
Details of these two cases can be found in the appendix.
We will focus on the NTK here since Schoenholz et al. ; Xiao et al.
(2018) contains a detailed description of the NNGP in this case.
As in sec. 2, we will be concerned with the ﬁxed points of Θ as well as the linearization of Equation 7 about its ﬁxed point.
Recall that the ﬁxed point structure is invariant within a phase so it sufﬁces to consider the ordered phase, the chaotic phase, and the critical line separately.
In cases where a stable ﬁxed point exists, we will describe how Θ converges to the ﬁxed point.
We will see that in the chaotic phase and on the critical line, Θ has no stable ﬁxed point and in that case we will describe its divergence.
As above, in each case the ﬁxed points of Θ have a simple structure with Θ ∗ =
p ∗ ((1 − ˆ c ∗ )
Id + ˆ c ∗ 11 T ) .
To simplify the forthcoming analysis, without a loss of generality, we assume the inputs are normalized to have variance q ∗ 3 .
As such, we can treat T and ˙ T , restricted on {K ( l ) } l , as a point-wise functions, since
[Equation]
Since the off-diagonal elements approach the same ﬁxed point at the same rate, we use q ( l ) ab ≡ K ( l ) ( x, x (cid:48) ) and p ( l ) ab ≡ Θ ( l ) ( x, x (cid:48) ) to denote any off diagonal entry of K ( l ) and Θ ( l ) respectively.
We will similarly use q ∗ ab and p ∗ ab to denote the limits, lim l →∞ q ( l ) ab = q ∗ ab = c ∗ q ∗ and lim l →∞ p ( l ) ab = p ∗ ab = ˆ c ∗ p ∗ . Using the above notation, Equation 7 and Equation 2 become
[Equation] where p ( l ) ≡ Θ ( l ) ( x, x ) and q ( l ) = K ( l ) ( x, x ) . In what follows, we split the discussion into three parts according to the values of χ 1 ≡ σ 2 ω ˙ T ( q ∗ ) recalling that in Poole et al.
; Schoenholz et al.
(2017) it was shown that χ 1 controls the ﬁxed point structure. 4. T HE C HAOTIC P HASE χ 1 > 1 :
The chaotic phase is so-named because q ∗ ab /q ∗ < 1 so that similar inputs become more uncorrelated as they pass through the network.
In this phase, the diagonal entries of Θ ( l ) grow exponentially and the off-diagonal entries converge to a ﬁxed value.
Indeed, Equation 14 implies,
[Equation]
which diverges exponentially.
To ﬁnd the limit of the off-diagonal terms, deﬁne χ c = σ 2 ω ˙ T ( q ∗ ab ) which was shown to control convergence of the q ( l ) ab and is always less than 1 (Schoenholz et al., 2017; Xiao et al., 2018). Let l → ∞ in Equation 13, we ﬁnd that
[Equation] The rate of convergence of p ∗ ab is O ( lχ lc ) (see Section A in the appendix).
Since the diagonal terms diverge and the off-diagonal terms are ﬁnite it follows that in very deep networks in the chaotic phase, ( p ( l ) ) − 1 Θ ( l ) → Id .
Thus, in the chaotic phase, the spectrum of the NTK for very deep networks approaches the diverging constant multiplying the identity.
From Equation 4 this implies that optimization in the chaotic phase should be easy since κ ( l ) → 1 (provided numerical precision issues from the prefactor do not become problematic).
However, computing the mean prediction on test points {+and noticing that P (Θ ∗ ) Y train = 
+}
we ﬁnd (see Section B for the derivation),
[Equation]
It follows that in the chaotic phase the networks predictions on unseen data to converge to 0 exponentially quickly in the depth.
Since Equation 17 decays like O ( l ( p ( l ) ) − 1 χ lc ) , we expect the network fails to generalize after O ( ξ ∗ ) layers, where ξ ∗ = − 1 / (log χ c − log χ 1 ) 4 . In summary, for wide networks, in the chaotic phase as the depth increases optimization becomes increasingly easy but the generalization performance degrades and eventually the network fails completely away from the training set after O ( ξ ∗ ) layers. Therefore, in the chaotic phase, deep network memorizes the training data.
We will conﬁrm this prediction for both kernel prediction and neural network training in the experimental results; see Fig 3. 4. T HE O
[Equation]
The ordered phase is deﬁned by the stable ﬁxed point with q ∗ ab /q ∗ = 1 ; in this case, disparate inputs will end up converging to the same output at the end of the network.
In the ordered phase, Equation 14 implies that all the diagonal entries of Θ converge to the same value,
[Equation]
However, as with the NNGP kernel, the off-diagonal terms of the NTK, p ( l ) ab , will also converge to the value on the diagonal, p ∗ .
It follows that the limiting kernels have the form Θ ∗ =
p ∗ 11 T and K ∗ =
q ∗ 11 T .
Thus, the limiting kernels are highly singular and feature only one non-zero eigenvalue.
Since the limit is singular, we must linearize the dynamics about the ﬁxed point to gain insight into the limiting behavior of the network.
To compute the corrections, let
[Equation]
The diagonal correction can be obtained directly from Equation 18 and we ﬁnd that (cid:15) ( l ) =
and δ ( l ) =
χ l +11 1 − χ 1 q ∗ .
To compute correction of the off-diagonals, we linearize the equation around the ﬁxed point to ﬁnd that asymptotically (see Section A),
[Equation]
where χ 2 = σ 2 ω ¨ T ( q ∗ ) .
While the NNGP and NTK feature the same exponential rate of convergence set by χ 1 , we see that terms in the off-diagonal terms of the NTK feature polynomial corrections.
Θ ( l ) has (approximately) two eigenspaces.
The ﬁrst eigenspace comes from the single non-zero eigenvalue at the ﬁxed point and it is very close to the DC mode (i.e. all entries of the eigenvector are equal to 1) with eigenvalue
[Equation]
(22) i.e. is the sum of one row, where m is the size of the dataset.
The second eigenspace comes from lifting the degenerate zero-modes when l < ∞ and it has dimension ( m − 1) with eigenvalue λ ( l ) rest
≈ − δ ( l ) + δ ( l ) ab
= O ( lχ l 1 ) → 0 , which goes to zero exponentially over depth l .
The eigenvalues of K ( l ) have a similar distribution with λ ( l )max ≈ mq ∗ − ( m − 1) (cid:15) ( l ) ab and λ ( l ) rest
= O ( χ l 1 ) .
Thus the conditioning number, κ ( l ) , of both Θ ( l ) and K ( l ) diverges exponentially as O ( χ − l 1 l − 1 ) and O ( χ − l 1 ) respectively.
As discussed above, there is a polynomial correction in the conditioning number of the
NTK that slightly improves its conditioning.
Since Θ ∗ is singular, we insert a diagonal regularization term σ Id into Θ train, train of the linear predictor Equation 6, where σ is a positive constant independent from l and χ 1 .
We ﬁnd ∆ ( l ) =
O σ ( lχ l 1 ) ; see Section B for the derivation.
In summary, in the ordered phase, ξ 1 = − 1 / log χ 1 (for simplicity, we ignore the polynomial correction) governs both trainability and generalizability of the predictor.
[Equation]
On the critical line both the diagonal and the off-diagonal terms of Θ ( l ) diverge linearly in the depth while K ( l ) converges to q ∗ 11 T .
From Equation 14 we see immediately that the diagonal terms are given by q ( l ) =
q ∗ and p ( l ) = lq ∗ .
To compute the correction of the off-diagonals, we keep the [Figure] deﬁnition of (cid:15) ( l ) ab unchanged but deﬁne δ ( l ) ab slightly differently to the above as δ ( l ) ab = p ( l ) ab − lq ∗ to take into account the linear divergence at large depths.
Taylor expanding to second order
we ﬁnd,
[Equation]
Thus for large l , Θ ( l ) has the following form p ( l ) = lq ∗ and p ( l ) ab = 13 lq ∗ + O (1) .
As in the ordered phase, for large l it follows that Θ ( l ) essentially has two eigenspaces: one has dimension one and the other has dimension ( m − 1) with
[Equation]
and the condition number κ ( l ) = m
+ [Equation]
→ [Equation] as l → ∞ .
Unlike the chaotic and ordered phases, κ ( l ) converges with rate O
( l − 1 ) .
The K ( l ) has λ ( l )max
= mq ∗ + m O
( l − 1 ) and λ ( l ) rest ≈ 2 χ 2 l − 1 and the condition number κ ( l ) diverges linearly with slope mχ 2 / 2 .
A similar calculation gives ∆ ( l ) =
O ( l − 1 ) on the critical line.
In summary, κ ( l ) converges to a ﬁnite number and the network ought to be trainable for arbitrary depth but the residual predictor ∆ ( l ) decays polynomially, explaining why critically initialized networks with thousands of layers could still generalize (Xiao et al., 2018).
101010 1 ρ 0.
a cc u r a c y depth 101010 1 ρ 0.
a cc u r a c y depth 101010 1 ρ 0.
a cc u r a c y depth 101010 1 ρ 0.
a cc u r a c y depth
4. R EMARKS
We end this section with a couple remarks.
The above theory holds for CNNs; see Section D.
In the large depth setting, the NTK of CNNs without pooling is essentially the same as the NTK of FCNs.
In the ordered phase, Adding a dropout layer could signiﬁcantly improve trainability of a network.
For example, adding dropout to the penultimate layer, the condition number κ ( l ) will converge to a ﬁnite number rather than diverge exponentially; see (f) in Figure 1 and Equation 85 in the appendix.
E XPERIMENTS
In this section, we provide empirical results to support the theoretical results in Section 4.
Figureis generated using synthetic data and all other plots are generated using CIFAR-10 with MSE as the loss function.
Evolution of κ ( l ) (Figure 1).
We randomly sample inputs with shapes ( m, k 2 × 3) for FCN and ( m, k, k, 3) for CNN-F/CNN-P, where m ∈ { 12 , 20 } and k ∈ { 6 , 10 } .
We compute the exact
NTK with activation function Erf using the Neural Tangents library (Novak et al., 2019a).
We see excellent agreement between the theoretical calculation of κ ( l ) in Section 4 (summarized in Table 1) and the experimental results Figure 1.
Maximum Learning Rates (Figure 2 top).
In practice, given a set of hyper-parameters of a network, knowing the range of feasible learning rates is extremely valuable. As discussed above, in the inﬁnite width setting, Equation 4 implies the maximal convergent learning rate is given by η theory ≡ 2 /λ ( l )max . We argue that η theory is a good prediction for the maximal convergent learning rate for wide network.
To test this statement, we apply SGD to train a collection of fully-connected networks on CIFAR-10 using 1 k training samples with the following conﬁgurations: (1) width: 2048 (2) σ b = 0 .
43 ﬁxed, (3) depths: l = 5 , 10 , 20 , 40 , (4) 10 different values of σ ω moving from the ordered phase (blue) to the chaotic phase (red) (5) 10 different learning rates η = ρη theory , with ρ ∈ [10 − 1 , 10 1 ] . Overall, we see excellent agreement for depths less or equal to 20 and reasonable good agreement for depth 40. We point out that the degradation of the agreement for larger depth may due to the fact that the ﬁnite width NTK becomes more stochastic as the ratio between depth and width increases (Hanin & Nica, 2019). Note that Table 1 tells that, as depth increases, η theory should decays exponentially and linearly in the chaotic and critical phases resp.
and remain roughly a constant in the ordered phase.
Trainability vs Generalization (Figure 3 top).
Our theoretical result suggests that in the deep chaotic regime ( χ l 1 is large) training becomes easier but the network can not generalize. On the other hand, the network can generalize but training becomes much more difﬁcult as one moves towards the deep ordered region because κ ( l ) blows up exponentially. To conﬁrm this claim, we conduct an experiment using 16k training samples from CIFAR-10 with 20 × 20 different ( σ ω , l ) conﬁgurations. We train each network using SGD with batch size b = 1024 and learning rate η = 0 .
3 η theory .
Deep in the chaotic phase we see that all conﬁgurations reach perfect training accuracy but the network
d e p t h 12 ξ 1 4 ξ ∗ , 10 ξ ∗ 0.5 1.0 1.5 2.0 2.5 3.0 3.5 4.0 σ 2 w 12 ξ 1 4 ξ ∗ , 10 ξ ∗ 0.
0.5 1.0 1.5 2.0 2.5 3.0 3.5 4.0 σ 2 w
CNN-P 12 ξ 1 (12+log64) ξ 1 (10+log64)
ξ ∗ 10 ξ ∗ 0.5 1.0 1.5 2.0 2.5 3.0 3.5 4.0 σ 2 w
CNN-F 12 ξ 1 (12+log64) ξ 1 (10+log64)
ξ ∗ 10 ξ ∗ 0.5 1.0 1.5 2.0 2.5 3.0 3.5 4.0 σ 2 w
Difference 12 ξ 1 (12+log64) ξ 1 (10+log64) ξ ∗ 10 ξ ∗ 0.
CNN-P, due to the fact that pooling increases the spectra gap by a factor of d . completely fails to generalize in the sense test accuracy approaches 10% .
However, in the ordered phase although the training accuracy degrades, generalization improves.
The network eventually becomes untrainable after O ( ξ 1 ) layers.
In both phases we see that the depth scales, ξ 1 and ξ ∗ respectively, perfectly capture the transition from generalizing to overﬁtting.
CNN-P v.s. CNN-F: spatial correction (Figure 3 bottom).
We compute the test accuracy using the analytic NTK predictor Equation 5, which corresponds to the test accuracy of ensemble of gradient descent trained neural networks taking the width to inﬁnity.
We choose 1 k training points, ﬁx σ 2 b , and choose 20 × 20 different ( σ ω , l ) conﬁgurations.
We plot the test performance of CNN-P and
CNN-F and the performance difference in Fig 3.
Remarkably, the performance of both CNN-P and
CNN-F are captured by ξ 1 = − 1 / log( χ 1 ) in the ordered phase and by ξ ∗ = − 1 / (log ξ c − log ξ 1 ) in the chaotic phase.
We see that the test performance difference between CNN-P and CNN-F exhibits a region in the ordered phase (a blue strip) where CNN-F outperforms CNN-P by a large margin. This performance difference is due to the correction term d as predicted by the ∆ ( l ) -row of Table 1.
C ONCLUSION AND F UTURE WORK
In this work, we identify several quantities ( λ max , λ rest , κ , and ∆ ( l ) ) related to the spectrum of the NTK that control trainability and generalization of deep networks.
We offer a precise characterization of these quantities and provide substantial experimental evidence supporting theoretical results.
In future work, we would like to extend our framework to other architectures, e.g., ResNet (with batch-norm), attention model.
Understanding the implication of the sub-Fourier modes in the NTK to the test performance of CNN is also an important research direction. A S IGNAL PROPAGATION OF NNGP AND NTK Recall that
[Equation] A. C ORRECTION OF THE OFF - DIAGONALS IN THE CHAOTIC / ORDERED PHASE
Applying Taylor’s expansion to the ﬁrst equation of 25 gives
[Equation]
With χ c = σ 2 ω ˙ T ( q ∗ ab )
[Equation]
we have [Equation] [Equation] Similarly
[Equation]
applying Taylor’s expansion to the second equation of 25 gives
[Equation] 31) where χ c, 2 = σ 2 ω ¨ T ( q ∗ ab ) .
This implies
[Equation] [Equation]
Note that δ ( l ) ab contains a polynomial correction term and decays like lχ lc . The correction to the ﬁxed points in the ordered phase could be obtained using the same calculation:
[Equation] A. C ORRECTION OF THE OFF - DIAGONALS ON THE CRITICAL LINE .
We have χ 1 = 1 on the critical line.
We need to expand the ﬁrst equation of 25 to the second order [Equation] Here we assume T has a continuous third derivative. The above equation implies
[Equation]
Then
[Equation]
[Equation]
[Equation]
Plugging Equation 37 into the above equation gives
[Equation]
A. R ELU
We only consider the critical initialization σ 2 ω
= 2 and σ 2 b = 0 .
Using the equations in Appendix C of (Lee et al., 2019) gives
[Equation]
and taking the derivative w.r.t.
[Equation]
This is enough to conclude (similar to the above calculation)
[Equation]
Recall that q ( l ) = 1 and p ( l ) =
l for Relu network with σ 2 ω = 2 .
and σ b = 0 .
Therefore
[Equation]
A. R ESIDUAL R ELU
We consider the following “continuum” residual network
[Equation]
47) where t denotes the number of layer and dt > 0 is sufﬁciently small and W and b are the weights and biases.
We also set σ 2 ω = 2 (i.e. E [ WW T ] = 2 Id ) and σ 2 b = 0 (i.e. b = 0 ).
The NNGP and
NTK have the following form
[Equation]
[Equation]
[Equation]
Taking the limit dt → 0 gives
[Equation]
[Equation]
Using the fact that q 0 = 1 (i.e. the inputs have unit variance), we can compute the diagonal terms q ( t ) =
e t and p ( t ) =
te t .
Let q ( t ) ab =
e t c ( t ) ab , applying the above fractional Taylor expansion to T and ˙ T , we have
[Equation]
[Equation]
Solving this gives
[Equation]
Thus the condition number is m/ 3 + 1 .
This is the same as the non-residual Relu case discussed above.
B A SYMPTOTIC OF ∆ ( l )
To keep the notation simple, we denote X d = X train , Y d = Y train , Θ td = Θ test, train , Θ dd =
Θ train, train .
Recall that
[Equation]
We split our calculation into three parts.
B. C HAOTIC PHASE
In this case the diagonal p ( l ) diverges exponentially and the off-diagonals p ( l ) ab converges to a bounded constant p ∗ ab .
Thus P (Θ ∗ ) Y d = 0 .
We expand Θ ( l ) about its ﬁxe point
[Equation]
In the last equation
[Equation]
we have used the fact 11 T Y d = 0 and Θ ∗ td Y d = 0 since Y d is balanced (i
[Equation]
e. containing the same number of positive (+1) and negative (-1) labels.)
Therefore
[Equation]
[Equation]
B. O RDER - TO - CHAOS PHASE
Note that in this phase, both the diagonals and the off-diagonals diverge linearly.
In this case
[Equation]
Here we use 1 d to denote the all ‘1’ (column) vector with length equal to the number of training points in X d and 1 t is deﬁned similarly.
Note that the constant matrix B is invertible.
By Equation
[Equation]
[Equation]
[Equation]
[Equation]
Thus
[Equation]
[Equation]
Ordered Phase
In this case Θ ∗ is a rank one matrix.
We add a diagonal regularization term and deﬁned
[Equation]
69) where σ > 0 is a positive constant independent of the hyper-parameters ( σ w , σ b , l ) .
Let B σ =
Θ ∗ + σ Id .
Then
[Equation] [Equation] [Equation] [Equation]
C D ROPOUT
In this section, we investigate the effect of adding a dropout layer to the penultimate layer.
Let 0 < ρ ≤ 1 and γ ( L ) j ( x ) be iid random variables
[Equation]
For 0 ≤ l ≤ L − 1 ,
[Equation] i [Equation]
and for the output layer,
[Equation]
[Equation]
i
[Equation]
where W ( l ) ij and b ( l ) i are iid Gaussians N (0 , 1) .
Since no dropout is applied in the ﬁrst L layers, the NNGP kernel K ( l ) and Θ ( l ) can be computed using Equation 2 and Equation 7.
Let K ( L +1) ρ and Θ ( L +1)
ρ denote the NNGP and NTK of the ( L +1) -th layer.
Note that when ρ = 1 , K ( L +1)= K ( L +1) and Θ ( L +1)= Θ ( L +1) .
We will compute the correction induced by ρ < 1 .
The fact [Equation] implies that the NNGP kernel K ( L +1) ρ [Equation]
Schoenholz et al., 2017) is K ( L + [Equation] ρ
[Equation]
Now we compute the NTK Θ ( L +1) ρ , which is a sum of two terms Θ ( L +1)
ρ
[Equation]
∂θ
[Equation]
≤ L
[Equation]
(cid:33) T   .
Here θ ( L +1) denote the parameters in the ( L + 1) layer, namely, W ( L +1) ij and b ( L +1)
i and θ ( ≤ L ) the remaining parameters.
Note that the ﬁrst term is in Equation 78 is equal to K ( L +1) ρ ( x, x (cid:48) ) .
Using the chain rule, the second term is equal to σ 2 ω
[Equation]
j ( x (cid:48) )
∂z ( L ) j
( x ) ∂θ ( ≤ L ) (cid:32) ∂z ( L ) k
( x (cid:48) )
∂θ ( ≤ L ) (cid:33) T 
 (79) =
σ 2 ω ρ 2 N ( L ) E  
N ( L ) (cid:88)
j ˙ φ ( z ( L ) j
( x )) γ ( L ) j ( x ) ˙
φ ( z ( L ) j
( x (cid:48) ))
γ ( L ) j ( x (cid:48) )
∂z ( L ) j
( x ) ∂θ ( ≤ L ) (cid:32) ∂z ( L )
j ( x (cid:48) )
∂θ ( ≤ L ) (cid:33) T 
 (80) =
σ 2 ω ρ 2 E (cid:104) γ ( L ) j ( x ) γ ( L ) j ( x (cid:48) )
(cid:105) E [ ˙ φ ( z ( L ) j
( x )) ˙
φ ( z ( L ) j
( x (cid:48) ))]
E   ∂z ( L ) j
( x ) ∂θ ( ≤ L ) (cid:32) ∂z ( L )
j ( x (cid:48) )
∂θ ( ≤ L ) (cid:33) T 
 (81) =
  σ 2 ω ˙ T ( K ( L ) ( x, x (cid:48) ))
Θ ( L ) ( x, x (cid:48) )
if x (cid:54) = x (cid:48)
1 ρ σ 2 ω ˙
T ( K ( L ) ( x, x ))
Θ ( L ) ( x, x )
if x = x (cid:48) .
In sum, we see that dropout only modiﬁes the diagonal
[Equation]
In the ordered phase, we see
[Equation]
and the condition number
[Equation]
D C ONVOLUTIONS
General setup.
For simplicity of presentation we consider 1D convolutional networks with circular padding as in Xiao et al.
We will see that this reduces to the fully-connected case introduced above if the image size is set to one and as such we will see that many of the same concepts and equations carry over schematically from the fully-connected case.
The theory of two-dimensional convolutions proceeds identically but with more indices.
Random weights and biases.
The parameters of the network are the convolutional ﬁlters and biases, ω lij,β and µ li , respectively, with outgoing (incoming) channel index i ( j ) and ﬁlter relative spatial location β ∈ [ ± k ] ≡ {− k, . . .
. , k } .
5 As above, we will assume a Gaussian prior on both the ﬁlter weights and biases,
[Equation]
As above, σ 2 ω and σ 2 b are hyperparameters that control the variance of the weights and biases respectively.
n l is the number of channels (ﬁlters) in layer l , 2 k + 1 is the ﬁlter size.
Inputs, pre-activations, and activations.
Let X denote a set of input images.
The network has activations y l ( x ) and pre-activations z l ( x ) for each input image x ∈ X ⊂ R n 0 d , with input channel count n 0 ∈ N , number of pixels d ∈ N , where
[Equation]
φ : R → R is a point-wise activation function.
Since we assume circular padding for all the convolutional layers, the spacial size d remains constant throughout the networks until the readout layer.
For each l > 0 , as min { n 1 . . . , n l − 1 } → ∞ , for each i ∈ N , the pre-activation converges in distribution to d -dimensional Gaussian with mean 0 and covariance matrix K ( l ) , which can be computed recursively (Novak et al., 2019b; Xiao et al., 2018)
[Equation]
Here K ( l ) ≡ [ K ( l ) α,α (cid:48)
( x, x (cid:48) )]
α,α (cid:48) ∈
[ d ] ,x,x (cid:48) ∈X , T is a non-linear transformation related to its fullyconnected counterpart, and A a convolution coupled with a shift term acting on X d × X d PSD matrices
[Equation]
[Equation]
D. T HE N EURAL T ANGENT K ERNEL
To understand how the neural tangent kernel evolves with depth, we deﬁne the NTK of the l -th hidden layer to be ˆΘ ( l )
[Equation]
where θ ≤ l denotes all of the parameters in layers at-or-above the l ’th layer.
It does not matter which channel index i is used because as the number of channels approach inﬁnity, this kernel will also converge in distribution to a deterministic kernel Θ ( l +1)
(Yang, 2019), which can also be computed recursively in a similar manner to the NTK for fully-connected networks as,
[Equation]
[Equation] where ˙ T is given by Equation 89 with φ replaced by its derivative φ (cid:48) .
We will also normalize the variance of the inputs to q ∗ and hence treat T and ˙ T as pointwise functions.
We will only present the treatment in the chaotic phase to showcase how to deal with the operator A .
The treatment of other phases are similar.
Note that the diagonal entries of K ( l ) and Θ ( l ) are exactly the same as the fully-connected setting.
We only need to consider the off-diagonal terms.
Letting l → ∞ in Equation 92 we see that all the off-diagonal terms also converge .
Let ¯
A = σ − 2 w ( A − σ 2 b ) , be the normalized convolution operator.
Note that A does not mix terms from different diagonals and it sufﬁces to handle each off-diagonal separately.
Let (cid:15) ( l ) ab and δ ( l ) ab denote the correction of the j -th diagonal of K ( l ) and Θ ( l ) to the ﬁxed points.
Linearizing Equation 88 and Equation 92 gives
[Equation]
Next let { ρ α } α be the eigenvalues of ¯ A and (cid:15) ( l ) ab,α and δ ( l +1) ab,α be the projection of (cid:15)
( l +1) ab and δ ( l ) ab onto the α -th eigenvector of ¯ A .
Then for each α ,
[Equation]
which gives
[Equation]
Therefore, the correction Θ ( l ) − Θ ∗ propagates independently through different Fourier modes.
In each mode, up to a scaling factor ρ α , the correction is the same as the correction of its FC counterpart.
Since the subdominant modes (with | ρ α | < 1 ) decay exponentially faster than the dominant mode (with ρ α = 1 ), for large depth, the NTK of CNN is essentially the same as that of its FC counterpart. D. T HE EFFECT OF POOLING AND FLATTENING OF CNN S With the bulk of the theory in hand, we now turn our attention to CNNs.
We show in the appendix that the dominant mode in CNNs behaves exactly like the fully-connected case, however we will see that the readout can signiﬁcantly affect the spectrum. The NNGP and NTK of the l -th hidden layer CNN are 4D tensors K ( l ) α,α (cid:48)
( x, x (cid:48) ) and Θ ( l ) α,α (cid:48) ( x, x (cid:48) ) , where α, α (cid:48) ∈ [ d ] ≡ [0 , 1 , . . . , d − 1] denote the pixel locations.
To perform tasks like image classiﬁcation or regression, “ﬂattening” and “pooling” (more precisely, global average pooling) are two popular readout strategies that transform the last convolution layer into the logit layer. The former strategy “ﬂattens” an image of size ( d, N ) into a vector in R dN and stacks a fully-connected layer on top. The latter projects the ( d, N ) image into a vector of dimension N via averaging out the spatial dimension and then stacks a fully-connected layer on top. The actions of “ﬂattening” and “pooling” on the image correspond to computing the mean of the trace and the mean of the pixel-to-pixel covariance on the NNGP/NTK, respectively, i.e.,
[Equation] where Θ ( l )ﬂatten ( Θ ( l )pool ) denotes the NTK right after ﬂattening (pooling) the last convolution. We will also use Θ ( l ) fc to denote the NTK of FC. K ( l )ﬂatten and K ( l )pool are deﬁned similarly.
As discussed above, in the large depth setting, all the diagonals Θ ( l ) α,α ( x, x ) = p ( l ) (since the inputs are normalized to have variance q ∗ for each pixel) and similar to Θ ( l ) fc , all the off-diagonals Θ ( l ) α (cid:48) ,α ( x, x (cid:48) ) are almost equal (in the sense they have the same order of correction to p ∗ ab if exists.) Without loss of generality, we assume all off-diagonals are the same and equal to p ( l ) ab (the leading correction of q ( l ) ab for CNN and FCN are of the same order.) Applying ﬂattening and pooling, the
NTKs become [Equation] respectively. As we can see, Θ ( l ) ft is essentially the same as its FCN counterpart Θ ( l ) fc up to subdominant Fourier modes which decay exponentially faster than the dominant Fourier modes.
Therefore the spectrum properties of Θ ( l ) ft and Θ ( l ) fc are essentially the same for large l .
However, pooling alters the NTK/NNGP spectrum in an interesting way.
On the critical line, asymptotically, λ ( l )max ≈ ( md + 2) q ∗ l/ (3 d ) and λ ( l ) rest ≈ 2 q ∗ l/ (3 d ) , and κ ( l ) = md + + md O ( l − 1 ) .
Here we use blue color to indicate the changes of such quantities against their Θ ( l )ﬂatten counterpart.
Thus pooling decreases λ ( l ) rest roughly by a factor of d and increases the condition number by a factor of d comparing to ﬂattening.
In the chaotic phase, pooling does not change the off-diagonals q ( l ) ab
= O (1) but does slow down the growth of the diagonals by a factor of d , i.e. p ( l ) =
O ( χ l 1 /d ) .
This suggests, in the chaotic phase, there exists a transient regime of depths, where CNN-F hardly perform while CNN-P performs well.
In the ordered phase, the pooling does not affect λ ( l )max much but does decrease λ rest by a factor of d and the condition number grows approximately like dχ − l 1 , d times bigger than its ﬂattening and fully-connected network counterparts.
This suggests the existence of a transient regime of depths, in which CNN-F outperforms CNN-P. This might be surprising because it is commonly believed CNN-P usually outperforms CNN-F.