Title: GENERALIZATION OF TWO-LAYER NEURAL NET-WORKS: AN ASYMPTOTIC VIEWPOINT

Abstract: This paper investigates the generalization properties of two-layer neural networks in high-dimensions, i.e. when the number of samples n, features d, and neurons h tend to infinity at the same rate. Specifically, we derive the exact population risk of the unregularized least squares regression problem with two-layer neural networks when either the first or the second layer is trained using a gradient flow under different initialization setups. When only the second layer coefficients are optimized, we recover the double descent phenomenon: a cusp in the population risk appears at h ≈ n and further overparameterization decreases the risk. In contrast, when the first layer weights are optimized, we highlight how different scales of initialization lead to different inductive bias, and show that the resulting risk is independent of overparameterization. Our theoretical and experimental results suggest that previously studied model setups that provably give rise to double descent might not translate to optimizing two-layer neural networks.
In modern neural networks, the number of parameters can easily exceed the number of training samples, yet in many circumstances, there is little sign of overfitting even in the absence of explicit regularization (Zhang et al., 2016) . This phenomenon is usually explained by the interplay between the model architecture and the optimization method. Existing works have analyzed the implicit regularization of gradient descent on simple models (Gunasekar et al., 

:

features d, and neurons h tend to infinity at the same rate, under which overparameterization corresponds to increasing the limit of h/n (network "width"). This regime is particularly interesting because even though n → ∞, the empirical risk is not equivalent to the population risk. In addition, the joint scaling of n, d, h is parallel to the practical choice of model architectures, where it is common to train a larger network when the number of samples and input features are larger. Following Hastie et al. (2019) , we assume unit Gaussian input and noisy linear observations, and analytically derive the population risk of the solution of gradient flow on either the first or the second layer parameters when the flow is initialized close to zero. Our findings can be summarized as follows (see Figure 1 ):

• When only the second layer is optimized, we derive the risk in its bias-variance decomposition and demonstrate the presence of the double descent phenomenon.

• When the first layer is optimized, we compare two solutions of gradient flow from different scales of initialization, which we term as vanishing and non-vanishing initialization, and show in both cases the population risk is independent to overparameterization.

• For the vanishing initialization, we show that the risk of the gradient flow solution is asymptotically close to that of a rank-1 model. For non-vanishing initialization, we show that the gradient flow solution is wellapproximated by a kernel model and derive the risk.

1.1. RELATED WORKS:

Global Convergence of Two-layer Networks. A plethora of recent works have explored the global convergence of shallow neural networks. Mei et al. (2018; 2019) ; Chizat and Bach (2018a) ; Rotskoff and Vanden-Eijnden (2018); Sirignano and Spiliopoulos (2018) ; Nitanda and Suzuki (2017) studied the mean-field limit where the number of neurons h → ∞ and the second layer scaled by 1/h, and established correspondence between the main-particle limit of gradient descent and Wasserstein gradient flow to demonsrate global convergence. On the other hand, Jacot et al. (2018) ; Du et al. (2018) ; Oymak and Soltanolkotabi (2019) ; Allen-Zhu et al. (2018b) ; Song and Yang (2019) considered a different scaling and showed that gradient descent on overparameterized models converges to global minimizer at a linear rate; key to these results is an observation that optimization via gradient descent is asymptotically equivalent to kernel regression with respect to the neural tangent kernel.

Active vs. Lazy Training. Following Chizat and Bach (2018b) , we refer to the two aforementioned scalings as the active and lazy (kernel) regime. It has been observed that different regimes lead to contrasting inductive biases. Williams et al. (2019) ; Woodworth et al. (2019) ; Li et al. (2017) showed that for certain two-layer network or overparameterized linear model, the scale of initialization controls the implicit regularization of gradient descent (from sparse to smooth solution). In the student-teacher setup (Tian, 2017; Zhong et al., 2017) , Ghorbani et al. (2019b; a) showed that kernel models in high dimensions perform no better than low-degree polynomials on the input or fullytrained two-layer network. Additionally, Suzuki (2018) ; Allen-Zhu and Li (2019) ; Yehudai and Shamir (2019) ; Wei et al. (2018) demonstrated that neural network outperforms linear estimators (including kernel method) in learning various target functions. The difference between fixed bases and adaptive bases mirrors the difference in optimizing the first or second layer in our setup.

Double Descent. The term double descent refers to the phenomenon that the population risk of an empirical risk minimizer manifests a "cusp" at the interpolation threshold, and further overparameterization decreases the risk. First observed in Krogh and Hertz (1992) , the phenomenon has been recently connected to the benefit of overparameterization (Belkin et al., 2018; Geiger et al., 2018; Spigler et al., 2018; Advani and Saxe, 2017) , and can be precisely characterized for certain simple models (Hastie et al., 2019; Belkin et al., 2019; Bartlett et al., 2019; Xu and Hsu, 2019) . Our work is inspired by Hastie et al. (2019) which uses random matrix theory to derive the asymptotic risk for linear and random feature models. Concurrent to our work, Mei and Montanari (2019) analyzed the random features model and derived its population risk for which double descent occurs both in bias and variance. This aligns with our results on optimizing the second layer in Section 4 although we do not derive the bias component explicitly. Compared to Hastie et al. (2019) ; Mei and Montanari (2019) , the focus of this work is to highlight the different generalization property of models obtained from optimizing different layers of the network and from different initialization.

Random Matrix Theory. High-dimensional models, including kernel models and neural networks, can be analyzed by studying the properties of random matrices. El Karoui et al. (2010) ; Cheng and Singer (2013) ; Fan and Montanari (2019) studied the spectral properties of kernel matrix via decomposing the nonlinearity with Taylor series or Hermite polynomials, which in turn explains the generalization of high-dimensional kernel ridgeless interpolators (Liang and Rakhlin, 2018) .

In addition, similar tools have been used to study two-layer neural networks (Louart et al., 2018; Pennington and Worah, 2017) and related quantities such as the Fisher information matrix (Karakida et al., 2018; Pennington and Worah, 2018) .

2. PRELIMINARIES: TWO-LAYER NEURAL NETWORK:

Consider the following bias-free two-layer neural network f : R d → R with h hidden units

EQUATION

where x ∈ R d is the input, w i ∈ R d is the weights corresponding to neuron i, a i ∈ R is the i-th coefficient of the second layer, and φ : R → R is a Lipschitz continuous activation function with bounded Gaussian moments, i.e. E[φ(G) k ] < ∞, ∀k ∈ Z + for G ∼ N (0, 1). For concise notation, we write W = [w 1 , ...w h ] ∈ R d×h for the weight matrix, a = [a 1 , ...a h ] ∈ R h for the coefficient vector, X = [x 1 , ...x n ] ∈ R d×n for the data matrix, y ∈ R n for the corresponding vector of labels, and Φ = φ(W X) ∈ R h×n for the feature matrix at the first layer. We omit arguments of f when they are clear from the context.

We consider a student-teacher setup, in which data is generated by a teacher model F : R d → R with additive noise, and the student model aims to minimize the squared loss:

EQUATION

where

E[x i ] = 0, Cov(x i ) = Σ, E[ε i ] = 0, Var(ε i ) = σ 2 . We are interested in the population risk R(f ) = E Px [(F (x) -f (x)) 2 ]

. Our analysis will be made under the proportional asymptotics:

n, d, h → ∞; d/n → γ 1 , h/n → γ 2 ; γ 1 , γ 2 ∈ (0, ∞),

in which overparameterization corresponds to increasing γ 2 . Thus the characteristics of double descent considered in this work are: 1) large population risk as γ 2 → 1; 2) decrease in the risk for γ 2 > 1. While the empirical risk can be minimized in various ways, we analyze the solution of gradient flow, in which we update either the first layer W or the second layer a:

