G ENERALIZATION OF T WO - LAYER N EURAL N ETWORKS : A N A SYMPTOTIC V IEWPOINT
A BSTRACT
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 inﬁnity at the same rate.
Speciﬁcally, we derive the exact population risk of the unregularized least squares regression problem with two-layer neural networks when either the ﬁrst or the second layer is trained using a gradient ﬂow under different initialization setups.
When only the second layer coefﬁcients are optimized, we recover the double descent phenomenon: a cusp in the population risk appears at h ≈ n and further overparameterization decreases the risk.
In contrast, when the ﬁrst layer weights are optimized, we highlight how different scales of initialization lead to different inductive bias, and show that the resulting risk is independent of overparameterization.
Our theoretical and experimental results suggest that previously studied model setups that provably give rise to double descent might not translate to optimizing two-layer neural networks.
I NTRODUCTION
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 overﬁtting even in the absence of explicit regularization (Zhang et al., 2016).
This phenomenon is usually explained by the interplay between the model architecture and the optimization method.
Existing works have analyzed the implicit regularization of gradient descent on simple models (Gunasekar et al., 2018;
Ji and Telgarsky, 2018), and provided generalization guarantees (Arora et al., 2018; Bartlett et al., 2017; Dziugaite and Roy, 2017) that align with the empirical observations.
Recently, a series of works highlighted the implicit regularization of interpolators in the overparameterized regime (Belkin et al., 2018; Spigler et al., 2018; Geiger et al., 2018; Advani and Saxe, 2017).
Speciﬁcally, a second decrease in the population risk is observed when the model is further overparameterized beyond the interpolation limit, i.e. when the model achieves zero training error.
This phenomenon is known as double descent , and can be precisely quantiﬁed 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.
and Mei and Montanari (2019) explicitly derived the population risk of linear regression and random features regression models in high dimensions using tools from random matrix theory.
However, there is still a gap between the practical beneﬁt of overparameterization and the recently proved double descent phenomenon, which is typically established under models that exhibits the following structure: the trained model solves a linear inverse problem, and the “cusp” in the risk arises from the instability of the inverse at the interpolation threshold.
Moreover, given a dataset or ﬁxed n, d , the number of parameters in the linear regression model is also ﬁxed, i.e. the level of overparameterization cannot be altered.
It is therefore unclear if the trend persists in the optimization of more complex models, for instance in two-layer neural networks where overparameterization can be controlled simply by adding more neurons.
In this work, we analyze the generalization properties of two-layer neural networks in the unregularized least squares regression setting and examine the presence/absence of the double descent phenomenon.
We consider the proportional asymptotic limit where the number of samples n , input features d , and neurons h tend to inﬁnity at the same rate, under which overparameterization corresponds to increasing the limit of h/n (network “width”).
This regime is particularly interesting because even though n → ∞ , the empirical risk is not equivalent to the population risk.
In addition, the joint scaling of n, d, h is parallel to the practical choice of model architectures, where it is common to train a larger network when the number of samples and input features are larger.
Following Hastie et al.
(2019), we assume unit Gaussian input and noisy linear observations, and analytically derive the population risk of the solution of gradient ﬂow on either the ﬁrst or the second layer parameters when the ﬂow is initialized close to zero.
[Figure]
Our ﬁndings can be summarized as follows (see Figure 1):
• When only the second layer is optimized, we derive the risk in its bias-variance decomposition and demonstrate the presence of the double descent phenomenon.
• When the ﬁrst layer is optimized, we compare two solutions of gradient ﬂow from different scales of initialization, which we term as vanishing and non-vanishing initialization, and show in both cases the population risk is independent to overparameterization.
• For the vanishing initialization, we show that the risk of the gradient ﬂow solution is asymptotically close to that of a rank-1 model.
For non-vanishing initialization, we show that the gradient ﬂow solution is wellapproximated by a
kernel model and derive the risk.
1. R ELATED W ORKS
Global Convergence of Two-layer Networks.
A plethora of recent works have explored the global convergence of shallow neural networks.
Mei et al.
(2018; 2019); Chizat and Bach (2018a); Rotskoff and Vanden-Eijnden (2018); Sirignano and Spiliopoulos (2018); Nitanda and Suzuki (2017) studied the mean-ﬁeld limit where the number of neurons h → ∞ and the second layer scaled by 1 /h , and established correspondence between the main-particle limit of gradient descent and Wasserstein gradient ﬂow to demonsrate global convergence.
On the other hand, Jacot et al.
Du et al.
(2018); Oymak and Soltanolkotabi (2019); Allen-Zhu et al.
(2018b)
; Song and Yang (2019) considered a different scaling and showed that gradient descent on overparameterized models converges to global minimizer at a linear rate; key to these results is an observation that optimization via gradient descent is asymptotically equivalent to kernel regression with respect to the neural tangent kernel .
Active vs. Lazy Training.
Following Chizat and Bach (2018b), we refer to the two aforementioned scalings as the active and lazy (kernel) regime.
It has been observed that different regimes lead to contrasting inductive biases.
Williams et al.
; Woodworth et al.
; Li et al.
(2017) showed that for certain two-layer network or overparameterized linear model, the scale of initialization controls the implicit regularization of gradient descent (from sparse to smooth solution).
In the student-teacher setup (Tian, 2017; Zhong et al., 2017), Ghorbani et al.
(2019b;a) showed that kernel models in high dimensions perform no better than low-degree polynomials on the input or fullytrained two-layer network.
Additionally, Suzuki (2018); Allen-Zhu and Li (2019)
; Yehudai and
Shamir (2019); Wei et al.
(2018) demonstrated that neural network outperforms linear estimators (including kernel method) in learning various target functions.
The difference between ﬁxed bases and adaptive bases mirrors the difference in optimizing the ﬁrst or second layer in our setup.
Generalization of Overparameterized Models.
It is often observed that overparameterization does not result in overﬁtting (Neyshabur et al., 2014).
In the lazy regime, generalization guarantees can be derived from the distance traveled by the parameters (Neyshabur et al., 2018; Nagarajan and
Kolter, 2019), which becomes small if the model is sufﬁciently overparameterized (Arora et al., 2019b;
Li and Liang, 2018; Allen-Zhu et al., 2018a; Cao and Gu, 2019).
Compared to these guarantees that usually require signiﬁcant overparameterization, our result relies on stronger data assumptions, but consequently we obtain the exact population risk instead of a vacuous upper-bound.
Beyond the kernel regime, Advani and Saxe (2017); Goldt et al.
(2019) analyzed the generalization dynamics of overparameterized models in the student-teacher setup.
Double Descent.
The term double descent refers to the phenomenon that the population risk of an empirical risk minimizer manifests a "cusp" at the interpolation threshold, and further overparameterization decreases the risk.
First observed in Krogh and Hertz (1992), the phenomenon has been recently connected to the beneﬁt of overparameterization (Belkin et al., 2018; Geiger et al., 2018;
Spigler et al., 2018; Advani and Saxe, 2017), and can be precisely characterized for certain simple models (Hastie et al., 2019; Belkin et al., 2019; Bartlett et al., 2019; Xu and Hsu, 2019).
Our work is inspired by Hastie et al.
(2019) which uses random matrix theory to derive the asymptotic risk for linear and random feature models.
Concurrent to our work, Mei and Montanari (2019) analyzed the random features model and derived its population risk for which double descent occurs both in bias and variance.
This aligns with our results on optimizing the second layer in Section 4 although we do not derive the bias component explicitly.
Compared to Hastie et al.
; 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.
; 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).
P RELIMINARIES : T WO - LAYER N EURAL N ETWORK
Consider the following bias-free two-layer neural network f : R d → R with h hidden units
[Equation]
where x ∈ R
d is the input,
w i ∈ R
d is the weights corresponding to neuron i , a i ∈ R is the i -th coefﬁcient of the second layer, and φ : R → R is a Lipschitz continuous activation function with bounded Gaussian moments, i.e. E [ φ ( G ) k ] < ∞ , ∀ k ∈ Z + for G ∼ N (0 , 1) .
For concise notation, we write W =
[ w 1 , ...
w h ] ∈ R d × h for the weight matrix, a =
[ a 1 , ...a h ] ∈ R h for the coefﬁcient vector, X =
[ x 1 , ... x n ] ∈ R d × n for the data matrix, y ∈ R n for the corresponding vector of labels, and Φ = φ ( W (cid:62) X ) ∈
R h × n for the feature matrix at the ﬁrst layer.
We omit arguments of f when they are clear from the context.
We consider a student-teacher setup, in which data is generated by a teacher model F : R d → R with additive noise, and the student model aims to minimize the squared loss:
[Equation]
where E [ x i ] = 0 , Cov( x i ) =
Σ , E [ ε i ] = 0 , Var( ε i ) =
σ 2 .
We are interested in the population risk
R ( f ) =
E P x
[( F ( x ) − f ( x )) 2 ] .
Our analysis will be made under the proportional asymptotics: n, d, h
[Equation]
n
[Equation]
∈ (0 , ∞ ) , in which overparameterization corresponds to increasing γ 2 .
Thus the characteristics of double descent considered in this work are: 1) large population risk as γ 2 → 1 ; 2) decrease in the risk for γ 2 > 1 .
While the empirical risk can be minimized in various ways, we analyze the solution of gradient ﬂow, in which we update either the ﬁrst layer W or the second layer a :
[Equation]
from small initialization.
The rest of the paper is organized as follows.
In Section 3, we start with a simple example of two-layer linear network as warm-up.
In Section 4, we consider optimizing the second layer coefﬁcients (ﬂow over a ) of a non-linear two-layer neural network under ﬁxed Gaussian ﬁrst layer, which is a random feature model.
Section 5 considers optimizing the ﬁrst layer weights (ﬂow over W ) of such network under ﬁxed Rademacher second layer.
We defer all proofs and details on experiments to appendix.
W ARM - UP : L INEAR N ETWORK
We begin with a simple linear model with φ ( x ) = x , i.e. Φ = W (cid:62) X .
We remark that although the model is linear, the solution obtained by gradient ﬂow on the two-layer model can be different than that from directly solving the linear regression problem on input features.
Training the Second Layer.
Following Hastie et al.
(2019), we ﬁx the ﬁrst layer parameters to be randomly drawn from a unit Gaussian and optimize the coefﬁcients a by minimizing (cid:13)(cid:13) a (cid:62) Φ − y (cid:13)(cid:13) 2 2 .
The following lemma characterizes the solution of the gradient ﬂow.
Lemma 1 (Least squares solution) .
Given data matrix X, response vector y and model f ( x ) = (cid:104) φ ( x (cid:62) W ) , a (cid:105) with ﬁxed ﬁrst layer coefﬁcients W , gradient ﬂow on the coefﬁcients a starting from zero initialization converges to ˆ a =
Φ † y , where † stands for the Moore-Penrose inverse.
We make two assumptions on the data and the teacher model to simplify the computation.
[Equation]
Denote the linear student network as f ( x ) = (cid:104) x , ˆ β (cid:105) , where ˆ β = W ˆ a and ˆ a is the least-square solution deﬁned by Lemma 1.
We write the population risk in its bias-variance decomposition.
[Equation] [Equation]
where (cid:107) x (cid:107) 2Σ
= x (cid:62) Σ x .
We compute the bias and the variance separately to obtain the risk.
Theorem 2.
Given (A1)(A2) and let w i
i .
i .
d . ∼ N (0 , d − 1 I d ) , at n, d, h → ∞ we have
R ( γ 1 < 1) →
[Equation]
where d/n → γ 1 , h/n → γ 2 , g 1 = | γ 1 − 1 | , and g 2 = | γ 2 − 1 | .
We observe that when d > n (i.e. γ 1 > 1 ), we obtain the double descent risk curve, i.e., the population risk achieves its maximum at γ 2 → 1 and further overparameterization ( γ 2 > 1 ) reduces both the bias and the variance.
Conversely when n > d and h > d (i.e. γ 1 < min(1 , γ 2 ) ), the population risk becomes constant and equals to that of the minimum-norm solution ˆ β min =
X † y on the input features.
Training the First Layer.
When the ﬁrst layer of a linear network is optimized via gradient ﬂow and the second layer is ﬁxed, the following holds for zero-initialization of W .
Proposition 3.
Given W (0) = 0 and ﬁxed a init (cid:54) = 0 , at any time t > 0 of the gradient ﬂow on
W , W ( t ) is rank-1.
Further, ˆ β = (cid:99) W a converges to the least squares solution of y =
X (cid:62) ˆ β , the population risk of which is given in (Hastie et al., 2019, Thm. 1 & 3)) as
[Equation]
In this case, overparameterization by increasing γ 2 does not inﬂuence the population risk.
In addition, since the obtained two-layer linear model is equivalent to the minimum-norm solution on the input features ˆ β min , optimizing the ﬁrst layer always results in smaller or equal population risk compared to optimizing the second layer.
In this simple scenario for two-layer linear networks, double descent is observed only when the second layer is optimized, which reduces the objective to least squares regression on the intermediate features.
On the other hand, training the ﬁrst layer from zero-initialization always yields the same solution that is independent to overparameterization.
One natural question to ask is: does this phenomenon generalize to nonlinear two-layer neural networks?
The following sections answer this question in the afﬁrmative under certain conditions.
[Figure]
N ONLINEAR M ODEL : O PTIMIZING THE S ECOND L AYER
In this section we analyze the case when the second layer a is learned under ﬁxed W and a nonlinear activation function φ (a random feature model).
We ﬁrst observe that by Lemma the gradient ﬂow ﬁnds the solution ˆ a = Φ † y .
We again consider the following bias-variance decomposition.
[Equation] [Equation]
We highlight that the variance term does not depend on the target function.
The following result characterizes the variance of the random feature model.
Theorem 4.
Given (A1) and w i i .
i .
d . ∼ N (0 , d − 1 I d ) , when n, d, h → ∞ , we have
[Equation]
in which m 1 ( ξ, ρ, τ ) , m 2 ( ξ, ρ, τ ) is the unique solution in {| m 1 | , | m 2 | < 1 / (cid:61) ξ } of
[Equation]
where variables ξ, ρ, τ satisﬁes (cid:61) ξ > 0 or ξ < 0 , ρ > τ > 0 , and constants c 1 , c 2 deﬁned as,
[Equation]
for G ∼ N (0 , 1) and (cid:61) ξ denoting the imaginary part of ξ .
Remark .
c 1 ≥ c 2 and the equality holds iff φ is linear.
Corollary 5.
If we let γ 1 → ∞ , the variance is equal to the lowest value of the variance of the linear model V ( γ 1 →∞ ) =
σ 2 min { γ 2 , 1 } / | 1 − γ 2 | .
The proof of Theorem 4 largely follows from Hastie et al.
with techniques similar to Cheng and Singer (2013)
, but with modiﬁcations in otder to handle unnormalized and uncentered activation functions.
The above theorem holds irrespective of the underlying teacher model, and is consistent with the double descent risk curve as it suggests that for all γ 1 , variance of the random feature model peaks at h = n then drops as γ 2 further increases.
Note that as γ 1 → ∞ , a linear and nonlinear network would have the same asymptotic variance.
Since double descent is observed in the variance term, we do not derive the bias for all γ 1 , γ 2 .
Instead, we show that for linear teacher, the bias also becomes unbounded as γ 2 → 1 .
Proposition 6.
Given (A1)(A2) and w i i .
i .
d . ∼ N (0 , I d ) , then B → ∞ as γ 2 → 1 .
Furthermore, B is ﬁnite when γ 2 > 1 .
Thus we have shown that a “cusp” in the population risk appears at h = n , which aligns with the double descent phenomenon.
Empirically as γ 1 → ∞ the nonlinear model also shares the same asymptotic bias with the linear model, as shown in Figure 2.
We note that (Mei and Montanari, 2019,
Thm. 1 & 3) analytically solved the risk of random feature model for a larger class of target functions than ours and conﬁrmed that double descent appears in both the bias and the variance.
N ONLINEAR M ODEL : O PTIMIZING THE F IRST L AYER
Having observed the double descent phenomenon in optimizing the second layer, in the sequel we consider a two-layer neural network with ﬁxed second layer coefﬁcients initialized from a Rademacher distribution a i ∼ Unif {− 1 / √ h, 1 √ h } , and the ﬁrst layer W is optimized with the following update
[Equation] [Equation]
which potentially has different stationary solutions with no explicit form, depending on the initialization.
We denote the solution of this ﬂow at time t started from designated initialization by W init ( t ) , its stationary solution by (cid:99) W , and the corresponding network by ˆ f .
Remark .
Although we let n → ∞ , this dynamics does not corresponds to the population gradient ﬂow considered in Tian (2017).
For instance when x ∼ N (0 , I/d ) , the spectrum of the data covariance is
Mar ˇ cenko–Pastur, whereas the population covariance is identity.
As we cannot characterize the gradient ﬂow solution from all possible initializations, we consider two speciﬁc scales of initialization:
[Equation]
Note that neither of the two initializations correspond to the “mean-ﬁeld” regime (e.g. analyzed in
Mei et al.
due to the 1 / √ h scaling of the second layer.
In other words, as h increases, we expect the distance traveled by each parameter to decrease under both initializations.
The difference, however, is the “relative” amount the parameters traveled compared to their initialized magnitude, which leads to solutions with contrasting properties.
As we will see, under (A1)(A2) and vanishing initialization we have (cid:107) W (0) − (cid:99) W (cid:107) F / (cid:107) W (0) (cid:107) F (cid:29) 1 , i.e. the contribution of initialization vanishes at the end of training, whereas for non-vanishing initialization the inequality is in the opposite direction, i.e. (cid:99) W “barely moves” and resembles the initialization W (0) .
5. V ANISHING I NITIALIZATION
As d, h → ∞ , the vanishing initialization becomes arbitrarily close to zero-initialization.
We thus expect the gradient ﬂow under vanishing initialization to "resemble" that of starting from exactly zero if the ﬂow converges sufﬁciently 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 φ (cid:48) (0) (cid:54) = 0; | φ (cid:48) ( ± x ) − φ (cid:48) ( ±∞ ) |
=
O ( e − x ) .
The above assumption requires that the derivative of the nonlinearity φ saturates beyond a O(1) region, which holds true for the commonly-used smooth activations such as sigmoid and SoftPlus .
In addition, the choice of scaling ensures that the gradient ﬂow converges sufﬁciently fast.
We thus have the following characterization of the population risk:
Theorem 7.
Given (A1-3).
Let T
=
O (log log h ) and ˆ f ( · ) =
f van ( · , W ( T )) , then as n, d, h → ∞ , the gradient ﬂow reaches a o (1) ﬁrst-order stationary point at time T, i.e. (cid:107) ∂W ( T )
/∂t
(cid:107) F ∈ o (1) , at which point the population risk is given as
[Equation]
The expression above is the same as the risk of the least squares solution on input ˆ β
= X † y ; therefore the risk is independent to overparameterization (increasing γ 2 ).
The intuition is that when the weights are initialized sufﬁciently small and travel inﬁnitesimally, then the activation can be linearized around 0 and thus the model is equivalent to a two-layer linear network.
Note that this result does not apply to the non-smooth ReLU activation.
Instead, in Appendix E we heuristically show that under the additional assumption that the data is symmetric, the risk of ReLU network is also independent to γ 2 .
5. N ON - VANISHING I NITIALIZATION
When initialization is sufﬁciently large, the amount each parameter travels to minimize the empirical risk becomes asymptotically negligible compared to the magnitude of initialization.
In this case we establish under (A1-3) that (11) is asymptotically equivalent to the kernel gradient ﬂow on the tangent kernel :
k ( x , y )
= (cid:104)∇ W init f ( x ) , ∇ W init f ( y ) (cid:105) .
The converged parameters under this linearized dynamics has the following closed-form:
[Equation]
[Equation]
where J ∈ R n × ( d × h ) is the Jacobian matrix w.r.t.
to the model parameters.
We remark that in contrast to most NTK-type global convergence results that require the width of the model to grow faster than the number of data points (e.g. Du et al.
, our result is not built upon the overparameterization on width, but instead an anti-concentration that relies on the scale of initialization.
Consequently the above initialization is larger than the scale that is commonly used in practice.
One may naturally expect the double descent phenomenon to appear in this kernel solution, as ∆ exhibits the form of a least squares solution which contains a pseudo-inverse.
However, we show that this is not the case under the same assumptions in Section 4; in fact, the risk is also independent to γ 2 .
An obstacle in computing the risk of the kernel model is the potentially non-zero f init ( X ) .
We thus adopt the "doubling-trick" from Chizat and Bach (2018b) to ensure f init ( · ) = 0 , i.e. we assume the following symmetric property on the initialized weights:
[Equation]
Theorem 8.
Given (A1-4) and let n, d, h → ∞ , the stationary solution ˆ f has the following risk
[Equation]
where m =
b 21 /b 20 , b 20 = E [ φ (cid:48) ( G )] 2 , and b 21 =
E [ φ (cid:48)
( G ) 2 ] − b 20 , G ∼ N (0 , 1) .
Note that the population risk is again independent to γ 2 , and thus double descent does not appear for this initialization.
Roughly speaking, the reason that the risk does not become unbounded at some point is that in the asymptotic limit the pseudo-inverse ( JJ (cid:62) ) † is stable due to the nonlinearity and dh (cid:29) n .
We make two additional observations:
• the stability of the inverse at γ 1 → 1 depends on the lowest eigenvalue of the tangent kernel matrix (smaller λ min ( K ) entails larger variance), which is determined by the nonlinearity; • While our result only holds for network with zero initial output, for non-symmetric
(i
i.d.)
initialization we also observe that the risk is independent to
γ
but
the
bias
is higher than that
in
symmetric
initialization
as shown in Figure 8.
We comment that the non-
zero f init ( X )
in (13) behaves as zero-mean Gaussian due to central limit theorem; therefore, in the kernel regime the function output at initialization is equivalent to additive noise to the labels y , and we thus expect the magnitude of f init ( X ) to negatively inﬂuence the model generalization.
[Figure]
5. C OMPARING THE I NITIALIZATIONS
Figure 3 shows the agreement between theoretical prediction and experimental results.
Although in both cases the risk is independent to overparameterization ( γ 2 ), the two initializations lead to models with contrasting properties, as demonstrated by the following comparison on the risk.
Corollary 9.
For any γ 1 ∈ (0 , ∞ ) and nonlinearity φ , B ( ˆ f Van ) ≤ B ( ˆ f NV ) ≤ 1 .
On the other hand, for all m > 0 , V ( ˆ f NV ) =
O (1)
[Figure]
whereas V ( ˆ f Van ) can be arbitrarily large as γ 1 → 1 .
Remark .
m ≥ 0 for all smooth activations φ , and the equality holds if φ is linear.
Intuitively
[Figure]
small initialization enables the model "evolve" more during optimization and better align with the data and target function.
This potentially results in a lower bias, at the expense of overﬁtting more to the noise (high variance).
In contrast, with sufﬁciently large initialization the ﬁnal model becomes close to the initialized model, and thus we may expect it to be less “aligned” to the target (high bias) but is more stable (lower variance).
In illustrate the different inductive bias of the two initializations, we plot the trajectory of neurons in
Appendix A Figure 4.
Observe that for vanishing initialization the neurons stay close to one another throughout the trajectory, which results in a low-rank weight matrix, as predicted by Theorem 7.
In contrast, for non-vanishing initialization the neurons stay close to initialization (therefore full-rank), which validates the kernel approximation.
Last but not least, although the derived risk is only for learning a linear target function, we empirically observe that when the teacher is also a two-layer network, the population risk follows the same trend, i.e. double descent occurs when only the second layer is optimized, as shown in Figure 7.
D ISCUSSION AND F UTURE W ORKS
We derived the exact population risk of high-dimensional two-layer neural networks in learning a linear target function over Gaussian data with additive label noise, and showed that optimizing the ﬁrst or the second layer via gradient ﬂow results in solutions with contrasting properties.
Speciﬁcally, double descent is present when the second layer coefﬁcients are optimized, but not when the ﬁrst layer weights are optimized under certain initializations.
Moreover, we highlight that the scale of initialization leads to different inductive bias in optimizing the ﬁrst layer.
It should be noted that our analysis only applies to the unregularized objective: it has been shown that explicit regularization (such as (cid:96) 2 penalty) stabilizes the singularity at γ 2 → 1 (Mei and
Montanari, 2019), and algorithmic regularization (Li et al., 2019a; Dong et al., 2019) also provides robustness against noisy observations.
We further remark that our ﬁndings do not directly contradict the experimental double descent phenomenon, nor the practical beneﬁt of overparameterization.
In particular, the interpolation limit could occur at γ 2 → 0 which is beyond the regime we consider (such as Figure 4 in (Belkin et al., 2018)).
Thus what we conclude is that under the studied proportional asymptotics, the mechanism that provably gives rise to double descent from previous works on least squares regression might not translate to neural networks trained with gradient descent.
To simplify the computation, we rely on a set of strong assumptions similar to those in Hastie et al.
(2019), some of which we believe can be relaxed in future works, such as isotropic Gaussian input and linear target.
Importantly, the two speciﬁc scales of initialization studied in Section 5 are by no means exhaustive, and thus one would expect that under a different initialization of the ﬁrst layer, or the mean-ﬁeld 1 /h scaling of the second layer, the risk of the trained model can be very different.
Changing the loss function may also alter the generalization behavior of the network.
Another challenging problem is to extend the current analysis to beyond two layers.
Last but not least, our result characterizes gradient ﬂow which resembles gradient descent with small stepsize, and thus it would be interesting to study the effect of learning rate schedule, which has a known impact on generalization (Smith and Le, 2017; Li et al., 2019b).
A CKNOWLEDGEMENT
We thank Xiuyuan Cheng, Xuechen Li, Yiping Lu, Atsushi Nitanda, Shengyang Sun and anonymous reviewers for helpful comments and feedback.
JB and DW were partially funded by LG Electronics and NSERC.
JB and MAE were supported by the CIFAR AI Chairs program.
TS was partially supported by JSPS Kakenhi (26280009, 15H05707 and 18H03201), Japan Digital Design and
JST-CREST.
A A DDITIONAL F IGURES AND P LOTS
[Figure]
[Figure]
S UMMARY OF THE P RESENCE / A BSENCE OF D OUBLE D ESCENT
Singularity in 2nd Layer Trained (RF)
Vanishing Init.
Non-vanishing Init.
Bias γ 1 : No; γ 2 : Yes γ 1 : No; γ 2 : No γ 1 : No; γ 2 : No
Variance γ 1 : No; γ 2 : Yes γ 1 :
Yes; γ 2 : No γ 1 : No; γ 2 : No
[Figure]
[Figure] [Figure]
[Figure]
B B ACKGROUND B. R OTATIONAL I NVARIANCE
The rotational invariance of Gaussian distribution is crucial in our analysis throughout this paper.
A basic observation is that for a random Gaussian matrix X and any ﬁxed unitary matrix U , the distribution of X and UX are the same.
Lemma 10 (Rotational Invariance) .
Denote A ( X ) ∈ R
d × d a matrix function of X ∈ R d × n .
If A ( X ) satisﬁes that A ( UX ) =
UA ( X ) U T for all unitary U , then
[Equation]
for any ﬁxed nonzero β ∈ R d and random matrix X with each entry i.i.d.
X ij ∼ N (0 , σ 2 ) .
Proof.
Choose a set of Unitary matrices { U i } di =1 such that U (cid:62) i β
= (cid:107) β (cid:107) e i , where e i is the i -th canonical vector in R d .
Since U i X ∼ X , we have E
[ A ( X )]
=
E [ A ( U i X )]
=
U i E
[ A ( X )]
U (cid:62) i and hence
[Equation]
Note that the property also holds for matrix function A ( X ) that satisﬁes A ( XU ) =
UA ( X ) U T , and can be extended to matrix function A that takes multiple matrices as input.
For brevity we will refer to rotational invariance instead of equation (15).
[Equation]
For a real symmetric random matrix A ∈ R p × p , deﬁne its empirical spectral density as
[Equation]
, where δ a ( x )
=
δ ( x − a ) is the Dirac delta function.
Assume A = W p ∼ W p
( I, n ) is a Wishart matrix , i.e. W p =
X (cid:62) X/n and X ∈ R n × p is random Gaussian matrix with each column i.i.d.
X i ∼ N ( 0 , I ) .
Mar ˇ cenko and Pastur (1967) showed that as n, p → ∞ and p/n = γ ∈ (0 , ∞ ) , the empirical spectral density µ W p ( dλ ) converges weakly to a limiting density µ MP ( γ ) ( λ ) :
[Equation]
[Equation]
We say µ MP ( γ ) is the density of the Marˇcenko–Pastur distribution with support S =
[(1 − √ γ ) 2 , (1 + √ γ ) 2 ] (for 0 < γ < 1 ) or S = { 0 } ∪ [(1 − √ γ ) 2 , (1 + √ γ ) 2 ] (for γ ≥ 1 ).
Note that this implies that the smallest non-zero eigenvalue of A is bounded away from 0 a.s. for γ (cid:54) = 1 .
The explicit form of Mar ˇ cenko–Pastur distribution allows us to investigate the asymptotic properties of random matrices.
Generally speaking, by Pormanteau theorem one can translate any bounded continuous function on the empirical spectral density to the one on Mar ˇ cenko–Pastur distribution, i.e. for any bounded f ( λ ) ∈ C ( S ) , as n, p → ∞ with p/n = γ , almost surely
[Equation]
One implication is the following trace concentration on the inverse Wishart matrix for γ < 1 ,
[Equation]
We remark that instability of the trace of the invert Wishart matrix as γ → 1 plays an important role in the double descent phenomenon.
B. O RTHOGONAL P OLYNOMIALS
Orthogonal polynomials are useful in the analysis of nonlinear random matrices.
Suppose φ is a function in L 2 ( R , µ G ) , where G ∼ N (0 , 1) , µ G ( dx ) =
( √ 2 π ) − 1 e − x 2 / 2 dx is the Gaussian measure.
For n ≥ 0 , deﬁne the Hermite polynomials
[Equation]
Note that orthogonality can be easily veriﬁed:
[Equation]
Since { H i ( x ) } ∞ i =0 forms a set of orthogonal basis in L 2 ( R , µ G ) , the function φ ( x ) can be expanded under the Hermite basis as
[Equation]
Following Cheng and Singer (2013), we mainly focus on the ﬁrst few terms of the Hermite expansion.
One can check that H 0 ( x ) = 1 , H 1 ( x ) = x , and φ ( x ) can be expanded as
[Equation]
) where c 0 =
E [ φ ( G )] , c 1 = E
[ Gφ ( G )] , and terms in the RHS are orthogonal to one another in
L 2 ( R , µ G ) .
Taking square and expectation over both sides yields
[Equation]
24) which indicates E [ φ ( G ) 2 ] − E [ φ ( G )] 2 − E [ Gφ ( G )]
2 = E
[ φ ⊥ ( G ) 2 ] ≥ 0 .
Note that the equality holds if and only if E [ φ ⊥ ( G ) 2 ] = 0 , i.e. φ ( x ) =
c
+ c 1 x is linear.
One application of this decomposition is to “linearize” a non-linear matrix in high dimensions, which will be useful for the following sections.
C P ROOF OF M AIN R ESULTS C.
P ROOF OF L EMMA
Given features X ∈ R d × n , labels y ∈ R d and model parameters θ , the gradient ﬂow of θ on the squared loss (cid:13)(cid:13)
y − X
[Equation]
θ
[Equation]
2 2 can be written as
[Equation]
Thus with initialization θ 0 , the solution of this ODE at time t can be written in explicit form
[Equation]
Since θ 0 = 0 , taking t → ∞ yields the desired result.
[Equation]
C. P ROOF OF T HEOREM
We compute the bias and variance for different cases of γ 1 , γ 2 .
We ﬁrst discuss the case where the random feature Φ X is not full rank (Case I).
Otherwise when Φ X is full rank, we discuss whether it is full column rank (Case II) or full row rank (Case III).
Case I: W (cid:62) X is not full rank, i.e. γ 1 < 1 , γ 2 > γ 1 .
In this case rank (Φ X ) =
d < min( n, h ) , and thus by taking the Moore-Penrose inverse we obtain
[Equation]
It is clear that the mean and variance is identical to the underparameterized regime in Hastie et al.
(2019), i.e. when n, d, h → ∞ ,
[Equation]
→ 0;
[Equation]
Case II: W (cid:62) X has full column rank,
i
[Equation]
e. γ 2 < 1 , γ 1 > γ 2 .
By Lemma 1, the solution of the second layer coefﬁcients is
[Equation]
Denote W = U Σ V (cid:62) the singular value decomposition.
We perform the block decomposition:
[Equation]
[Equation]
where Σ 0 ∈ R h × h , X 0 ∈ R h × n , X 1 ∈ R ( d − h ) × n
[Equation]
and notice that X 0 , X 1 are independent.
By a concentration of measure argument (e
[Equation]
g. Tao (2012); Hastie et al.
one can show that the quantity below tightly concentrates at its expectation.
For the variance we have
[Equation]
where the last equality follows from Appendix B
[Equation]
Similarly for the bias term we have
[Equation]
where symmetric arguments are omitted as ( · · · )
[Equation]
and (i) follows from the rotational invariance argument introduced in Lemma 10 and that β (cid:62)
β
= r 2 .
By the block decomposition (30),
[Equation]
Therefore the bias term simpliﬁes to
[Equation]
Thus we have obtained that as n, d, h → ∞
[Equation] [Equation]
Case III: W (cid:62) X has full
row rank, i.e. γ 1 >
, γ 2 >
Similarly, the least squares solution is
[Equation]
[Equation]
Simplifying the variance:
[Equation]
where we applied the SVD of X
= U Σ V (cid:62) and the rotational invariance argument.
Using a similar block decomposition on W :
[Equation]
where Σ 0 ∈ R n × n
[Equation]
W 0 ∈ R n × h , W 1 ∈ R ( d − n ) × h
[Equation]
and W
[Equation]
W 1 independent.
We thus simplify V as
[Equation]
[Equation]
Hence we obtain the following expression on the variance
V
[Equation]
We omit the derivation of bias
[Equation]
which follows a similar derivation:
[Equation]
Combining Case I, II, III yields theorem 2.
[Equation]
C.
P ROOF OF P ROPOSITION
Given the squared loss, one can derive the dynamics of W with ﬁxed second layer a w.r.t the loss:
[Equation]
[Equation]
Note that the update of W can be written as a linear combination of a .
Since W (0) = 0 , we can write W ( t ) =
ˆ w ( t ) a (cid:62) for some ˆ w .
The corresponding ﬂow on ˆ w is
[Equation]
which gives the following solution
[Equation] [Equation]
Thus gradient ﬂow on the ﬁrst layer leads to the minimum-norm solution on the input features.
[Equation]
C. P ROOF OF T HEOREM
Following the bias-variance decomposition (7)
[Equation]
the variance term can be written as ( σ 2 omitted)
[Equation]
where we deﬁne the expected non-linear Gram matrix K W ∈ R
[Equation]
as
[Equation]
and for each entry we have ( K W )
[ i
[Equation]
Random matrix in the form of covariance matrix of nonlinear features has been studied in many works Hastie et al.
; Mei and Montanari (2019); Liao and Couillet (2018); Louart et al.
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.
with modiﬁcations mentioned below.
In contrast to the linear network in Section C.2, the Gram matrix of a nonlinear activation is almost surely full-rank, as speciﬁed in the following lemma from Pennington and Worah (2017):
Lemma 11.
Suppose φ is not linear.
Then the smallest singular value of Φ = φ ( W (cid:62) X ) ∈ R h × n is of order O ( √ n ) .
To be precise, consider the empirical spectral density µ ΦΦ (cid:62) /n
( dλ ) , when n, d, h → ∞ with d/n → γ 1 and h/n → γ 2 , µ ΦΦ
[Equation]
n ( dλ ) converges weakly to
[Equation]
where µ + ( dλ ) has non-negative support [ ρ, ∞ ) with ρ > 0 .
We therefore consider two scenarios: Φ is full column rank (Case I) or full row rank (Case II).
Case 1.
h < n .
In this case (45) simpliﬁes into
[Equation]
where the continuity and boundness of V ξ at ξ = 0 − is guaranteed by Lemma 11 when γ 2 = h/n (cid:54) = 1 .
A ridge ξ is added make use of (Louart et al., 2018, Theorem 1), which derived the asymptotic equivalent of the resolvent (cid:0)
[Equation]
It follows that as n, d, h → ∞ ,
[Equation]
47) where we have used the inequality tr ( AB ) ≤ (cid:107) A (cid:107) F (cid:107) B (cid:107)
[Equation]
and equivalently by taking ξ → 0 we get
[Equation]
Therefore lim ξ → 0 lim n,d,h →∞ n/h · tr (cid:16)(cid:0) ΦΦ (cid:62) − ξI (cid:1) − 1 K W (cid:17) =
γ 2 / (1 − γ 2 ) .
Note that ∂V ξ /∂ξ = tr (cid:0) ξ (ΦΦ (cid:62) − ξI )
− 2 K W (cid:1) is bounded around the neighbourhood of ξ = 0 , hence following the same argument as (Hastie et al., 2019, Theorem 4), we exchange the limit of ξ → 0 and n, d, h → 0 :
V → γ1 − γ 2 .
Case 2.
h > n .
Techniques used in the current proof are largely borrowed from Hastie et al.
; Cheng and Singer (2013), and we include the full proof for completeness.
It should be noted that compared to Hastie et al.
we handle the non-zero expectation of the nonlinearity under
Gaussian distribution, i.e. the off-diagonal entries of the kernel matrix is no longer zero-centered.
For simplicity we mainly adhere to the notations in Hastie et al.
We brieﬂy summarizes the procedure for deriving V .
Instead of calculating the variance directly, we analyze a modiﬁed quantity V ξ and then take ξ → 0 , which can be connected to the trace of the resolvent of matrix ˜
A deﬁned in (56); this translates the calculation of V ξ into the calculation of the
Stieltjes transform of ˜ A (57), (61), (62).
C
D ERIVING THE VARIANCE V FOR h > n
Step 1.
An equivalent expression.
For notational simplicity we omit the magnitude σ :
[Equation]
and due to the same continuity argument as in Case 1 we have
[Equation]
where S = φ ( W (cid:62) X ) / √ n = Φ / √ n , ξ ∈ C and (cid:61) ξ > 0 or ξ < 0 .
We decompose the normalized feature matrix S
= φ ( W (cid:62) X ) / √ n as S = U Σ V (cid:62) , where Σ = diag h × n ( φ 1 , · · · , φ n ) ∈
R h × n is a tall diagonal matrix, and U =
[ u 1 , · · · , u h ] ∈ R h × h is the set of orthogonal eigenvectors of SS (cid:62) =
φ ( W (cid:62) X ) φ ( X (cid:62) W ) /n , and V ∈ R n × n is the set of orthogonal eigenvectors of S (cid:62) S
=
φ
( X (cid:62) W ) φ ( W (cid:62) X )
/n .
Now the variance can be written as
[Equation]
By the same argument as in Lemma 13 of Hastie et al.
(2019) one can show that when n, d, h → ∞
[Equation]
52) in which ˜ K W is the approxmation of K W deﬁned in Lemma 16.
Writing the trace explicitly (denote eigenvalues λ
i =
φ 2 i , and φ n +1 = · · · = φ h = 0 ), we have
[Equation]
Since the positive support of spectrum λ is lower bounded and the density at 0 is 1 − γ − 1 2 , we have
V ξ →
[Equation]
where we deﬁne µ n ( x ) = 1 h (cid:80) hi =
δ λ i ( x ) u (cid:62) i ˜ K W u
i and its positive part µ + n ( x ) .
Hence we have
[Equation]
We deﬁne the following matrix ˜
A n ( ρ, ς, τ ) ∈ R N × N where N = n + h :
[Equation]
And denote the Stieltjes transform of ˜ A n as
[Equation]
Then following the deﬁnition of ˜ K W one can show that
[Equation]
and taking matrix derivative gives − ∂
[Equation]
Denote the limit ˜ m ( ξ, ρ, ς, τ ) =
lim n,h,d →∞ ˜ m n ( ξ, ρ, ς, τ ) , and the derivative is given as −
[Equation]
For simplicity we deﬁne the following function on ξ :
[Equation]
[Equation]
also denote q + ( ξ ) =
q ( ξ ) − γ 2 − 1 ξ 2 = γ 2 (cid:82) λ>ρ 1 λ + ξ 2 µ + ∞ ( dλ ) .
Comparing with (55) yields:
[Equation] [Equation]
Step 2.
Calculating q ( ξ ) and m n ( ξ, ρ, ς, τ ) This subsection aims to calculate ˜ m ( ξ, ρ, ς, τ ) and q ( ξ ) =
− ˜ m (cid:48)
x ( ξ, rx, sx, tx )
| x =0 , from which the variance can be computed from (57)(61)(62).
We deﬁne A n by subtracting the off-diagonal entries of the upper-left block of ˜ A n :
[Equation]
where S =
˜ S − a 0 I p × n , i.e. S ik = φ ( w (cid:62) i x k ) − a 0 = ϕ
( w (cid:62) i x k ) , a 0 = E [ φ ( x )] .
The Stieltjes transform of A n given by m n ( ξ, ρ, τ ) = 1 n tr (cid:0) ( A n ( ρ, τ ) − ξI N ) − 1 (cid:1) .
The following Lemma shows that ˜ m n and m n have the same limit:
Lemma 12.
when n → ∞ and for (cid:61) ξ > 0 or ξ < 0 , we have m n ( ξ, ρ, τ ) →
˜ m n ( ξ, ρ, ς, τ ) .
Proof.
By deﬁnition of ¯ A n and A n ,
[Equation]
which is a rank-2 matrix.
By theorem A.43 from Bai and Silverstein (2010), which characterizes the effect of ﬁnite-rank perturbation on the e.s.d.
of random matrices:
[Equation]
where F M is the empirical spectral distribution of M ∈ R n × n .
The claim follows from the Stieltjes continuity theorem (e.g. Section 2.4 in Tao (2012)).
To calculate the Stieltjes transform m n , we take advantage of the block structure of A n :
[Equation]
One can observe that m n ( ξ, ρ, τ ) =
γ 2 m 1 ,n ( ξ, ρ, τ ) + m 2 ,n ( ξ, ρ, τ ) .
In the following equations we omit the subscript n , as well as dependency on ρ, ς, τ .
Following
Hastie et al.
(2019), we rewrite the matrix A n =
A = (cid:20)
A ∗ a a (cid:62) (cid:21) , where A ∗ is a ( N − 1) × ( N − 1) matrix with last column and row of A removed and s the activation vector:
[Equation]
[Equation]
Hence by the block matrix inverse formula
[Equation]
Plugging this back in the Stieltjes transform we have
[Equation]
To obtain an asymptotic description of m
[Equation]
n
[Equation]
we perform the orthonormal decomposition on the nonlinearity ϕ introduced in Cheng and Singer (2013): ϕ ( x )
= a 1 x + ϕ ⊥ ( x ) , where a 1 = E x
∼N (0 , 1) [ xϕ ( x )] .
For each w i (1 ≤ i ≤ h )
[Equation]
we perform the following orthonormal decomposition (along the direction of x n and the direction of ˜ w i perpendicular to x n ):
[Equation]
We can thus simplify the activation vector a as
[Equation]
and for 1 ≤ i (cid:54)
=
j ≤ h , 1 ≤ k ≤ n
[Equation] [Equation]
Similarly we decompose the activation function in S ,
[Equation]
We thus have an equivalent expression of matrix A ∗
[Equation]
By argument similar to (Hastie et al., 2019, B.1.2)
[Equation]
E diminishes to 0 as n → ∞ with respect to the
Frobenius norm, therefore by the Woodbury’s identity and the expression of m 2 ,n in (70) m 2 ,n (
[Equation]
We bound each term u, v , S to compute m 2 ,n .
For u
[Equation] [Equation]
where b = E x ∼N (
[Equation] [Equation]
[ ϕ
[Equation]
x
[Equation]
=
E x ∼N
[Equation]
E
[Equation]
From a standard concentration of measure argument we have that as n, h, d → ∞
[Equation]
[Equation]
And for v (note that U is dependent on a )
[Equation]
0 (cid:105) .
(79) where v = tr (cid:16) ( ˜ A ∗ − ξI N − 1 ) − 1 [1 ..h,
1 ..h ] E s [ ηs (cid:62) ] (cid:17) =
[Equation]
1 (cid:112) γ 22 /
[Equation]
1 m 1 ,n ( ξ, ρ, τ ) .
We thus have
[Equation]
[Equation] [Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table] [Equation] [Table] [Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
[Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
And hence as n, d, h → ∞ ,
[Equation]
Therefore by combining (78)
[Equation] [Equation]
we arrive at the following expression on m 2 ,n ,
[Equation]
Similarly we can calculate m 1 ,n
[Equation]
ξ, ρ, τ ) as
[Equation]
Uniqueness in (84)(83) follows from (Hastie et al., 2019, Sec B.1) and is omitted.
(cid:3)
C. P ROOF OF C OROLLARY
In this section we take the limit γ 1 → ∞ .
In this case (8), (9) simplify to
[Equation]
[Equation] [Equation]
Recall that m ( ξ, ρ, τ ) =
γ 2 m 1 ( ξ, ρ, τ ) + m 2 ( ξ, ρ, τ ) .
By taking the derivative we have − q
[Equation]
[Equation]
Observe that (85), (86) constitutes a set of implicit functions.
Thus differentiating the two functions with respect to τ, ρ and then substitute by ρ = τ = 0 gives
[Equation]
Hence by (62) we obtain the asymptotic variance:
[Equation]
Combining the case where γ 2 < 1 in Theorem 4 completes the proof.
[Equation]
Remark .
For φ ( x ) = ReLU( x ) , c 1 = 1 / 2 − 1 / (2 π ) , c 2 = 1 / 4 .
For φ ( x ) =
SoftPlus( x ) =
log(1 + e x ) , numerical integration yields c 1 ≈ 0 .
2715 , c 2 = 1 / 4 .
C. P ROOF OF C OROLLARYC.7. U NBOUNDED B IAS
WHEN h = n
From the bias-variance decomposition (7), the bias B is written as B = r 2 d tr ( Q 1 + Q 2 + I d ) , where
[Equation]
[Equation]
When h = n , due to the nonlinearity of φ , we have φ ( W (cid:62) X ) is full rank a.s., and hence [ φ ( X (cid:62) W )] † =
[ φ ( X (cid:62) W )] − 1 .
We have the following bound for Q
[Equation]
Since W and X/ √ d are R d × n = R d × p follows the same distribution where each entry i.i.d.
N (0 , 1 /d ) ,
[Equation]
where in Section C.5 we have showed (91) is unbounded when n → ∞ .
Moreover, by (154) and
Weyl’s theorem we have λ min ( K W ) =
O (1) , and thus d − 1 tr ( Q 1 ) is unbounded.
For d − 1 tr
[Equation]
To sum up, for n → ∞ and γ 2 → 1 we have B → ∞ .
[Equation]
C.7. B OUNDED B IAS WHEN
γ >
Since h > n
[Equation]
the two terms in the expression of the bias can be written as
[Equation]
Therefore we have
[Equation]
[Equation]
[Table]
we used
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
where we applied tr ( AB ) ≤ (tr (cid:0) A (cid:62) A (cid:1)
+ tr (cid:0) B (cid:62) B (cid:1) ) / 2 .
We therefore conclude that B is bounded when h > n and h, n → ∞ .
[Equation]
Remark .
Concurrent to this work, Mei and Montanari (2019) provides a complete characterization of the bias term and conﬁrms our observations above.
C. P ROOF OF T HEOREM
For simplicity we assume n, d, h to be even and let d 0 = d/ 2 , n 0 = n/
2 and h 0 = h/ 2 .
Since the second layer is ﬁxed a i ∼ Unif {− 1 / √ h, 1 / √ h } , we let a i = 1 / √ h and a i + h
= − 1 / √ h for all 1 ≤ i ≤ h 0 .
We therefore write a (cid:62) =
h − 1 / 2 [ 1 h 0 , − 1 h 0 ] (cid:62) and W =
[ W + , W − ] :
[Equation]
The empirical risk can thus be written as
[Equation]
C.8. D EFINING G RADIENT F LOWS
In this section we deﬁne three gradient ﬂows and show that the three ﬂows are similar in some sense.
GF-Original is the original gradient ﬂow (11), i.e.
[Equation]
[Equation] starting from the vanishing initialization w Oi (0) ∼ N ( 0 , I/dh 1+ (cid:15) ) .
Note that the gradient for the negative part W O − can be similarly deﬁned.
We now deﬁne the ﬂow under the same objective but from exact zero initialization w Di (0) =termed GF-Double .
Due to zero initialization, a basic observation is that the solution [ W D + , W D − ] is at most rank-2, and more precisely, the parameters in the ﬂow takes the form of W D ± ( t ) =
w D ± ( t
[Equation]
1 (cid:62) where w D ± ( t ) admits the following dynamics:
[Equation] [Equation]
Lastly, we deﬁne the GF-Single with solution denoted as w ± =
w S ± ( t ) :
[Equation]
from zero initialization w D ± (0) = 0 .
This can be seen as replacing the nonlinearity φ with its ﬁrst-order Taylor expansion at the origin.
C.8. F ROM GF- DOUBLE TO GF- SINGLE
Step 1.
Solution of GF-single.
Among the three ﬂows deﬁned above, only GF-single an explicit form at any time t .
Speciﬁcally, the solution can be written as:
[Equation]
when d < n , or otherwise
[Equation]
when d > n .
For simplicity we elaborate the proof only for d < n .
We provide a condition under which the difference between the trajectories can be controlled
[Equation]
and we ﬁrst that this condition holds for the linearized ﬂow GF-single for all t in Lemma 17.
[Equation]
Step 2.
Bounding the Difference in Gradient Flow Trajectory.
Due to low rank property of
GF-single and GF-double, in this subsection we slightly abuse the notation and deﬁne
[Equation]
We now show that the difference between the two trajectories deﬁned above is asymptotically vanishing for w + ( w − follows the same argument).
Compare the two trajectories up to time T
[Equation]
106) where we have deﬁned the error term as
[Equation]
[Table]
(108) where (i) follows from the error of Taylor expansion on φ and (ii) from Condition A.
Therefore, by
Equation (106) and Gronwall’s inequality,
[Equation]
for T ∈ O (log log h ) .
This shows that up to time T, the difference between the trajectories of
GF-single and GF-double vanishes under the the assumption that Condition A holds for w D up to
T. Importantly, note that w S (0) =
w D (0) = 0 , and that Condition A holds for w S for all T > 0 .
Therefore, by a standard contradiction argument (e.g. (Du et al., 2018, Lemma 3.4)), one can show that this closeness between the two trajectories implies that Condition A also holds for w D up to time T. With Equation (108) we bound the difference in population risk between the two ﬂows.
Step 3.
Bounding the Difference in Risk.
Now we consider the difference of the population risk
R S and R D of two models with parameters w S ± and w D ± .
[Equation]
+ x ) − φ ( w D (cid:62) + x ) | + | φ ( w S (cid:62)− x ) − φ ( w D (cid:62)− x ) | ]· (cid:113) E
x [ | β (cid:62) x − f ( x ; w S ± ) | + | β (cid:62) x − f ( x ; w D ± ) | ]( ii ) ≤ 2 (cid:113)
h
E x [ | φ ( w S (cid:62) + x ) − φ ( w D (cid:62) + x ) | 2 + | φ ( w S (cid:62)− x ) − φ ( w D (cid:62)− x ) | 2 ] · (cid:113)
E
x [ | β (cid:62) x − f ( x ; w S ± ) | 2 + | β (cid:62) x − f ( x ; w D ± ) | 2 ] ( iii ) ≤ 2 (cid:113)
h Lip( φ ) E x [ | w S (cid:62) + x − w D (cid:62) + x | 2 + | w S (cid:62)− x − w D (cid:62)− x | 2 ] · (cid:112) R S + R D
[Equation]
(cid:19) , (110) where (i) is due to Cauchy-Schwarz inequality on norm, (ii) from Jenson’s inequality on squares and
Young’s inequality, and (iii) from the Lipschitz assumption on the activation, and the observation that both R S and R D are ﬁnite for γ 1 (cid:54) = 1 due to the justiﬁed Condition A above.
Therefore the difference between R S and R D vanishes for T
=
O (log log h ) .
Step 4.
Bounding the Difference from Stationarity We compute the difference in risk between the model at some ﬁnite time t and the stationary point i.e. t = ∞ ,
[Equation]
for constants C 3 , C 4 > 0 .
Combining (110) and 111) yields for
[Equation]
The result in (112) for case d > n follows a similar proof and is omitted.
C.8. F ROM GF- ORIGINAL TO GF- DOUBLE
In this section we compare GF-original with GF-double.
Note that the two ﬂows differ only at initialization: vanishing initialization w Oi (0) ∼ N ( 0 , I/dh 1+ (cid:15) )
v.s. zero initialization w Di (0)
= 0 .
Note that at initialization (cid:107)
y − f ( X ) (cid:107)
= O ( √ n ) , and gradient ﬂow decreases the empirical risk; therefore the condition for Lemma 18 is satisﬁed
[Equation]
and by the Lipschitz condition on the empirical gradient we have
[Equation]
And hence by Gronwall’s lemma one obtains:
[Equation]
Therefore the difference in the function output can be bounded as
[Equation]
Taking T = log log h , together with the same argument in Step 1 yields
[Equation]
C.8. P UTTING THINGS TOGETHER
By (112) and (116), we know that for T = log log h ,
[Equation]
Finally the proof is completed by observing that R S ( ∞ ) is the risk of the minimum-norm solution on the input discussed in Section 3.
In addition, note that at T = log log h , from (109)(114) one obtains that (cid:13)(cid:13) W O ( t ) − W S ( t ) (cid:13)(cid:13) F → 0 , and therefore by Lemma 18 we have (cid:13)(cid:13) ∂L
( X ; W O ( t )) /∂W
(cid:13)(cid:13) F → 0 , i.e. the ﬂow on the original objective reaches a o (1) ﬁrst-order stationary point.
(cid:3)
C. P ROOF OF T HEOREM
Denote ω = vec( W ) =
vec([ W + , W − ]) , and ω 0 = vec( W init ) .
Deﬁne
[Equation]
[Equation]
which is the kernel matrix of the neural tangent kernel Jacot et al.
Du et al.
In the following sections we show that under the non-vanishing initialization, the trained two-layer network is well-approximated by the regression model on the NTK, for which we derive the population risk.
C.9. T HE K ERNEL L INEARIZATION
Write y NN ( t ) =
f NN ( X, t ) ∈ R n and its evolution:
[Equation]
(119) and the corresponding linearized ﬂow:
[Equation]
Previous works (e.g. Du et al.
; Oymak and Soltanolkotabi (2019)) have proved (nonasymptotically) that the two trajectories (119) and (120) are close if the model is overparameterized, i.e. h = poly( n ) , under no assumptions on the teacher model.
In our asymptotic setup (together with assumptions (A1-3)), we argue that similar conclusion holds without signiﬁcant overparameterization.
We ﬁrst show the global convergence of the training of two-layer neural network f NN .
We employ an argument similar to (Du et al., 2018, Theo. 3.2) by ﬁrst identifying the condition under which training converges at linear rate:
[Equation]
From Corollary 15 we know that at initialization the lowest eigenvalue of the NTK matrix satisﬁes λ min ( K (0))
=
O ( d ) , and thus the kernel regression on the NTK enjoys linear convergence.
By
Lemma 19, we know that for (cid:107) W ( t ) − W (0) (cid:107)
=
O
( d 1 − (cid:15) ) , the order of λ min ( K ( t ))
=
O ( d ) remains unchanged.
Therefore, if we assume that up to time T the weights satisfy (cid:107) w i ( t ) − w i (0) (cid:107) 2 =
O ( d − 1 / 2 ) , then Condition B is satisﬁed for y NN and from (Chizat and Bach, 2018b, Lemma B1)
we have the following linear convergence
[Equation]
[Equation]
Since (cid:107) y NN (0) − y (cid:107) 2 = O ( √ n ) at initialization, setting T =
O (log d ) ensures that the training loss 1 n (cid:107) y NN ( T ) − y (cid:107) 22 → 0 as n → ∞ .
Consequently it is easy to check that (cid:13)(cid:13)(cid:13) ∂L ( X ; W ( T ))
∂W
[Equation]
and thus at time T the gradient ﬂow reaches an o(1)
ﬁrst order stationary point.
We now verify that each weight vector w i travels at most O ( d − 1 / 2 ) from initialization.
The norm of gradient for w i can be bounded as:
[Equation]
where (i) follows from the boundedness of φ (cid:48) .
Integrating the gradient yields (cid:107) w i [Equation] t ) − w i [Equation]
) (cid:107) 2 = O ( d − 1 / 2 ) , i.e. (cid:107) W ( t ) − W (0) (cid:107)
= O (1) .
Thus the distance traveled by W indeed satisﬁes the norm assumption above, and following the same argument as (Du et al., 2018, Lemma 3.4) we conclude that Condition B holds true for the gradient ﬂow of f NN for t > 0 .
Next we show that for any input ˆ x with (cid:107) ˆ x (cid:107)
= O ( √ d ) , the prediction of the neural network is uniformly close to the prediction of the prediction of the kernel model, a result similar to (Arora et al., 2019a, Lemma F.1).
Following the notation of Arora et al. (2019a), we write the time derivative of the prediction as
[Equation]
where u NN ( x , t ) =
∂
f ( X ; ω ( t ))
∂ ω ( t )
[Equation]
∈ R n and u NTK ( x , t ) similarly deﬁned on the initialized weights ω (0) .
We bound the difference between the predictions on ˆ x up to terminal time T as
[Equation]
For the ﬁrst term have
[Equation]
where (i) is due to Corollary 15, Lemma 19 for some (cid:15) (cid:48) >
0 and the linear convergence of y NN , and (ii) is due to Gronwall’s inequality.
Note that the log factor in T is omitted.
Similarly, for the second term we have
[Equation]
[Equation]
where we used Lemma 19 and the linear convergence of y NN in (i).
Combining the two cases yields
[Equation]
where we utilized Corollary 14 in (i).
Thus we know that the difference between the population risk of f NN and f NTK is also asymptotically vanishing (note that the derivation above is independent of the target function as long as (cid:107) y (cid:107)
= O ( √ n ) ).
Therefore, in the following subsection we compute the risk of the linearized (kernel) model f NTK .
C.9. C OMPUTING THE K ERNEL R ISK
Given input X ∈ R d × n and label y
= β (cid:62) X + ε , gradient ﬂow on the tangent kernel solves the following equation of the parameters ω :
[Equation]
where ∂ f ( X ; ω 0 )
/∂
ω is a dh × n matrix with column ∂f ( x i ; ω 0 ) /∂
ω .
Note that for n → ∞ and γ 1 , γ 2 ∈ (0 , ∞ ) , dh > n trivially holds, and by Corollary 15 we know that solution is given by
[Equation]
And the population risk can be written as (note that there is a factor of 2 due to the "doubling trick" at initialization to ensure f init ( · ) = 0 ):
[Equation]
where a bias-variance decomposition is made here, and for simplicity we deﬁne
[Equation]
C.9. A PPROXIMATING THE K ERNEL M ATRIX
In this section we drop the negligible (cid:15) in the initialization.
Following Cheng and Singer (2013) we utilize the orthonormal decomposition of φ (cid:48)
( x ) in L 2 ( R , µ G ) .
Denote b 0 =
E [ φ (cid:48) ( G )] , and b 21 =
E [ φ (cid:48)
( G ) 2 ] − b 20 .
We have the orthogonal decomposition of φ (cid:48)
[Equation]
131) where E [ φ (cid:48)⊥ ( G )] = 0 .
We develop the following lemmas to approximate the kernel matrix.
Lemma 13 (Approximation of ( ˆ K X ) ij ) .
There exist constants c, c (cid:48) > 0 such that for i (cid:54) =
j with probability 1 − e − cnε 2 we have
[Equation]
[Equation]
and with probability 1 − e − c (cid:48) nε 2 , (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
ε.
Proof.
When i (cid:54) =
j (i
[Equation]
e. Equation (132)), we have
[Equation]
[Equation]
The matrix H ( x i , x j ) =
x (cid:62) i x j E w (cid:104)
[Equation]
w (cid:62)
[Equation]
φ
[Equation]
( w (cid:62) x j ) (cid:105) can be seen as the expected tangent kernel of nonlinear activation function studied in Du et al.
Arora et al. (2019b).
Moreover, due to the assumed boundedness of φ (cid:48) ( x ) (A3), by Hoeffding’s inequality we have
[Equation]
[Equation]
In addition, by the concentration of x (cid:62) i x j and (cid:107) x i (cid:107) 22 , i.e. Pr x (cid:62) i x j /d > ε < 1 − e − c 2 dε 2 and
Pr | x (cid:62) i x i /d − 1 | < ε > 1 − e
[Equation]
c 3 dε 2 , the orthonormal decomposition φ (cid:48) ( x ) =
b 0 x + φ (cid:48)⊥ ( x ) leads to the following linear approximation of the matrix H
[Equation] [Equation]
and by taking ε = x (cid:62) i x j /d under the joint event we can show that
[Equation]
with probability 1 − e − cdε 2 .
The same argument follows for the case where i = j .
Corollary 14 (Approximation of ˆ u ) .
For large enough l > 0 , with probability 1 − de − c log l d
[Equation]
where ˜ u
= b 20 x (cid:62) X .
Proof.
Taking ε = log l d/d together with Lemma 13 yields the desired result.
Corollary 15 (Approximation of ˆ K X ) .
With probability 1 − de − c log l d
[Equation]
where ˜ K X = b 20 X (cid:62) X + b 21 dI .
Proof.
Also by directly applying Lemma 13.
Remark .
For initialization larger or equal to w i (0) ∼ N (0 , I d /d ) , the above approximation does not depend on the scale of initialization.
Remark .
For φ ( x ) =
SoftPlus( x ) , b 20 = 1 / 4 , b 21 = 0 .
For φ ( x ) = sigmoid( x ) =
(1 + e − x ) − 1 , b 20 = 0 .
042692 , b 21 = 0 .
Note that b 1 ≥ 0 for all smooth activations φ , and the equality holds (i.e. b 1 = 0 )
if and only if φ is linear.
We comment that smaller b 1 entails larger variance as γ 1 → 1 , and vice versa, as shown in the following section.
C.9. T HE B IAS T ERM
With these approximation above we proceed to calculating (129)
[Equation]
We ﬁrst bound the error in substituting ˆ u with ˜ u :
[Equation]
where (i) is due to the fact that (cid:13)(cid:13)(cid:13) ˆ K − 1 X (cid:13)(cid:13)(cid:13) 2 = λ − 1min ( ˆ K X )
=
O (1 /d ) and (cid:107) X (cid:107)
= O ( √ d ) .
Therefore we have as n, d, h → ∞
[Equation]
By taking expectation over x and the rotational invariance argument similar to Hastie et al.
[Equation]
In addition
[Equation]
we bound the error in substituting ˆ K X by ˜ K X deﬁned in (139):
[Equation]
and similarly,
[Equation]
Combining these two formulas in (143) yields
[Equation]
Utilizing the Marˇcenko–Pastur law from Section B.2 we obtain
[Equation]
147) where m = b − 2 0 b 2 1 .
C.9. T HE V ARIANCE T ERM
Similarly, for the variance we utilize the approximation
[Equation]
[Equation]
Speciﬁcally, we bound the approximation error
[Equation]
and similarly
[Equation]
By combining the two approximations above we know that as n, d, p → ∞
[Equation]
Therefore the variance is given as
[Equation]
where m is deﬁned in the derivation of the bias term.
C.9. P UTTING T HINGS T OGETHER
Recall the population risk is the sum of the bias and variance
[Equation]
Observe that the population risk is independent of γ 2 , i.e. double descent does not occur when the network is overparameterized via changing the width.
In addition, the bias is monotonically increasing and upper-bounded by the null risk r 2 and lower-bounded by the bias of the least squares solution on the input features ˆ β = X † y , whereas the variance remains bounded for all γ 1 ∈ (0 , ∞ ) as long as m > 0 , i.e. φ is nonlinear.
[Equation]
D U SEFUL L EMMAS
Lemma 16.
Given K W from (46) , deﬁne
[Equation]
where Q ∈ R h × h with Q
i (cid:54) =
j = w (cid:62) i w j and Q i,i = 0 , and r =
E [ φ ( G ) 2 ] − E
[ φ ( G )] 2 , s =
E [ φ ( G )] 2 , t = E
[ Gφ ( G )] 2 .
Then as n, d, h → ∞ , (cid:13)(cid:13)(cid:13) K W − ˜ K W (cid:13)(cid:13)(cid:13) F ≤ log c d a.s..
Proof.
Consider the event
where
A (cid:15) =
(cid:8) | (cid:107) w i (cid:107) 2 − 1 | < (cid:15), | w (cid:62) i w j | < (cid:15) (cid:9) .
Under event A (cid:15) , for the diagonal term of the kernel matrix we have
[Equation]
And for off-diagonal term, by the decomposition introduced in Section B.3 we have
[Equation]
[Equation]
[Equation]
hence | [ K W ] ij − [ ˜ K W ] ij | < ( w (cid:62) i w j ) 2 .
Notice that A (cid:15) holds a.s. for (cid:15)
= log c d/ √ d and large enough c > 0 ; we therefore have (cid:13)(cid:13)(cid:13) K W − ˜ K W (cid:13)(cid:13)(cid:13) F ≤ log c d .
Lemma 17.
Let ˆ β ( t ) be the solution to the gradient ﬂow at time t deﬁned in (103)(104) .
Then as n, d → ∞ and γ 1 (cid:54) = 1 the following holds.:
[Equation]
Proof.
We consider d < n for simplicity, and result for the other case follows in similar fashion.
In this case ˆ β ( t ) = (cid:0) I − exp( − tn XX (cid:62) ) (cid:1) ( XX (cid:62) ) − 1 X y .
From (Hastie et al., 2019, Corollary 1) we know that (cid:107) ˆ β ( ∞ ) (cid:107)
=
(cid:13)(cid:13)
( XX (cid:62) ) − 1 X y (cid:13)(cid:13)
= O (1) for γ 1 < 1 .
Note that (cid:13)(cid:13)
I − exp( − tn XX (cid:62) )
(cid:13)(cid:13) 2 =
O (1) for t ≥ 0 ; it follows that (cid:107) ˆ β ( t ) (cid:107)
=
O (1) .
For the second part, we utilize the SVD X
= U Σ V (cid:62) , where Σ = (cid:2) ˆΣ; 0 (cid:3) , ˆΣ ∈ R d × d and ˆΣ ii
=
λ i .
We have (cid:0) I − exp( − tn XX (cid:62) ) (cid:1) =
U ¯Σ U (cid:62) where ¯Σ i,i = 1 − exp( − tλ 2 i /n )
if i ≤ d and 0 otherwise.
[Equation]
where (i) follows from the concentration of the Gaussian maxima, and the fact that the law of V is the
Haar measure on SO ( n ) , and thus for any unit vector z independent to V , V z is uniform on sphere and (cid:107) V z
[Equation]
O (log d/ √ d ) .
We wherefore have
[Equation]
Note that this result also implies that (cid:107) y − X (cid:62) β ( t ) (cid:107) ∞
=
O P (poly log d ) .
Lemma 18.
For weight matrices W , W (cid:48) satisfying (cid:107) y − f ( X ) (cid:107)
= O ( √ n ) , where f ( X ) =
φ ( XW ) a with ﬁxed a i ∼ Unif {− 1 / √ h, 1 / √ h } , given (A1)-(A3), the gradient of the empirical risk deﬁned in (11) is Lipschitz w.r.t.
W in the Frobenius norm, i.e.
[Equation]
Proof.
Denote y 1 =
φ ( X (cid:62) W 1 ) a and y 2 = φ ( X (cid:62) W 2 ) a for W 1 , W 2 satisfying the assumption above (which can be seen as a condition on the magnitude of training loss), we have
[Equation]
We upper bound the two terms separately:
[Equation]
where we applied the inequality (cid:107)
A
[Equation]
B (cid:107) F in (i), boundedness of φ (cid:48) in (ii) and
Lipschitzity of φ in (iii).
Similarly, for the second term
[Equation]
where we used the assumption on the training loss and the Lipscthizity of φ (cid:48) in (i).
Combining the two terms yields the desired result.
Lemma 19.
Under assumptions (A1-3) and the non-vanishing initialization, given that (cid:107) w i ( t ) − w i (0) (cid:107)
=
O
( d − 1 / 2 ) for all i , then we have (cid:107) K ( t ) − K (0) (cid:107)
= O ( d 1 / 2 − (cid:15) (cid:48) ) for some positive (cid:15) (cid:48) ∈ Θ(1) .
Proof.
Recall the deﬁnition of the NTK:
[Equation]
or equivalently the matrix form
[Equation]
At initialization, x i ∼ N (0 , I d ) and w k (0) ∼ N (0 , d (cid:15) I d ) .
Thus for ﬁxed x
i , by Gaussian anti-concentration we have Pr | x (cid:62) i w k | < log d ≤ O (1 /d 1 / 2+ (cid:15) 1 ) for some (cid:15) 1 > 0 .
In addition, note that (cid:107) w k ( t ) − w k (0) (cid:107)
= O ( d − 1 / 2 ) for all k , and therefore for i, j, k such that | x (cid:62) i w k (0) | >
O (log d ) and | x (cid:62) j w k (0) | >
O (log d ) , we know that | φ (cid:48) ( x (cid:62) i w k ( t )) φ (cid:48)
( x (cid:62) j w k ( t )) − φ (cid:48)
( x (cid:62) i w k (0)) φ (cid:48)
( x (cid:62) j w k (0))
| =
O ( d − 2 ) .
Given ﬁxed x i , deﬁne y k
= 1 {| x (cid:62) i w k | < log d } as the indicator variable that the k -th neuron does not saturate.
We know that E
[ y k ]
= O
(1 /d 1 / 2+ (cid:15) 1 ) , and Var[ y k ]
=
E [ y 2 k ] − E
[ y k ] 2 = O (1 /d 1 / 2+ (cid:15) 1 ) .
By Bernstein’s inequality
[Equation]
[Equation]
Setting ε = (cid:113) c log h
h 1+ (cid:15) 2 , we know that with probability at least 1 − h − c , h h (cid:88) k =
[Equation]
(cid:18) poly log h h
[Equation] [Equation] (cid:15)(cid:19) . Therefore, given x i and x j , for large enough c 1 with probability at least 1 − h − 3 we have (cid:12) (cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
=
[Equation]
167) in which we utilized the boundedness of φ (cid:48) .
Taking union bound over d 2 elements in the random feature matrix yields
[Equation]
Using the exact same argument, one can derive that (cid:107) u NN (ˆ x ) − u NTK (ˆ x ) {+(cid:107) 
+}
= O ( d 1 / 2 − (cid:15) (cid:48) ) , the proof of which we omit.
E A DDITIONAL R ESULTS E. R ISK OF ReLU N ETWORK UNDER S YMMETRIC D ATA
If the dataset is symmetric, that is
[Equation]
A5)
Symmetric [Equation]
then population risk of the gradient ﬂow solution can be given explicitly for certain nonlinearities:
Proposition 20.
Given (A1-3)(A5), if the nonlinearity satisﬁes φ (cid:48) ( x ) + φ (cid:48) ( − x ) =
C for constant C, then as n, d, h → ∞
[Equation]
Note that the requirement on the nonlinearity holds for ReLU and SoftPlus .
This expression is again independent to γ 2 and aligns with the experimental results in Figure 9 (we only plot the bias component for veriﬁcation).
In addition, the bias is upper-bounded by the null risk for all γ 1 .
We remark that the symmetry assumption does not hold for i.i.d.
samples from symmetric distributions, and Figure 9 demonstrates that the additional condition alters the risk. [Figure]
Proof.
Without loss of generality assume X =
[ X 0 , − X 0 ] .
Then by (100) we have
[Equation]
and the ﬂow for w − follows from symmetry.
In this case one can show that from exact zero initialization, for nonlinearity satisfying φ ( x ) − φ ( − x ) = x , such as ReLU and SoftPlus , [Equation]
[Equation]
And therefore the gradient ﬂow of w + is
[Equation]
[Equation]
The ﬂow of w − follows from symmetry.
Solving for the stationary points (i
[Equation]
e. gradient becomes zero), it the clear that
[Equation]
And hence the asymptotic risk is
[Equation]
The same conclusion holds for vanishing initialization if we assume that the trajectory stays close to that of exact zero initialization.
Note that although the prediction aligns well with the experimental results, the argument in Theorem 7 does not directly apply due to the undeﬁned derivative of ReLU at the origin, and thus this result is not rigorously justiﬁed.
F E XPERIMENT S ETUP
Optimizing the Second Layer.
We compute the minimum-norm solution by directly solving the pseudo-inverse.
We set n = 1000 and vary γ 1 , γ 2 from 0.1 to 3.
The linear teacher model
F ( x ) = x (cid:62) β is ﬁxed as β = − 1 d / √ d .
For each ( γ 1 , γ 2 ) we average across 50 random draws of data.
Optimizing the First Layer.
For both initializations, we use gradient descent with small step size ( η = 0 . 1 ) and train the model for minimally 25000 steps and till (cid:107)∇ W f ( X, W ) (cid:107) 2 F < 10 − 6 .
We ﬁx n = 320 and vary γ 1 , γ 2 from 0.1 to 3 with the same linear teacher model β = − 1 d / √ d .
The risk is averaged across 20 models trained from different initializations.