Title: DISENTANGLING TRAINABILITY AND GENERALIZATION IN DEEP LEARNING

Abstract: 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 simplifies 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 simplifies 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 simplifies 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 find 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 finding excellent agreement on real datasets.

1. INTRODUCTION:

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. (2018) . 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, significant theoretical progress has been made on several fronts that have shown promise in making neural network design more systematic. In particular, in the infinite 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 infinite 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 infinite-width limit are Gaussian Processes (GPs). Recently, there has been renewed interest in studying random, infinite, networks starting with concurrent work on "conjugate kernels" (Daniely et al., 2016; Daniely, 2017) and "mean-field 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 infinite-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-defined, 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 sufficiently 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 fisher (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 significant 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, significant 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 infinite-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:

1. 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 finite depth and infinite depth NTK predictions, which is related to the model's ability to generalize: the network fails to generalize if ∆ (l) is too small.

3. We show that the ordered and chaotic phases identified 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.

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

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

FC/CNN-F, CNN-P NTK Ordered χ 1 < 1 Critical χ 1 = 1 Chaotic χ 1 > 1 λ (l) max mq * -mO(lχ l 1 ) md+2 3d lq * + mO(1) O(χ l 1 )/d λ (l) rest O(lχ l 1 )/d 2 3d q * l + 1 d O(1) O(χ l 1 )/d κ (l) dmq * O(χ -l 1 /l) md+2 2 + dmO(l -1 ) → 1 ∆ (l) O(lχ l 1 )/d dO(l -1 ) dO(l(χ c * /χ 1 ) l )

Table 1 : Evolution of the NTK spectra and ∆ (l) . The NTKs of FCN and CNN-F are essentially the same and the scaling of λ

EQUATION

rest , κ (l) , and ∆ (l) for these networks is written in black. Corrections to these quantities due to the addition of an average pooling layer with window size d is written in blue.

network Saxe et al. (2013) . Here, we note that this is true only for networks that are initialized critically, i.e. on the order-to-chaos phase boundary.

2. BACKGROUND:

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 specifics -such as the phase diagram -can vary substantially. For simplicity, we will take the width of the hidden layers to infinity 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 ) = E[z (l) i (x)z (l) i (x )].

Inspecting Equation 75, we see that K (l+1) (x, x ) can be computed in terms of K (l) (x, x ) as

K (l+1) (x, x ) = σ 2 w E (z,z )∼N (0,K (l) (x,x )) [φ(z)φ(z )] + σ 2 b ≡ σ 2 w T (K (l) (x, x )) + σ 2 b .

(2) for T , an appropriately defined operator from the space of positive semi-definite matrices to itself.

Equation 2 describes a dynamical system on positive semi-definite matrices K(x, x ). It was shown in Poole et al. (2016) that fixed points, K * (x, x ), of these dynamics exist such that lim l→∞ K (l) (x, x ) = K * (x, x ) with K * (x, x ) = q * [δ x,x + c * (1 -δ x,x )] independent of the inputs x and x . The values of q * and c * are determined by the hyperparameters, σ w and σ b . However Equation 2 admits multiple fixed points (e.g. c * = 0, 1) and the stability of these fixed points plays a significant role in determining the properties of the network. Generically, there are large regions of the (σ w , σ b ) plane in which the fixed-point structure is constant punctuated by curves, called phase transitions, where the structure changes.

The rate at which K(x, x ) approaches or departs K * (x, x ) can be determined by expanding Equation 2 about its fixed point, δK(x, x ) = K(x, x ) -K * (x, x ) to find1 

EQUATION

which exhibits exponential convergence to -or divergence from -the fixed-point as δK (l) (x, x ) ∼ χ(x, x ) l where χ(x, x ) = σ 2 w Ṫ (K * (x, x )). Since K * (x, x ) does not depend on x or x it follows that χ(x, x ) will take on a single value, χ c * , whenever x = x . If χ c * > 1 then the fixed point is unstable and, as discussed above, there will be another fixed point that becomes stable, if χ c * < 1 then the fixed point is stable, and if χ c * = 1 then the hyperparameters lie at a phase transition. As was shown in Poole et al. (2016) , there is always a fixed-point at c * = 1 whose stability is determined by χ 1 . This defines the order-to-chaos transition. Note, that χ c * can be used to define 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 finite width NTK is defined by Θ(x, x ) = J(x)J(x ) T where J iα (x) = ∂ θα z L i (x) is the Jacobian evaluated at a point x. The main result in Jacot et al. (2018) was to show that in the infinite-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 infinitely wide network,