EQUATION

from small initialization. The rest of the paper is organized as follows. In Section 3, we start with a simple example of two-layer linear network as warm-up. In Section 4, we consider optimizing the second layer coefficients (flow over a) of a non-linear two-layer neural network under fixed Gaussian first layer, which is a random feature model. Section 5 considers optimizing the first layer weights (flow over W ) of such network under fixed Rademacher second layer. We defer all proofs and details on experiments to appendix.

3. WARM-UP: LINEAR NETWORK:

We begin with a simple linear model with φ(x) = x, i.e. Φ = W X. We remark that although the model is linear, the solution obtained by gradient flow on the two-layer model can be different than that from directly solving the linear regression problem on input features.

Training the Second Layer. Following Hastie et al. (2019) , we fix the first layer parameters to be randomly drawn from a unit Gaussian and optimize the coefficients a by minimizing a Φ -y 2 2 . The following lemma characterizes the solution of the gradient flow. Lemma 1 (Least squares solution). Given data matrix X, response vector y and model f (x) = φ(x W ), a with fixed first layer coefficients W , gradient flow on the coefficients a starting from zero initialization converges to â = Φ † y, where † stands for the Moore-Penrose inverse.

We make two assumptions on the data and the teacher model to simplify the computation.

(A1) Gaussian Features:

x i ∼ N (0, I d ); (A2) Linear Teacher: F (x) = x, β , β = r.

Denote the linear student network as f (x) = x, β , where β = W â and â is the least-square solution defined by Lemma 1. We write the population risk in its bias-variance decomposition.

EQUATION

where x 2 Σ = x Σx. We compute the bias and the variance separately to obtain the risk. Theorem 2. Given (A1)(A2) and let w i

i.i.d. ∼ N (0, d -1 I d ), at n, d, h → ∞ we have R (γ1<1) →          γ 1 -γ 2 γ 1 g 2 r 2 + γ 2 g 2 σ 2 , γ 2 < γ 1 , γ 1 g 1 σ 2 , γ 2 > γ 1 , R (γ1>1) →          γ 1 -γ 2 γ 1 g 2 r 2 + γ 2 g 2 σ 2 , γ 2 < 1, γ 2 g 1 γ 1 g 2 r 2 + g 1 + g 2 g 1 g 2 σ 2 , γ 2 > 1.

(5) where d/n → γ 1 , h/n → γ 2 , g 1 = |γ 1 -1|, and g 2 = |γ 2 -1|.

We observe that when d > n (i.e. γ 1 > 1), we obtain the double descent risk curve, i.e., the population risk achieves its maximum at γ 2 → 1 and further overparameterization (γ 2 > 1) reduces both the bias and the variance. Conversely when n > d and h > d (i.e. γ 1 < min(1, γ 2 )), the population risk becomes constant and equals to that of the minimum-norm solution βmin = X † y on the input features.

Training the First Layer. When the first layer of a linear network is optimized via gradient flow and the second layer is fixed, the following holds for zero-initialization of W . Proposition 3. Given W (0) = 0 and fixed a init = 0, at any time t > 0 of the gradient flow on W , W (t) is rank-1. Further, β = W a converges to the least squares solution of y = X β, the population risk of which is given in (Hastie et al., 2019, Thm. 1 & 3) ) as

EQUATION

)

In this case, overparameterization by increasing γ 2 does not influence the population risk. In addition, since the obtained two-layer linear model is equivalent to the minimum-norm solution on the input features βmin , optimizing the first layer always results in smaller or equal population risk compared to optimizing the second layer.

In this simple scenario for two-layer linear networks, double descent is observed only when the second layer is optimized, which reduces the objective to least squares regression on the intermediate features.

On the other hand, training the first layer from zero-initialization always yields the same solution that is independent to overparameterization. One natural question to ask is: does this phenomenon generalize to nonlinear two-layer neural networks?

The following sections answer this question in the affirmative under certain conditions. (c) bias of network with ReLU activation. Black line corresponds to γ1 → ∞ for linear network, which is empirically observed as an upper-bound. Note that as γ2 → 1 both bias and variance becomes unbounded.

4. NONLINEAR MODEL: OPTIMIZING THE SECOND LAYER:

In this section, we analyze the case when the second layer a is learned under fixed W and a nonlinear activation function φ (a random feature model). We first observe that by Lemma 1, the gradient flow finds the solution â = Φ † y. We again consider the following bias-variance decomposition.

EQUATION

= E x E[φ(x W )â X, W ] -F (x) 2 2 B=bias + E x φ(x W )â -E[φ(x W )â] 2 2 X, W V =variance .

We highlight that the variance term does not depend on the target function. The following result characterizes the variance of the random feature model.

Theorem 4. Given (A1) and w i i.i.d.

∼ N (0, d -1 I d ), when n, d, h → ∞, we have

EQUATION

EQUATION

)

where variables ξ, ρ, τ satisfies ξ > 0 or ξ < 0, ρ > τ > 0, and constants c 1 , c 2 defined as,

EQUATION

)

for G ∼ N (0, 1) and ξ denoting the imaginary part of ξ.

Remark. c 1 ≥ c 2 and the equality holds iff φ is linear.

Corollary 5. If we let γ 1 → ∞, the variance is equal to the lowest value of the variance of the linear model

V (γ1→∞) = σ 2 min{γ 2 , 1}/|1 -γ 2 |.

The proof of Theorem 4 largely follows from Hastie et al. (2019) with techniques similar to Cheng and Singer (2013) , but with modifications in otder to handle unnormalized and uncentered activation functions. The above theorem holds irrespective of the underlying teacher model, and is consistent with the double descent risk curve as it suggests that for all γ 1 , variance of the random feature model peaks at h = n then drops as γ 2 further increases. Note that as γ 1 → ∞, a linear and nonlinear network would have the same asymptotic variance.

Since double descent is observed in the variance term, we do not derive the bias for all γ 1 , γ 2 . Instead, we show that for linear teacher, the bias also becomes unbounded as γ 2 → 1.

Proposition 6. Given (A1)(A2) and w i i.i.d.

∼ N (0, I d ), then B → ∞ as γ 2 → 1. Furthermore, B is finite when γ 2 > 1.

Thus we have shown that a "cusp" in the population risk appears at h = n, which aligns with the double descent phenomenon. Empirically as γ 1 → ∞ the nonlinear model also shares the same asymptotic bias with the linear model, as shown in Figure 2 . We note that (Mei and Montanari, 2019, Thm. 1 & 3) analytically solved the risk of random feature model for a larger class of target functions than ours and confirmed that double descent appears in both the bias and the variance.

5. NONLINEAR MODEL: OPTIMIZING THE FIRST LAYER:

Having observed the double descent phenomenon in optimizing the second layer, in the sequel we consider a two-layer neural network with fixed second layer coefficients initialized from a Rademacher distribution a i ∼ Unif{-1/ √ h, 1 √ h}, and the first layer W is optimized with the following update

EQUATION

