E ARLY S TOPPING IN D EEP N ETWORKS : D OUBLE D ESCENT AND H OW TO E LIMINATE IT
Anonymous authors
Paper under double-blind review
A BSTRACT
Over-parameterized models, such as large deep networks, often exhibit a double descent phenomenon, where as a function of model size, error ﬁrst decreases, increases, and decreases at last.
This intriguing double descent behavior also occurs as a function of training epochs and has been conjectured to arise because training epochs control the model complexity.
In this paper, we show that such epoch-wise double descent occurs for a different reason: It is caused by a superposition of two or more bias-variance tradeoffs that arise because different parts of the network are learned at different epochs, and mitigating this by proper scaling of stepsizes can signiﬁcantly improve the early stopping performance.
We show this analytically for i) linear regression, where differently scaled features give rise to a superposition of bias-variance tradeoffs, and for ii) a wide two-layer neural network, where the ﬁrst and second layers govern bias-variance tradeoffs.
Inspired by this theory, we study two standard convolutional networks empirically and show that eliminating epoch-wise double descent through adjusting stepsizes of different layers improves the early stopping performance.
I NTRODUCTION
Most machine learning algorithms learn a function that predicts a label from features.
This function lies in a hypothesis class, such as a neural networks parameterized by its weights.
Learning amounts to ﬁtting the parameters of the function by minimizing an empirical risk over the training examples.
The goal is to learn a function that performs well on new examples, which are assumed to come from the same distribution as the training examples.
Classical machine learning theory says that the test error or risk as a function of the size of the hypothesis class is U-shaped: a small hypothesis class is not sufﬁciently expressive to have small error, and a large one leads to overﬁtting to spurious patterns in the data.
The superposition of those two sources of errors, typically referred to as bias and variance, yields the classical U-shaped curve.
However, increasing the model size beyond the number of training examples can decrease the error again.
This phenomena, dubbed “double descent” by Belkin et al. (2019a) has been observed as early as 1995 by Opper (1995), and is relevant today because most modern machine learning models, in particular deep neural networks, operate in the over-parameterized regime, where the error often decreases again as a function of model size, and where the model is sufﬁciently expressive to describe any data, even noise.
Interestingly, this double descent behavior also occurs as a function of training time, as observed by Nakkiran et al.
(2020a)
and as illustrated in Figure 1.
The left panel of Figure 1 shows that as a function of training epochs, the test error ﬁrst decreases, increases, and then decreases again.
It is important to understand this so-called epoch-wise double descent behavior to determine the early stopping time that gives the best performance.
Early stopping, or other regularization techniques, are critical for learning from noisy labels (Arpit et al., 2017;
Yilmaz & Heckel, 2020).
Nakkiran et al.
(2020a) conjectured that epoch-wise double descent occurs because the training time controls the “effective model complexity”.
This conjecture is intuitive, because the model-size, and thus the size of the hypothesis class, can be controlled by regularizing the empirical risk via early stopping the gradient descent iterations, as formalized in the under-parameterized regime by Yao et al.
; Raskutti et al.
; B¨uhlmann & Yu (2003).
Speciﬁcally, limiting the number
[Figure]
t r a i n /t e s t
e rr o r test train bias variance [-
 [Figure]-]{+10101010100 .0 .
+}
.t iterations r
{+i s k bias-variancebias-variancerisk=
+}
+ of gradient descent iterations ensures that the functions parameters in around the parameters.
While this conjecture might be true for certain problem setups, it is not consistent with our empirical observation for the 5-layer CNN studied by Nakkiran et (2020a) empirically measured in Figure model would imply on this experiment).
In this paper, we show empirically and theoretically that epoch-wise double descent—at least in the setups we observed it— arises for a different reason: It is explained by a superposition of biasvariance tradeoffs, as illustrated for a toy-regression example in the right panel of Figure 1.
If the risk can be decomposed into two U-shaped bias-variance tradeoffs with minima at different epochs/iterations, then the overall risk/test error has a double descent behavior.
We also note that epoch-wise double descent is not a phenomena tied to over-parameterization.
Both under- and overparameterized models can have epoch-wise double descent as we show in this paper.
1. C ONTRIBUTIONS
The goal of this paper is to understand the epoch-wise double descent behavior.
Our main ﬁnding is that epoch-wise double descent can be explained as a superposition of bias variance tradeoffs, and arises naturally in some standard neural networks because parts of the network are learned faster than others.
Our contributions are as follows:
First, we consider a linear regression model and theoretically characterize the risk of early stopped least squares.
We show that if features have different scales, then the early stopped least squares estimate as a function of the early stopping time is a superposition of bias-variance tradeoffs, which yields a double descent like curve (see Figure 1, right panel).
Second, we characterize the early stopped risk of a two-layer neural network theoretically and show that it is upper bounded by a curve consisting of over-lapping bias-variance tradeoffs that are governed by the initializations and stepsizes of the two layers.
The initialization scales and stepsizes of the weights in the ﬁrst and second layer determine whether double descent occurs or not.
We provide numerical examples showing how epoch-wise double descent occurs when training such a two-layer network on data, and how it can be eliminated by scaling the stepsizes of the layers accordingly.
Third, we study a standard 5-layer convolutional network as well as ResNet-18 empirically.
For the 5-layer convolutional network we ﬁnd—similarly as for the two-layer model—epoch-wise double descent occurs because the convolutional layers (representation layers) are learned faster than the ﬁnal, fully connected layer.
Similarly, for ResNet-18, we ﬁnd that later layers are learned faster than early layers, which again results in double descent.
In both cases, epoch-wise double descent can be eliminated through adjusting the stepsizes of different coefﬁcients or layers.
In summary, we provide new examples on when epoch-wise double descent occurs, as well as analytical results explaining epoch-wise double descent theoretically.
Our theory is constructive in that it suggests a simple and effective mitigation strategy: scaling stepsizes appropriately.
We also note that epoch-wise double descent should be eliminated by adjusting the stepsizes and/or the initialization, because this often translates to better overall performance.
1. R ELATED WORKS
There is a large number of works that have studied early stopping theoretically.
Intuitively, each step of an iterative algorithm reduces the bias but increases variance.
Thus early stopping can ensure that neither bias nor variance are too large.
A variety of papers Yao et al.
; Raskutti et al.
B¨uhlmann & Yu (2003);
Wei et al.
(2019) formalized this intuition and developed theoretically sound early stopping rules.
Those works do not, however, predict when a double descent curve can occur.
A second, more recent line of works, studies early stopping from a different perspective, namely that of gradient descent ﬁtting different components of a signal or different labels at different speeds.
For linear least squares, the data in the direction of singular vectors associated with large singular values is ﬁtted faster than that in the direction of singular vectors associated with small singular values.
Advani & Saxe (2017) have shown this for a linear least squares problem or stated differently, a linear neural network with a single layer.
Li et al.
; Arora et al.
(2019) have shown that this view explains why neural network often ﬁt clean labels before noisy ones, and Heckel & Soltanolkotabi (2020b) have used this view to prove that convolutional neural networks provably denoise images.
Our theoretical results for neural networks build on a line of works that relate the dynamics of gradient descent to those of an associated linear model or a kernel method in the highly overparameterized regime Jacot et al.
; Lee et al.
Arora et al.
Du et al.
; Oymak & Soltanolkotabi (2020); Oymak et al.
; Heckel & Soltanolkotabi (2020b).
We use the same proof strategy as those papers to characterize the early stopping performance of a simple two-layer neural network, but in contrast to those earlier works, we develop early stopping results and optimize over both the weights in the ﬁrst and second layer, as opposed to only optimizing over the weights in the ﬁrst layer.
That is important, because we want to demonstrate that initialization and stepsize choices of different layers lead to different bias-variance tradeoffs.
Next, we note that there is an emerging line of works that theoretically establishes double descent behavior as a function of the model complexity (e.g., measured by the number of parameters) for linear regression Hastie et al.
; Belkin et al. (2019b), for random feature regression Mei &
Montanari (2019); d’Ascoli et al.
(2020), and for binary linear regression Deng et al.
A number of recent theoretical double-descent works Jacot et al.
; Yang et al.
(2020); d’Ascoli et al.
(2020) have decomposed the risk into bias and variance terms, and studied their behavior.
Those works demonstrate that the bias typically decreases as a function of the model size, and the variance ﬁrst increases, and then decreases, which can yield a double-descent behavior.
As we demonstrate in Appendix B.2, the epoch-wise double descent phenomena for the standard CNN studied by Nakkiran et al.
(2020a) cannot be explained with this observation: The variance is increasing as a function of training epochs, as opposed to being unimodal.
Finally, our suggestion to mitigate epoch-wise double descent with step-size adaption and early stopping is a form of regularization.
Related work for model-wise double descent shows that model-wise double descent can be mitigated with ( (cid:96) 2 ) regularization Nakkiran et al.
(2020b), and (cid:96) 2 regularization and early stopping are strongly related ( ? ).
E ARLY - STOPPED GRADIENT DESCENT FOR LINEAR LEAST SQUARES
We start by studying the risk of early stopped gradient descent for ﬁtting a linear model to data generated by a Gaussian linear model.
Our main ﬁnding is that the risk as a function of the early stopping time is characterized by a superposition of U-shaped bias-variance tradeoffs, and if the features of the Gaussian linear model have different scales, those bias-variance tradeoff curves add up to a double descent shaped risk curve.
We also show that the performance of the estimator can be improved through double descent elimination by scaling the stepsizes associated with the features.
2. D ATA MODEL AND RISK
Consider a regression problem, and suppose data is generated from a Gaussian linear model as y = (cid:104) x , θ ∗ (cid:105)
+ z, where x ∈ R d is a zero-mean Gaussian feature vector with diagonal co-variance matrix
Σ = diag ( σ 21 , . . .
, σ 2 d ) , and z is independent, zero-mean Gaussian noise with variance σ 2 .
We are given a training set D
= { ( x 1 , y 1 ) , . . .
, ( x n , y n ) } consisting of n data points drawn iid from this Gaussian linear model.
We consider the class of linear estimators parameterized by a vector
ˆ θ ∈ R d , which we estimate based on the training data D .
The linear estimator predicts the label associated with a feature vector x as ˆ y =
x
T ˆ θ .
The (mean-squared) risk of this estimator is
[Equation]
where expectation is over an example ( x , y ) drawn independently (of the training set) from the underlying linear model.
2. E ARLY - STOPPED LEAST SQUARES ESTIMATE
We consider the estimate based on early stopping gradient descent applied to the empirical risk
[Equation]
We initialize gradient descent with θ 0 = 0 and iterate, for t = 1 , 2 , . . .
, with updates θ t +1 =
θ t − 12 diag ( η ) ∇ ˆ R ( θ t ) , where diag ( η ) is a diagonal matrix containing the stepsizes η i > 0 associated with each of the features as entries.
Note that we allow for different stepsizes for all of the features.
In the following, we study the properties of the iterates t , i.e., θ t .
2. R ISK OF EARLY STOPPED LEAST SQUARES
The main result of this section is that in the underparameterized regime, where d (cid:28) n , the risk of gradient descent after t iterations, R ( θ t ) , is very close to a risk expression deﬁned as
[Equation]
as formalized by the theorem below.
We focus on the underparameterized regime here, because in the over-parameterized regime our estimator cannot achieve small risk in general.
In Section 3 we study a more general setting in the overparameterized regime.
Theorem 1.
Suppose that the stepsizes obey η i ≤ 1 σ 2 i , for all i = 1 , . . .
, d .
With probability at least 1 − 2 d − 5 − 2 de − n/ 8 − e − d − 2 e − 32 over the random training set generated by a linear Gaussian model with parameters θ ∗ and Σ , the difference of the early stopped risk and the risk expression in (1) at iteration t is at most
[Equation]
Here, c is a numerical constant.
Theorem 1 guarantees that with high probability the risk R ( θ t ) is well approximated by the risk expression ¯ R (˜ θ t ) , provided the model is sufﬁciently underparameterized (i.e., d/n is small).
As a consequence, the risk of early stopped least-squares is a superposition of U-shaped biasvariance tradeoffs, and if the features are differently scaled, this can give rise to epoch-wise double descent.
To see this, ﬁrst note that the terms
U i ( t ) in the risk expression (1) are U-shaped as a function of the early stopping time t , because σ 2 i ( θ ∗ i ) 2 (1 − η i σ 2 i ) 2 t decreases in t and σ 2 n (1 − (1 − η i σ 2 i ) t )increases in t ; see Figure 2a for an example.
The minima of the individual U-shaped curves U
i
( t ) depend on the product of the stepsize and the i -th features’ variance, η i σ 2 i ; the larger this product, the earlier (as a function of the number of iterations, t ) the respective U-shaped curve reaches its minimum.
Therefore, if we add up two (or more) such U-shaped curves with minima at different iterations, the resulting risk curve can have a double descent shape (again, see Figure 2a).
This establishes our claim that differently scaled features can give rise to epoch-wise double descent.
Finally we note that the reason why we refer to the U-shaped curves as bias-variance tradeoffs, is that the terms (cid:80) di =1 σ 2 i ( θ ∗ i ) 2 (1 − η i σ 2 i ) 2 t and (cid:80) di =1 σ 2 n (1 − (1 − η i σ 2 i ) t ) 2 in the risk expression (1) are approximately equal to the bias and the variance of the model θ t in the standard textbook biasvariance decomposition of the risk, see Appendix A.2 for a detailed discussion.
[Figure]
Improving performance by eliminating double descent: Epoch-wise double descent can be eliminated by properly scaling the stepsizes associated with each of the features, so that the minima of the individual bias-variance tradeoffs overlap at the same iteration ˜ t :
Proposition 1.
Pick an optimal early stopping time ˜ t ≥ 1 .
The minimum of the risk expression min η 1 ,...,η d min t ¯ R
(˜ θ t ) is achieved at iteration ˜ t by choosing the stepsizes pertaining to the features as η
[Equation]
[Table]
Elimination of double descent is illustrated in Figure 2b.
By eliminating double descent optimally so that all the individual bias-variance tradeoffs U i ( t ) achieve their minima at the same early stopping point t , we achieve the lowest overall risk at the optimal early stopping point.
Thus eliminating double descent is important for optimal performance.
In practice we typically do not know the variances of the features and therefore may not be able to optimally choose the stepsizes.
However, we may be able to mitigate double descent sub-optimally by treating the stepsizes as hyperparameters.
E ARLY STOPPING IN TWO LAYER NEURAL NETWORKS
In this section, we establish a bound on the risk of a two-layer neural network and show that this bound can be interpreted as a super-position of U-shaped bias-variance tradeoffs, similar to the expression governing the risk of the linear model from the previous section.
The risk of the twolayer network is governed by two associated kernels pertaining to the ﬁrst and second layer, and the initialization scale and stepsizes of the weights in the ﬁrst and second layer determine whether double descent occurs or not.
We also show in an experiment that if double descent occurs, it can be eliminated by adapting the stepsizes of the two layers.
Network model: We consider a two-layer neural network with ReLU activation functions and k neurons in the hidden layer: f W , v ( x ) = 1 √ k relu( x T W ) v . .
Here, x ∈ R d is the input of the network, W ∈ R d × k and v ∈ R k are the weights of the ﬁrst and second layer.
Moreover, relu( z ) = max( z, 0) is the rectiﬁed linear unit, applied elementwise.
Data model: We assume that we are given a training set D
= { ( x 1 , y 1 ) , . . .
, ( x n , y n ) } with examples ( x i , y i ) drawn iid from some joint distribution.
For convenience, we assume that the datapoints are normalized, i.e., (cid:107) x
i (cid:107) 2 = 1 , and the
Training with early stopped gradient descent: We train the network with early stopped and randomly initialized gradient descent on a quadratic loss.
We choose the weights at initialization as
[Equation]
Here, ω and ν are parameters that trade off the magnitude of the weights of the ﬁrst and second layer.
Note that with this initialization, for a ﬁxed unit norm feature vector x , we have f W 0 , v 0 ( x ) =
O ( νω ) .
We apply gradient descent to the mean-squared loss
[Equation]
The gradient descent updates are v t
[Equation]
and
[Equation]
t ) , where η is a constant learning rate.
We study the risk of the network as a function of the iterations t .
Evaluation and performance metric: Our goal is to bound the test error as a function of the iterations of gradient descent.
Let (cid:96) : R × R → [0 , 1] be a loss function that is 1 -Lipschitz in its ﬁrst argument and obeys (cid:96) ( y, y )
= 0 ; a concrete example is the loss (cid:96) ( z, y ) =
| z − y | for arguments z, y ∈ [0 , 1] .
The test error or risk is deﬁned, as before, as R ( f ) =
E [ (cid:96) ( f ( x ) , y )] , where expectation is over examples ( x , y ) drawn from the unknown joint distribution from which the training set is drawn as well.
3. R ISK OF EARLY STOPPED NEURAL NETWORK TRAINING
Our main result is a bound on the test error of the two layer neural network trained for t iterations, in the regime where the network is very wide.
The result depends on the Gram matrix Σ ∈ R n × n determined by two kernels associated with the ﬁrst and second layer of the network.
The ( i, j ) -th entry of the Gram matrix as a function of the training examples is deﬁned as
[Equation]
with kernels
[Equation]
where ρ ij = (cid:104) x i , x j (cid:105) (recall that we assume (cid:107)
x i (cid:107) 2 = 1 , for all i ).
Our result depends on the singular values and vectors of this Gram matrix: Σ = (cid:80) ni
=1 σ 2 i u i u
Ti .
We are now ready to state our result.
Theorem 2.
Let α > 0 be the smallest eigenvalue of the Gram matrix Σ , suppose that the network is sufﬁciently wide, i.e., k ≥ Ω (cid:16) n 10 α 15 min( ν,ω ) (cid:17) , and suppose the initialization scale parameters obey νω ≤ α/ (cid:112) 32 log(2 n/δ ) and ν + ω ≤ 1 for some δ ∈ (0 , 1) .
Then, with probability at least 1 − δ , the risk of the network trained with gradient descent for t iterations is at most
[Equation]
Regarding the assumptions of the theorem, we remark that while the exponent of n and α in the width-condition ( k ≥ Ω (cid:16) n 10 α 15 min( ν,ω ) (cid:17) ) can be improved, the width condition ensures that the network is sufﬁciently wide so that the network operates in the kernel regime where the network behaves similar to an associated linear model.
Regarding the assumption that the smallest eigenvalue of the Gram matrix obeys α > 0 , Theorem 3.1 by Du et al.
(2019) shows that if no two x i , x j are parallel, then α > 0 , for a very related Gram matrix (speciﬁcally, the Gram matrix only consisting of the kernel K 1 deﬁned above).
As argued in that work, for most real-world datasets no two inputs are parallel, therefore this assumption is rather mild.
The risk bound established by Theorem 2 can be interpreted as a superposition of n -many U-shaped bias variance tradeoffs, similar to the expression (1) governing the risk of early stopped linear least squares.
Speciﬁcally, the i -th “bias” term (cid:104) u
i , y (cid:105) 2 (1 − ησ 2 i ) 2 t decreases in the number of gradient descent iterations t , while the i -th “variance” term (cid:104) u i , y (cid:105) 2 (1 − (1 − ησ 2 i ) t ) 2 σ 2 i increases in the number of gradient descent iterations.
The speed at which the two terms increase and decrease, respectively, is determined by the singular value σ 2 i .
Those singular values, in turn, are determined by the kernels
K 1 and K 2 , the random initialization (in particular the scale parameters ν, ω ), and the distribution of the examples.
Whether epoch-wise double descent occurs or not depends on those singular values and therefore on the kernels, the initialization, and the distribution of the examples, as illustrated with the following numerical example.
Numerical example to illustrate the theorem: We draw data from the linear model speciﬁed in
Section 2.1 with geometrically decaying diagonal co-variance entries and zero additive noise.
We then train the network for different initialization scale parameters ω, ν once with the same stepsize for both layers ( η = 8e-5), and once with a smaller stepsize for the second layer, i.e., η W = 8eand
η v = 1e-6.
In the top row of Figure 3, it can be seen that the empirical risk has a double-descent behavior if both layers are initialized at the same scale (
i.e., ω = ν = 1 ).
To understand the relation to the theorem better, we also plot in Figure 3 the extent to which the singular values are associated with the parameters in the ﬁrst and second layer.
To capture this, we ﬁrst comment on the relation of the parameters of the ﬁrst and second layer to the singular values σ 2 i and vectors u 2 i of the Gram matrix:
In the wide-network regime in which the theorem applies, the networks output is well approximated by its linearization around around the initialization.
With this, the networks predictions for the training examples are approximately
[Equation]
where J ∈ R n × dk + k is (approximately) the Jacobian of the network at initialization and J = (cid:80) ni
=
σ i u i v iT is its singular value decomposition.
Here, we denote by v i, W ∈ R dk and v i, v ∈ R k
the parts of the right-singular vectors of the Jacobian associated with the weights in the ﬁrst and second layer, respectively.
The norm of those vectors measures to what extent the singular value σ i is associated with the weights in the ﬁrst and second layer.
Returning to the numerical example, as the bottom row of Figure 3 shows, if we initialize both layers at the same scale ( ω = ν = 1 ), then most of the large singular values are associated, for the most part, with the weights in the second layer.
This leads to double descent, that can be mitigated by choosing a smaller stepsize associated with the weights in the second layer.
Improving performance by eliminating double descent: Similarly as for the linear least squares problem studied in the previous section, it is possible to shape the bias variance tradeoffs by adapting the stepsizes (or through initialization of the layers).
In Figure 3, we illustrate this behavior: Double descent is eliminated by choosing a smaller stepsize for the second layer, or by choosing a smaller initialization for the ﬁrst layer, as suggested by our theoretical results, and similar to the linear least squares setup as discussed in the previous section.
Also note that, not only does choosing a smaller stepsize for the second layer eliminate double descent, it also gives a better overall risk.
To understand the relation to the kernels, suppose we choose the initialization equally, i.e., ω =and ν = 1 .
If we update the variables of the second layer (i.e., v ) with a much larger stepsize than that of the ﬁrst layer (i.e., W ), then the kernel associated with the second layer dominates and the network behaves like a random feature model Rahimi & Recht (2008).
Similarly, if we update the variables of the ﬁrst layer with a much larger stepsize than that of the second layer, then the network behaves like a network with the ﬁnal layer weights v ﬁxed.
Thus, the stepsizes trade off the impact of the two kernels, and this tradeoff yields a double descent curve.
E ARLY STOPPING IN CONVOLUTIONAL NEURAL NETWORKS
We ﬁnally study the training of a standard 5-layer convolutional neural network (CNN) and a standard ResNet-18 model on the (10 class classiﬁcation) CIFAR-10 dataset.
Both networks were studied in Nakkiran et al.
(2020a).
As shown in that paper, the risk has a double descent behavior if the [Figure] [Figure] network is trained on a dataset with label noise, and we consider the same setup with 20% random label noise.
While we have no theoretical results for those two complicated neural network models, we demonstrate—inspired by our theory—that epoch-wise double descent can be eliminated and the early stopping performance can be improved by adjusting the stepsizes/learning rates.
5-layer CNN:
The 5-layer CNN consists for 4 convolutional layers followed by a fully connected layer.
Figure 4 shows that, just like for the two-layer network from the previous section, double descent can be eliminated by changing stepsizes, this time by decreasing the stepsize of the ﬁnal fully connected layer.
The intuition behind this is that large singular values of the Jacobian of the network at initialization are mostly associated with the last fully connected layer, measured in the same way as in the previous section.
This causes the convolutional layers to be learned slower than the fully connected layer which results in double descent.
Analogously as before, decreasing the stepsize pertaining to the fully connected layer eliminates double descent.
ResNet-18: We next consider the popular ResNet-18 model.
ResNet-18 has a double descent behavior when trained on the noisy CIFAR-10 problem Nakkiran et al.
(2020a).
Inspired by our theory, we again hypothesize that the double descent behavior occurs because some layer(s) of the
ResNet-18 model are ﬁtted at a faster rate than others.
If that hypothesis is true, then scaling the learning rates of some layers should eliminate double descent.
Indeed, Figure 6 in the appendix shows that when scaling the stepsizes of the later half of the layers of the network mitigates double descent.
A S UPPORTING M ATERIAL FOR : E ARLY - STOPPED GRADIENT DESCENT FOR LINEAR LEAST SQUARES A. I NTUITION FOR THE RISK EXPRESSION (1)
Below, we provide a proof of Theorem 1.
Here we provide intuition why the risk is governed by the risk expression (1).
First, note that the risk of the estimator can be written as a function of the variances of the features, σ 2 i , and of the coefﬁcients of the underlying true linear model, θ ∗ =
[ θ ∗ 1 , . . .
, θ ∗ d ] , as
[Equation]
where we used that z and x are drawn independently.
Next, recall that we consider the estimate based on early stopping gradient descent applied to the empirical risk
[Equation]
Here, the matrix X ∈ R n × d contains the scaled training feature vectors 1 √ n x 1 , . .
. , 1 √ n x n as rows, and y = 1 √ n [ y 1 , . .
, y n ] are the corresponding scaled responses.
The gradient descent iterates obey
[Equation]
where z =
[ z 1 , . .
, z n ] is the noise.
As we formalize below, in the under-parameterized regime where n (cid:29) d , we have that X T X ≈ Σ 2 .
Therefore the original iterates are close to the proximal iterates ˜ θ t deﬁned by
[Equation]
The proximal iterates are, up to the extra term diag ( η ) X T z , equal to the iterates of gradient descent applied to the population risk R ( θ ) .
Note that in contrast to the literature where it is common to bound the deviation of the original iterates from the iterates on the population risk Raskutti et al.
(2014), here we control the deviation of the original iterates to the proximal iterates ˜ θ t .
The iterates ˜ θ t can easily be written out in closed form.
To do so, ﬁrst note that for the recursion θ t +1 =
αθ t
+ γ we have ˜ θ t = α t θ 0 + γ (cid:80) t − 1 i =
α i = α t θ 0 + γ 1 − α t 1 − α , where we used the formula for a geometric series.
Using this relation, and that we are starting our iterations at θ 0 i = 0 , we obtain for the i -th entry of ˜ θ t that
[Equation]
where ˜ x i is the i -th column of X (not the i -th example/feature vector!).
Next note that, E (cid:2) (˜ x
Ti z ) 2 (cid:3)
≈ σ 2 σ 2 i because the entries of z are N (0 , σ 2 ) distributed, and the entries of ˜ x
i are 1 / √ n N (0 , σ 2 i ) distributed.
Using this expectation in the iterates ˜ θ t , and evaluating the risk of those iterates via the formula for the risk given by (7) yields the risk expression (1).
The proof of Theorem 1 in the appendix makes this intuition precise by formally bounding the difference of the proximal iterates to the original iterates.
A. M OTIVATION FOR CALLING THE U- SHAPED CURVES BIAS - VARIANCE TRADEOFFS
Let ˆ θ
=
ˆ θ ( D ) be the parameter obtained based on the training data (for example by early stopping).
The textbook bias-variance decomposition of the risk of ˆ θ is
[Equation]
[Table]
[Equation]
[Figure]
[Figure]
t iterations r i s k (a) n = 5 d
R ( θ t ) ¯ R (˜ θ t ) 10 {+
 [Figure] [Figure] [Figure]+}
