Generalization of Two-layer Neural Networks: An Asymptotic Viewpoint
Authors: Jimmy Ba
Murat Erdogdu
Taiji Suzuki
Denny Wu
Tianzong Zhang
Venue:  (2020)
Keywords: Neural Networks
Generalization
High-dimensional Statistics

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 \textit{double descent} phenomenon: a cusp in the population risk appears at $h\approx 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 \textit{independent} of overparameterization. Our theoretical and experimental results suggest that previously studied model setups that provably give rise to \textit{double descent} might not translate to optimizing two-layer neural networks.

Published as a conference paper at ICLR 2020
GENERALIZATION OF TWO-LAYER NEURAL NET- WORKS: AN ASYMPTOTIC VIEWPOINT
Jimmy Ba Murat A. Erdogdu 12 Taiji Suzuki Denny Wu Tianzong Zhang University of Toronto', Vector Institute". University of Tokyo', RIKEN AIP", Tsinghua University (iba, erdogdu, dennywu) @cs. toronto. edu, ta taijiemi jiiimiiss.......m w-t.kyo. ac. jp, zt mamammmmlllll.... wwwtttiingghu edu.cr cn
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 differrnt 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 z and further overparameterization decreases the risk. In contrast, when the first layer weights are optimized, we highlight how different scales of initialization lead 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 double descent might not translate to optimizing two-layer neural networks.
INTRODUCTION
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 al., 2018; Ji and Telgarsky, 2018), and provided generalization guarantees (Arora al., 2018; Bartlett et al., 2017; Dziugaite and Roy, 2017) that align with the empirical observations.
Recently, series of works highlighted the implicit regularization of interpolators in the overpa- rameterized regime (Belkin et al., 2018; Spigler et al., 2018; Geiger et al., 2018; Advani and Saxe, 2017). Specifically, 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 double descent, and 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 gap between the practical benefit of overparameterizzation and the recently proved double descent phenomenon, which is typically established under models that exhibits the following structure: the trained model solves linear inverse problem, and the "cusp" in the risk arises from the instability of the inverse at the interpolation threshold. Moreover, given dataset or fixed d, the number of parameters in the linear regression model is also fixed, i.c. the level of overparameterization cannot be altered. is therefore unclear the trend persists in the optimization of more complex models, for instance in two-layer neural networks where overparameterizztton can be controlled simply by adding more neurons.
In this work, analyze the generalization properties of two-layer neural networks in the unreg- ularized least squares regression setting and examine the presence/absence the double descent phenomenon. We consider the proportional asymptotic limit where the number of samples input

---

Published as a conference paper at ICLR 2020
features d, and neurons h tend to infinity at the same rate, under which overparameterrizaton cor- responds to increasing the limit of h/n (network "width"). This regime is particularly interesting because even though n + oc, 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 I):
3.0
2.5
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.
2.0
1.5
When the first layer is optimized, we compare two so- lutions of gradient flow from different scales of initial- ization, which we term as vanishing and non-vanishing initialization, and show in both cases the population risk is independent to overparameterization.
empirical
prediction
0.5
0.5
1.0
1.5 h/io
2.0
2.5
Figure 1: Illustration of the double descent risk curve in two-layer linear networks (SNR= 16). Brighter color indicates larger 71 = 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).
For the vanishing initialization, we show that the risk of the gradient flow solution is asymptotically close to that of rank-1 model. For non-vanishing initializa- tion, we show that the gradient flow solution is well- approximated 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-Eijndes (2018); Sirignano and Spiliopoulos (2018), 1): Nitanda and Suzuki (2017) studied the mean-field limit where the number of neurons h + Do 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-Zhur et al. (b); Song and Yang (2019) considered a different scaling and showed that gradient descent on overparameterized models converges to global minimizer a linear rate; key 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 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 fully- trained 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; 8; 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., 201 8a; 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 mudent-teechhe setup.
2

---

Published as a conference paper at ICLR 2020
Double Descent. The term double descent refers the phenomenon that the population risk of an empirical risk minimizer manifests a "cusp" at the interpolation threshold, and further overparame- terization decreases the risk. First observed in Krogh and Hertz (1992), the phenomenon has been recently connected to the benefit of overparameterivation (Belkin et al., 2018; Geiger et al., 2018, 8; Spigler et al., 2018; 8; Advani and Saxe, 2017), and can be precisely characterized for certain simple models (Hastic et al., 2019; Belkin et al., 2019; Bartlett et al., 2019; Xu and Hsu, 2019). Our work is inspired by Hastie 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: IRd + E with / hidden units
f(x) Σ a, wi)),
where E IRd is the input, w; E Rd is the weights corresponding to neuron i, E R is the i-th coefficient of the second layer, and 6 R + R is a Lipschitz continuous activation function with bounded Gaussian moments, i.. [[o(G)*] co, vk E E+ for G N/0. 1). For concise notation, we write W = [11, u.. E Tadx for the weight matrix, a = [a],...an] an) E R2 for the coefficient vector, X = [1 L. Rdxx for the data matrix, y T for the corresponding vector of labels, and क- = <W X) E TRxx for the feature matrix at the first layer. We omit arguments of f when they are clear from the context.
We consider 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:
(Izte ei) i.d. Pa x Ps. V = F(x.) Ees
L(X:f)
1 1 2n
f
(v:
where E[z,] = O. Cov (x: = E, E[es] = O, Var(e;) = 02. We are interested in the population risk
R(f) Ep, (F(x) f(. Our analysis will be made under the proportional asymptotics:
n, d, h oc; d/n- + 7, h/n + 725 71 E (0, oo),
in which overparameterizzaion corresponds to increasing 72. Thus the characteristics of double descent considered in this work are: 1) large population risk 72 - 1; 1;2) decrease in the risk for 12 > 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:
(t) = V LX;; dt or da(t) (X:
(3)
from small initialization. The rest of the paper is organized as follows. 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 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

---

Published as a conference paper at ICLR 2020
3 WARM-UP: LINEAR NETWORK
We begin with simple linear model with ((x)) = æ. i.e. + = W X. We remark that although the model is linear, the solution obtained by gradient flow 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 TaTT y/12- The following lemma characterizes the solution of the gradient flow.
Lemma 1 LLeast squares solution). Given data matrix X, response vector y and model f(:) (Ø(x) T W), a) with fixed first layer coefficients W, gradient flow on the coefficients a starting from zero initialization converges to 4 y, where stands for the Mloore-Penrose inverse.
We make two assumptions on the data and the teacher model to simplify the computation.
(A1) Gaussian Features: 2. ~0, Ia); (A2) Linear Teacher: F(x)= (æ, (3), |1
Denote the linear student network as f(x) (x, 3), where 3 W and the least-square
solution defined by Lemma 1. We write the population risk in its vlas-variance decomposition.
R = Ex P |1| - 31 w|
|| W - 313 + tr (Cov(B)X. W)).
(4)
B=bins
V=variance
where ||~:1/2 = T Ex.. We compute the bias and the variance separately to obtain the risk.
Theorem 2. Given (AI) and let w;
Lid. N(0,d='Ia) n. d, h
have
Y1 22,2 2 + 1202 72 71
71
72 2 + 72 02
72
1,
71922
7192
92
R((m<))
R(
71 02
7291 y2 +
91 + 92
12 71
0²
72 >
91
792
9192
where d/n + 71, h/n + 72. = In 11. and = 1h2 1/.
We observe that when d > (i.e. 71 > 1), we obtain the double descent risk curve, i.e., the population
risk achieves its maximum at 7/2 + 1 and further overparameterizztion (72 > 1) reduces both the bias and the variance. Conversely when n d and h > d (i.e. ?1 < min 72)). the population risk becomes constant and equals to that of the minimum-nome solution Busn X y on the input features.
Training the First Layer. When the first layer of linear network is optimized via gradient flow and the second layer is fixed, the following holds for zero- aaro-intiialization of W. Proposition 3. Given (O) 0 and fixed 'nir + O. at any time W. (t) is rank-1. I. Further, / Wa converges to the least squares solution of y = XT 3, the population risk of which is given in (Hastie et al., 2019, Thm. 1 & 3)) as
0 of the gradient flow
71 o²: R(
71 1 2
R( 1 7
(6)
7
71
In this case, overparameterization by increasing 722 does not influence the population risk. In addition, since the obtained two-layer linear model is equivalent the minimum norm solution on the input features Bmin 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 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-ing zero-initiann iroo-iitiaaizationn always yields the same solution that is independent to overparameteriizaton.. 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.
4

---

Published as a conference paper at ICLR 2020
10
10
1.0
emprical prediction
empiniall lower bound
empincal
bound
0.89
VDl varfancca
0.6
0.2
8:50
(a) risk (711 < 1).
(b) variance (ReLU).
(c) bias (ReLU).
Figure 2: Population risk of two-layer neural networks with optimized second layer under (A (A1).2).. Brighter color indicates larger ?1. (a) risk of lincar network with r²/r² //² = 16 and "1 1. (yn > 1 is shown in Figure 1) (b) variance of network with ReLU activation. Black line corresponds 71 oo predicted by Corollary 5.
(c) bias network with ReLU activation. Black line corresponds to " -> oo for linear network, which empirically observed as an upper-bound. Note that as 72 + both bias and variance becomes unbounded.
4 NONLINEAR MODEL: OPTIMIZING THE SECOND LAYER
In this section, analyze the case when the second layer a is learned under fixed W and nonlinear activation function & (a random feature model). We first observe that by Lemma 1, the gradient flow finds the solution à Øi We again consider the following Slas-variance decomposition.
R Ea-ra [1 W)å F(=) X, w] W) x, w| F
Ex
EE ||0( [lo(=TT))
w)å Bld( w)a
bias
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.
i..dd Theorem 4. Given (Al) and w; N(0, d d')a when n, d, h + oo, we have
7/2
12 < 1,
V
a
a mz (E, ci 022)
12 1
a² lim <-+
1/2 Dx
mı m1 012,
72
==
Ox
z=0
82
which mn (. ma(E, P.T T) is the unique solution {1mil, |m2| < 1/$5)mr of
r'm "2mi(ceme 2yc +3
1 mj - 727 2 m. -
(8)
m (1(0mm2-77) -T) -7122
720mmf(o22221 T)
m,y + mi
(9)
-T) (c)ma-r))-111 -7)
where variables P. satisfies 36 > or < 0, 00.0p>7333 0, and constants C1, C2 defined as,
C1 = -
C2 = E[Go()))
(10)
for NN0 1) and 3 denoting the imaginary part S.
Remark. Ci 2 C2 and the equality holds iff o is linear.
Corollary 5. we let 71 + oo, the variance is equal the lowest value of the variance of the linear
model Vi,, -c)) = o' miinn(2: 1}/11 72).
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 consistent with the double descent risk curve it suggests that for all 71, variance of the random feature model
5

---