EQUATION

for train and test points respectively; see Section 2 in Lee et al. (2019) . Here Θ test, train denotes the NTK between the test inputs X test and training inputs X train and Θ train, train is defined similarly. Since Θ converges to Θ, the gradient flow 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 infinity 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 , Ṫ , and the NNGP as,

EQUATION

)

where Θ (l) is the NTK for the pre-activations at layer-l.

3. METRICS FOR TRAINABILITY AND GENERALIZATION AT LARGE DEPTH:

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

EQUATION

where λ i are the eigenvalues of Θ train,train and μt (X train ) = U µ t (X train ), Ỹtrain = U Y 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 8we see that the smallest eigenvalue will converge exponentially at a rate given by κ = λ M /λ 0 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 defined by Equation 5 completely fails to generalize. We define the finite depth correction to the infinite depth predictor2 ,

EQUATION

By the triangle inequality, the generalization error is lower bounded by

EQUATION

)

P (Θ * )Y train -Y test 2 is a constant independent of the test inputs and Equation 10 is large if ∆ (l) Y Train 2

2 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. ( 2017); 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 λ

EQUATION

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 flattening) 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 first 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, specifically, 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 definition, in the chaotic phase χ c * < 1 and χ 1 > 1 it follows that ∆ (l) → 0 over a depth ξ = -1/ log(χ c * /χ 1 ). In the chaotic phase networks fail to generalize at a finite depth but remain trainable indefinitely. 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.

4. LARGE-DEPTH ASYMPTOTICS 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. ( 2017); Xiao et al. (2018) contains a detailed description of the NNGP in this case. As in sec. 2, we will be concerned with the fixed points of Θ as well as the linearization of Equation 7 about its fixed point. Recall that the fixed point structure is invariant within a phase so it suffices to consider the ordered phase, the chaotic phase, and the critical line separately. In cases where a stable fixed point exists, we will describe how Θ converges to the fixed point. We will see that in the chaotic phase and on the critical line, Θ has no stable fixed point and in that case we will describe its divergence. As above, in each case the fixed points of Θ have a simple structure with Θ * = p * ((1 -ĉ * )Id + ĉ * 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 Ṫ , restricted on {K (l) } l , as a point-wise functions, since

EQUATION

)

Since the off-diagonal elements approach the same fixed point at the same rate, we use q

(l) ab ≡ K (l) (x, x ) and p (l)

ab ≡ Θ (l) (x, x ) 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 = ĉ * p * . Using the above notation, Equation 7and Equation 2become

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 ω Ṫ (q * ) recalling that in Poole et al. (2016) ; Schoenholz et al. ( 2017) it was shown that χ 1 controls the fixed point structure.

4.1. THE CHAOTIC PHASE χ 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 fixed value. Indeed, Equation 14 implies,

EQUATION

which diverges exponentially. To find the limit of the off-diagonal terms, define χ c = σ 2 ω Ṫ (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 find that

EQUATION

)

The rate of convergence of p * ab is O(lχ l c ) (see Section A in the appendix). Since the diagonal terms diverge and the off-diagonal terms are finite 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 = 0 we find (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 χ l c ), 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 confirm this prediction for both kernel prediction and neural network training in the experimental results; see Fig 3 . 4.2 THE ORDERED PHASE χ 1 = σ 2 ω Ṫ (q * ) < 1:

The ordered phase is defined by the stable fixed 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,

p (l) = q * χ l+1 1 -1 χ 1 -1 l→∞ ---→ p * = q * 1 1 -χ 1 < ∞ (18)

