Year: 2020

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.

heading: 1 INTRODUCTION
text: 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., 2018; Ji and Telgarsky, 2018), and provided generalization guarantees (Arora et al., 2018; Bartlett et al., 2017; Dziugaite and Roy, 2017) that align with the empirical observations.
Recently, a series of works highlighted the implicit regularization of interpolators in the overparameterized regime (Belkin et al., 2018; Spigler et al., 2018; Geiger et al., 2018; Advani and Saxe, 2017). Specifically, a second decrease in the population risk is observed when the model is further overparameterized beyond the interpolation limit, i.e. when the model achieves zero training error. This phenomenon is known as double descent, and can be precisely quantified for certain linear models (Hastie et al., 2019; Mei and Montanari, 2019; Belkin et al., 2019; Bartlett et al., 2019; Xu and Hsu, 2019). Among the recent works, Hastie et al. (2019) and Mei and Montanari (2019) explicitly derived the population risk of linear regression and random features regression models in high dimensions using tools from random matrix theory.
However, there is still a gap between the practical benefit of overparameterization and the recently proved double descent phenomenon, which is typically established under models that exhibits the following structure: the trained model solves a linear inverse problem, and the “cusp” in the risk arises from the instability of the inverse at the interpolation threshold. Moreover, given a dataset or fixed n, d, the number of parameters in the linear regression model is also fixed, i.e. the level of overparameterization cannot be altered. It is therefore unclear if the trend persists in the optimization of more complex models, for instance in two-layer neural networks where overparameterization can be controlled simply by adding more neurons.
In this work, we analyze the generalization properties of two-layer neural networks in the unregularized least squares regression setting and examine the presence/absence of the double descent phenomenon. We consider the proportional asymptotic limit where the number of samples n, input
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.
0.5 1.0 1.5 2.0 2.5 2 = h/n
0.5
1.0
1.5
2.0
2.5
3.0
ris k
empirical prediction
1 2 3 0.0
0.5
Our findings can be summarized as follows (see Figure 1):

heading: 1.1 RELATED WORKS
text: 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.
Generalization of Overparameterized Models. It is often observed that overparameterization does not result in overfitting (Neyshabur et al., 2014). In the lazy regime, generalization guarantees can be derived from the distance traveled by the parameters (Neyshabur et al., 2018; Nagarajan and Kolter, 2019), which becomes small if the model is sufficiently overparameterized (Arora et al., 2019b; Li and Liang, 2018; Allen-Zhu et al., 2018a; Cao and Gu, 2019). Compared to these guarantees that usually require significant overparameterization, our result relies on stronger data assumptions, but consequently we obtain the exact population risk instead of a vacuous upper-bound. Beyond the kernel regime, Advani and Saxe (2017); Goldt et al. (2019) analyzed the generalization dynamics of overparameterized models in the student-teacher 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).

heading: 2 PRELIMINARIES: TWO-LAYER NEURAL NETWORK
text: Consider the following bias-free two-layer neural network f : Rd → R with h hidden units
f(x) = h∑ i=1 aiφ(〈x,wi〉), (1)
where x ∈ Rd is the input, wi ∈ Rd is the weights corresponding to neuron i, ai ∈ 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 = [w1, ...wh] ∈ Rd×h for the weight matrix, a = [a1, ...ah] ∈ Rh for the coefficient vector, X = [x1, ...xn] ∈ Rd×n for the data matrix, y ∈ Rn for the corresponding vector of labels, and Φ = φ(W>X) ∈ Rh×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 : Rd → R with additive noise, and the student model aims to minimize the squared loss:
(xi, εi) i.i.d.∼ Px × Pε, yi = F (xi) + εi, L(X; f) =
1
2n n∑ i=1 (yi − f(xi))2 , (2)
where E[xi] = 0, Cov(xi) = Σ, E[εi] = 0, Var(εi) = σ2. We are interested in the population risk R(f) = EPx [(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:
dW (t) = −∇WL(X; f) dt or da(t) = −∇aL(X; f) dt, (3) 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.

heading: 3 WARM-UP: LINEAR NETWORK
text: 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 â = Φ†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: xi∼N (0, Id); (A2) Linear Teacher: F (x)=〈x,β〉, ‖β‖=r.
Denote the linear student network as f(x) = 〈x, β̂〉, where β̂ = W â and â is the least-square solution defined by Lemma 1. We write the population risk in its bias-variance decomposition.
R = Ex∼Px [‖β̂ − β‖2Σ|X,W ] = ‖E[β̂|X,W ]− β‖22︸ ︷︷ ︸ B=bias
+ tr ( Cov(β̂|X,W ) )
︸ ︷︷ ︸ V=variance , (4)
where ‖x‖2Σ = x>Σx. We compute the bias and the variance separately to obtain the risk.
Theorem 2. Given (A1)(A2) and let wi i.i.d.∼ N (0, d−1Id), at n, d, h→∞ we have
R(γ1<1) →  γ1 − γ2 γ1g2 r2 + γ2 g2 σ2, γ2 < γ1,
γ1 g1 σ2, γ2 > γ1,
R(γ1>1) →  γ1 − γ2 γ1g2 r2 + γ2 g2 σ2, γ2 < 1,
γ2g1 γ1g2 r2 + g1 + g2 g1g2 σ2, γ2 > 1.
(5) where d/n→ γ1, h/n→ γ2, g1 = |γ1 − 1|, and g2 = |γ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 ainit 6= 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
R(γ1<1) → γ1
1− γ1 σ2; R(γ1>1) → γ1 − 1 γ1 r2 + 1 γ1 − 1 σ2. (6)
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.

heading: 4 NONLINEAR MODEL: OPTIMIZING THE SECOND LAYER
text: 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 â = Φ†y. We again consider the following bias-variance decomposition.
R = Ex∼Px [∥∥φ(x>W )â− F (x)∥∥2
2 ∣∣X,W] (7) = Ex [∥∥E[φ(x>W )â∣∣X,W ]− F (x)∥∥2 2 ] ︸ ︷︷ ︸
B=bias
+Ex [∥∥φ(x>W )â− E[φ(x>W )â]∥∥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 wi i.i.d.∼ N (0, d−1Id), when n, d, h→∞, we have
V =  σ2 γ2 1− γ2 , γ2 < 1, σ2 lim ξ→0 − [ γ2 ∂ ∂x m1(ξ, c1x, c2x) ∣∣∣ x=0 + ∂ ∂x m2(ξ, c1x, c2x) ∣∣∣ x=0 + γ2 − 1 ξ2 ] γ2 > 1.
in which m1(ξ, ρ, τ), m2(ξ, ρ, τ) is the unique solution in {|m1|, |m2| < 1/=ξ} of
m−11 =−ξ−ρ−γ −1 1 γ2τ 2m1−c1m2+ τ2γ−11 γ2m 2 1(c2m2−τ)−2τc2m1m2+c22m1m22 m1(c2m2−τ)−γ1γ−12 , (8)
m−12 =−ξ−rγ2m1+ γ2c2m
2 1(c2m2− τ)
m1(c2m2−τ)−γ1γ−12 , (9)
where variables ξ, ρ, τ satisfies =ξ > 0 or ξ < 0, ρ > τ > 0, and constants c1, c2 defined as,
c1 = E[φ(G)2]− E[φ(G)]2, c2 = E[Gφ(G)]2, (10)
for G ∼ N (0, 1) and =ξ denoting the imaginary part of ξ. Remark. c1 ≥ c2 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→∞) = σ 2min{γ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) andwi i.i.d.∼ N (0, Id), 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.

heading: 5 NONLINEAR MODEL: OPTIMIZING THE FIRST LAYER
text: 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 ai ∼ Unif{−1/ √ h, 1 √ h}, and the first layer W is optimized with the following update
∂W (t) ∂t = −∂L(X;W ) ∂W = 1 n n∑ i=1 [ yi − a>φ(W>xi) ] xi[φ ′(x>i W ) ◦ a], (11)
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 Marč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: wVani (0) ∼ N (0, Id/dh1+ ); Non-vanishing: wNVi (0) ∼ N (0, Id/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).

heading: 5.1 VANISHING INITIALIZATION
text: 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) 6= 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
R(f̂)→ max {
0, γ1 − 1 γ1
} r2 +
min{γ1, 1} |1− γ1| σ2. (12)
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.

heading: 5.2 NON-VANISHING INITIALIZATION
text: 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 initf(x),∇W initf(y)〉. The converged parameters under this linearized dynamics has the following closed-form:
vec(W ∗) ≈ vec(W init) + ∆; ∆ = J†(y − f init(X)); J[i,j] = ∇vec(W init)jf init(xi), (13) where J ∈ Rn×(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. aiwiniti = −ajwinitj . Theorem 8. Given (A1-4) and let n, d, h→∞, the stationary solution f̂ has the following risk
R(f̂)→
( γ1 − 1
2γ1 + γ1(γ1 + γ1m+m− 2) + 1 2γ1 √ γ1(γ1 +m(γ1(m+ 2) + 2)− 2) + 1
) r2
+
( γ1 + γ1m+ 1
4 √ γ1(γ1 +m(γ1(m+ 2) + 2)− 2) + 1 − 1 4
) σ2, (14)
where m = b21/b 2 0, b 2 0 = E[φ′(G)]2, and b21 = E[φ′(G)2]− b20, 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.

heading: 5.3 COMPARING THE INITIALIZATIONS
text: 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.

heading: 6 DISCUSSION AND FUTURE WORKS
text: 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).

heading: ACKNOWLEDGEMENT
text: We thank Xiuyuan Cheng, Xuechen Li, Yiping Lu, Atsushi Nitanda, Shengyang Sun and anonymous reviewers for helpful comments and feedback. JB and DW were partially funded by LG Electronics and NSERC. JB and MAE were supported by the CIFAR AI Chairs program. TS was partially supported by JSPS Kakenhi (26280009, 15H05707 and 18H03201), Japan Digital Design and JST-CREST.

heading: A ADDITIONAL FIGURES AND PLOTS

heading: SUMMARY OF THE PRESENCE / ABSENCE OF DOUBLE DESCENT
text: Singularity in 2nd Layer Trained (RF) Vanishing Init. Non-vanishing Init. Bias γ1: No; γ2: Yes γ1: No; γ2: No γ1: No; γ2: No
Variance γ1: No; γ2: Yes γ1: Yes; γ2: No γ1: No; γ2: No

heading: B BACKGROUND

heading: B.1 ROTATIONAL INVARIANCE
text: 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 UX are the same.
Lemma 10 (Rotational Invariance). Denote A(X) ∈ Rd×d a matrix function of X ∈ Rd×n. If A(X) satisfies that A(UX) = UA(X)UT for all unitary U , then
EX [βTA(X)β] = 1
d βTβEX [tr (A(X))]. (15)
for any fixed nonzero β ∈ Rd and random matrix X with each entry i.i.d. Xij ∼ N (0, σ2).
Proof. Choose a set of Unitary matrices {Ui}di=1 such that U>i β = ‖β‖ ei, where ei is the i-th canonical vector in Rd. Since UiX ∼ X , we have E[A(X)] = E[A(UiX)] = UiE[A(X)]U>i and hence
E[βTA(X)β] = 1
d d∑ i=1 E[βTUiA(X)U>i β] = 1 d d∑ i=1 eTi E[A(X)]ei = βTβ d E[tr (A(X))]. (16)
Note that the property also holds for matrix function A(X) that satisfies A(XU) = UA(X)UT , 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).