Published as a conference paper at ICLR 2020
peaks at h then drops as "2 further increases. Note that " oo, 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 71, 72. Instead, we show that for linear teacher, the bias also becomes unbounded as 7/2 + 1.
i..d. Proposition 6, Given (AT)(A2) and wi ~ N(0 Ia), then B - ao as 72 - Furthermore, B is finite when 72 1.
Thus have shown that "cusp" in the population risk appears at h n, which aligns with the double descent phenomenon. Empirically as " - oo 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. I & 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 two-layer neural network with fixed second layer coefficients initialized from Rademacher distribution a; (nif--1/Vh, h, IVh), and the first layer W is optimized with the following update
OW (t) Øt
aL(X; W) aw
1 Σ y a TOW 2;)] a. 1 (2x W) o al,,
(11)
which potentially has different stationary solutions with no explicit form, depending on the initializa- tion. We denote the solution of this flow at time f started from designated initialization by Wini (t),
its stationary solution by W, and the corresponding network by f.
Remark. Although we let - oo, this dynamics does not corresponds the population gradient flow considered in Tian (2017). For instance when æ AO 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: va)an(0) (N Ta/dh
Non-vanishing: wxY(0) (N Ia/d
Note that neither of the two initializations correspond the "mean-field" regime (e.g. analyzed in Mei et al. (2018)) due the 1/ l/vh scaling of the second layer. In other words, as h increases, expect the distance traveled by each parameter to decrease under both initializations. The difference, however, is the "relative" amount the parameters traveled compared their initialized magnitude, which leads to solutions with contrasting properties. As we will see, under (A1)(A2) and vanishing initialization we have [0) (0)-WWll= / [IW(O))) > 1, i.c. 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 W(0).
5.1 VANISHIGG INITIALIZATION
As d, h -> oc, 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 d'(0) + 0; lo'(=z) o' o'(+))) = O(e
The above assumption requires that the derivative of the nonlinearity 6 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:

---

Published as a conference paper at ICLR 2020
Theorem 7. Given (AI-3). Let T O and f(-) fm W(T),, then as n, d, co, the gradient flow reaches a o(1) first-order stationary point at time T, ie. (T) E o(1), at which point the population risk is given
R(f) + max (0. 0..21 1, +
71 10
min (n, 13 0²
(12)
[] 11-71 71
The expression above is the same as the risk of the least squares solution on input 3 = X Xly; therefore the risk is independent to overparameterization (increasing 72). The intuition is that when the weights initialized sufficiently small and travel infinitesimally, then the activation can be linearized around O 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- NON-VANISHING INITIALIZATION
When initialization 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) asymptotically equivalent to the kernel gradient flow on the tangent kernel: k(xx. y) = Wai fxx) V Wim f fz)). The converged parameters under this linearized dynamics has the following closed-form:
voc(W") vec(Wiii)) A: A= J' fial(() Jua V (e ft (x:).
(13)
where J E (xtxh)) the Jacobian matrix w.r.t to the model parameters. We remark that 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. 22) our result is not built upon the overparameterization on width, but instead anticconcentration that relies 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, A 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 72-
An obstacle in computing the risk of the kernel model is the potentially non-zero finil (X). We thus adopt the "doubling-triee" trick"" from Chizat and Bach (2018b) to ensure (-) = 0, i.c. assume the following symmetric property on the initialized weights:
(A4) Symmetric Initialization: Vi E [1. E [1, s.t. a.wo 0, a.w.jit.
Theorem 8. Given (A1-4) and let n, d, h + xc, the stationary solution f has the following risk
R(}-
71 11-1-
7 (rI + 7um + 2) +
,2
271
2n mnvvnnn + mm1 + +2 + 2) - 2) +1
+
71 + 7 +1
1
0²
(14)
41/1 + m ( (m + +) - 2) - 4
where = 17/02. b2 = E()' and 13 (G) - 12 G G~
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 Jt is stable due the nonlinearity and dh > n. We make two additional observations:
the stability of the inverse at 71 + depends on the lowest eigenvalue of the tangent kernel matrix (smaller Amin (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 72, but the bias is higher than that in symmetric initialization shown in Figure 8. We comment that the non finit in (13) behaves as zero-meat Gaussian due to central limit theorem; therefore, in the kernel regime the function output initialization is equivalent to additive noise to the labels y, and we thus expect the magnitude of finit(X) (X) to negatively influence the model generalization.
7

---

Published as a conference paper at ICLR 2020
0.6
vanishing (emprical)
1.4
vanishing (emprical) non-vanishing (empirical) prediction
non-vanishing (empirical)
0.5
prediction
1.2
1.0
0.4
Varancca
0.8
0.3
0.6
0.2
0.4
0.1
0.0
0.50
0.75
1.00
1.25 V1 din
1.50
1.75
2.00
0.50
0.75
1.00
1.25
1.50
1.75
2.00
= d/n
(a) bias.
(b) variance.
Figure 3: Bias and variance of two-layer sigmoid network with optimized first layer under (A1)A22) Individual dotted lines correspond different (from 0.2 to 2) which independent the risk. The bias and variance for both initializations is well-aligned with Theorem and Theorem 8.
5.3 COMPARING THE INITIALIZATIONS
Figure shows the agreement between theoretical prediction and experimental results. Although in both cases the risk is independent to overparameterization (72), the two initializations lead to models with contrasting properties, demonstrated by the following comparison the risk.
Corollary 9. For any 7/1 E (0, oo) and nonlinearity o. B(f & B(f*) < 1. On the other hand, for all m > 0, V ) O(), whereas V (ffiin) can be arbitrarily large as "1 + 1. Remark. m 2 for all smooth activations Φ and the equality holds if Φ 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, 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 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, 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 low-rank weight matrix, as predicted by Theorem 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-dimensiona two-layer neural networks in learning 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 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 l penalty) stabilizes the singularity at 72 + 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, particular, the interpolation limit could occur at f2 + 0 which is beyond the regime we consider (such Figure 4 in (Belkin et al., 2018)). Thus what conclude is that under the studied proportional
the practical benefit of overparameterizattin. In
8

---

Published as a conference paper at ICLR 2020
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, rely (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
set of strong assumptions similar to those in Hastie et al.
extend the current analysis to beyond two layers. Last but least, our result characterizes gradient flow which resembles gradient descent with small stepsize, and thus 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).
ACKNOWLEDGEMENT
We thank Xiuyuan Cheng, Xuechen Li, Yiping Lu, Atsushi Nitanda, Shengyang Sun and anonymous reviewers for helpful comments and feedback. and DW were partially funded by LG Electronics and NSERC. JB and MAE were supported by the CIFAR Al Chairs program. TS was partially supported by JSPS Kakenhi (26280009, 15H05707 and 18H03201), Japan Digital Design and JST-CREST.
REFERENCES
Madhu S Advani and Andrew M Saxe. High-dimensional dynamics of generalization error in neural networks. arXiv preprint 17X0::1710.06667,
Zeyuan Allen-Zhu and Yuanzhi Li. What can resnet learn efficiently, going beyond kernels? arXiv preprint arXiv:1905.10337 2019.
Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. arXiv preprint arXiv:1811.04911 2018a.
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. convergence theory for deep learning via over- parameterization. arXiv preprint arXiv: /8/ 1.0962, 2018b.
Sanjeev Arora, Rong Ge, Behnam Neyshabur, and Zhang. Stronger generalization bounds for 2018. 8.
deep nets via a compression approach. arXiv preprint
Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Ruslan Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. arXiv preprint arXiv: //955, 2019a.
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. arXiv preprint /rXiv:1901.00884 2019b.
Zhidong Bai and Jack W Silverstein. Spectral analysis of large dimensional random matrices, volume 20. Springer, 2010.
Peter L Bartlett, Dylan J Foster, and Matus J Telgarsky. Epectrally-normaaline margin bounds for neural networks. In Advances in Neural Information Processing Systems, pages 6240-6249, 2017.
Peter I. Bartlett, Philip M Long, Gábor Lugosi, and Alexander Tsigler. Benign overfitting in linear regression. arXiv preprint arXiv:1906.1 300 2019.
Mikhail Belkin, Daniel Hsu, Siyuan Ma, and Soumik Mandal. Reconciling modern machine learning and the bias-variance trade-off. arXiv preprint arXiv: arXiiv/8/2.111 18/2.11118 2018
Mikhail Belkin, Daniel Hsu, and Ji Xu. Two models of double descent for weak features. arXiv preprint 1rXiv: 2019.

---

Published as a conference paper at ICLR 2020
Yuan Cao and Quanquan Gu. A generalization theory of gradient descent for learning parameterized deep relu networks. arXiv preprint arXiv:190..01384, 2019.
Xiuyuan Cheng and Amit Singer. The spectrum of random inner-product kernel matrices. Random Matrices: Theory and Applications, 2(04):1350014) 2013.
Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over-parameterized models using optimal transport. In Advances in neural information processing systems, pages 3036-3046, 2018a.
Lenaic Chizat and Francis Bach. A note on lazy training in supervised differentiable programming. arXiv preprint 11Xvv1811..07950 2018b.
Bin Dong, Jikai Hou, Yiping Lu, and Zhihua Zhang. Distillation early stopping? harvesting dark knowledge utilizing anisotropic information retrieval for overparameterized neural network. arXiv preprint arXiv:1910.0123: 2019.
Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterriee neural networks. arXiv preprint arXiv:1810.02054. 2018.
Gintare Karolina Dziugaite and Daniel Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. arXiv preprint arXiv:1703. 2017.
Noureddine El Karoui et al. The spectrum of kernel random matrices. The Annals of Statistics, 38(1): 1-50, 2010.
Zhou Fan and Andrea Montanari. The spectral norm of random inner-product kernel matrices. Probability Theory and Related Fields, 173(1-2):27-85, 2019.
Mario Geiger, Stefano Spigler, Stéphane Ascoli, Levent Sagun, Marco Baity-Jesi, Giulio Biroli, and Matthieu Wyart. The jamming transition paradigm to understand the loss landscape of deep neural networks. arXiv preprint arXiv: 1809,0934 2018.
Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Limitations of lazy training of two-layers neural networks. arXiv preprint arXiv:1906.000 2019a.
Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Linearized two-layers neural networks in high dimension. arXiv preprint arXiv:/90k /2/9/, 2019b.
Sebastian Goldt, Madhu S Advani, Andrew M Saxe, Florent Krzakala, and Lenka Zdeborová. Generalisation dynamics of online learning in over-parameterised neural networks. arXiv preprint arXiv:1901.09085, 2019.
Suriya Gunasekar, Jason D Lee, Daniel Soudry, and Nati Srebro. Implicit bias of gradient descent on linear convolutional networks. In Advances in Neural Information Processing Systems, pages 9461-9471, 2018
Trevor Hastic, Andrea Montanari, Saharon Rosset, and Ryan Tibshirani. Surprises in high- dimensional ridgeless least squares interpolation. arXiv preprint /903vv/903.008 2019.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in neural information processing systems, pages 8571-8580, 2018
Ziwei Ji and Matus Telgarsky. Gradient descent aligns the layers of deep linear networks. arXiv preprint
2018.
Ryo Karakida, Shotaro Akaho, and Shun-ichi Amari. Universal statistics of fisher information in deep neural networks: Mean field approach. arXiv preprint 1rXiv:1806.01311 2018
Anders Krogh and John Hertz. A simple weight decay can improve generalization, pages 950-957, 1992.
10

---