However, as with the NNGP kernel, the off-diagonal terms of the NTK, p

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 fixed 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 18and we find that (l) = 0 and δ (l) = χ l+1 1 1-χ1 q * . To compute correction of the off-diagonals, we linearize the equation around the fixed point to find 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 first eigenspace comes from the single non-zero eigenvalue at the fixed point and it is very close to the DC mode (i.e. all entries of the eigenvector are equal to 1) with eigenvalue

EQUATION

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) (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 find ∆ (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.

4.3. THE CRITICAL LINE χ:

1 = σ 2 ω Ṫ (q * ) = 1

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 14we 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 For example, in the yellow "12-6" , "12" represents the size of the dataset and "6" represents the dimension number (6 * 6 * 3 for FCNand (6, 6, 3) for CNN) (a) In the chaotic phase, κ (l) converges to 1 for all architectures. (b) We plot χ l 1 κ (l) , confirming κ explodes with rate χ l 1 /l in the ordered phase. In (c) and (d), the dashed lines representing the condition number κ (l) and solid lines the ratio between first and second eigenvalues. We see that, on the order-to-chaos transition, these two numbers converge to m+2 2 and dm+2 2 (horizontal lines) for FC/CNN-F and CNN-P respectively. In (e), we plot rates of convergence for CNN-P (solid) and CNN-F (dashed), confirming that pooling slows down the convergence of κ (l) by a factor of d. (f) Adding dropout to the penultimate layer prevents κ (l) from divergence in the ordered phase. The legends indicate the rate of the mask with ρ = 1 meaning keeping all activations. Horizontal lines are the limit of κ (l) computed in Equation 85(here m = 20 for all curves.) , definition of (l) ab unchanged but define δ (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 find,

EQUATION

Thus for large l, Θ (l) has the following form p (l) = lq * and p (l) ab = 1 3 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+2 2 + mO(l -1 ) → m+2 2 as l → ∞. Unlike the chaotic and ordered phases, κ (l) converges with rate O(l -1 ). The K (l) has λ (l) max = mq * +mO(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 finite 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) . ) is moving from the order phase to the chaotic phase. ρ = 1 is the theoretical prediction.

4.4. REMARKS:

We end this section with a couple remarks. (1) 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.

(2) In the ordered phase, Adding a dropout layer could significantly improve trainability of a network. For example, adding dropout to the penultimate layer, the condition number κ (l) will converge to a finite number rather than diverge exponentially; see (f) in Figure 1 and Equation 85 in the appendix.

5. EXPERIMENTS:

In this section, we provide empirical results to support the theoretical results in Section 4. Figure 1 is 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 infinite width setting, Equation 4implies 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 1k training samples with the following configurations: (1) width: 2048 (2) σ b = 0.43 fixed, (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 finite 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 difficult as one moves towards the deep ordered region because κ (l) blows up exponentially. To confirm this claim, we conduct an experiment using 16k training samples from CIFAR-10 with 20×20 different (σ ω , l) configurations. 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 configurations reach perfect training accuracy but the network 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 overfitting.

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 infinity. We choose 1k training points, fix σ 2 b , and choose 20 × 20 different (σ ω , l) configurations. 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 .

6. CONCLUSION AND FUTURE 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 SIGNAL PROPAGATION OF NNGP AND NTK:

Recall that

EQUATION

A.1 CORRECTION OF THE OFF-DIAGONALS IN THE CHAOTIC/ORDERED PHASE Applying Taylor's expansion to the first equation of 25 gives

EQUATION

2 ) (28)

EQUATION

With χ c = σ 2 ω Ṫ (q * ab ), we have

EQUATION

Similarly, applying Taylor's expansion to the second equation of 25 gives

EQUATION

where χ c,2 = σ 2 ω T (q * ab ). This implies

EQUATION

EQUATION

Note that δ

(l)

ab contains a polynomial correction term and decays like lχ l c . The correction to the fixed points in the ordered phase could be obtained using the same calculation:

EQUATION

A.2 CORRECTION OF THE OFF-DIAGONALS ON THE CRITICAL LINE.

We have χ 1 = 1 on the critical line. We need to expand the first 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

Plugging Equation 37 into the above equation gives

EQUATION

B.1 CHAOTIC 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 fixe point

EQUATION

In the last equation, we have used the fact 11 T Y d = 0 and Θ * td Y d = 0 since Y d is balanced (i.e. containing the same number of positive (+1) and negative (-1) labels.) Therefore

EQUATION

B.2 ORDER-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 defined similarly. Note that the constant matrix B is invertible. By Equation 41

EQUATION

Thus

EQUATION

Ordered Phase In this case Θ * is a rank one matrix. We add a diagonal regularization term and defined

EQUATION

where σ > 0 is a positive constant independent of the hyper-parameters (σ w , σ b , l). Let B σ = Θ * + σId. Then

EQUATION

C DROPOUT

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

γ (L) j (x) =

1, with probability ρ 0, with probability 1 -ρ.

(73)

For 0 ≤ l ≤ L -1, z (l+1) i (x) = σ w √ N (l) j W (l+1) ij φ(z (l) j (x)) + σ b b (l+1) i (74)

and for the output layer,

EQUATION

where

W (l) ij and b (l)

i are iid Gaussians N (0, 1). Since no dropout is applied in the first L layers, the NNGP kernel K (l) and Θ (l) can be computed using Equation 2 and Equation 7. Let K = Θ (L+1) . We will compute the correction induced by ρ < 1. The fact

EQUATION

implies that the NNGP kernel K (L+1) ρ (Schoenholz et al., 2017) is

K (L+1) ρ (x, x ) ≡ E[z (L+1) i (x)z (L+1) i (x )] =      σ 2 w T (K (L) (x, x )) + σ 2 b , if x = x 1 ρ σ 2 w T (K (L) (x, x)) + σ 2 b if x = x .

(77)

Now we compute the NTK Θ (L+1) ρ

, which is a sum of two terms

Θ (L+1) ρ (x, x ) = E   ∂z (L+1) i (x) ∂θ (L+1) ∂z (L+1) i (x ) ∂θ (L+1) T   + E   ∂z (L+1) i (x) ∂θ (≤L) ∂z (L+1) i (x ) ∂θ (≤L) T   . (78)

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 first term is in Equation 78 is equal to K (L+1) ρ (x, x ). Using the chain rule, the second term is equal to

EQUATION

In sum, we see that dropout only modifies the diagonal

     Θ (L+1) ρ (x, x ) = Θ (L+1) (x, x ) Θ (L+1) ρ (x, x) = 1 ρ Θ (L+1) (x, x) + (1 -1/ρ)σ 2 b (83)

In the ordered phase, we see

EQUATION

)

and the condition number

EQUATION

D CONVOLUTIONS

General setup. For simplicity of presentation we consider 1D convolutional networks with circular padding as in Xiao et al. (2018) . 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 filters and biases, ω l ij,β and µ l i , respectively, with outgoing (incoming) channel index i (j) and filter relative spatial location β ∈ [±k] ≡ {-k, . . . , 0, . . . , k}.5 As above, we will assume a Gaussian prior on both the filter 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 (filters) in layer l, 2k + 1 is the filter 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) α,α (x, x )] α,α ∈[d],