101 .1 .1 .1 .t
iterations (b) n = 10 d
R ( θ t ) ¯ R (˜ θ t )
The ﬁrst term above is the bias of the hypothesis ˆ h ( x ) =
(cid:68) x , ˆ θ (cid:69) .
It measures how well the average function can estimate the true underlying function h ( x ) =
(cid:104)
x , θ ∗ (cid:105) .
A low bias means that the hypothesis accurately estimates the true underlying function (cid:104) x , θ ∗ (cid:105) .
The second term is the variance of the method.
The variance of the method measures the variance of the hypothesis over the training sets.
Recall from the previous paragraph that the estimate ˜ θ t approximates the original iterations θ t well provided that the model is sufﬁciently underparameterized, i.e., d/n is small.
It is straightforward to verify that Bias(˜ θ t )
=
(cid:80) di
=1 ( θ ∗ i ) 2 (1 − η i σ 2 i ) 2 t and Variance(˜ θ t )
= (cid:80) di
=1 σ 2 n (1 − (1 − η i σ 2 i ) t ) 2 , exactly equal to the bias and variance terms in the risk expression (1).
It follows that the bias and variance of the original gradient descent iterates θ t are also approximately equal to the terms in the risk expression (1).
The U-shaped curves are then the bias-variance terms pertaining to the i -th feature; this formally establishes the U-shaped curves as bias-variance tradeoffs.
A. N UMERICAL RESULTS FOR LINEAR LEAST SQUARES
In this section we provide further numerical results for linear least squares.
We consider a linear model with d = 700 features, and with n = 6 d examples.
We let a fraction 6 / 7 of the features have singular value σ i = 1 and associated model coefﬁcient θ i = 1 , and the rest, 1 / 7 of the features, have singular value σ i = 0 .
1 and θ i = 10 .
In Figure 5(a) we show the risk obtained by simulating the risk empirically along with the risk expression ¯ R (˜ θ t ) given by equation (1).
It can be seen that the risk expression ¯ R ( ˜ θ t ) slightly under-estimates the true risk.
The quality of the estimate becomes better as we increase n ; in Figure 5(b) we show simulations for the same conﬁguration but with n = 10 d .
B S UPPORTING M ATERIAL FOR : E ARLY STOPPING IN CONVOLUTIONAL NEURAL NETWORKS B. R ES N ET -18 TRAINING DOUBLE DESCENT ELIMININATION
In Figure 6 we provide test and train error curves for ResNet-18 trained with different stepsizes on noisy CIFAR-10.
The results show that, as mentioned in the main body, that double descent is eliminated by choosing the stepsizes appropriately.
Note that, while in the case of the 5-layer convolutional network the weights of the linear layer is learned faster than the convolutional layers and decreasing the stepsize of the linear layer eliminates the resulting double descent, whereas
ResNet-18 has a different architecture than the simple convolutional network as it features residual blocks with residual connections.
For ResNet-18, we found that the latter half of the network layers, which include the latter two residual blocks together with the ﬁnal, fully-connected layer, is learned faster than the initial convolutional layer and the ﬁrst two residual blocks.
Similarly to the 5-layer [Figure] e rr o r i: same stepsize ii: smaller stepsize last half layers
[Figure]
[Figure] [Figure] [Figure] [Figure]
0 .epoch
[Figure]
[Figure]
convolutional network decreasing the stepsizes of these layers eliminates the double descent and improves the early stopping performance.
As discussed before, classical machine learning theory for the underparameterized regime establishes the bias-variance tradeoff as a result of the bias decreasing and the variance increasing as a function of the model size (complexity).
In the over-parameterized regime, the bias often continues to decrease, while the variance also decreases.
This has been established in a number of recent works Jacot et al.
; Yang et al.
(2020); d’Ascoli et al.
(2020), and provides a bias-variance decomposition of the model-wise double-descent shaped risk curve.
In this paper, we demonstrated that epoch-wise double descent occurs for a different reason than the model-wise double descent.
Namely, epoch-wise double descent can be explained as a temporal superposition of multiple bias-variance tradeoff curves rather than with a unimodal variance curve.
That also means that the overall bias (i.e., the sum of the individual bias terms) might not be decreasing and the overall variance might not be uni-modal like in the model-wise case.
To demonstrate that it is in fact not, in Figure 7, we plot the numerically computed bias and variance terms (computed as proposed in Yang et al.
for the CNN experiment from Section 4 along with the risk, which shows that in fact the overall bias is increasing, while the variance has a double-descent like shape. [Figure]
C M ULTI - DESCENT
We note that in principle we can also observe multiple descents as a function of training time.
Specifically, recall the risk expression for the linear case, equation 1.
It consists of d many bias-variance tradeoffs, so in principle those curves might give rise not only to epoch-wise double descent, but can give rise to multiple descent.
See Figure 8, left panel, in which we show an example of three bias-variance tradeoffs that add up to a multi-descent curve.
The same can, of course happen for neural networks.
In Figure 8, right panel, we demonstrate this for the two-layer network introduced in Section 3.
We note that it is much harder to observe multi-descent in neural networks as it requires a very particular data structure.
Speciﬁcally, for the two-layer neural network we consider, we found that the existence of the multiple descents depends heavily on the noise in the data generation process.
Speciﬁcally, for the two-layer neural network, multi-descent only occurs for a particular range of noise levels where σ noise (cid:29) σ max i .
Intuitively, multiple descents could be observed in other empirical scenarios and for other architectures based on our theoretical and experimental ﬁndings for the linear case and the two-layer neural network.
However, we did not observe multi-descent in a practical setup (such as for training
CIFAR-10 with a convolutional network), as it requires a very particular setup (i.e., combination of underlying data distribution, network architecture, and training).
Image classiﬁcation datasets are considered to be minimally noisy and highly structured and this particular setup does not seem to occur in practice even with the artiﬁcially injected label noise, at least we didn’t observe it when training standard networks on CIFAR-10, and it hasn’t been reported elsewhere.
D P ROOF OF T HEOREM
The difference of the risk and risk expression can be bounded by
[Equation]
We bound the two terms on the righ-hand-side separately.
We start with bounding the ﬁrst term by applying the lemma below.
Lemma 1.
Deﬁne ˜
X
so that X = ˜ XΣ .
Suppose that (cid:13)(cid:13) (cid:13) I − ˜ X T ˜ X
(cid:13)(cid:13) (cid:13) ≤ (cid:15) , with (cid:15) ≤ min i η i σ
i 2max i η i σ 2 i .
Then (cid:12)(cid:12)(cid:12) R ( θ t ) − R ( ˜ θ t ) (cid:12)(cid:12)(cid:12) ≤ (1 − (
[Equation]
In order to apply the lemma, we start by verifying its condition.
Towards this goal, consider the matrix X =
˜
XΣ and note that the entries of ˜ X are iid N (0 , 1 /n ) .
A standard concentration inequality from the compressive sensing literature (speciﬁcally (Foucart & Rauhut, Holger, 2013, Chapter 9)) states that, for any β ∈ (0 , 1) ,
[Equation]
With β = (cid:113) 75 dn we obtain that, with probability at least 1 − e − d ,
[Equation]
Next, we bound the term on the RHS of in (10), with the following lemma.
Lemma 2.
Provided that η i σ 2 i ≤ 1 for all i , with probability at least 1 − 2 d ( e − β 2 / 2 + e − n/ 8 ) ,
[Equation]
Applying the lemma with β 2 = 10 log( d ) , we obtain that with probability at least 1 − 2 d − 5 − 2 de − n/ 8 − e − d we have
[Equation]
We are now ready to bound the second term in (9):
Lemma 3.
With probability at least 1 − 4 e − β 2 8 , we have that
[Equation]
with ¯ R (˜ θ t ) as deﬁned in (1) .
Applying the two bounds (11) and (12) to the RHS of the bound (9) concludes the proof.
The remainder of the proof is devoted to proving the three lemmas above.
D. P ROOF OF L EMMA
Recall that the iterates of the original and closely related problem are given by
[Equation]
Note that X =
˜
XΣ , where we deﬁned ˜ X which has iid Gaussian entries N (0 , 1 /n ) .
With this notation [Equation] and using that Σ is diagonal and therefore commutes with diagonal matrices
[Equation]
we obtain the following expressions for the residuals of the two iterates:
[Equation]
The difference between the residuals is
[Equation]
where the last equality follows by adding and subtracting diag ( η ) Σ
˜ X T ˜ X
( Σ ˜ θ t − Σ θ ∗ ) and rearranging the terms.
It follows that
[Equation]
Here, we used the bound
[Equation]
Here, we used that η i σ 2 i ≤ 1 , by assumption, and the last inequality follows by the assumption (cid:15) ≤ min i η i σ 2 i 2max η i σ 2 i .
Iterating the bound (13) yields
[Equation]
which concludes the proof.
D. P ROOF OF L EMMA
Recall that
[Equation]
With η i σ 2 i ≤ 1 , by assumption, it follows that
[Equation]
[Equation]
Conditioned on z , the random variable ˜ x
Ti z is zero-mean Gaussian with variance (cid:107)
z (cid:107) 2 /n .
Thus,
P
[Equation]
| [Equation] z |
[Equation]
β
[Equation]
≤ 2 e − β 2 / 2 . Moreover, as used previously in (16), with probability at least 1 − 2 e − n/ 8 , (cid:107) z (cid:107) 22 ≤ 2 σ 2 .
Combining the two with the union bound, we obtain
[Equation]
Using this bound in inequality (14), we have that, with probability at least 1 − [Equation] e
[Equation]
β
[Equation]
[Equation]
e − n/ 8 )
that
[Equation]
By the union bound over all i we therefore get that
[Equation]
with probability at least 1 − 2 d ( e − β 2 / 2 + e − n/ 8 ) .
D. P ROOF OF L EMMA
We have
R
[Equation]
i . The random variable Z i , conditioned on z , is a squared Gaussian with variance upper bounded by (cid:107) z (cid:107) 2 √ n and has expectation
[Equation]
By a standard concentration inequality of sub-exponential random variables (see e
[Equation]
g. (Wainwright, 2019, Chapter 2, Equation 2.21)), we get, for β ∈ (0 , √ d ) and conditioned on z , that the event
[Equation]
occurs with probability at least 1 − 2 e − β 2 8 .
With the same standard concentration inequality for sub-exponential random variables, we have that the event
[Equation]
also occurs with probability at least 1 − 2 e − β 2 8 .
By the union bound, both events hold simultaneously with probability at least 1 − 4 e − β 2 8 .
On both events, we have that
[Equation]
concluding the proof of our lemma.
E P ROOF OF P ROPOSITION
By equation (1), the risk expression is a sum of U-shaped curves: R (˜ θ t )
=
σ 2 + (cid:80) di
=1 U i ( t ) .
We start by considering one such U-shaped curve
[Equation]
and ﬁnd its minimum as a function of the number of iterations, t .
Towards this end, we set the derivative of one such U-shaped curve, given by ∂
[Equation]
to zero
[Equation]
which gives that the minimum occurs when
[Equation]
For the iteration t which satisﬁes this equation [Equation] we get
[Equation]
thus this minimum is independent of the iteration t and independent of the stepsize, provided their relation is as described in (17) above.
F P ROOF AND STATEMENTS FOR NEURAL NETWORKS
In this section, we prove the following result, which is a slightly more formal version of our main result for neural networks, Theorem 2.
Theorem 3.
Draw a dataset D = { ( x 1 , y 1 ) , . . .
, ( x n , y n ) } consisting of n examples i.i.d.
from a distribution with (cid:107) x i (cid:107) 2 = 1 and | y i | ≤ 1 .
Let Σ ∈ R n × n be the corresponding Gram matrix deﬁned in (4) , and suppose its smallest singular value obeys α > 0 .
Pick an error parameter ξ ∈ (0 , 1) and a failure probability δ ∈ (0 , 1) , and consider the two-layer neural network f W , v ( x ) = 1 √ k relu( x T W ) v , with parameters W d × k , v ∈ R k initialized according to (3) with initialization scale parameters ν, ω obeying νω ≤ ξ/ (cid:112) 32 log(2 n/δ ) and ν + ω ≤ 1 .
Suppose that the network is sufﬁciently overparameterized, i.e.,
[Equation]
Then, the risk of the network trained with gradient descent with constant stepsize η for t iterations obeys, with probability at least 1 − δ ,
[Equation]
Theorem 2 directly follows by choosing the error parameter as ξ
= O ( α ) .
F. P ROOF OF T HEOREM
In this section, we provide a proof of Theorem 2.
Our proof relies on the observation that highly overparameterized neural networks behave as associated linear models, as established in a large number of prior works (Arora et al., 2019; Du et al., 2018; Oymak & Soltanolkotabi, 2020; Oymak et al., 2019; Heckel & Soltanolkotabi, 2020b).
The proof consists of two parts.
First, we control the empirical risk as a function of the number of gradient descent steps, t .
Second, we control the generalization error, i.e., the gap between the population risk and the empirical risk by bounding the Rademacher complexity of the function class consisting of two-layer networks trained with t iterations of gradient descent.
Recall that our result depends on the singular values and vectors of the gram matrix of kernels associated with the twolayer network.
The Gram matrix is given as the expectation of the outer product of the Jacobian of the network at initialization:
[Equation]
Here, expectation is with respect to the random initialization W 0 , [Equation]
Bound on the training error: We start with a results that controls the training error and ensures that the coefﬁcients of the neural network move little from its initialization.
Theorem 4.
Pick an error parameter ξ ∈ (0 , 1) and any failure probability δ ∈ (0 , 1) , and choose ν, ω so that they satisfy νω ≤ ξ/ (cid:112) 32 log(2 n/δ ) .
Suppose that the network is sufﬁciently overparameterized, i.e.,
[Equation]
i)
Then, with probability at least 1 − δ , the mean squared loss after t iterations of gradient descent obeys (cid:118)(cid:117)(cid:117)(cid:116) n (cid:88) i =
( y i − f W t ,
[Equation]
ii) Moreover, the coefﬁcients overall deviate little from its initialization,
i
e.,
[Equation]
Here,
(cid:107)·(cid:107) F denotes
the Frobenius norm.
In addition each of the coefﬁcients changes only little,
i
e., for all iterations
t
[Equation]
Here, w t,r is the r -th row of W t , and v t,r is the r -th entry of v t .
Bound on the empirical risk: Because we train with respect to the (cid:96) 2 -loss but deﬁne the risk with respect to the (generic Lipschitz) loss (cid:96) , the empirical risk and training loss are not the same.
Nevertheless, we can upper bound the empirical risk computed over the training set at iteration t with the training loss at iteration t :
[Equation]
where (i) follows from (cid:96) ( z, y )
= (cid:96) ( z, y ) − (cid:96) ( y, y ) ≤
| z − y | because the loss is 1 -Lipschitz.
Equation (ii) is the most interesting one, and follows from Theorem 4, equation (21), and holds with probability at least 1 − δ .
This bound is proven by showing that, provided the network is sufﬁciently wide, the training loss behaves as gradient descent applied to a linear least-squares problem with dynamics governed by the gram matrix Σ .
Bound on the generalization error: Next, we bound the generalization error R ( f ) − ˆ R ( f ) by bounding the Rademacher complexity of the functions that gradient descent can reach with t gradient descent iterations.
Let F be a class of functions f : R d → R .
Let (cid:15) 1 , . . .
, (cid:15) n be iid Rademacher random variables, i.e., random variables that are chosen uniformly from {− 1 , 1 } .
Given the dataset D , deﬁne the empirical Rademacher complexity of the function class
F as
[Equation]
Here, D = { ( x 1 , y 1 ) , . . .
, ( x n , y n ) } is the training set, consisting of n points drawn iid from the example generating distribution.
By a standard result from statistical learning theory, a bound on the Radermacher complexity directly gives a bound on the generalization error for each predictor in a class of predictors.
Theorem 5 ( (Mohri et al., 2012, Thm. 3.1) ) .
Suppose (cid:96) ( · , · ) is bounded in [0 , 1] and 1 -Lipschitz in its ﬁrst argument.
With probability at least 1 − δ over the random dataset D consisting of n iid examples, we have that
[Equation]
We consider the class of neural networks with weights close to the random initialization W 0 , v 0 , deﬁned as:
[Equation]
with
[Equation]
The Rademacher complexity of this class of functions is controlled with the following result.
Lemma 4.
Let W 0 be drawn from a Gaussian distribution with N (0 , ω 2 ) entries, and suppose the entries of v 0 are draw uniformly from {− ν, ν } .
Assume the ( x i , y i ) are drawn iid from some distribution with (cid:107)
x
i (cid:107) 2 = 1 and | y i | ≤ 1 .
With probability at least 1 − δ over the random training set, provided that (cid:112) log(2 n/δ ) / 2 k ≤ 1 / 2 , the empirical Rademacher complexity of F Q,M is, simultaneously for all Q , bounded by
[Equation]
We set M =
O ( ξ α k − 1 / 4 ) .
With this choice, the term on the right hand side above is bounded by
[Equation]
where we used νω ≤ 1 and √ log(2 /δ )
/ 2 k 1 /≤ 1 , by assumption (18).
Note that by (23) and by (24) combined with the assumption (18) we have that (cid:107) w r − w 0 ,
[Equation]
2 ≤ ωM and | v r − v 0 ,r | ≤ νM , as desired.
Let Q i
=
i for i = 1 , 2 , . . . .
Simultaneously for all i , by the lemma above, for this choice of M , the function class F Q i ,M has Rademacher complexity bounded by
[Equation]
We next choose the radius Q as deﬁned in (22).
Let i ∗ be the smallest integer such that Q ≤ Q i ∗ , so that Q i ∗ ≤ Q + 1 .
We have that i ∗ ≤
O ( (cid:112) n/α ) and
[Equation] [Equation]
by the assumption of the theorem on k being sufﬁciently large, and by ν + ω ≤ 1 .
Next, from a union bound over the ﬁnite set of integers i = 1 , . . .
, i ∗ , we obtain
[Equation]
as desired.
Final bound on the risk: Combining the bound on the training with the generalization bound yields the upper bound (19) on the risk of the network trained for t iterations of gradient descent.
The remainder of the proof is devoted to proving Theorem 4 and Lemma 4.
F. P RELIMINARIES
We start with introducing some useful notation.
First note that the prediction of the neural network for the n training data points as a function of the parameters are
[Equation]
where X n × d is the feature matrix and
W ∈ R d × k and v ∈ R k are the trainable weights of the network.
The transposed Jacobian of the function f is given by
[Equation]
where we deﬁned the Jacobians corresponding to the weights of the ﬁrst layer, W , and the second layer, v , respectively as
[Equation]
Here, relu (cid:48)
( x ) = 1 { x ≥ 0 } is the derivative of the relu activation function, which is the step function.
Our results depend on the singular values and vectors of the expected Jacobian at initialization:
[Equation]
where (cid:12) is the Hadamard product, and where we used that the entries of v 0 are choosen iid uniformly from {− ν, ν } .
Expectation is over the weights W 0 at initialization, which are iid N (0 , ω 2 ) .
This yields
[Equation]
where K 1 and K 2 are two kernels associated with the ﬁrst and second layers of the network and are given by
[Equation]
with ρ ij =
[Equation]
i ,
[Equation]
2 and by
[Equation]
For both of those expressions [Equation] we used the calculations from (Daniely et al., 2016, Sec. 4.2) for the ﬁnal expressions of the kernels.
Also note that, by assumption (cid:107) x i (cid:107) 2 = 1 .
F
P ROOF OF T HEOREM 4 ( BOUND ON THE TRAINING ERROR )
In this subsection, we prove Theorem 4.
F.3.
T
HE DYNAMICS OF LINEAR AND NONLINEAR LEAST - SQUARES
Theorem 4 relies on approximating the trajectory of gradient descent applied to the training loss with an associated linear model that approximates the non-linear neural network in the highly-overparameterized regime.
This strategy has been used in a number of recent publications (Arora et al., 2019; Du et al., 2018; Oymak & Soltanolkotabi, 2020; Oymak et al., 2019;
Heckel & Soltanolkotabi, 2020b); in order to avoid repetition, we rely on a statement (Heckel &
Soltanolkotabi, 2020a, Theorem 4), which bounds the error between the true trajectory of gradient descent and the trajectory of an associated linear problem.
Let f : R N → R n be a non-linear function with parameters θ ∈ R N , and consider the non-linear least squares problem
[Equation]
The gradient descent iterations starting from an initial point θ 0 are given by
[Equation]
∇L ( θ
[Equation]
J
[Equation]
θ
[Equation]
y
[Equation]
where J ( θ ) ∈ R n × N is the Jacobian of f at θ (i.e., [ J ( θ )] i,j = ∂
f i ( θ ) ∂ θ j ).
The associated linearized least-squares problem is deﬁned as
[Equation]
Here, J ∈ R n × N , refered to as the reference Jacobian, is a ﬁxed matrix independent of the parameter θ that approximates the Jacobian mapping at initialization, J ( θ 0 ) .
Starting from the same initial point θ 0 , the gradient descent updates of the linearized problem are
[Equation]
To show that the non-linear updates (33) are close to the linearized iterates (35), we make the following assumptions:
i)
We assume that the singular values of the reference Jacobian obey for some α, β
[Equation]
Furthermore, we assume that the norm of the Jacobian associated with the nonlinear model f is bounded in a radius R around the random initialization
[Equation]
for all θ ∈
[Equation]
Here, B R ( θ 0 ) :
= { θ : (cid:107) θ − θ 0 (cid:107) ≤ R } is the ball with radius R around θ 0 .
ii)
We assume the reference Jacobian and the Jacobian of the nonlinearity at initialization J ( θ 0 ) are (cid:15)
-close:
[Equation]
iii)
We assume that within a radius R around the initialization, the Jacobian varies by no more than (cid:15) :
[Equation]
Under these assumptions the difference between the non-linear residual
[Equation]
and the linear residual
[Equation]
are close throughout the entire run of gradient descent.
Theorem 6 ((Heckel & Soltanolkotabi, 2020a, Theorem 4), Closeness of linear and nonlinear least- -squares problems) .
Assume the Jacobian J ( θ ) ∈ R n × N associated with the function f ( θ ) obeys
Assumptions (36a) , (36b) , (36c) , and (36d) around an initial point θ 0 ∈ R N with respect to a reference Jacobian J ∈ R n × N and with parameters α, β, (cid:15) 0 , (cid:15) , obeying 2 β ( (cid:15)
+ (cid:15) ) ≤ α 2 , and R .
Furthermore, assume the radius R is given by
[Equation]
Here, J † is the pseudo-inverse of J .
We run gradient descent with stepsize η ≤ 1 β 2 on the linear and non-linear least squares problem, starting from the same initialization θ 0 .
Then, for all iterations t ,
i) the non-linear residual converges geometrically
[Equation]
ii) the residuals of the original and the linearized problems are close
[Equation]
iii) the parameters of the original and the linearized problems are close
[Equation]
iv) and the parameters are not far from the initialization
[Equation]
Theorem 6 above formalizes that in a (small) radius around the initialization, the non-linear problem behaves very similar to its associated linear problem.
As a consequence, to characterize the dynamics of the nonlinear problem, it sufﬁces to characterize the dynamics of the linearized problem.
This is the subject of our next theorem, which is a standard result on the gradient iterations of a least squares problem, see for example (Heckel & Soltanolkotabi, 2020b, Thm. 5) for the proof.
Theorem 7 (E
[Equation]
g. Theorem 5 in Heckel & Soltanolkotabi (2020b)) .
Consider a linear least squares problem (34) and let J = (cid:80) ni
=
σ i u i v Ti be the singular value decomposition of the matrix J .
Then the linear residual ˜ r t after t iterations of gradient descent with updates (35) is
[Equation]
Moreover, using a step size satisfying η ≤ 1 σ 21 , the linearized iterates (35) obey
[Equation]
F.3. P ROVING T HEOREM 4 BY APPLYING T HEOREM
We are now ready to prove Theorem 4.
We apply Theorem 6 to the predictions of the network given by f ( W , v ) deﬁned in (30) with parameter θ =
( W , v ) .
As reference Jacobian we choose a matrix J ∈ R n × dk
+ k that satisﬁes JJ T =
E (cid:2) J ( W 0 , v 0 )
J T ( W 0 , v 0 ) (cid:3) (where expectation is over the random initialization ( W 0 , v 0 ) ), and at the same time is very close to the Jacobian of f at initialization, i.e., to J ( W 0 , v 0 ) .
Towards this goal, we apply Theorem 6 with the following choices of parameters:
[Equation]
Note that assumption (20) guarantees that (cid:15) 0 ≤ (cid:15) , a fact we used later.
We now verify that the conditions of Theorem 4 are satisﬁed for this choice of parameters with probability at least 1 − δ .
Speciﬁcally we show that each of the conditions holds with probability at least 1 − δ .
By a union bound, the success probability is then at least 1 − Ω( δ ) , and by rescaling δ by a constant, the conditions are satisﬁed with probability at least 1 − δ .
Bound on residual: We need a bound on the network outputs at initialization as well as on the initial residual to verify the conditions of the theorem.
We start with the former:
[Equation]
45) where the inequality holds with probability at least 1 − δ , by Gaussian concentration (see Lemmain Heckel & Soltanolkotabi (2020b) and recall that W 0 has iid N (0 , ω 2 ) entries).
Moreover, the last equality follows from (cid:107) x
i (cid:107) 2 = 1 .
It follows that, with probability at least 1 − δ , the initial residual is bounded by
[Equation] [Equation]
where the ﬁrst inequality holds by the triangle inequality and using the assumption | y i | ≤ 1 , and the second inequality by νω (cid:112) 8 log(2 n/δ ) ≤ 1 , again by assumption.
Radius in the theorem:
In order to verify the condition of the theorem, we need to control the radius in the theorem, which we do next.
With our assumptions and the choices of parameters above, the radius in the theorem, deﬁned in equation (37), obeys
[Equation]
Here, (i) follows from the fact that (cid:13)(cid:13)
[Equation]
1 √ 2 α (cid:107) r 0 (cid:107) 2 , (ii) from (cid:15) 0 + (cid:15) ≤ 2 (cid:15) =
ξ α 3 β 2 (by deﬁnition of (cid:15) 0 and (cid:15)
, and from the bound
For we used
) in the theorem.
Verifying Assumptions (36a) and (36b) :
By deﬁnition JJ T = Σ , thus the lower bound in assumption (36a) holds by the deﬁnition of α as σ n ( Σ ) ≥ √ 2 α .
Regarding the upper bound of (36a), note that
[Equation]
where K 1 ∈ R n × n and K 2 ∈ R n × n are the kernel matrix with entries K 1 ( x i , x j ) and K 2 ( x i , x j ) .
It follows that (cid:107) J (cid:107) ≤ 2( ω + ν ) √ n ≤ β , as desired.
This concludes the veriﬁcation of (36a).
To verify assumption (36b) note that
[Equation]
For the last inequality, we used that (cid:107) v 0 (cid:107) 2 = ν √ k , and that (cid:107) v − v 0 (cid:107) 2 + (cid:107) W − W 0 (cid:107) F ≤ R ≤ √ k min( ω, ν ) , by the bound on the radius in (47), and ﬁnally that (cid:107) XW 0 (cid:107) F ≤ ω 6 √ k with probability at least 1 − δ provided that k ≥ log( n/δ ) , which holds by assumption.
For this inequality we used that x Ti W 0 is a Gaussian vector with iid N (0 , ω 2 ) entries.
It follows that assumption (36b) holds with probability at least 1 − δ , as desired.
Verifying Assumption (36c) : We start with stating a concentration lemma from Heckel &
Soltanolkotabi (2020b).
Lemma 5 (Concentration lemma (Heckel & Soltanolkotabi, 2020b, Lemma 3)) .
Consider the partial Jacobian J 1 ( W ) , and let W ∈ R n × k be generated at random with i.i.d.
N (0 , ω 2 ) entries, and suppose the v (cid:96) are drawn from a distribution with | v (cid:96) | ≤ ν .
Then, with probability at least 1 − δ ,
[Equation]
Lemma 6.
Let J 2 ( W ) = 1 √ k relu( XW ) , with W generated at random with i.i.d.
N (0 , ω 2 ) entries.
With probability at least 1 − δ ,
[Equation]
Combining the statements of the two lemmas, it follows that, with probability at least 1 − 2 δ ,
[Equation]
To show that (49) implies the condition in (36c), we use the following lemma.
Lemma 7 ((Oymak et al., 2019, Lem. 6.4)) .
Let J 0 ∈ R n × N , N ≥ n and let Σ be n × n psd matrix obeying (cid:13) (cid:13)
J 0 J T 0 − Σ (cid:13) (cid:13) ≤ ˜ (cid:15) 2 , for a scalar ˜ (cid:15) ≥ 0 .
Then there exists a matrix J ∈ R n × N obeying Σ
=
JJ T such that
[Equation]
From Lemma 7 combined with equation (49), there exists a matrix J ∈ R n × N that obeys
[Equation]
This part of the proof also speciﬁes our choice of the matrix J as a matrix that is (cid:15)
close to the
Jacobian at initialization, J ( W 0 ) , and that exists by Lemma 7 above.
Verifying Assumption (36d) :
We control the perturbation around the random initialization.
Lemma 8.
Let W 0 have iid N (0 , ω 2 ) entries and let v 0 have (arbitrary) entries in {− ν, + ν } entries.
Then for all W and v obeying, for some ˜ R ≤ 12 √ k ,
[Equation]
the Jacobian in (31) obeys
[Equation]
with probability at least 1 − ne − 12 ˜ R
[Equation]
Recall the deﬁnition ˜ R = √ k (cid:16) (cid:15) (cid:107) X (cid:107) (
[Equation]
+3 ω ) (cid:17) 3 from (47).
From ˜ R ≤ ( kR ) 1 / 3 for ˜ R ≤ √ k , the bound provided by lemma 8 guarantees that
[Equation]
where the second inequality follows by choosing ˜ R = √ k (cid:16)
[Equation]
(cid:17) 3 .
This holds with probability at least
[Equation]
where in (i) we used (20).
Therefore, Assumption (36d) holds with high probability by our choice of (cid:15)
=
ξ α 3 β 2 .
Concluding the proof of Theorem 4: By the previous paragraphs, the assumptions of Theoremare satisﬁed with probability at least 1 − O ( δ ) .
Therefore we can bound bound the training error and the deviation of the coefﬁcients from the initialization as follows.
Training error:
We bound the training error in (20).
The training error at iteration t is bounded by
[Equation]
where inequality (i) follows from bounding the linear residual (cid:107) ˜ r t (cid:107) 2 with theorem 7, as well as bounding the distance between the linear residual and the non-linear one with (39).
Inequality (ii) follows from r
= f
( W 0 , v 0 )
− y , and ﬁnally (iii) follows from (cid:107) f
( W 0 , v 0 ) (cid:107) 2 ≤ νω (cid:112) 8 log(2 n/δ ) (cid:107) y
(cid:107) 2 ≤ ξ 2 (cid:107) y (cid:107) 2 , by (45), and βα 2 ( (cid:15)
+ (cid:15) ) ≤ βα
[Equation]
Distance from initialization: We next bound the distance from the initialization, i
[Equation]
e., we establish (22).
Combining equation (43) in theorem 7 with equation (39) in theorem 6, we obtain
[Equation]
where (i) follows from r
= f
( W 0 , v 0 )
− y and (ii) follows from 2 .
5 β 2 α 4 ( (cid:15)
+ (cid:15) )
(cid:107) r 0 (cid:107) 2 ≤ 516 α (cid:107) r 0 (cid:107) 2 ≤ 516 α
( (cid:107) f
( W 0 , v 0 ) (cid:107)
+ (cid:107) y (cid:107) 2 ) , where we used (cid:15)
+ (cid:15) ≤ 2 (cid:15) =
ξ α 3 β 2 , by deﬁnition of (cid:15), (cid:15) 0 , combined with (cid:107) f
( W 0 , v 0 ) (cid:107) 2 ≤ ξ/ 4 √ n .
Bound on change of coefﬁcients: Finally, we establish the bounds on the change of the individual coefﬁcients (23) and (24).
We start with the weights in the ﬁrst layer, w r .
The gradient with respect to w r is given by ∇ w r L
( W , v ) =
[ J T 1 ( W , v )] r r , where [ J T 1 ( W , v )] r is the submatrix of the
Jacobian multiplying with the weight w r , and r is the residual.
Therefore, we obtain
[Equation]
Here, (i) follows from
[Equation]
where the last inequality follows from | v τ,r − v 0 ,r | ≤ (cid:107) v τ,r − v 0 ,r (cid:107) 2 ≤ R ≤ 4 α √ n , by (47).
Moreover, for (ii) we used that, by (38), the non-linear residuals converge geometrically, and (iii) follows from the formula for a geometric series.
This conclude the proof of the bound (23).
Analogously, we obtain
[Equation]
where we used that
[Equation]
Here, the last inequality follows by using that the entries of Xw 0 ,r are not independent but N (0 , ω 2 ) distributed, and by taking an union bound over all entries of that vector and over all Xw 0 ,r .
This concludes the proof of the bound (24).
F.3. P ROOF OF LEMMA
First note that
[Equation]
In the reminder of the proof we bound the three terms above
[Equation]
We start by bounding the third term in (50) as:
[Equation]
We proceed with bounding the ﬁrst term in (50) as:
[Equation]
Next we establish below that with probability at least 1 − ne − kq 2 / 2 , the second term in in (50) is bounded as
[Equation]
51) provided that
[Equation]
where the last inequality follows from setting q = (2 k ˜ R ) 2 / 3 (note that the assumption ˜ R ≤ 12 √ k ensures q ≤ k ).
Putting those three bounds together in (50) yields
[Equation]
which established the claim.
It remains to prove (51).
Towards this goal, ﬁrst note that
[Equation]
We proceed with bounding the second term in the RHS of (53) as:
[Equation] [Equation]
Because relu (cid:48) is the step function, we have to bound the number of sign ﬂips between the matrices
XW and XW (cid:48) .
For this we use the lemma below:
Lemma 9.
Let | v | π ( q ) be the q -th smallest entry of v in absolute value.
Suppose that, for all i , and q ≤ k ,
[Equation]
Then
[Equation]
For the entries W (cid:48) being iid N (0 , ω 2 ) , we note that, with probability at least 1 − ne − kq 2 / 2 , the q -th smallest entry of x Ti W (cid:48)
∈ R k obeys
[Equation]
for all i =
[Equation] [Equation]
n.
[Equation] [Equation]
We are now ready to conclude the proof of the lemma.
By equation (54),
[Equation]
provided that
[Equation]
with probability at least 1 − ne − kq 2 / 2 .
F. P ROOF OF L EMMA 4 ( BOUND ON THE R ADEMACHER COMPLEXITY )
Our proof follows that of a related result, speciﬁcally (Arora et al., 2019, Lem. 6.4) which pertains to a two-layer ReLU network where only the ﬁrst layer is trained and the second layers’ coefﬁcients are ﬁxed.
Our goal is to bound the empirical Rademacher complexity
[Equation]
where expectation is over the iid Rademacher random variables (cid:15) i , and where { x 1 , . .
, x n } are the training examples.
The derivation of the bound on the Rademacher complexity is based on the intuition that if the parameter M of the constraint (cid:107) w r − w 0 ,r (cid:107) 2 ≤ ωM is sufﬁciently small, then relu (cid:48) ( w Tr x i ) is constant for most r , because | w T 0 ,r x i | is bounded away from ωM with high probability, by anticoncentration of a Gaussian.
For those coefﬁcient vectors r for which relu (cid:48) ( w Tr x i ) is constant, we have relu( w Tr x i ) =
relu (cid:48) ( w T 0 ,r x i ) w
Tr x i .
For the other coefﬁcients, we can bound the difference of those two values as
[Equation]
where the last inequality holds for W ∈ W .
It follows that
[Equation]
where J 1 is the Jacobian deﬁned in (31)
[Equation]
and where we use w = vect( W ) ∈
R dk for the vectorized version of the matrix W with a slight abuse of notation.
With this notation, we can bound the ﬁrst term in (56) by
[Equation]
Here, equality (i) follows because (cid:15)
J 1 ( W 0 , v 0 ) w 0 has zero mean, inequality (ii) follows from the
Cauchy-Schwarz inequality as well as from
[Equation]
, and inequality (iii) follows from E [ (cid:107) A (cid:15) (cid:107) 2 ] ≤ (cid:114) E (cid:104)
[Equation]
(cid:105) = (cid:107) A (cid:107) F , by Jensen’s inequality, and from the bounds (cid:107)J 1 ( W 0 , v 0 ) (cid:107) F ≤ ν and (cid:107)J 2 ( W 0 ) (cid:107) F ≤ ω , which holds with probability at least 1 − δ provided that (cid:112) log(2 n/δ ) /
2 k ≤ 1 / 2 , which in turn holds by assumption.
We next upper bound the second term in (56).
Following the argument in (Arora et al., 2019, Proof of Lemma 5.4), we get
[Equation]
with probability at least 1 − δ .
Here, we used that v r ≤ | v 0 ,r | + | v r − v 0 ,r | ≤ 2 ν .
Putting the bounds on the ﬁrst and second term in (56) (given by inequality (57) and inequality (58)) together, we get that, with probability at least 1 − δ , the Rademacher complexity is upper bounded by
[Equation]
which concludes our proof.