which potentially has different stationary solutions with no explicit form, depending on the initialization. We denote the solution of this flow at time t started from designated initialization by W init (t), its stationary solution by W , and the corresponding network by f . Remark. Although we let n → ∞, this dynamics does not corresponds to the population gradient flow considered in Tian (2017) . For instance when x ∼ N (0, I/d), the spectrum of the data covariance is Marčenko-Pastur, whereas the population covariance is identity.

As we cannot characterize the gradient flow solution from all possible initializations, we consider two specific scales of initialization:

Vanishing: w Van i (0) ∼ N (0, I d /dh 1+ ); Non-vanishing:

w NV i (0) ∼ N (0, I d /d -).

Note that neither of the two initializations correspond to the "mean-field" regime (e.g. analyzed in Mei et al. (2018) ) due to the 1/ √ h scaling of the second layer. In other words, as h increases, we expect the distance traveled by each parameter to decrease under both initializations. The difference, however, is the "relative" amount the parameters traveled compared to their initialized magnitude, which leads to solutions with contrasting properties. As we will see, under (A1)(A2) and vanishing initialization we have W (0)-W F / W (0) F 1, i.e. the contribution of initialization vanishes at the end of training, whereas for non-vanishing initialization the inequality is in the opposite direction, i.e. W "barely moves" and resembles the initialization W (0).

5.1. VANISHING INITIALIZATION:

As d, h → ∞, the vanishing initialization becomes arbitrarily close to zero-initialization. We thus expect the gradient flow under vanishing initialization to "resemble" that of starting from exactly zero if the flow converges sufficiently fast and the gradient being Lipschitz. The Lipschitz condition (Lemma 18) can be established under the following assumption on the activation.

(A3): φ is smooth, Lipschitz and monotone with φ (0) = 0; |φ (±x) -φ (±∞)| = O(e -x ).

The above assumption requires that the derivative of the nonlinearity φ saturates beyond a O(1) region, which holds true for the commonly-used smooth activations such as sigmoid and SoftPlus. In addition, the choice of scaling ensures that the gradient flow converges sufficiently fast. We thus have the following characterization of the population risk:

Theorem 7. Given (A1-3). Let T = O(log log h) and f (•) = f van (•, W (T )), then as n, d, h → ∞,

the gradient flow reaches a o(1) first-order stationary point at time T, i.e. ∂W (T )/∂t F ∈ o(1), at which point the population risk is given as

EQUATION

)

The expression above is the same as the risk of the least squares solution on input β = X † y; therefore the risk is independent to overparameterization (increasing γ 2 ). The intuition is that when the weights are initialized sufficiently small and travel infinitesimally, then the activation can be linearized around 0 and thus the model is equivalent to a two-layer linear network. Note that this result does not apply to the non-smooth ReLU activation. Instead, in Appendix E we heuristically show that under the additional assumption that the data is symmetric, the risk of ReLU network is also independent to γ 2 .

5.2. NON-VANISHING INITIALIZATION:

When initialization is sufficiently large, the amount each parameter travels to minimize the empirical risk becomes asymptotically negligible compared to the magnitude of initialization. In this case we establish under (A1-3) that ( 11) is asymptotically equivalent to the kernel gradient flow on the tangent kernel:

k(x, y) = ∇ W init f (x), ∇ W init f (y) .

The converged parameters under this linearized dynamics has the following closed-form:

EQUATION

where J ∈ R n×(d×h) is the Jacobian matrix w.r.t. to the model parameters. We remark that in contrast to most NTK-type global convergence results that require the width of the model to grow faster than the number of data points (e.g. Du et al. ( 2018)), our result is not built upon the overparameterization on width, but instead an anti-concentration that relies on the scale of initialization. Consequently the above initialization is larger than the scale that is commonly used in practice.

One may naturally expect the double descent phenomenon to appear in this kernel solution, as ∆ exhibits the form of a least squares solution which contains a pseudo-inverse. However, we show that this is not the case under the same assumptions in Section 4; in fact, the risk is also independent to γ 2 .

An obstacle in computing the risk of the kernel model is the potentially non-zero f init (X). We thus adopt the "doubling-trick" from Chizat and Bach (2018b) to ensure f init (•) = 0, i.e. we assume the following symmetric property on the initialized weights:

(A4) Symmetric Initialization: ∀i ∈ [1, h], ∃!j ∈ [1, h] s.t. a i w init i = -a j w init j .

Theorem 8. Given (A1-4) and let n, d, h → ∞, the stationary solution f has the following risk

EQUATION

)

where m = b 2 1 /b 2 0 , b 2 0 = E[φ (G)] 2 , and b 2 1 = E[φ (G) 2 ] -b 2 0 , G ∼ N (0, 1

). Note that the population risk is again independent to γ 2 , and thus double descent does not appear for this initialization. Roughly speaking, the reason that the risk does not become unbounded at some point is that in the asymptotic limit the pseudo-inverse (JJ ) † is stable due to the nonlinearity and dh n. We make two additional observations:

• the stability of the inverse at γ 1 → 1 depends on the lowest eigenvalue of the tangent kernel matrix (smaller λ min (K) entails larger variance), which is determined by the nonlinearity; • While our result only holds for network with zero initial output, for non-symmetric (i.i.d.) initialization we also observe that the risk is independent to γ 2 , but the bias is higher than that in symmetric initialization as shown in Figure 8 . We comment that the non-zero f init (X) in ( 13) behaves as zero-mean Gaussian due to central limit theorem; therefore, in the kernel regime the function output at initialization is equivalent to additive noise to the labels y, and we thus expect the magnitude of f init (X) to negatively influence the model generalization. dotted lines correspond to different γ2 (from 0.2 to 2) which is independent to the risk. The bias and variance for both initializations is well-aligned with Theorem 7 and Theorem 8.

5.3. COMPARING THE INITIALIZATIONS:

Figure 3 shows the agreement between theoretical prediction and experimental results. Although in both cases the risk is independent to overparameterization (γ 2 ), the two initializations lead to models with contrasting properties, as demonstrated by the following comparison on the risk.

Corollary 9. For any

γ 1 ∈ (0, ∞) and nonlinearity φ, B( f Van ) ≤ B( f NV ) ≤ 1. On the other hand, for all m > 0, V ( f NV ) = O(1), whereas V ( f Van ) can be arbitrarily large as γ 1 → 1.

Remark. m ≥ 0 for all smooth activations φ, and the equality holds if φ is linear.

Intuitively, small initialization enables the model "evolve" more during optimization and better align with the data and target function. This potentially results in a lower bias, at the expense of overfitting more to the noise (high variance). In contrast, with sufficiently large initialization the final model becomes close to the initialized model, and thus we may expect it to be less "aligned" to the target (high bias) but is more stable (lower variance).

In illustrate the different inductive bias of the two initializations, we plot the trajectory of neurons in Appendix A Figure 4 . Observe that for vanishing initialization the neurons stay close to one another throughout the trajectory, which results in a low-rank weight matrix, as predicted by Theorem 7. In contrast, for non-vanishing initialization the neurons stay close to initialization (therefore full-rank), which validates the kernel approximation. Last but not least, although the derived risk is only for learning a linear target function, we empirically observe that when the teacher is also a two-layer network, the population risk follows the same trend, i.e. double descent occurs when only the second layer is optimized, as shown in Figure 7 .

6. DISCUSSION AND FUTURE WORKS:

We derived the exact population risk of high-dimensional two-layer neural networks in learning a linear target function over Gaussian data with additive label noise, and showed that optimizing the first or the second layer via gradient flow results in solutions with contrasting properties. Specifically, double descent is present when the second layer coefficients are optimized, but not when the first layer weights are optimized under certain initializations. Moreover, we highlight that the scale of initialization leads to different inductive bias in optimizing the first layer.

It should be noted that our analysis only applies to the unregularized objective: it has been shown that explicit regularization (such as 2 penalty) stabilizes the singularity at γ 2 → 1 (Mei and Montanari, 2019), and algorithmic regularization (Li et al., 2019a; Dong et al., 2019 ) also provides robustness against noisy observations. We further remark that our findings do not directly contradict the experimental double descent phenomenon, nor the practical benefit of overparameterization. In particular, the interpolation limit could occur at γ 2 → 0 which is beyond the regime we consider (such as Figure 4 in (Belkin et al., 2018) ). Thus what we conclude is that under the studied proportional asymptotics, the mechanism that provably gives rise to double descent from previous works on least squares regression might not translate to neural networks trained with gradient descent.

To simplify the computation, we rely on a set of strong assumptions similar to those in Hastie et al. (2019) , some of which we believe can be relaxed in future works, such as isotropic Gaussian input and linear target. Importantly, the two specific scales of initialization studied in Section 5 are by no means exhaustive, and thus one would expect that under a different initialization of the first layer, or the mean-field 1/h scaling of the second layer, the risk of the trained model can be very different.

Changing the loss function may also alter the generalization behavior of the network. Another challenging problem is to extend the current analysis to beyond two layers. Last but not least, our result characterizes gradient flow which resembles gradient descent with small stepsize, and thus it would be interesting to study the effect of learning rate schedule, which has a known impact on generalization (Smith and Le, 2017; Li et al., 2019b) .

A ADDITIONAL FIGURES AND PLOTS Note that bias under i.i.d. initialization is also independent to overparameterization (γ2), but is higher than the bias under symmetric initialization ("doubling trick") and not always upper-bounded by the null risk r 2 .

B BACKGROUND B.1 ROTATIONAL INVARIANCE:

The rotational invariance of Gaussian distribution is crucial in our analysis throughout this paper. A basic observation is that for a random Gaussian matrix X and any fixed unitary matrix U , the distribution of X and U X are the same.

Lemma 10 (Rotational Invariance). Denote A(X) ∈ R d×d a matrix function of X ∈ R d×n . If A(X) satisfies that A(U X) = U A(X)U T for all unitary U , then

EQUATION

)

for any fixed nonzero β ∈ R d and random matrix X with each entry i.i.d. X ij ∼ N (0, σ 2 ).

Proof. Choose a set of Unitary matrices

{U i } d i=1 such that U i β = β e i

, where e i is the i-th

EQUATION

)

Note that the property also holds for matrix function A(X) that satisfies A(XU ) = U A(X)U T , and can be extended to matrix function A that takes multiple matrices as input.

For brevity we will refer to rotational invariance instead of equation ( 15).

B.2 MAR ČENKO-PASTUR LAW:

For a real symmetric random matrix A ∈ R p×p , define its empirical spectral density as Marčenko and Pastur (1967) showed that as n, p → ∞ and p/n = γ ∈ (0, ∞), the empirical spectral density µ Wp (dλ) converges weakly to a limiting density µ MP(γ) (λ):

µ A (dλ) = 1 p p i=1 δ λi(A) (λ)dλ, where δ a (x) = δ(x -a) is the Dirac delta function. Assume A = W p ∼ W p (I, n) is a Wishart matrix, i.e. W p = X X/n and X ∈ R n×p is random Gaussian matrix with each column i.i.d. X i ∼ N (0, I).

EQUATION

)

We say µ MP(γ) is the density of the Marčenko-Pastur distribution with support

S = [(1 - √ γ) 2 , (1 + √ γ) 2 ] (for 0 < γ < 1) or S = {0} ∪ [(1 - √ γ) 2 , (1 + √ γ) 2

] (for γ ≥ 1). Note that this implies that the smallest non-zero eigenvalue of A is bounded away from 0 a.s. for γ = 1.

The explicit form of Marčenko-Pastur distribution allows us to investigate the asymptotic properties of random matrices. Generally speaking, by Pormanteau theorem one can translate any bounded continuous function on the empirical spectral density to the one on Marčenko-Pastur distribution, i.e. for any bounded f (λ) ∈ C(S), as n, p → ∞ with p/n = γ, almost surely

EQUATION

)

One implication is the following trace concentration on the inverse Wishart matrix for γ < 1,

tr X X = tr 1 p W -1 p = 1 p p i=1 1 λ i (W p ) = S 1 λ µ Wp (dλ) → S 1 λ µ MP(γ) (dλ) = 1 1 -γ .

(19) We remark that instability of the trace of the invert Wishart matrix as γ → 1 plays an important role in the double descent phenomenon.

B.3 ORTHOGONAL POLYNOMIALS:

Orthogonal polynomials are useful in the analysis of nonlinear random matrices. Suppose φ is a function in L 2 (R, µ G ), where G ∼ N (0, 1), µ G (dx) = ( √ 2π) -1 e -x 2 /2 dx is the Gaussian measure. For n ≥ 0, define the Hermite polynomials

EQUATION

Published as a conference paper at ICLR 2020 Note that orthogonality can be easily verified:

EQUATION

Since {H i (x)} ∞ i=0 forms a set of orthogonal basis in L 2 (R, µ G ), the function φ(x) can be expanded under the Hermite basis as

EQUATION

Following Cheng and Singer (2013) , we mainly focus on the first few terms of the Hermite expansion.

One can check that H 0 (x) = 1, H 1 (x) = x, and φ(x) can be expanded as

EQUATION

where

c 0 = E[φ(G)], c 1 = E[Gφ(G)],

and terms in the RHS are orthogonal to one another in L 2 (R, µ G ). Taking square and expectation over both sides yields

EQUATION

which indicates

E[φ(G) 2 ] -E[φ(G)] 2 -E[Gφ(G)] 2 = E[φ ⊥ (G) 2 ] ≥ 0.

Note that the equality holds if and only if E[φ ⊥ (G) 2 ] = 0, i.e. φ(x) = c 0 + c 1 x is linear. One application of this decomposition is to "linearize" a non-linear matrix in high dimensions, which will be useful for the following sections.

C PROOF OF MAIN RESULTS:

C.1 PROOF OF LEMMA 1

Given features X ∈ R d×n , labels y ∈ R d and model parameters θ, the gradient flow of θ on the squared loss y -X θ 2 2 can be written as

EQUATION

Thus with initialization θ 0 , the solution of this ODE at time t can be written in explicit form

EQUATION

Since θ 0 = 0, taking t → ∞ yields the desired result.

C.2 PROOF OF THEOREM 2:

We compute the bias and variance for different cases of γ 1 , γ 2 . We first discuss the case where the random feature Φ X is not full rank (Case I). Otherwise when Φ X is full rank, we discuss whether it is full column rank (Case II) or full row rank (Case III).

Case I: W X is not full rank, i.e. γ 1 < 1, γ 2 > γ 1 . In this case rank (Φ X ) = d < min(n, h), and thus by taking the Moore-Penrose inverse we obtain

EQUATION

)

It is clear that the mean and variance is identical to the underparameterized regime in Hastie et al. (2019) , i.e. when n, d, h → ∞,