x,x ∈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

D.1 THE NEURAL TANGENT KERNEL:

To understand how the neural tangent kernel evolves with depth, we define 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 infinity, 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

where Ṫ is given by Equation 89with φ replaced by its derivative φ . We will also normalize the variance of the inputs to q * and hence treat T and Ṫ 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 92we see that all the off-diagonal terms also converge . Let Ā = σ -2 w (A -σ 2 b ), be the normalized convolution operator. Note that A does not mix terms from different diagonals and it suffices to handle each off-diagonal separately. Let ab denote the correction of the j-th diagonal of K (l) and Θ (l) to the fixed points. Linearizing Equation 88and Equation 92gives

EQUATION

Next let {ρ α } α be the eigenvalues of Ā and 

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.2 THE EFFECT OF POOLING AND FLATTENING OF CNNS:

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 significantly affect the spectrum. The NNGP and NTK of the l-th hidden layer CNN are 4D tensors K 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 Θ 

ft is essentially the same as its FCN counterpart Θ 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 flattening 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.

:

More precisely, one needs to consider the Jacobian of T as an operator from positive semi-definite matrices to positive semi-definite matrices. We refer the readers to Section B ofXiao et al. (2018) for more details.

If Θ (l) diverges to infinity, we define P (Θ * ) = lim l→∞ P (Θ(l) ). If Θ * train, train is singular, we will add a diagonal regularizer σId into Θ * train, train .