heading: B.2 MARČENKO–PASTUR LAW
text: For a real symmetric random matrix A ∈ Rp×p, define its empirical spectral density as
µA(dλ) = 1
p p∑ i=1 δλi(A)(λ)dλ,
where δa(x) = δ(x − a) is the Dirac delta function. Assume A = Wp ∼ Wp(I, n) is a Wishart matrix, i.e. Wp = X>X/n and X ∈ Rn×p is random Gaussian matrix with each column i.i.d. Xi ∼ N (0, I). Marč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(γ)(λ):
µMP(γ)(dλ) = [1− γ−1]+δ0(λ)dλ+ 1
2πγλ
√ ((1 + √ γ)2 − λ)(λ− (1−√γ)2)dλ. (17)
We say µMP(γ) is the density of the Marč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 γ 6= 1. The explicit form of Marč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 Marčenko–Pastur distribution, i.e. for any bounded f(λ) ∈ C(S), as n, p→∞ with p/n = γ, almost surely∫
S
f(λ) · µW (dλ)→ ∫ fS(λ) · µMP (dλ). (18)
One implication is the following trace concentration on the inverse Wishart matrix for γ < 1,
tr ( X>X ) = tr
( 1
p W−1p
) = 1
p p∑ i=1
1
λi(Wp) = ∫ 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.

heading: B.3 ORTHOGONAL POLYNOMIALS
text: Orthogonal polynomials are useful in the analysis of nonlinear random matrices. Suppose φ is a function in L2(R, µG), where G ∼ N (0, 1), µG(dx) = ( √ 2π)−1e−x
2/2dx is the Gaussian measure. For n ≥ 0, define the Hermite polynomials
Hn(x) = (−1)ne−x 2/2 ∂
n
∂xn e−x 2/2, (20)
Note that orthogonality can be easily verified: E[Hj(G)Hk(G)] = ∫ R Hj(x)Hk(x)µG(dx) = j! · δjk. (21)
Since {Hi(x)}∞i=0 forms a set of orthogonal basis in L2(R, µG), the function φ(x) can be expanded under the Hermite basis as
φ(x) = ∞∑ i=0 ciHi(x) = ∞∑ i=0 ( 1 k! ∫ R φ(a)Hi(a)µG(da) ) Hi(x). (22)
Following Cheng and Singer (2013), we mainly focus on the first few terms of the Hermite expansion. One can check that H0(x) = 1, H1(x) = x, and φ(x) can be expanded as
φ(x) = c0 + c1x+ φ⊥(x), (23)
where c0 = E[φ(G)], c1 = E[Gφ(G)], and terms in the RHS are orthogonal to one another in L2(R, µG). Taking square and expectation over both sides yields
E[φ(G)2] = E[φ(G)]2 + E[Gφ(G)]2 + E[φ⊥(G)2], (24)
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) = c0 +c1x 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.

heading: C PROOF OF MAIN RESULTS

heading: C.1 PROOF OF LEMMA 1
text: Given features X ∈ Rd×n, labels y ∈ Rd and model parameters θ, the gradient flow of θ on the squared loss ∥∥y −X>θ∥∥2 2 can be written as
∂θ(t)
∂t =
1 n X(y −X>θ(t)). (25)
Thus with initialization θ0, the solution of this ODE at time t can be written in explicit form
θ(t) = e− t nXX > θ0 + (XX >)† ( I − e− tnXX > ) Xy. (26)
Since θ0 = 0, taking t→∞ yields the desired result.

