Generalization of Two-layer Neural Networks: An Asymptotic Viewpoint
Anonymous authors
Paper under double-blind review
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 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 diﬀerent initialization setups.
When only the second layer coeﬃ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 diﬀerent scales initialization lead to diﬀerent 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 phenomenon 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 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 (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).
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.
However, there is still a gap between the practical beneﬁt of overparameterization and the recently proved double descent , which is rigorously established under models that exhibits the same 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 double descent phenomenon.
We consider the 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 features are larger.
Following Hastie et al.
(2019), we assume Gaussian data 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 double descent risk curve.
• When the ﬁrst layer is optimized, we compare two solutions of gradient ﬂow from diﬀerent scales of initialization, which we term as vanishing and nonvanishing initialization, and show in both cases the risk is independent to overparameterization.
• For the vanishing initialization, we prove 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 well-approximated by a kernel model and derive the population risk.
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); Rotskoﬀ and Vanden-Eijnden (2018); Sirignano and Spiliopoulos (2018) 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 partial diﬀerential equations or Wasserstein gradient ﬂow.
On the other hand, Jacot et al.
Du et al.
Arora et al. (2019a); Oymak and Soltanolkotabi (2019); Cao and
Gu (2019); Allen-Zhu et al.
(2018b)
; Song and Yang (2019) considered a diﬀerent scaling and showed that gradient descent converges to global minimizer at a linear rate; key to these results is that in the overparameterization limit, training is asymptotically equivalent to kernel regression with respect to a speciﬁc kernel termed the 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 diﬀerent 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.
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 two-layer neural network.
Suzuki (2018); Allen-Zhu and Li (2019); Yehudai and
Shamir (2019) demonstrated that neural network outperforms linear estimators (including kernel methods) in learning various target functions.
The diﬀerence between ﬁxed bases and adaptive bases mirrors the diﬀerence in optimizing the ﬁrst or second layer in our setup.
Generalization of Overparameterized Models.
It is often observed that overparameterization does not lead to 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 may decrease as the model capacity increases (Arora et al., 2019b; Li and Liang, 2018; Allen-Zhu et al., 2018a; Oymak et al., 2019).
Compared to these generalization bounds that require signiﬁcant overparameterization, our result relies on stronger data assumptions, but consequently we obtain the exact population risk.
Beyond the kernel regime, Advani and Saxe (2017) and Goldt et al.
(2019) analyzed the generalization dynamics of overparameterized networks 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 (Belkin et al., 2018; Geiger et al., 2018;
Spigler et al., 2018; Advani and Saxe, 2017).
This phenomenon can be rigorously proved for certain simple models (Hastie et al., 2019; Belkin et al., 2019; Muthukumar et al., 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) focused on 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 only 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 diﬀerent generalization property of models obtained from optimizing diﬀerent layers of the network and from diﬀerent initialization.
Random Matrix Theory.
High-dimensional models, including kernel methods and neural networks, can be analyzed by studying the properties of random matrices.
For instance,
El Karoui et al.
; Cheng and Singer (2013); Fan and Montanari (2019) studied the spectral properties of kernel matrix via Taylor series or Hermite polynomials on the nonlinearity, 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 matrix (Karakida et al., 2018; Pennington and Worah, 2018).
Preliminaries: Two-layer Neural Network
Consider the following bias-free two-layer neural network f : R d → R with h hidden units f ( x ) = h (cid:88) i = a i φ (
[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 coeﬃcient of the second layer, and φ : R → R is a Lipschitz continuous activation function with bounded Gaussian moments, i.e. E [ φ ( z ) k ] < ∞ , ∀ k ∈ Z + for z ∼ 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 coeﬃcient vector, X =
[ x 1 , ... x n ] ∈ R d × n for the data matrix, y ∈ R n for the corresponding vector of labels, and Φ X = φ ( X (cid:62) W ) ∈
R n × h 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 in the following regime:
[Equation]
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 gradient ﬂow solution, 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 (ﬂ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 (ﬂow over W ) of such network under ﬁxed Rademacher second layer.
We defer all proofs and details on experiments to appendix.
Warm-up: Linear network
We begin with the simple model with φ ( x ) = x , i.e. Φ X =
X (cid:62) W .
We remark that although the model is linear, the solution obtained by gradient ﬂow on the two-layer model can be diﬀerent 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 coeﬃcients a by minimizing (cid:107) Φ X a − y (cid:107) 22 .
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 coeﬃcients W , gradient ﬂow on the second-layer coeﬃcients a converges to ˆ a =
(Φ X ) † 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]
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 , 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, increasing γ 2 does not inﬂuence the population risk.
In addition, since the model is equivalent to the minimum-norm solution on the input ˆ β min , training the ﬁrst layer always results in smaller or equal population risk compared to training 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 aﬃrmative under certain conditions.
[Figure]
Nonlinear Model: Optimizing the Second Layer
In this section
[Figure]
we analyze the case when the second layer a is learned under ﬁxed W and a nonlinear activation function φ .
We ﬁrst observe that by Lemma 1, the gradient ﬂow ﬁnds the solution ˆ a = Φ † X y .
We again consider the following bias-variance decomposition.
[Equation] [Equation]
We highlight that the variance term does not depend on the teacher model.
Theorem 4.
Given (A1) and w i i .
i .
d . ∼ N (0 , 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]
[Equation] [Equation]
for z ∼ N (0 , 1) and (cid:61) ξ denoting the imaginary part of ξ .
Remark .
If φ is absolutely continuous with bounded Var[ φ ( z )] and φ (cid:48) , then c 1 ≥ c 2 and the equality holds iﬀ φ is linear (Cacoullos et al., 1982).
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 follows from similar techniques used in Cheng and Singer (2013);
Hastie et al.
; but with two modiﬁcations: (i) we handle unnormalized and uncentered activation functions; (ii) we compute the variance of the entire model rather than that of the second layer coeﬃcients.
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.
We empirically found that 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 term.
Nonlinear Model: Optimizing the First Layer
Having observed the double descent phenomenon in optimizing the second layer, in the sequel we consider a two-layer neural network with ﬁxed second layer coeﬃcients satisfying a i ∼ Unif {− 1 / √ h, 1 √ h } , and the ﬁrst layer W is trained using the gradient ﬂow given as
[Equation] [Equation]
which potentially has diﬀerent 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čenko–Pastur, whereas the population covariance is identity.
We consider two diﬀerent initialization settings:
[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 second layer.
In other words, as h increases, the distance traveled by each parameter decreases under both initializations.
The diﬀerence, 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 resembles the initialization W (0) .
5. Vanishing Initialization
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 suﬃciently fast and the gradient being Lipschitz.
The
Lipschitz condition (Lemma 13) can be established under the following assumption.
[Equation]
We thus have the following characterization of the population risk:
Theorem 7.
Given (A1-3) and φ (cid:48) (0) (cid:54) = 0 .
Let T =
log log
h and deﬁne ˆ f = f van ( T ) , then as n, d, h → ∞ we have
[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 suﬃ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 D we heuristically show that under the additional assumption that the data is symmetric, the risk of ReLU network can also be derived and is independent to γ 2 .
5. Non-vanishing Initialization
When initialization is suﬃciently large, the amount each parameter travels to minimize the empirical risk becomes asymptotically negligible compared to 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 solution to the linearized dynamics as the following closed-form:
[Equation]
where J ∈ R n × ( d × h ) is the Jacobian matrix w.r.t.
to the model parameters.
One may naturally expect the double descent phenomenon to appear in this kernel solution, as ∆ exhibits the form of a least squares solution.
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 .
One 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 on the initialization:
(A4)
Symmetric Initialization
: ∀ i ∈ [1 , h ] , ∃ ! j ∈ [1 , h ] s.t. a i w init
i
[Equation]
Theorem 8.
Given (A1-4) and let n, d, h → ∞ ,
[Equation]
where m =
b 21 /b 20 , b 20 = E N (0 , 1) [ φ (cid:48) ( x )] 2 , and b 21 = E N (0 , 1) [ φ (cid:48)
( x ) 2 ] − b 20 .
Note that the population risk is again independent to γ 2 , and thus double descent does not appear for this initialization.
The reason that the risk does not become unbounded is that in the asymptotic limit the pseudo-inverse ( JJ (cid:62) ) † is stable due to the nonlinearity and dh (cid:29) n .
For non-symmetric (i.i.d.) initialization we also observe that the risk is independent to γ 2 , but the bias is no longer upper-bounded by the null risk, as shown in Figure 7.
5. Comparing the Two Initializations.
Figure 3 demonstrates the agreement between theoretical prediction and experimental results.
Observe that in both cases the risk is independent to overparameterization ( γ 2 ), yet the two models have contrasting properties:
Corollary 9.
For any γ 1 ∈ (0 , ∞ ) and activation φ , B ( ˆ f Van ) ≤ B ( ˆ f NV ) ≤ 1 .
On the other hand for all m > 0 , V ( ˆ f NV ) ∈
O (1) , whereas V ( ˆ f Van ) can be arbitrarily large as γ 1 → 1 .
Remark .
m ≥ 0 for all smooth activations φ , and the equality holds if φ is linear.
Intuitively, small initialization enables the model "evolve" more during optimization and align with the data or the target function.
This potentially results in a lower bias, at the expense of overﬁtting more to the noise (high variance).
In contrast, with suﬃciently large initialization the ﬁnal model becomes close to the initialized model, and thus we may expect it to be less "aligned" to the teacher (high bias) but is more stable (lower variance).
[Figure]
In illustrate the diﬀerence in inductive bias
[Figure]
we plot the trajectory of neurons in Appendix A
Figure 5.
Observe that for vanishing initialization the neurons stay close to one another throughout the trajectory
[Figure]
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 our theorems are for linear teacher model
[Figure]
we empirically observe that when the target function is also a two-layer network, the population risk follows the same trend,
i
[Figure]
e. double descent occurs when only the second layer is optimized, as shown in Figure 6.
Discussion and Future Works
We derived the exact population risk of two-layer neural networks in learning a linear target function over Gaussian data with additive noise in high dimensions ( n, d, h tend to inﬁnity at the same rate), and showed that optimizing the ﬁrst or the second layer via gradient ﬂow results solutions with contrasting properties.
Speciﬁcally, the double descent phenomenon is present when the second layer coeﬃcients are optimized, but not when the ﬁrst layer weights are optimized under certain initializations.
In addition, we highlight that the scale of initialization also leads to diﬀerent inductive bias when optimizing the ﬁrst layer.
Note that our analysis only applies to the unregularized least squares objective, and it has been shown that explicit regularization (such as L 2 ) provably stabilizes the singularity at γ 2 → 1 Mei and Montanari (2019); Hastie et al.
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 →which is beyond the regime we consider.
For instance, in two-layer neural networks trained on MNIST, the "cusp" was empirically found at h
= O ( n/d ) Belkin et al.
Thus what we conclude is that in our studied asymptotic limit, the mechanism that provably gives rise to double descent from previous works Hastie et al.
; Belkin et al.
(2019) might not translate to optimizing two-layer neural networks.
To simplify the computation, we rely on a set of assumptions similar to those appeared in
Hastie et al.
(2019), some of which we believe can be relaxed in future work.
For instance, the Gaussian data assumption can be generalized to the standard model with given mean and covariance, and the teacher model may be relaxed to a more general class of functions.
Importantly, the scales of initialization studied in Section 5 are by no means exhaustive, and it remains a challenging problem to characterize the stationary points of the mean-ﬁeld PDE ( 1 /h scaling) under empirical risk.
Another interesting direction would be to understand the inductive bias of discretization (stepsize of gradient descent).
A Additional Figures and Plots
[Figure] [Figure]
[Figure] [Figure]
Summary of the Presence / Absence of Double Descent Singularity in
2nd Layer Trained (RF)
Vanishing Init.
Non-vanishing Init.
Bias γ 1 : No; γ 2 : Yes γ
[Table]
B Background B. Rotational Invariance
The rotational invariance of Gaussian distribution is crucial in our analysis throughout this paper.
The intuition is that for a random Gaussian matrix X and any ﬁxed unitary matrix
U , the distribution of X and UX is the same.
Lemma 10 (Rotational Invariance) .
Denote A ( X ) ∈ R
d × d a matrix function of X ∈ R d × n .
Then if U , A ( UX ) =
UA ( X ) U T for all unitary U, then
[Equation]
[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 i β
= (cid:107) β (cid:107) e
i , where e i is the i -th canonical vector in R d .
Then since U i X ∼ X , we have
[Equation]
Notice that Equation (15) also holds for matrix function A ( X ) that satisﬁes A ( XU ) =
UA ( X ) U T , and moreover, such rotational invariance can be extended to matrix function
A that takes multiple matrices as input.
For example, if for any unitary U , f ( UX, UW ) =
Uf ( X, W ) U T (with valid dimension of U, W, X ), then Equation 15 still holds in the following form:
[Equation]
Marčenko–Pastur Law
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.
When A is Wishart matrix , i.e. A =
X (cid:62) X/n and X ∈ R n × p is random Gaussian matrix with each entry i.i.d.
X ij ∼ N (0 , 1) .
Marčenko and Pastur (1967) showed that as n, p → ∞ and p/n = γ ∈ (0 , ∞ ) , the empirical spectral density µ W p ( dλ ) converges to a limiting density µ MP ( γ ) ( λ ) almost surely:
[Equation]
We say µ MP ( γ ) is the density of the Marčenko–Pastur distribution with support S =
[(1 − √ γ ) 2 , (1 + √ γ ) 2 ] (for γ ≥ 1 ) or S = { 0 } ∪ [(1 − √ γ ) 2 , (1 + √ γ ) 2 ] (for 0 < γ < 1 ).
Note that this implies that the smallest non-zero eigenvalue of A is bounded away from 0 asymptotically.
The explicit form of MP distribution allows us to investigate the asymptotic properties of random matrices.
Generally speaking, due to the almost sure convergence of empirical spectrum density and the Pormanteau theorem, one can translate any bounded continuous function on the empirical spectral density to one on the MP distribution, i.e. almost surely
[Equation]
When γ < 1 , we have the following trace concentration on the inverse Wishart matrix:
[Equation]
We remark that instability of the trace of the invert Wishart matrix as γ → 1 plays an important role in the double descent phenomenon.
C Proof of Main Results C. Proof of Lemma
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 (cid:62) θ (cid:13)(cid:13) 2 2 can be written as
[Equation]
[Equation]
Thus with initialization θ 0 , the solution of this ODE at time t can be written in explicit form
[Equation]
Taking t → ∞ yields the desired result.
[Equation]
C. Proof of Theorem
We compute the bias and variance for diﬀerent 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] [Equation]
Case II: W (cid:62) X has full column rank, i.e. γ 2 < 1 , γ 1 > γ 2 .
Since rank (Φ X )
= h , the minimum-norm solution is
[Equation] [Equation]
Denote W = U Σ V (cid:62) the singular value decomposition of W .
We perform the following block decomposition on X, Σ : Σ
[Equation]
(27) where Σ 0 ∈ R h × h , X 0 ∈ R h × n , X 1 ∈ R ( d − h ) × n .
Note 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] [Equation]
where the last equality follows from Appendix B
[Equation]
Similarly for the bias term we have
[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]
where symmetric arguments are omitted as ( · )
[Equation]
and (i) follows from the rotational invariance argument used in Hastie et al.
and introduced in Lemma
[Equation]
By block decomposition (27),
[Equation]
Therefore
[Equation]
Thus we have obtained that as n, d, h → ∞
[Equation]
Case III
: W (cid:62) X has full row rank, i.e. γ 1 > 1 , γ 2 > 1 .
Similarly, rank (Φ X ) = n and we have the following solution:
[Equation]
Simplifying the variance:
[Equation] [Equation]
where the SVD decomposition of X = U Σ V (cid:62) and the rotational invariance argument is used here.
Using the same block decomposition trick:
[Equation]
[Equation]
where Σ 0 ∈ R n × n , W 0 ∈ R n × h , W 1 ∈ R ( d − n ) × h
[Equation]
and W 0 , W 1 independent.
[Equation]
Hence
[Equation]
We omit the derivation of bias, which can be done in similar fashion.
Combining the three cases yields theorem 2.
[Equation]
C. Proof of Proposition
Given the MSE loss, one can easily 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]
39) which gives the following solution
[Equation]
Thus the gradient ﬂow solution on training the ﬁrst layer is the same as the minimum-norm solution on the input features.
[Equation]
C. Proof of Theorem
Following the bias-variance decomposition (7),
[Equation]
where we deﬁne the expected non-linear Gram matrix K W ∈ R
h × h as
[Equation] [Equation]
and for each entry of K W
[Equation] [Equation]
Random matrix in the form of covariance matrix of nonlinear features has been the object 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 similar to that for nonlinear features in Hastie et al.
with modiﬁcations mentioned below.
Note that in contrast to the linear network derivation in Section C.2, the Gram matrix of a nonlinear activation function is almost surely full-rank, which is speciﬁed in the following lemma from Pennington and Worah (2017)
Lemma 11.
The least singular value of φ ( W (cid:62) X ) is positively bounded from a nonzero constant.
Therefore we only discuss whether Φ is full column rank (Case I) or full row rank (Case II).
Case 1.
h < n .
In this case (41) simpliﬁes into
[Equation]
where the continuity of V ξ at ξ = 0 − is guaranteed by Lemma 11.
From Theorem 1 in Louart et al.
it follows that as n, d, h → ∞
[Equation]
Or equivalently,
[Equation]
where C = tr (cid:0) (ΦΦ (cid:62) − ξI ) − 1 (cid:1) is bounded also due to Lemma 11.
By taking the limit ξ → 0 ,
V = σ
[Equation]
− γ
[Equation]
Case 2.
h > n .
Techniques used in the current proof are largely borrowed from Cheng and Singer (2013); Hastie et al.
(2019), but with one modiﬁcation: our non-linearity function does not have zero-expectation under Gaussian distribution, i.e. E x ∼N (0 , 1) [ φ ( x )]
= 0 , which complicates the analysis since the oﬀ-diagonal entries of the activation Gram matrix is no longer zero-centered.
For simplicity we mainly adhere to the notations in Hastie et al.
and rescale w i ∼ N ( 0 d , I d /d ) .
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 (59); this translates the calculation of V ξ into the calculation of the Stieltjes transform of ˜ A (60), (65), (67).
C. Deriving variance V for h > n
Step 1.
An equivalent expression for variance using Stieltjes transform For notational simplicity we omit the magnitude of noise σ :
[Equation] [Equation]
and due to the same continuity argument as in Case 1 we have
[Equation]
where S = φ ( W (cid:62) X ) / √ n , ξ ∈ C and (cid:61) ξ > 0 or ξ < 0 .
We decompose the normalized feature matrix S
= φ ( W (cid:62) X ) / √ n as
[Equation] [Equation]
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]
We deﬁne ˜ K W as
[Equation]
where Q ∈ R h × h with Q ij = w (cid:62) i w j for oﬀ-diagonal terms and 0 for diagonal terms, and r = E
[ φ ( x ) 2 ] − E [ φ ( x )] 2 , s =
E [ φ ( x )] 2 , t = E
[ xϕ ( x )] 2 are constants of the non-linear activation φ .
By the same argument as in Lemma 13 of Hastie et al.
(2019) one can show that when n, d, h → ∞
[Equation]
Writing the trace explicitly (denote eigenvalues λ i
=
φ 2 i , and φ n +1 = · · · = φ h = 0 ), we have
[Equation] [Equation]
Since the positive support of spectrum λ is lower bounded by a constant by Lemma 11 and the density at 0 is 1 − n/h = 1 − γ − 1 2 , we have
[Equation]
[Equation] [Equation]
where we deﬁne µ n ( x ) and its positive part as
[Equation]
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 matrix derivative gives − ∂
[Equation]
Denote the limit
[Equation] [Equation]
we have
[Equation]
For simplicity we deﬁne the following function on ξ :
[Equation] [Equation]
and denote
[Equation]
Comparing (58) and (66) yields the following equivalence:
[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 (60)(65)(67).
We deﬁne a modiﬁed matrix A n , which corresponds to subtracting the oﬀ-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
[Equation]
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 ,
[Equation]
Proof.
By deﬁnition
[Equation]
which is a rank-2 matrix.
By theorem A.43 from Bai and Silverstein (2010), which characterizes the eﬀect of ﬁnite-rank perturbation on the e.s.d.
of random matrices:
[Equation] [Equation]
where F M is the empirical spectral distribution of M ∈ R n × n , i.e. F M ( x ) = 1 n n (cid:88)
i =
I λ i ( M ) <x .
The claim follows from the Stieltjes continuity theorem in Section 2.4 in Tao (2012).
(cid:3)
To calculate the Stieltjes transform m n of matrix A n , we take advantage of the block structure and deﬁne the following terms:
[Equation] [Equation]
It’s easy to see that
[Equation]
In the following equations we omit the subscript n , as well as the dependency on ρ, ς, τ .
Following Hastie et al.
(2019), we rewrite A =
A n as a [( N − 1) + 1] ×
[( N − 1) + 1] matrix, i.e.
A
[Equation]
(76) where A ∗ is a ( N − 1) ×
( N − 1) matrix with last column and row of A removed:
[Equation]
and
[Equation]
Hence by the block matrix inverse formula
[Equation]
We have
[Equation]
We perform the orthonormal decomposition of the non-linear function ϕ (under Gaussian density), introduced in Cheng and Singer (2013).
[Equation]
where a 1 = E x ∼N (0 , 1) [ xϕ ( x )] .
Orthonormal decomposition of the function together with the decomposition of vectors introduced below gives an asymptotic description of m 2 ,n expressed in (80).
Speciﬁcally, for each w i (1 ≤ i ≤ h ) , we perform the following orthonormal decomposition (along the direction of x n and the direction of ˜ w i which is perpendicular to x n ,
i
[Equation]
e. x (cid:62) n ˜ w i = 0 ):
[Equation]
We thus have
[Equation]
and for
[Equation]
In addition
[Equation]
S ik
[Equation]
Hence
[Equation]
By similar argument as in (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 (80) we have
[Equation]
We bound each term u, v , S to compute m 2 ,n .
For u
[Equation]
where b = E x ∼N (0 , 1) [ ϕ ( x ) 2 ] =
E x ∼N (0 , 1) [( φ ( x ) − E φ ( x )) 2 ]
= r .
Due to concentration of measure we have that as n, h, d → ∞
[Equation]
E a
[Equation]
ξ
[Equation]
[Equation] [Equation]
And for v (note that U is dependent on a )
[Equation]
where
[Equation]
Hence by similar argument we have when n
[Equation]
d
[Equation]
p → ∞ ,
[Equation]
And ﬁnally for S ,
[Equation]
And hence as n
[Equation]
d
[Equation]
h → ∞ ,
[Equation]
Therefore by combining (89), (92), (94)
[Equation]
m 2 ,n ( ξ, ρ, τ )
[Equation]
Similarly we can calculate m 1 ,n ( ξ, ρ, τ ) as m 1 ,n ( ξ, ρ, τ ) →
[Equation]
The proof of the uniqueness of solution by (96), (95) follows from Section B.1 of Hastie et al.
and is omitted.
[Equation]
C. Proof of Corollary
In this section we take the limit γ 1 → ∞ and check the correctness of Theorem 4.
In this case (8), (9) simplify to
[Equation]
Recall that
[Equation]
by taking the derivative we have
[Equation]
Observe that (97), (98) constitutes a set of implicit functions.
Thus diﬀerentiating the two functions with respect to τ, ρ and then substitute by ρ = τ = 0 gives
[Equation]
Hence by (66), (67) we obtain the asymptotic variance:
[Equation]
Combining the case where γ 2 < 1 in Theorem 4 completes the proof.
(cid:3)
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. Proof of CorollaryC.7.
Unbounded Bias when h = n
From the bias-variance (7), the bias B is derived as
[Equation]
where
[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 ) , we exchange W and X/ √ d and get
[Equation]
where in Section C.5 we have showed (107) is unbounded when n → ∞ .
Moreover, by (52) and Weyl’s theorem we have λ min ( K W ) =
O (1) , and thus d − 1 tr ( Q 1 ) is unbounded.
For d − 1 tr ( Q 2 ) , we have
[Equation]
To sum up, for n → ∞ and γ 2 → 1 we have B → ∞ .
[Equation]
C.7. Bounded Bias when γ 2 >
Recall the expression of the bias term when h > n
[Equation]
) where
[Equation]
Therefore we have
[Equation]
(cid:62)
[Equation]
= O (1) ·
σ
[Equation]
2min (
[Equation]
( X (cid:62)
[Equation]
))tr ( K W ) =
O (1) · O ( n − 1 ) · O ( n ) =
O (1)
[Equation]
and similarly
[Equation]
Therefore B =
O (1) < ∞ when 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. Proof of TheoremC.8. Verifying the Lipschitz Condition
We ﬁrst show that the gradient of the empirical risk is Lipschitz in the Frobenius norm.
Lemma 13.
Given (A1-3), the gradient of the empirical risk deﬁned in (11) is L f -Lipschitz w.r.t.
[Equation]
i.
[Equation]
Proof.
For simplicity we assume n, d, h to be even and let d 0 = d/ 2 , n 0 = n/
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] [Equation]
We show the Lipschitz condition for W + .
Write the derivative of W + as ∂L
[Equation]
Or weight-wise (let w be a column of W + ),
[Equation]
We upper-bound the spectral norm of the derivative of each term.
[Equation]
Note that the spectral norm of the gradient matrix satisﬁes
[Equation]
where we have utilized the fact that λ max ( X (cid:62) X ) =
O ( d ) and | y | = | β (cid:62) x
+ (cid:15) | ∈ O ( √ d ) , as well as the boundness assumption on activation derivative (A3).
Hence we have
[Equation]
With a similar claim, for the second term g 2 ( w , W + ) we have
[Equation]
Therefore
[Equation]
Deﬁne w (cid:48) ∈ W + \{ w } as a diﬀerent column of W + , observe that
[Equation]
Combining the two equations above and sum over all w (cid:48)
∈ W + \{ w } yields
[Equation]
[Equation]
Similarly for G 3 ( w , W − ) we have
[Equation]
Applying the results above weight-wise
[Equation]
we get
[Equation] [Equation]
Note that the same argument holds for W − .
We have thus shown that for some constant
L f > 0 such that
[Equation]
C.8. Defining Gradient Flows
Recall the simpliﬁed deﬁnition of the network in the previous section:
[Equation]
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]
with vanishing initialization w Oi (0) ∼ N ( 0 , I/dh 1+ (cid:15) ) .
We now deﬁne the same ﬂow (123) from exact zero initialization w Di (0) = 0 termed GFdouble .
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 ) 1 (cid:62) where w D ± ( t ) admits the following dynamics:
[Equation]
Lastly, we deﬁne the GF-single with solution denoted as w ± = w S ± ( t ) :
∂ w S +
[Equation]
x i
[Equation] [Equation]
i
[Equation]
127) from zero initialization w Di (0) = 0 .
This can be seen as the 1st order Taylor expansion of the non-linearity φ at the origin of GF-double.
C.8. Step 1. From GF-double to GF-single
Step 1.1.
Bounding the Diﬀerence in Risk.
Due to low rank property of solution of
GF-single and GF-double, in this subsection we slightly abuse the notation and deﬁne
[Equation]
Consider the diﬀerence of the population risk of two models with parameters w S ± and w D ± :
[Table]
[Equation]
in which (i) is due to the fact that both R S and R D are bounded at initialization
[Equation]
and both ﬂows do not result in unbounded risk for γ 1 ∈ (0 , 1) ∪ (1 , ∞ ) as shown in the following section.
We now show that this diﬀerence between the trajectories is asymptotically vanishing for w + , and w − follows the same argument.
[Equation]
+
[Equation]
=
[Equation]
= O
[Equation]
Given that up to time t the norm of the weights (cid:13)(cid:13) w D (cid:13)(cid:13) 2 =
O (1 / √ h ) , the error from Taylor expansion E + ≤
O (1 /h ) t is derived below.
[Equation]
t
[Equation] [Equation]
=
[Equation]
O
[Equation]
(cid:48) (
[Equation]
(cid:62)
[Equation]
i )
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
(cid:88) i =
[Equation]
O (cid:0) (
[Equation]
(cid:62)
[Equation]
(cid:1) · x i (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
ds
[Equation]
[Equation]
(cid:13)(cid:13)(cid:13) (cid:13) (cid:13)
[Equation]
(cid:88)
i
[Equation]
O (cid:0)
[Equation]
(cid:62)
[Equation]
(cid:1)
[Equation] [Equation]
(cid:62)
[Equation]
i
[Equation]
(cid:62) i (cid:13)(cid:13)(cid:13) (cid:13) (cid:13)
[Equation]
ds
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
[Equation]
By Gronwall’s theorem we have
[Equation]
A similar inequality can also derived for w − .
Hence we obtain an upper-bound on the diﬀerence between the empirical risk of the two models at time t :
[Equation]
Now we have established that the diﬀerence in the parameters of GF-Single and GF-double asymptotically vanishes if the norm of the weights is of order O (1 / √ h ) .
In the following subsection we show that the closeness of the two ﬂows implies that (cid:13)(cid:13)
[Equation]
O (1 / √ h ) , hence the assumption of the norm of w D is satisﬁed.
Step 1.2.
Solution of GF-single.
A basic observation of GF-single is that w S + ( t ) = − w S − ( t ) for all t .
Therefore, The solution of GF-single can be written in analytical form:
[Equation]
when d < n , or otherwise
[Equation]
when d > n .
We ﬁrst assume that d < n , under which the stationary point is
[Equation]
Note that the stationary point satisﬁes the norm assumption above.
We compute the diﬀerence 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 .
If we take t = log log h → ∞ , combining (133) and (137) yields
[Equation]
→ 0 .
Equation (138) for case d > n follows a similar proof.
C.8. Step 2. From GF-original to GF-double
In this section we compare GF-original with GF-double.
Note that the two ﬂows diﬀer only at initialization: vanishing initialization w O i (0) ∼ N ( 0 , I/dh 1+ (cid:15) )
v.s. zero initialization w Di
[Equation]
= 0 .
By the Lipschitz condition on the empirical gradient (Lemma 13) we have
[Equation]
And hence by Gronwall’s lemma one obtains:
[Equation]
Therefore
[Equation]
Taking
= log h , together with the same argument in Step 1 yields
[Equation]
C.8. Putting things together
By (138) and (142), we know that taking t = log log h ,
[Equation]
Finally observe that R S ( ∞ ) is the risk of the minimum-norm solution on the input features discussed in Section 3.
[Equation]
C. Proof of TheoremC.9. The Kernel Linearization
Denote ω = vec( W ) =
vec([ W + , W − ]) , and ω 0 = vec( W init ) .
Deﬁne
[Equation]
144) which is the kernel matrix of the neural tangent kernel Jacot et al.
Du et al.
Write ˜ y ( t ) =
f ( X, t ) ∈ R n and its evolution:
[Equation]
145) and the corresponding linearized ﬂow:
[Equation]
with corresponding parameters ˜ ω
=
Vec( ˜ W ) and ¯ ω =
Vec( ¯ W ) , respectively.
Previous works (e.g. Du et al.
; Oymak and Soltanolkotabi (2019)) have proved (non-asymptotically) that the two paths (145) and (146) are close if the model is overparameterized,
i.e. h ∈ poly( n ) , under no assumptions on the teacher model.
In our asymptotic setup and assumptions (A1)(A2), we show that similar conclusion holds without signiﬁcant overparameterization.
We employ an argument similar to (Du et al., 2018, Theo. 3.2) to bound the distance traveled by the parameters.
From Corollary 16 (i.e. λ min ( K (0))
=
O ( d ) ) and Lemma 13 we know that for (cid:107) W ( t ) − W (0) (cid:107) 2 F
= O (1) , λ min ( K ( t ))
=
O ( d ) .
Therefore, assume that up to time
T the parameters satisfy (cid:107) W ( t ) − W (0) (cid:107) 2 F
= O (1) , from (Chizat and Bach, 2018b, Lemma
B1,B.2) and Lemma 13 we have
[Equation]
for 0 ≤ t ≤ T and some constants C 1 , C 2 > 0 .
Since (cid:107) ˜ y (0) − y (cid:107)
= O ( n ) , setting T = n (cid:15) − 1 log( n ) for small (cid:15) ensures that (cid:107) ˜
y ( t ) − y (cid:107) 2 → 0 as n → ∞ .
The Frobenius norm of the gradient can be bounded as follows:
[Equation]
Integrating the gradient yields (cid:13)(cid:13)(cid:13) ˜ W ( t ) − ˜ W (0)
(cid:13)(cid:13)(cid:13) F = O (1) .
Thus the distance traveled by ˜ W from initialization to the optimum indeed satisﬁes the assumption above.
The same argument trivially applies to (cid:13)(cid:13) ¯ W ( t ) − ¯ W (0) (cid:13)(cid:13) F .
Diﬀerence between parameters of the two trajectories can also be obtained:
[Equation]
¯ W ( s )) ∂W (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
F ds ( i )
[Equation]
O ( t ) , (149) where (i) is due to Lemma 13 and that in both trajectories the distance traveled by the parameters is O (1) .
We now bound the diﬀerence of ˜ f and ¯ f on Gaussian input:
[Equation]
Take T
= O ( n (cid:15) − 1 log( n )) at which the empirical risk tends to 0 (i.e. the ﬂow converges), we have (cid:13)(cid:13) (cid:13) ˜
f
( x ) − ¯ f ( x ) (cid:13)(cid:13) (cid:13) =
O ( √ dt )
=
O ( n (cid:15) − 1 / 2 log n ) → 0 .
Thus the diﬀerence between the population risk of ˜ f
and ¯ f is also asymptotically vanishing.
In the following subsection we compute the risk of the linearized (kernel) model ¯ f .
C.9. Computing the Kernel Risk
Given input X ∈ R d × n and label y
= β (cid:62) X + ε ∈ R 1 × n , gradient ﬂow on the tangent kernel solves the following equation of ω :
[Equation] [Equation]
where ∂ f ( X ; ω 0 )
/∂
ω is a dh × n matrix with each column ∂f ( x i ; ω 0 ) /∂
ω .
This corresponds to the kernel linearization.
Note that for n → ∞ and γ 1 , γ 2 ∈ (0 , ∞ ) , dh > n trivially holds, and by Corollary 16 we know that the least squares solution is given by
[Equation]
And the population risk can be written as
[Equation]
Where a bias-variance decomposition has been made here
[Equation]
and for simplicity we deﬁne
[Equation]
C.9. Approximating the Kernel Matrix
Following Cheng and Singer (2013) we utilize the orthonormal decomposition of φ (cid:48)
( x ) in
L 2 ( R , dµ G ) .
Denote b 20 = E x ∼N (0 , 1) [ φ (cid:48) ( x )] 2 , and b 21 =
E x ∼N (0 , 1) [ φ (cid:48) ( x ) 2 ] − b 20 .
We have the orthogonal decomposition of φ (cid:48)
[Equation]
) where E [ φ (cid:48)⊥ ( x )] = 0 .
We have the following lemmas.
Lemma 14 (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ε
[Equation]
(155) and with probability 1 − e − c (cid:48) nε
[Equation]
Proof.
When i (cid:54) =
j (i.e. Equation (155)).
We have
[Equation]
The expectation H ( x i , x j ) =
x (cid:62) i x j E w (cid:104) φ (cid:48) ( w (cid:62) x i ) φ (cid:48) ( w (cid:62) x j ) (cid:105) can be seen as the expected tangent kernel of nonlinear activation function studied in Arora et al.
(2019b).
Moreover, due to the assumed boundedness of φ (cid:48) ( x ) , by Hoeﬀding’s inequality we have
[Equation]
In addition, by the concentration of x (cid:62) i x j and (cid:107) x i (cid:107) 22 in high dimensions, 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 − c 3 dε 2 , the orthonormal decomposition of φ (cid:48) ( x ) =
b 0 x + φ (cid:48)⊥ ( x ) in L 2 ( R , µ G ) leads to the following linear approximation of the kernel matrix
[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 .
(cid:3)
Corollary 15 (Approximation of ˆ u ) .
For arbitrary m > 0 , with probability 1 − e − c log m d
[Equation]
where ˜ u
= b 20 x (cid:62) X .
Proof.
Taking ε = log m d/d together with Lemma 14 yields the desired result.
[Equation]
Corollary 16 (Approximation of ˆ K X ) .
With probability 1 − e − c log m d
[Equation]
where ˜ K X = b 20 X (cid:62) X + b 21 dI .
Proof.
Also by directly applying Lemma 14.
[Equation]
Remark .
˜
K X is always positive deﬁnite and from Vershynin (2010) one can show that λ min ( ˆ K X )
=
O ( d ) .
Remark .
For φ ( x ) = ReLU( x ) , b 20 = b 21 = 1 / 4 .
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 if and only if φ is linear.
C.9. The Bias Term
With these approximation above we proceed to calculating (152)
[Equation]
We ﬁrst bound the error in substituting ˆ u with ˜ u :
[Equation]
where (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 ) by Vershynin (2010).
Therefore we have as n, d, h → ∞
[Equation]
By taking expectation over x and the rotational invariance argument similar to Hastie et al.
[Equation]
we have
[Equation]
In addition
[Equation]
we bound the error in substituting ˆ K X by ˜ K X deﬁned in (162):
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
Combining these two formulas in (166) yields
[Equation]
By the same technique on Marčenko–Pastur showed in Section B
[Equation]
2 we obtain
[Equation]
where m =
b − 2 0 b 21 .
[Equation]
C.9. The Variance Term
Similarly, for the variance we utilize the approximation
[Equation]
Speciﬁcally, we have
[Equation]
and
[Equation]
By combining the two approximations above we know that as n, d, p → ∞
[Equation]
Therefore
[Equation]
where m =
b − 2 0 b 21 .
C.9. Putting Things Together
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
[Figure]
D Additional Results D. Risk of ReLU Network under Symmetric Data
If the dataset is symmetric and the label is noiseless, that is
Symmetric Data: ∀ i ∈
[Equation]
, n ] ,
[Equation] [Equation]
then the population risk of the gradient ﬂow solution can be given explicitly for certain nonlinearities:
Proposition 17.
Given the same conditions as in Proposition 7, if the data is symmetric and the nonlinearity satisﬁes φ (cid:48) ( x ) + φ (cid:48) ( − x ) =
C for constant C, then as n, d, h → ∞
R
[Equation]
( ˆ f ) → 0;
[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 8.
In addition, the risk is upper-bounded by the null risk for all γ 1 .
We remark that the symmetry assumption does not hold in general for i.i
[Equation]
d. samples from symmetric distributions, and
Figure 8 demonstrates that the additional condition alters the dependence of the risk on γ 1 .
Proof.
Without loss of generality assume X =
[ X 0 , − X 0 ] .
Then by (124), (125) one have
[Equation]
Therefore from exact zero initialization one can show that
[Equation]
for nonlinear activations satisfying φ ( x ) − φ ( − x ) = x , such as ReLU and SoftPlus ,
[Equation]
= 0 .
[Equation] [Equation]
And therefore the gradient ﬂow of w + is
[Equation]
The ﬂow of w − follows from symmetry.
It is easy to show that the stationary point satisﬁes
[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.
E Experiment Setup
Optimizing the Second Layer.
We compute the minimum-norm solution by directly solving the linear inverse problem.
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 set of ( γ 1 , γ 2 ) we average across 50 draws to obtain the risk.
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 diﬀerent initializations.