It has been observed in previous works Poole et al. (2016); Schoenholz et al. (2017) that the diagonals converge much faster than the off-diagonals for tanhor erf-networks.

For simplicity, we ignore the polynomial correction in l.

We will use Roman letters to index channels and Greek letters for spatial location. We use letters i, j, i , j , etc to denote channel indices, α, α , etc to denote spatial indices and β, β , etc for filter indices.

Figures and Tables:

Figure1: Condition numbers of NTKs and their rate of convergence. Different colors represent images of different size. For example, in the yellow "12-6" , "12" represents the size of the dataset and "6" represents the dimension number (6 * 6 * 3 for FCNand (6, 6, 3) for CNN) (a) In the chaotic phase, κ (l) converges to 1 for all architectures. (b) We plot χ l 1 κ(l) , confirming κ explodes with rate χ l 1 /l in the ordered phase. In (c) and (d), the dashed lines representing the condition number κ (l) and solid lines the ratio between first and second eigenvalues. We see that, on the order-to-chaos transition, these two numbers converge to m+2

Figure2: Maximal learning rate can be calculated via the λ max . y-axis: accuracy and x-axis: multiples of η theory . Each point on the solid (dashed) lines represents the best training (test) accuracy throughout training of one configuration. From blue to purple to red, (σ ω , σ b ) is moving from the order phase to the chaotic phase. ρ = 1 is the theoretical prediction.

Figure 3: Top: training (left) and test accuracy of FCN using SGD. Bottom: test accuracy of CNN-P, CNN-F and the difference. In the blue strip, CNN-F significantly outperforms CNN-P, due to the fact that pooling increases the spectra gap by a factor of d.

ρ denote the NNGP and NTK of the (L+1)-th layer. Note that when ρ = 1, K

th eigenvector of Ā. Then for each α,(l+1) ab,α ≈ (ρ α χ c ) (l+1) (0)

(x, x ) and Θ (l) α,α (x, x ), where α, α ∈ [d] ≡ [0, 1, . . . , d -1] denote the pixel locations. To perform tasks like image classification or regression, "flattening" 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 "flattens" 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 "flattening" 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, ithe NTK right after flattening (pooling) the last convolution. We will also use Θ (l) fc to denote the NTK of FC. K

(x, x )  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

ab for CNN and FCN are of the same order.) Applying flattening and pooling, the NTKs becomeΘ (l) flatten (x, x ) = 1 d α Θ (l) α,α (x, x ) = 1 x=x p (l) + 1 x =x p we can see, Θ

subdominant Fourier modes which decay exponentially faster than the dominant Fourier modes. Therefore the spectrum properties of Θ 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/(3d) and λ (l) rest ≈ 2q * l/(3d), and κ (l) = md+2 2 + mdO(l -1 ). Here we use blue color to indicate the changes of such quantities against their Θ roughly by a factor of d and increases the condition number by a factor of d comparing to flattening. 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

Peer review (iclr2020):
This paper studies the relation between trainability and generalization ability in deep neural networks. In the theoretical analysis, the authors used the Neural network Gaussian process (NNGP) kernel and Neural Tangent kernel (NTK). The paper clarified that the spectrum of the NTK and NNGP has an important role in investigating the generalization and trainability, i.e., the condition number of the NTK. Some numerical experiments showed an agreement of theory with the practical behavior of learning algorithms. In this paper, some existing theoretical results on deep neural networks were combined to extract new insight. Thought the attempt of this paper is interesting, the readability of the paper is not necessarily high. - In equation (2), the operator T is defined as the kernel K(x,x*). However, the definition seems different from that in equation (8). The authors need to make clear the definition of T. - What is the *DC* mode in the sentence above the equation (15)? - Is the derivation of the left part in equation (9) straightforward? How was the second term, chi_1 q^* p^(ell), derived? I*m not sure how the dot{T} was dealt with. The argument below equation (3) should be used?