heading: C.2 PROOF OF THEOREM 2
text: 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
β̂ = W (X>W )†y = (XX>)−1Xy. (27)
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→∞,
B → 0; V → γ1 1− γ1 σ2. (28)
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
β̂ = W (W>XX>W )−1W>Xy. (29)
Denote W = UΣV > the singular value decomposition. We perform the block decomposition:
Σ = [ Σ0 0 ] , X = [ X0 X1 ] , (30)
where Σ0 ∈ Rh×h, X0 ∈ Rh×n, X1 ∈ R(d−h)×n, and notice that X0, X1 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
V = tr ( W ( W>XX>W )−1 W>X>σ2XW ( (W>XX>W )−1 W> ) = σ2tr ( W>W ( W>XX>W )−1) = σ2tr ( Σ>Σ ( Σ>U>XX>UΣ
)−1) ∼ σ2tr ( Σ>Σ ( Σ>XX>Σ )−1) = σ2tr ( (X0X > 0 ) −1)→ σ2 γ2
1− γ2 , (31)
where the last equality follows from Appendix B.2. Similarly for the bias term we have B = ∥∥∥W (W>XX>W )−1W>XX>β − β∥∥∥2
2 = β> ( W (W>XX>W )−1W>XX> − Id )>( W (W>XX>W )−1W>XX> − Id ) β
(i) = r2
d tr
(( W (W>XX>W )−1W>XX> − Id )>( W (W>XX>W )−1W>XX> − Id )) = r2
d tr
(( UΣV >(V Σ>U>XX>UΣV >)−1V Σ>U>XX> − Id )>( · · · ))
= r2
d tr
(( UΣ(Σ>XX>Σ)−1Σ>XX>U> − UU> )>( · · · ))
= r2
d tr
(( Σ(Σ>XX>Σ)−1Σ>XX> − Id )>( · · · )) , (32)
where symmetric arguments are omitted as (· · · ), and (i) follows from the rotational invariance argument introduced in Lemma 10 and that β>β = r2. By the block decomposition (30),
Σ ( Σ>XX>Σ )−1 Σ>XX> − Id = [ 0 (X0X > 0 ) −1X0X > 1
0 −Id−h
] . (33)
Therefore the bias term simplifies to
B = r2
d tr
(( Σ(Σ>XX>Σ)−1Σ>XX> − Id )>( Σ(Σ>XX>Σ)−1Σ>XX> − Id )) = r2
d
( tr ( (X0X > 0 ) −1X0X > 1 X1X > 0 (X0X > 0 ) −1)+ (d− h))
→ r 2
d ( (d− h)h n− h− 1 + d− h ) → γ1 − γ2 γ1(1− γ2) r2. (34)
Thus we have obtained that as n, d, h→∞
B → γ1 − γ2 γ1(1− γ2) r2. (35)
Case III: W>X has full row rank, i.e. γ1 > 1, γ2 > 1. Similarly, the least squares solution is
β̂ = WW>X(X>WW>X)−1y, (36)
Simplifying the variance: V = tr ( WW>X ( X>WW>X )−1 σ2 ( X>WW>X )−1 X>WW> ) = σ2tr ( WW>UΣV > ( V Σ>U>WW>UΣV > )−2 V Σ>U>WW>
) ∼ σ2tr ( WW>Σ ( Σ>WW>Σ )−2 Σ>WW> ) . (37)
where we applied the SVD of X = UΣV > and the rotational invariance argument. Using a similar block decomposition on W :
Σ = [ Σ0 0 ] , W = [ W0 W1 ] , (38)
where Σ0 ∈ Rn×n,W0 ∈ Rn×h,W1 ∈ R(d−n)×h, and W0,W1 independent. We thus simplify V as
V = σ2tr ( WW>Σ(Σ>WW>Σ)−2Σ>WW> ) = σ2tr ([ W0W > 0 Σ0(Σ > 0 W0W > 0 Σ0) −2Σ0W0W > 0 · · ·
· · · W1W>0 Σ0(Σ>0 W0W>0 Σ0)−2Σ0W0W>1 ]) = σ2 ( tr ( Σ−T0 Σ −1 0 ) + tr ( W1W > 0 ( W0W > 0 )−1 Σ−T0 Σ −1 0 ( W0W > 0 )−1 W0W > 1
)) = σ2tr (( X>X )−1) + σ2tr ( W>1 W1W > 0 ( W0W > 0 )−1 ( Σ>0 Σ0 )−1 ( W0W > 0 )−1 W0 ) . (39)
Hence we obtain the following expression on the variance
V → σ2 1 γ1 − 1
+ σ2(d− n)EW,XV tr (( W0W > 0 )−1 ( Σ>0 Σ0 )−1)→ σ2( 1 γ1 − 1 + 1 γ2 − 1 ) .
(40)
We omit the derivation of bias, which follows a similar derivation:
B → γ2(γ1 − 1) γ1(γ2 − 1) r2. (41)
Combining Case I, II, III yields theorem 2.

heading: C.3 PROOF OF PROPOSITION 3
text: Given the squared loss, one can derive the dynamics of W with fixed second layer a w.r.t the loss:
∂W (t) ∂t = − 1 n X(y −X>W (t)a)a>. (42)
Note that the update of W can be written as a linear combination of a. Since W (0) = 0, we can write W (t) = ŵ(t)a> for some ŵ. The corresponding flow on ŵ is
∂ŵ(t) ∂t = − 1 n X(y −X>ŵ(t) ‖a‖22), (43)
which gives the following solution
ŵ∗ = 1
‖a‖22 X†y ⇒ β̂ = W ∗a = X†y. (44)
Thus gradient flow on the first layer leads to the minimum-norm solution on the input features.

heading: C.4 PROOF OF THEOREM 4
text: Following the bias-variance decomposition (7), the variance term can be written as (σ2 omitted) V = Ex,ε [ ‖a>φ(W>x)− Eεa>φ(W>x)‖22 ] = Ex [∥∥∥[φ(W>X)]†φ(W>x)∥∥∥2 2
] = tr ([ φ(X>W ) ]†[ φ(W>X)
]†Ex[φ(W>x)φ(W>x)>]) = tr ([ φ(X>W ) ]†[ φ(W>X) ]† KW ) , (45)
where we define the expected non-linear Gram matrix KW ∈ Rh×h as KW = Ex [ φ(W>x)φ(W>x)> ] . (46)
and for each entry we have (KW )[i,j] = Ex [ φ(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) ∈ Rh×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− γ−12 ]+δ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 KW ) = lim ξ→0− tr (( ΦΦ> − ξI )−1 KW ) = lim ξ→0− Vξ,
where the continuity and boundness of Vξ at ξ = 0− is guaranteed by Lemma 11 when γ2 = h/n 6= 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→∞,
tr h−1(ΦΦ> − ξI)−1 n h
KW 1 + h−1tr ( h ( ΦΦ> − ξI )−1 KW
) − h−1I  ≤ 1 h ∥∥∥∥∥∥∥ ( ΦΦ>−ξI )−1 − n h KW 1 + h−1tr ( h ( ΦΦ>−ξI )−1 KW ) −1 ∥∥∥∥∥∥∥ F
· ∥∥∥∥∥∥nh KW1 + h−1tr(h(ΦΦ>−ξI)−1KW) ∥∥∥∥∥∥
2
≤ 1 h O(n−1/2+ )O(n1/2)→ 0. (47)
where we have used the inequality tr (AB) ≤ ‖A‖F ‖B‖2, and equivalently by taking ξ → 0 we get
lim ξ→0 lim n,d,h→∞
n
h
tr (( ΦΦ> − ξI )−1
KW ) 1 + tr (( ΦΦ> − ξI )−1 KW ) = 1. (48)
Therefore limξ→0 limn,d,h→∞ n/h · tr (( ΦΦ> − ξI )−1 KW ) = γ2/(1− γ2). Note that ∂Vξ/∂ξ =
tr ( ξ(ΦΦ> − ξI)−2KW ) 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:
V → γ2 1− γ2 . (49)
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 Ã defined in (56); this translates the calculation of Vξ into the calculation of the Stieltjes transform of Ã (57), (61), (62).
C.5 DERIVING THE VARIANCE V FOR h > n
Step 1. An equivalent expression. For notational simplicity we omit the magnitude σ:
V = tr ( φ(W>X) ( φ(X>W )φ(W>X) )−2 φ(X>W )KW ) . (50)
and due to the same continuity argument as in Case 1 we have
V = lim ξ→0
1
n
[ tr ( S(S>S − ξIn)−2S>KW ) ] = 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 Σ = diagh×n(φ1, · · · , φn) ∈ Rh×n is a tall diagonal matrix, and U = [u1, · · · ,uh] ∈ Rh×h is the set of orthogonal eigenvectors of SS> = φ(W>X)φ(X>W )/n, and V ∈ Rn×n is the set of orthogonal eigenvectors of S>S = φ(X>W )φ(W>X)/n. Now the variance can be written as
Vξ = 1 n tr ( UΣ(Σ>Σ + ξIn) −2ΣU>KW ) . (51)
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Σ(Σ>Σ + ξIn) −2ΣU>KW ) ] → 1
n
[ tr ( UΣ(Σ>Σ + ξIn) −2ΣU>K̃W ) ] , (52)
in which K̃W is the approxmation of KW defined in Lemma 16. Writing the trace explicitly (denote eigenvalues λi = φ2i , and φn+1 = · · · = φh = 0), we have
Vξ → 1 n tr ( UΣ(Σ>Σ + ξIn) −2ΣU>K̃W ) = γ2 1 h h∑ i=1 λi (λi + ξ)2 u>i K̃Wui. (53)
Since the positive support of spectrum λ is lower bounded and the density at 0 is 1− γ−12 , we have
Vξ → γ2 1
h lim h,d,n→∞ h∑ i=1 λi (λi + ξ)2 u>i K̃Wui = γ2 ∫ λ (λ+ ξ)2 µ∞(dλ) = γ2 ∫ λ>ρ
λ
(λ+ ξ)2 µ+∞(dλ),
(54) where we define µn(x) = 1h ∑h i=1 δλi(x)u > i K̃Wui and its positive part µ + n (x). Hence we have
V = lim ξ→0 Vξ = lim ξ→0 γ2 ∫ λ>ρ
λ
(λ+ ξ)2 µ+∞(dλ) = γ2 ∫ λ>ρ 1 λ µ+∞(dλ). (55)
We define the following matrix Ãn(ρ, ς, τ) ∈ RN×N where N = n+ h:
Ãn(ρ, ς, τ) =
[ ρIh + ς1h1 > h + τQ S
S> 0n
] . (56)
And denote the Stieltjes transform of Ãn as
m̃n(ξ, ρ, ς, τ) = 1 n tr ( (Ãn(ρ, ς, τ)− ξIN )−1 ) . (57)
Then following the definition of K̃W one can show that
m̃n(ξ, rx, sx, tx) = 1
n tr
([ K̃Wx− ξIh S
S> −In
]−1) , (58)
and taking matrix derivative gives
− ∂ ∂x m̃n(ξ, rx, sx, tx) ∣∣∣ x=0 = 1 n tr ([ −ξIh S S> −In ]−1 [ K̃W 0 0 0 ] [ −ξIh S S> −In ]−1)
= 1
n tr
([ ξ2Ih + SS
> 0 0 In + S >S
]−1 [ K̃W 0
0 0
])
= 1 n tr ( U(ΣΣ> + ξ2Ih) −1U>K̃W ) = 1 n h∑ i=1 1 λ+ ξ2 u>i K̃Wui.
(59)
Denote the limit m̃(ξ, ρ, ς, τ) = limn,h,d→∞ m̃n(ξ, ρ, ς, τ), and the derivative is given as
− ∂ ∂x m̃(ξ, rx, sx, tx) ∣∣∣ x=0 = lim h,d,n→∞ 1 n h∑ i=1 1 λ+ ξ2 u>i K̃Wui
= γ2 ∫ λ≥0 1 λ+ ξ2 µ∞(dλ) = γ2 − 1 ξ2 + γ2 ∫ λ>ρ 1 λ+ ξ2 µ+∞(dλ). (60)
For simplicity we define the following function on ξ:
q(ξ) = − ∂ ∂x m̃(ξ, rx, sx, tx) ∣∣∣ x=0 , (61)
also denote q+(ξ) = q(ξ)− γ2−1ξ2 = γ2 ∫ λ>ρ 1 λ+ξ2µ + ∞(dλ). Comparing with (55) yields:
V = lim ξ→0 q+(ξ). (62)
Step 2. Calculating q(ξ) and mn(ξ, ρ, ς, τ) 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 An by subtracting the off-diagonal entries of the upper-left block of Ãn:
An(ρ, τ) =
[ ρIh + τQ S
S> 0n
] , (63)
where S = S̃ − a0Ip×n, i.e. Sik = φ(w>i xk) − a0 = ϕ(w>i xk), a0 = E[φ(x)]. The Stieltjes transform of An given by mn(ξ, ρ, τ) = 1n tr ( (An(ρ, τ)− ξIN )−1 ) . The following Lemma shows that m̃n and mn have the same limit: Lemma 12. when n→∞ and for =ξ > 0 or ξ < 0, we have mn(ξ, ρ, τ)→ m̃n(ξ, ρ, ς, τ).
Proof. By definition of Ān and An, Ãn(ρ, ς, τ)−An(ρ, τ) = [
1h 0h 0n 1n ] [ ξ a0 a0 0 ] [ 1h 0h 0n 1n ]> (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:
sup x |F Ãn(x)− FAn(x)| ≤ O
( n−1 ) , (65)
where FM is the empirical spectral distribution of M ∈ Rn×n. The claim follows from the Stieltjes continuity theorem (e.g. Section 2.4 in Tao (2012)).
To calculate the Stieltjes transform mn, we take advantage of the block structure of An:
m1,n(ξ, ρ, τ) = 1 p tr ( (An(ρ, τ)− ξIN )−1[1..p,1..p] ) , (66)
m2,n(ξ, ρ, τ) = 1 n tr ( (An(ρ, τ)− ξIN )−1[p+1..p+n,p+1..p+n] ) . (67)
One can observe that mn(ξ, ρ, τ) = γ2m1,n(ξ, ρ, τ) +m2,n(ξ, ρ, τ).
In the following equations we omit the subscript n, as well as dependency on ρ, ς, τ . Following Hastie et al. (2019), we rewrite the matrixAn = A = [ A∗ a a> 0 ] , whereA∗ is a (N−1)× (N−1)
matrix with last column and row of A removed and s the activation vector:
A∗ =
[ ρIh + τQ S∗
S>∗ 0n−1
] ; a> = [φ(W>xn) > 0>n−1] = [s > 0>n−1]. (68)
Hence by the block matrix inverse formula (A− ξIN )−1 = [ ∗ ∗ ∗ [−ξ − a>(A∗ − ξIN−1)−1a]−1 ] . (69)
Plugging this back in the Stieltjes transform we have
m2,n(ξ, ρ, τ) = 1 n tr ( (An(ρ, ς, τ)− ξIN )−1[h+1,N ] ) = Ea [ (An(ρ, τ)− ξIN )−1 ] NN
= Ea [( − ξ − a>(A∗ − ξIN−1)−1a )−1] . (70)
To obtain an asymptotic description of m2,n, we perform the orthonormal decomposition on the nonlinearity ϕ introduced in Cheng and Singer (2013): ϕ(x) = a1x + ϕ⊥(x), where a1 = Ex∼N (0,1)[xϕ(x)]. For each wi(1 ≤ i ≤ h), we perform the following orthonormal decomposition (along the direction of xn and the direction of w̃i perpendicular to xn):
wi = w > i xn︸ ︷︷ ︸ ηi xn ‖xn‖ + w̃i = ηi xn ‖xn‖ + w̃i. (71)
We can thus simplify the activation vector a as
a> =  1√ n ϕ(‖xn‖η1) · · · 1√ n ϕ(‖xn‖ηh) 0 · · · 0︸ ︷︷ ︸
n−1  =  1√ n a1‖xn‖η1 · · · 1√ n a1‖xn‖ηh 0 · · · 0︸ ︷︷ ︸
n−1  ︸ ︷︷ ︸
a>1 =[s > 1 ,0 >]
+  1√ n ϕ⊥(‖xn‖η1) · · · 1√ n ϕ⊥(‖xn‖ηh) 0 · · · 0︸ ︷︷ ︸
n−1  ︸ ︷︷ ︸
a>2 =[s > 2 ,0 >]
. (72)
and for 1 ≤ i 6= j ≤ h, 1 ≤ k ≤ n− 1,
Qij = ( ηi xn ‖xn‖ + w̃i )>( ηj xn ‖xn‖ + w̃j ) = ηiηj + w̃
> i w̃j︸ ︷︷ ︸ Q̃ij . (73)
Similarly we decompose the activation function in S,
Sik = 1√ n ϕ ( ηi x>nxk ‖xn‖ + w̃>i xk ) = 1√ n a1ηi x>nxk ‖xn‖ + 1√ n a1w̃ > i xk + 1√ n ϕ⊥ ( ηi x>nxk ‖xn‖ + w̃>i xk ) =
1√ n ϕ(w̃>i xk)︸ ︷︷ ︸ S̃ik + 1√ n a1ηi x>nxk ‖xn‖︸ ︷︷ ︸ a1ηiuk + 1√ n
[ ϕ⊥ ( ηi x>nxk ‖xn‖ + w̃>i xk ) − ϕ⊥(w̃>i xk) ] ︸ ︷︷ ︸
Eik
. (74)
We thus have an equivalent expression of matrix A∗
A∗ =
[ ρIh + τQ̃ S̃∗
S̃>∗ 0n−1
] + [ tηη> a1ηu >
a1uη > 0n−1
] + [ E0 E1 E>1 0n−1 ] = [ ρIh + τQ̃ S̃∗
S̃>∗ 0n−1 ] ︸ ︷︷ ︸
Ã∗
+
[ η 0h
0n−1 u ] ︸ ︷︷ ︸
U
[ τ a1 a1 0 ] ︸ ︷︷ ︸
C
[ η 0h
0n−1 u ]> ︸ ︷︷ ︸
U>
+ [ E0 E1 E>1 0n−1 ] ︸ ︷︷ ︸
E
= Ã∗ + UCU > + E. (75)
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 m2,n in (70)
m2,n(ξ, ρ, τ) = Ea [( − ξ − a>(A∗ − ξIN−1)−1a )−1] → Ea [( − ξ − a>(Ã∗ − ξIN−1 + UCU>)−1a
)−1] = Ea [( − ξ − a>(Ã∗ − ξIN−1)−1a︸ ︷︷ ︸
u
+
a>(Ã∗ − ξIN−1)−1U︸ ︷︷ ︸ v>
( C−1 + U>(Ã∗ − ξIN−1)−1U )−1 ︸ ︷︷ ︸
S
U(Ã∗ − ξIN−1)−1a︸ ︷︷ ︸ v
)−1] .
(76)
We bound each term u,v, S to compute m2,n. For u Eau = Ea [ a>(Ã∗ − ξIN−1)−1a ] = tr ( Es[ss>](Ã∗ − ξIN−1)−1[1..h,1..h] ) = bγ2m1,n(ξ, ρ, τ),
(77)
where b = Ex∼N (0,1)[ϕ(x)2] = Ex∼N (0,1)[(φ(x)− Eφ(x))2] = r. From a standard concentration of measure argument we have that as n, h, d→∞
u→ Eau = rγ2m1,n(ξ, ρ, τ). (78) And for v (note that U is dependent on a)
Eav> = Ea [ a>(Ã∗ − ξIN−1)−1U ] = Ea [ [s>, 0>n−1](Ã∗ − ξIN−1)−1 [ η 0h
0n−1 u ] ] = [ Es[s>(Ã∗ − ξIN−1)−1[1..h,1..h]η] 0 ] = [ tr ( (Ã∗ − ξIN−1)−1[1..h,1..h]Es[ηs >] )
︸ ︷︷ ︸ v
0 ] .
(79) where v = tr ( (Ã∗ − ξIN−1)−1[1..h,1..h]Es[ηs >] ) = a1 √ γ22/γ1m1,n(ξ, ρ, τ). We thus have
v → Eav = [ a1 √ γ22/γ1m1,n(ξ, ρ, τ) 0 ] . (80)
as n, d, p→∞. And finally for S, EaS−1 = Ea [ C−1 + U>(Ã∗ − ξIN−1)−1U ] = [ τ a1 a1 0 ]−1 + Ea [[ η 0h 0n−1 u ]> (Ã∗ − ξIN−1)−1 [ η 0h 0n−1 u ]]
=
[ 0 1/a1
1/a1 −τ/a21
] + Ea [ η>(Ã∗ − ξIN−1)−1[1..h]η 0
0 u>(Ã∗ − ξIN−1)−1[h+1..h+n−1]u
]
=
[ 0 1/a1
1/a1 −τ/a21
] +  tr((Ã∗ − ξIN−1)−1[1..h]Ea[ηη>]) 0 0 tr ( (Ã∗ − ξIN−1)−1[h+1..h+n−1]Ea[uu >] ) 
=
[ 0 1/a1
1/a1 −τ/a21
] + [ γ2/γ1m1,n(ξ, ρ, τ) 0
0 m2,n(ξ, ρ, τ) ] = [ γ−11 γ2m1,n(ξ, ρ, τ) 1/a1 1/a1 m2,n(ξ, ρ, τ)− τ/a2 ] . (81)
And hence as n, d, h→∞, S−1 → EaS−1 = [ γ−11 γ2m1,n(ξ, ρ, τ) 1/a1 1/a1 m2,n(ξ, ρ, τ)− τ/a2 ] . (82)
Therefore by combining (78), (80), (82), we arrive at the following expression on m2,n, m2,n(ξ, ρ, τ)→ Ea [( − ξ − u+ v>Sv )−1] → ( − ξ − u+ v>Sv )−1 → ( − ξ − rγ2m1,n + ( a1 √ γ22/γ1m1,n )2 [ γ−11 γ2m1,n(ξ, ρ, τ) 1/a1 1/a1 m2,n(ξ, ρ, τ)− τ/a2 ]−1 [1,1]
)−1 = ( −ξ − rγ2m1,n + γ2a 2 1m 2 1,n(a 2 1m2,n − τ)
m1,n(a21m2,n − τ)− γ1γ −1 2
)−1 . (83)
Similarly we can calculate m1,n(ξ, ρ, τ) as
m1,n(ξ, ρ, τ)→( −ξ−ρ−γ−11 γ2τ2m1,n−rm2,n+ τ2γ−11 γ2m 2 1,n(a 2 1m2−τ)−2τa21m1,nm2,n+a41m1,nm22,n
m1,n(a21m2,n − τ)− γ1γ −1 2
)−1 .
(84)
Uniqueness in (84)(83) follows from (Hastie et al., 2019, Sec B.1) and is omitted.

heading: C.6 PROOF OF COROLLARY 5
text: In this section we take the limit γ1 →∞. In this case (8), (9) simplify to
m2 = (−ξ − rγ2m1)−1 , (85) m1 = ( −ξ − ρ− γ2τ2m1 − rm2 )−1 . (86)
Recall that m(ξ, ρ, τ) = γ2m1(ξ, ρ, τ) +m2(ξ, ρ, τ). By taking the derivative we have
−q(ξ) = ∂ ∂x m(ξ, rx, tx) ∣∣∣ x=0 = r ∂ ∂ρ m(ξ, ρ, 0) ∣∣∣ ρ=0 + t ∂ ∂τ m(ξ, 0, τ) ∣∣∣ τ=0
= rγ2 ∂
∂ρ m1(ξ, ρ, 0) ∣∣∣ ρ=0 + r ∂ ∂ρ m2(ξ, ρ, 0) ∣∣∣ ρ=0 + tγ2 ∂ ∂τ m1(ξ, 0, τ) ∣∣∣ τ=0 + t ∂ ∂τ m2(ξ, 0, τ) ∣∣∣ τ=0 .
(87)
Observe that (85), (86) constitutes a set of implicit functions. Thus differentiating the two functions with respect to τ, ρ and then substitute by ρ = τ = 0 gives
q(ξ) =
( r(γ2 − 1)− ξ2 )(√ (r(γ2 − 1) + ξ2)2 − 4rγ2ξ2 + r(γ2 − 1) + ξ2 ) 2ξ2 √ (r(γ2 − 1) + ξ2)2 − 4rγ2ξ2 . (88)
Hence by (62) we obtain the asymptotic variance:
V(γ1→∞) = lim ξ→0 q+(ξ) = lim ξ→0
( q(ξ)− γ2 − 1
ξ2
) = 1
γ2 − 1 . (89)
Combining the case where γ2 < 1 in Theorem 4 completes the proof.
Remark. For φ(x) = ReLU(x), c1 = 1/2 − 1/(2π), c2 = 1/4. For φ(x) = SoftPlus(x) = log(1 + ex), numerical integration yields c1 ≈ 0.2715, c2 = 1/4.

heading: C.7 PROOF OF COROLLARY 6
text: C.7.1 UNBOUNDED BIAS WHEN h = n
From the bias-variance decomposition (7), the bias B is written as B = r 2
d tr (Q1 +Q2 + Id), where
Q1 = X[φ(W >X)]†KW [φ(X >W )]†X>; Q2 = X[φ(W >X)]†W>. (90)
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 Q1
1 d tr (Q1) = 1 d tr ( X[φ(W>X)]†KW [φ(X >W )]†X> ) = 1 d tr ( KW [φ(X >W )]−1X>X[φ(W>X)]−1 )
≥ 1 d λmin(KW )tr
( [φ(X>W )]−1X>X[φ(W>X)]−1 ) = λmin(KW )
n tr
( (SS>)−1 · 1
d X>X
) .
Since W and X/ √ d are Rd×n = Rd×p follows the same distribution where each entry i.i.d. N (0, 1/d),
1
dλmin(KW ) tr (Q1) ≥
1 n tr
( (SS>)−1 · 1
d X>X ) ∼ 1 n tr (( S>S )−1 ·W>W) = 1 n tr (( S>S )−1 · (I +Q)) = lim ξ→0 − ∂ ∂x m̃n(ξ, x, 0, x)→∞, (91)
where in Section C.5 we have showed (91) is unbounded when n → ∞. Moreover, by (154) and Weyl’s theorem we have λmin(KW ) = O(1), and thus d−1tr (Q1) is unbounded. For d−1tr (Q2),
1 d tr (Q2) = 1 d tr ( W>X[φ(W>X)]−1 ) ≤ 1 d λmax(φ(W >X)−1)tr ( W>X ) = O(1). (92)
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 Q1 = X ( φ(X>W )φ(W>X) )−1 φ(X>W )KWφ(W >X) ( φ(X>W )φ(W>X) )−1 X>, (93)
Q2 = X ( φ(X>W )φ(W>X) )−1 φ(X>W )W>. (94)
Therefore we have
2 d tr (Q1) = 2tr
( X>X
d
( φ(X>W )φ(W>X) )−1 φ(X>W )KWφ(W >X) ( φ(X>W )φ(W>X) )−1)
≤ 2λmax ( X>X
d
) tr (( φ(X>W )φ(W>X) )−1 φ(X>W )KWφ(W >X) ( φ(X>W )φ(W>X) )−1) = O(1) · tr ( φ(W>X) ( φ(X>W )φ(W>X) )−2 φ(X>W )KW
) ≤ O(1) · λmax ( φ(W>X) ( φ(X>W )φ(W>X) )−2 φ(X>W ) ) · tr (KW )
= O(1) · σ−2min(φ(X >W ))tr (KW ) = O(1) ·O(n−1) ·O(n) = O(1), (95)
in which we used Lemma 11 and tr (AB) ≤ λmax(A)tr (B) for positive semi-definite A,B. Similarly
2 d tr (Q2) = 2tr
(( φ(X>W )φ(W>X) )−1 φ(X>W ) · 1
d W>X ) ≤ tr ((( φ(X>W )φ(W>X) )−1 φ(X>W ) ) (...)> ) + tr ( d−2X>WW>X
) ≤ n · σmin ( φ(X>W ) )−2 + 1
d λmax
( 1
d XX>
) tr ( WW> ) = O(1), (96)
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.

heading: C.8 PROOF OF THEOREM 7
text: For simplicity we assume n, d, h to be even and let d0 = d/2, n0 = n/2 and h0 = h/2. Since the second layer is fixed ai ∼ Unif{−1/ √ h, 1/ √ h}, we let ai = 1/ √ h and ai+h0 = −1/ √ h for all 1 ≤ i ≤ h0. We therefore write a> = h−1/2[1h0 ,−1h0 ]> and W = [W+,W−]:
f(x;W−,W+) = a >φ(W>x) = 1√ h 1>φ(W>+ x)− 1√ h 1>φ(W>−x). (97)
The empirical risk can thus be written as
L(X;W+,W−) = 1
n ∑ x∈X L(x;W+,W−) = 1 n ∑ x∈X [ y − 1√ h 1>φ(W>+ x) + 1√ h 1>φ(W>−x) ]2 .
(98)

heading: C.8.1 DEFINING GRADIENT FLOWS
text: 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.
∂WO+ ∂t = 1
2n0 2n0∑ i=1 [ 1√ h ( yi − 1√ h 1>φ(WO>+ xi) + 1√ h 1>φ(WO>− xi) ) xiφ ′(x>i W O + ) ] (99)
= 1
2n0
[ X(y − yO(t)) 1√
h 1> ◦ φ′(XWO+ )
] , (100)
starting from the vanishing initialization wOi (0) ∼ N (0, I/dh1+ ). Note that the gradient for the negative part WO− can be similarly defined.
We now define the flow under the same objective but from exact zero initialization wDi (0) = 0 termed GF-Double. Due to zero initialization, a basic observation is that the solution [WD+ ,WD− ] is
at most rank-2, and more precisely, the parameters in the flow takes the form of WD± (t) = w D ±(t)1
>
where wD±(t) admits the following dynamics:
∂wD+ ∂t = gD+(w D +) = 1
2n0 2n0∑ i=1 [ 1√ h ( yi − √ hφ(wD>+ xi) + √ hφ(wD>− xi) ) φ′(wD>+ xi)xi ] .
(101)
Lastly, we define the GF-Single with solution denoted as w± = wS±(t) :
∂wS+ ∂t = gS+(w S +) = 1
2n0 2n0∑ i=1 [ 1√ h ( yi − √ hφ′(0)wS>+ xi + √ hφ′(0)w1>− xi ) φ′(0)xi ] . (102)
from zero initialization wD±(0) = 0. This can be seen as replacing the nonlinearity φ with its first-order Taylor expansion at the origin.

heading: C.8.2 FROM GF-DOUBLE TO GF-SINGLE
text: 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:
wS+(t) = −wS−(t) = 1
2 √ h
( I − e− φ′(0)2 n0 XX>t ) (XX>)−1Xy, (103)
when d < n, or otherwise
wS+(t) = −wS−(t) = 1
2 √ h X
( I − e− φ′(0)2 n0 X>Xt ) (X>X)−1y. (104)
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
f(x;w±) = f(x;w+1 >,w−1 >) = √ hφ(w>+x)− √ hφ(w>−x). (105)
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∥∥wD+(T )−wS+(T )∥∥2 = ∥∥∥∥∥ ∫ T
0
gD+(w D +(t))− gS+(wS+(t))dt ∥∥∥∥∥ 2
≤ ∥∥∥∥∥ ∫ T
0
gS+(w S +(t))− gS+(wD+(t)) dt ∥∥∥∥∥ 2 + ∥∥∥∥∥ ∫ T 0 gD+(w D +(t))− gS+(wD+(t)) dt ∥∥∥∥∥ 2
= ∥∥∥∥∥ ∫ T
0
1
2n0 2n0∑ i=1 [( − φ′(0)(wD+(t)−wS+(t))>xi + φ′(0)(wD−(t)−wS−(t))>xi ) φ′(0)xi ] dt ∥∥∥∥∥ 2 + E+
=O(1) ∫ T 0 (∥∥wD+(t)−wS+(t)∥∥2 + ∥∥wD−(t)−wS−(t)∥∥2) dt+ E+, (106) where we have defined the error term as
E+ = ∥∥∥∥∥ ∫ T
0
gD+(w D +(s))− gS+(wD+(s)) dt ∥∥∥∥∥ 2 . (107)
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
E+ = ∥∥∥∥∥ ∫ T
0
gD+(w D +(t))− gS+(wD+(t)) dt ∥∥∥∥∥ 2
≤ ∥∥∥∥∥ ∫ T
0
1
n n∑ i=1 [ 1√ h ( yi− √ hφ(wD>+ xi)+ √ hφ(wD>− xi) )( φ′(wD>+ xi)−φ′(0) ) xi ] dt ∥∥∥∥∥ 2
+ ∥∥∥∥∥ ∫ T
0
1
n n∑ i=1 [( −φ(wD>+ xi)+φ′(0)wD>+ xi+φ(0)+φ(wD>− xi)−φ′(0)wD>− −φ(0) ) φ′(0)xi ] dt ∥∥∥∥∥ 2
(i) ≤ ∥∥∥∥∥ ∫ T
0
1
n √ h X(y−yD(t)) ◦ (φ′(X>wD+)−φ′(0)1) dt ∥∥∥∥∥ 2
+ ∥∥∥∥∥ ∫ T
0
1
n n∑ i=1 [( O(wD>+ xi) 2−O(wD>− xi)2 ) xi ] dt ∥∥∥∥∥ 2
(i) ≤ 1 n √ h ∫ T 0 ‖X‖2 ∥∥y − yD(t)∥∥∞ ∥∥O(φ′(X>wD+))∥∥2 dt+O(n−1)∫ T 0 ‖X‖2 √ n ∥∥X>wD±∥∥2∞ dt
(ii)
≤ 1 n √ h ∫ T 0 O( √ d)O(logc d)O(1) dt+ 1√ n ∫ T 0 O( √ d)O ( logc d d ) dt = O ( logc d d ) T.
(108)
where (i) follows from the error of Taylor expansion on φ and (ii) from Condition A. Therefore, by Equation (106) and Gronwall’s inequality,∥∥wD+(T )−wS+(T )∥∥2 ≤C1 · logc hh eC2T = O ( poly log h h ) → 0. (109)
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 wD up to T. Importantly, note that wS(0) = wD(0) = 0, and that Condition A holds for wS 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 wD 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 RS and RD of two models with parameters wS± and w D ± .
|2(RS −RD)| = ∣∣∣Ex (x>β − f(x;wS±))2 − Ex (x>β − f(x;wD±))2∣∣∣
(i) ≤ √ Ex[f(x;wS±)− f(x;wD±)]2Ex[β >x− f(x;wS±) + β >x− f(x;wD±)]2
≤ √ hEx[|φ(wS>+ x)−φ(wD>+ x)|+|φ(wS>− x)−φ(wD>− x)|]2
· √ Ex[|β>x−f(x;wS±)|+|β >x−f(x;wD±)|]2
(ii) ≤ 2 √ hEx[|φ(wS>+ x)−φ(wD>+ x)|2+|φ(wS>− x)−φ(wD>− x)|2]
· √ Ex[|β>x−f(x;wS±)|2+|β
>x−f(x;wD±)|2] (iii) ≤ 2 √ hLip(φ)Ex[|wS>+ x−wD>+ x|2 + |wS>− x−wD>− x|2] · √ RS +RD
≤O( √ h) √∥∥wS+ −wD+∥∥22 + ∥∥wS− −wD−∥∥22 ≤ O(poly log h√h ) , (110)
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 RS and RD are finite for γ1 6= 1 due to the justified Condition A above. Therefore the difference between RS and RD 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.e. t =∞,
|RS(t)−RS(∞)| ≤ C √ h · ∥∥wS+(t)−wS+(∞)∥∥ = C√h∥∥∥∥ 12√he−φ ′(0)2 n0 XX>t(XX>)−1Xy ∥∥∥∥ 2
=C ∥∥∥∥e−φ′(0)2n0 XX>t(XX>)−1XX>β∥∥∥∥ 2 ≤ C exp ( −φ′(0)2 ∥∥∥∥ 1nXX> ∥∥∥∥ 2 t ) ‖β‖2 = C3e −C4t.
(111)
for constants C3, C4 > 0. Combining (110) and (111) yields for T = log log h,
|RD(T )−RS(∞)| ≤ |RS(T )−RD(T )|+ |RS(T )−RS(∞)|
= O ( poly log h√
h
) +O ( 1
poly log h
) → 0. (112)
The result in (112) for case d > n follows a similar proof and is omitted.

heading: C.8.3 FROM GF-ORIGINAL TO GF-DOUBLE
text: In this section we compare GF-original with GF-double. Note that the two flows differ only at initialization: vanishing initialization wOi (0) ∼ N (0, I/dh1+ ) v.s. zero initialization wDi (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∥∥WO(T )−WD(T )∥∥2
F = ∥∥WO(0)−WD(0)∥∥2
F + ∫ T 0 ∂ ∥∥WO(t)−WD(t)∥∥2 F ∂t dt
= ∥∥WO(0)∥∥2
F + ∫ T 0 ∥∥WO(t)−WD(t))∥∥ F ∥∥∥∥∂L(X;WO(t))∂W − ∂L(X;WD(t))∂W ∥∥∥∥ F dt
≤ ∥∥WO(0)∥∥2
F + ∫ T 0 ∥∥WO(t)−WD(t)∥∥2 F + ∥∥∥∥∂L(X;WO(t))∂W − ∂L(X;WD(t))∂W ∥∥∥∥2 F dt
≤ ∥∥WO(0)∥∥2
F + (1 + Lf ) ∫ T 0 ∥∥WO(t)−WD(t)∥∥2 F dt, (113)
And hence by Gronwall’s lemma one obtains:∥∥WO(T )−WD(T )∥∥ F ≤ ∥∥WO(0)∥∥ F e(1+Lf )T/2 = O(d−(1+ )/2)eCT . (114) Therefore the difference in the function output can be bounded as∥∥f(x,WD(T ))− f(x,WO(T ))∥∥ 2 = ∥∥φ(x>WD(T ))a− φ(x>WO(T ))a∥∥ 2
≤ ∥∥φ(x>WD(T ))− φ(x>WO(T ))∥∥
2 ‖a‖2 ≤ Lφ ‖x‖2 ∥∥WD(T )−WO(T )∥∥ F ‖a‖2
=O( √ d) · ∥∥WD(T )−WO(T )∥∥ F = O(d− /2)eCT . (115)
Taking T = log log h, together with the same argument in Step 1 yields |RO(T )−RD(T )| = O ( poly log h
d /2
) → 0. (116)

heading: C.8.4 PUTTING THINGS TOGETHER
text: By (112) and (116), we know that for T = log log h,
|RO(T )−RS(∞)| ≤ |RO(T )−RD(T )|+ |RD(T )−RS(∞)| → 0. (117)
Finally the proof is completed by observing that RS(∞) is the risk of the minimum-norm solution on the input discussed in Section 3. In addition, note that at T = log log h, from (109)(114) one obtains that ∥∥WO(t)−WS(t)∥∥ F → 0, and therefore by Lemma 18 we have ∥∥∂L(X;WO(t))/∂W∥∥ F → 0, i.e. the flow on the original objective reaches a o(1) first-order stationary point.

heading: C.9 PROOF OF THEOREM 8
text: Denote ω = vec(W ) = vec([W+,W−]), and ω0 = vec(W init). Define
K(t) = ∂f(X;ω(t))
∂ω(t)
> ∂f(X;ω(t))
∂ω(t) , (118)
which is the kernel matrix of the neural tangent kernel Jacot et al. (2018); Du et al. (2018). In the following sections we show that under the non-vanishing initialization, the trained two-layer network is well-approximated by the regression model on the NTK, for which we derive the population risk.

heading: C.9.1 THE KERNEL LINEARIZATION
text: Write yNN (t) = fNN (X, t) ∈ Rn and its evolution:
dyNN (t) = 1
n K(t)(y − yNN (t)) dt, (119)
and the corresponding linearized flow:
dyNTK(t) = 1
n K(0)(y − yNTK(t)) dt, (120)
Previous works (e.g. Du et al. (2018); Oymak and Soltanolkotabi (2019)) have proved (nonasymptotically) that the two trajectories (119) and (120) are close if the model is overparameterized, i.e. h = poly(n), under no assumptions on the teacher model. In our asymptotic setup (together with assumptions (A1-3)), we argue that similar conclusion holds without significant overparameterization.
We first show the global convergence of the training of two-layer neural network fNN . We employ an argument similar to (Du et al., 2018, Theo. 3.2) by first identifying the condition under which training converges at linear rate:
Condition B: λmin(K(t)) = O(d).
From Corollary 15 we know that at initialization the lowest eigenvalue of the NTK matrix satisfies λmin(K(0)) = O(d), and thus the kernel regression on the NTK enjoys linear convergence. By Lemma 19, we know that for ‖W (t)−W (0)‖2 = O(d1− ), the order of λmin(K(t)) = O(d) remains unchanged. Therefore, if we assume that up to time T the weights satisfy ‖wi(t)−wi(0)‖2 = O(d−1/2), then Condition B is satisfied for yNN and from (Chizat and Bach, 2018b, Lemma B1) we have the following linear convergence
‖yNN (t)− y‖2 ≤ C1 ‖yNN (0)− y‖2 e −C2t. (121)
Since ‖yNN (0)− y‖2 = O( √ n) at initialization, setting T = O(log d) ensures that the training loss 1n ‖yNN (T )− y‖ 2 2 → 0 as n→∞. Consequently it is easy to check that ∥∥∥∂L(X;W (T ))∂W (T ) ∥∥∥ 2 → 0 and thus at time T the gradient flow reaches an o(1) first order stationary point.
We now verify that each weight vector wi travels at most O(d−1/2) from initialization. The norm of gradient for wi can be bounded as:∥∥∥∥∂L(X;wi(t))∂wi(t) ∥∥∥∥ 2 = ∥∥∥∥ 1nX [ 1√ h (y − yNN (t)) ◦ φ′(X>wi(t)) ]∥∥∥∥ 2
(i) ≤O(1) 1 d1.5 ‖X‖2 ‖y − yNN (t)‖2 ≤ O(1)d −1 ‖y − yNN (0)‖2 e −t, (122)
where (i) follows from the boundedness of φ′. Integrating the gradient yields ‖wi(t)−wi(0)‖2 = O(d−1/2), i.e. ‖W (t)−W (0)‖2 = O(1). Thus the distance traveled by W indeed satisfies the norm
assumption above, and following the same argument as (Du et al., 2018, Lemma 3.4) we conclude that Condition B holds true for the gradient flow of fNN for t > 0. Next we show that for any input x̂ with ‖x̂‖2 = O( √ d), the prediction of the neural network is uniformly close to the prediction of the prediction of the kernel model, a result similar to (Arora et al., 2019a, Lemma F.1). Following the notation of Arora et al. (2019a), we write the time derivative of the prediction as d
dt fNN (x̂, t) =
1 n uNN (x̂, t) >(y − yNN (t)); d dt fNTK(x̂, t) = 1 n uNTK(x̂, t) >(y − yNTK(t)),
where uNN (x, t) = ∂f(X;ω(t)) ∂ω(t) > ∂f(x;ω(t)) ∂ω(t) ∈ R
n and uNTK(x, t) similarly defined on the initialized weights ω(0). We bound the difference between the predictions on x̂ up to terminal time T as
|fNN (x̂, t)− fNTK(x̂, t)| = | ∫ T
0
[ d
dt fNN (x̂, t)−
d dt fNTK(x̂, t)
] dt|
= 1 n | ∫ T
0
[ uNN (x̂, t) >(y − yNN (t))− uNTK(x̂, t)>(y − yNTK(t)) ] dt|
≤ 1 n | ∫ T
0
uNTK(x̂, t) >(yNN (t)− yNTK(t)) dt|
+ 1 n | ∫ T
0
(uNN (x̂, t)− uNTK(x̂, t))>(y − yNN (t)) dt|
≤ 1 n ‖uNTK(x̂, t)‖2 ∫ T 0 ‖yNN (t)−yNTK(t)‖2 dt
+ 1
n max t<T ‖uNN (x̂, t)− uNTK(x̂, t)‖2 ∫ T 0 ‖y−yNN (t)‖2 dt. (123)
For the first term have
‖yNN (T )− yNTK(T )‖2 ≤ 1
n ∫ T 0 ‖K(t)(y − yNN (t))−K(0)(y − yNTK(t))‖2 dt
≤ 1 n max 0<t<T
‖K(t)−K(0)‖2 ∫ T
0
‖y − yNN (t)‖2 dt+ 1
n ‖K(0)‖2 ∫ T 0 ‖yNN (t)− yNTK(t)‖2 dt
(i) ≤ 1 n O(d1/2− ′ )O( √ d) +O(1) ∫ T 0 ‖yNN (t)− yNTK(t)‖2 dt (ii) ≤ O(d− ′ ), (124) where (i) is due to Corollary 15, Lemma 19 for some ′ > 0 and the linear convergence of yNN , and (ii) is due to Gronwall’s inequality. Note that the log factor in T is omitted. Similarly, for the second term we have 1
n max t<T ‖uNN (x̂, t)− uNTK(x̂, t)‖2 ∫ T 0 ‖y−yNN (t)‖2 dt (i) ≤ 1 n O(d1/2− ′ )O( √ d) = O(d− ′ ),
(125) where we used Lemma 19 and the linear convergence of yNN in (i). Combining the two cases yields
|fNN (x̂, t)− fNTK(x̂, t)| ≤ 1
n ‖uNTK(x̂, t)‖2O(d
− ′) +O(d− ′ ) (i) = O(d− ′ ), (126)
where we utilized Corollary 14 in (i). Thus we know that the difference between the population risk of fNN and fNTK is also asymptotically vanishing (note that the derivation above is independent of the target function as long as ‖y‖2 = O( √ n)). Therefore, in the following subsection we compute the risk of the linearized (kernel) model fNTK .

heading: C.9.2 COMPUTING THE KERNEL RISK
text: Given input X ∈ Rd×n and label y = β>X + ε, gradient flow on the tangent kernel solves the following equation of the parameters ω:
y = f(X;ω) = ∂f(X;ω0)
∂ω
> (ω − ω0), (127)
where ∂f(X;ω0)/∂ω is a dh× n matrix with column ∂f(xi;ω0)/∂ω. Note that for n→∞ and γ1, γ2 ∈ (0,∞), dh > n trivially holds, and by Corollary 15 we know that solution is given by
ω1 = ω0 + ∂f(X;ω0)
∂ω
( ∂f(X;ω0)
∂ω
> ∂f(X;ω0)
∂ω
)−1 (X>β + ε). (128)
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):
2R = Ex,ε [( x>β − f(x;ω1) )2] = Ex,ε (x>β − ∂f(x;ω0) ∂ω > (ω1 − ω0)
)2 = Ex,ε  x>β − ∂f(x;ω0)
∂ω
> ∂f(X;ω0)
∂ω
( ∂f(X;ω0)
∂ω
> ∂f(X;ω0)
∂ω
)−1 (X>β + ε) 2 
= Ex [( x>β − û>K̂−1X X >β )2]
︸ ︷︷ ︸ 2B
+Ex [ ûK̂−1X K̂ −1 X û > ] σ2︸ ︷︷ ︸
2V
, (129)
where a bias-variance decomposition is made here, and for simplicity we define
û = ∂f(X;ω0)
∂ω
> ∂f(x;ω0)
∂ω , K̂X =
∂f(X;ω0)
∂ω
> ∂f(X;ω0)
∂ω . (130)