EQUATION

)

Case II: W X has full column rank, i.e. γ 2 < 1, γ 1 > γ 2 . By Lemma 1, the solution of the second layer coefficients is

EQUATION

Denote W = U ΣV the singular value decomposition. We perform the block decomposition:

EQUATION

where

Σ 0 ∈ R h×h , X 0 ∈ R h×n , X 1 ∈ R (d-h

)×n , and notice that X 0 , X 1 are independent. By a concentration of measure argument (e.g. Tao (2012) ; Hastie et al. (2019) ) one can show that the quantity below tightly concentrates at its expectation. For the variance we have

EQUATION

)

where the last equality follows from Appendix B.2. Similarly for the bias term we have

EQUATION

)

where symmetric arguments are omitted as (• • • ), and (i) follows from the rotational invariance argument introduced in Lemma 10 and that β β = r 2 . By the block decomposition (30),

EQUATION

)

Therefore the bias term simplifies to

EQUATION

)

Thus we have obtained that as n, d, h → ∞

EQUATION

)

Case III: W X has full row rank, i.e. γ 1 > 1, γ 2 > 1. Similarly, the least squares solution is

EQUATION

Simplifying the variance:

EQUATION

)

where we applied the SVD of X = U ΣV and the rotational invariance argument. Using a similar block decomposition on W :

EQUATION

where

Σ 0 ∈ R n×n , W 0 ∈ R n×h , W 1 ∈ R (d-n)×h

, and W 0 , W 1 independent. We thus simplify V as

EQUATION

)

Hence we obtain the following expression on the variance

EQUATION

)

We omit the derivation of bias, which follows a similar derivation:

EQUATION

)

Combining Case I, II, III yields theorem 2.

C.3 PROOF OF PROPOSITION 3:

Given the squared loss, one can derive the dynamics of W with fixed second layer a w.r.t the loss:

EQUATION

)

Note that the update of W can be written as a linear combination of a. Since W (0) = 0, we can write W (t) = ŵ(t)a for some ŵ. The corresponding flow on ŵ is

EQUATION

which gives the following solution

EQUATION

)

Thus gradient flow on the first layer leads to the minimum-norm solution on the input features.

C.4 PROOF OF THEOREM 4:

Following the bias-variance decomposition (7), the variance term can be written as (σ 2 omitted)

EQUATION

where we define the expected non-linear Gram matrix K W ∈ R h×h as

EQUATION

)

and for each entry we have

(K W ) [i,j] = E x φ(w i x)φ(w j x) .

Random matrix in the form of covariance matrix of nonlinear features has been studied in many works Hastie et al. (2019) ; Mei and Montanari (2019) ; Liao and Couillet (2018) ; Louart et al. (2018) ; Pennington and Worah (2017) . We note that our setup for the variance term is very similar to that for nonlinear features in Hastie et al. (2019) with modifications mentioned below.

In contrast to the linear network in Section C.2, the Gram matrix of a nonlinear activation is almost surely full-rank, as specified in the following lemma from Pennington and Worah (2017): Lemma 11. Suppose φ is not linear. Then the smallest singular value of Φ = φ(W X) ∈ R h×n is of order O( √ n). To be precise, consider the empirical spectral density µ ΦΦ /n (dλ), when n, d, h → ∞ with d/n → γ 1 and h/n → γ 2 , µ ΦΦ /n (dλ) converges weakly to

µ ΦΦ /n (dλ) → [1 -γ -1 2 ] + δ 0 (λ)dλ + µ + (dλ)

, where µ + (dλ) has non-negative support [ρ, ∞) with ρ > 0.

We therefore consider two scenarios: Φ is full column rank (Case I) or full row rank (Case II).

Case 1. h < n. In this case (45) simplifies into

V = tr φ(W X)φ(X W ) -1 K W = lim ξ→0 - tr ΦΦ -ξI -1 K W = lim ξ→0 - V ξ ,

where the continuity and boundness of V ξ at ξ = 0 -is guaranteed by Lemma 11 when γ 2 = h/n = 1. A ridge ξ is added make use of (Louart et al., 2018, Theorem 1) , which derived the asymptotic equivalent of the resolvent ΦΦ -ξI -1 . It follows that as n, d, h → ∞,

EQUATION

)

where we have used the inequality tr (AB) ≤ A F B 2 , and equivalently by taking ξ → 0 we get

EQUATION

Therefore lim ξ→0 lim n,d,h→∞ n/h • tr ΦΦ -ξI -1 K W = γ 2 /(1 -γ 2 ). Note that ∂V ξ /∂ξ = tr ξ(ΦΦ -ξI) -2 K W is bounded around the neighbourhood of ξ = 0, hence following the same argument as (Hastie et al., 2019 , Theorem 4), we exchange the limit of ξ → 0 and n, d, h → 0:

EQUATION

)

Case 2. h > n. Techniques used in the current proof are largely borrowed from Hastie et al. (2019) ; Cheng and Singer (2013) , and we include the full proof for completeness. It should be noted that compared to Hastie et al. (2019) we handle the non-zero expectation of the nonlinearity under Gaussian distribution, i.e. the off-diagonal entries of the kernel matrix is no longer zero-centered. For simplicity we mainly adhere to the notations in Hastie et al. (2019) .

We briefly summarizes the procedure for deriving V . Instead of calculating the variance directly, we analyze a modified quantity V ξ and then take ξ → 0, which can be connected to the trace of the resolvent of matrix Ã defined in (56); this translates the calculation of V ξ into the calculation of the Stieltjes transform of Ã (57), ( 61), (62).

C.5 DERIVING THE VARIANCE V FOR h > n:

Step 1. An equivalent expression. For notational simplicity we omit the magnitude σ:

EQUATION

)

and due to the same continuity argument as in Case 1 we have

V = lim ξ→0 1 n tr S(S S -ξI n ) -2 S K W = lim ξ→0 V ξ .

where S = φ(W X)/ √ n = Φ/ √ n, ξ ∈ C and ξ > 0 or ξ < 0.

We decompose the normalized feature matrix S = φ(W X)/ √ n as S = U ΣV , where

Σ = diag h×n (φ 1 , • • • , φ n ) ∈ R h×n is a tall diagonal matrix, and U = [u 1 , • • • , u h ] ∈ R h×h

is the set of orthogonal eigenvectors of SS = φ(W X)φ(X W )/n, and V ∈ R n×n is the set of orthogonal eigenvectors of S S = φ(X W )φ(W X)/n. Now the variance can be written as

EQUATION

)

By the same argument as in Lemma 13 of Hastie et al. (2019) one can show that when n, d, h → ∞

V ξ = 1 n tr U Σ(Σ Σ + ξI n ) -2 ΣU K W → 1 n tr U Σ(Σ Σ + ξI n ) -2 ΣU KW , (52)

in which KW is the approxmation of K W defined in Lemma 16. Writing the trace explicitly (denote eigenvalues λ i = φ 2 i , and

EQUATION

)

Since the positive support of spectrum λ is lower bounded and the density at 0 is 1 -γ -1 2 , we have

EQUATION

where we define µ n (x) = 1 h h i=1 δ λi (x)u i KW u i and its positive part µ + n (x). Hence we have

EQUATION

)

We define the following matrix Ãn (ρ, ς, τ ) ∈ R N ×N where N = n + h:

EQUATION

)

And denote the Stieltjes transform of Ãn as

EQUATION

)

