U NDERSTANDING O VERFITTING IN R EWEIGHTING A LGORITHMS FOR W ORST - GROUP P ERFORMANCE
Anonymous authors
Paper under double-blind review
A BSTRACT
Prior work has proposed various reweighting algorithms to improve the worstgroup performance of machine learning models for fairness.
However, Sagawa et al.
(2020a) empirically found that these algorithms overﬁt easily in practice under the overparameterized setting, where the number of model parameters is much greater than the number of samples.
In this work, we provide a theoretical backing to the empirical results above, and prove the pessimistic result that reweighting algorithms always overﬁt.
Speciﬁcally we prove that with reweighting, an overparameterized model always converges to the same ERM interpolator that ﬁts all training samples, and consequently its worst-group test performance will drop to the same level as ERM in the long run.
That is, we cannot hope for reweighting algorithms to converge to a different interpolator than ERM with potentially better worst-group performance.
Then, we analyze whether adding regularization helps ﬁx the issue, and we prove that for regularization to work, it must be large enough to prevent the model from achieving small training error.
Our results suggest that large regularization (or early stopping) and data augmentation are necessary for reweighting algorithms to achieve high worst-group test performance.
I NTRODUCTION
It has been well established by prior work that overparameterized models, whose number of parameters is much larger than the number of training samples, can empirically achieve high test performance on a variety of tasks, in contrast to the theory that models with too many parameters could have large generalization error.
This high performance however is on average ; a large body of prior work (Hovy & Søgaard, 2015;
Blodgett et al., 2016; Tatman, 2017) showed that these models tend to learn spurious features , such as learning the background in image classiﬁcation instead of the object, and learning keywords like “not” in language sentiment analysis instead of really understanding the sentences.
Consequently, these models are unfair , i.e. they fail on certain minority groups (such as positive sentences containing “not”) while still having high average-case performance.
To solve this problem, people have proposed various reweighting algorithms to improve the model’s worst-group performance, such as upweighting the minority groups or using distributionally robust optimization (DRO) based methods (Shimodaira, 2000; Hashimoto et al., 2018; Duchi & Namkoong, 2018;
Sagawa et al., 2020a).
While reweighting algorithms in principle can improve the worst-group performance compared to vanilla empirical risk minimization (ERM), previous work empirically found that when applied to modern overparameterized models, these methods could overﬁt very easily, so that they have poor test worst-group performance.
For example, Sagawa et al.
(2020a) studied a reweighting algorithm called group DRO.
They found that compared to ERM, group DRO does improve the worst-group test accuracy by a large margin at the early stage of training.
However, if no regularization is applied, then as training goes on, the worst-group test accuracy of group DRO will drop signiﬁcantly and eventually to a level almost the same as ERM.
Some previous work tried to explain why reweighting algorithms
can overﬁt so easily.
For instance, Sagawa et al.
(2020b) argued that with these algorithms, an overparameterized model would typically memorize all training samples in the minority groups while still learning the spurious features from the majority groups.
In this work, we aim to understand the overﬁtting phenomenon in reweighting algorithms by studying their implicit biases.
Speciﬁcally, we prove for a family of overparameterized neural networks that for almost all reweighting algorithms, the model always converges to the same interpolator that ﬁts all training samples, no matter the reweighting.
Since ERM is a special case of such reweighting algorithms (where each sample receives the same weight), this means that the implicit biases of all reweighting algorithms are equivalent to that of ERM.
Consequently, the model trained by any reweighting algorithm always overﬁts to the ERM interpolator, so we cannot hope for its worstgroup test performance to be better than ERM.
In short, reweighting algorithms always overﬁt .
Given this pessimistic result, we analyze whether regularization can help mitigate overﬁtting, as proposed by Sagawa et al.
(2020a).
We ﬁnd that a necessary condition for regularization to work is that it considerably lowers the training performance.
Speciﬁcally, we prove that if the overparameterized model trained by a reweighting algorithm with regularization can still perform almost perfectly on the training set, then overﬁtting is still inevitable.
This explains why in practice we need very large regularization that prevents the model from achieving nearly zero training error to avoid overﬁtting.
Our results have two important consequences for practice: (i) We should always use large regular-
ization
or early stopping when optimizing for worst-group performance; (ii) We should always try to obtain more training samples, e.g. with strong data augmentation or semi-supervised learning.
1. R ELATED WORK
Group fairness.
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.
Later,
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 problem we study in this paper is most closely related to Rawlsian max-min fairness.
A large body of recent work in machine learning have studied how to improve this worstgroup performance (Duchi & Namkoong, 2018; Oren et al., 2019; Xu et al., 2020; Liu et al., 2021;
Zhai et al., 2021).
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., 2013) and counterfactual fairness (Kusner et al., 2017), which we do not study in this work.
Implicit 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 training loss minimizer the model parameters will converge to (Du et al., 2019; Allen-Zhu et al., 2019).
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.
P RELIMINARIES
Consider a data domain X × Y ⊆
R d × R that consists of K groups (subdomains) 1 , where each data point belongs to one of the groups
We assume that the input space X is a subset of the unit ball of R d , such that any x ∈ X satisﬁes (cid:107) x (cid:107) 2 ≤ 1 .
We are given a training set { ( x i , y i ) } ni =1 i.i.d.
sampled from some underlying distribution P over X ×Y .
Let the K groups be D 1 , · · · , D K where each D i is a subset of X × Y .
Let P k
( z ) =
P ( z | z ∈ D k ) be the conditional data distribution over D k , where z = ( 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 , we overload notation and use g ( X ) =
( g ( x 1 ) , · · · , g ( x n )) .
Let the loss function be (cid:96) :
Y × Y → [0 , 1] .
In vanilla training, the goal is to minimize the expected risk denoted by R ( f ; P ) =
E z ∼ P
[ (cid:96) ( f ( x ) , y )] , which is done by minimizing the empirical risk
ˆ R
[Equation]
For tasks requiring high worst-group performance, the goal is to train a model f : X → Y that performs well over every P k , which can be achieved by minimizing the worst-group risk deﬁned as
[Equation]
2. R EWEIGHTING A LGORITHMS
Most existing methods that minimize the worst-group risk are reweighting algorithms that assign each sample with a weight during training and minimize the weighted average risk.
At time t , we assign a weight q ( t ) i to sample z i , and minimize the weighted empirical risk:
[Equation]
where q ( t ) = ( q ( t ) 1 , · · · , q ( t ) n ) and q ( t ) 1 + · · · + q ( t ) n = 1 .
A static reweighting algorithm 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 famous example is Importance Weighting (IW, Shimodaira (2000)), in which if z i ∈
D k and the size of D k is n k , then q i =
( Kn k ) − 1 .
Under IW, each group has the same weight, and the reweighted empirical risk is a simple (unweighted) average of the empirical risk over each group, so that each group has an equal contribution to the overall risk objective.
Note that ERM is also a special case of static reweighting algorithms: by assigning q 1 = · · · = q n = 1 /n .
On the other hand, in a dynamic reweighting algorithm , q ( t ) changes with t .
Speciﬁcally, it upweights samples over which the model has a high risk in order to help the model learn “hard” samples.
A popular dynamic reweighting algorithm is Group DRO (Sagawa et al., 2020a).
Denote the empirical risk over group k by ˆ R k ( f ) , and the model at time t by f ( t ) .
Group DRO 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 by
[Equation]
for some ν > 0 , and then normalized so that q ( t ) 1 + · · · + q ( t ) n = 1 .
Sagawa et al.
(2020a) proved a convergence rate theorem (their Proposition 2) showing that in the convex setting, the worst-group training risk of Group DRO converges to the global minimum with the rate O ( t − 1 / 2 ) .
There are many other reweighting algorithms.
Particularly, all variants of DRO and DRO-based methods like CVaR and χ 2 -DRO are reweighting algorithms.
See Appendix A for more examples.
2. R EWEIGHTING ALGORITHMS CAN EASILY OVERFIT
In this section, we will empirically demonstrate that while IW and Group DRO can achieve higher worst-group test performances than ERM at the early stage of training, they can easily overﬁt after a number of training epochs.
Following Sagawa et al.
(2020a), we conduct the experiment on two datasets: Waterbirds and
CelebA. Each dataset contains a binary confounding variable a and a binary target variable y , dividing the dataset into four groups (four combinations of ( a, y ) ).
In Waterbirds y is the type of the bird and a is the background; In CelebA y is whether the person has blond hair and a is whether the person is male.
On each dataset, a model trained by ERM always exhibits a very strong empirical correlation between y and a , so its performance on one of the groups is extremely poor.
The goal is to make the model perform well on every group.
See Appendix C.1 of Sagawa et al.
(2020a) for detailed information of these datasets.
On each dataset, we use the ResNet18 model as the classiﬁer and optimize it with momentum SGD.
We run each of the three algorithms: ERM, IW and group DRO (GDRO), for 500 epochs on Waterbirds and 200 epochs on CelebA, and plot the average training/test and worst-group (WG) training/test accuracy curves throughout training in Figure 1.
From the plots we can conclude that:
• All algorithms can achieve and maintain high average training/test accuracy throughout training, i.e. there is almost no overﬁtting in the average test accuracy.
[Figure]
[Figure]
• Regarding the worst-group test accuracy, while the two reweighting algorithms outperform ERM by a large margin at the early epochs, they overﬁt very quickly.
On CelebA after roughly 100 epochs, the worst-group test accuracies of the two reweighting algorithms become the same as ERM.
On Waterbirds, the worst-group test performances of IW and Group DRO drop signiﬁcantly after around 30 epochs though they are still better than ERM.
I MPLICIT BIASES OF REWEIGHTING ALGORITHMS
In the previous section, we empirically demonstrated that the worst-group test performances of reweighting algorithms converge to the same level as ERM.
To theoretically understand why this happens in practice, we analyze the implicit biases of reweighting algorithms.
Our main theorem (Theorem 6) states that almost all reweighting algorithms (including ERM) have equivalent implicit biases, in the sense that they converge to the same interpolator.
Meanwhile, it is observed in practice that the ERM interpolator has a poor worst-group test performance.
This leads to the pessimistic result that reweighting algorithms always overﬁt .
All proofs can be found in Appendix B.
3. L INEAR MODELS
We ﬁrst demonstrate this pessimistic result on simple linear models to provide our readers with a key intuition, and later we will apply this same intuition to neural networks.
Let the linear model be f ( x ) =
(cid:104) θ, x (cid:105) , where θ ∈ R d .
In the overparameterized setting, we have d > n .
Consider using the squared loss (cid:96) (ˆ y, y )
=
12 (ˆ y − y ) 2 , and minimizing the weighted empirical risk with gradient descent:
[Equation]
where η > 0 is the learning rate.
For a linear model with the squared loss, the update rule is
[Equation]
It is a well known result that under the overparameterization setting where d > n , if x 1 , · · · , x n are linearly independent, then with a sufﬁciently small η , a linear model trained by ERM can always converge to an interpolator which ﬁts all training samples (i.e. θ ( t ) → θ ∗ such that (cid:104) θ ∗ , x i (cid:105) =
y i for all i ).
Here the linear independence is necessary, because otherwise in the extreme case where x 1 = x 2 but y 1 (cid:54) =
y 2 , the model cannot ﬁt ( x 1 , y 1 ) and ( x 2 , y 2 ) simultaneously.
In this section, we aim to extend this ERM convergence analysis to general reweighting algorithms.
Our results require the following assumption:
Assumption 1.
There exist constants q 1 , · · · , q n such that for all i , q ( t ) i → q i as t → ∞ .
And min i q i =
q ∗ > 0 . This assumption avoids the scenario where there is some i such that q ( t ) i
≈
for all t , in which case the model could never ﬁt z i .
Assumption 1 empirically holds for Group DRO on Waterbirds and CelebA (see Appendix D.2).
Under this assumption, we can prove that the model always converges to an interpolator:
Theorem 1.
For any reweighting algorithm satisfying Assumption 1, if x 1 , · · · , x n are linearly independent, then there exists an η 0 > 0 such that for any η ≤ η 0 , as t → ∞ , θ ( t ) converges to some interpolator θ ∗ such that for all i , (cid:104) θ ∗ , x i (cid:105) =
y i .
We now make the following key observation regarding the update rule (5): θ ( t +1) − θ ( t ) is a linear combination of x 1 , · · · , x n for all t , and thus θ ( t ) − θ (0) always lies in the linear subspace span ( x 1 , · · · , x n ) .
Note that this is an n -dimensional linear subspace if x 1 , · · · , x n are linearly independent, and by Cramer’s rule, there is exactly one ˜ θ in this subspace such that
(cid:104) ˜ θ
+ θ (0) , x i (cid:105) =
y i for all i , which implies that θ ∗ =
˜ θ
+ θ (0) is unique.
Together with Theorem 1, this leads to:
Theorem 2.
If x 1 , · · · , x n are linearly independent, then there exists η 0 > 0 such that for any reweighting algorithm satisfying Assumption 1, and any η ≤ η 0 , θ ( t ) converges to the same interpolator θ ∗ that does not depend on q ( t ) i .
Note that ERM is also a reweighting algorithm satisfying Assumption 1.
Therefore, we have essentially proved the following result:
The implicit bias of any reweighting algorithm satisfying Assumption 1 is equivalent to ERM, so reweighting algorithms always overﬁt 3 .
The key intuition here is that no matter what reweighting algorithm we use, θ ( t ) − θ (0) always lies in a low-dimensional subspace, in which the interpolator is unique.
Therefore, as long as a model trained by the algorithm converges to some interpolator, it must converge to that unique interpolator, which means that the implicit bias of the algorithm is equivalent to ERM.
3. L INEARIZED NEURAL NETWORKS
Now we prove the same result for neural networks.
Of course it would be very hard to prove it for all neural networks.
However, we can prove the result for a family of overparameterized neural networks that can be approximated by their linearized counterparts Lee et al.
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.
The linearized neural network of f ( t ) ( x ) is deﬁned as
[Equation]
(6) where we use the shorthand ∇ θ f (0) ( x ) :
= ∇ θ f ( x ; θ ) (cid:12)(cid:12) θ
= θ 0 .
Consider training f ( t ) lin ( x ) via gradient descent on the reweighted risk (as in (4)) using the squared loss.
Given a training set { ( x i , y i ) } ni =1 , we can construct a new training set { (cid:0) ∇ θ f (0) ( x i ) , y i − f (0) ( x i ) (cid:1) } ni =1 , so that training a linearized neural network on the original training set is equivalent to training a linear model on the new training set.
Based on this observation, we have the following corollary of Theorem 2:
Corollary 3.
If ∇ θ f (0) ( x 1 ) , · · · , ∇ θ f (0) ( x n ) are linearly independent, then there exists η 0 >such that for any reweighting algorithm satisfying Assumption 1, and any η ≤ η 0 , θ ( t ) converges to the same interpolator θ ∗ that does not depend on q i .
Here we are still using the key intuition: θ ( t ) − θ (0) always lies in the n -dimensional linear subspace span (cid:0) ∇ θ f (0) ( x 1 ) , · · · , ∇ θ f (0) ( x n ) (cid:1) .
By Cramer’s rule, there is a unique interpolator θ ∗ such that θ ∗ − θ (0) ∈ span (cid:0) ∇ θ f (0) ( x 1 ) , · · · , ∇ θ f (0) ( x n ) (cid:1) , and θ ( t ) always converges to that θ ∗ .
Thus, we have essentially proved that for linearized neural networks, reweighting algorithms always overﬁt.
Now let us delve deeper into the training dynamics of a linearized neural network.
Note that ∇
[Equation]
lin [Equation] X ) = ∇ θ f (0) ( X ) ∈ R p × n , so the change in the training function value vector is
[Equation]
where Q ( t ) =
diag ( q ( t ) 1 , · · · , q ( t ) n ) .
The function value vector moves along the kernel gradient with respect to Θ (0) q ( t ) = ∇ θ f (0) ( X ) (cid:62) ∇ θ f (0) ( X ) Q ( t ) .
Meanwhile, the neural tangent kernel (NTK, Jacot et al. (2018)) 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 ) , so Θ (0) q ( t ) =
Θ (0) Q ( t ) .
We can thus extend our result for gradient descent on linearized neural networks to a kernel gradient descent algorithm as above.
3. W IDE FULLY - CONNECTED NEURAL NETWORKS
Now we prove the result for sufﬁciently wide fully-connected neural networks, which can be approximated by the linearized neural networks.
First we deﬁne a fully-connected neural network with L hidden layers (we always assume L ≥
so there is at least one hidden layer).
Let h l and x l be the pre- and post-activation outputs of layer l , and d l be the width of layer l .
Let x 0 = x and d 0 = d .
Deﬁne the neural network as
[Equation]
where σ is a non-linear activation function,
W l ∈ R d l +1 ×
d l and W L ∈ R 1 × d L .
The parameters θ consist of W 0 , · · · , W L and b 0 , · · · , b L ( θ is the concatenation of all ﬂattened weights and biases).
The ﬁnal output of the neural network is f ( x ) =
h L +1 .
And let the neural network be initialized as
[Equation] [Equation]
, L
[Equation] [Equation]
and (cid:40) W L (0) i,j =
[Equation]
L
[Equation] [Equation]
j
[Equation] [Equation] We also need the following assumption for our approximation theorem:
Assumption 2.
σ is differentiable everywhere, and both σ and ˙ σ are Lipschitz.
Difference from Jacot et al.
Our initialization (9) is different from the original one in
Jacot et al.
(2018) in the last (output) layer.
For the output layer, 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 enables us to accurately approximate the neural network with its linearized counterpart (6), as we notice that the proofs in Lee et al.
(particularly the proofs of their Theorem 2.1 and their Lemma 1 in
Appendix G) are ﬂawed.
In Appendix C we will explain what goes wrong in their proofs and how we manage to ﬁx the proofs with our modiﬁcation.
For our new initialization, we still have the following NTK theorem:
Theorem 4.
If σ is Lipschitz and d l → ∞ for l = 1 , · · · , L sequentially, then Θ (0) ( x , x (cid:48) ) converges in probability to a non-degenerated 5 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-degenerated, so we assume that λ min > 0 (which holds almost surely in the overparameterized setting where d L (cid:29) n ).
Then, we can prove the following approximation theorem:
Theorem 5 (Approximation Theorem) .
Let η ∗ =
( λ min + λ max ) − 1 .
For a fully-connected neural network f ( t ) that satisﬁes Assumption 2 and is trained by any reweighting algorithm satisfying Assumption 1, let f ( t ) lin be its linearized neural network which is trained by the same reweighting algorithm (i.e. ∀ i, t , q ( t ) i are the same for both networks).
If d 1 = d 2 = · · · = d L = ˜
d and λ min > 0 , then for any δ > 0 , there exists ˜ D > 0 and a constant C such that as long as η ≤ η ∗ and ˜ d ≥ ˜ D , for any test point x ∈ R d such that (cid:107) x (cid:107) 2 ≤ 1 , with probability at least 1 − δ over random initialization,
[Equation]
[Equation]
[Equation]
Remark.
We can easily extend this theorem to the case where there exists α l > 0 for each of l = 2 , · · · , L such that d l /d 1 → α l and d 1 → ∞ .
Combining all the above results altogether, we achieve our main theorem:
Theorem 6.
Under the conditions of Theorem 5, there exists an η 1 > 0 such that if η ≤ η 1 and ∇ θ f (0) ( x 1 ) , · · · , ∇ θ f (0) ( x n ) are linearly independent, then as ˜ d → ∞ , for any test point x ∈ R d such that (cid:107) x (cid:107) 2 ≤ 1 , with probability close to 1 over random initialization,
[Equation]
where f ( t ) is trained by the reweighting algorithm and f ( t ) ERM is trained by ERM.
The main theorem shows that at any test point x , the gap between the function values of the two models converges to an inﬁnitely small term, so the worst-group test performance of the reweighting algorithm will converge to the same level as ERM.
Therefore, we have proved that for sufﬁciently wide fully-connected neural networks, reweighting algorithms always overﬁt.
Our key intuition tells us that the change in the model parameters always lies in an n -dimensional subspace.
Thus, one possible way to improve the worst-group test performance is to enlarge this subspace by adding more training samples, e.g. via data augmentation or semi-supervised learning.
However, even if we have more training samples, as long as the model is still overparameterized, and all ∇ θ f (0) ( x i ) are linearly independent, then our result still says that no reweighting algorithm can do better than ERM in the long run (though the performance of ERM itself might be improved).
Moreover, our theoretical results can explain the surprising empirical observation in Sagawa et al.
(2020b) that removing some samples from the majority groups to match the group sizes can sometimes achieve even higher worst-group test performance than reweighting even though it wastes lots of data (see their Section 6).
When training samples are removed, the model will converge to an interpolator of the smaller training set which is different from the interpolator of the original training set, so there is a chance that the performance of the new interpolator is actually higher.
D OES REGULARIZATION REALLY HELP ?
In the previous section, we proved the pessimistic result that reweighting algorithms always overﬁt, i.e. in the long run their worst-group test performances always drop to the same level as ERM.
And even if we use strong data augmentation or semi-supervised learning, reweighting algorithms still cannot outperform ERM if the training set is not sufﬁciently enlarged.
Sagawa et al.
(2020a) proposed to tackle the overﬁtting problem of reweighting algorithms via regularization.
In particular, they empirically demonstrated with experiments that large regularization is required to prevent reweighting algorithms such as group DRO from overﬁtting.
With a large regularization, the model can maintain a high test worst-group performance, but it cannot obtain perfect training accuracy, in contrast to the case where no regularization is applied.
In this section, we study the necessary conditions for regularization to maintain high worst-group test performance.
Speciﬁcally, we will show that regularization will not work if it is not large enough to prevent the model from obtaining nearly zero training error.
In other words, lowering the training performance is the key to keeping a high worst-group test performance.
Note that the results in this section do not require Assumption 1 , so the results hold for all reweighting algorithms.
4. T HEORETICAL ANALYSIS
Consider a reweighting algorithm with sample weights q ( t ) i .
Following Sagawa et al.
(2020a), we consider adding L 2 penalty to the weighted empirical risk (2):
[Equation]
[Equation]
Given that sufﬁciently wide neural networks can be approximated by linearized ones, we ﬁrst focus on linearized neural networks.
We will use the subscript “reg” to refer to a regularized model (which is trained trained by minimizing the regularized risk (12)).
Let f ( t ) linreg be a regularized linearized neural network trained by some reweighting algorithm, and f ( t ) linERM be an unregularized linearized neural network trained by ERM.
As before, we consider training the models with gradient descent under the squared loss (cid:96) (ˆ y, y )
=
12 (ˆ y − y ) 2 .
The following result shows that these two models are very close if f ( t ) linreg can achieve low training error:
Theorem 7.
If there is a constant M 0 > 0 such that (cid:13)(cid:13) ∇ θ f (0) ( x ) (cid:13)(cid:13) 2 ≤ M 0 for all (cid:107) x (cid:107) 2 ≤ 1 , ∇ θ f (0) ( x 1 ) , · · · , ∇ θ f (
[Equation]
( x n ) are linearly independent, and the empirical training risk of f ( t ) linreg satisﬁes
[Equation]
for some (cid:15) > 0 , then for any test point x such that (cid:107) x (cid:107) 2 ≤ 1 we have
[Equation]
The proof of this theorem also follows the key intuition: we can show that even with the L 2 penalty added, θ ( t ) − θ (0) is still limited in a low-dimensional subspace.
And although we cannot prove that θ ( t ) always converges to the ERM interpolator, we can prove that it can get very close to that interpolator if its training error is very low, so the resulting model is very close to the ERM model.
Then, we can extend this result to sufﬁciently wide fully-connected neural networks:
Theorem 8.
If λ min > 0 and µ > 0 , then let η ∗ = ( µ + λ min + λ max ) − 1 .
For a wide fully-connected neural network f ( t ) reg deﬁned by (8) and (9) and satisfying Assumption 2, and any reweighting algorithm, 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 as ˜ d → ∞ , with probability close to 1 over random initialization, for any test point x such that (cid:107) x (cid:107) 2 ≤ 1 we have
[Equation] [Equation]
The result shows that a regularized model trained by any reweighting algorithm will get very close to an unregularized ERM model at any test point x if the training error of the former is nearly zero.
Thus, regularization only helps when it is large enough to keep the training error of the model away from zero by a margin.
Our results explain the empirical observation of Sagawa et al.
(2020a) that by using large regularization, the model can maintain a high worst-group test performance, but it cannot achieve perfect training accuracy.
If smaller regularization is applied and the model can achieve nearly perfect training accuracy, then its worst-group test performance will still signiﬁcantly drop.
4. E MPIRICAL STUDY
In this section, we validate our theoretical results above with experiments on Waterbirds and CelebA.
We run ERM, IW and group DRO under different levels of weight decay for 500 epochs on Waterbirds and 250 epochs on CelebA. Note that we do not strictly follow our L 2 penalty formulation (12), but we study the L 2 weight decay regularization which is most widely used in practice.
We repeat each experiment ﬁve times with different random seeds and report the 95% conﬁdence interval of the mean average training and worst-group test accuracies of the last 10 training epochs in Table 1.
To compare with early stopping, we also report the mean accuracies of epochs 11-20 with no regularization in blue.
Moreover, we plot the average training and worst-group test accuracy curves throughout training for IW and Group DRO with one of the random seeds in Figure 2.
On both datasets, early stopping achieve the best performances.
Particularly, on Waterbirds, there is no clear sign that regularization could help prevent overﬁtting.
When the regularization is small,
[Table]
[Figure]
the training accuracy is still 100% and the algorithm continues to overﬁt.
However, when the regularization is large enough to lower the training accuracy, the worst-group test accuracy drops more because the model cannot learn the samples well under such a large regularization.
Thus, perhaps not surprisingly, a lower training performance is only a necessary condition but not sufﬁcient.
On CelebA, regularization does help mitigate overﬁtting, but a useful regularization must be large enough to lower the training accuracy.
We observe that Group DRO overﬁts more slowly than IW, as it still has over 70% worst-group test accuracy after 70 epochs.
However, as Figure 2d clearly shows, its worst-group test accuracy will still drop to the ERM level at 200 epochs.
We also notice that Group DRO requires a smaller regularization than IW: for IW we need the weight decay level to be as large as 0.1 to achieve a similar performance as early stopping, but for Group DRO it only needs to be 0.01, and using 0.1 is actually harmful.
Overall, we ﬁnd that early stopping achieves a markedly better performance.
On the other hand, using large regularization could result in training instability, as well as a loss in overall performance, and there may or may not be a small band for the regularization parameter where the worst-group test performance is better.
C ONCLUSION
In this work, we theoretically studied why reweighting algorithms overﬁt in practice by analyzing their implicit biases.
Speciﬁcally, we proved the pessimistic result that reweighting algorithms always overﬁt.
Our proof was based on the key intuition that the change in model parameters always lies in a low-dimensional subspace, so that even with reweighting, the model still converges to the same unique interpolator.
When regularization is applied, we proved that the regularization must be large enough to keep the model from achieving nearly zero training error in order to prevent overﬁtting.
We empirically validated our theoretical results on real datasets, and our results can also explain the empirical observations in previous work.
Our results are especially important for large-scale machine learning tasks, where early stopping is not always possible in order to achieve high performances.
Practitioners shooting for high worst-group performances in those tasks must be very careful about to what extent overﬁtting affects reweighting algorithms.
R EPRODUCIBILITY STATEMENT
To guarantee the reproducibility of all our empirical results, in all our experiments we use a ﬁxed set of random seeds, and we run some of the experiments twice with the same random seed to make sure that the outputs are the same.
See Appendix D.1 for experiment details.
After this paper is deanonymized we will provide a GitHub repository that contains all the codes, datasets, hyperparameters, random seeds, machine speculations and anaconda environment speculations that are sufﬁcient to exactly reproduce our empirical results
A O THER REWEIGHTING ALGORITHMS
In this section, we will review some other previously proposed reweighting algorithms.
First, we will look at DRO-based methods, where DRO stands for Distributionally Robust Optimization.
DRO is designed for tasks with distributional shift , where the training distribution and the test distribution are different, and there are some constraints on the distance between these two distributions (typically described by a divergence function D ).
Since the real test distribution is unknown, DRO minimizes the model’s risk over the worst distribution that satisﬁes the distance constraints, which is an upper bound of the model’s real test error.
Formally speaking, given a training distribution P ,
DRO minimizes the expected risk over the worst-case distribution Q in a ball w.r.t.
divergence D around the training distribution P .
For group shift problems which require high worst-group performance, Q also needs to be absolutely continuous with respect to P , i.e. Q (cid:28) P .
Overall, DRO minimizes the following expected DRO risk :
[Equation]
The expected DRO risk is typically minimized in the following way: for each epoch t , we ﬁrst ﬁnd the worst Q that maximizes E Q [ (cid:96) ( θ ; Z )] and satisﬁes D ( Q (cid:107) P ) ≤ ρ, Q (cid:28) P , and then minimize the model’s expected risk over this Q with gradient descent.
The rationale behind this algorithm is the famous Danskin’s Theorem, which says that if F ( x ) is the maximum of a family of functions, then its gradient at point x is equal to the gradient of the function that attains the maximum value at x .
Note that in practice we only have a ﬁnite set of training samples { z 1 , · · · , z n } , so P is always chosen as the empirical distribution, i.e. uniform distribution over z 1 , · · · , z n .
Then, note that
Q (cid:28) P , which implies that the support of Q must be a subset of the support of P , which is { z 1 , · · · , z n } .
This means that Q must be a distribution over z 1 , · · · , z n , i.e. it is a reweighting over the training samples.
Thus, we have essentially showed that DRO is a reweighting algorithm, and in fact almost all methods based on DRO are reweighting algorithms.
Two widely used variants of DRO are CVaR (Conditional Value at Risk) and χ 2 -DRO.
In CVaR, for a ﬁxed α ∈ (0 , 1) , we let D ( Q (cid:107) P ) =
sup log dQdP and ρ = − log α .
As a result, suppose that αn is an integer, then CVaR will assign weight 1 αn to αn training samples that incur the highest losses, and weight 0 to the rest of the samples, so we can easily see that CVaR is a reweighting algorithm.
χ 2 -DRO was ﬁrst used in Hashimoto et al.
(2018) to deal with fairness tasks where the group labels are unknown, where D ( Q (cid:107) P ) =
12 (cid:82) ( dQ/dP − 1) 2 dP and ρ = 12 ( 1 α − 1) 2 .
χ 2 -DRO is also a reweighting algorithm.
There are many other previously proposed methods of maximizing the worst-group performance that are also reweighting algorithms.
For instance, Xu et al.
(2020) studied the imbalanced class problem where a standard trained model always has high performance over classes with many training samples and low performance over minority classes.
They proposed to balance the classes with
Label CVaR, which is based on DRO and is a reweighting algorithm.
Liu et al.
(2021) proposed a two-stage training process called JTT: in the ﬁrst identiﬁcation stage they trained a model with ERM to identify training samples that are hard to learn, and in the second upweighting stage they trained a new model with the hard samples upweighted, so that the model could learn all samples equally well.
As the process itself suggests, JTT is a reweighting algorithm.
Finally, Zhai et al.
(2021) argued that DRO-based methods are very sensitive to outliers in the training set because they upweight training samples with high losses and outliers tend to incur high losses.
They proposed the DORO algorithm which at each iteration removes the samples with the highest losses, and then performs DRO on the rest of the samples.
DORO is a reweighting algorithm.
B P ROOFS
Notations.
In all of the proofs, for a 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.
B. P ROOF OF T HEOREM
To help our readers understand the proof more easily, we will ﬁrst prove the result for static reweighting algorithms where q ( t ) i
=
q i for all t , and then we will prove the result for dynamic reweighting algorithms that satisfy q ( t ) i
→ q i as t → ∞ .
B.1. S TATIC REWEIGHTING ALGORITHMS
We ﬁrst prove the result for all static reweighting algorithms such that min i q i = q ∗ > 0 .
We will use a standard optimization proof technique called smoothness .
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 (cid:105) is
[Equation]
[Equation]
whose Hessian is
[Equation]
So for any unit vector v ∈ R d , we have (since q i ∈ [0 , 1] )
[Equation]
[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 ) )
=
√ Q ( X (cid:62) θ ( t ) − Y ) ∈
R n where √ Q = diag ( √ q 1 , · · · , √ q n ) .
We can see that (cid:13)(cid:13) g
( θ ( t ) )
(cid:13)(cid:13) 2 2 = F ( θ ( t ) ) , so that ∇ F ( θ ( t ) )
=
2 X √ Q g ( θ ( t ) ) .
The update rule of a static reweighting algorithm with gradient descent and the squared loss is:
[Equation]
Substituting θ 1 and θ 2 in (21) 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) 2 =
(cid:112) q ∗ F ( θ ( t ) ) , so we have g ( θ ( t ) )
(cid:62) √ Q (cid:62) X (cid:62) X √ Q 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 .
Moreover, (cid:112) F ( θ ( t +1) ) ≤ (1 − ηq ∗ λ min) (cid:112) F ( θ ( t ) )
due to √ 1 − x ≤ 1 − x/ 2 .
The convergence in F ( θ ) implies the convergence in θ .
This is because
[Equation]
which implies that for any η ≤ η 0 ,
[Equation]
Therefore, lim T →∞ (cid:80) ∞ t = T (cid:13)(cid:13) θ ( t +1) − θ ( t ) (cid:13)(cid:13) 2 = 0 , which means that θ ( t ) converges, and it converges to some interpolator.
B.1. D YNAMIC REWEIGHTING ALGORITHMS
Now we prove the result for all dynamic reweighting algorithms satisfying Assumption1.
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]
[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 .
Select and a ﬁx an (cid:15) such that 0 < (cid:15) < max { 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 reweighting algorithm with gradient descent and the squared loss is:
[Equation]
(28) 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 for all (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]
where Q ( t [Equation] ( t
[Equation] Q (29)
Again, substituting θ 1 and θ 2 in (21) 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 as proved before, we also have
[Equation] [Equation]
Since (cid:15) ≤ ( q ∗ λ min ) 2 12 λ max2 , we have
[Equation]
t
[Equation]
Thus,
[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 .
As before, we can prove that the convergence in F ( θ ) implies the convergence in θ .
Thus, θ converges to some interpolator.
B. P ROOF OF T HEOREM
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 9 (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] [Equation]
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]
and
[Equation]
Then we can achieve (37) by the law of large numbers.
B. P ROOF OF T HEOREM
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 (29), 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 reweighting algorithm with gradient descent and the squared loss for the wide neural network is:
[Equation] [Equation]
and for t ≥ t (cid:15) , it can be rewritten as
[Equation]
First, we will prove the following theorem:
Theorem 10.
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 11 (Local Lipschitzness of the Jacobian) .
Under Assumption
[Equation]
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 can be found in Appendix B.4.
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]
And by Theorem 4, 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 11.
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) R 0 q ∗ λ min .
By Lemma 11, there exists D 1 > 0 such that with probability at least (1 − δ/ 3) , for any ˜ d ≥ D 1 , (47) is true for all θ, ˜ θ ∈ B ( θ (0) , C 0 ) .
By union bound, with probability at least (1 − δ ) , (47), (48) and (49) are all true.
Now we assume that all of them are true, and prove (43) and (44) by induction.
(43) is true for t =
due to (48), and (44) is always true for t = 0 .
Suppose (43) and (44) are true for t , then for t +
we have
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation] [Equation]
So ( is
[Equation]
Therefore, (43) and (44) 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 (45) is true for t = t (cid:15) .
And (46) 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 (42).
If (45) and (46) are true for t , then for t + 1 , there is
[Equation]
g ( θ
[Table]
t ) )
(cid:13)(cid:13)(cid:13)
[Equation]
η (cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
[Equation]
( t ) )
(cid:113)
[Equation]
[Table]
t
[Table]
[Equation]
(cid:15) (cid:13)(cid:13)(cid:13)(cid:13)
F (cid:13)(cid:13)(cid:13)(cid:112)
[Equation]
[Equation]
( t )
[Equation]
(cid:13)(cid:13)(cid:13)
[Equation]
(cid:13)(cid:13)(cid:13)
[Equation]
(cid:13)(cid:13)(cid:13)
[Equation]
(cid:18)
[Equation]
(cid:19)
[Equation]
(cid:15)
[Equation]
(cid:15)
[Equation]
[Table]
[Equation]
+
[Equation] [Equation]
where ˜ θ ( t ) is some linear interpolation between θ ( t ) and θ ( t +1) .
Now we prove that
[Equation]
For any unit vector v ∈ R n , we have
[Equation]
55) (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:62)
[Equation] [Equation]
[Equation] [Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table] [Table]
for
[Table]
, which implies that
[Table]
[Equation]
(due to (32)) (57) for all (cid:15) ≤ (cid:15) 0 .
Thus, (45) is also true for t + 1 .
In conclusion, (45) and (46) are true with probability at least (1 − δ ) for all ˜ d
[Equation]
Returning back to the proof of Theorem 5.
Choose and ﬁx an (cid:15) such that (cid:15) < min { (cid:15) 0 , 13 (cid:16) q ∗ λ min 3 λ max + q ∗ λ min (cid:17) 2 } , where (cid:15) 0 is deﬁned by Theorem 10.
Then, t (cid:15) is also ﬁxed.
There exists ˜ D ≥ 0 such that for any ˜ d ≥ ˜ D , with probability at least (1 − δ ) , Theorem 10 and Lemmaare true and
[Equation]
58) which immediately implies that
[Equation]
We still denote B = 1 + η ∗ M 2 and C 0 =
MB t(cid:15) R 0 B −+ 3 √
[Equation]
3 (cid:15)MB t(cid:15) R 0 q
[Equation]
λ
[Equation]
Theorem 10 ensures that for all t , θ ( t )
∈ B ( θ (0) , C 0 ) .
Then we have
[Equation]
so it follows that
[Equation]
√ 3 (cid:15) (61)
Thus, for all (cid:15) < 13 (cid:16)
[Equation]
, there is
[Equation]
The update rule of the reweighting algorithm for the linearized neural network is:
[Equation] [Equation]
where we use the subscript “lin” to denote the linearized neural network
[Equation]
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] [Equation]
where ˜ θ ( t ) is some linear interpolation between θ ( t ) and θ ( t +1) .
Thus,
[Equation] [Equation]
By Lemma
[Equation]
we have
[Equation]
which implies that for all t < t (cid:15) ,
[Equation]
(cid:13)(cid:13)(cid:13)
[Equation]
(cid:13)(cid:13)(cid:13)
[Equation]
( t
[Equation]
(cid:13)(cid:13)(cid:13)
[Equation]
Therefore, we have
[Equation]
Since ∆ 0 = 0 , it follows that for all t ≤ t (cid:15) ,
[Equation]
69) and particularly we have
[Equation] [Equation]
For t ≥ t (cid:15) , we have the alternative update rule (42).
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 (cid:112) Q ∆ t +1 =
A (cid:112) Q ∆ t + η (cid:112) Q (cid:104)
J
[Equation]
˜ θ
[Equation]
t ) ) (cid:62) J ( θ ( t ) )
− J ( θ (0) ) (cid:62) J
( θ (0) )
(cid:105) (cid:113) Q ( t ) 3 (cid:15) (cid:16)(cid:112) Q g ( θ ( t ) ) (cid:17) (72)
Let γ = 1 − ηq ∗ λ min< 1 .
Combining with Theorem 10 and (62), the above leads to
[Equation]
This implies that
[Equation]
Combining with (70), it implies that for all t ≥ t
[Equation] [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
[Equation]
f ( t ) ( x ) .
Then we have
[Equation]
) which yields
[Equation]
For t ≤ t (cid:15) , we have
[Equation]
(cid:13)(cid:13)(cid:13) g
[Equation]
( s ) ) (cid:13)(cid:13)(cid:13)+
[Equation]
(cid:62) J
[Equation]
(0) ) Q ( s ) (cid:13)(cid:13)(cid:13)
[Equation]
So we can see that there exists a constant C 1 such that (cid:107) δ t
[Table]
(cid:107) 2 ≤ C 1 ˜
d − 1 / 4 .
Then, for t > t (cid:15) , we have
[Table]
[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 .
B. P ROOF OF L EMMA
We will use the following theorem regarding the eigenvalues of random Gaussian matrices:
Theorem 12 (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
[Equation]
and
[Equation]
as well as
[Equation]
Now we assume that (81) and (82) are true.
Then, for any x such that (cid:107) x (cid:107) 2 ≤ 1 ,
[Equation]
[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 ≤
√ ˜ d (cid:13)(cid:13) x l (cid:13)(cid:13)
(cid:13)(cid:13)
α l +1 (cid:13)(cid:13) 2 ≤ 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 (81) and (82) 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]
[Table]
h l − ˜
h l (cid:13)(cid:13)(cid:13)
[Equation]
∀
[Equation] [Equation]
[Equation]
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 .
For α l , we have α L +1 =
˜
α L +1 = 1 , and for all l ≥ 1 ,
[Equation]
(85) 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
[Equation]
2 ≤ 3 ˜ d 1 / 4 , (85) is still true.
By induction, there exists an M 6 > 0 such that (cid:13)(cid:13)
[Equation]
˜
[Equation]
for all l ≥ 1 (note that this is also true for l = L + 1 ).
Thus, if (81) and (82) are true, then for all θ, ˜ θ ∈ B ( θ (0) , C 0 ) , any x such that (cid:107) x (cid:107) 2 ≤ 1 , we have
[Equation]
≤
[Equation]
and for l = 1 , · · · , L , we have
[Equation]
Moreover, for any l = 0 , · · · , L , there is
[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 7 √ n · 4 √ ˜ d (cid:13) (cid:13)(cid:13)
θ − ˜ θ
(cid:13) (cid:13)(cid:13) 2 , so that (cid:13) (cid:13)(cid:13)
J ( θ ) − J (˜ θ ) (cid:13) (cid:13)(cid:13)
F ≤ M
[Equation]
4 √ ˜ d (cid:13) (cid:13)(cid:13) θ − ˜ θ (cid:13) (cid:13)(cid:13)
[Equation]
B. P ROOF OF T HEOREM
Let η 1 = min { η 0 , η ∗ } , where η 0 is deﬁned in Corollary 3 and η ∗ is deﬁned in Theorem 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
Theorem 5, for any δ > 0 , there exists ˜ D > 0 and a constant C such that
[Equation]
By Corollary 3, we have
[Equation]
Summing the above yields
[Equation]
which is the result we want.
B. P ROOF OF T HEOREM
To minimize the regularized risk (12) with gradient descent, the update rule is
[Equation]
We can see that under the new 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 1 2 n (cid:80) n i =1 (cid:104)
θ − θ ∗ , ∇ θ f (0) ( x i ) (cid:105) 2 = 1 2 n (cid:13) (cid:13) ∇ θ f (0) ( X ) (cid:62) ( θ − θ ∗ ) (cid:13) (cid:13) 2 2 .
Thus, there exists T >such that for any t ≥ T ,
[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
[Equation]
Hv (cid:13)(cid:13) 2 =
(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]
By Corollary 3, 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
[Equation]
which implies (14).
B. P ROOF OF T HEOREM
First of all, with some simple linear algebra analysis, we can prove the following proposition:
Proposition 13.
For any positive deﬁnite symmetric matrix H ∈ R n × n , denote its largest and smallest eigenvalues by λ max and λ min .
Then, for any q ∈ R n + and Q = diag ( q 1 , · · · , q n ) , HQ has n positive eigenvalues that are all 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 and vice versa, 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) = 0 as λ v (cid:54) = 0 .
We can prove the other direction similarly.
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 ] .
Now return back to the proof of Theorem 8.
We still use the shorthand (40).
With L 2 penalty, the update rule of the reweighting algorithm for the neural network is:
[Equation]
And the update rule for the linearized neural network is:
[Equation]
First
[Equation]
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 (96) we have
[Equation]
which implies
[Equation]
(99) where ˜ θ ( t ) is some linear interpolation between θ ( t ) and θ (0) .
Our choice of η ensures that ηµ < 1 .
Similar to (48), we can show that 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 11.
Denote
A = ηMR 0 , and let C
= 4 Aηµ in Lemma 11 6 .
By Lemma 11, there exists D 1 such that for all ˜ d ≥ D 1 , with probability at least (1 − δ/ 3) , (47) is true.
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 (99), we have
[Equation]
Like what we have done before, we can show that all non-zero eigenvalues of J ( θ (0) )
Q ( t ) J ( θ (0) )
[Equation]
are eigenvalues of J ( θ (0) )
(cid:62) J ( θ (0) )
Q ( t ) .
This is because for any λ (cid:54) =
0 , if J ( θ (0) )
Q ( t ) J ( θ (0) )
(cid:62) v = λ v , then J ( θ (0) ) (cid:62) J ( θ (0) )
Q ( t ) ( 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 ) .
On the other hand, by Theorem 4,
J ( θ (0) )
(cid:62) J ( θ (0) )
Q ( t ) converges in probability to Θ Q ( t ) whose eigenvalues are all in [0 , λ max ] by Proposition 13.
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 ) J ( θ (0) )
(cid:62) are all in [0 , λ max + λ min ] for all t .
Since η/ (1 − ηµ ) ≤ ( λ min + λ max ) − 1 by our choice of η , we have
[Equation]
On the other hand, we can use (47) 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 the second term, we have (cid:13)(cid:13)(cid:13)
[Equation]
g ( θ ( t ) )
− g ( θ (0) )
(cid:13)(cid:13)(cid:13) 2 + (cid:13)(cid:13)(cid:13) g
( θ (0) )
(cid:13)(cid:13)(cid:13)≤ (cid:13)(cid:13)(cid:13) J (˜ θ ( t ) )
(cid:13)(cid:13)(cid:13) 2 (cid:13)(cid:13)(cid:13) θ ( t ) − θ (0) (cid:13)(cid:13)(cid:13) 2 + R 0 ≤ MC 0 + R(103)
And for the third term, we have η (cid:13)(cid:13)(cid:13) [Equation]
( θ (0) ) (cid:13)(cid:13)(cid:13) F (cid:13)(cid:13)(cid:13) g (
[Equation]
[Equation] ≤ ηMR
[Equation]
= A (104)
Thus, we have
[Equation]
−
[Equation]
(cid:17) (cid:107) a t
[Equation]
So there exists D 4 such that for all ˜ d ≥ D 4 , (cid:107) a t +1 (cid:107) 2 ≤ (cid:0) 1 − ηµ 2 (cid:1) (cid:107) a t (cid:107) 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, by union bound, we have proved that for any δ > 0 , with probability at least (1 − δ ) 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 (103).
Second, let ∆ t
= θ ( t ) lin − θ ( t ) .
Then we have
[Equation]
which implies
[Equation]
By (102), 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] [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]
Given that ˆ R ( f ( t ) linreg ) < (cid:15) for sufﬁciently large t , this also 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] [Equation]
By Theorem 5, we have
[Equation]
Combining Theorem 7 with (113) and (116) derives
[Equation]
Letting ˜
d
→ ∞ leads to the result we need
[Equation]
C A NOTE ON THE PROOFS 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.
C. T HEIR PROBLEMS
Lee et al.
(2019) claimed in their Theorem 2.1 that under the conditions of our Theorem 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]
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 Theorem 5 and also their learning rate in their Theorem 2.1, then (118) 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]
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 (121) 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 (121), 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.
With that being said, in (121) 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 (122) for the NTK formulation.
However, their proof of (122) in incorrect.
Speciﬁcally, the right-hand side of their inequality (S86) is incorrect.
Using the notations in our Appendix B.4, 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 (123).
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 (123).
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)(cid:13) θ − ˜ θ (cid:13)(cid:13)(cid:13) 2 = (cid:12)(cid:12)(cid:12) b l 1 − ˜ 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 + −
[Equation]
=
[Equation]
β ( b l 1 − ˜
[Equation]
l 1 ) (cid:12)(cid:12)(cid:12) .
Moreover, we have W l +
=
˜
W l +1 , so there is α
[Equation]
Then we can lower bound (cid:13)(cid:13)
α l +1 − ˜
α l +1 (cid:13)(cid:13) 2 by
[Equation]
α l +1 −
[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
[Equation]
On the other hand, suppose that claim (123) is true, then (cid:13)(cid:13)
α
[Table]
[Equation]
− ˜ α
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
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
[Equation]
(cid:13)(cid:13)
= 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ﬁnitely small term compared to the ﬁrst term, so we can only prove that
[Equation]
which is different from (123) because it lacks a critical ˜ d − 1 / 2 and thus leads to a contradiction.
Hence, we cannot prove (123) with the ˜ d − 1 / 2 factor, and consequently we cannot prove (122) 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 .
C. O UR FIXES
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 (81) 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]
(129) for any r ∈ (0 , 1 / 2) , though we cannot really prove the O ( ˜ d − 1 / 2 ) bound as originally claimed in (118).
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 Theorem 4 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 smaller bound.
D E XPERIMENT DETAILS AND ADDITIONAL EXPERIMENTS D. E XPERIMENT DETAILS
All experiments are conducted on a Ubuntu 18.04.6 machine with NVIDIA Geforce GTX 1080ti
GPUs.
Each model is trained with one GPU.
On each of Waterbirds and CelebA, we use a ResNet as the model.
The model is trained with SGD with momentum = 0.9.
On Waterbirds the learning rate is 10 − 4 , and on CelebA it is 10 − 3 .
For Group DRO, ν is selected as 0.01 (see the deﬁnition of ν in (3)).
The batch size used for Waterbirds is 128, and for CelebA it is 400.
Data augmentation including random cropping, random horizontal ﬂip and normalization is performed on both datasets.
D. S AMPLE WEIGHTS CONVERGE IN G ROUP DRO
The results in Section 3 require Assumption 1 which states that each sample weight q ( t ) i converges to some positive value as t → ∞ .
Our readers might wonder how strong this assumption is, and whether reweighting algorithms satisfy this assumption in practice.
In this section we empirically demonstrate that for Group DRO, the dynamic reweighting algorithm we experiment on, this assumption is satisﬁed on Waterbirds and CelebA.
Recall that in Section 2.2 we empirically showed that reweighting algorithms could easily overﬁt without regularization.
Here using the same experimental settings, we keep track of the weight of each group g k during training, and we plot the group weight curves in Figure 3.
We also train the models longer (1000 epochs on Waterbirds and 300 epochs on CelebA).
Clearly we can see that as the training accuracy converges to 100%, the group weights also converge to an equilibrium.
Note that q ( t ) i =
g ( t )
k /n k
for all z i ∈ D k , so the sample weights also converge.
[Figure]