heading: C.9.3 APPROXIMATING THE KERNEL MATRIX
text: In this section we drop the negligible in the initialization. Following Cheng and Singer (2013) we utilize the orthonormal decomposition of φ′(x) in L2(R, µG). Denote b0 = E[φ′(G)], and b21 = E[φ′(G)2]− b20. We have the orthogonal decomposition of φ′
φ′(x) = b0 + φ ′ ⊥(x), (131)
where E[φ′⊥(G)] = 0. We develop the following lemmas to approximate the kernel matrix.
Lemma 13 (Approximation of (K̂X)ij). There exist constants c, c′ > 0 such that for i 6= j with probability 1− e−cnε2 we have∣∣∣∣∣1d ∂f(xi;ω0)∂ω > ∂f(xj ;ω0)∂ω − 1db20x>i xj
∣∣∣∣∣ < ε2, (132) and with probability 1− e−c′nε2 ,∣∣∣∣∣1d ∂f(xi;ω0)∂ω > ∂f(xi;ω0)∂ω − (b20 + b21)
∣∣∣∣∣ < ε. (133) Proof. When i 6= j (i.e. Equation (132)), we have
1 d [K̂X ]ij = 1 d ∂f(xi;ω0) ∂ω
> ∂f(xj ;ω0)
∂ω
= 1
d h∑ k=1 ∂f(xi;ω0) ∂wk > ∂f(xj ;ω0) ∂wk = 1 dh h∑ k=1 x>i xjφ ′(w>k xi)φ ′(w>k xj)
→1 d x>i xjEw
[ φ′(w>xi)φ ′(w>xj) ] = 1
d H(xi,xj). (134)
The matrixH(xi,xj) = x>i xjEw [ φ′(w>xi)φ ′(w>xj) ]
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
Pr ∣∣∣∣∣1d ∂f(xi;ω0)∂ω > ∂f(xj ;ω0)∂ω − 1dH(xi,xj) ∣∣∣∣∣ < 1dx>i xjε > 1− e−c1hε2 . (135)
In addition, by the concentration of x>i xj and ‖xi‖22, i.e. Prx>i xj/d > ε < 1 − e−c2dε 2
and Pr |x>i xi/d− 1| < ε > 1− e−c3dε 2
, the orthonormal decomposition φ′(x) = b0x+ φ′⊥(x) leads to the following linear approximation of the matrix H
1 d H(xi,xj) = b 2 0 1 d x>i xj +O
( (x>i xj/d) 2 ) . (136)
and by taking ε = x>i xj/d under the joint event we can show that∣∣∣∣∣1d ∂f(xi;ω0)∂ω > ∂f(xj ;ω0)∂ω − 1db20x>i xj ∣∣∣∣∣ < ε2 (137)
with probability 1− e−cdε2 . The same argument follows for the case where i = j.
Corollary 14 (Approximation of û). For large enough l > 0, with probability 1− de−c logl d
1 d ‖û− ũ‖2 = ∥∥∥∥∥1d ∂f(x;ω0)∂ω > ∂f(X;ω0)∂ω − 1d ũ ∥∥∥∥∥
2
< logl d
d , (138)
where ũ = b20x >X .
Proof. Taking ε = logl d/d together with Lemma 13 yields the desired result.
Corollary 15 (Approximation of K̂X ). With probability 1− de−c log l d
1
d ∥∥∥K̂X − K̃X∥∥∥ F = ∥∥∥∥∥1d ∂f(X;ω0)∂ω > ∂f(X;ω0)∂ω − 1dK̃X ∥∥∥∥∥ F < logl d, (139)
where K̃X = b20X >X + b21dI .
Proof. Also by directly applying Lemma 13.
Remark. For initialization larger or equal towi(0) ∼ N(0, Id/d), the above approximation does not depend on the scale of initialization. Remark. For φ(x) = SoftPlus(x), b20 = 1/4, b21 = 0.043379. For φ(x) = sigmoid(x) = (1 + e−x)−1, b20 = 0.042692, b 2 1 = 0.002144. Note that b1 ≥ 0 for all smooth activations φ, and the equality holds (i.e. b1 = 0) if and only if φ is linear. We comment that smaller b1 entails larger variance as γ1 → 1, and vice versa, as shown in the following section.