Then following the definition of KW one can show that

EQUATION

)

and taking matrix derivative gives

EQUATION

)

Denote the limit m(ξ, ρ, ς, τ ) = lim n,h,d→∞ mn (ξ, ρ, ς, τ ), and the derivative is given as

EQUATION

)

For simplicity we define the following function on ξ:

EQUATION

) also denote q + (ξ) = q(ξ) -γ2-1 ξ 2 = γ 2 λ>ρ 1 λ+ξ 2 µ + ∞ (dλ).

Comparing with (55) yields:

EQUATION

)

Step 2. Calculating q(ξ) and m n (ξ, ρ, ς, τ ) This subsection aims to calculate m(ξ, ρ, ς, τ ) and q(ξ) = -m x (ξ, rx, sx, tx)| x=0 , from which the variance can be computed from (57)(61)(62).

We define A n by subtracting the off-diagonal entries of the upper-left block of Ãn :

EQUATION

)

where S = S -a 0 I p×n , i.e.

S ik = φ(w i x k ) -a 0 = ϕ(w i x k ), a 0 = E[φ(x)]. The Stieltjes transform of A n given by m n (ξ, ρ, τ ) = 1 n tr (A n (ρ, τ ) -ξI N ) -1 .

The following Lemma shows that mn and m n have the same limit: Lemma 12. when n → ∞ and for ξ > 0 or ξ < 0, we have m n (ξ, ρ, τ ) → mn (ξ, ρ, ς, τ ).

Proof. By definition of Ān and A n ,

Ãn (ρ, ς, τ ) -A n (ρ, τ ) = 1 h 0 h 0 n 1 n ξ a 0 a 0 0 1 h 0 h 0 n 1 n (64)

which is a rank-2 matrix. By theorem A.43 from Bai and Silverstein (2010) , which characterizes the effect of finite-rank perturbation on the e.s.d. of random matrices:

EQUATION

)

where F M is the empirical spectral distribution of M ∈ R n×n . The claim follows from the Stieltjes continuity theorem (e.g. Section 2.4 in Tao ( 2012)).

To calculate the Stieltjes transform m n , we take advantage of the block structure of A n :

EQUATION

EQUATION

)

One can observe that m n (ξ, ρ, τ

) = γ 2 m 1,n (ξ, ρ, τ ) + m 2,n (ξ, ρ, τ ).

In the following equations we omit the subscript n, as well as dependency on ρ, ς, τ . Following Hastie et al. (2019) , we rewrite the matrix

A n = A = A * a a 0 , where A * is a (N -1) × (N -1)

matrix with last column and row of A removed and s the activation vector:

EQUATION

)

Hence by the block matrix inverse formula

EQUATION

)

Plugging this back in the Stieltjes transform we have

EQUATION

)

To obtain an asymptotic description of m 2,n , we perform the orthonormal decomposition on the nonlinearity ϕ introduced in Cheng and Singer (2013

): ϕ(x) = a 1 x + ϕ ⊥ (x), where a 1 = E x∼N (0,1) [xϕ(x)].

For each w i (1 ≤ i ≤ h), we perform the following orthonormal decomposition (along the direction of x n and the direction of wi perpendicular to x n ):

EQUATION

)

We can thus simplify the activation vector a as

EQUATION

)

and for

EQUATION

)

Similarly we decompose the activation function in S,

EQUATION

)

We thus have an equivalent expression of matrix

EQUATION

By argument similar to (Hastie et al., 2019, B.1.2 ), E diminishes to 0 as n → ∞ with respect to the Frobenius norm, therefore by the Woodbury's identity and the expression of m 2,n in ( 70)

EQUATION

We bound each term u, v, S to compute m 2,n . For u

EQUATION

EQUATION

And for v (note that U is dependent on a)

EQUATION

where

EQUATION

as n, d, p → ∞. And finally for S,

EQUATION

And hence as n, d, h → ∞,

EQUATION

)

Therefore by combining ( 78), ( 80), ( 82), we arrive at the following expression on m 2,n ,

EQUATION

)

Similarly we can calculate m 1,n (ξ, ρ, τ ) as

EQUATION

Uniqueness in (84)(83) follows from (Hastie et al., 2019, Sec B.1 ) and is omitted.

C.6 PROOF OF COROLLARY 5:

In this section we take the limit γ 1 → ∞. In this case (8), (9) simplify to

EQUATION

EQUATION

Recall that m(ξ, ρ, τ ) = γ 2 m 1 (ξ, ρ, τ ) + m 2 (ξ, ρ, τ ). By taking the derivative we have

EQUATION

Observe that ( 85), ( 86) constitutes a set of implicit functions. Thus differentiating the two functions with respect to τ, ρ and then substitute by ρ = τ = 0 gives

EQUATION

)

Hence by ( 62) we obtain the asymptotic variance:

EQUATION

)

Combining the case where γ 2 < 1 in Theorem 4 completes the proof.

Remark. For φ(x) = ReLU(x), c 1 = 1/2 -1/(2π), c 2 = 1/4. For φ(x) = SoftPlus(x) = log(1 + e x ), numerical integration yields c 1 ≈ 0.2715, c 2 = 1/4. C.7 PROOF OF COROLLARY 6 C.7.1 UNBOUNDED BIAS WHEN h = n

From the bias-variance decomposition (7), the bias B is written as

EQUATION

When h = n, due to the nonlinearity of φ, we have φ(W X) is full rank a.s., and hence

[φ(X W )] † = [φ(X W )] -1 . We have the following bound for Q 1 1 d tr (Q 1 ) = 1 d tr X[φ(W X)] † K W [φ(X W )] † X = 1 d tr K W [φ(X W )] -1 X X[φ(W X)] -1 ≥ 1 d λ min (K W )tr [φ(X W )] -1 X X[φ(W X)] -1 = λ min (K W ) n tr (SS ) -1 • 1 d X X .

Since W and X/ √ d are R d×n = R d×p follows the same distribution where each entry i.i.d.

EQUATION

where in Section C.5 we have showed ( 91) is unbounded when n → ∞. Moreover, by ( 154) and Weyl's theorem we have λ min (K W ) = O(1), and thus

EQUATION

To sum up, for n → ∞ and γ 2 → 1 we have B → ∞.

C.7.2 BOUNDED BIAS WHEN γ 2 > 1

Since h > n, the two terms in the expression of the bias can be written as

EQUATION

EQUATION

Therefore we have

EQUATION

in which we used Lemma 11 and tr (AB) ≤ λ max (A)tr (B) for positive semi-definite A, B. Similarly

EQUATION

where we applied tr (AB) ≤ (tr A A + tr B B )/2. We therefore conclude that B is bounded when h > n and h, n → ∞.

Remark. Concurrent to this work, Mei and Montanari (2019) provides a complete characterization of the bias term and confirms our observations above.

C.8 PROOF OF THEOREM 7

For simplicity we assume n, d, h to be even and let d 0 = d/2, n 0 = n/2 and h 0 = h/2. Since the second layer is fixed a i ∼ Unif{-1/ √ h, 1/ √ h}, we let a i = 1/ √ h and a i+h0 = -1/ √ h for all 1 ≤ i ≤ h 0 . We therefore write a = h -1/2 [1 h0 , -1 h0 ] and W = [W + , W -]:

EQUATION

The empirical risk can thus be written as

EQUATION

C.8.1 DEFINING GRADIENT FLOWS