Published as a conference paper at ICLR 2020
Mingchen Li, Mahdi Soltanolkotabi, and Samet Oymak. Gradient descent with early stopping is prov- ably robust to label noise for overparameterized neural networks. arXiv preprint arXiv:1903.11680 2019a.
Yuanzhi Li and Yingyu Liang. Learning overparameterrzee neural networks via stochastic gradient descent on structured data. In Advances in Neural Information Processing Systems, pages 8157- 8166, 2018.
Yuanzhi Li, Tengyu Ma, and Hongyang Zhang. Algorithmic regularization in over-parameterized matrix sensing and neural networks with quadratic activations. arXiv preprint arXiv: 2017.
Yuanzhi Li, Colin Wei, and Tengyu Ma. Towards explaining the regularization effect of initial large learning rate in training neural networks. arXiv preprint arXiv:1907, 04595, 2019b.
Tengyuan Liang and Alexander Rakhlin. Just interpolate: Kernel" ridgeless" regression can generalize. arXiv preprint
het::00080000000 2018.
Zhenyu Liao and Romain Couillet. On the spectrum of random features maps of high dimensional data. arXiv preprint arXiv: / /9/6, 2018.
Cosme Louart, Zhenyu Liao, Romain Couillet, et al. A random matrix approach to neural networks. The Annals of Applied Probability, 28(2):1100-1244 2018.
V.A. Marčenko and Leonid Pastur. Distribution of eigenvalues for some sets of random matrices. Math USSR Sb, 1:457-483, 01 1967.
Song Mei and Andrea Montanari. The generalization error of random features regression: Precise asymptotics and double descent curve. arXiv preprint /rXiv:1908.05:0 05.555555,, 2019.
Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two- layer neural networks. Proceedings of the National Academy of Sciences, 115(33) 2018.
Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Mean-field theory of two-layers neural networks: dimension- free bounds and kernel limit. arXiv preprint arXiv:1902.06015, 2019.
Vaishnavh Nagarajan and Zico Kolter. Generalization in deep networks: The role of distance from initialization. arXiv preprini arXiv:/901.0f1
2019.
Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. In search of the real inductive bias: On the role of implicit regularization in deep learning. arXiv preprint urXiv:1442.6600 2014.
Behnam Neyshabur, Zhiyuan Li, Srinadh Bhojanapalli, Yann LeCun, and Nathan Srebro. Towards understanding the role of over per-parramttizationn in generalization of neural networks. arXiv preprint 1r2.:::805./2076, 2018
Atsushi Nitanda and Taiji Suzuki. Stochastic particle gradient descent for infinite ensembles. arXiv preprint 1712://772.055333 2017.
Samet Oymak and Mahdi Soltanolkotabi. Towards moderate overparameterivation: global con- vergence guarantees for training shallow neural networks. arXiv preprint arXiv:1902.04674 2019.
Jeffrey Pennington and Pratik Worah. Nonlinear random matrix theory for deep learning. In Advances in Neural Information Processing Systems, pages 2637-2646, 2017.
Jeffrey Pennington and Pratik Worah. The spectrum of the fisher information matrix of a single- hiyden-layer neural network. In Advances Neural Information Processing Systems, pages 5410-5419, 2018.
Grant M Rotskoff and Eric Vanden- Vanden-Eijnden. Neural networks interacting particle systems: Asymptotic convexity of the loss landscape and universal scaling of the approximation error. arXiv preprint arXiv: 1805.0800000000 009 2018.
11

---

Published as a conference paper at ICLR 2020
Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks: A central limit theorem. arXiv preprint arXiv: 09372, 2018.
Samuel L Smith and Quoc V Le. A bayesian perspective on generalization and stochastic gradient descent. arXiv preprint arXiv:1710.0645/1 2017.
Zhao Song and Xin Yang. Quadratic suffices for over-paramttrization via matrix chernoff bound. arXiv preprint arXiv: wrXvv:/96.00359; 2019.
Stefano Spigler, Mario Geiger, Stéphane d' Ascoli, Levent Sagun, Giulio Biroli, and Matthieu Wyart. A jamming transition from under-to parammetttttiiiionn affects loss landscape and generalization. arXiv preprint arXiv: 1/X.099//0 2018.
Taiji Suzuki. Adaptivity of deep relu network for learning in besov and mixed smooth besov spaces: optimal rate and curse of dimensionality. arXiv preprint arXiv:1810.08033. 2018.
Terence Tao. Topics in random matrix theory, volume 132. American Mathematical Soc., 2012.
Yuandong Tian. An analytical formula of population gradient for two-layered relu network and its applications convergence and critical point analysis. In Proceedings of the 34th International Conference Machine Learning- Volume 70, pages
JMLR. org, 2017.
Colin Wei, Jason D Lee, Qiang Liu, and Tengyu Ma. On the margin theory of feedforward neural networks. arXiv preprint arXiv: 180000566, 2018.
Francis Williams, Matthew Trager, Claudio Silva, Daniele Panozzo, Denis Zorin, and Joan Bruna. Gradient dynamics of shallow univariate relu networks. arXiv preprint arXiv: /906. 1906607842 2019.
Blake Woodworth, Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro. Kernel and deep regimes in overparametrized models. arXiv preprint vrXiv:190.005877 2019.
Ji Xu and Daniel Hsu. On the number of variables
in principal component regression. 2019.
Gilad Ychudai and Ohad Shamir. On the power and limitations of random features for understanding neural networks. arXiv preprint 19Xiv:1904.00687 2019.
Chiyuan Zhang, Samy Bengio, Moritz. Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. arXiv preprint arXiv: /61 1.03530, 2016.
Kai Zhong, Zhao Song. Prateek Jain, Peter L Bartlett, and Inderjit Dhillon. Recovery guarantees for one-hidden-layer neural networks. In Proceedings of the 34th International Conference Machine Learning- Volume 70, pages 41 4140-4149, JMLRR org, 2017.
12

---

Published as a conference paper at ICLR 2020
A
ADDITIONAL FIGURES AND PLOTS
Initialization
0.15
0.004
optimum
0.10
0.000
0.00
-0.05
0.002
-0.10
initialization
-0.004
optimum
-.000
0.004 -0.002 0.000 0.002 0.004 0.006
-0.10
-0.05
0.00
0.05
0.10
(a) vanishing initialization.
(b) non-vanishing initialization.
Figure 4: trajectory of neurons from initialization (dark blue) optimum (orange) the first dimensions (two-layer Plus student and linear teacher, SNR= 1/4). For vanishing initialization the neurons stay close one another throughout trajectory, whereas for non-vanishing initialization neurons stay close to initialization.
3.0
6
empirical prediction
empirical
prediction
5
Age
2.0
19
1.5
risk
3
1.0
0.5
0.2
0.4
0.6
0.8
0.5
1.0
1.5
2.0
2.5
V2 h/n
Y2 =
(a) 11 <
(b) m > 1
Figure 5: Population risk of two-layer linear network with fixed random 1st layer with SNR=2516 under Gaussian input and linear teacher. Brighter color indicates larger ?1-
SUMMARY OF THE PRESENCE ABSENCE OF DOUBLE DESCENT
Singularity in
2nd Layer Trained (RF)
Vanishing Init.
Non-vanishing Init.
Bias
71 No; 72: Yes
71 No; 72: No
71 No; 72
Variance
?1 No; 72: Yes
71 Yes; 72: No
71 No; ?2! No
13

---