heading: C.9.4 THE BIAS TERM
text: With these approximation above we proceed to calculating (129) 2B = Ex [( x>β − û>K̂−1X X >β )2] . (140)
We first bound the error in substituting û with ũ:∥∥∥û>K̂−1X X>β − ũ>K̂−1X X>β∥∥∥ 2 ≤ ‖û− ũ‖2 ∥∥∥K̂−1X ∥∥∥ 2 ‖X‖2 ‖β‖2
(i) = O ( logl d · d−1 · √ d · 1 ) = O ( logl d√ d ) , (141)
where (i) is due to the fact that ∥∥∥K̂−1X ∥∥∥
2 = λ−1min(K̂X) = O(1/d) and ‖X‖2 = O(
√ d). Therefore
we have as n, d, h→∞ 2B = Ex [( x>β − û>K̂−1X X >β )2] → Ex [( x>β − ũ>K̂−1X X >β )2]
= Ex [( x>β − b20x>XK̂−1X X >β )2] . (142)
By taking expectation over x and the rotational invariance argument similar to Hastie et al. (2019), Ex [( x>β − b20x>XK̂−1X X >β )2] =Ex [ β> ( I − b20XK̂−1X X > )2 β ] = β>β d tr (( I − b20XK̂−1X X > )( I − b20XK̂−1X X > )) .
(143)
In addition, we bound the error in substituting K̂X by K̃X defined in (139):∣∣∣∣1d tr(XK̂−1X X> −XK̃−1X X>) ∣∣∣∣ = ∣∣∣∣1d tr(X>XK̂−1X (K̂X − K̃X)K̃−1X ) ∣∣∣∣ < 1
d
∥∥X>X∥∥ 2 ∥∥∥K̂−1X ∥∥∥ 2 ∥∥∥K̂X − K̃X∥∥∥ F ∥∥∥K̃−1X ∥∥∥ 2
= O ( d−1 · d · d−1 · d logl d · d−1 ) = O
( logl d
d
) , (144)
and similarly,∣∣∣∣1d tr(XK̂−1X X>XK̂−1X X> −XK̃−1X X>XK̃−1X X>) ∣∣∣∣
= ∣∣∣∣1d tr(X>X(K̂−1X − K̃−1X )X>X(K̂−1X + K̃−1X )) ∣∣∣∣
< 1
d
∥∥X>X∥∥ 2 ∥∥∥K̂−1X ∥∥∥ 2 ∥∥∥K̂X − K̃X∥∥∥ F ∥∥∥K̃−1X ∥∥∥ 2 ∥∥X>X∥∥ 2 ∥∥∥K̂−1X + K̃−1X ∥∥∥ 2
=O ( d−1 · d · d−1 · d logl d · d−1 · d · d−1 ) = O
( logl d
d
) . (145)
Combining these two formulas in (143) yields 2B → Ex [( x>β − b20x>XK̂−1X X >β )2]
→ β >β d tr (( I − b20XK̃−1X X > )( I − b20XK̃−1X X > )) = β>β d tr (( I − b20X(b20X>X + b21dI)−1X> )2) . (146)
Utilizing the Marčenko–Pastur law from Section B.2 we obtain
B = β>β
( γ1 − 1
2γ1 + γ1(γ1 + γ1m+m− 2) + 1 2γ1 √ γ1(γ1 +m(γ1(m+ 2) + 2)− 2) + 1
) , (147)
where m = b−20 b 2 1.