In this section we define three gradient flows and show that the three flows are similar in some sense. GF-Original is the original gradient flow (11), i.e.

EQUATION

starting from the vanishing initialization w O i (0) ∼ N (0, I/dh 1+ ). Note that the gradient for the negative part W O -can be similarly defined. We now define the flow under the same objective but from exact zero initialization w D i (0) = 0 termed GF-Double. Due to zero initialization, a basic observation is that the solution [W D + , W D -] is at most rank-2, and more precisely, the parameters in the flow takes the form of W D ± (t) = w D ± (t)1 where w D ± (t) admits the following dynamics:

EQUATION

Lastly, we define the GF-Single with solution denoted as w ± = w S ± (t) :

EQUATION

)

from zero initialization w D ± (0) = 0. This can be seen as replacing the nonlinearity φ with its first-order Taylor expansion at the origin.

C.8.2 FROM GF-DOUBLE TO GF-SINGLE:

Step 1. Solution of GF-single. Among the three flows defined above, only GF-single an explicit form at any time t. Specifically, the solution can be written as:

EQUATION

when d < n, or otherwise

EQUATION

when d > n. For simplicity we elaborate the proof only for d < n. We provide a condition under which the difference between the trajectories can be controlled, and we first that this condition holds for the linearized flow GF-single for all t in Lemma 17.

Condition A: w(t) 2 = O 1 √ d ; X w(t) ∞ = O poly log d √ d .

Step 2. Bounding the Difference in Gradient Flow Trajectory. Due to low rank property of GF-single and GF-double, in this subsection we slightly abuse the notation and define

EQUATION

We now show that the difference between the two trajectories defined above is asymptotically vanishing for w + (w -follows the same argument). Compare the two trajectories up to time T

EQUATION

where we have defined the error term as

EQUATION

)

To bound the error term, we note that at t = 0 Condition A holds for GF-double. Assume that for some 0 ≤ t ≤ T , Condition A also holds for GF-double, we have

EQUATION

where (i) follows from the error of Taylor expansion on φ and (ii) from Condition A. Therefore, by Equation ( 106) and Gronwall's inequality,

EQUATION

for T ∈ O(log log h). This shows that up to time T, the difference between the trajectories of GF-single and GF-double vanishes under the the assumption that Condition A holds for w D up to T. Importantly, note that w S (0) = w D (0) = 0, and that Condition A holds for w S for all T > 0.

Therefore, by a standard contradiction argument (e.g. (Du et al., 2018, Lemma 3.4 )), one can show that this closeness between the two trajectories implies that Condition A also holds for w D up to time T. With Equation ( 108) we bound the difference in population risk between the two flows.

Step 3. Bounding the Difference in Risk. Now we consider the difference of the population risk R S and R D of two models with parameters w S ± and w D ± .

EQUATION

where (i) is due to Cauchy-Schwarz inequality on norm, (ii) from Jenson's inequality on squares and Young's inequality, and (iii) from the Lipschitz assumption on the activation, and the observation that both R S and R D are finite for γ 1 = 1 due to the justified Condition A above. Therefore the difference between R S and R D vanishes for T = O(log log h).

Step 4. Bounding the Difference from Stationarity We compute the difference in risk between the model at some finite time t and the stationary point i.

EQUATION

for constants C 3 , C 4 > 0. Combining ( 110) and ( 111) yields for T = log log h,

EQUATION

The result in (112) for case d > n follows a similar proof and is omitted.

C.8.3 FROM GF-ORIGINAL TO GF-DOUBLE:

In this section we compare GF-original with GF-double. Note that the two flows differ only at initialization: vanishing initialization

w O i (0) ∼ N (0, I/dh 1+ ) v.s. zero initialization w D i (0) = 0. Note that at initialization y -f (X) 2 = O( √ n)

, and gradient flow decreases the empirical risk; therefore the condition for Lemma 18 is satisfied, and by the Lipschitz condition on the empirical gradient we have

EQUATION

And hence by Gronwall's lemma one obtains:

EQUATION

)

Therefore the difference in the function output can be bounded as

EQUATION

Taking T = log log h, together with the same argument in Step 1 yields

EQUATION

C.8.4 PUTTING THINGS TOGETHER By ( 112) and ( 116), we know that for T = log log h,

EQUATION

)

where ∂f (X; ω 0 )/∂ω is a dh × n matrix with column ∂f (x i ; ω 0 )/∂ω. Note that for n → ∞ and γ 1 , γ 2 ∈ (0, ∞), dh > n trivially holds, and by Corollary 15 we know that solution is given by

EQUATION

And the population risk can be written as (note that there is a factor of 2 due to the "doubling trick" at initialization to ensure f init (•) = 0):

EQUATION

where a bias-variance decomposition is made here, and for simplicity we define

EQUATION

C.9.3 APPROXIMATING THE KERNEL MATRIX

In this section we drop the negligible in the initialization. Following Cheng and Singer (2013) we utilize the orthonormal decomposition of φ

EQUATION

where E[φ ⊥ (G)] = 0. We develop the following lemmas to approximate the kernel matrix.

Lemma 13 (Approximation of ( KX ) ij ). There exist constants c, c > 0 such that for i = j with probability 1 -e -cnε 2 we have

EQUATION

)

and with probability 1 -e -c nε 2 ,

EQUATION

Proof. When i = j (i.e. Equation ( 132)), we have

EQUATION

The matrix H(x i , x j ) = x i x j E w φ (w x i )φ (w x j ) can be seen as the expected tangent kernel of nonlinear activation function studied in Du et al. ( 2018); Arora et al. (2019b) . Moreover, due to the assumed boundedness of φ (x) (A3), by Hoeffding's inequality we have

EQUATION

)

In addition, by the concentration of x i x j and x i 2 2 , i.e.

Pr x i x j /d > ε < 1 -e -c2dε 2 and Pr |x i x i /d -1| < ε > 1 -e -c3dε 2

, the orthonormal decomposition φ (x) = b 0 x + φ ⊥ (x) leads to the following linear approximation of the matrix H

EQUATION

)

and by taking ε = x i x j /d under the joint event we can show that

1 d ∂f (x i ; ω 0 ) ∂ω ∂f (x j ; ω 0 ) ∂ω - 1 d b 2 0 x i x j < ε 2 (137)

with probability 1 -e -cdε 2 . The same argument follows for the case where i = j.

Corollary 14 (Approximation of û). For large enough l > 0, with probability 1 -de

EQUATION

where ũ = b 2 0 x X.

Proof. Taking ε = log l d/d together with Lemma 13 yields the desired result.

Corollary 15 (Approximation of KX ). With probability 1 -de -c log l d

EQUATION

where KX = b 2 0 X X + b 2 1 dI.

Proof. Also by directly applying Lemma 13.

Remark. For initialization larger or equal to w i (0) ∼ N (0, I d /d), the above approximation does not depend on the scale of initialization.

Remark. For φ(x) = SoftPlus(x), b 2 0 = 1/4, b 2 1 = 0.043379. For φ(x) = sigmoid(x) = (1 + e -x ) -1 , b 2 0 = 0.042692, b 2 1 = 0.002144. Note that b 1 ≥ 0 for all smooth activations φ, and the equality holds (i.e. b 1 = 0) if and only if φ is linear. We comment that smaller b 1 entails larger variance as γ 1 → 1, and vice versa, as shown in the following section. C.9.4 THE BIAS TERM With these approximation above we proceed to calculating (129)