Published as a conference paper at ICLR 2020
0.6
vanishing (empirical)
1.2
venishing (empirical) non-vanishing (empirical) prodiction
non-vanishing (empiricall
0.5
prediction
1.0
0.4
Varlance
0.3
0.6
0.2
0.2
0.0
0.50
0.75
1.00
1.50 y1
1.25
1.75
2.00
0.50
0.75
1.00
1.50
1.75
2.00
yi
(a) bias.
(b) variance.
Figure 6: Bias and variance of two-layer Plft network with optimized first layer under (AI)(A2). Individual dotted lines correspond to different (from 0.2 to 2) which is independent to the risk. The bias and variance for both initializations is well-aligned with Theorem and Theorem 8, respectively.
500
120
100
400
300
risk
60
200
40
100
20
0.5
1.0
1.5
2.0
0.5
1.0
1.5
2.0
h/n
hin
trained 2nd layer.
trained 1st layer.
Figure 7: Population risk (scaled by 1/d) of two-layer ReLU network trained to fit two-layer ReLU teacher model with h d neurons. Brighter color corresponds to larger ni Similar to the lincar teacher casc, double descent is observed when the second layer is optimized (a) but when the first layer is optimized (b).
0.8
i.d data (empirical)
iid. init (empirical)
symmetric (ompirical)
symmetric init (empirical)
symmetric data (prediction)
symmetric init (prediction)
0.6
0.4
0.4
0.2
0.2
0.0
0.0
0.5
1.0
1.5
2.0
0.5
1.0
1.5
2.0
din (a) SoftPlus activation.
Y: din (b) sigmoid activation.
Figure 8: Bias of (a) Plut and (b) sigmoid two-layer network with optimized first layer under (A (A2).... Note that bias under i.i.d. initialization is also independent overparamcterization (72), but is higher than the bias under symmetric initialization ("doubling trick") and not always upper-bounded by the null risk ?
B BACKGROUND
ROTATIONAL. INVARIANCE
The rotational invariance of Gaussian distribution is crucial in our analysis throughout this paper. basic observation that for random Gaussian matrix X and any fixed unitary matrix U, the distribution of X and U X are the same.
14

---

Published as a conference paper at ICLR 2020
Lemma 10 (Rotational Invariance), Denote A(X) Raxd matrix of X E Rdx 1f A(X) satisfies that A(U = =UA(X)U" for all unitary U, then
(xx13"A(X))) = / g (A(X))]
(15)
for any fixed nonzero 3 E IR and random matrix X with each entry ii.d. Xzj ~ ((
Choose a set of Unitary matrices {U:] such that U, 3 = || e, where e, the i-th canonical vector in Rd. Since U,X X, we have B[A(X)] B[A(U,X)] U[E[A(X)]U, and hence
Proof.
A[6"A(XX)) Σ U 3 E [[[[(()],; 3² Eftr (A(X))). (16)
Note that the property also holds for matrix function A(X) that satisfies A(XU) = -(:))UU((A UA(X)UT and 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 real symmetric random matrix A E RPxp, define its empirical spectral density as
(a(d) = Σ (A)da)(AdXx,
1
p i=1
where da (x)) = 5(x a) is the Dirac delta function. Assume A = Wp Wp(T,n) is Wishart matrix, i.e. W X X/n and X E R=x is random Gaussian matrix with each column i..d. X; N0, I). Marčenko and Pastur (1967) showed that as n. p + oo and p/n = 7 E (0, oc), the empirical spectral density jw, (dX) converges weakly to limiting density (MMP(>)
(MP())
- (A)) + 2 V ((1 v A)(A (1 Vd)
(17)
We say /MP(s) is the density of the Marčenko- Parttnnk-PPustur distribution with support S S= [(1 111
v (for y > 1). Note that this implies that the smallest non-zero eigenvalue of A is bounded away from O a.s. for 7 * 1.
77) (for 0 < 1) S = (0) [([1 - 2 (1+
The explicit form of Marčenko-Pastu distribution allows
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 Marčenko-Pastur distribution, i.e. for any bounded fA) E C(S), as n, with p/n 7, almost surely
[ f(A) pw
/ fs(A) (dd(
(18)
One implication is the following trace concentration on the inverse Wishart matrix for y < 1,
tr (xTx)== = tr p -
(.W - [. (dA) + [ { (MMP())(A)) =
1
(19)
We remark that instability of the trace of the invert Wishart matrix as 7 + 1 plays an important role in the double descent phenomenon.
B.3 ORTHOGONAT. POLYNOMIALS
Orthogonal polynomials are useful in the analysis of nonlinear random matrices. Suppose & is a function in (2, PG), where G ~ (0, 1), no(dx) = (V2m)-11e 1=2222222 i is the Gaussian measure. For 2 0. define the Hermite polynomials
x2 0n
Hx(x) = (-1)"e-=*
x2/2
(20)
Drn
15

---

Published as a conference paper at ICLR 2020
Note that orthogonality be easily verified:
([H((G)H&(())] = H)(+)H(())c((dd) B jt. -8kk
(21)
Since {H,( forms set of orthogonal basis in (R, pc), the function d(x) can be expanded under the Hermite basis as
1 (())H(())cc(aa)) Hz(x). k! 3
<(x) (H,(x)= Σ
(22)
1=0
Following Cheng and Singer (2013), we mainly focus on the first few terms of the Hermite expansion. One can check that Ho(x) 1, H1 (z) æ, and Φ(x) can be expanded as
<(x) co + c + Φ1 (x),
(23)
where E[ø(G)), c1 E[G)(G)), and terms in the RHS are orthogonal to one another in (RR, pG). Taking square and expectation over both sides yields
which indicates E[o(G)) = [ + [GG Bi (G)²),
(24)
[ 2 0. Note that the equality holds and only E[ (G)* 0, ic. ((x) = co + is linear. One application of this decomposition is "linearize" non-linear matrix in high dimensions, which will be useful for the following sections.
C PROOF OF MAIN RESULTS
C.1 PROOF OF LEMMA
Given features X gdxn, labels y Rd and model parameters 0, the gradient flow of 0 on the squared loss y X TTol can written as
a0(1) 1 --y-xy-XXT x(ly-xT((t))
(25)
Ōt
Thus with initialization On. the solution of this ODE time can be written in explicit form
(tt) = - xxx Oo + (XXT) (1-e- Xy.
(26)
Since 0, 0, taking + oo yields the desired result.
C.2 PROOF OF THEOREM 2
We compute the bias and variance for different cases of 71, 72. We first discuss the case where the random feature & not full rank (Case I). Otherwise when 4x is full rank, we discuss whether is full column rank (Case II) or full row rank (Case III).
Case I: W X is not full rank, i.e. 71 < 1, 71. In this case rank (Φx) min(n, h). and thus by taking the Moore-Penross, inverse we obtain
B = w( wwXTTw)l = (XxT)-
(27)
It is clear that the mean and variance is identical to the underparamettrrized regime in Hastie et al. (2019), ie. when n, d, oc,
7 0².
B - 0: V -
(28)
7
16

---

Published as a conference paper at ICLR 2020
Case II: WTX X has full column rank, ie. 72 < 1, > 72. By Lemma 1, the solution of the second layer coefficients is
B = 3- Xy-
(29)
Denote W UEVT the singular value decomposition. We perform the block decomposition:
Σ-
X X1
Σ
X
(30)
where En Rixh, Xo E R X2 E (d=A)xxn and notice that Xo. X1 are independent. By concentration of measure argument (e.g. Tao (2012); Hastie al. (2019)) one can show that the quantity below tightly concentrates at its expectation. For the variance we have
V = ur ( (W(WTXXT))" X ((W
(WTXXX
W) U (
xx
xx
WT)
ortr ( W (W xx xx o'tr MΣ
XXTUE)
2er E E (xXxxX =
ortr ((xox))
0² 72
(31)
7/
where the last equality follows from Appendix B.2. Similarly for the bias term we have
X. X X B $1²
=3T (V x -Ta) T
- -aa)3
@ y2 ((W(TTXXT))-WWXXX ( -Ia)
(USX
(.)().
d tr
- (
tr J
(IITTXXTE--XXXX (..).
(32)
where symmetric arguments are omitted as (...) and (i) follows from the rotational invariance
argument introduced in Lemma 10 and that 3 B = ? By the block decomposition (30),
Σ x - Ia = [0
0 (XoX.) ) ].
(33)
Ia-h
Therefore the bias term simplifies to
2 B = tr ((Σ(E) x - xxT -1) (S(ET xx( xxxT - -a1)
(
2
(tr ((XoXo) - Xoxt X1X Xox + (d (d-1)) h))
d
2 (d (d-h)/h
hi nn-h- + d d- h) h)
71 - 72 p
(34)
71(1 - 72)
Thus have obtained that n, h
- n1 7/2 ,
(35)
71(1 - 72)
Case III: WT X has full row rank, i.c. 71 > 1, > 1. Similarly, the least squares solution is
B-WWT x(X WW x)-1y. y.
(36)
17

---

Published as a conference paper at ICLR 2020
Simplifying the variance:
V = tr (WWT x (x (X'"xXx""""""""""
x) x xTwTT) oÂur (WWT (VE W -? UTT ~ ortr (WWTE (Σ
ww
(37)
where we applied the SVD of X UEVT and the rotational invariance argument. Using a similar block decomposition on W:
E
Wo
E
W
(38)
Wi
where Eg E Ruxn Wo E Ruxh Wi E (dd- and Wo, W1 independent. We thus simplify V as
V o2 (WiWWWWESE" WW Σ) - WWT)
= o ([ WoW E EoO5 Wo 2
WiW 2o (2 Woh. 20) -22-22oo 1) ṭ² (ur (En _60) + tr (WiWo (Wow))-'>>) "2020' (WoWo) wol,T
ortr ((xT)) ter (W7 W. (Wow)) -1 (Eg
(E) En) (WoWo) -Wo) (39)
Hence we obtain the following expression on the variance
Det.:xVir(WWWWWy) ((Wow)) (2) (EF (1 - 0²
m1 - + 02 - ((dd-n)Ewe Ew.
11
7/2 - (40)
We omit the derivation of bias, which follows similar derivation:
B - 72(71 1)
(41)
71
Combining Case I, II, III yields theorem 2.
C.3 PROOF OF PROPOSITION
Given the squared loss, can derive the dynamics of W with fixed second layer a w.r.t the loss:
@w
Öt - 1
x x W(t)aaaT
(42)
Note that the update of W can be written as a linear combination of a. Since W(0) 0, we can write (tt)= (t)a for some The corresponding flow on t
Pir(() X(y X ((a))t)aa (|))))
(43)
Øt
n
which gives the following solution
X =
(44)
la 2
Thus gradient flow on the first layer leads to the minimum-norm solution on the input features.
18

---

Published as a conference paper at ICLR 2020
C.4 PROOF OF THEOREM 4
Following the bias-variance decomposition (7), the variance term can be written as (o² omitted)
V Ez,s ||a <(W x) - EEa' "(W x>
Ex [|(o(WTT x)] ((W a0) 10
tr (Lx(X [T)]TW]] (o T x)] 'Es (oW x)((WTx)) ")
tr (0(XTW)]] (W [x) kw),
(45)
where we define the expected non-linear Gram matrix Kw E R as
Kw - En @WW x)o(W x) T
(46)
and for each entry we have (Kw [..5] Eu [6 (w m)
Random matrix 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, specified in the following lemma from Pennington and Worah (2017): Lemma 11. Suppose @ is not linear. Then the smallest singular value of $ @( X) E TRAx is of order O OVvn). To be precise, consider the empirical spectral density Poot /n (dA)), when n, d, h with d/n + 7/1 and h/n + ?/, 1.b-- /n (dx) converges weakly to peo0 /a(d)) [1 1/ + 50 (A)dA 1xt (dX),
where et (dA) has non-negative support [p, oo) 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 - ( X)O(T w)) Kw
lim tr -E) -x1))<<w) E-0
lim Ve, E-
where the continuity and boundness of V&3 & = 0- guaranteed by Lemma 11 when 72 h/n + 1. A ridge & is added make use of (Louart et al., 2018, Theorem 1), which derived the asymptotic equivalent of the resolvent (ΦD -5---) {1) -1 follows that n, d, h oc.
Kw
h h-1( -EI) -1
h h-
+ h-lir (1(TTT - ( 'Kw)
(Φ T -ET) -1_
n
Kw
5 h
h 1 h 1-ttt (IfODT-El 'Kw)
11
Kw
h + h 1t
(h( -1 "Kw)
O( 11/2t )0 +0
(47)
h
where we have used the inequality tr (AB) < || r |B1| and equivalently by taking g - 0 we get
ir (Φ0T -ET) 'Kw)
lim lim
1. h + ir (#OT -E)) -1 Kw
(48)
8 8 d
19

---

Published as a conference paper at ICLR 2020
Therefore limg-+0 di n/h-tt (00 -E1) -
-s1)~kw> Kw) = 72/ 72). Note that aV./0g
ir (5(#0TT (w) is bounded around the neighhourhood of g = 0, hence following the same argument (Hastie et al., 2019, Theorem 4), exchange the limit of g and n. d, h 0:
72
V
(49)
72
Case 2. > n. Techniques used in the current proof are largely borrowed from Hastic et al. (2019); Cheng and Singer (2013), and include the full proof for completeness. should be noted that compared to Hastie et al. (2019) we handle the non-zero expectation the nonlinearity under Gaussian distribution, i.e. the off-diagonal entries of the kernel matrix 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 caleulating the variance directly, we analyze a modified quantity Ve and then take ş - 0, which can be connected to the trace of the resolvent of matrix Å defined in (56); this translates the calculation of Ve into the calculation of the Stieltjes transform of Á (57), (61), (62).
C.5 DERIVING THE VARIANCE V FOR
Step 1. An equivalent expression.
For notational simplicity we omit the magnitude a:
= ( (( (O) (W) W)((T x)) "(x Ww)
(50)
and due to the same continuity argument as in Case we have
_ tr (S(S s (ln) Kw) 8
V = lim
lim Ve.
where S = @( XV/VX)/V==00 v/V C and o > 0 or & < 0.
We decompose the normalized feature matrix S = o(W X)/ n as S = UEVT where E
diagıxm (01,, , 2x)) E *x* " is tall diagonal matrix, and U [u,,... un Rhxh is the set of orthogonal eigenvectors of SS <(W X)Ø(XT)) W)/n, and V E RR is the set of orthogonal cigenvectors of S S = o(X W W)o(WT x)/n. Now the variance can be written as
1 Ve
(UE(E -UU Kw)
(51)
n
By the same argument as in Lemma 13 of Hastic et al. (2019) one can show that when n, d, h
+
oo
1 tr (UΣ(E Σ + -Z Kw)
frr (UN(E (Im) Rw
V8
(52)
in which Kw is the approxmation of Kw defined in Lemma 16. Writing the trace explicitly (denote eigenvalues A of and Φn+1 Φn 0), we have
VS tr (U n
Rww
12
1 (As + E)2 in Kwu,
Σ
X
(53)
Since the positive support of spectrum À is lower bounded and the density at 0 is - -- we have
1 lim
A (A: + Eyz u. Kwu = 72 / (a+ 72 /
A
A
VS + 12 h +00
usD (A (A+ 1) (dA),
(54)
where we define un (x) = E/= Sx. (I) Kw u; and its positive part 1 (Ir). Hence we have
(d2) (d)) =1 ( 7 1 (d2).
1
V = lim
VS= = lim
A>p (A 43
(55)
8
20

---

Published as a conference paper at ICLR 2020
We define the following matrix (,, E ERNXM where N = + h:
(,, [
pp pla + s1,1, +rQ S
(56)
5
On
And denote the Stieltjes transform of An as
(5, tr -
(57)
Then following the definition of Kw show that
mn (5, ræ, Sz, tx) = tr (1
KWI EIh
5 Inn
1 ]
(58)
ST
and taking matrix derivative gives
1
-EIn
S
Rw
S
0x
mm rz, sx, tz)
==
-
1
ST
-In
ST
-In
1
{ + SST
[
Kw 01
1)
0
In + STSS
0 0
1
(U(Σ + U Kw)
= Σ
A x <2 Un Kwu.
n
(59)
Denote the limit vi(S, 5, litm,, mn (E, Ps T, and the derivative is given as
①
mE rx, sx, sx)
lim
-
Σ
1
u, Kwu,
Ox
x=0 h,d,n +20 +20 71
x
52
1
72 72-1
/
1
=72
+
A>p + te (dA). (60)
A +
g2
A20
For simplicity we define the following function on E:
a(E) = dr
mi rw, sx, tx)
(61)
2=0
also denote a+ (5) - ((E) 22-1 72 JAZP x+.* to (dX). Comparing with (55) yields:
= lim (E). 5-+
(62)
Step 2. Calculating q(E) and m, p,s,(
T, This subsection aims to calculate mi{E, s. and q{E) = -m, rz, sæ, tx (z=o, from which the variance can be computed from (57)()1)(62).
We define A, by subtracting the off-diagonal entries of the upper-left block of A,;
An(p.T) pln +7Q ST
S
(63)
On
where S S aolpxxs i.e. Sik = blw xa.) - a0 = =(ww] xx). ) ap = E[d(x)). The Stieltjes
transform of A, given by mn(E.p.7 = 1 tr ((An(P,T) - -Tx)) The following Lemma shows that m, and mn have the same limit:
Lemma 12. when n and for Si > 0 or & < 0, we have mn mn (E.P.T) - m., (S,
Proof. By definition of An and Ans
An(p.s,7) An(0,)==
1i
On
ag
1h
Oh
(64)
On
In
ag
On
1n
21

---

Published as a conference paper at ICLR 2020
which is 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. random matrices:
sup F (x) PMn (x:) s0(n
(65)
where F* is the empirical spectral distribution of M E R=x", The claim follows from the Stieltjes continuity theorem (e.g. Section 2.4 in Tao (2012)).
To calculate the Stieltjes transform Mas we take advantage of the block structure of An:
m1,n(5,P.T) tr 1
- -5T)))
(66)
m2 (5. -tr
- pp.
(67)
n
One can observe that mn (5: P; 7) = (, p, (2, In the following equations we omit the subscript as well as dependency P55,T. Following Hastie et al. (2019), we rewrite the matrix A, = A A, a where A is a (N - 1) x (N -)
a
matrix with last column and row of A removed and the activation vector:
A, pIn + TQ
S.
]: T = xT) 0
0 T_-1..
(68)
S
On-1
Hence by the block matrix inverse formula
* - 11- = * I-5- a a (IN-1)
(69)
Plugging this back in the Stieltjes transform we have
m2,n(E,p tr -ET)))t)) = =Ba -
n
= =Ea ( - a (A TAA" (IN- 7)
(70)
To obtain an asymptotic description of m2,,, we perform the orthonormal decomposition on the nonlincarity introduced in Cheng and Singer (2013): sfx) a₁ + (1),),, where
a1 (o) [=]](=)]. For each w,,1 composition (along the direction of xn and the direction of w; perpendicular to x,,):
S h). we perform the following orthonormal de-
En
En
W.
w: En
w:
(71)
[
We can thus simplify the activation vector a as
1
n p( ((
Vn ((1(ona111)))
n-
1 a1 x
a1 |2 1776
n
a [[==(,,000 .0T]
+
pi (1
(1120.1171))
(72)
vn
カ-1
a .0T]
22

---

Published as a conference paper at ICLR 2020
and for 1 1 k
In
In
Qij
?/
+
w
11
+
w
+
w
(73)
[.
||
1115
Similarly we decompose the activation function in S.
Sik
1
1/6 F
Lk
1
01
7/
2k
1
a ūo. a2
1
41
1/1
2k
uo!
vn
x
vni
nn
vni
|
1
4
xxx)
+
a1
x
1
41
+ wn _ arkk
10 (20 xx)
(74)
vn
vn
n2
[nn]
Sik
019
Eit
We thus have equivalent expression of matrix A.
pla
+
S.
Imm"
T
Eg
E1
A
+
ST
On-1
0uum
On-1
E
Oz
plp + +
S.
0g
a
7
R2
Eg
E
ST
On-1
0
u
@1
0
0
E
0
U
E
= =A UCUT E.
(75)
By argument similar to (Hastie et al., 2019, B.1.2), E d'mminishhe to 0 + oo with respect the Frobenius norm, therefore by the Woodbury's identity and the expression of m2,, in (70)
- a T (A. 5IN (Ix-1)=t J
m2.1 PT B.
( - - a TT (A. - --IN-1 + +UCUT)-Ia) "
+ Ea
Ea
( - f a T (Ā. - EIN 1) a+
a+
u
a (A. - ELN-1) U (0-l + U (A, - ((N-1)= -1
U(A. - ((N-1) )r
(76)
We bound each term u, v, S to compute m For
Egu- = Ea a (Å, - (IN-1) a = tr Es(ss "((A. &IN-1) = by.m1, P.T).
(77)
where b EE- of measure argument we have that as n, h, +
E.- Ea- [
From a standard concentration
" - E r 11 (5.
(78)
And for v (note that U is dependent on a)
Ea aa (Å. (IN-1) - = Eg [s Ts 1 -(IN-1) " On
0,-1 ]
Eav
=aa
=E (Ā. &IN- (....) 0] tr ((A, EI 1) D B. Bains T
0].
(79)
where tr ((A. EIN-1) EIN-1) Eolns" 1) a1 / We thus have
1 - F = / (6173/mmm 0)
(80)
23

---

Published as a conference paper at ICLR 2020
n. d 20.
And finally for S.
Ea
Sass
Eg
C'I - (A,
-(N-1))-UU
T
a1
n 0₂
(Å. 124--110 ((N-1)
n 0₂
01
On-1 u
0n-1 u
1/a1
17 (A.
ElN
f1N-1) 11.19)
0
1/a1
-T/09
0
u T (Å, B/N-1/btt...nn
((A, EIN
pp Ea Inm
0
1/a1
1/a1
-T/af
0
ur ((A. - 11N
Ba uu
0
1/a1
12 I (2,(E.
1/a1
-T/ap
0
V1
12 (E:P.T)
1/al
(81)
1/01
m2,n (5: T/
And hence n, d, h - x0.
1/a1 Therefore by combining (78), (80), (82), we arrive at the following expression on m2,,
71
(82)
1/a1
m2,n P, 7/a²
(2n(E,P;T) + Ea [( ((- -u- o Sw
uu ss) -1
1 +
-( (-5 - - m2mmin (aiv / m
2
711
1/a1 (2, -
1/01
[1.1]
72amm3 (afmmi, T)
-5-
(83)
m1,n(a (afm2,, -T) 7172
Similarly we can calculate m1,1(E,p,7 P) as
11,1 (S
22y 1 (affm-aamm-- -) (errr))aaaaaa.... 2raffm + at m m
-1
- 5-p--pp p-- - 2 m mm22
2mm
mami(amm2.n 7) 7172
(84)
Uniqueness (84)(83) follows from (Hastic et al., 2019, See B.1) and is omitted.
C.6 PROOF OF COROLLARY 5
In this section we take the limit 71 - oc. In this case (8), (9) simplify
(-E - 1/2m1)
(85)
m1 = (-5 - p - 7/2 2m rmg)
(86)
72mm (5, P,) + m2(E, p, T). By taking the derivative we have Recall that mmS.p.7)
(EE)
m( rxr, tx))
3 mm p, 0)
+ 5 + mm(, 0,))
dx
==0
p=o
r=
= -172 Dp
m2 (E. p, 0)
+ Dp
m2 (S, p, 0)
+ t/2 dr
m, (E. 0,
+ - "
m2 0,7)
p=0
1=0
1=0
(87)
24

---

Published as a conference paper at ICLR 2020
Observe that (85), (86) constitutes set of implicit functions. Thus differentiating the two functions with respect to T, p and then substitute by r gives
(r(1-- - 1) -2) (Virtn 1) +2)² (Vir(ma1)++ q(<)
-4 r 1) +
(88)
2 (r(12 + E²)² - 42722
Hence by (62) obtain the asymptotic variance:
(()) lim q+ (5) k-0
lim {-+
((5)
72
(89)
2²
72
Combining the case where 72 < in Theorem 4 completes the proof.
Remark. For ((x) = ReLU(x),
1/2 - 1/(2mr, 1/4. For ((x) = SoftPlus(z) =
log(1 + e"), numerical integration yields e₁ ~ 0.2715, 02 = 1/4.
C.7 PROOF OF COROLLARY
C.7.1 UNBOUNDED BIAS WHEN h
From the bias-variance decomposition (7), the bias B written B tr (Q1 + Q2 + Ia), where
Q: [( X))' Kw w)j'x:: Q2 = =[o(WT x)) WT
(90)
When n. due to the nonlinearity of o, we have o(W X) is full rank a.s., and hence 16(X W)] [6(XT))]* = 1 We have the following bound for Q1
d tr (Q1) = jor (Xlo(WT x)]'
x = for (Kwlo(X T 'x= x) [[]
2 d Amin (Kw) )r (lo(x x (x)-)= Amin(Kw) tr (SST)-1. xxxx x (Kw
n
Since W and X/v are Rdx== Raxy follows the same distribution where each entry ii.d. NO 1/d).
1 ddmin (Kw) tr (Q1) 2 tr (SSTT)-
g . xxTx) ~ -tr (S
tr (ST -1 (I+ 2)) = lim 0
mn x, 0, x)) + oc,
(91)
dx
where in Section C.5 we have showed (91) is unbounded when Weyl's theorem we have Amin (Kw) = O(1), and thus 't (Q1) is unbounded. For d d tr (Q2).
Moreover, by (154) and
1 dtt (W x[o(WTX)) x) < d Amaxx (o (c(WT)))))) x)- tw (WT [WX]] x) O(1.
d tr (Q2)
(92)
To sum up, for n and 72 + we have B oc.
C.7.2 BOUNDED BIAS WHEN 722
Since h n, the two terms in the expression of the bias can be written
Q = x x((xxx (Xx((TT)o(WW
"6(x "6(XT (W (W)))//XX)(O) (oo(xTT)oo( T))((T))) x)) x)))"xx" XT, (93)
Q2 x x(x(X (((XTT))(WT ()( 6(xTT))
(94)
Therefore we have
XTX (0 (((XTM)((TX))
() Woo (orroW/T)(((TT)))) (o(xTMO((Tx)) -1
- tr (Q1) = -2tr
d
25

---

Published as a conference paper at ICLR 2020
xx"x ir ((oXTTMO(WTX))) x))xx)))/ ( o(x WWW x)
-1
< 2Amax
+1() tr (o(W ()x x)) 6(x (W)WW)
< 0(1 Amax
x)) o(xTT) tr (K()
0(1) min (6(X ( (Kw) -00))((m0))))))) 0n O(n)
0(1),
(95)
in which used Lemma 11 and tr (AB) S Amax (A)tr (B) for positive semi-definite A, B. Simi- larly
7 1rr (rr()2)= = 21r (d(XT x)) -1 6(x W) (O( x)) - ((xT)) ( x
2
/ _WTx)
-
S n min ((XXWW) + + (2xxT) ir (WWT) 0(1,
(96)
where applied tr (AB) < (tr (ATA) tr (BT B))/2. We therefore conclude that B is bounded when h > n and h, n + co.
Remark. Concurrent to this work, Mei and Montanari (2019) provides complete characterization of the bias term and confirms our observations above.
C.8 PROOF OF THEOREM
For simplicity assume n, h be and let do d/2, no n/2 and hg h/2. Since the second layer is fixed a, Unif(-1/v/ 1/ 1/Vh), we let a, 1/vh and 0;tho
1///h for all
1 i ho. We therefore write a² h [1 1no] and W [W+, W-]:
f(x; W-, W+) = a @(W x)) = Jn1
$( x)
1 vh
(WW[.)
(97)
Vh
The empirical risk can thus be written as
L(X; W- = 1 Σ L(x; W W
Σ
y
1 T
x)) +
1 "oW )=)²
vh
vh
ØEX
xEX
(98)
C.8.1 DEFINING GRADIENT FLOWS
In this section we define three gradient flows and show that the three flows are similar some sense. GF-Original is the original gradient flow (11),
awo
Ø1
2no
2no
Vh
Vi
Vh
1 OWTT
x.)
J T(WOT ((WTTTa:) 2. ((( (,,)))
Vh
1 1
(99)
= 2no
X(y - (()
Vh
I
o (XW9)
(100)
starting from the vanishing initialization w ( NO I/dh)+s), Note that the gradient for the negative part 0 can be similarly defined.
We now define the flow under the same objective but from exact zero initialization w(((0) termed GF-Double. Due initialization, basic observation is that the solution WWP WP] is
26

---

Published as a conference paper at ICLR 2020
most rank-2, and more precisely, the parameters in the flow takes the form of WP(() (t)1
where w (t) admits the following dynamics:
2no
1
g- (2
Vholw 2T xs)) + vho(w PT mi) d'(w" DT a: )a:)
Ui
2nd
(101)
Lastly, we define the GF-Single with solution denoted w w1 (t)
Ow 01
2no
(w
w xs + (0)w"" vz) a;)0 (0)a.. (102)
2no
from initialization (0) 0.
This can be seen 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 any time Specifically, the solution can be written as:
w = (t)
1
(I-e
s'(o? XXTe
(XXT)-
(103)
2
when d < n, or otherwise
a'ion2 x x X1 (X X) X)-
w (1) (1)
2 X
(104)
when d > For simplicity elaborate the proof only for d < n. We provide 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 in Lemma 17.
|X wo(l)l =0
coooooy Booyyoogd log d' Vd
poly
Condition A: I/ O
Vol
Step 2. Bounding the Difference in Gradient Flow Trajectory. Due low rank property of GF-single and GF-double, in this subsection we slightly abuse the notation and define
s(x; w (x; w T w
1 1 Vho(wIa>) vho((wIx)
(105)
We show that the difference between the two trajectories defined above is asymptotically vanishing for wt (w_ follows the same argument). Compare the two trajectories up to time T
(T - F (T) - 1//" of (l()((()) (f
1- T
<
g ( (t)
(t) - (w) dt
+
((())(()) g (w dt
(00) (t)) +,00((2))))000 (t)) (0)0))) m dt
E+
2no
=0(1) (' ([w - w (t) 1, + - -w
dt + ++.,
(106)
where we have defined the error term as
E+ L' (s) of (w) de
(107)
27

---

Published as a conference paper at ICLR 2020
To bound the error term, note that at t = 0 Condition A holds for GF-double. Assume that for some 0 I S T. Condition A also holds for GF-double, we have
% gh (l(w(()) - 9 (w)(() dl
E+- =
V/ Vh PT x.) - - wholw æ.)) (d'(woTTx:) x;) a;) xi dt
1=1
o( PT x.) x.)+ + PT x0. oo)) -010) o' d'(0)x. -
dt
1
X y (()) (d(("" di
nvh
(0(w0 PT -00
0w (x:) dt
+
n2
(
1
/''
||x||2 1epllly'l PO 10 w wfill12 di + O(1 L-² |1x1|2 vn xx dl
n/h.
1 nvh olVā)O((og ())(1) dt J/ I' olvajo ( log d
/'
T
e dt= (l0g d
d T.
(108)
where (i) follows from the error of Taylor expansion on &i and (ii) from Condition A. Therefore, by Equation (106) and Gronwall's inequality,
[ (T)())l2 SCI log" h e C2T = =o c
poly log h (poly
-
(109)
h
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 w (0) (0) O, and that Condition A holds for w for all T > 0. Therefore, by 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 Bounding the Difference in Risk.
Now we consider the difference of the population risk R and RD of two models with parameters wi. and wh.
12 - R²) = Eu x T B f(z; w)))² En (x: T - f(x; w w 2
Ex ( wi) f (x; w 2, x f(x; wi) ) + B f(x w!)
NEW ( x) T x) bow 2T 2T
Eun ||BT x f(; s(:; 2)12
(61)
hE hEE ol)
x) p x
Bull3 x [(x; (xx 13 x f(x; 2
(384
$ 2 hLip(0) Ex m x - w x + ST x - D DT x|?| VRS + RDD
poly log h vh.
<o(Vh) /. 211 - w D
(110)
where (i) is due to Cauchy- Schwarz inequality on norm, (ii) from Jenson's 's inequality on squares and Young's inequality, and (iii) from the Lipschitz assumption the activation, and the observation
28

---

Published as a conference paper at ICLR 2020
that both R and RD are finite for " * 1 due the justified Condition A above. Therefore the difference between R and RD vanishes for = (log log h).
Step 4. Bounding the Difference from Stationarity We compute the difference in risk between the model at some finite time and the stationary point i.e. = 00,
"'o xx ((xx)) X3
(t) (t)
R
(o=)) & Cvh. (wf w Cvh
2/h
=0
2 xx (XX xx 3 Cexp (-o(0)²
- xxx
1) [131|2 - -se-Cst
Cst
(111)
for constants C3, C4 > 0. Combining (110) and (111) yields for T = log log h,
R² (T) R2 (oo) S | (T) - (T)(T)) + 12 (T) (xx))
poly log h
= 0
+ +
- 0.
(112)
Vh
poly log h
The result in (112) for case d > n follows a similar proof and 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 (0) ~ //hhI/dhh"" v.s. zero initialization w (0) = Note that at initialization |yy f(X)]l, = O(yr), and gradient flow decreases the empirical risk; therefore the condition for Lemma 18 is satisfied, and by the Lipschitz condition on the empirical gradient have
(T) - ((T))-P(TT) WD(())]
/-'
T a (
- de
=1 10)0
ōt
=||W"(0)|| /' || (W)))()))]1) OL(X;WO(B) aL(X;WP(I)
dp
OW
OW
<||wo /' -
(000-PPOU0 WPIl + OL(X:WO(1) (t) OL(X:WD(1)) (t))
2
OW
OW
< < + L' (t) ( dt,
(113)
And hence by Gronwall's lemma one obtains:
(DTT)-P((T)]+ s || e 04400))) (4
(114)
Therefore the difference in the function output can be bounded as
[f(z, (()) - f(x (T(T)) = w - o(x 6 (
< |o( W T o(x w |all] L. |0l|
w (
=0(Vd)- || (T) O(d 0 (T
(115)
Taking T log log h. together with the same argument in Step 1 yields
poly log
R =
+ 0.
(116)
d'/2
C.8.4 PUTTING THINGS TOGETHER
By (112) and (116), we know that for log log h,
Rº (T)))) Roo)] < RºO RD( IR (T) (R)(T)
R$ (20)) - 0.
(117)
29

---

Published as a conference paper at ICLR 2020
Finally the proof is completed by observing that R (po) is the risk of the minimum- -orm solution the input discussed in Section 3. In addition, note that at T = log log h, from (109)(114) one obtains that [t - W &t) F O, and therefore by Lemma 18 we have [|ā[(X; (t))))/W F + 0.
Le. the flow on the original objective reaches ((1) first-order stationary point.
C.9 PROOF OF THEOREM 8
Denote w = vec(W) = veo([W. W- and wi = (eiiii)))) Define
(f(X:w(I) (f(Xx
(t) (w)(t
(118)
(w)
which is the kernel matrix of the neural tangent kernel Jacot et al. (2018); Du et al. (2018). In the following sections 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.
C.9.1 THE KERNEL LINEARIZATION
Write (t) = (X,( e R" and its evolution:
Ny = k()) NNN((t) dt,
(119)
and the corresponding linearized flow:
dYNTK (t) K(() (t)x((t) dt,
(120)
Previous works (c.g. Du et al. (2018); Oymak and Soltanolkotabi (2019)) have proved (non- asymptotically) that the two trajectories (119) and (120) are close if the model is overparameterized. iee. h = poly(n), (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 al., 2018, Theo. 3.2) by first identifying the condition under which training converges at linear rate:
Condition B: Amin (K())) O(d).
From Corollary 15 we know that at initialization the lowest eigenvalue of the NTK matrix satisfies
Amin (K (0)) (dd), and thus the kernel regression on the NTK enjoys linear convergence. By
Lemma 19, we know that for [tW - ( O(d' the order of Amin (ain(K(P)) (K(t)) = O(d) re-
mains unchanged. Therefore, if we assume that up time T the weights satisfy || [w, (t) w,(l))l
Od-1/)), then Condition satisfied for nNN and from (Chizat and Bach, 2018b, Lemma BI) we have the following linear convergence
|1 (1) Ull, C1 (0) Ull2 C22e
(121)
Since N (0) - 2|1, m at initialization, setting T = Olog d) ensures that the training loss nn N (T) 2112 - 0 as n + xc. Consequently it is easy to check that (T)
OW(T)
+
2
and thus at time T the gradient flow reaches an o(1) first order stationary point.
We now verify that each weight vector un travels at most (dd11/2 from initialization. The norm of gradient for un; can be bounded as:
OX( (t))
(NN(t)) &' (x (T.((t))
(t)
=
72
X
Vh
ly
(8)
1
00())
|
ly UNN (t)l2 S 00 (NN(2) (122)
di.5
where (i) follows from the boundedness of ø'. Integrating the gradient yields (1) - w. w()ll
(d-1//2), i.e. [[(()) W(O))l2 0(1). Thus the distance traveled by W indeed satisfies the norm
30

---

Published as a conference paper at ICLR 2020
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 with | = O(vd), the prediction of the neural network is uniformly close the prediction of the prediction of the kernel model, 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
d (N 1) 1 u ( (NN(()); (NTK( NNTT (a, 1) - YNTK (1)),
n
where uNN( (xr af(xw(t)) (f(m((()
Owi() E Rh and uNTk(ar, 1) similarly defined on the initial- ized weights w(0). We bound the difference between the predictions on ø up terminal time" as
ow(!)
()) [NTK( =1 [ - d - (NTK((, del
d
1 I'" (uNN((, -xxx(() - (NT (o-K)NT((()] del
12
1 n
S'" uNTK(EE (UNN(t)) (UNN(() (N) dt
|
1 L// (uNN( (NTK((, ()) (t) dtl
n
( ('" [l (NNK dt
1
max (2, (xTx(2 (' 11 (ly))Na(t)lla dt.
(123)
n KTT
For the first term have
1 (' |K(()(y VNN( - K(0) UNTK (1)ll2 dt
[T) - UNTK (T)Il2 I n
<- max [K(() - K(0)ll2 / 1y UNN(t)]2 dt 1 |I((0)l22 L'" [l (tlxx de (i) 2 10 + 0(1) L' [l - (1)x dt S O(d
NOCKET
(124)
where (i) is due to Corollary 15, Lemma 19 for some e' > 0 and the linear convergence of yN and (ii) is due to Gronwall's inequality. Note that the log factor is omitted. Similarly, for the second term we have
(0) I max (ū) uNTK (āe, 1)12 I' 11 11-- de S 0 = o(d n
n KT
(125)
where we used Lemma 19 and the linear convergence of yNN in (i). Combining the two cases yields
( fNTK (#, t)l < |INTK (#t, t)ll, Old O 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 | O()R)) Therefore, in the following subsection we compute the risk of the linearized (kernel) model SNTK-
C.9.2 COMPUTING THE KERNEL RISK
Given input X E [dx= and label y ,37 X + gradient flow on the tangent kernel solves the following equation of the parameters wi
y = f(;:)) Wf(X:wo) (w - wp).
(127)
dw
31

---

Published as a conference paper at ICLR 2020
where (f(X:wo) wo) is dh matrix with column (x;; wo)/ Note that for - and Y, 72 E (0, 00), dh > n trivially holds, and by Corollary 15 we know that solution is given by
of(X:wo) @f((X:w) Of(X: dw
w1 = wa +
(xT
3 E)..
(128)
dw
dw
And the population risk can be written as (note that there is factor of 2 due to the "doubling trick" at initialization to ensure fumll (-) 0):
2R = Eo,,
(x: T [(x;w1))
0f(x; wo)
Eon
2 T B
(wi wo)
dw
27
-1
0f(x:wo) Ow
@f(X:wo) dw
@ff(X:wo) dw
Of(X:ww) Ow
E,,,e
x B
(x B + E)
27
x
B ú T kx x xT3
Eun
k
0²
(129)
2B
2V
where decomposition is made here, and for simplicity define
i @f(X:wo) af(x:wo)
Of(X:wo)
T Of((:oo)
Kx
(130)
Ow
Ow
Ow
Ow
C.9.3 APPROXIMATING THE KERNEL MATRIX
In this section we drop the negligible e the initialization. Following Cheng and Singer (2013) utilize the orthonormal decomposition of o' 6'(x) in (R, Denote bo [[)) and
0f E/ [G)) - by We have the orthogonal decomposition of
o'(x) = bo + &1 (z),
(131)
where E[ø' (G) 0. We develop the following lemmas to approximate the kernel matrix.
Lemma 13 (Approximation of (Kx)is). There exist constants e, d' > 0 such that for i + j with probability 1 enn? have
1 0 (f
wf( dw
1 2 2j <
(132)
d
dw
d
and with probability e'ne²
(f((:w))
T Of(x;ww))
- (69 + b7) <s.
(133)
dw
Ow
Proof. When i * (i.e. Equation (132)), we have
10f(z;w))) a Ow
- 1x = d
Ow
af(x;ww)) 'f(lz;:wo) Σ x. ag (w) (n:))((w)) my))
Own
Oer
dh
d a ar, (0(ww" (0'(waa))'' ax))((w" a;)) d H(x,,
(134)
32

---

Published as a conference paper at ICLR 2020
The matrix (z,, 2,) x; F 6 ( ),)o(( x;)] 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 Ø3' (xx) (A3), by Hoeffding's inequality we have
T (x():w))) @f(=zww) d
d (((x 23) < a 2FE > 1- -
Pr
(135)
dw
dw
In addition, by the concentration of a. 2j and || i.. Pro/aa//d> < 1 e cade? and Pr lor æ;/d - 1 <
cadee the orthonormal decomposition (xxx box + o1 (x) leads
to the following linear approximation of the matrix H
2 (2, 2,)
by j x, O( 2z/d/²)
(136)
and by taking = xx 2z/d under the joint event we can show that
θ ff Of(
baja Dj <
(137)
dw
dw
with probability 1-e eode
ode?
The same argument follows for the case where = j.
Corollary 14 (Approximation of ú.. For large enough I O, with probability 1 - de eoloo
1 d
1 1ff wo)
Of(X:wo)
log's d
|| u/12
d i
(138)
d
dw
dw
where u =bxx X.
Proof. Taking & = log' d/d together with Lemma 13 yields the desired result.
Corollary 15 (Approximation of Rx).
With probability 1 de - log'
1 à 1x Kx
1 Of(X;wo) T Of(X:wo) 1 d Kx log' d,
(139)
F
d
dw
dw
where Kx = b2X X + bRdI.
Proof. Also by directly applying Lemma 13.
Remark. For initialization larger or equal to w0,()) ~ (0, Ia/d), the above approximation does not depend on the scale of initialization.
Soft Remark. For d(x) Plus(z.) 63 1/4, b 0.043379. For ((x) sigmoid(x)) (1
--=)- 6g = 0.042692, 6 = 0.002144. Note that b, > 0 for all smooth activations Φ. and the equality holds (i.e. bj 0) if and only if á is linear. We comment that smaller by entails larger variance as 71 + 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)
2B Exg (x (x T xTB) }
(140)
We first bound the error in substituting á with iz:
á x x B á
K
x
< û - u|ا2 xxx 11/ |x11 11311
log d Vd
0 log' d. d vd. vd.1) 1) -
(141)
33

---

Published as a conference paper at ICLR 2020
where (i) is due to the fact that | 112 = Amia = O(1/d) and |X|| = O(vd). Therefore
we have as n, d, + 20
2B Es ( 3 á R x 8) })
x: 3 ú kikx x ( })
2)
- Ex
T B 1y x x 3 2
)..
Ex
x
(142)
By taking expectation over æ and the rotational invariance argument similar to Hastie et al. (2019),
Ex
xx B
bij
x x B)
= =Bx B (r (1-66XRx |]
BB 3 tr ((I- )1
d
(143)
In addition, we bound the error in substituting Kx by Kx defined in (139):
2t1 (xRxIX= x x
1
- (
x
=
j
1
x
x
Ax 2
Ax Rx
1xx
2
llo''
0
d. d d d. d-1)
(144)
and similarly,
tr [K X x X -
d
tr X K X + kx'1)
|| x 1|2
xx* |g| hx Kx || {xx || | 1xTx112 Rxx
Rz'
=0 a d- log' d. d= d d=1) = llog d
(145)
Combining these two formulas in (143) yields
x T - 1ij T XKx xxxxxxTpx xT)
2B + Eg
3 B
((I- 12 x (I (
B B
Ir
13 x -
(146)
d
Utilizing the Marčenko- Pastur law from Section B.2 we obtain
1-1
71 + mm m - 2) +
B B
(147)
271
2 21 + m
+ 2) 2) 2) +
where m = 2b.
C.9.5 THE VARIANCE TERM
Similarly, for the variance we utilize the approximation
2V
= Em
R
0²
(148)
x
x
34

---

Published as a conference paper at ICLR 2020
Specifically, we bound the approximation enror
ük x xx it
uky
Kx a
- | u11
1kxx |
| + u
K
=0
log' -0(10g'' d ā d 2/) = 0
log' d
(149)
d
and similarly
Ex
ük x Kx 4
-ūk x K
=tr [xx] kxx)
(Rx! + kxx
Egg uuu
-tr [Kx² [xx {x)
Rx²
K
x"x)
[ x
[g
[xx= Kx ||, K
Rx Kx 1 ||xTx|12
2
(log' log'
=0 dlog'
d. ddoog
d= 1, d- 1. d) 0
(150)
By combining the two approximations above we know that n, d, p -
log' d d
2V Ex
x
ú
0² =
- 0.
(151)
Therefore the variance is given
2V + R
it X R x i
= o L
life
X (b T X + bldll) -2 xxx
XXX (2xTx+ b6 b52b91
tr d
1
71 + mmm + 2/7(111 + m(n (m + 2) + 2) 2) +
= o²
(152)
where is defined in the derivation of the bias term.
C.9.6 PUTTING THINGS TOGETHER
Recall the population risk the sum of the bias and variance
R 2
11-1 11
+
71 + 71m m - 2) +
2711
2 71 m(71 (m + 2) 2) 2) +
71 + 7/m +1
+ 0
+ 4/n(m + mmm (m + 2) + 2) - +
(153)
Observe that the population risk is independent of 72, iee. double descent does not occur when the network is overparameterized via changing the width. In addition, the bias monotonically increasing and upper-bounded by the null risk 7 and lower-bounded by the bias of the least squares solution on the input features 3 Xty, y. whereas the variance remains bounded for all 71 E (0, oo) as long as m > 0, i.e. é is nonlinear.
35

---

Published as a conference paper at ICLR 2020
D USEFUL LEMMAS
Lemma 16. Given Kw from (46), define
Kw = rT + 81 + IQ,
(154)
where Q E Raxch with Qifs = w| w and Qii = 0, and E[o(G) - B[o(G)]
E[o(G)) B[G Then n, d, + 0c, Kw Kw
I log" d a.s..
Proof. Consider the event where A, = (11]/ - < k,, w)l < e). Under event Ac, for the
diagonal term the kernel matrix we have
[Kwl [Kw]a| = Eo ( x) x) [
|B(((| w0 [ O(s)
(155)
And for off-diagonal term, by the decomposition introduced in Section B.3 we have
[Kwly- = Ex ((w] () x)] = ||w,112 ||,, + Ex L (w) (w) x)],
(156)
hence [[Kw]y - [ellll (w: wy)2. Notice that A, holds a.s. for log" d/ vd and large enough 0; therefore have Kw Kw log d.
Lemma 17. Lei B(1) be the solution to the gradient flow time defined in (103)(104). Then as n, d + 20 and 7 # the following holds.:
[1b(t)]12 = Op((); ||x B(t)ll 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 exp( 1 _XXT)) -XXT)- From (Hastie et al., 2019, Corollary 1) we know that [13((x) = [| - O(1) for t 2 0; it follows that [|3(t)]2 0
0(1) for 71 < 1. Note that exp(=x 112
For the second part, we utilize the SVD U.V where [Σ: E E dxd and Σn A We have - (I-exp(-3 (xx(-=xxx]] xXX = U where ΣS, = - exp( tx? /n) if < and 0 otherwise.
x - (XXT)-1 -'y
|X [(t)ll
Xy
VETUTUSU US 2 UEVT| ||3///0
(1XT B1 + +|1ello))
S|I To mン Eン 32
|| ao (I|x
Bl [|=|loo) Op 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 (n))., (n). and thus for any unit vector z independent to V. Vz is uniform on sphere
and | V 1/21.0 Ollog d/ Vd). We wherefore have
|V O
log dl Vd
|||||2
V sup
O d),
(158)
|
Note that this result also implics that 1/l - X (Tt((t) loo
Op(poly log d).
36

---

Published as a conference paper at ICLR 2020
Lemma 18. For weight matrices W. W satisfying 11y f(X)]l2 O(vn), where f(X) 6(XW)a a with fixed a, ~ Unif(-1/Vh,, 1/Vhj, given (A1)-(A3), the gradient of the empirical risk defined in (11) is Lipschitz w.E.. W in the Frobenius norm, i.e.
OL(X;W) ( oW
S WW
(159)
OW
Proof. Denote y1 <(x Wi)a and y/ - o(X Waja for W1,W2 satisfying the assumption above (which can be seen as condition on the magnitude training loss), we have
AL(Wi) 2
OW1
OW2
1 x (1 "1)a" y1)a ('(xTWi)) _x [yy y2)a o o0(X of( Wa)]
|1l (1y - vi)a T c Wi) - -y- - yz)a T o (xTW)))lF
n
50
[[022 vija o o'( Will
0
(V U2)a T d
(160)
We upper bound the two terms separately:
(4)
(12 vija T o ( (PTTTiW)ls v max (xx 1122 - yı 1111221aa1
S 0(1) | (XTWW) wi)a 6'x (0TxW)))al
W2)al1
(18)
- 01) || W W2ll = O(Va) [[1 W2llF (161)
where we applied the inequality ||A BlI I max (|4:,]) | in (i), boundedness of ø' in (ii) and Lipschitzity of & in (ii). Similarly, for the second term
( v2)a (0'(X - o'(X (x(x Wa))ll
< <max((ai) by - y2 |X d' (x Wa)l
19
_0(1) || ||W, = O(Vd) || W211F
(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 ([w,(t) (,)ll] = (dd11/2) for all i, then we have [t((t) K ((0)l22 = O(d-/2-c)) for some positive O
Proof. Recall the definition of the NTK:
(1)(1) af (()) af(x;;;(()) x. a.) @w)() aw(()
&' (((ww(t)) (o.)T=)))) (=))((unx())) Tny), (163)
or equivalently the matrix form
K(t)= = X X ( ((1))'(W(1)) x)).
(164)
At initialization, æ; NO,, Ta) and unx(0) N(O, d Ta). Thus for fixed æ;, by Gaussian
anti-concentratton have Pr T wxl < log d < (11/d//+a)) for some e 0. In addi- tion, note that || (1) = (0)l22 0(2-1/2) for all k and therefore for i,j, k: such that
37

---

Published as a conference paper at ICLR 2020
| wx((0] (log d) and [z] w,( (log d)) we know that / (z) wx(t)))'(zz) (a wr (0)) (a] ur (x)) (
(x] wx(())
Given fixed 2, define = 1 [ wkl < log df as the indicator variable that the k-th neuron does not saturate. We know that E[ya = (1/dd and Var[ys] = ElyEl Elya] = (1/d//2+e1) By Bernstein's inequality
>es eexp (-
hs²
Pr
Σ yu
E[ya]
(165)
20 + 25/3
Setting 8
elog h Teg know that with probability least h-c
h
(poly log
Uk Se =
(166)
h1/2+cs
Therefore, given æ; and zy, for large enough with probability at least 1 - h- we have
Σ (z)) 22)
Σ ('(wx xx a.)o'(wx(00) (wx( x;) 0(hh/2-44, (167)
in which we utilized the boundedness of Φ' Taking union bound over dl elements in the random feature matrix yields
(tt) K(0)l
1 - X X T
I||||
X x11 max (W0
= (22
(168)
Using the exact same argument, one derive that [ uNTx(#)]l2 O(d2/2-e'), the proof of which we omit.
E ADDITIONAL RESULTS
E. RISK OF ReLU NETWORK UNDER SYMMETRIC DATA
If the dataset is symmetric, that is
(A5) Symmetric Data:
Vi E [L, Blj E [1,n] s.t. 2, + + = =0,
then population risk the gradient flow solution can be given explicitly for certain nonlinearities:
Proposition 20. Given (AI-3)AA5), if the nonlinearity satisfies d'(x) + o'(-2) C for constant then as n, d, + 30
2m =;; R (1 (1- 2
R R + 1- 2n1
2 + 2711
(169)
Note that the requirement the nonlinearity holds for ReLU and Soft.Plus. This expression is again independent to "/2 and aligns with the experimental results in Figure (we only plot the bias component for verification). In addition, the bias is upper-bounded by the null risk for all 71- 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.
38

---

Published as a conference paper at ICLR 2020
0.8
1.0
0.8
risk
risk
0.6
0.4
i..d data (empirical) symmetric (empirical)
L..d. (empirical)
0.2
symmetric init (empirical)
symmetric data (prediction)
symmetric init (prediction)
0.0
0.0
0.5
1.0
1.5 = din
2.0
2.5
3.0
0.5
1.0
1.5
2.5
= din
(a) bias of vanishing initialization.
(b) bias of non- 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 different "/2 (from to 2) which 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.
Proof. Without loss of generality assume X [Xo, -Xo]. Then by (100) we have
Ow++ 1 Öt
2no
-vi hoo(w Ix,) x;) + hoo( (oo(wJan xx)) o' (w (w.[x;)xz
(170)
2no
and the flow for w- follows from symmetry. In this case one can show that from exact zero initialization, for nonlinearity satisfying o(x) o(-x) z, such ReLU and SoftPlus.
((w O(w-) 0t 0t
Σ
2no
(u hoo + hoo(wIa.) (6'(wIIx) o' (wIx;))<s =
2no
(171)
And therefore the gradient flow of u. is
ow+
2no
(vi /oo(w]x + Aoo(-[[x.) 0 (2(w) xx;)
OL
2n10
i=1
na
(v: ho(w]ar.) +
(0' (w) x.) + (-w]Ix,) )), a;
2no
i=1
ni
1
(v: how Joo
Xoyo
No w+.
(172)
2no
2no
2no
The flow w_ follows from symmetry. Solving for the stationary points (i.e. gradient becomes zero), it the clear that
71 <<0.0. 0.5, (XXT)-IXX)
(t=>))
(t=oo)
(173)
X(X -1 y, m > 0.5. h10
And hence the asymptotic risk
2 a2:
<0.50
<0.5)
R(
2 +
(174)
2n1
2
27
The same conclusion holds for vanishing initialization we assume that the trajectory stays close to that of exact zero initialization. Note that although the prediction aligns well with the experimental results, argument in Theorem does not directly apply due to the undefined derivative of ReLU at the origin, and thus this result is not rigorously justified.
39

---

Published as a conference paper at ICLR 2020
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 71, 72 from 0.1 to 3. The linear teacher model F(a:) x: fixed as = - vd. For each (71; 72) we average across 50 random draws of data.
Optimizing the First Layer. For both initializations, we use gradient descent with small step size (1) = 0.1) and train the model for minimally 25000 steps and till V f(X, W)ll < 10 We fix 320 and vary 71: from 0.1 to 3 with the same linear teacher model / = -11/ /d. The risk is averaged across 20 models trained from different initializations.
40


Reviews:
[{"review_id": "ryxeE1cjYH", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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$.\n\nThis paper considers the following scenarios of training described below, where the data is generated from a linear model on Gaussian inputs and with a zero-mean noise. The emphasis of the results is on understanding when a \"double descent\" type phenomenon occurs (\"Double descent\" is a recently coined phenomenon in literature where the risk, as a function of the \"complexity of the model\", initially has a classical U-shape behavior, but eventually decreases again once the complexity of the model exceeds the number of training points.)\n\n1. Training only the second layer: The risk is first decomposed into a bias and a variance term. An exact bound on the variance term of the risk is obtained. While the exact nature of the bound is rather complex to parse, the takeaway is that a double descent phenomenon is observed in terms of $\\gamma_2$, namely, the risk blows up when $h \\approx n$, but decreases as $h$ is increased beyond $n$.\n\n2. Training only the first layer: Two different regimes are considered here, depending on the scale of initialization, called \"vanishing\" and \"non-vanishing\" initializations. In both regimes, the risk is independent of $\\gamma_2$, that is, the risk does not depend on number of hidden units (although the risk bounds are different and there is an additional assumption in the case of non-vanishing initialization to ensure that the initialized network computes the zero function). In other words, a \"double descent\" phenomenon is not observed in this setting.\n\nRecommendation:\nI recommend \"weak acceptance\". The paper extends prior works that obtain asymptotic risk bounds on linear models to the setting of two-layer neural networks (where only one layer is trained).  However, I am unable to assess the technical novelty of this work as it seems to heavily rely on prior work which in turn use techniques from random matrix theory.\n\nTechnical Comments:\n- I felt that while it is valuable to have exact bounds on the risk, the form of the bounds are quite complex and hard to parse (especially in Thm 4, case of training only the second layer). Moreover, these bounds are just in the case where the teacher model is linear and while it is claimed that this could be relaxed to a more general class of functions, the specific bounds might change drastically. So any insights on the nature of these bounds will be valuable, especially with some comments on how these bounds change if the teacher model is itself realized as a 2-layer neural network.\n- The parameter count of a 2-layer network with $h$ hidden units and input dimension $d$ is $O(dh)$. So perhaps it makes sense to study an asymptotic regime where $dh/n$ approaches $\\gamma$, instead of both d and h growing linearly in n. While this issue is hinted at in the discussion section, I don't understand the statement \"the mechanism that provably gives rise to double descent from previous works Hastie et al. (2019); Belkin et al. (2019) might not translate to optimizing two-layer neural networks.\"\n- Another future direction that could be included in discussions is the setting where both layers are trained simultaneously.", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I did not assess the derivations or theory.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper."}, "extra_text": {}}, {"review_id": "rylsXynatr", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nSummary: The work shows that in two layered neural networks with non-linearity\n1) 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.\n2) the bias-variance decomposition does not exhibit double descent when optimizing only the first layer with both vanishing and non-vanishing initialization of weights.\n3) 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.\n4) 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.\n\nI believe this is an interesting work that needs to be accepted.", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I did not assess the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "Bke6usYzKr", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors study the generalization error of two-layer neural nets, where an asymptotic point of view is taken. Their main results can be summarized as follows.\n1. If only the second layer is optimized, they observe the double-descent phenomenon.\n2. However, if only the first layer is optimized, the double-descent is not observed.\nThis shows that recent results for certain linear models (e.g. Song, Montanari 2019) do not directly transfer to neural networks. As the authors point out, however, if a different scaling is used in the asymptotics, double descent might still be observed.\n\nI see the following strengths of the paper. \n-This is a very well-written paper with a clear message.\n-The result is important and gives new insights into the generalization properties of neural networks.\n\nIn my view, this is an interesting contribution, which should be accepted. \n\n---------\n\nThank you for your response. I will leave the rating unchanged.", "extra_scores": {"experience_assessment": "I do not know much about this area.", "review_assessment:_checking_correctness_of_experiments": "I did not assess the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I did not assess the derivations or theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper focuses on studying the double descent phenomenon in a one layer neural network training in an asymptotic regime where various dimensions go to infinity together with fixed ratios. The authors provide precise asymptotic characterization of the risk and use it to study various phenomena. In particular they characterize the role of various scales of the initialization and their effects. The reviewers all agree that this is an interesting paper with nice contributions. I concur with this assessment.  I think this is a solid paper with very precise and concise theory. I recommend acceptance.
Decision: Accept (Spotlight)