heading: C.9.5 THE VARIANCE TERM
text: Similarly, for the variance we utilize the approximation 2V = Ex [ ûK̂−1X K̂ −1 X û > ] σ2 (148)
Specifically, we bound the approximation error∣∣∣ûK̂−1X K̂−1X û> − ũK̂−1X K̂−1X ũ>∣∣∣ ≤‖û− ũ‖2 ∥∥∥K̂−1X ∥∥∥ 2 ‖û+ ũ‖2 ∥∥∥K̂−1X ∥∥∥ 2
=O ( logl d · 1
d · d · 1 d
) = O ( logl d
d
) , (149)
and similarly ∣∣∣Ex [ũK̂−1X K̂−1X ũ> − ũK̃−1X K̃−1X ũ>]∣∣∣ =tr (( K̂−1X − K̃ −1 X )( K̂−1X + K̃ −1 X ) Exũũ>
) =tr ( K̂−1X ( K̂X − K̃X ) K̃−1X ( K̂−1X + K̃ −1 X ) XTX
) ≤ ∥∥∥K̂−1X ∥∥∥
2 ∥∥∥K̂X − K̃X∥∥∥ F ∥∥∥K̃−1X ∥∥∥ 2 ∥∥∥K̂−1X + K̃−1X ∥∥∥ 2 ∥∥XTX∥∥ 2
=O(d−1 · d logl d · d−1 · d−1 · d) = O
( logl d
d
) , (150)
By combining the two approximations above we know that as n, d, p→∞∣∣∣2V − Ex [ũK̃−1X K̃−1X ũ>]σ2∣∣∣ = O ( logl d
d
) → 0. (151)
Therefore the variance is given as 2V → σ2Ex [ ũK̃−1X K̃ −1 X ũ > ]
= σ2Ex [ b40x TX ( b20X >X + b21dI )−2 XTx ]
= σ2 1
d tr
( 1
d XTX ·
( 1
d XTX + b−20 b 2 1I
)−2)
= σ2 ( −1
2 +
γ1 + γ1m+ 1 2 √ γ1(γ1 +m(γ1(m+ 2) + 2)− 2) + 1
) , (152)
where m is defined in the derivation of the bias term.