EQUATION

)

We first bound the error in substituting û with ũ:

EQUATION

Lemma 18. For weight matrices W ,W satisfying y -f (X) 2 = O( √ n), where f (X) = φ(XW )a with fixed a i ∼ Unif{-1/ √ h, 1/ √ h}, given (A1)-(A3), the gradient of the empirical risk defined in (11) is Lipschitz w.r.t. W in the Frobenius norm, i.e.

∂L(X; W:

EQUATION

)

Proof. Denote y 1 = φ(X W 1 )a and y 2 = φ(X W 2 )a for W 1 ,W 2 satisfying the assumption above (which can be seen as a condition on the magnitude of training loss), we have

EQUATION

)

We upper bound the two terms separately:

EQUATION

where we applied the inequality A • B F ≤ max{|A ij |} B F in (i), boundedness of φ in (ii) and Lipschitzity of φ in (iii). Similarly, for the second term

EQUATION

where we used the assumption on the training loss and the Lipscthizity of φ in (i). Combining the two terms yields the desired result.

Lemma 19. Under assumptions (A1-3) and the non-vanishing initialization, given that w i (t)w i (0) 2 = O(d -1/2 ) for all i, then we have K(t) -K(0) 2 = O(d 1/2-) for some positive ∈ Θ(1).

Proof. Recall the definition of the NTK:

EQUATION

or equivalently the matrix form Proof. Without loss of generality assume X = [X 0 , -X 0 ]. Then by (100) we have

EQUATION

EQUATION

and the flow for w -follows from symmetry. In this case one can show that from exact zero initialization, for nonlinearity satisfying φ(x) -φ(-x) = x, such as ReLU and SoftPlus,

EQUATION

And therefore the gradient flow of w + is ∂w + ∂t = 1 2n 0 2n0 i=1 y i -h 0 φ(w + x i ) + h 0 φ(-w + x i ) φ (w + x i )x i = 1 2n 0 n0 i=1 y i -hφ(w + x i ) + h 0 φ(-w + x i ) φ (w

EQUATION

The flow of w -follows from symmetry. Solving for the stationary points (i.e. gradient becomes zero), it the clear that

w (t=∞) + = -w (t=∞) - =          1 h 0 (XX ) -1

Xy, γ 1 < 0.5, 1 h 0 X(X X) -1 y, γ 1 > 0.5.

(173)

And hence the asymptotic risk is

EQUATION

)

The same conclusion holds for vanishing initialization if we assume that the trajectory stays close to that of exact zero initialization. Note that although the prediction aligns well with the experimental results, the argument in Theorem 7 does not directly apply due to the undefined derivative of ReLU at the origin, and thus this result is not rigorously justified.

F EXPERIMENT SETUP:

Optimizing the Second Layer. We compute the minimum-norm solution by directly solving the pseudo-inverse. We set n = 1000 and vary γ 1 , γ 2 from 0.1 to 3. The linear teacher model F (x) = x β is fixed as β = -1 d / √ d. For each (γ 1 , γ 2 ) we average across 50 random draws of data.

Optimizing the First Layer. For both initializations, we use gradient descent with small step size (η = 0.1) and train the model for minimally 25000 steps and till ∇ W f (X, W ) 2 F < 10 -6 . We fix n = 320 and vary γ 1 , γ 2 from 0.1 to 3 with the same linear teacher model β = -1 d / √ d. The risk is averaged across 20 models trained from different initializations.



Figures and Tables:

Figure 1: Illustration of the double descent risk curve in two-layer linear networks (SNR= 16). Brighter color indicates larger γ1 = d/n. Double descent is observed when the second layer coefficients are optimized (main figure), but not when the first layer weights are optimized (subfigure).

Figure 2: Population risk of two-layer neural networks with optimized second layer under (A1)(A2). Brighter color indicates larger γ1. (a) risk of linear network with r 2 /σ 2 = 16 and γ1 < 1. (γ1 > 1 is shown in Figure 1) (b) variance of network with ReLU activation. Black line corresponds to γ1 → ∞ predicted by Corollary 5.(c) bias of network with ReLU activation. Black line corresponds to γ1 → ∞ for linear network, which is empirically observed as an upper-bound. Note that as γ2 → 1 both bias and variance becomes unbounded.

Figure 3: Bias and variance of two-layer sigmoid network with optimized first layer under (A1)(A2). Individual

Figure4: trajectory of neurons from initialization (dark blue) to optimum (orange) on the first two dimensions (two-layer SoftPlus student and linear teacher; SNR= 1/4). For vanishing initialization the neurons stay close to one another throughout the trajectory, whereas for non-vanishing initialization the neurons stay close to initialization.

Figure6: Bias and variance of two-layer SoftPlus network with optimized first layer under (A1)(A2). Individual dotted lines correspond to different γ2 (from 0.2 to 2) which is independent to the risk. The bias and variance for both initializations is well-aligned with Theorem 7 and Theorem 8, respectively.

Figure 7: Population risk (scaled by 1/d) of two-layer ReLU network trained to fit a two-layer ReLU teacher model with h = d neurons. Brighter color corresponds to larger γ1. Similar to the linear teacher case, double descent is observed when the second layer is optimized (a) but not when the first layer is optimized (b).

.d data (empirical) symmetric data (empirical) symmetric data (prediction) (a) SoftPlus activation. .d. init (empirical) symmetric init (empirical) symmetric init (prediction) (b) sigmoid activation.

Figure 8: Bias of (a) SoftPlus and (b) sigmoid two-layer network with optimized first layer under (A1)(A2).

At initialization, x i ∼ N (0, I d ) and w k (0) ∼ N (0, d I d ). Thus for fixed x i , by Gaussian anti-concentration we have Pr|x i w k | < log d ≤ O(1/d 1/2+ 1 ) for some 1 > 0. In addition, note that w k (t) -w k (0) 2 = O(d -1/2) for all k, and therefore for i, j, k such that bias of non-vanishing initialization.

Figure 9: Bias of two-layer ReLU networks with optimized first layer under Gaussian data and linear teacher. Individual dotted lines correspond to different γ2 (from 0.2 to 2) which is independent to the risk. (a) Vanishing initialization. The bias under symmetric data is predicted by Proposition 20. (b) Non-vanishing initialization.The red and blue lines represent models optimized from i.i.d. and symmetric initialization, respectively. The bias for symmetric initialization is predicted by Theorem 8.

Peer review (iclr2020):
Overview: This work is an interesting work to understand the generalization capabilities of a two layered neural network in a high dimensional setting (samples, features and neurons tend to infinity). It studies the conditions under which the *double descent phenomenon* may be observed. Summary: The work shows that in two layered neural networks with non-linearity 1) the double descent phenomenon of the bias-variance decomposition may be observed when the second layer weights are optimized assuming that the first layer weights are constant. 2) the bias-variance decomposition does not exhibit double descent when optimizing only the first layer with both vanishing and non-vanishing initialization of weights. 3) For vanishing initalization of weights for the first layer with non-linear activation , the gradient flow solution is asymptotically close to a two layered linear network. It is independent of overparametrization. However, the condition for this is smooth activation and the result does not hold for ReLU activation. 4) For non-vanishing initilization of the weights for the first layer with non-linear activation, the gradient flow solution is well approximated by a kernel model. However, the risk is independent of overparametrization. I believe this is an interesting work that needs to be accepted.