U NDERSTANDING W HY G ENERALIZED R EWEIGHTING D OES N OT I MPROVE O VER ERM
Anonymous authors
Paper under double-blind review
A BSTRACT
Empirical risk minimization (ERM) is known to be non-robust in practice to distributional shift where the training and the test distributions are different.
A suite of approaches, such as importance weighting, and variants of distributionally robust optimization (DRO), have been proposed to solve this problem.
But a line of recent work has empirically shown that these approaches do not signiﬁcantly improve over ERM in real applications with distribution shift.
The goal of this work is to obtain a comprehensive theoretical understanding of this intriguing phenomenon.
We ﬁrst posit the class of Generalized Reweighting (GRW) algorithms, as a broad category of approaches that iteratively update model parameters based on iterative reweighting of the training samples.
We show that when overparameterized models are trained under GRW, the resulting models are close to that obtained by ERM.
We also show that adding small regularization which does not greatly affect the empirical training accuracy does not help.
Together, our results show that a broad category of what we term GRW approaches are not able to achieve distributionally robust generalization.
Our work thus has the following sobering takeaway: to make progress towards distributionally robust generalization, we either have to develop non-GRW approaches, or perhaps devise novel classiﬁcation/regression loss functions that are adapted to GRW approaches.
I NTRODUCTION
It has now been well established that empirical risk minimization (ERM) can empirically achieve high test performance on a variety of tasks, particularly with modern overparameterized models where the number of parameters is much larger than the number of training samples.
This strong performance of ERM however has been shown to degrade under distributional shift , where the training and test distributions are different (Hovy & Søgaard, 2015; Blodgett et al., 2016; Tatman, 2017).
There are two broad categories of distribution shift: domain generalization , deﬁned as the scenario where the test distribution contains samples from new domains that did not appear during training; and subpopulation shift , deﬁned as the scenario where the training set contains several subgroups and the testing distribution weighs these subgroups differently, like in fair machine learning.
People have proposed various approaches to learn models robust to distributional shift.
The most classical one is importance weighting (IW) (Shimodaira, 2000; Fang et al., 2020), which reweights training samples; for subpopulation shift these weights are typically set so that each subpopulation has the same overall weight in the training objective.
The approach most widely used today is Distributional Robust Optimization (DRO) (Duchi & Namkoong, 2018; Hashimoto et al., 2018), which assumes that the test distribution belongs to a certain uncertainty set of distributions that are close to the training distribution, and train on the worst distribution in that set.
Many variants of DRO have been proposed and are used in practice (Sagawa et al., 2020a; Zhai et al., 2021a;b).
While these approaches have been developed for the express purpose of improving ERM for distribution shift, a line of recent work has empirically shown the negative result that when used to train overparameterized models, these methods do not improve over ERM.
For IW, Byrd & Lipton (2019) observed that its effect under stochastic gradient descent (SGD) diminishes over training epochs, and ﬁnally does not improve over ERM.
For variants of DRO, Sagawa et al.
(2020a) found that these methods overﬁt very easily, i.e. their test performances will drop to the same low level as ERM after sufﬁciently many epochs if no regularization is applied.
Gulrajani & Lopez-Paz (2021); Koh et al.
compared these methods with ERM on a number of real-world applications, and found that in most cases none of these methods improves over ERM.
This line of empirical results has also been bolstered by some recent theoretical results.
Sagawa et al.
(2020b) constructed a synthetic dataset where a linear model trained with IW is provably not robust to subpopulation shift.
Xu et al.
(2021) further proved that under gradient descent (GD) with a sufﬁciently small learning rate, a linear classiﬁer trained with either IW or ERM converges to the same max-margin classiﬁer, and thus upon convergence, are no different.
These previous theoretical results are limited to linear models and speciﬁc approaches such as IW where sample weights are ﬁxed during training.
They are not applicable to more complex models, and more general approaches where the sample weights could iteratively change, including most DRO variants.
Towards placing the empirical results on a stronger theoretical footing, we deﬁne the class of generalized reweighting (GRW), which dynamically assigns weights to the training samples, and iteratively minimizes the weighted average of the sample losses.
By allowing the weights to vary with iterations, we cover not just static importance weighting, but also DRO approaches outlined earlier; though of course, the GRW class is much broader than just these instances.
We prove that GRW and ERM have (almost) equivalent implicit biases, in the sense that the points they converge to are very close to each other , under a much more general setting than those used in previous work.
Thus, GRW cannot improve over ERM because it does not yield a signiﬁcantly different model.
We are the ﬁrst to extend this line of theoretical results (i) to wide neural networks, (ii) to reweighting methods with dynamic weights, (iii) to regression tasks, and (iv) to methods with L 2 regularization.
We note that these extensions are non-trivial technically as they require the result that wide neural networks can be approximated by their linearized counterparts to hold uniformly throughout the iterative process of GRW algorithms.
Moreover, we ﬁx the proof in a previous paper (Lee et al., 2019)
(see Appendix E) which is also a great contribution.
Overall, the important takeaway is that distributionally robust generalization (DRG) cannot be directly achieved by the broad class of GRW algorithms (which includes popular approaches such as importance weighting and most DRO variants).
Progress towards this important goal thus requires either going beyond GRW algorithms, or devising novel loss functions that are adapted to GRW approaches.
In Section 6 we will discuss some promising future directions and the case with nonoverparameterized models and early stopping.
Finally, we want to emphasize that while the models we use in our results (linear models and wide neural networks) are different from practical models, they are general models most widely used in existing theory papers, and our results based on these models provide explanations to the bafﬂing observations made in previous empirical work, as well as valuable insights into how to improve distributionally robust generalization.
P RELIMINARIES
Let the input space be X ⊆ R d and the output space be Y ⊆ R .
1 We assume that X is a subset of the unit L 2 ball of R d , so that any x ∈ X satisﬁes (cid:107) x (cid:107) 2 ≤ 1 .
We have a training set { z i =
( x i , y i ) } ni = i.i.d.
sampled from an underlying distribution P over X × Y .
Denote X
=
( x 1 , · · · , x n ) ∈ R d × n , and Y =
( y 1 , · · · , y n ) ∈ R n .
For any function g : X (cid:55)→ R m , we overload notation and use g ( X ) =
( g ( x 1 ) , · · · , g ( x n ))
∈
R m × n (except when m = 1 , g ( X ) is deﬁned as a column vector).
Let the loss function be (cid:96) :
Y × Y → [0 , 1] .
ERM trains a model by minimizing its expected risk
R ( f ; P ) =
E z ∼ P
[ (cid:96) ( f ( x ) , y )] via minimizing the empirical risk
ˆ R ( f )
= 1 n (cid:80) ni
=1 (cid:96) ( f ( x i ) , y i ) .
In distributional shift, the model is evaluated not on the training distribution P , but a different test distribution P test , so that we care about the expected risk R ( f ; P test ) .
A large family of methods designed for such distributional shift is distributionally robust optimization (DRO), which minimizes the expected risk over the worst-case distribution Q (cid:28) P 2 in a ball w.r.t.
divergence D around the training distribution P .
Speciﬁcally, DRO minimizes the expected DRO risk deﬁned as:
[Equation]
for ρ > 0 .
Examples include CVaR, χ 2 -DRO (Hashimoto et al., 2018), and DORO (Zhai et al., 2021a), among others.
A common category of distribution shift is known as subpopulation shift.
Let the data domain contain K groups D 1 , · · · , D K .
The training distribution P is the distribution over all groups, and the test distribution P test is the distribution over one of the groups.
Let P k ( z )
=
P ( z | z ∈ D k ) be the conditional distribution over group k , then P test can be any one of P 1 , · · · , P k .
The goal is to train a model f that performs well over every group.
There are two common ways to achieve this goal: one is minimizing the balanced empirical risk which is an unweighted average of the empirical risk over each group, and the other is minimizing the worst-group risk deﬁned as
[Equation]
G ENERALIZED R EWEIGHTING (GRW)
Various methods have been proposed towards learning models that are robust to distributional shift.
In contrast to analyzing each of these individually, we instead consider a large class of what we call
Generalized Reweighting (GRW) algorithms that includes the ones mentioned earlier, but potentially many others more.
Loosely, GRW algorithms iteratively assign each sample a weight during training (that could vary with the iteration) and iteratively minimize the weighted average risk.
Speciﬁcally, at iteration t , GRW assigns a weight q ( t ) i to sample z i , and minimizes the weighted empirical risk:
[Equation]
where q ( t ) = ( q ( t ) 1 , · · · , q ( t ) n ) and q ( t ) 1 + · · · + q ( t ) n = 1 .
Static GRW assigns to each z i =
( x i , y i ) a ﬁxed weight q i that does not change during training, i.e. q ( t ) i ≡ q i .
A classical method is importance weighting (IW) (Shimodaira, 2000), where if z i ∈
D k and the size of D k is n k , then q i =
( Kn k ) − 1 .
Under IW, (3) becomes the balanced empirical risk in which each group has the same weight.
Note that ERM is also a special case of static GRW.
On the other hand, in dynamic GRW , q ( t ) changes with t .
For instance, any approach that iteratively upweights samples with high losses in order to help the model learn “hard” samples, such as DRO, is an instance of GRW.
When estimating the population DRO risk R D,ρ ( f ; P ) in Eqn.
(1), if P is set to the empirical distribution over the training samples, then Q (cid:28) P implies that Q is also a distribution over the training samples.
Thus, DRO methods belong to the broad class of GRW algorithms.
There are two common ways to implement DRO.
One uses Danskin’s theorem and chooses Q as the maximizer of E Q [ (cid:96) ( f ( x ) , y )] in each epoch.
The other one formulates DRO as a bi-level optimization problem, where the lower level updates the model to minimize the expected risk over Q , and the upper level updates Q to maximize it.
Both can be seen as instances of GRW.
As one popular instance of the latter, Group DRO was proposed by Sagawa et al.
(2020a) to minimize (2).
Denote the empirical risk over group k by ˆ R k ( f ) , and the model at time t by f ( t ) .
Group DRO iteratively sets q ( t ) i = g ( t )
k /n k
for all z i ∈ D k where g ( t ) k is the group weight that is updated as
[Equation]
(4) for some ν > 0 , and then normalized so that q ( t ) 1 + · · · + q ( t ) n = 1 .
Sagawa et al.
(2020a) then showed (in their Proposition 2) that for convex settings, the Group DRO risk of iterates converges to the global minimum with the rate O ( t − 1 / 2 ) if ν is sufﬁciently small.
T HEORETICAL R ESULTS FOR R EGRESSION
In this section, we will study GRW for regression tasks that use the squared loss
[Equation]
We will prove that for both linear models and sufﬁciently wide fully-connected neural networks, the implicit bias of GRW is equivalent to ERM, which means that starting from the same initial point,
GRW and ERM will converge to the same point when trained for an inﬁnitely long time. Thus,
GRW cannot improve over ERM as it produces the exact same model as ERM.
We will further show that while regularization can affect this implicit bias, it must be large enough to signiﬁcantly lower the training performance , or the ﬁnal model will still be similar to the unregularized ERM model.
4000 Epochs 0.
N o r m o f
D i ff e r e n c e |θ iw −θ erm | 2 |θ gdro −θ erm |(a) Weight Difference  4000 Epochs 0.
N o r m o f
E R M M o d e l (b) Norm of ERM Model  4000 Epochs 0.
T r a i n i n g
L o ss
ERM IW GDRO (c) Training Loss  4000 Epochs 0.
G r o u p
W e i g h t (
[Figure]
Group Weights
4. L INEAR M ODELS
We ﬁrst demonstrate our result on simple linear models to provide our readers with a key intuition which we will later apply to neural networks.
This key intuition draws from results of Gunasekar et al.
Let the linear model be denoted by f ( x ) =
(cid:104) θ, x (cid:105) , where θ ∈ R d .
We consider the overparameterized setting where d > n .
The weight update rule of GRW under GD is the following:
[Equation]
where η > 0 is the learning rate.
For a linear model with the squared loss, the update rule is
[Equation]
For this training scheme, we can prove that if the training error converges to zero, then the model converges to an interpolator θ ∗ (
s.t. ∀
i , (cid:104) θ ∗ , x i (cid:105) =
y i ) independent of q ( t ) i (proofs in Appendix D):
Theorem 1.
If x 1 , · · · , x n are linearly independent, then under the squared loss, for any GRW such that the empirical training risk ˆ R ( f ( t ) )
→ 0 as t → ∞ , it holds that θ ( t ) converges to an interpolator θ ∗ that only depends on θ (0) and x 1 , · · · , x n , but does not depend on q ( t ) i .
The proof is based on the following key intuition regarding the update rule (7): θ ( t +1) − θ ( t ) is a linear combination of x 1 , · · · , x n for all t , so θ ( t ) − θ (0) always lies in the linear subspace span { x 1 , · · · , x n } , which is an n -dimensional linear subspace if x 1 , · · · , x n are linearly independent.
By Cramer’s rule, there is exactly one ˜ θ in this subspace such that we get interpolation of all the data (cid:104) ˜ θ
+ θ (0) , x i (cid:105) =
y i for all i ∈ { 1 , . . . , n } .
In other words, the parameter θ ∗ =
˜ θ
+ θ (0) in this subspace that interpolates all the data is unique.
Thus the proof would follow if we were to show that θ ( t ) − θ (0) , which lies in the subspace, also converges to interpolating the data.
Moreover, this proof works for any ﬁrst-order optimization method such that the training risk converges to 0.
We have essentially proved the following sobering result: any GRW algorithm that achieves zero training error exactly produces the ERM model, so it does not improve over ERM .
While the various distributional shift methods discussed in the introduction have been shown to satisfy the precondition of convergence to zero training error with overparameterized models and linearly independent inputs (Sagawa et al., 2020a), we provide the following theorem that shows this for the broad class of GRW methods.
Speciﬁcally, we show this result for any GRW method that satisﬁes the following assumption with a sufﬁciently small learning rate:
Assumption 1.
There are constants q 1 , ·
· · , q n s.t. ∀
i , q ( t ) i
→ q i as t → ∞ .
And min i q i =
q ∗ > 0 .
Theorem 2.
If x 1 , · · · , x n are linearly independent, then there exists η 0 > 0 such that for any
GRW satisfying Assumption 1 with the squared loss, and any η ≤ η 0 , the empirical training risk ˆ
R ( f ( t ) )
→ 0 as t → ∞ .
Finally, we use a simple experiment to demonstrate the correctness of this result.
The experiment is conducted on a training set of six MNIST images, ﬁve of which are digit 0 and one is digit 1.
We use a 784-dimensional linear model and run ERM, importance weighting and group DRO.
The results are presented in Figure 1, and they show that the training loss of each method converges to 0, and the gap between the model weights of importance weighting, Group DRO and ERM converges to 0, meaning that all three model weights converge to the same point, whose L 2 norm is about 0.63.
Figure 1d also shows that the group weights in Group DRO empirically satisfy Assumption 1.
4. W IDE N EURAL N ETWORKS (W IDE NN S )
Now we study sufﬁciently wide fully-connected neural networks .
We extend the analysis in Lee et al.
(2019) in the neural tangent kernel (NTK) regime (Jacot et al., 2018).
In particular we study the following network:
[Equation]
and
[Equation]
where σ is a non-linear activation function,
W l ∈ R d l +1 ×
d l and W L ∈ R 1 × d L .
Here d 0 = d .
The parameter vector θ consists of W 0 , · · · , W L and b 0 , · · · , b L ( θ is the concatenation of all ﬂattened weights and biases).
The ﬁnal output is f ( x ) =
h L +1 .
And let the neural network be initialized as
[Equation]
and
[Equation]
We also need the following assumption on the wide neural network:
Assumption 2.
σ is differentiable everywhere.
Both σ and its ﬁrst-order derivative ˙
σ are Lipschitz.
Difference from Jacot et al.
Our initialization (9) differs from the original one in Jacot et al.
(2018) in the last (output) layer, where we use the zero initialization W L (0) i,j = 0 instead of the
Gaussian initialization W L (0) i,j ∼ N (0 , 1) .
This modiﬁcation permits us to accurately approximate the NN with its linearized counterpart (11), as we notice that the proofs in Lee et al.
(2019) (particularly the proofs of their Theorem 2.1 and their Lemma 1 in Appendix G) are ﬂawed.
In Appendix E we will explain what went wrong in their proofs and how we ﬁx it with this modiﬁcation.
Denote the neural network at time t by f ( t ) ( x ) =
f
( x ; θ ( t ) ) which is parameterized by
θ ( t ) ∈
R p where p is the number of parameters.
We use the shorthand ∇ θ f (0) ( x ) := ∇ θ f ( x ; θ ) (cid:12)(cid:12) θ
= θ 0 .
The neural tangent kernel (NTK) of this model is Θ (0) ( x , x (cid:48) )
= ∇ θ f (0) ( x ) (cid:62) ∇ θ f (0) ( x (cid:48) ) , and the Gram matrix is Θ (0) =
Θ (0) ( X , X ) ∈ R n × n .
For this wide NN, we still have the following NTK theorem:
Lemma 3.
If σ is Lipschitz and d l → ∞ for l = 1 , · · · , L sequentially, then Θ (0) ( x , x (cid:48) ) converges in probability to a non-degenerate 4 deterministic limiting kernel Θ( x , x (cid:48) ) .
The kernel Gram matrix Θ = Θ( X , X ) ∈ R n × n is a positive semi-deﬁnite symmetric matrix.
Denote its largest and smallest eigenvalues by λ max and λ min .
Note that Θ is non-degenerate, so we can assume that λ min > 0 (which is almost surely true when d L (cid:29) n ).
Then we have:
Theorem 4.
Let f ( t ) be a wide fully-connected neural network that satisﬁes Assumption 2 and is trained by any GRW satisfying Assumption 1 with the squared loss.
Let f ( t ) ERM be the same model trained by ERM from the same initial point.
If d 1 = · · · = d L = ˜ d , ∇ θ f (0) ( x 1 ) , · · · , ∇ θ f (0) ( x n ) are linearly independent, and λ min > 0 , then there exists a constant η 1 > 0 such that: if η ≤ η 1 5 , then for any δ > 0 , there exists ˜ D > 0 such that as long as ˜ d ≥ ˜ D , with probability at least (1 − δ ) over random initialization we have: for any test point x ∈ R d such that (cid:107) x (cid:107) 2 ≤ 1 , as ˜ d → ∞
[Equation]
Essentially this theorem implies that on any test point x in the unit ball, the GRW model and the
ERM model produce almost the same output, so they have almost the same performance.
Note that for simplicity, we only prove for d 1 = · · · = d L = ˜
d → ∞ , but the result can be very easily extended to the case where d l /d 1 → α l for l = 2 , · · · , L for some constants α 2 , · · · , α L , and d 1 → ∞ .
The key to proving this theorem is to consider the linearized neural network of f ( t ) ( x ) : [Equation]
[Equation]
which is a linear model w.r.t. ∇ θ f (0) ( x ) .
If ∇ θ f (0) ( x 1 ) , · · · , ∇ θ f (0) ( x n ) are linearly independent (which is almost surely true when the model is overparameterized so that θ has a very high dimension), then our key intuition tells us that the linearized NN will converge to the unique interpolator.
Then we show that the wide NN can be approximated by its linearized counterpart uniformly throughout training , which is considerably more subtle in our case due to the GRW dynamics.
Here we prove the upper bound O ( ˜ d − 1 / 4 ) , but in fact the upper bound can be O ( ˜ d − 1 / 2+ (cid:15) ) for any (cid:15) > 0 :
Lemma 5 (Approximation Theorem) .
For a wide fully-connected neural network f ( t ) satisfying
Assumption 2 and is trained by any GRW satisfying Assumption 1 with the squared loss, let f ( t ) lin be its linearized neural network trained by the same GRW (i.e. q ( t ) i are the same for both networks for any i and t ).
Under the conditions of Theorem 4, with a sufﬁciently small learning rate, for any δ > 0 , there exist constants ˜ D > 0 and C > 0 such that as long as ˜ d ≥ ˜ D , with probability at least (1 − δ ) over random initialization we have: for any test point x ∈ R d such that (cid:107) x (cid:107) 2 ≤ 1 ,
[Equation] [Equation]
This lemma essentially says that throughout the GRW training process, on any test point x in the unit ball, the linear NN and the wide NN produce almost the same output.
So far, we have shown that in a regression task, for both linear and wide NNs, GRW does not improve over ERM.
4. W IDE N EURAL N ETWORKS , WITH L 2 R EGULARIZATION
Previous work such as Sagawa et al.
(2020a) proposed to improve GRW by adding L 2 penalty to the objective function.
In this section, we thus study adding L 2 regularization to GRW algorithms:
[Equation]
At ﬁrst sight, adding regularization seems to be a natural approach and should make a difference.
Indeed, from the outset, we can easily show that with L 2 regularization, the GRW model and the ERM model are different unlike the case without regularization.
As an concrete example, when f is a linear model, (cid:96) is convex and smooth, then ˆ R µ q ( t ) ( f ) with static GRW is a convex smooth objective function, so under GD with a sufﬁciently small learning rate, the model will converge to the global minimizer (see Appendix D.1).
Moreover, the global optimum θ ∗ satisﬁes ∇ θ ˆ R µ q ( t ) ( f ( x ; θ ∗ ))
= 0 , solving which yields θ ∗ =
θ (0) + ( XQX (cid:62) + µ I ) − 1 XQ ( Y − f (0) ( X )) , which depends on Q = diag ( q 1 , · · · , q n ) , so adding L 2 regularization at least seems to yield different results from ERM (albeit whether it improves over ERM might depend on q 1 , · · · , q n ).
However, the following result shows that this regularization must be large enough to signiﬁcantly lower the training performance , or the ﬁnal model would still be close to the unregularized ERM model.
We still denote the largest and smallest eigenvalues of the kernel Gram matrix Θ by λ max and λ min .
We use the subscript “reg” to refer to a regularized model (trained by minimizing (13)).
Theorem 6.
Suppose there exists M 0 > 0 s.t.
(cid:13)(cid:13)
∇ θ f (0) ( x ) (cid:13)(cid:13) 2 ≤ M 0 for all (cid:107) x (cid:107) 2 ≤ 1 .
If λ min >and µ > 0 , then for a wide NN satisfying Assumption 2, and any GRW minimizing the squared loss with a sufﬁciently small learning rate η , if d 1 = d 2 = · · · = d L = ˜ d , ∇ θ f (0) ( x 1 ) , · · · , ∇ θ f (0) ( x n ) are linearly independent, and the empirical training risk of f ( t ) reg satisﬁes
[Equation]
for some (cid:15) > 0 , then with a sufﬁciently small learning rate, as ˜ d → ∞ , with probability close toover random initialization, for any x such that (cid:107) x (cid:107) 2 ≤ 1 we have
[Equation]
where f ( t ) reg is trained by regularized GRW and f ( t ) ERM by unregularized ERM from same initial points.
This theorem essentially says that if the regularization is too small and the training error is close to zero, then the regularized GRW model still produces an output very close to that of the ERM model on any test point x in the unit ball.
Thus, a small regularization makes little difference.
The proof again starts from analyzing linearized NNs, and showing that regularization does not help there (Appendix D.4.2).
Then, we prove a new approximation theorem for L 2 regularized GRW connecting wide NNs to linearized NNs uniformly throughout training (Appendix D.4.1).
With regularization, we no longer need Assumption 1 to prove the new approximation theorem which was used to prove the convergence of GRW, because with regularization GRW naturally converges.
4000 Epochs 0.
N o r m o f
D i ff e r e n c e |θ iw −θ erm | 2 |θ gdro −θ erm |(a) Weight Difference  4000 Epochs 0.
T r a i n i n g
L o ss
IW GDRO ERM (b) Training Loss  4000 Epochs 0.
N o r m o f
D i ff e r e n c e |θ iw −θ erm | 2 |θ gdro −θ erm |
[Figure]
Weight Difference  4000 Epochs 0.
T r a i n i n g
L o ss
IW GDRO
ERM
[Figure]
To empirically demonstrate this result, we run the same experiment as in Section 4.1 but with Lregularization.
The results are presented in Figure 2.
We can see that when the regularization is small, the training losses still converge to 0, and the three model weights still converge to the same point.
On the contrary, with a large regularization, the training loss does not converge to 0, and the three model weights converge to different points.
This shows that the regularization must be large enough to lower the training performance to make a signiﬁcant difference to the implicit bias.
T HEORETICAL R ESULTS FOR C LASSIFICATION
Now we consider classiﬁcation where
Y =
+1 , − 1 } .
The big difference is that classiﬁcation losses don’t have ﬁnite minimizers .
A classiﬁcation loss converging to zero means that the model weight “explodes” to inﬁnity instead of converging to a ﬁnite point.
We focus on the canonical logistic loss:
[Equation]
5. L INEAR M ODELS
We ﬁrst consider training the linear model f ( x ) =
(cid:104) θ, x (cid:105) with GRW under gradient descent with the logistic loss.
As noted earlier, in this setting, Byrd & Lipton (2019) made the empirical observation that importance weighting does not improve over ERM.
Then, Xu et al.
(2021) proved that for importance weighting algorithms, as t → ∞ , (cid:107) θ ( t ) (cid:107) 2 → ∞ and θ ( t ) / (cid:107) θ ( t ) (cid:107) 2 converges to a unit vector that does not depend on the sample weights, so it does not improve over ERM.
To extend this theoretical result to the broad class of GRW algorithms, we will prove two results.
First, in Theorem 7 we will show that for the logistic loss and any GRW algorithm satisfying the weaker assumption:
Assumption 3.
For all i , lim inf t →∞ q ( t ) i
[Equation]
if the training error converges to 0, and the direction of the model weight converges to a ﬁxed unit vector, then this unit vector must be the max-margin classiﬁer deﬁned as
[Equation]
Second, Theorem 8 shows that for any GRW satisfying Assumption 1, the training error converges to 0 and the direction of the model weight converges, so it does not improve over ERM.
Theorem 7.
If x 1 , · · · , x n are linearly independent, then for the logistic loss, we have: for any
GRW satisfying Assumption 3, if as t → ∞ the empirical training risk ˆ R ( f ( t ) ) converges to 0 and θ ( t ) / (cid:107) θ ( t ) (cid:107) 2 → u for some unit vector u , then u
= ˆ θ MM .
This result is an extension of Soudry et al.
(2018), and says that all GRW methods (including ERM) make the model converge to the same point ˆ θ MM that does not depend on q ( t ) i . In other words, the samples weights do not affect the implicit bias.
Thus, for any GRW method that only satisﬁes the weak Assumption 3, as long as the training error converges to 0 and the model weight direction converges, GRW does not improve over ERM.
We next show that any GRW satisfying Assumption 1 does have its model weight direction converge, and its training error converge to 0.
Theorem 8.
For any loss (cid:96) that is convex, L -smooth in ˆ y and strictly monotonically decreasing to zero as y ˆ
y
→ + ∞ , and GRW satisfying Assumption 1, denote F ( θ )
= (cid:80) n i =
1 q i (cid:96) ( (cid:104) θ, x i (cid:105) , y i ) .
If x 1 , · · · , x n are linearly independent, then with a sufﬁciently small learning rate η , we have: F ( θ ( t ) ) → 0 as t → ∞ . (i)
[Equation]
Let θ R = arg min θ { F ( θ ) : (cid:107) θ (cid:107) 2 ≤ R } .
θ R is unique for any R such that min (cid:107) θ (cid:107) 2 ≤ R F ( θ ) <
min i q i (cid:96) (0 , y i ) .
And if lim R →∞ θ R R exists, then lim t →∞
[Equation]
t ) (cid:107) θ ( t ) (cid:107) 2 also exists and they are equal.
(iii)
This result is an extension of Theorem 1 of Ji et al.
For the logistic loss, it is easy to show that it satisﬁes the conditions of the above theorem and lim R →∞ θ R R
=
ˆ θ MM .
Thus, Theoremsand 7 together imply that all GRW satisfying Assumption 1 (including ERM) have the same implicit bias (see Appendix D.5.3).
We also have empirical veriﬁcation for these results (see Appendix C).
Remark.
It is impossible to extend these results to wide NNs like Theorem 4 because for a neural network, if (cid:107) θ ( t ) (cid:107) 2 goes to inﬁnity, then (cid:107)∇ θ f (cid:107) 2 will also go to inﬁnity.
However, for a linear model, the gradient is a constant.
Consequently, the gap between the neural networks and its linearized counterpart will “explode” under gradient descent, so there can be no approximation theorem like
Lemma 5 that can connect wide NNs to their linearized counterparts.
Thus, we consider regularized
GRW, for which θ ( t ) converges to a ﬁnite point and there is an approximation theorem.
5. W IDE N EURAL N ETWORKS , WITH L 2 R EGULARIZATION
Consider minimizing the regularized weighted empirical risk (13) with (cid:96) being the logistic loss.
As in the regression case, with L 2 regularization, GRW methods have different implicit biases than
ERM for the same reasons as in Section 4.3.
And similarly, we can show that in order for GRW methods to be sufﬁciently different from ERM, the regularization needs to be large enough to signiﬁcantly lower the training performance.
Speciﬁcally, in the following theorem we show that if the regularization is too small to lower the training performance, then a wide neural network trained with regularized GRW and the logistic loss will still be very close to the max-margin linearized neural network :
[Equation]
Note that f MM does not depend on q ( t ) i .
Moreover, using the result in the previous section we can show that a linearized neural network trained with unregularized ERM will converge to f MM :
Theorem 9.
Suppose there exists M 0 > 0 such that (cid:13)(cid:13) ∇ θ f (0) ( x ) (cid:13)(cid:13) 2 ≤ M 0 for all test point x .
For a wide NN satisfying Assumption 2, and for any GRW satisfying Assumption 1 with the logistic loss, if d 1 = d 2 = · · · = d L = ˜ d and ∇ θ f (0) ( x 1 ) , · · · , ∇ θ f (0) ( x n ) are linearly independent and the learning rate is sufﬁciently small, then for any δ > 0 there exists a constant C > 0 such that: with probability at least (1 − δ ) over random initialization, as ˜
d → ∞ we have: for any (cid:15) ∈ (0 , 14 ) , if the empirical training error satisﬁes lim sup t →∞ ˆ R ( f ( t ) reg ) < (cid:15) , then for any test point x such that | f MM ( x ) | > C · ( − log 2 (cid:15) ) − 1 / 2 , f ( t ) reg ( x ) has the same sign as f MM ( x ) when t is sufﬁciently large.
This result says that at any test point x on which the max-margin linear classiﬁer classiﬁes with a margin of Ω(( − log 2 (cid:15) ) − 1 / 2 ) , the neural network has the same prediction.
And as (cid:15) decreases, the conﬁdence threshold also becomes lower.
Similar to Theorem 6, this theorem provides the scaling of the gap between the regularized GRW model and the unregularized ERM model w.r.t.
(cid:15) .
This result justiﬁes the empirical observation in Sagawa et al.
(2020a)
that with large regularization, some GRW algorithms can maintain a high worst-group test performance, with the cost of suffering a signiﬁcant drop in training accuracy.
On the other hand, if the regularization is small and the model can achieve nearly perfect training accuracy, then its worst-group test performance will still signiﬁcantly drop.
D ISCUSSION 6. D ISTRIBUTIONALLY R OBUST G ENERALIZATION AND F UTURE D IRECTIONS
A large body of prior work focused on distributionally robust optimization, but we show that these methods have (almost) equivalent implicit biases as ERM.
In other words, distributionally robust optimization (DRO) does not necessarily achieve better distributionally robust generalization (DRG).
Our results pinpoint a critical bottleneck in the current distribution shift research, and we argue that a deeper understanding in DRG is crucial for developing better distributionally robust training algorithms.
Here we discuss three promising future directions to improving DRG.
The ﬁrst approach is data augmentation and pretraining on large datasets.
Our theoretical ﬁndings suggest that the implicit bias of GRW is determined by the training samples and the initial point, but not the sample weights.
Thus, to improve DRG, we can either obtain more training samples, or start from a better initial point, as proposed in two recent papers (Wiles et al., 2022; Sagawa et al., 2022).
The second approach (for classiﬁcation) is to go beyond the class of (iterative) sample reweighting based GRW algorithms, for instance via logit adjustment (Menon et al., 2021), which makes a classiﬁer have larger margins on smaller groups to improve its generalization on smaller groups.
An early approach by Cao et al.
(2019) proposed to add an O ( n − 1 / 4 k ) additive adjustment term to the logits output by the classiﬁer.
Following this spirit, Menon et al.
(2021) proposed the LA-loss which also adds an additive adjustment term to the logits.
Ye et al.
(2020) proposed the CDT-loss which adds a multiplicative adjustment term to the logits by dividing the logits of different classes with different temperatures.
Kini et al.
(2021) proposed the VS-loss which includes both additive and multiplicative adjustment terms, and they showed that only the multiplicative adjustment term affects the implicit bias, while the additive term only affects optimization, a fact that can be easily derived from our Theorem 8.
Finally, Li et al.
(2021a) proposed AutoBalance which optimizes the adjustment terms with a bi-level optimization framework.
The third approach is to stay within the class of GRW algorithms, but to change the classiﬁcation/regression loss function to be suited to GRW.
A recent paper (Wang et al., 2022) showed that for linear classiﬁers, one can make the implicit bias of GRW dependent on the sample weights by replacing the exponentially-tailed logistic loss with the following polynomially-tailed loss :
[Equation]
And this result can be extended to GRW satisfying Assumption 1 using our Theorem 8.
The reason why loss (19) works is that it changes lim R →∞ θ R R , and the new limit depends on the sample weights.
6. L IMITATIONS
Like most theory papers, our work makes some strong assumptions.
The two main assumptions are:
(i)
The model is a linear model or a sufﬁciently wide fully-connected neural network.
(ii)
The model is trained for sufﬁciently long time,
i
e. without early stopping.
Regarding (i), Chizat et al.
(2019) argued that NTK neural networks fall in the “lazy training” regime and results might not be transferable to general neural networks.
However, this class of neural networks has been widely studied in recent years and has provided considerable insights into the behavior of general neural networks, which is hard to analyze otherwise.
Regarding (ii), in some easy tasks, when early stopping is applied, existing algorithms for distributional shift can do better than ERM (Sagawa et al., 2020a).
However, as demonstrated in Gulrajani & Lopez-Paz (2021); Koh et al.
(2021), in real applications these methods still cannot signiﬁcantly improve over ERM even with early stopping, so early stopping is not the ultimate universal solution.
Thus, though inevitably our results rely on some strong assumptions, we believe that they provide important insights into the problems of existing methods and directions for future work, which are signiﬁcant contributions to the study of distributional shift problems.
C ONCLUSION
In this work, we posit a broad class of what we call Generalized Reweighting (GRW) algorithms that include popular approaches such as importance weighting, and Distributionally Robust Optimization (DRO) variants, that were designed towards the task of learning models that are robust to distributional shift.
We show that when used to train overparameterized linear models or wide NN models, even this very broad class of GRW algorithms does not improve over ERM, because they have the same implicit biases.
We also showed that regularization does not help if it is not large enough to signiﬁcantly lower the average training performance.
Our results thus suggest to make progress towards learning models that are robust to distributional shift, we have to either go beyond this broad class of GRW algorithms, or design new losses speciﬁcally targeted to this class.
A R ELATED WORK A. S UBPOPULATION SHIFT
In this work, we mainly focus on the subpopulation shift problem, which has two main applications: group fairness and long-tailed learning (learning with class imbalance).
In both applications, the dataset can be divided into several subgroups, and this work considers minimizing the worst-group risk, deﬁned as the maximum risk over any group.
Group Fairness.
Group fairness refers to the scenario where the dataset contains several “groups” (such as demographic groups), and a model is considered fair if its per-group performances meet certain criteria (a “fairness” notion).
Group fairness in machine learning was ﬁrst studied in Hardt et al.
(2016) and Zafar et al.
(2017), where they required the model to perform equally well over all groups.
Many previous papers proposed a number of fairness notions, such as equal opportunity, statistical parity, etc.
Among them, Hashimoto et al.
(2018) studied another type of group fairness called Rawlsian max-min fairness Rawls (2001), which does not require equal performance but rather requires high performance on the worst-off group.
The subpopulation shift problem we study in this paper is naturally connected to the Rawlsian max-min fairness.
A large body of recent work have studied how to improve this worst-group performance Duchi & Namkoong (2018);
Oren et al.
Liu et al.
; Zhai et al.
(2021a).
Recent work however observe that these approaches, when used with modern overparameterized models, easily overﬁt Sagawa et al.
(2020a;b).
Apart from group fairness, there are also other notions of fairness, such as individual fairness Dwork et al.
(2012); Zemel et al.
and counterfactual fairness Kusner et al.
(2017), which we do not study in this work.
Long-tailed Learning. Long-tailed learning refers to the scenario where different classes have different sizes, and usually there are some “minority classes” with extremely few samples that are much more difﬁcult to learn than the other classes. Using GRW such as importance weighting for long-tailed learning is a very old idea which dates back to Xie & Manski (1989). However, recently
Byrd & Lipton (2019) found that the effect of importance weighting for long-tailed learning diminishes as training proceeds, which leads to a line of recent work on how to improve the generalization in long-tailed learning (Cao et al., 2019; Menon et al., 2021;
Ye et al., 2020; Kim & Kim, 2020; Kini et al., 2021).
Most of these papers share a common idea:
Forcing the model to have larger margins on smaller groups, so that its generalization on smaller groups can be better.
Self-supervised learning is also used in long-tailed learning.
For instance, Liu et al.
(2022) found that self-supervised learning can achieve good performances in long-tailed learning, Wang et al.
(2021) used contrastive learning for long-tailed learning, and Li et al.
(2021b) used self-distillation.
A. D OMAIN GENERALIZATION
Domain generalization is the second common type of distribution shift.
In domain generalization and the related domain adaptation, a model is tested on a different domain than what it is trained on.
The most common idea in domain generalization is invariant learning , which learns a feature extractor that is invariant across domains, usually by matching the feature distribution of different domains.
Since we have no access to the target domain, in invariant learning we assume that we have access to multiple domains in the training set, and we learn a feature extractor with a small variance across these domains.
Examples include CORAL (Sun & Saenko, 2016), DANN (Ganin et al., 2016), MMD (Li et al., 2018) and IRM (Arjovsky et al., 2019).
However, Gulrajani & LopezPaz (2021); Koh et al. (2021) empirically showed that most of these methods cannot do better than standard ERM, and Rosenfeld et al.
(2021) theoretically proved that IRM cannot do better than ERM unless the number of training domains is greater than the number of independent features.
One problem of invariant learning methods is that they do not necessarily align the classes.
For a source domain P and a target domain Q , even if we have successfully learned a feature extractor Φ such that Φ( P ) ≈ Φ( Q ) , there is no guarantee that Φ can map the samples in P and Q from the same class to the same location in the feature space.
In the worst case, Φ can map the positive samples in P and the negative samples in Q to the same location and vice versa, in which case 100% accuracy over P means 0% accuracy over Q .
The goal of class alignment is to make sure that samples from the same class are mapped together, and far away from the other classes.
For example, Tzeng et al.
(2015) used soft labels to align the classes, Long et al.
(2016) minimized the class-based cross entropy on the target domain while keeping the source and target classiﬁers close with a residual block, and Motiian et al.
(2017) adopted a similarity penalty to keep samples from different classes away from each other.
A. I MPLICIT BIAS UNDER THE OVERPARAMETERIZED SETTING
For overparameterized models, there could be many model parameters which all minimize the training loss.
In such cases, it is of interest to study the implicit bias of speciﬁc optimization algorithms such as gradient descent i.e. to what minimizer the model parameters will converge to Du et al.
(2019); Allen-Zhu et al.
Our results use the NTK formulation of wide neural networks Jacot et al.
(2018), and speciﬁcally we use linearized neural networks to approximate such wide neural networks following Lee et al.
There is some criticism of this line of work, e.g. Chizat et al.
(2019) argued that inﬁnitely wide neural networks fall in the “lazy training” regime and results might not be transferable to general neural networks.
Nonetheless such wide neural networks are being widely studied in recent years, since they provide considerable insights into the behavior of more general neural networks, which are typically intractable to analyze otherwise.
A. C OMPARISON WITH H U ET AL . (2018)
A prior work Hu et al.
(2018) also proved that GRW is equivalent to ERM under certain conditions.
However, we would like to point out that this work is substantially different from Hu et al.
Hu et al.
(2018) proved that in classiﬁcation that uses the zero-one loss, GRW methods such as DRSL are equivalent to ERM, in the sense that the minimizer of the DRSL risk is also the minimizer of the average risk. However, this does not mean that DRSL and ERM will always converge to the same point , as there could be multiple minimizers. Their result relies on the zero-one loss, which leads to a monotonic linear relationship between the DRSL risk and the average risk. Moreover, their result is only about the relationship between two minimizers, and they did not prove that DRSL and ERM can actually converge to these global minima.
On the other hand, in our results, we ﬁrst show that without regularization, GRW and ERM will converge to the exact same point , so that they have equivalent implicit biases, which is a much stronger result.
Then we show that even with regularization, if the regularization is not large enough,
GRW will still converge to a point that is very close to the point ERM converges to.
Our results do not depend on the loss function, and work for both the squared loss for regression and the logistic loss for classiﬁcation (and can be extended to other losses).
Instead, our results depend on the optimization method (must be ﬁrst-order or gradient-based) as well as the model architecture (linear or wide NN), since we need to explicitly prove that both GRW and ERM can reach the global minima if trained under a small learning rate for sufﬁciently long.
In a word, Hu et al.
(2018) proves the equivalence between GRW and ERM under the zero-one loss with a monotonic relationship between the two risk functions, while our results focus on the optimization and training dynamics, and prove that GRW and ERM have almost equivalent implicit biases.
B E XTENSION TO M ULTI - DIMENSIONAL R EGRESSION / M ULTI - CLASS C LASSIFICATION
In our results, we assume that f : R d → R for simplicity, but our results can be very easily extended to the case where f : R d → R k .
For most of our results, the proof consists of two major components: (i) The linearized neural network will converge to some point (interpolator, max-margin classiﬁer, etc.); (ii) The wide fully-connected neural network can be approximated by its linearized counterpart.
For both components, the extension is very simple and straightforward.
For (i), the proof only relies on the smoothness of the objective function and the upper quadratic bound it entails, and the function is still smooth when its output becomes multi-dimensional;
For (ii), we can prove that sup t (cid:107) f
( x ) − f lin ( x ) (cid:107)
= O ( ˜ d − 1 / 4 ) in exactly the same way.
Thus, all of our results hold for multi-dimensional regression and multi-class classiﬁcation.
Particularly, for the multi-class cross-entropy loss, using Theorem 8 we can show that under any
GRW satisfying Assumption 1, the direction of the weight of a linear classiﬁer will converge to the [Figure] following max-margin classiﬁer:
[Equation]
which is still independent of q i .
C M ORE E XPERIMENTS
We run ERM, importance weighting and Group DRO on the training set with 6 MNIST images which we used in Section 4.1 with the logistic loss and the polynomially-tailed loss (Eqn. (19), with α = 1 , β = 0 and (cid:96) left being the logistic loss shifted to make the overall loss function continuous) on this dataset for 10 million epochs (note that we run for much more epochs because the convergence is very slow).
The results are shown in Figure 3.
From the plots we can see that:
• For either loss function, the training loss of each method converges to 0.
• In contrast to the theory that the norm of the ERM model will go to inﬁnity and all models will converge to the max-margin classiﬁer, the weight of the ERM model gets stuck at some point, and the norms of the gaps between the normalized model weights also get stuck.
The reason is that the training loss has got so small that it becomes zero in the ﬂoating number representation, so the gradient also becomes zero and the training halts due to limited computational precision.
However, we can still observe a fundamental difference between the logistic loss and the polynomially-tailed loss.
For the logistic loss, the norm of the gap between importance weighting (or Group DRO) and ERM will converge to around 0.06 when the training stops, while for the polynomially-tailed loss, the norm will be larger than 0.22 and will keep growing, which shows that for the polynomially-tailed loss the normalized model weights do not converge to the same point.
• For either loss, the group weights of Group DRO still empirically satisfy Assumption
D P ROOFS
In this paper, for any matrix A , we will use (cid:107) A (cid:107) 2 to denote its spectral norm and (cid:107) A (cid:107) F to denote its Frobenius norm.
D. B ACKGROUND ON S MOOTHNESS
A ﬁrst-order differentiable function f over D is called L -smooth for L > 0 if
[Equation]
which is also called the upper quadratic bound .
If f is second-order differentiable and D is a convex set, then f is L -smooth is equivalent to
[Equation]
A classical result in convex optimization is the following:
Theorem 10.
If f ( x ) is convex and L -smooth with a unique ﬁnite minimizer x ∗ , and is minimized by gradient descent x t +
= x
t − η ∇ f ( x t ) starting from x 0 where the learning rate η ≤ 1 L , then we have
[Equation]
23) which also implies that x T converges to x ∗ as T → ∞ .
D. P ROOFS FOR S UBSECTION 4. D.2. P ROOF OF T HEOREM
Using the key intuition, the weight update rule (7) implies that θ ( t +1) − θ ( t ) ∈ span { x 1 , · · · , x n } for all t , which further implies that θ ( t ) − θ (0) ∈ span { x 1 , · · · , x n } for all t .
By Cramer’s rule, in this n -dimensional subspace there exists one and only one θ ∗ such that θ ∗ − θ (0) ∈ span { x 1 , · · · , x n } and (cid:104) θ ∗ , x i (cid:105) for all i .
Then we have
[Equation]
θ
[Equation] − Y [Equation] 2 →(24) because (cid:13)(cid:13) X (cid:62) θ − Y (cid:13)(cid:13) 2 2 = 2 n ˆ R ( f ( x ; θ )) .
On the other hand, let s min be the smallest singular value of X .
Since X is full-rank, s min > 0 , and (cid:13)(cid:13)
X (cid:62) ( θ ( t ) − θ ∗ ) (cid:13)(cid:13) 2 ≥ s min (cid:13)(cid:13) θ ( t ) − θ ∗ (cid:13)(cid:13) 2 .
This shows that (cid:13)(cid:13) θ ( t ) − θ ∗ (cid:13)(cid:13) 2 → 0 .
Thus, θ ( t ) converges to this unique θ ∗ .
D.2. P ROOF OF T HEOREM
To help our readers understand the proof more easily, we will ﬁrst prove the result for static GRW where q ( t ) i
=
q i for all t , and then we will prove the result for dynamic GRW that satisfy q ( t ) i
→ q i as t → ∞ .
Static GRW.
We ﬁrst prove the result for all static GRW such that min i q i = q ∗ > 0 .
We will use smoothness introduce in Appendix D.1.
Denote A = (cid:80) ni =1 (cid:107) x
i (cid:107) 22 .
The empirical risk of the linear model f ( x ) =
(cid:104) θ, x [Equation] is [Equation]
[Equation]
whose Hessian is
[Equation]
[Equation] So for any unit vector v ∈ R d , we have (since q i ∈ [0 , 1] )
[Equation]
which implies that F ( θ ) is 2 A -smooth.
Thus, we have the following upper quadratic bound: for any θ 1 , θ 2 ∈ R d ,
[Equation]
Denote g ( θ ( t ) ) =
( X (cid:62) θ ( t ) − Y ) ∈
R n .
We can see that (cid:13)(cid:13) √ Q g ( θ ( t ) )
(cid:13)(cid:13) 2 2 = F ( θ ( t ) ) , where where √
Q = diag ( √ q 1 , · · · , √ q n ) .
Thus, ∇ F ( θ ( t ) )
=
2 XQ g ( θ ( t ) ) .
The update rule of a static GRW with gradient descent and the squared loss is:
[Equation]
Substituting θ 1 and θ 2 in (28) with θ ( t ) and θ ( t +1) yields
[Equation]
Since x 1 , · · · , x n are linearly independent, X (cid:62) X is a positive deﬁnite matrix.
Denote the smallest eigenvalue of X (cid:62) X by λ min > 0 .
And (cid:13) (cid:13) Q g ( θ ( t ) )
(cid:13) (cid:13) 2 ≥ √ q ∗ (cid:13) (cid:13) g
( θ ( t ) )
{+(cid:13) (cid:13) 
+} =
(cid:112) q ∗ F ( θ ( t ) ) , so we have g ( θ ( t ) )
(cid:62) Q (cid:62) X (cid:62)
XQ g ( θ ( t ) ) ≥ q ∗
λ min F ( θ ( t ) ) .
Thus,
[Equation]
Let η 0 = q ∗ λ min
A 2 .
For any η ≤ η 0 , we have F ( θ ( t +1) ) ≤ (1 −
ηq ∗ λ min ) F ( θ ( t ) ) for all t , which implies that lim t →∞ F ( θ ( t ) )
= 0 .
This implies that the empirical training risk must converge to 0.
Dynamic GRW.
Now we prove the result for all dynamic GRW satisfying Assumption 1.
By
Assumption 1, for any (cid:15) > 0 , there exists t (cid:15) such that for all t ≥ t (cid:15) and all i ,
[Equation]
This is because for all i , there exists t i such that for all t ≥ t i , q ( t ) i ∈ ( q i − (cid:15), q i + (cid:15) ) .
Then, we can deﬁne t (cid:15) =
max { t 1 , · · · , t n } .
Denote the largest and smallest eigenvalues of X (cid:62) X by λ max and λ min , and because X is full-rank, we have λ min > 0 .
Deﬁne (cid:15) =
min { q ∗ 3 , ( q ∗ λ min ) 2 12 λ max2 } , and then t (cid:15) is also ﬁxed.
We still denote Q = diag ( q 1 , · · · , q n ) .
When t ≥ t (cid:15) , the update rule of a dynamic GRW with gradient descent and the squared loss is:
[Equation]
[Equation] where Q ( t ) (cid:15) =
Q ( t ) , and we use the subscript (cid:15) to indicate that (cid:13)(cid:13)(cid:13) Q ( t ) (cid:15) − Q (cid:13)(cid:13)(cid:13) 2 < (cid:15) .
Then, note that we can rewrite Q ( t ) (cid:15) as Q ( t ) (cid:15)
= (cid:113) Q ( t ) 3 (cid:15) ·√ Q as long as (cid:15) ≤ q ∗ / 3 .
This is because q i + (cid:15) ≤ (cid:112) ( q i + 3 (cid:15) )
q i and q i − (cid:15) ≥ (cid:112) ( q i − 3 (cid:15) )
q i for all (cid:15) ≤ q i / 3 , and q i ≥ q ∗ .
Thus, we have
[Equation]
= θ ( t ) − η X (cid:113) Q ( t ) 3 (cid:15) (cid:112) Q g ( θ ( t ) )
where Q ( t ) (cid:15) =
(cid:113) Q ( t ) 3 (cid:15) · (cid:112) Q (34)
Again, substituting θ 1 and θ 2 in (28) with θ ( t ) and θ ( t +1) yields
[Equation]
[Equation]
Then, note that
[Equation]
where the last step comes from the following fact: for all (cid:15) <
q i / 3 , [Equation]
[Equation]
and √ q i − (cid:112) [Equation] [Equation]
And as proved before, we also have
[Equation]
[Equation]
Since (cid:15) ≤ ( q ∗ λ min ) 2 12 λ max2 , we have
[Equation]
Thus,
[Equation]
≤
[Equation]
[Equation]
for all (cid:15) ≤ 1 / 3 .
Let η 0 = q ∗
λ min 4 A 2 .
For any η ≤ η 0 , we have F ( θ ( t +1) ) ≤ (1 −
ηq ∗ λ min / 2) F ( θ ( t ) ) for all t ≥ t (cid:15) , which implies that lim t →∞ F ( θ ( t ) )
= 0 .
Thus, the empirical training risk converges to 0.
D. P ROOFS FOR S UBSECTION 4. D.3. P ROOF OF L EMMA
Note that the ﬁrst l layers (except the output layer) of the original NTK formulation and our new formulation are the same, so we still have the following proposition:
Proposition 11 (Proposition 1 in Jacot et al.
If σ is Lipschitz and d l → ∞ for l = 1 , · · · , L sequentially, then for all l = 1 , · · · , L , the distribution of a single element of h l converges in probability to a zero-mean Gaussian process of covariance Σ l that is deﬁned recursively by:
[Equation]
41) where f is sampled from a zero-mean Gaussian process of covariance Σ ( l − 1) .
Now we show that for an inﬁnitely wide neural network with L ≥ 1 hidden layers, Θ (0) converges in probability to the following non-degenerated deterministic limiting kernel
[Equation]
Consider the output layer
h L +
=
W L √ ˜ d σ ( h L ) + β b L .
We can see that for any parameter θ i before the output layer,
[Equation]
And for W L and b L , we have
[Equation]
Then we can achieve (42) by the law of large numbers.
D.3. P ROOF OF L EMMA
We will use the following short-hand in the proof:
[Equation]
For any (cid:15) > 0 , there exists t (cid:15) such that for all t ≥ t (cid:15) and all i , q ( t ) i ∈ ( q i − (cid:15), q i + (cid:15) ) .
Like what we have done in (34), we can rewrite Q ( t ) =
Q ( t ) (cid:15) =
(cid:113) Q ( t ) 3 (cid:15) · √ Q , where Q = diag ( q 1 , · · · , q n ) .
The update rule of a GRW with gradient descent and the squared loss for the wide neural network is:
[Equation]
46) and for t ≥ t (cid:15) , it can be rewritten as
[Equation]
First, we will prove the following theorem:
Theorem 12.
There exist constants M > 0 and (cid:15) 0 > 0 such that for all (cid:15) ∈ (0 , (cid:15) 0 ] , η ≤ η ∗ and any δ > 0 , there exist R 0 > 0 , ˜ D > 0 and B > 1 such that for any ˜ d ≥ ˜ D , the following (i) and (ii) hold with probability at least 1 − δ ) over random initialization when applying gradient descent with learning rate η :
(i) For all t ≤ t (cid:15) , there is
[Equation]
(ii)
For all t ≥ t (cid:15) , we have
[Equation]
Proof.
The proof is based on the following lemma:
Lemma 13 (Local Lipschitzness of the Jacobian) .
Under Assumption 2, there is a constant M >such that for any C 0 > 0 and any δ > 0 , there exists a ˜ D such that: If ˜ d ≥ ˜ D , then with probability at least (1 − δ ) over random initialization, for any x such that (cid:107) x (cid:107) 2 ≤ 1 ,
[Equation]
where B ( θ (0) , R ) = { θ : (cid:13)(cid:13) θ − θ (0) (cid:13)(cid:13) 2 < R } .
The proof of this lemma can be found in Appendix D.3.3.
Note that for any x , f (0) ( x ) =
β b L where b L is sampled from the standard Gaussian distribution.
Thus, for any δ > 0 , there exists a constant R 0 such that with probability at least (1 − δ/ 3) over random initialization,
[Equation] [Equation]
And by Proposition 3, there exists D 2 ≥ 0 such that for any ˜ d ≥ D 2 , with probability at least (1 − δ/ 3) ,
[Equation]
Let M be the constant in Lemma 13.
Let (cid:15)
= ( q ∗ λ min ) 2 108 M 4 .
Let B = 1 + η ∗ M 2 , and C 0 =
MB t(cid:15) R 0 B −+ 3 √ 1+3 (cid:15)MB t(cid:15)
[Equation]
[Equation]
λ min .
By Lemma 13, there exists D 1 > 0 such that with probability at least (1 − δ/ 3) , for any ˜ d ≥ D 1 , (52) is true for all θ, ˜ θ ∈ B ( θ (0) , C 0 ) .
By union bound, with probability at least (1 − δ ) , (52), (53) and (54) are all true.
Now we assume that all of them are true, and prove (48) and (49) by induction.
(48) is true for t =
due to (53), and (49) is always true for t = 0 .
Suppose (48) and (49) are true for t , then for t +
we have
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Equation]
So (49) is also true for t
[Table]
And we also have
[Equation]
t )
[Equation]
(cid:13)(cid:13)(cid:13)
[Equation]
(cid:13)(cid:13)(cid:13)
[Equation]
[Equation] ( t ) ) (cid:13)(cid:13)(cid:13)
[Equation]
(cid:17)
[Equation] [Equation] ( t )
[Equation]
(cid:13)(cid:13)(cid:13)
[Equation]
Therefore, (48) and (49) are true for all t ≤ t (cid:15) , which implies that (cid:13)(cid:13) √ Q g ( θ ( t (cid:15) ) )
(cid:13)(cid:13) 2 ≤ (cid:13)(cid:13) g
( θ ( t (cid:15) ) )
(cid:13)(cid:13) 2 ≤
B t (cid:15) R 0 , so (50) is true for t = t (cid:15) .
And (51) is obviously true for t =
t (cid:15) .
Now, let us prove (ii) by induction.
Note that when t ≥ t (cid:15) , we have the alternative update rule (47).
If (50) and (51) are true for t , then for t + 1 , there is
[Table]
[Equation]
where ˜ θ ( t ) is some linear interpolation between θ ( t ) and θ ( t +1) .
Now we prove that
[Equation]
θ
[Equation]
[Equation]
θ
[Equation]
For any unit vector v ∈ R n , we have
[Equation]
(cid:13)(cid:13) √
Qv (cid:13)(cid:13) 2 ∈ [ √ q ∗ , 1] , so for any η ≤ η ∗ , v (cid:62) ( I − η √ Q Θ √ Q ) v ∈ [0 , 1 − ηλ min q ∗ ] , which implies that (cid:13)(cid:13)
I − η √ Q Θ √ Q (cid:13)(cid:13) 2 ≤ 1 − ηλ min q ∗ .
Thus,
[Equation]
≤ (cid:13)(cid:13)(cid:13) I − η (cid:112) Q Θ (cid:112) Q (cid:13)(cid:13)(cid:13) 2 + η (cid:13)(cid:13)(cid:13)(cid:112) Q (Θ − Θ (0) ) (cid:112) Q (cid:13)(cid:13)(cid:13) 2 + η (cid:13)(cid:13)(cid:13)(cid:112) Q ( J ( θ (0) ) (cid:62) J ( θ (0) )
− J (˜ θ ( t ) ) (cid:62) J ( θ ( t ) )) (cid:112) Q (cid:13)(cid:13)(cid:13)≤ 1 − ηλ min q ∗ + η (cid:13)(cid:13)(cid:13)(cid:112) Q (Θ − Θ (0) ) (cid:112) Q (cid:13)(cid:13)(cid:13) F + η (cid:13)(cid:13)(cid:13)(cid:112) Q ( J ( θ (0) )
(cid:62) J ( θ (0) )
− J (˜ θ ( t ) ) (cid:62) J ( θ ( t ) )) (cid:112) Q (cid:13)(cid:13)(cid:13) F ≤ 1 − ηλ min q ∗ + η (cid:13)(cid:13) (cid:13) Θ − Θ (0) (cid:13)(cid:13) (cid:13) F + η (cid:13)(cid:13) (cid:13) J ( θ (0) ) (cid:62) J
( θ (0) ) − J (˜ θ ( t ) ) (cid:62) J
( θ ( t ) )
{+(cid:13)(cid:13) (cid:13) F ≤ 1 − ηλ min q ∗ + ηq ∗ λ min + ηM4 (cid:112) ˜ d (cid:16)(cid:13)(cid:13)(cid:13) θ ( t ) − θ (0) (cid:13)(cid:13)(cid:13) 2 + (cid:13)(cid:13)(cid:13) ˜ θ ( t ) − θ (0) (cid:13)(cid:13)(cid:13) 2 (cid:17) ≤ 1 − ηq ∗ λ min(61) for all ˜ d ≥ max (cid:26) D 1 , D 2 , (cid:16) 12 M 2 C 
+}
q ∗ λ min (cid:17) 4 (cid:27) , which implies that (cid:13)(cid:13)(cid:13)(cid:13) I − η (cid:112) Q J (˜ θ ( t ) ) (cid:62)
[Equation]
(due to (37)) (62) for all (cid:15) ≤ (cid:15) 0 .
Thus, (50) is also true for t + 1 .
In conclusion, (50) and (51) are true with probability at least (1 − δ ) for all ˜ d ≥ ˜ D
[Equation] max (cid:26)[-D 1 , D 2 , (cid:16) 12 M 2 C 
-]
[Equation]
Returning back to the proof of Lemma 5.
Choose and ﬁx an (cid:15) such that (cid:15) < min { [Equation] [Equation] (cid:17) 2 } , where (cid:15) 0 is deﬁned by Theorem 12.
Then, t (cid:15) is also ﬁxed.
There exists ˜ D ≥ 0 such that for any ˜ d ≥ ˜ D , with probability at least (1 − δ ) , Theorem 12 and Lemmaare true and
[Equation]
63) which immediately implies that
[Equation]
We still denote B = 1 + η ∗ M 2 and C 0 =
MB t(cid:15) R 0 B [Equation] [Equation] √
[Equation]
Theorem 12 ensures that for all t , θ ( t )
∈ B ( θ (0) , C 0 ) .
Then we have
[Equation]
so it follows that
[Equation]
Thus, for all (cid:15) < 13 (cid:16) q ∗ λ min 3 λ max + q ∗
{+λ min (cid:17) 
+}
, there is
[Equation]
The update rule of the GRW for the linearized neural network is:
[Equation]
[Equation]
where we use the subscript “lin” to denote the linearized neural network, and with a slight abuse of notion denote g lin ( θ ( t ) )
=
g ( θ ( t ) lin ) .
First, let us consider the training data X .
Denote ∆ t = g lin ( θ ( t ) )
− g ( θ ( t ) ) .
We have
[Equation]
where ˜ θ ( t ) is some linear interpolation between θ ( t ) and θ ( t +1) .
Thus,
[Equation]
By Lemma 13, we have
[Equation]
[Equation] [Table] [Equation] [Table]
[Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
[Equation] [Table] [Equation] [Table]
[Equation] [Table] [Equation] [Table]
[Equation]
[Table]
[Equation]
Therefore
[Table]
we have
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
Since ∆ 0 =
[Equation]
it follows that for all t ≤ t (cid:15) ,
[Equation]
and particularly we have
[Equation]
For t ≥ t (cid:15) , we have the alternative update rule (47).
Thus,
[Equation]
Let A =
I − η √ Q J ( θ (0) ) (cid:62) J ( θ (0) )
(cid:113) Q ( t ) 3 (cid:15)
=
I − η √ Q Θ (0) (cid:113) Q ( t ) 3 (cid:15) .
Then, we have
[Equation]
=
[Equation]
[Equation] Let γ = 1 − ηq ∗ λ min< 1 .
Combining with Theorem 12 and (67), the above leads to [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]
[Equation]
This implies that
[Equation]
Combining with (75), it implies that for all t ≥ t (cid:15) ,
[Equation]
Next [Equation] we consider an arbitrary test point x such that (cid:107) x (cid:107) 2 ≤ 1 .
Denote δ t = f ( t ) lin ( x ) − f ( t ) ( x ) .
Then we have
[Equation]
which yields
[Equation]
For t ≤ t (cid:15) , we have
[Table]
[Equation]
So we can see that there exists a constant C 1 such that (cid:107) δ t
[Equation]
2 ≤ C 1 ˜ d − 1 / 4 .
Then, for t > t (cid:15)
[Table]
we have
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)(cid:104)
[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]
[Equation]
Note that (cid:80) ∞ t =
tγ t is ﬁnite as long as γ ∈ (0 , 1) .
Therefore, there is a constant C such that for any t , (cid:107) δ t (cid:107) 2 ≤ C ˜ d − 1 / 4 with probability at least (1 − δ ) for any ˜ d ≥ ˜ D .
D.3. P ROOF OF L EMMA
We will use the following theorem regarding the eigenvalues of random Gaussian matrices:
Theorem 14 (Corollary 5.35 in Vershynin (2010)) .
If A ∈ R p × q is a random matrix whose entries are independent standard normal random variables, then for every t ≥ 0 , with probability at least 1 − 2 exp( − t 2 / 2) ,
[Equation]
By this theorem, and also note that W L is a vector, we can see that for any δ , there exist ˜ D > 0 and
M 1 > 0 such that if ˜ d ≥ ˜ D , then with probability at least (1 − δ ) , for all θ ∈ B ( θ (0) , C 0 ) , we have (cid:13)(cid:13) W l (cid:13)(cid:13) 2 ≤ 3 (cid:112) ˜
d ( ∀ 0 ≤ l ≤ L − 1) and (cid:13)(cid:13)
[Equation]
as well as
[Equation]
Now we assume that (86) and (87) are true.
Then, for any x such that (cid:107) x (cid:107) 2 ≤ 1 ,
[Equation]
where L 0 is the Lipschitz constant of σ and σ ( 0 l )
=
( σ (0) , · · · , σ (0))
∈ R d l .
By induction, there exists an M 2 > 0 such that (cid:13)(cid:13) x l (cid:13)(cid:13) 2 ≤ M 2 (cid:112)
˜
d and (cid:13)(cid:13)
h l (cid:13)(cid:13) 2 ≤ M 2 (cid:112) ˜
d for all l = 1 , · · · , L .
Denote α l = ∇ h l f ( x ) = ∇ h l h L +1 .
For all l = 1 , · · · , L , we have α l = diag ( ˙ σ ( h l ))
W l (cid:62) √
˜
d α l + where ˙ σ ( x ) ≤ L 0 for all x ∈ R since σ is L 0 -Lipschitz, α L +
= 1 and (cid:13)(cid:13)
α L (cid:13)(cid:13) 2 = (cid:13)(cid:13)(cid:13)(cid:13) diag ( ˙ σ ( h L ))
W L (cid:62) √ ˜ d (cid:13)(cid:13)(cid:13)(cid:13) ≤ 3 4 √ ˜ d L 0 .
Then, we can easily prove by induction that there exists an
M 3 > 1 such that (cid:13)(cid:13)
α l (cid:13)(cid:13) 2 ≤ M 3 / 4 (cid:112) ˜
d for all l = 1 , · · · , L (note that this is not true for L +because α L +1 = 1 ).
For l = 0 , ∇ W 0 f ( x ) = 1 √ d 0 x 0 α 1 (cid:62) , so (cid:107)∇ W l f ( x ) (cid:107) 2 ≤ 1 √ d 0 (cid:13)(cid:13) x 0 (cid:13)(cid:13) 2 (cid:13)(cid:13)
α 1 (cid:13)(cid:13) 2 ≤ 1 √ d 0 M 3 / 4 (cid:112) ˜
d .
And for any l = 1 , · · · , L , ∇ W l f ( x ) = [-1 √ ˜ d x l α l +1 , so (cid:107)∇ W l f ( x ) (cid:107) 2 ≤ 
-][Equation]
√ ˜ [Equation] x l [-(cid:13)(cid:13) 
-]α l +1 , so (cid:107)∇
[Equation]
≤ M 2 M 3 . (Note that if M 3 > 1 , then (cid:13)(cid:13) α L +1 (cid:13)(cid:13) 2 ≤ M 3 ; and since ˜ d ≥ 1 , there is (cid:13)(cid:13)
α l (cid:13)(cid:13) 2 ≤ M 3 for l ≤ L .)
Moreover, for l = 0 , · · · , L , ∇ b l f ( x ) = β
α l +1 , so (cid:107)∇ b l f ( x ) (cid:107) 2 ≤ βM 3 .
Thus, if (86) and (87) are true, then there exists an M 4 > 0 , such that (cid:107)∇ θ f ( x ) (cid:107) 2 ≤ M 4 / √ n .
And since (cid:107) x i (cid:107) 2 ≤ 1 for all i , so (cid:107) J ( θ ) (cid:107) F ≤ M 4 .
Next, we consider the difference in ∇ θ f ( x ) between θ and ˜ θ .
Let ˜ f, ˜ W, ˜ b , ˜ x , ˜ h , ˜
α be the function and the values corresponding to ˜ θ .
There is
[Equation]
[Table]
[Equation]
( ∀ l ≥ 1) (89)
By induction, there exists an M 5 >
such that (cid:13)(cid:13) x
l − ˜ x l (cid:13)(cid:13) 2 ≤ M 5 (cid:13)(cid:13)(cid:13) θ − ˜ θ (cid:13)(cid:13)(cid:13) 2 for all l .
[Table]
[Equation]
where L 1 is the Lipschitz constant of ˙ σ .
Particularly, for l = L , though ˜ α L +1 = 1 , since (cid:13)(cid:13)(cid:13) ˜ W L (cid:13)(cid:13)(cid:13) 2 ≤ 3 ˜ d 1 / 4 , (90) is still true.
By induction, there exists an M 6 > 0 such that (cid:13)(cid:13)
[Equation]
(cid:13)(cid:13)
[Equation]
− ˜ θ
[Equation]
for all l ≥ 1 (note that this is also true for l = L + 1 ).
Thus, if (86) and (87) are true, then for all θ, ˜ θ ∈ B ( θ (0) , C 0 ) , any x such that (cid:107) x (cid:107) 2 ≤ 1 , we have
[Equation]
91) and for l = 1 , · · · , L , we have
[Equation]
Moreover, for any l = 0 , · · · , L , there is
[Equation]
[Equation]
Overall, we can see that there exists a constant M 7 > 0 such that (cid:13)(cid:13)(cid:13) ∇ θ f ( x ) − ∇ ˜ θ ˜ f ( x ) (cid:13)(cid:13)(cid:13) 2 ≤ M [Table]
n
[Table]
4 √
[Table]
d
[Table]
θ − ˜ θ (cid:13)(cid:13)(cid:13) 2 , so that (cid:13)(cid:13)(cid:13)
J ( θ ) − J (˜ θ )
[Table]
F ≤ M 7 4 √ ˜
d
[Table]
θ −
[Table]
D.3. P ROOF OF T HEOREM
First of all, for a linearized neural network (11), if we view {∇ θ f (0) ( x i ) } ni =1 as the inputs and { y i − f (0) ( x i ) + (cid:104) θ (0) , ∇ θ f (0) ( x i ) (cid:105)} ni =1 as the targets, then the model becomes a linear model.
So by Theorem 2 we have the following corollary:
Corollary 15.
If ∇ θ f (0) ( x 1 ) , · · · , ∇ θ f (0) ( x n ) are linearly independent, then there exists η 0 > 0 such that for any GRW satisfying Assumption 1, and any η ≤ η 0 , θ ( t ) converges to the same interpolator θ ∗ that does not depend on q i .
Let η 1 = min { η 0 , η ∗ } , where η 0 is deﬁned in Corollary 15 and η ∗ is deﬁned in Lemma 5.
Let f ( t ) lin ( x ) and f ( t ) linERM ( x ) be the linearized neural networks of f ( t ) ( x ) and f ( t ) ERM ( x ) , respectively.
By
Lemma 5, for any δ > 0 , there exists ˜ D > 0 and a constant C such that
[Equation]
By Corollary 15, we have
[Equation]
Summing the above yields
[Equation]
96) which is the result we want.
D. P ROOFS FOR S UBSECTION 4. D.4. A N EW A PPROXIMATION T HEOREM
Lemma 16 (Approximation Theorem for Regularized GRW) .
For a wide fully-connected neural network f , denote J ( θ )
= ∇ θ f
( X ; θ ) ∈ R p × n and g ( θ ) =
∇ ˆ y (cid:96) ( f ( X ; θ ) , Y ) ∈ R n .
Given that the loss function (cid:96) satisﬁes: ∇ θ g ( θ ) =
J ( θ ) U ( θ ) for any θ , and U ( θ ) is a positive semi-deﬁnite diagonal matrix whose elements are uniformly bounded, we have: for any GRW that minimizes the regularized weighted empirical risk (13) with a sufﬁciently small learning rate η , there is: for a sufﬁciently large ˜ d , with high probability over random initialization, on any test point x such that (cid:107) x (cid:107) 2 ≤ 1 ,
[Equation]
(97) where both f ( t ) linreg and f ( t ) reg are trained by the same regularized GRW and start from the same initial point.
First of all, with some simple linear algebra analysis, we can prove the following proposition:
Proposition 17.
For any positive deﬁnite symmetric matrix H ∈ R n × n , denote its largest and smallest eigenvalues by λ max and λ min .
Then, for any positive semi-deﬁnite diagonal matrix Q = diag ( q 1 , · · · , q n ) , HQ has n eigenvalues that all lie in [min i q i · λ min , max i q i · λ max ] .
Proof.
H is a positive deﬁnite symmetric matrix, so there exists A ∈ R n × n such that H
=
A (cid:62) A , and A is full-rank.
First, any eigenvalue of AQA (cid:62) is also an eigenvalue of A (cid:62) AQ , because for any eigenvalue λ of AQA (cid:62) we have some v (cid:54) =
such that AQA (cid:62) v = λ v .
Multiplying both sides by A (cid:62) on the left yields A (cid:62) AQ ( A (cid:62) v )
=
λ ( A (cid:62) v ) which implies that λ is also an eigenvalue of
A (cid:62) AQ because A (cid:62) v (cid:54)
=
as λ v (cid:54) = 0 .
Second, by condition we know that the eigenvalues of A (cid:62) A are all in [ λ min , λ max ] where λ min > 0 , which implies for any unit vector v , v (cid:62) A (cid:62)
Av ∈ [ λ min , λ max ] , which is equivalent to (cid:107) Av (cid:107) 2 ∈ [ √ λ min , √ λ max ] .
Thus, we have v (cid:62) A (cid:62) QAv ∈
[ λ min min i q i , λ max max i q i ] , which implies that the eigenvalues of A (cid:62) QA are all in [ λ min min i q i , λ max max i q i ] .
Thus, the eigenvalues of HQ = A (cid:62) AQ are all in [ λ min min i q i , λ max max i q i ] .
Proof of Lemma
By the condition (cid:96) satisﬁes, without loss of generality, assume that the elements of U ( θ ) are in [0 , 1] for all θ .
Then, let η ≤ ( µ + λ min + λ max ) − 1 .
(If the elements of U ( θ ) are bounded by [0 , C ] , then we can let η ≤ ( µ + Cλ min + Cλ max ) − 1 and prove the result in the same way.)
With L 2 penalty, the update rule of the GRW for the neural network is:
[Equation]
And the update rule for the linearized neural network is:
[Equation]
By Proposition 11, f ( x ; θ ) converges in probability to a zero-mean Gaussian process.
Thus, for any δ > 0 , there exists a constant R 0 > 0 such that with probability at least (1 − δ/ 3) , (cid:13)(cid:13) g ( θ (0) )
(cid:13)(cid:13) 2 < R 0 .
Let M be as deﬁned in Lemma 13.
Denote A = ηMR 0 , and let C
= 4 Aηµ in Lemma 13 6 .
By Lemma 13, there exists D 1 such that for all ˜ d ≥ D 1 , with probability at least (1 − δ/ 3) , (52) is true.
Similar to the proof of Proposition 17, we can show that for arbitrary ˜ θ , all non-zero eigenvalues of J ( θ (0) )
Q ( t ) U
(˜ θ ) J ( θ (0) )
(cid:62) are eigenvalues of J ( θ (0) )
(cid:62) J ( θ (0) )
Q ( t ) U (˜ θ ) .
This is because for any λ (cid:54) =
0 , if J ( θ (0) )
Q ( t ) U
(˜ θ ) J ( θ (0) )
(cid:62) v = λ v , then J ( θ (0) ) (cid:62) J ( θ (0) )
Q ( t ) U (˜ θ )
( J ( θ (0) ) (cid:62) v )
=
λ ( J ( θ (0) ) (cid:62) v ) , and J ( θ (0) ) (cid:62) v (cid:54) = 0 since λ v (cid:54) = 0 , so λ is also an eigenvalue of
J ( θ (0) )
(cid:62) J ( θ (0) )
Q ( t ) U (˜ θ ) .
On the other hand, by Proposition 3, J ( θ (0) ) (cid:62) J
( θ (0) )
Q ( t ) U
(˜ θ ) converges in probability to Θ Q ( t ) U (˜ θ ) whose eigenvalues are all in [0 , λ max ] by Proposition 17.
So there exists D 2 such that for all ˜ d ≥ D 2 , with probability at least (1 − δ/ 3) , the eigenvalues of
J ( θ (0) ) Q ( t ) U (˜ θ ) J ( θ (0) )
(cid:62) are all in [0 , λ max + λ min ] for all t .
By union bound, with probability at least (1 − δ ) , all three above are true, which we will assume in the rest of this proof.
First, we need to prove that there exists D 0 such that for all ˜ d ≥ D 0 , sup t ≥ 0 (cid:13)(cid:13) θ ( t ) − θ (0) (cid:13)(cid:13) 2 is bounded with high probability.
Denote a t = θ ( t ) − θ (0) .
By (98) we have
[Equation]
which implies
[Equation]
[Equation]
where ˜ θ ( t ) is some linear interpolation between θ ( t ) and θ (0) .
Our choice of η ensures that ηµ < 1 .
Now we prove by induction that (cid:107) a t (cid:107) 2 < C 0 .
It is true for t = 0 , so we need to prove that if (cid:107) a t (cid:107) 2 < C 0 , then (cid:107) a t +1 (cid:107) 2 < C 0 .
For the ﬁrst term on the right-hand side of (101), we have
[Equation]
J ( θ (0) )
(cid:13)(cid:13)(cid:13) F (cid:13)(cid:13)(cid:13) J
(˜ θ ( t ) )
− J ( θ (0) ) (cid:13)(cid:13)(cid:13) F (102)
Since η/ (1 − ηµ ) ≤ ( λ min + λ max ) − 1 by our choice of η , we have (cid:13)(cid:13)(cid:13)(cid:13)
I −
[Equation]
(cid:62) (cid:13)(cid:13)(cid:13)(cid:13) ≤(103)
On the other hand, we can use (52) since (cid:107) a t (cid:107) 2 < C 0 , so (cid:13)(cid:13)
J ( θ (0) )
(cid:13)(cid:13) F (cid:13)(cid:13)(cid:13)
J
(˜ θ ( t ) )
− J ( θ (0) ) (cid:13)(cid:13)(cid:13) F ≤
M 2 4 √ ˜
d
C 0 .
Therefore, there exists D 3 such that for all ˜ d ≥ D 3 ,
[Equation]
For , we have
[Equation]
And for the third term, we have
[Equation]
Thus, we have
[Equation]
So there exists D 4 such that for all ˜ d ≥ D 4 , (cid:107) a t + [Equation] 2 ≤ [Equation]
1 − [Equation] [Equation] a
[Equation] 2 + 2 A . This shows that if (cid:107) a t (cid:107) 2 < C 0 is true, then (cid:107) a t +1 (cid:107) 2 < C 0 will also be true. In conclusion, for all ˜ d ≥ D 0 = max { D 1 , D 2 , D 3 , D 4 } , (cid:13)(cid:13) θ ( t ) − θ (0) (cid:13)(cid:13) 2 < C 0 is true for all t .
This also implies that for C 1 = MC 0 + R 0 , we have (cid:13)(cid:13) g
( θ ( t ) )
(cid:13)(cid:13) 2 ≤ C 1 for all t by (105).
Similarly, we can prove that (cid:107) θ ( t ) lin − θ (0) (cid:107) 2 < C 0 for all t .
Second, let ∆ t
= θ ( t ) lin − θ ( t ) .
Then we have
[Equation] [Equation]
which implies
[Equation]
where ˜ θ ( t ) is some linear interpolation between θ ( t ) and θ ( t ) lin .
By (104), with probability at least (1 − δ ) for all ˜ d ≥ D 0 , we have
[Equation]
Again, as ∆ 0 = 0 , we can prove by induction that for all t ,
[Equation]
For any test point x such that (cid:107) x (cid:107) 2 ≤ 1 , we have
[Equation]
f ( x ; θ (0) ) (cid:13)(cid:13)(cid:13) 2 (cid:13)(cid:13)(cid:13) θ ( t )
[Equation]
( t ) lin (cid:13)(cid:13)(cid:13)
[Equation]
For the ﬁrst term, note that
[Equation]
where ˜ θ ( t ) is some linear interpolation between θ ( t ) and θ (0) .
Since f ( x ; θ (0) ) = f lin ( x ; θ (0) ) ,
[Equation]
Thus, we have shown that for all ˜ d ≥ D 0 , with probability at least (1 − δ ) for all t and all x ,
[Equation]
which is the result we need.
D.4. R ESULT FOR L INEARIZED N EURAL N ETWORKS
Lemma 18.
Suppose there exists M 0 > 0 such that (cid:13)(cid:13) ∇ θ f (0) ( x ) (cid:13)(cid:13) 2 ≤ M 0 for all test point x .
If the gradients ∇ θ f (0) ( x 1 ) , · · · , ∇ θ f (0) ( x n ) are linearly independent, and the empirical training risk of f ( t ) linreg satisﬁes
[Equation]
for some (cid:15) > 0 , then for x such that (cid:107) x (cid:107) 2 ≤ 1 we have
[Equation]
First, we can see that under the new weight update rule, θ ( t ) − θ (0) ∈ span {∇ θ f (0) ( x 1 ) , · · · , ∇ θ f (0) ( x n ) } is still true for all t .
Let θ ∗ be the interpolator in span ( ∇ θ f (0) ( x 1 ) , · · · , ∇ θ f (0) ( x n )) , then the empirical risk of θ is 12 n (cid:80) ni
=1 (cid:104) θ − θ ∗ , ∇ θ f (0) ( x i ) (cid:105) [Equation] =
[Equation]
∇
[Equation]
Thus, there exists T > 0 such that for any t ≥ T ,
[Equation]
[Equation] Let the smallest singular value of 1 √ n ∇ θ f (0) ( X ) be s min , and we have s min > 0 .
Note that the column space of ∇ θ f (0) ( X ) is exactly span ( ∇ θ f (0) ( x 1 ) , · · · , ∇ θ f (0) ( x n )) .
Deﬁne H ∈ R p × n such that its columns form an orthonormal basis of this subspace, then there exists G ∈ R n × n such that ∇ θ f (0) ( X ) =
HG , and the smallest singular value of 1 √ n G is also s min .
Since θ ( t ) − θ (0) is also in this subspace, there exists v ∈ R n such that θ ( t ) − θ ∗ =
Hv .
Then we have √ 2 n(cid:15) ≥ (cid:13)(cid:13) G (cid:62) H (cid:62)
Hv (cid:13)(cid:13)
= (cid:13)(cid:13) G (cid:62) v (cid:13)(cid:13) 2 .
Thus, (cid:107) v (cid:107) 2 ≤ √ 2 (cid:15) s min , which implies
[Equation]
We have already proved in previous results that if we minimize the unregularized risk with ERM, then θ always converges to the interpolator θ ∗ .
So for any t ≥ T and any test point x such that (cid:107) x (cid:107) 2 ≤ 1 , we have
[Equation]
which implies (117).
D.4. P ROOF OF T HEOREM
Given that ˆ R ( f ( t ) linreg ) < (cid:15) for sufﬁciently large t , Lemma 16 implies that
[Equation]
So for a ﬁxed (cid:15) , there exists D > 0 such that for all d ≥ D , for sufﬁciently large t ,
[Equation]
By Lemma 5 and Lemma 16, we have
[Equation]
Combining Lemma 18 with (123) derives
[Equation]
Letting ˜
d
→ ∞ leads to the result we need.
Remark.
One might wonder whether (cid:107)∇ θ f (0) ( x ) (cid:107) 2 will diverge as ˜ d → ∞ .
In fact, in Lemma 13, we have proved that there exists a constant M such that with high probability, for any ˜
d there is (cid:107)∇ θ f (0) ( x ) (cid:107) 2 ≤ M for any x such that (cid:107) x (cid:107) 2 ≤ 1 .
Therefore, it is ﬁne to suppose that there exists such an M 0 .
D. P ROOFS FOR S UBSECTION 5. D.5. P ROOF OF T HEOREM
First we need to show that ˆ θ MM is unique.
Suppose both θ 1 and θ 2 maximize min i =1 , ··· ,n y i · (cid:104) θ, x i (cid:105) and θ 1 (cid:54)
= θ 2 , (cid:107) θ 1 (cid:107)
= (cid:107) θ 2 (cid:107) 2 = 1 .
Then consider θ 0 = θ/ (cid:107) θ (cid:107) 2 where θ = ( θ 1 + θ 2 ) / 2 .
Obviously, (cid:107) θ (cid:107) 2 < 1 , and for any i , y i · (cid:104) θ, x i (cid:105) =
( y i · (cid:104) θ 1 , x i (cid:105)
+ y i · (cid:104) θ 2 , x i (cid:105) ) / 2 , so y i · (cid:104) θ 0 , x i (cid:105) > min { y i · (cid:104) θ 1 , x i (cid:105) , y i · (cid:104) θ 2 , x i (cid:105)} , which implies that min i =1 , ··· ,n y i · (cid:104) θ 0 , x i (cid:105) > min { min i =1 , ··· ,n y i · (cid:104) θ 1 , x i (cid:105) , min i =1 , ··· ,n y i · (cid:104) θ 2 , x i (cid:105)} , contradiction!
Now we start proving the result.
Without loss of generality, let ( x 1 , y 1 ) , · · · , ( x m , y m ) be the samples with the smallest margin to u , i.e.
[Equation]
And denote y 1 · (cid:104) u , x 1 (cid:105) = · · · = y m · (cid:104) u , x m (cid:105) =
γ u .
Since the training error converges to 0, γ u > 0 .
Note that for the logistic loss, if y i · (cid:104) θ, x i (cid:105) <
y j · (cid:104) θ, x j (cid:105) , then for any M > 0 , there exists an R M > 0 such that for all R ≥ R M ,
[Equation]
which can be shown with some simple calculation.
And because the training error converges to 0, we must have (cid:13)(cid:13) θ ( t ) (cid:13)(cid:13) → ∞ .
Then, by Assumption 3 this means that when t gets sufﬁciently large, the impact of ( x j , y j ) to θ ( t ) where j > m is an inﬁnitesimal compared to ( x i , y i ) where i ≤ m (because there exists a positive constant δ such that q ( t ) i > δ for all sufﬁciently large t by
Assumption 3).
Thus, we must have u ∈ span { x 1 , · · · , x m } .
Let u = α 1 y 1 x 1 + · · · + α m y m x m .
Now we show that
α i ≥ 0 for all i = 1 , · · · , m .
This is because when t is sufﬁciently large such that the impact of ( x j , y j ) to θ ( t ) where j > m becomes inﬁnitesimal, we have
[Equation]
and since (cid:107) θ ( t ) (cid:107) → ∞ as t → ∞ , we have
[Equation]
where T 0 is sufﬁciently large.
Here the notion α i ∝ lim T →∞ α i ( T ) means that lim T →∞ α
i ( T ) α j ( T ) =
α i α j for any pair of i, j and α j (cid:54) = 0 .
Note that each term in the sum is non-negative.
This implies that all α 1 , · · · , α m have the same sign (or equal to 0).
On the other hand,
[Equation]
Thus, α i ≥ 0 for all i and at least one of them is positive.
Now suppose u (cid:54)
=
ˆ θ MM , which means that γ u is smaller than the margin of ˆ θ MM .
Then, for all i = 1 , · · · , m , there is y i ·(cid:104) u , x i (cid:105) <
y i ·(cid:104) ˆ
θ MM , x i (cid:105) .
This implies that
[Equation]
which is a contradiction.
Thus, we must have u = ˆ θ MM .
D.5. P ROOF OF T HEOREM
Denote the largest and smallest eigenvalues of X (cid:62) X by λ max and λ min , and by condition we have λ min > 0 .
Let (cid:15) =
min { q ∗ 3 , ( q ∗ λ min ) 2 192 λ max2 } .
Then similar to the proof in Appendix D.2.2, there exists t (cid:15) such that for all t ≥ t (cid:15) and all i , q ( t ) i ∈ ( q i − (cid:15), q i + (cid:15) ) .
Denote Q = diag ( q 1 , · · · , q n ) , then for all t ≥ t (cid:15) , Q ( t ) :
=
Q ( t ) (cid:15) = √ Q (cid:113) Q ( t ) 3 (cid:15) , where we use the subscript (cid:15) to indicate that (cid:13)(cid:13)(cid:13) Q ( t ) (cid:15) − Q (cid:13)(cid:13)(cid:13) 2 < (cid:15) .
First, we prove that F ( θ ) is L -smooth as long as (cid:107)
x i (cid:107) 2 ≤ 1 for all i .
The gradient of F is
[Equation]
Since (cid:96) (ˆ y, y ) is L -smooth in ˆ y , we have for any θ 1 , θ 2 and any i ,
[Equation]
Thus, we have
[Equation]
133) which implies that F ( θ ) is L -smooth.
Denote ˜ g ( θ ) =
∇ ˆ y (cid:96) ( f ( X ; θ ) , Y ) ∈ R n , then ∇ F ( θ ( t ) )
=
XQ ˜
g
( θ ( t ) ) , and the update rule is
[Equation]
So by the upper quadratic bound
[Equation]
we have
[Equation]
t ) )
(cid:13)(cid:13)(cid:13) 2(135) Let η 1 = q [Equation]
) λ max .
Similar to what we did in Appendix D.2.2 (Eqn.
, we can prove that for all η ≤ η 1 , (135) implies that for all t ≥ t (cid:15) , there is
[Equation]
[Table] [Equation] [Table] [Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table] [Equation] [Equation] This shows that F ( θ ( t ) ) is monotonically non-increasing.
Since F ( θ ) ≥ 0 , F ( θ ( t ) ) must converge as t → ∞ , and we need to prove that it converges to 0.
Suppose that F ( θ ( t ) ) does not converge to 0, then there exists a constant C > 0 such that F ( θ ( t ) ) ≥ 2 C for all t .
On the other hand, it is easy to see that there exists θ ∗ such that (cid:96) ( (cid:104) θ ∗ , x i (cid:105) , y i ) <
C for all i .
(136) also implies that (cid:13)(cid:13)
˜ g
( θ ( t ) )
(cid:13)(cid:13) 2 →as t → ∞ because we must have F ( θ ( t ) )
− F ( θ ( t +1) )
→ 0 .
Note that from (134) we have
[Equation] [Equation]
Denote
[Equation]
Then F t is convex because (cid:96) is convex and q ( t ) i are non-negative, and ∇ F t ( θ ( t [Equation]
=
XQ ( [Equation]
( θ ( t ) ) .
By the lower linear bound F t ( y ) ≥ F t ( x ) + (cid:104)∇ F t ( x ) , y − x (cid:105) , we have for all t , (cid:104) XQ ( t )
˜ g ( θ ( t ) ) , θ ∗ − θ ( t ) (cid:105) ≤ F t ( θ ∗ ) − F t ( θ ( t ) ) ≤ F t ( θ ∗ ) − 2 3 F ( θ ( t ) )
≤ C − 4 C 3 = − C 3 (139) because q ( t ) i ≥ q[-i − (cid:15) ≥ 
 q-]
i − (cid:15) ≥ 23 q i and (cid:80) ni
=1 q ( t ) i = 1 .
Since (cid:13)(cid:13) ˜ g
( θ ( t ) )
(cid:13)(cid:13) 2 → 0 , there exists T > 0 such that for all t ≥ T and all η ≤ η 0 ,
[Equation]
(140) which means that (cid:13)(cid:13) θ ( t ) − θ ∗ (cid:13)(cid:13) 2 2 → −∞ because ηC 3 is a positive constant.
This is a contradiction!
Thus, F ( θ ( t ) ) must converge to 0, which is result (i).
(i) immediately implies (ii) because (cid:96) is strictly decreasing to 0 by condition.
Now let’s prove (iii).
First of all, the uniqueness of θ R can be easily proved from the convexity of
F ( θ ) .
The condition implies that y i (cid:104) θ R , x i (cid:105) > 0 , i.e. θ R must classify all training samples correctly.
If there are two different minimizers θ R and θ (cid:48) R in whose norm is at most R , then consider θ (cid:48)(cid:48) R = 12 ( θ R + θ (cid:48) R ) .
By the convexity of F , we know that θ (cid:48)(cid:48)
R must also be a minimizer, and (cid:107) θ (cid:48)(cid:48) R (cid:107) 2 < R .
Thus, F ( R (cid:107) θ (cid:48)(cid:48) R (cid:107) 2 θ (cid:48)(cid:48) R ) <
F ( θ (cid:48)(cid:48) R ) and (cid:107) R (cid:107) θ (cid:48)(cid:48)
R (cid:107) 2 θ (cid:48)(cid:48) R (cid:107)
= R , which contradicts with the fact that θ (cid:48)(cid:48) R is a minimizer.
To prove the rest of (iii), the key is to consider (135).
On one hand, similar to (36) we can prove that for all t ≥ t (cid:15) , there is
[Equation]
Since we choose (cid:15) =
min { q ∗ 3 , ( q ∗ λ min ) 2 192 λ max2 } , this inequality implies that
[Equation]
[Table]
On the other hand, if η ≤ η 2 = 12 L , we will have
[Equation]
Combining all the above with (135) yields
[Equation]
Denote u = lim
R →∞ θ R (cid:107) θ R (cid:107) 2 .
Similar to Lemma 9 in Ji et al.
, we can prove that: for any α > 0 , there exists a constant ρ ( α ) > 0 such that for any θ subject to (cid:107) θ (cid:107) 2 ≥
ρ ( α ) , there is
[Equation]
for any t .
Let t α ≥ t (cid:15) satisfy that for all t ≥ t α , (cid:107) θ ( t )
[Equation]
2 ≥ max { ρ
[Equation]
α ) , 1 } .
By the convexity of F t , for all t ≥ t α ,
[Equation]
[Equation]
Thus, we have
[Equation]
=
[Equation]
[Equation]
θ
[Equation]
t
[Equation]
(cid:13)(cid:13)(cid:13)
[Equation]
−
[Equation]
θ ( t +1) − θ ( t ) (cid:13)(cid:13)(cid:13) 2 2 (cid:19) (1 + α ) (cid:107) θ ( t ) (cid:107)
[Equation]
By 12 ( (cid:107) θ ( t +1) (cid:107) 2 −(cid:107) θ ( t ) (cid:107) 2 ) 2 ≥ 0 , we have ( 12 (cid:13)(cid:13)
{+θ ( t +1) (cid:13)(cid:13) 2 
+}
− (cid:13)(cid:13)
θ ( t ) (cid:13)(cid:13) 2 2 ) / (cid:107) θ ( t ) (cid:107) 2 ≥ (cid:13)(cid:13) θ ( t +1) (cid:13)(cid:13) 2 − (cid:13)(cid:13) θ ( t ) (cid:13)(cid:13) 2 .
Moreover, by (144) we have (cid:13)(cid:13) θ ( t +1) − θ ( t ) (cid:13)(cid:13) 22(
[Equation]
Summing up (147) from t =
t α to t − 1 , we have (cid:104) [Equation] t
[Equation]
t
[Equation]
) , u
[Equation]
which implies that
[Equation]
Since lim t →∞ (cid:107) θ ( t ) (cid:107) 2 = ∞ , we have lim inf t →∞
[Equation]
[Equation] Since α is arbitrary, we must have lim t →∞
[Equation]
u as long as η ≤ min { η 1 , η 2 } .
D.5. C OROLLARY OF T HEOREM
We can show that for the logistic loss, it satisﬁes all conditions of Theorem 8 and lim R →∞ θ R R
=
ˆ θ MM .
First of all, for the logistic loss we have ∇ 2 ˆ y (cid:96) (ˆ y, y ) =
y 2 e y ˆ y + e − y
ˆ y +2 ≤ max i
{+y 
+}
i 4 , so (cid:96) is smooth.
Then, we prove that lim R →∞ θ R R exists and is equal to ˆ θ MM .
For the logistic loss, it is easy to show that for any ˆ θ (cid:48) (cid:54)
=
ˆ θ MM , there exists an R
(ˆ θ (cid:48) ) > 0 and an δ (ˆ θ (cid:48) )
> 0 such that F ( R · θ ) >
F ( R · ˆ θ MM ) for all R ≥ R (ˆ θ (cid:48) ) and θ ∈ B (ˆ θ (cid:48) , δ (ˆ θ (cid:48) )) .
Let S
=
{ θ : (cid:107) θ (cid:107) 2 = 1 } .
For any (cid:15) > 0 , S − B (ˆ θ MM , (cid:15) ) is a compact set.
And for any θ ∈
S − B (ˆ θ MM , (cid:15) ) , there exist R ( θ ) and δ ( θ ) as deﬁned above.
Thus, there must exist θ 1 , · · · , θ m ∈
S − B (ˆ θ MM , (cid:15) ) such that S − B
(ˆ θ MM , (cid:15) ) ⊆ ∪ mi =1 B
( θ i , δ ( θ i )) .
Let R ( (cid:15) )
=
max { R ( θ 1 ) , · · · , R ( θ m ) } , then for all R ≥ R ( (cid:15) ) and all θ ∈ S − B (ˆ θ MM , (cid:15) ) , F ( R · θ ) >
F ( R · ˆ θ MM ) , which means that θ
R R ∈ B (ˆ θ MM , (cid:15) ) for all R ≥ R ( (cid:15) ) .
Therefore, lim R →∞ θ R R exists and is equal to ˆ θ MM .
Therefore, by Theorem 8, any GRW satisfying Assumption 1 makes a linear model converge to the max-margin classiﬁer under the logistic loss.
D.
P ROOF OF T HEOREM
We ﬁrst consider the regularized linearized neural network f ( t ) linreg .
Since by Proposition 11 f (0) ( x ) is sampled from a zero-mean Gaussian process, there exists a constant M > 0 such that | f (0) ( x i ) | <
M for all i with high probability.
Deﬁne
[Equation]
Denote ˜ θ R = arg min θ { F ( R · θ ) : (cid:107) θ (cid:107) 2 ≤ 1 } .
when the linearized neural network is trained by a GRW satisfying Assumption 1 with regularization, since this is convex optimization and the objective function is smooth, we can prove that with a sufﬁciently small learning rate, as t → ∞ , θ ( t ) → R · ˜ θ R + θ (0)
where
R = lim t →∞ (cid:107) θ {+( t ) − θ (0) (cid:107) 
+}
(which is the minimizer).
And deﬁne
[Equation]
First, we derive the lower bound of R .
By Theorem 16, with a sufﬁciently large ˜ d , with high probability ˆ R ( f ( t ) reg ) < (cid:15) implies
ˆ R ( f ( t ) linreg ) < 2 (cid:15) .
By the convexity of (cid:96) , we have 2 (cid:15) > 1 n n (cid:88) i =
[Equation]
≥ log (cid:32) [Equation] 154) which implies that R = Ω( − log 2 (cid:15) ) for all (cid:15) ∈ (0 , 14 ) .
Denote δ = (cid:107) ˆ θ MM − ˜ θ R (cid:107) 2 .
Let θ (cid:48)
=
ˆ θ MM + ˜ θ R, then we can see that (cid:107) θ (cid:48) (cid:107) 2 = (cid:113) 1 − δ 2 4 .
Let ˜ θ (cid:48)
= θ (cid:48) (cid:107) θ (cid:48) (cid:107) 2 .
By the deﬁnition of ˆ θ MM , there exists j such that y j · (cid:104) ˜ θ (cid:48) , ∇ θ f (0) ( x j ) (cid:105) ≤ γ , which implies
[Equation]
Thus, we have
[Equation]
( since √ 1 − x ≤ 1 − x 2 )
[Equation]
On the other hand, we have
[Equation]
[Equation]
which implies that
[Equation]
and this leads to
[Equation]
which is equivalent to
[Equation]
Thus, we have
[Equation]
So for any test point x [Equation] since (cid:107)∇ θ f (0) ( x ) (cid:107) 2 ≤ M 0 , we have
[Equation]
Combined with Theorem [Equation] we have: with high probability,
[Equation]
So there exists a constant C > 0 such that: As ˜ d → ∞ , with high probability, for all (cid:15) ∈ (0 , 14 ) , if | f MM ( x ) | > C · ( − log 2 (cid:15) )
− 1 / 2 , then f ( t ) reg ( x ) will have the same sign as f MM ( x ) for a sufﬁciently large t .
Note that this C only depends on n , q ∗ , γ , M and M 0 , so it is a constant independent of (cid:15) .
Remark.
Note that Theorem 9 requires Assumption 1 while Theorem 6 does not due to the fundamental difference between the classiﬁcation and regression.
In regression the model converges to a ﬁnite point.
However, in classiﬁcation, the training loss converging to zero implies that either (i)
The direction of the weight is close to the max-margin classiﬁer or (ii)
The norm of the weight is very large.
Assumption 1 is used to eliminate the possibility of (ii).
If the regularization parameter µ is sufﬁciently large, then a small empirical risk could imply a small weight norm.
However, in our theorem we do not assume anything on µ , so Assumption 1 is necessary.
E A N OTE ON THE P ROOFS IN L EE ET AL . (2019)
We have mentioned that the proofs in Lee et al.
(2019), particularly the proofs of their Theorem 2.
and Lemma 1 in their Appendix G, are ﬂawed.
In order to ﬁx their proof, we change the network initialization to (9).
In this section, we will demonstrate what goes wrong in the proofs in Lee et al.
(2019), and how we manage to ﬁx the proof.
For clarity, we are referring to the following version of the paper: https://arxiv.org/pdf/1902.06720v4.pdf .
To avoid confusion, in this section we will still use the notations used in our paper.
E. T HEIR P ROBLEMS
Lee et al.
(2019) claimed in their Theorem 2.1 that under the conditions of our Lemma 5, for any δ > 0 , there exist ˜ D > 0 and a constant C such that for any ˜ d ≥ ˜ D , with probability at least (1 − δ ) , the gap between the output of a sufﬁciently wide fully-connected neural network and the output of its linearized neural network at any test point x can be uniformly bounded by
[Equation]
where they used the original NTK formulation and initialization in Jacot et al.
[Equation]
and
[Equation]
where x 0 = x and f ( x ) =
h L +1 .
However, in their proof in their Appendix G, they did not directly prove their result for the NTK formulation, but instead they proved another result for the following formulation which they called the standard formulation : [Equation] and
[Equation]
See their Appendix F for the deﬁnition of their standard formulation.
In the original formulation, they also included two constants σ w and σ b for standard deviations, and for simplicity we omit these constants here.
Note that the outputs of the NTK formulation and the standard formulation at initialization are actually the same.
The only difference is that the norm of the weight W l and the gradient of the model output with respect to W l are different for all l .
In their Appendix G, they claimed that if a network with the standard formulation is trained by minimizing the squared loss with gradient descent and learning rate η (cid:48) =
η/ ˜ d , where η is our learning rate in Lemma 5 and also their learning rate in their Theorem 2.1, then (164) is true for this network, so it is also true for a network with the NTK formulation because the two formulations have the same network output.
And then they claimed in their equation (S37) that applying learning rate η (cid:48) to the standard formulation is equivalent to applying the following learning rates [Equation] and
[Equation]
) to W l and b l of the NTK formulation, where d max = max { d 0 , · · · , d L } .
To avoid confusion, in the following discussions we will still use the NTK formulation and initialization if not stated otherwise.
Problem 1.
Claim (167) is true, but it leads to two problems.
The ﬁrst problem is that η l b =
O ( d − 1max ) since η = O (1) , while their Theorem 2.1 needs the learning rate to be O (1) .
Nevertheless, this problem can be simply ﬁxed by modifying their standard formulation as h l +
=
W l x l + β √ d l b l where b l (0) i ∼ N (0 , d − 1 l ) .
The real problem that is non-trivial to ﬁx is that by (167), there is η 0 W = d 0 d max η .
However, note that d 0 is a constant since it is the dimension of the input space, while d max goes to inﬁnity.
Consequently, in (167) they were essentially using a very small learning rate for the ﬁrst layer W 0 but a normal learning rate for the rest of the layers, which deﬁnitely does not match with their claim in their Theorem 2.1.
Problem 2.
Another big problem is that the proof of their Lemma 1 in their Appendix G is erroneous, and consequently their Theorem 2.1 is unsound as it heavily depends on their Lemma 1.
In their Lemma 1, they claimed that for some constant M > 0 , for any two models with the parameters θ and ˜ θ such that θ, ˜ θ ∈ B ( θ (0) , C 0 ) for some constant C 0 , there is
[Equation]
Note that the original claim in their paper was (cid:13)(cid:13)(cid:13)
J ( θ ) − J (˜ θ )
(cid:13)(cid:13)(cid:13) F ≤ M (cid:112) ˜
d (cid:13)(cid:13)(cid:13) θ − ˜ θ (cid:13)(cid:13)(cid:13) 2 .
This is because they were proving this result for their standard formulation.
Compared to the standard formulation, in the NTK formulation θ is (cid:112) ˜
d times larger, while the Jacobian J ( θ ) is (cid:112) ˜ d times smaller.
This is also why here we have θ, ˜ θ ∈ B ( θ (0) , C 0 ) instead of θ, ˜ θ ∈ B ( θ (0) , C
˜
d − 1 / 2 ) for the NTK formulation.
Therefore, equivalently they were claiming (168) for the NTK formulation.
However, their proof of (168) in incorrect.
Speciﬁcally, the right-hand side of their inequality (S86) is incorrect.
Using the notations in our Appendix D.3.3, their (S86) essentially claimed that
[Equation]
for any θ, ˜ θ ∈ B ( θ (0) , C 0 ) , where α l = ∇ h l h L +1 and ˜
α l is the same gradient for the second model.
Note that their (S86) does not have the (cid:112)
˜ d in the denominator which appears in (169).
This is because for their standard formulation, θ is (cid:112) ˜
d times smaller than the original NTK formulation, while (cid:13)(cid:13)
α l (cid:13)(cid:13) 2 has the same order in the two formulations because all h l are the same.
However, it is actually impossible to prove (169).
Consider the following counterexample: Since θ and ˜ θ are arbitrarily chosen, we can choose them such that they only differ in b l 1 for some 1 ≤ l < L . Then, (cid:13)(cid:13) [Equation] −
[Equation]
= [Equation] ˜ b l 1 (cid:12)(cid:12) (cid:12) .
We can see that h l +1 and ˜
h
l +1 only differ in the ﬁrst element, and (cid:12)(cid:12)(cid:12) h l + − ˜ h
l +
[Equation]
= (cid:12)(cid:12)(cid:12) β ( b l 1 − ˜ b
l 1 ) (cid:12)(cid:12)(cid:12) .
Moreover, we have W l +
=
˜
W l +1 , so there is
[Equation] [Equation]
Then we can lower bound (cid:13)(cid:13)
α l +1 − ˜
α l +1 (cid:13)(cid:13) 2 by (cid:13)(cid:13)
α l +1 − ˜
α l +1 (cid:13)(cid:13) 2 ≥ (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:104)
[Equation]
˙
[Equation]
+
[Equation]
) − [Equation] +
[Equation]
) (cid:105) W [Equation] +1 (cid:62) (cid:112) ˜
d
[Equation]
−
[Equation]
The ﬁrst term on the right-hand side is equal to (cid:12)(cid:12)(cid:12)(cid:104) ˙ σ ( h l +11 ) − ˙ σ (˜ h l +11 ) (cid:105) (cid:104) W l + / (cid:112)
˜
d, α l +2 (cid:105) (cid:12)(cid:12)(cid:12) where
W l +1is the ﬁrst row of W l +1 .
We know that (cid:13) (cid:13) W l +1(cid:13) (cid:13) 2 =
Θ (cid:16)(cid:112) ˜
d (cid:17) with high probability as its elements are sampled from N (0 , 1) , and in their (S85) they claimed that (cid:13)(cid:13)
α l +2 (cid:13)(cid:13)
= O (1) , which is true.
In addition, they assumed that ˙ σ is Lipschitz.
Hence, we can see that (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:104) [Equation] [Equation] d
[Equation]
l
[Equation]
=
[Equation]
− ˜ h l
[Equation]
(cid:12)(cid:12)(cid:12)(cid:17) =
O (cid:16)(cid:13)(cid:13)(cid:13) θ − ˜ θ (cid:13)(cid:13)(cid:13) 2 (cid:17) (172)
On the other hand, suppose that claim (169) is true, then (cid:13)(cid:13)
α l +2 − ˜
α l +2 (cid:13)(cid:13) 2 = O (cid:16) ˜
d − 1 / 2 (cid:13)(cid:13)(cid:13) θ − ˜ θ (cid:13)(cid:13)(cid:13) 2 (cid:17) .
Then we can see that the second term on the right-hand side is O (cid:16)
˜ d − 1 / 2 (cid:13)(cid:13)(cid:13) θ − ˜ θ (cid:13)(cid:13)(cid:13) 2 (cid:17) because (cid:13) (cid:13) W l +1 (cid:13) (cid:13) 2 =
O ( (cid:112) ˜ d ) and ˙ σ ( x ) is bounded by a constant as σ is Lipschitz.
Thus, for a very large ˜ d , the second-term is an inﬁnitesimal compared to the ﬁrst term, so we can only prove that
[Equation]
173) which is different from (169) because it lacks a critical ˜ d − 1 / 2 and thus leads to a contradiction.
Hence, we cannot prove (169) with the ˜ d − 1 / 2 factor, and consequently we cannot prove (168) with the (cid:112)
˜
d in the denominator on the right-hand side.
As a result, their Lemma 1 and Theorem 2. cannot be proved without this critical ˜ d − 1 / 2 .
Similarly, we can also construct a counterexample where θ and ˜ θ only differ in the ﬁrst row of some W l .
E. O UR F IXES
Regarding Problem 1, we can still use an O (1) learning rate for the ﬁrst layer in the NTK formulation given that (cid:107) x (cid:107) 2 ≤ 1 .
This is because for the ﬁrst layer, we have
[Equation]
For all l ≥ 1 , we have (cid:13)(cid:13) x l (cid:13)(cid:13)
=
O
( ˜ d 1 / 2 ) .
However, for l = 0 , we instead have (cid:13)(cid:13) x 0 (cid:13)(cid:13)
=
O (1) .
Thus, we can prove that the norm of ∇ W 0 f ( x ) has the same order as the gradient with respect to any other layer, so there is no need to use a smaller learning rate for the ﬁrst layer.
Regarding Problem 2, in our formulation (8) and initialization (9), the initialization of the last layer of the NTK formulation is changed from the Gaussian initialization W L (0) i,j ∼ N (0 , 1) to the zero initialization W L (0) i,j = 0 .
Now we show how this modiﬁcation solves Problem 2.
The main consequence of changing the initialization of the last layer is that (86) becomes different: instead of (cid:13)(cid:13) W L (cid:13)(cid:13) 2 ≤ 3 (cid:112)
˜ d , we now have (cid:13)(cid:13) W L (cid:13)(cid:13) 2 ≤ C 0 ≤ 3 4 (cid:112) ˜
d .
In fact, for any r ∈ (0 , 1 / 2) , we can prove that (cid:13)(cid:13) W L (cid:13)(cid:13) 2 ≤ 3 ˜ d r for sufﬁciently large ˜ d .
In our proof we choose r = 1 / 4 .
Consequently, instead of (cid:13)(cid:13)
α l (cid:13)(cid:13) 2 ≤ M 3 , we can now prove that (cid:13)(cid:13)
α l (cid:13)(cid:13) 2 ≤ M 3 ˜
d r − 1 / 2 for all l ≤ L by induction.
So now we can prove (cid:13)(cid:13)
α l − ˜
α l (cid:13)(cid:13) 2 = O (cid:16) ˜
d r − 1 / 2 (cid:13)(cid:13)(cid:13) θ − ˜ θ (cid:13)(cid:13)(cid:13) 2 (cid:17) instead of O (cid:16)(cid:13)(cid:13)(cid:13) θ − ˜ θ (cid:13)(cid:13)(cid:13) 2 (cid:17) , because
• For l < L , we now have (cid:13)(cid:13) α l +1 (cid:13)(cid:13)
= O ( ˜ d r − 1 / 2 ) instead of O (1) , so we can have the additional ˜
d
r − 1 / 2 factor in the bound.
• For l = L , although (cid:13)(cid:13)
α L +1 (cid:13)(cid:13) 2 = 1 , note that (cid:13)(cid:13) W L (cid:13)(cid:13) 2 now becomes O ( ˜ d r ) instead of O ( ˜ d 1 / 2 ) , so again we can decrease the bound by a factor of ˜ d r − 1 / 2 .
Then, with this critical ˜ d r − 1 / 2 , we can prove the approximation theorem with the form
[Equation]
for any r ∈ (0 , 1 / 2) , though we cannot really prove the O ( ˜ d − 1 / 2 ) bound as originally claimed in (164).
So this is how we solve Problem 2.
One caveat of changing the initialization to zero initialization is whether we can still safely assume that λ min > 0 where λ min is the smallest eigenvalue of Θ , the kernel matrix of our new formulation.
The answer is yes.
In fact, in our Proposition 3 we proved that Θ is non-degenerated (which means that Θ( x , x (cid:48) ) still depends on x and x (cid:48) ), and under the overparameterized setting where d L (cid:29) n , chances are high that Θ is full-rank.
Hence, we can still assume that λ min > 0 .
As a ﬁnal remark, one key reason why we need to initialize W L as zero is that the dimension of the output space (i.e. the dimension of h L +1 ) is ﬁnite, and in our case it is 1.
Suppose we allow the dimension of h L +1 to be ˜ d which goes to inﬁnity, then using the same proof techniques, for the NTK formulation we can prove that sup t (cid:13)(cid:13)(cid:13)
h
L +1( t ) − h L +1( t ) lin (cid:13)(cid:13)(cid:13) 2 ≤ C , i.e. the gap between two vectors of inﬁnite dimension is always bounded by a ﬁnite constant.
This is the approximation theorem we need for the inﬁnite-dimensional output space.
However, when the dimension of the output space is ﬁnite, sup t (cid:13)(cid:13)(cid:13)
h
L +1( t ) − h L +1( t ) lin (cid:13)(cid:13)(cid:13) 2 ≤ C no longer sufﬁces, so we need to decrease the order of the norm of W L in order to obtain a