heading: C.9.6 PUTTING THINGS TOGETHER
text: Recall the population risk is the sum of the bias and variance R→ r2 ( γ1 − 1
2γ1 + γ1(γ1 + γ1m+m− 2) + 1 2γ1 √ γ1(γ1 +m(γ1(m+ 2) + 2)− 2) + 1
)
+ σ2 ( −1
4 +
γ1 + γ1m+ 1 4 √ γ1(γ1 +m(γ1(m+ 2) + 2)− 2) + 1
) . (153)
Observe that the population risk is independent of γ2, i.e. double descent does not occur when the network is overparameterized via changing the width. In addition, the bias is monotonically increasing and upper-bounded by the null risk r2 and lower-bounded by the bias of the least squares solution on the input features β̂ = X†y, whereas the variance remains bounded for all γ1 ∈ (0,∞) as long as m > 0, i.e. φ is nonlinear.

heading: D USEFUL LEMMAS
text: Lemma 16. Given KW from (46), define
K̃W = rIh + s1h1 > h + tQ, (154)
where Q ∈ Rh×h with Qi6=j = w>i wj and Qi,i = 0, and r = E[φ(G)2] − E[φ(G)]2, s = E[φ(G)]2, t = E[Gφ(G)]2. Then as n, d, h→∞, ∥∥∥KW − K̃W∥∥∥ F ≤ logc d a.s..
Proof. Consider the event where A = { | ‖wi‖2 − 1| < , |w>i wj | < } . Under event A , for the diagonal term of the kernel matrix we have∣∣∣[KW ]ii − [K̃W ]ii∣∣∣ = ∣∣∣Ex[φ(w>i x)φ(w>i x)]− E[φ(G)2]∣∣∣ = ∣∣E[φ(‖wi‖2G)2]− E[φ(G)2]∣∣ = O( ). (155)
And for off-diagonal term, by the decomposition introduced in Section B.3 we have [KW ]ij = Ex [ φ(w>i x)φ(w > j x) ] = ‖wi‖2 ‖wj‖2 + Ex[φ⊥(w > i x)φ⊥(w > j x)], (156) hence |[KW ]ij − [K̃W ]ij | < (w>i wj)2. Notice that A holds a.s. for = log c d/ √ d and large
enough c > 0; we therefore have ∥∥∥KW − K̃W∥∥∥
F ≤ logc d.
Lemma 17. Let β̂(t) be the solution to the gradient flow at time t defined in (103)(104). Then as n, d→∞ and γ1 6= 1 the following holds.:
‖β̂(t)‖2 = OP (1); ‖X>β̂(t)‖∞ = OP (poly log d).
Proof. We consider d < n for simplicity, and result for the other case follows in similar fashion. In this case β̂(t) = ( I − exp(− tnXX >) )
(XX>)−1Xy. From (Hastie et al., 2019, Corollary 1) we know that ‖β̂(∞)‖2 = ∥∥(XX>)−1Xy∥∥ 2 = O(1) for γ1 < 1. Note that ∥∥I − exp(− tnXX>)∥∥2 = O(1) for t ≥ 0; it follows that ‖β̂(t)‖2 = O(1).
For the second part, we utilize the SVD X = UΣV >, where Σ = [ Σ̂; 0 ] , Σ̂ ∈ Rd×d and Σ̂ii = λi.
We have ( I − exp(− tnXX >) ) = U Σ̄U> where Σ̄i,i = 1− exp(−tλ2i /n) if i ≤ d and 0 otherwise.
‖X>β̂(t)‖∞ = ∥∥∥∥X>(I − exp(− tnXX>) ) (XX>)−1Xy ∥∥∥∥ ∞
≤ ∥∥∥V Σ>U>U Σ̄U>U Σ̂−2U>UΣV >∥∥∥
∞ ‖y‖∞ ≤ ∥∥∥V Σ>Σ̄Σ̂−2ΣV >∥∥∥
∞
( ‖X>β‖∞ + ‖ε‖∞ ) ≤‖V ‖∞ ∥∥∥Σ>Σ̄Σ̂−2Σ∥∥∥ ∞ ‖V >‖∞ ( ‖X>β‖∞ + ‖ε‖∞ ) (i) ≤ OP (poly log d), (157)
where (i) follows from the concentration of the Gaussian maxima, and the fact that the law of V is the Haar measure on SO(n), and thus for any unit vector z independent to V , V z is uniform on sphere and ‖V z‖∞ = O(log d/ √ d). We wherefore have
‖V ‖∞ = sup z ‖V z‖∞ ‖z‖∞ = O ( log d√ d ) ‖z‖2 ‖z‖∞ = O(log d), (158)
Note that this result also implies that ‖y −X>β(t)‖∞ = OP (poly log d).
Lemma 18. For weight matrices W ,W ′ satisfying ‖y − f(X)‖2 = O( √ n), where f(X) =
φ(XW )a with fixed ai ∼ 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 )∂W − ∂L(X;W ′)∂W ∥∥∥∥ F ≤ L ‖W −W ′‖F . (159)
Proof. Denote y1 = φ(X>W1)a and y2 = φ(X>W2)a for W1,W2 satisfying the assumption above (which can be seen as a condition on the magnitude of training loss), we have∥∥∥∥∂L(W1)∂W1 − ∂L(W2)∂W2 ∥∥∥∥ F
= ∥∥∥∥ 1nX [(y − y1)a> ◦ φ′(X>W1)]− 1nX [(y − y2)a> ◦ φ′(X>W2)] ∥∥∥∥ F ≤ 1 n ‖X‖2
∥∥(y − y1)a> ◦ φ′(X>W1)− (y − y2)a> ◦ φ′(X>W2)∥∥F ≤O
( 1√ d )∥∥(y2 − y1)a> ◦ φ′(X>W1)∥∥F +O
( 1√ d )∥∥(y − y2)a> ◦ (φ′(X>W1)− φ′(X>W2))∥∥F . (160) We upper bound the two terms separately:∥∥(y2 − y1)a> ◦ φ′(X>W1)∥∥2 (i)≤max{φ′(X>W1)ij} ‖y2 − y1‖2 ‖a‖2
(ii) ≤O(1) ∥∥φ′(X>W1)a− φ′(X>W2)a∥∥F
(iii) ≤ O(1) ‖X‖2 ‖W1 −W2‖F = O( √ d) ‖W1 −W2‖F , (161)
where we applied the inequality ‖A ◦B‖F ≤ max{|Aij |} ‖B‖F in (i), boundedness of φ′ in (ii) and Lipschitzity of φ in (iii). Similarly, for the second term∥∥(y − y2)a> ◦ (φ′(X>W1)− φ′(X>W2))∥∥F
≤max{|ai|} ‖y − y2‖2 ∥∥φ′(X>W1)− φ′(X>W2)∥∥F
(i) ≤O(1) ‖X‖2 ‖W1 −W2‖F = O( √ d) ‖W1 −W2‖F , (162)
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 ‖wi(t) − wi(0)‖2 = O(d−1/2) for all i, then we have ‖K(t)−K(0)‖2 = O(d1/2− ′ ) for some positive
′ ∈ Θ(1).
Proof. Recall the definition of the NTK:
Kij(t) = ∂f(xi;ω(t))
∂ω(t)
> ∂f(xj ;ω(t))
∂ω(t) = x>i xj
1
h h∑ k=1 φ′(wk(t) >xi)φ ′(wk(t) >xj), (163)
or equivalently the matrix form
K(t) = X>X ◦ 1 h [φ′(X>W (t))φ′(W (t)>X)]. (164)
At initialization, xi ∼ N(0, Id) and wk(0) ∼ N(0, d Id). Thus for fixed xi, by Gaussian anti-concentration we have Pr |x>i wk| < log d ≤ O(1/d1/2+ 1) for some 1 > 0. In addition, note that ‖wk(t)−wk(0)‖2 = O(d−1/2) for all k, and therefore for i, j, k such that
|x>i wk(0)| > O(log d) and |x>j wk(0)| > O(log d), we know that |φ′(x>i wk(t))φ′(x>j wk(t)) − φ′(x>i wk(0))φ ′(x>j wk(0))| = O(d−2).
Given fixed xi, define yk = 1{|x>i wk| < log d} as the indicator variable that the k-th neuron does not saturate. We know that E[yk] = O(1/d1/2+ 1), and Var[yk] = E[y2k]− E[yk]2 = O(1/d1/2+ 1). By Bernstein’s inequality
Pr ∣∣∣∣∣ 1h h∑ k=1 yk − E[yk] ∣∣∣∣∣ > ε ≤ 2 exp ( − hε 2 2σ2 + 2ε/3 ) . (165)
Setting ε = √
c log h h1+ 2 , we know that with probability at least 1− h−c,
1
h h∑ k=1 yk ≤ ε+ E[yk] = O ( poly log h h1/2+ 3 ) . (166)
Therefore, given xi and xj , for large enough c1 with probability at least 1− h−3 we have∣∣∣∣∣ h∑ k=1 φ′(wk(t) >xi)φ ′(wk(t) >xj)− h∑ k=1 φ′(wk(0) >xi)φ ′(wk(0) >xj) ∣∣∣∣∣ = O(h1/2− 4), (167) in which we utilized the boundedness of φ′. Taking union bound over d2 elements in the random feature matrix yields
‖K(t)−K(0)‖2
= ∥∥∥∥X>X ◦ 1h [φ′(X>W (t))φ′(W (t)>X)− φ′(X>W (0))φ′(W (0)>X)] ∥∥∥∥
≤ 1 h ∥∥X>X∥∥ 2 max {∣∣φ′(X>W (t))φ′(W (t)>X)− φ′(X>W (0))φ′(W (0)>X)∣∣ ij } ≤ 1 h O(d)O(d1/2− ′ ) = O(d1/2− ′ ). (168)
Using the exact same argument, one can derive that ‖uNN (x̂)− uNTK(x̂)‖2 = O(d1/2− ′ ), the proof of which we omit.

heading: E ADDITIONAL RESULTS
text: E.1 RISK OF ReLU NETWORK UNDER SYMMETRIC DATA
If the dataset is symmetric, that is
(A5) Symmetric Data: ∀i ∈ [1, n],∃!j ∈ [1, n] s.t. xi + xj = 0,
then population risk of the gradient flow solution can be given explicitly for certain nonlinearities: Proposition 20. Given (A1-3)(A5), if the nonlinearity satisfies φ′(x) + φ′(−x) = C for constant C, then as n, d, h→∞
R(γ1<0.5)(f̂)→ 2γ1
1− 2γ1 σ2; R(γ1≥0.5)(f̂) =
( 1− 1
2γ1
) r2 +
1
2γ1 − 1 σ2. (169)
Note that the requirement on the nonlinearity holds for ReLU and SoftPlus. This expression is again independent to γ2 and aligns with the experimental results in Figure 9 (we only plot the bias component for verification). In addition, the bias is upper-bounded by the null risk for all γ1. We remark that the symmetry assumption does not hold for i.i.d. samples from symmetric distributions, and Figure 9 demonstrates that the additional condition alters the risk.
Proof. Without loss of generality assume X = [X0,−X0]. Then by (100) we have
∂w+ ∂t = 1
2n0 2n0∑ i=1 [( yi − h0φ(w>+xi) + h0φ(w>−xi) ) φ′(w>+xi)xi ] , (170)
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,
∂(w+) ∂t + ∂(w−) ∂t = 1
2n0 2n0∑ i=1 [( yi − h0φ(w>+xi) + h0φ(w>−xi) ) (φ′(w>+xi)− φ′(w>−xi))xi ] = 0.
(171)
And therefore the gradient flow of w+ is
∂w+ ∂t = 1
2n0 2n0∑ i=1 [( yi − h0φ(w>+xi) + h0φ(−w>+xi) ) φ′(w>+xi)xi ] = 1
2n0 n0∑ i=1 [( yi − hφ(w>+xi) + h0φ(−w>+xi) )( φ′(w>+xi) + φ ′(−w>+xi) ) xi ] = 1
2n0 n0∑ i=1 [( yi − h0w>+xi ) xi ] = 1 2n0 X0y0 − 1 2n0 h0X0X > 0 w+. (172)
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 h0 (XX>)−1Xy, γ1 < 0.5, 1
h0 X(X>X)−1y, γ1 > 0.5.
(173)
And hence the asymptotic risk is
R(γ1<0.5) → 2γ1
1− 2γ1 σ2; R(γ1≥0.5) =
( 1− 1
2γ1
) r2 +
1
2γ1 − 1 σ2. (174)
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.

heading: F EXPERIMENT SETUP
text: 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 β = −1d/ √ 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 )‖2F < 10−6. We fix n = 320 and vary γ1, γ2 from 0.1 to 3 with the same linear teacher model β = −1d/ √ d. The risk is averaged across 20 models trained from different initializations.

Review summaries:
This paper provides exact bounds on the risk when training a two-layer neural network in an asymptotic regime. Namely, the paper considers training under the square-loss objective, a two-layer neural network with $h$ hidden units on inputs of dimension $d$ and training on $n$ samples. The asymptotic regime is considered by making all of $d$, $h$, $n$ go to $\infty$, in a way that the ratio $d/n$ approaches $\gamma_1$ and the ratio $h/n$ approaches $\gamma_2$.