summaries: The authors study the problem of improving the worst-group accuracy in the over-parametrized setting. They offer theoretical and empirical evidence that reweighting does not affect the final solution in this setting; in other words, the algorithm would converge to the same interpolators as that of basic ERM. The theoretical analysis is done for linear models, linearized networks, and wide fully connected networks.   The authors also argue that the solution does not change unless a considerable amount of regularization is incorporated, so much so that it affects the training error. Again, some theoretical and empirical evidence is provided.

summaries: This paper theoretically proves the empirical observation that has been made in a lot of recent works regarding the role of importance reweighing in overparameterized deep networks. For squared loss and independent data points, it theoretically proves that reweighing the data points has no effect on the final model learned for linear models and linearized neural networks. This work also analyzes the role of regularization in preventing this behavior and proves that regularization has to be large enough to prevent small training error to have any effect. This work also has supporting empirical results on two datasets- celeba and waterbirds.

summaries: This work studies the problem of training for worst-case subgroup performance. This has been a very prominent research avenue in the past few years and provides a natural approach towards learning fair and "causal" models (that do not perform too badly for any pre-specified subgroup). In practice however there have been some difficulties with existing approaches, and e.g., in the overparameterization regime, these methods do not perform well in terms of  worst-case performance on test sets. This paper provides a theoretical explanation for this fact (using linear models and linearized neural networks) and also illustrates it using numerical experiments on two datasets. Furthermore, the authors assess the role of $L_2$ regularization in training. The practical takeaway is that existing approaches based on reweighting require substantial regularization or early stopping to perform well in terms of test accuracy (in the overparamerized regime).

summaries: This paper mainly focuses on the reweighting algorithms (e.g. Importance Weighting, Group DRO) for the worst-group performance. It tries to theoretically explain why overfitting problems always appears in reweighting algorithms. Specifically, the authors prove that under several conditions (e.g. assumptions for algorithms, wide fully-connected neural networks, squared loss), the worst-group test performance of the reweighting algorithm will converge to the same level as ERM.

abstractText: 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 overfit 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 overfit. Specifically we prove that with reweighting, an overparameterized model always converges to the same ERM interpolator that fits 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 fix 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.

authors:
Zeyuan Allen-Zhu
Yuanzhi Li
Zhao Song

title: A convergence theory for deep learning via overparameterization

authors:
Su Lin Blodgett
Lisa Green
Brendan O’Connor

title: Demographic dialectal variation in social media: A case study of African-American English

venue: In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing,

authors:
Lénaı̈c Chizat
Edouard Oyallon
Francis Bach

title: On lazy training in differentiable programming

venue: Advances in Neural Information Processing Systems,

authors:
Simon Du
Jason Lee
Haochuan Li
Liwei Wang
Xiyu Zhai

title: Gradient descent finds global minima of deep neural networks

authors:
John Duchi
Hongseok Namkoong

title: Learning models with uniform performance via distributionally robust optimization

venue: arXiv preprint arXiv:1810.08750,

authors:
Cynthia Dwork
Moritz Hardt
Toniann Pitassi
Omer Reingold
Richard Zemel

title: Fairness through awareness

venue: In Proceedings of the 3rd innovations in theoretical computer science conference,

authors:
Moritz Hardt
Eric Price
Nati Srebro

title: Equality of opportunity in supervised learning

authors:
Tatsunori Hashimoto
Megha Srivastava
Hongseok Namkoong
Percy Liang

title: Fairness without demographics in repeated loss minimization

venue: International Conference on Machine Learning,

authors:
Dirk Hovy
Anders Søgaard

title: Tagging performance correlates with author age. In Proceedings of the 53rd annual meeting of the Association for Computational Linguistics and the 7th international joint conference on natural language processing

authors:
Arthur Jacot
Franck Gabriel
Clement Hongler

title: Neural tangent kernel: Convergence and generalization in neural networks

authors:
Matt J Kusner
Joshua Loftus
Chris Russell
Ricardo Silva

title: Counterfactual fairness

authors:
Jaehoon Lee
Lechao Xiao
Samuel Schoenholz
Yasaman Bahri
Roman Novak
Jascha SohlDickstein
Jeffrey Pennington

title: Wide neural networks of any depth evolve as linear models under gradient descent

venue: Advances in neural information processing systems,

authors:
Evan Z Liu
Behzad Haghgoo
Annie S Chen
Aditi Raghunathan
Pang Wei Koh
Shiori Sagawa
Percy Liang
Chelsea Finn

title: Just train twice: Improving group robustness without training group information

authors:
Yonatan Oren
Shiori Sagawa
Tatsunori Hashimoto
Percy Liang

title: Distributionally robust language modeling

venue: In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP),

title: Justice as fairness: A restatement

authors:
Shiori Sagawa
Pang Wei Koh
Tatsunori B. Hashimoto
Percy Liang

title: Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization

authors:
Shiori Sagawa
Aditi Raghunathan
Pang Wei Koh
Percy Liang

title: An investigation of why overparameterization exacerbates spurious correlations

venue: Proceedings of the 37th International Conference on Machine Learning,

title: Improving predictive inference under covariate shift by weighting the loglikelihood function

venue: Journal of statistical planning and inference,

title: Gender and dialect bias in youtube’s automatic captions

venue: In Proceedings of the First ACL Workshop on Ethics in Natural Language Processing,

title: Introduction to the non-asymptotic analysis of random matrices

venue: arXiv preprint arXiv:1011.3027,

authors:
Ziyu Xu
Chen Dan
Justin Khim
Pradeep Ravikumar

title: Class-weighted classification: Tradeoffs and robust approaches

authors:
Muhammad Bilal Zafar
Isabel Valera
Manuel Gomez Rodriguez
Krishna P Gummadi

title: Fairness beyond disparate treatment & disparate impact: Learning classification without disparate mistreatment

venue: In Proceedings of the 26th international conference on world wide web,

authors:
Rich Zemel
Yu Wu
Kevin Swersky
Toni Pitassi
Cynthia Dwork

title: Learning fair representations

authors:
Runtian Zhai
Chen Dan
Zico Kolter
Pradeep Ravikumar. Doro

title: Distributional and outlier robust optimization

venue: Proceedings of the 38th International Conference on Machine Learning,

authors: Finally Zhai

title: If σ is Lipschitz and dl →∞ for l = 1

venue: L sequentially,

title: B.4 PROOF OF LEMMA 11 We will use the following theorem regarding the eigenvalues of random Gaussian matrices: Theorem

venue: C A NOTE ON THE PROOFS IN LEE ET AL

heading: 1 INTRODUCTION

text:
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 classification 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 overfit 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 significantly and eventually to a level almost the same as ERM. Some previous work tried to explain why reweighting algorithms can overfit 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 overfitting phenomenon in reweighting algorithms by studying their implicit biases. Specifically, we prove for a family of overparameterized neural networks
that for almost all reweighting algorithms, the model always converges to the same interpolator that fits 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 overfits to the ERM interpolator, so we cannot hope for its worstgroup test performance to be better than ERM. In short, reweighting algorithms always overfit.
Given this pessimistic result, we analyze whether regularization can help mitigate overfitting, as proposed by Sagawa et al. (2020a). We find that a necessary condition for regularization to work is that it considerably lowers the training performance. Specifically, we prove that if the overparameterized model trained by a reweighting algorithm with regularization can still perform almost perfectly on the training set, then overfitting 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 overfitting.
Our results have two important consequences for practice: (i) We should always use large regularization 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.

heading: 1.1 RELATED WORK

text:
Group fairness. Group fairness in machine learning was first 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 overfit (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 specific 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 specifically we use linearized neural networks to approximate such wide neural networks following Lee et al. (2019). There is some criticism of this line of work, e.g. Chizat et al. (2019) argued that infinitely 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.

heading: 2 PRELIMINARIES

text:
Consider a data domain X × Y ⊆ Rd × R that consists of K groups (subdomains)1, where each data point belongs to one of the groups2. We assume that the input space X is a subset of the unit ball of Rd, such that any x ∈ X satisfies ‖x‖2 ≤ 1. We are given a training set {(xi, yi)}ni=1 i.i.d. sampled from some underlying distribution P over X ×Y . Let the K groups beD1, · · · ,DK where each Di is a subset of X × Y . Let Pk(z) = P (z|z ∈ Dk) be the conditional data distribution over Dk, where z = (x, y). Denote X = (x1, · · · ,xn) ∈ Rd×n, and Y = (y1, · · · , yn) ∈ Rn; for any function g : X 7→ R, we overload notation and use g(X) = (g(x1), · · · , g(xn)) . Let the loss function be ` : Y × Y → [0, 1]. In vanilla training, the goal is to minimize the expected risk denoted by R(f ;P ) = Ez∼P [`(f(x), y)], which is done by minimizing the empirical risk R̂(f) = 1n ∑n i=1 `(f(xi), yi).
1We prove our results for Y ⊆ R, but our results can be easily extended to the multi-class scenario Y ⊆ Rm. 2This is the non-overlapping setting. There is also the overlapping setting where groups can overlap with
each other. We focus on the non-overlapping setting in this paper.
For tasks requiring high worst-group performance, the goal is to train a model f : X → Y that performs well over every Pk, which can be achieved by minimizing the worst-group risk defined as
Rmax(f ;P ) = max k=1,··· ,K R(f ;Pk) = max k=1,··· ,K Ez∼P [`(f(x), y)|z ∈ Dk] (1)

heading: 2.1 REWEIGHTING ALGORITHMS

text:
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 zi, and minimize the weighted empirical risk:
R̂q(t)(f) = n∑ i=1 q (t) i `(f(xi), yi) (2)
where q(t) = (q(t)1 , · · · , q (t) n ) and q (t) 1 + · · ·+ q (t) n = 1.
A static reweighting algorithm assigns to each zi = (xi, yi) a fixed weight qi that does not change during training, i.e. q(t)i ≡ qi. A famous example is Importance Weighting (IW, Shimodaira (2000)), in which if zi ∈ Dk and the size of Dk is nk, then qi = (Knk)−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 q1 = · · · = qn = 1/n.
On the other hand, in a dynamic reweighting algorithm, q(t) changes with t. Specifically, 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 /nk for all zi ∈ Dk where g (t) k is the group weight that is updated by
g (t) k ∝ g (t−1) k exp ( νR̂k(f (t−1)) ) (∀k = 1, · · · ,K) (3)
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.

heading: 2.2 REWEIGHTING ALGORITHMS CAN EASILY OVERFIT

text:
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 overfit 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 classifier 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 overfitting in the average test accuracy.
• Regarding the worst-group test accuracy, while the two reweighting algorithms outperform ERM by a large margin at the early epochs, they overfit 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 significantly after around 30 epochs though they are still better than ERM.

heading: 3 IMPLICIT BIASES OF REWEIGHTING ALGORITHMS

text: 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 overfit. All proofs can be found in Appendix B.

heading: 3.1 LINEAR MODELS

text:
We first 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) = 〈θ,x〉, where θ ∈ Rd. In the overparameterized setting, we have d > n. Consider using the squared loss `(ŷ, y) = 12 (ŷ−y) 2, and minimizing the weighted empirical risk with gradient descent:
θ(t+1) = θ(t) − η n∑ i=1 q (t) i ∇θ`(f (t)(xi), yi) (4)
where η > 0 is the learning rate. For a linear model with the squared loss, the update rule is
θ(t+1) = θ(t) − η n∑ i=1 q (t) i xi(f (t)(xi)− yi) (5)
It is a well known result that under the overparameterization setting where d > n, if x1, · · · ,xn are linearly independent, then with a sufficiently small η, a linear model trained by ERM can always converge to an interpolator which fits all training samples (i.e. θ(t) → θ∗ such that 〈θ∗,xi〉 = yi for all i). Here the linear independence is necessary, because otherwise in the extreme case where x1 = x2 but y1 6= y2, the model cannot fit (x1, y1) and (x2, y2) 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 q1, · · · , qn such that for all i, q(t)i → qi as t → ∞. And mini qi = q ∗ > 0.
This assumption avoids the scenario where there is some i such that q(t)i ≈ 0 for all t, in which case the model could never fit zi. 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 x1, · · · ,xn 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, 〈θ∗,xi〉 = yi.
We now make the following key observation regarding the update rule (5): θ(t+1) − θ(t) is a linear combination of x1, · · · ,xn for all t, and thus θ(t) − θ(0) always lies in the linear subspace span(x1, · · · ,xn). Note that this is an n-dimensional linear subspace if x1, · · · ,xn are linearly independent, and by Cramer’s rule, there is exactly one θ̃ in this subspace such that 〈θ̃+θ(0),xi〉 = yi for all i, which implies that θ∗ = θ̃ + θ(0) is unique. Together with Theorem 1, this leads to:
Theorem 2. If x1, · · · ,xn 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 overfit3.
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.

heading: 3.2 LINEARIZED NEURAL NETWORKS

text:
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. (2019). Denote the neural network at time t by f (t)(x) = f(x; θ(t)) which is parameterized by θ(t) ∈ Rp where p is the number of parameters. The linearized neural network of f (t)(x) is defined as
f (t) lin (x) = f (0)(x) + 〈θ(t) − θ(0),∇θf (0)(x)〉 (6)
where we use the shorthand ∇θf (0)(x) := ∇θf(x; θ) ∣∣ θ=θ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 {(xi, yi)}ni=1, we can construct a new training set { ( ∇θf (0)(xi), yi − f (0)(xi) ) }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)(x1), · · · ,∇θf (0)(xn) 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 qi.
Here we are still using the key intuition: θ(t)−θ(0) always lies in the n-dimensional linear subspace span ( ∇θf (0)(x1), · · · ,∇θf (0)(xn) ) . By Cramer’s rule, there is a unique interpolator θ∗ such that
θ∗ − θ(0) ∈ span ( ∇θf (0)(x1), · · · ,∇θf (0)(xn) ) , and θ(t) always converges to that θ∗. Thus, we have essentially proved that for linearized neural networks, reweighting algorithms always overfit.
Now let us delve deeper into the training dynamics of a linearized neural network. Note that ∇θf (t)lin (X) = ∇θf (0)(X) ∈ Rp×n, so the change in the training function value vector is
f (t+1) lin (X)− f (t) lin (X) = −η ∇θf (0)(X)>∇θf (0)(X)Q(t) ∇ŷ`(f (t)lin (X),Y ) (7)
3By overfit, we are saying that the training error of the model trained by the reweighting algorithm will converge to zero, but the worst-group test performance will converge to the same low level as ERM.
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)>∇θf (0)(X)Q(t). Meanwhile, the neural tangent kernel (NTK, Jacot et al. (2018)) is Θ(0)(x,x′) = ∇θf (0)(x)>∇θf (0)(x′) , 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.

heading: 3.3 WIDE FULLY-CONNECTED NEURAL NETWORKS

text:
Now we prove the result for sufficiently wide fully-connected neural networks, which can be approximated by the linearized neural networks. First we define a fully-connected neural network with L hidden layers (we always assume L ≥ 1 so there is at least one hidden layer). Let hl and xl be the pre- and post-activation outputs of layer l, and dl be the width of layer l. Let x0 = x and d0 = d. Define the neural network ash l+1 = W l√ dl xl + βbl
xl+1 = σ(hl+1)
(l = 0, · · · , L) (8)
where σ is a non-linear activation function, W l ∈ Rdl+1×dl and WL ∈ R1×dL . The parameters θ consist of W 0, · · · ,WL and b0, · · · , bL (θ is the concatenation of all flattened weights and biases). The final output of the neural network is f(x) = hL+1. And let the neural network be initialized as{
W l(0) i,j ∼ N (0, 1)
b l(0) j ∼ N (0, 1)
(l = 0, · · · , L− 1) and
{ W
L(0) i,j = 0
b L(0) j ∼ N (0, 1)
(9)
We also need the following assumption for our approximation theorem: Assumption 2. σ is differentiable everywhere, and both σ and σ̇ are Lipschitz.4
Difference from Jacot et al. (2018). 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 modification enables us to accurately approximate the neural network with its linearized counterpart (6), 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 flawed. In Appendix C we will explain what goes wrong in their proofs and how we manage to fix the proofs with our modification.
For our new initialization, we still have the following NTK theorem:
Theorem 4. If σ is Lipschitz and dl →∞ for l = 1, · · · , L sequentially, then Θ(0)(x,x′) converges in probability to a non-degenerated5 deterministic limiting kernel Θ(x,x′).
The kernel Gram matrix Θ = Θ(X,X) ∈ Rn×n is a positive semi-definite 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 dL 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 satisfies 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 d1 = d2 = · · · = dL = 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 ∈ Rd such that ‖x‖2 ≤ 1, with probability at least 1 − δ over random initialization,
sup t≥0 ∣∣∣f (t)lin (x)− f (t)(x)∣∣∣ ≤ Cd̃−1/4 (10) 4f is Lipschitz if there exists a constant L > 0 such that for any x1,x2, |f(x1)−f(x2)| ≤ L ‖x1 − x2‖2. 5Non-degenerated means that Θ(x,x′) depends on x and x′ and is not a constant.
Remark. We can easily extend this theorem to the case where there exists αl > 0 for each of l = 2, · · · , L such that dl/d1 → αl and d1 →∞.
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)(x1), · · · ,∇θf (0)(xn) are linearly independent, then as d̃→∞, for any test point x ∈ Rd such that ‖x‖2 ≤ 1, with probability close to 1 over random initialization,
lim sup t→∞ ∣∣∣f (t)(x)− f (t)ERM(x)∣∣∣ = O(d̃−1/4)→ 0 (11) 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 infinitely 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 sufficiently wide fully-connected neural networks, reweighting algorithms always overfit.
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)(xi) 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.

heading: 4 DOES REGULARIZATION REALLY HELP?

text:
In the previous section, we proved the pessimistic result that reweighting algorithms always overfit, 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 sufficiently enlarged.
Sagawa et al. (2020a) proposed to tackle the overfitting 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 overfitting. 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. Specifically, 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.

heading: 4.1 THEORETICAL ANALYSIS

text:
Consider a reweighting algorithm with sample weights q(t)i . Following Sagawa et al. (2020a), we consider adding L2 penalty to the weighted empirical risk (2):
R̂µ q(t) (f) = n∑ i=1 q (t) i `(f(xi), yi) + µ 2 ∥∥∥θ − θ(0)∥∥∥2 2
(12)
Given that sufficiently wide neural networks can be approximated by linearized ones, we first 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 `(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 M0 > 0 such that ∥∥∇θf (0)(x)∥∥2 ≤ M0 for all ‖x‖2 ≤ 1, ∇θf (0)(x1), · · · ,∇θf (0)(xn) are linearly independent, and the empirical training risk of f (t)linreg satisfies
lim sup t→∞
R̂(f (t)linreg) < , (13)
for some > 0, then for any test point x such that ‖x‖2 ≤ 1 we have
lim sup t→∞ ∣∣∣f (t)linreg(x)− f (t)linERM(x)∣∣∣ = O(√ ). (14) The proof of this theorem also follows the key intuition: we can show that even with the L2 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 sufficiently 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 defined by (8) and (9) and satisfying Assumption 2, and any reweighting algorithm, if d1 = d2 = · · · = dL = d̃, η ≤ η∗,∇θf (0)(x1), · · · ,∇θf (0)(xn) are linearly independent, and the empirical training risk of f (t)reg satisfies
lim sup t→∞
R̂(f (t)reg ) < (15)
for some > 0, then as d̃ → ∞, with probability close to 1 over random initialization, for any test point x such that ‖x‖2 ≤ 1 we have
lim sup t→∞ ∣∣∣f (t)reg (x)− f (t)ERM(x)∣∣∣ = O(d̃−1/4 +√ )→ O(√ ) (16) 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 significantly drop.

heading: 4.2 EMPIRICAL STUDY

text:
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 L2 penalty formulation (12), but we study the L2 weight decay regularization which is most widely used in practice. We repeat each experiment five times with different random seeds and report the 95% confidence 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 overfitting. When the regularization is small,
the training accuracy is still 100% and the algorithm continues to overfit. 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 sufficient.
On CelebA, regularization does help mitigate overfitting, but a useful regularization must be large enough to lower the training accuracy. We observe that Group DRO overfits 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 find 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.

heading: 5 CONCLUSION

text: In this work, we theoretically studied why reweighting algorithms overfit in practice by analyzing their implicit biases. Specifically, we proved the pessimistic result that reweighting algorithms always overfit. 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 overfitting. 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 overfitting affects reweighting algorithms.

heading: REPRODUCIBILITY STATEMENT

text: To guarantee the reproducibility of all our empirical results, in all our experiments we use a fixed 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 sufficient to exactly reproduce our empirical results.

heading: A OTHER REWEIGHTING ALGORITHMS

text:
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 satisfies 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 P . Overall, DRO minimizes the following expected DRO risk:
RD,ρ(θ;P ) = sup Q P {EQ[`(θ;Z)] : D(Q ‖ P ) ≤ ρ} (17)
The expected DRO risk is typically minimized in the following way: for each epoch t, we first find the worst Q that maximizes EQ[`(θ;Z)] and satisfies D(Q ‖ P ) ≤ ρ,Q 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 finite set of training samples {z1, · · · , zn}, so P is always chosen as the empirical distribution, i.e. uniform distribution over z1, · · · , zn. Then, note that Q P , which implies that the support of Q must be a subset of the support of P , which is {z1, · · · , zn}. This means that Q must be a distribution over z1, · · · , zn, 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 fixed α ∈ (0, 1), we let D(Q ‖ 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 first used in Hashimoto et al. (2018) to deal with fairness tasks where the group labels are unknown, where D(Q ‖ P ) = 12 ∫ (dQ/dP − 1)2dP 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 first identification 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.

heading: B PROOFS

text: Notations. In all of the proofs, for a matrix A, we will use ‖A‖2 to denote its spectral norm and ‖A‖F to denote its Frobenius norm.

heading: B.1 PROOF OF THEOREM 1

text: To help our readers understand the proof more easily, we will first prove the result for static reweighting algorithms where q(t)i = qi for all t, and then we will prove the result for dynamic reweighting algorithms that satisfy q(t)i → qi as t→∞.

heading: B.1.1 STATIC REWEIGHTING ALGORITHMS

text:
We first prove the result for all static reweighting algorithms such that mini qi = q∗ > 0. We will use a standard optimization proof technique called smoothness. Denote A = ∑n i=1 ‖xi‖ 2 2. The empirical risk of the linear model f(x) = 〈θ,x〉 is
F (θ) = n∑ i=1 qi(x > i θ − yi)2 (18)
whose Hessian is
∇2θF (θ) = 2 n∑ i=1 qixix > i (19)
So for any unit vector v ∈ Rd, we have (since qi ∈ [0, 1])
v>∇2θF (θ)v = 2 n∑ i=1 qi(x > i v) 2 ≤ 2 n∑ i=1 qi ‖xi‖22 ≤ 2A (20)
which implies that F (θ) is 2A-smooth. Thus, we have the following upper quadratic bound: for any θ1, θ2 ∈ Rd,
F (θ2) ≤ F (θ1) + 〈∇θF (θ1), θ2 − θ1〉+A ‖θ2 − θ1‖22 (21)
Denote g(θ(t)) = √ Q(X>θ(t) − Y ) ∈ Rn where √ Q = diag( √ q1, · · · , √ qn). We can see that∥∥g(θ(t))∥∥2
2 = F (θ(t)), so that ∇F (θ(t)) = 2X
√ Qg(θ(t)). The update rule of a static reweighting
algorithm with gradient descent and the squared loss is:
θ(t+1) = θ(t) − η n∑ i=1 qixi(f (t)(xi)− yi) = θ(t) − ηX √ Qg(θ(t)) (22)
Substituting θ1 and θ2 in (21) with θ(t) and θ(t+1) yields
F (θ(t+1)) ≤ F (θ(t))− 2ηg(θ(t))> √ Q > X>X √ Qg(θ(t)) +A ∥∥∥ηX√Qg(θ(t))∥∥∥2 2
(23)
Since x1, · · · ,xn are linearly independent,X>X is a positive definite matrix. Denote the smallest eigenvalue of X>X by λmin > 0. And ∥∥√Qg(θ(t))∥∥ 2 ≥ √ q∗ ∥∥g(θ(t))∥∥ 2 = √ q∗F (θ(t)), so we have g(θ(t))> √ Q > X>X √ Qg(θ(t)) ≥ q∗λminF (θ(t)). Thus,
F (θ(t+1)) ≤ F (θ(t))− 2ηq∗λminF (θ(t)) +Aη2 ∥∥∥X√Q∥∥∥2
2 ∥∥∥g(θ(t))∥∥∥2 2
≤ F (θ(t))− 2ηq∗λminF (θ(t)) +Aη2 ∥∥∥X√Q∥∥∥2
F F (θ(t))
≤ F (θ(t))− 2ηq∗λminF (θ(t)) +Aη2 ‖X‖2F F (θ (t)) = (1− 2ηq∗λmin +A2η2)F (θ(t))
(24)
Let η0 = q ∗λmin A2 . For any η ≤ η0, we have F (θ (t+1)) ≤ (1 − ηq∗λmin)F (θ(t)) for all t, which implies that limt→∞ F (θ(t)) = 0. Moreover, √ F (θ(t+1)) ≤ (1 − ηq ∗λmin 2 ) √ F (θ(t)) due to√
1− x ≤ 1− x/2.
The convergence in F (θ) implies the convergence in θ. This is because∥∥∥θ(t+1) − θ(t)∥∥∥2 2 = η2 ∥∥∥X√Qg(θ(t))∥∥∥2 2 ≤ η2 ∥∥∥X√Q∥∥∥2 F ∥∥∥g(θ(t))∥∥∥2 2
≤ η2 ‖X‖2F ∥∥∥g(θ(t))∥∥∥2
2 = Aη2F (θ(t))
(25)
which implies that for any η ≤ η0, ∞∑ t=T ∥∥∥θ(t+1) − θ(t)∥∥∥ 2 ≤ √ Aη2 ∞∑ t=T √ F (θ(t)) ≤ 2A q∗λmin √ F (θ(T )) (26) Therefore, limT→∞ ∑∞ t=T ∥∥θ(t+1) − θ(t)∥∥ 2
= 0, which means that θ(t) converges, and it converges to some interpolator.

heading: B.1.2 DYNAMIC REWEIGHTING ALGORITHMS

text:
Now we prove the result for all dynamic reweighting algorithms satisfying Assumption1. By Assumption 1, for any > 0, there exists t such that for all t ≥ t and all i,
q (t) i ∈ (qi − , qi + ) (27)
This is because for all i, there exists ti such that for all t ≥ ti, q(t)i ∈ (qi − , qi + ). Then, we can define t = max{t1, · · · , tn}. Denote the largest and smallest eigenvalues of X>X by λmax and λmin, and because X is full-rank, we have λmin > 0. Select and a fix an such that 0 < < max{ q ∗
3 , (q∗λmin)2
12λmax 2 }, and then t is also fixed. We still denote Q = diag(q1, · · · , qn). When t ≥ t , the update rule of a dynamic reweighting algorithm with gradient descent and the squared loss is:
θ(t+1) = θ(t) − ηXQ(t) (X>θ(t) − Y ) (28) where Q(t) = Q(t), and we use the subscript to indicate that ∥∥∥Q(t) −Q∥∥∥
2 < . Then, note that
we can rewriteQ(t) asQ (t) = √ Q (t) 3 · √ Q for all < q∗/3. This is because qi+ < √ (qi + 3 )qi
and qi − > √ (qi − 3 )qi for all < qi/3, and qi ≥ q∗. Thus, we have
θ(t+1) = θ(t) − ηX √ Q
(t) 3 g(θ
(t)) whereQ(t) = √ Q
(t) 3 ·
√ Q (29)
Again, substituting θ1 and θ2 in (21) with θ(t) and θ(t+1) yields
F (θ(t+1)) ≤ F (θ(t))− 2ηg(θ(t))> √ Q > X>X √ Q
(t) 3 g(θ (t)) +A ∥∥∥∥ηX√Q(t)3 g(θ(t))∥∥∥∥2 2 (30)
Then, note that ∣∣∣∣g(θ(t))>√Q>X>X (√Q(t)3 −√Q) g(θ(t))∣∣∣∣ ≤ ∥∥∥∥√Q>X>X (√Q(t)3 −√Q)∥∥∥∥
2
∥∥∥g(θ(t))∥∥∥2 2
≤ ∥∥∥√Q∥∥∥
2
∥∥X>X∥∥ 2 ∥∥∥∥√Q(t)3 −√Q∥∥∥∥ 2 ∥∥∥g(θ(t))∥∥∥2 2
≤λmax √ 3 F (θ(t))
(31)
where the last step comes from the following fact: for all < qi/3,√ qi + 3 − √ qi ≤ √ 3 and √ qi − √ qi − 3 ≤ √ 3 (32)
And as proved before, we also have g(θ(t))> √ Q > X>X √ Qg(θ(t)) ≥ q∗λminF (θ(t)) (33)
Since ≤ (q ∗λmin)2
12λmax 2 , we have
g(θ(t))> √ Q > X>X √ Q
(t) 3 g(θ
(t)) ≥ ( q∗λmin − λmax √ 3 ) F (θ(t)) ≥ 1
2 q∗λminF (θ(t)) (34)
Thus,
F (θ(t+1)) ≤ F (θ(t))− ηq∗λminF (θ(t)) +Aη2 ∥∥∥∥X√Q(t)3 ∥∥∥∥2
2
∥∥∥g(θ(t))∥∥∥2 2
≤ (1− ηq∗λmin +A2η2(1 + 3 ))F (θ(t)) ≤ (1− ηq∗λmin + 2A2η2)F (θ(t))
(35)
for all < 1/3. Let η0 = q ∗λmin 4A2 . For any η ≤ η0, we have F (θ (t+1)) ≤ (1−ηq∗λmin/2)F (θ(t)) for all t ≥ t , which implies that limt→∞ F (θ(t)) = 0. As before, we can prove that the convergence in F (θ) implies the convergence in θ. Thus, θ converges to some interpolator.

heading: B.2 PROOF OF THEOREM 4

text:
Note that the first 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. (2018)). If σ is Lipschitz and dl →∞ for l = 1, · · · , L sequentially, then for all l = 1, · · · , L, the distribution of a single element of hl converges in probability to a zero-mean Gaussian process of covariance Σl that is defined recursively by:
Σ1(x,x′) = 1
d0 x>x′ + β2
Σl(x,x′) = Ef [σ(f(x))σ(f(x′))] + β2 (36)
where f is sampled from a zero-mean Gaussian process of covariance Σ(l−1).
Now we show that for an infinitely wide neural network with L ≥ 1 hidden layers, Θ(0) converges in probability to the following non-degenerated deterministic limiting kernel
Θ = Ef∼ΣL [σ(f(x))σ(f(x′))] + β2 (37)
Consider the output layer hL+1 = W L√ d̃ σ(hL) + βbL. We can see that for any parameter θi before the output layer,
∇θihL+1 = diag(σ̇(hL)) WL>√ dL ∇θihL = 0 (38)
And for WL and bL, we have
∇WLhL+1 = 1√ dL σ(hL) and ∇bLhL+1 = β (39)
Then we can achieve (37) by the law of large numbers.

heading: B.3 PROOF OF THEOREM 5

text:
We will use the following short-hand in the proof: g(θ(t)) = f (t)(X)− Y J(θ(t)) = ∇θf(X; θ(t)) ∈ Rp×n
Θ(t) = J(θ(t))>J(θ(t))
(40)
For any > 0, there exists t such that for all t ≥ t and all i, q(t)i ∈ (qi − , qi + ). Like what we have done in (29), we can rewriteQ(t) = Q(t) = √ Q (t) 3 · √ Q, whereQ = diag(q1, · · · , qn).
The update rule of a reweighting algorithm with gradient descent and the squared loss for the wide neural network is: θ(t+1) = θ(t) − ηJ(θ(t))Q(t)g(θ(t)) (41) and for t ≥ t , it can be rewritten as
θ(t+1) = θ(t) − ηJ(θ(t)) √ Q
(t) 3
[√ Qg(θ(t)) ] (42)
First, we will prove the following theorem: Theorem 10. There exist constants M > 0 and 0 > 0 such that for all ∈ (0, 0], η ≤ η∗ and any δ > 0, there exist R0 > 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 , there is∥∥∥g(θ(t))∥∥∥ 2 ≤ BtR0 (43)
t∑ j=1 ∥∥∥θ(j) − θ(j−1)∥∥∥ 2 ≤ ηMR0 t∑ j=1 Bj−1 < MBt R0 B − 1
(44)
(ii) For all t ≥ t , we have∥∥∥√Qg(θ(t))∥∥∥ 2 ≤ ( 1− ηq ∗λmin 3 )t−t Bt R0 (45)
t∑ j=t +1 ∥∥∥θ(j) − θ(j−1)∥∥∥ 2 ≤ η √ 1 + 3 MBt R0 t∑ j=t +1 ( 1− ηq ∗λmin 3 )j−t < 3 √
1 + 3 MBt R0 q∗λmin
(46)
Proof. The proof is based on the following lemma: Lemma 11 (Local Lipschitzness of the Jacobian). Under Assumption 2, there is a constant M > 0 such that for any C0 > 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 ‖x‖2 ≤ 1, ∥∥∥∇θf(x; θ)−∇θf(x; θ̃)∥∥∥ 2 ≤ M 4 √ d̃ ∥∥∥θ − θ̃∥∥∥ 2 ‖∇θf(x; θ)‖2 ≤M∥∥∥J(θ)− J(θ̃)∥∥∥ F ≤ M 4 √ d̃ ∥∥∥θ − θ̃∥∥∥ 2
‖J(θ)‖F ≤M
, ∀θ, θ̃ ∈ B(θ(0), C0) (47)
where B(θ(0), R) = {θ : ∥∥θ − θ(0)∥∥
2 < R}.
The proof can be found in Appendix B.4. Note that for any x, f (0)(x) = βbL where bL is sampled from the standard Gaussian distribution. Thus, for any δ > 0, there exists a constant R0 such that with probability at least (1− δ/3) over random initialization,∥∥∥g(θ(0))∥∥∥
2 < R0 (48)
And by Theorem 4, there existsD2 ≥ 0 such that for any d̃ ≥ D2, with probability at least (1−δ/3),∥∥∥Θ−Θ(0)∥∥∥ F ≤ q ∗λmin 3 (49)
Let M be the constant in Lemma 11. Let 0 = (q∗λmin)2 108M4 . Let B = 1 + η ∗M2, and C0 = MB t R0 B−1 + 3 √
1+3 MBt R0 q∗λmin . By Lemma 11, there exists D1 > 0 such that with probability at least (1 − δ/3), for any d̃ ≥ D1, (47) is true for all θ, θ̃ ∈ B(θ(0), C0). 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 = 0 due to (48), and (44) is always true for t = 0. Suppose (43) and (44) are true for t, then for t+ 1 we have∥∥∥θ(t+1) − θ(t)∥∥∥
2 ≤ η ∥∥∥J(θ(t))Q(t)∥∥∥ 2 ∥∥∥g(θ(t))∥∥∥ 2 ≤ η ∥∥∥J(θ(t))Q(t)∥∥∥ F ∥∥∥g(θ(t))∥∥∥ 2
≤ η ∥∥∥J(θ(t))∥∥∥
F ∥∥∥g(θ(t))∥∥∥ 2 ≤MηBtR0
(50)
So (44) is also true for t+ 1. And we also have∥∥∥g(θ(t+1))∥∥∥ 2 = ∥∥∥g(θ(t+1))− g(θ(t)) + g(θ(t))∥∥∥ 2
= ∥∥∥J(θ̃(t))>(θ(t+1) − θ(t)) + g(θ(t))∥∥∥
2 = ∥∥∥−ηJ(θ̃(t))>J(θ(t))Q(t)g(θ(t)) + g(θ(t))∥∥∥
2 ≤ ∥∥∥I − ηJ(θ̃(t))>J(θ(t))Q(t)∥∥∥
2 ∥∥∥g(θ(t))∥∥∥ 2
≤ ( 1 + ∥∥∥ηJ(θ̃(t))>J(θ(t))Q(t)∥∥∥
2 )∥∥∥g(θ(t))∥∥∥ 2
≤ ( 1 + η ∥∥∥J(θ̃(t))∥∥∥
F ∥∥∥J(θ(t))∥∥∥ F )∥∥∥g(θ(t))∥∥∥ 2
≤ (1 + η∗M2) ∥∥∥g(θ(t))∥∥∥
2 ≤ Bt+1R0
(51)
Therefore, (43) and (44) are true for all t ≤ t , which implies that ∥∥√Qg(θ(t ))∥∥ 2 ≤ ∥∥g(θ(t ))∥∥ 2 ≤ Bt R0, so (45) is true for t = t . And (46) is obviously true for t = t . Now, let us prove (ii) by induction. Note that when t ≥ t , we have the alternative update rule (42). If (45) and (46) are true for t, then for t+ 1, there is
∥∥∥θ(t+1) − θ(t)∥∥∥ 2 ≤ η ∥∥∥∥J(θ(t))√Q(t)3 ∥∥∥∥ 2 ∥∥∥√Qg(θ(t))∥∥∥ 2 ≤ η ∥∥∥∥J(θ(t))√Q(t)3 ∥∥∥∥ F ∥∥∥√Qg(θ(t))∥∥∥ 2
≤ η √ 1 + 3 ∥∥∥J(θ(t))∥∥∥
F ∥∥∥√Qg(θ(t))∥∥∥ 2 ≤Mη √ 1 + 3 ( 1− ηq ∗λmin 3 )t−t Bt R0
(52)
So (46) is true for t+ 1. And we also have∥∥∥√Qg(θ(t+1))∥∥∥ 2 = ∥∥∥√Qg(θ(t+1))−√Qg(θ(t)) +√Qg(θ(t))∥∥∥ 2
= ∥∥∥√QJ(θ̃(t))>(θ(t+1) − θ(t)) +√Qg(θ(t))∥∥∥
2 = ∥∥∥−η√QJ(θ̃(t))>J(θ(t))Q(t)g(θ(t)) +√Qg(θ(t))∥∥∥
2 ≤ ∥∥∥∥I − η√QJ(θ̃(t))>J(θ(t))√Q(t)3 ∥∥∥∥
2
∥∥∥√Qg(θ(t))∥∥∥ 2
≤ ∥∥∥∥I − η√QJ(θ̃(t))>J(θ(t))√Q(t)3 ∥∥∥∥
2
( 1− ηq ∗λmin
3
)t R0
(53)
where θ̃(t) is some linear interpolation between θ(t) and θ(t+1). Now we prove that∥∥∥∥I − η√QJ(θ̃(t))>J(θ(t))√Q(t)3 ∥∥∥∥ 2 ≤ 1− ηq ∗λmin 3 (54)
For any unit vector v ∈ Rn, we have v>(I − η √ QΘ √ Q)v = 1− ηv> √ QΘ √ Qv (55)∥∥√Qv∥∥
2 ∈ [ √ q∗, 1], so for any η ≤ η∗, v>(I − η √ QΘ √ Q)v ∈ [0, 1− ηλminq∗], which implies that ∥∥I − η√QΘ√Q∥∥
2 ≤ 1− ηλminq∗. Thus,
∥∥∥I − η√QJ(θ̃(t))>J(θ(t))√Q∥∥∥ 2
≤ ∥∥∥I − η√QΘ√Q∥∥∥
2 + η ∥∥∥√Q(Θ−Θ(0))√Q∥∥∥ 2 + η ∥∥∥√Q(J(θ(0))>J(θ(0))− J(θ̃(t))>J(θ(t)))√Q∥∥∥ 2
≤1− ηλminq∗ + η ∥∥∥√Q(Θ−Θ(0))√Q∥∥∥
F + η ∥∥∥√Q(J(θ(0))>J(θ(0))− J(θ̃(t))>J(θ(t)))√Q∥∥∥ F
≤1− ηλminq∗ + η ∥∥∥Θ−Θ(0)∥∥∥
F + η ∥∥∥J(θ(0))>J(θ(0))− J(θ̃(t))>J(θ(t))∥∥∥ F
≤1− ηλminq∗ + ηq ∗λmin 3 + ηM2 4 √ d̃ (∥∥∥θ(t) − θ(0)∥∥∥ 2 + ∥∥∥θ̃(t) − θ(0)∥∥∥ 2 ) ≤ 1− ηq ∗λmin
2 (56)
for all d̃ ≥ max { D1, D2, ( 12M2C0 q∗λmin )4} , which implies that∥∥∥∥I − η√QJ(θ̃(t))>J(θ(t))√Q(t)3 ∥∥∥∥
2
≤1− ηq ∗λmin
2 + ∥∥∥∥η√QJ(θ̃(t))>J(θ(t))(√Q(t)3 −√Q)∥∥∥∥ 2
≤1− ηq ∗λmin
2 + ηM2
√ 3 ≤ 1− ηq ∗λmin
3 (due to (32))
(57)
for all ≤ 0. Thus, (45) is also true for t+ 1. In conclusion, (45) and (46) are true with probability at least (1− δ) for all d̃ ≥ D̃ = max { D1, D2, ( 12M2C0 q∗λmin )4} .
Returning back to the proof of Theorem 5. Choose and fix an such that < min{ 0, 13 ( q∗λmin 3λmax+q∗λmin )2 }, where 0 is defined by Theorem 10. Then, t is also fixed. There
exists D̃ ≥ 0 such that for any d̃ ≥ D̃, with probability at least (1− δ), Theorem 10 and Lemma 11 are true and ∥∥∥Θ−Θ(0)∥∥∥
F ≤ q
∗λmin
3 (58)
which immediately implies that∥∥∥Θ(0)∥∥∥ 2 ≤ ‖Θ‖2 + ∥∥∥Θ−Θ(0)∥∥∥ F ≤ λmax + q ∗λmin 3 (59)
We still denote B = 1 + η∗M2 and C0 = MB t R0 B−1 + 3 √ 1+3 MBt R0 q∗λmin . Theorem 10 ensures that for all t, θ(t) ∈ B(θ(0), C0). Then we have∥∥∥I − η√QΘ(0)√Q∥∥∥ 2 ≤ ∥∥∥I − η√QΘ√Q∥∥∥ 2 + η ∥∥∥√Q(Θ−Θ(0))√Q∥∥∥ 2
≤ 1− ηλminq∗ + ηq ∗λmin
3 = 1− 2ηq
∗λmin
3
(60)
so it follows that∥∥∥∥I − η√QΘ(0)√Q(t)3 ∥∥∥∥ 2 ≤ ∥∥∥I − η√QΘ(0)√Q∥∥∥ 2 + ∥∥∥∥η√QΘ(0)(√Q(t)3 −√Q)∥∥∥∥ 2
≤ 1− 2ηq ∗λmin
3 + η(λmax +
q∗λmin 3 ) √ 3
(61)
Thus, for all < 13 (
q∗λmin
3λmax+q∗λmin )2 , there is∥∥∥∥I − η√QΘ(0)√Q(t)3 ∥∥∥∥
2
≤ 1− ηq ∗λmin
3 (62)
The update rule of the reweighting algorithm for the linearized neural network is:
θ (t+1) lin = θ (t) lin − ηJ(θ (0))Q(t)glin(θ (t)) (63)
where we use the subscript “lin” to denote the linearized neural network, and with a slight abuse of notion denote glin(θ(t)) = g(θ (t) lin ).
First, let us consider the training dataX . Denote ∆t = glin(θ(t))− g(θ(t)). We have{ glin(θ
(t+1))− glin(θ(t)) = −ηJ(θ(0))>J(θ(0))Q(t)glin(θ(t)) g(θ(t+1))− g(θ(t)) = −ηJ(θ̃(t))>J(θ(t))Q(t)g(θ(t))
(64)
where θ̃(t) is some linear interpolation between θ(t) and θ(t+1). Thus, ∆t+1 −∆t =η [ J(θ̃(t))>J(θ(t))− J(θ(0))>J(θ(0)) ] Q(t)g(θ(t))
− ηJ(θ(0))>J(θ(0))Q(t)∆t (65)
By Lemma 11, we have∥∥∥J(θ̃(t))>J(θ(t))− J(θ(0))>J(θ(0))∥∥∥ F
≤ ∥∥∥∥(J(θ̃(t))− J(θ(0)))> J(θ(t))∥∥∥∥
F
+ ∥∥∥J(θ(0))> (J(θ(t))− J(θ(0)))∥∥∥
F
≤2M2C0d̃−1/4
(66)
which implies that for all t < t , ‖∆t+1‖2 ≤ ∥∥∥[I − ηJ(θ(0))>J(θ(0))Q(t)]∆t∥∥∥ 2 + ∥∥∥η [J(θ̃(t))>J(θ(t))− J(θ(0))>J(θ(0))]Q(t)g(θ(t))∥∥∥ 2
≤ ∥∥∥I − ηJ(θ(0))>J(θ(0))Q(t)∥∥∥
F ‖∆t‖2 + η ∥∥∥J(θ̃(t))>J(θ(t))− J(θ(0))>J(θ(0))∥∥∥ F ∥∥∥g(θ(t))∥∥∥ 2
≤ (1 + ηM2) ‖∆t‖2 + 2ηM 2C0B tR0d̃ −1/4
≤ B ‖∆t‖2 + 2ηM 2C0B tR0d̃ −1/4
(67)
Therefore, we have
B−(t+1) ‖∆t+1‖2 ≤ B −t ‖∆t‖2 + 2ηM 2C0B −1R0d̃ −1/4 (68)
Since ∆0 = 0, it follows that for all t ≤ t ,
‖∆t‖2 ≤ 2tηM 2C0B t−1R0d̃ −1/4 (69)
and particularly we have∥∥∥√Q∆t ∥∥∥ 2 ≤ ‖∆t ‖2 ≤ 2t ηM 2C0B t −1R0d̃ −1/4 (70)
For t ≥ t , we have the alternative update rule (42). Thus,√ Q∆t+1 − √ Q∆t =η √ Q [ J(θ̃(t))>J(θ(t))− J(θ(0))>J(θ(0)) ]√ Q (t) 3 [√ Qg(θ(t)) ] − η √ QJ(θ(0))>J(θ(0)) √ Q
(t) 3 [√ Q∆t ] (71)
LetA = I − η √ QJ(θ(0))>J(θ(0)) √ Q
(t) 3 = I − η
√ QΘ(0) √ Q
(t) 3 . Then, we have
√ Q∆t+1 = A √ Q∆t+η √ Q [ J(θ̃(t))>J(θ(t))− J(θ(0))>J(θ(0)) ]√ Q (t) 3 (√ Qg(θ(t)) ) (72)
Let γ = 1− ηq ∗λmin
3 < 1. Combining with Theorem 10 and (62), the above leads to∥∥∥√Q∆t+1∥∥∥ 2 ≤ ‖A‖2 ∥∥∥√Q∆t∥∥∥ 2 + η ∥∥∥∥√Q [J(θ̃(t))>J(θ(t))− J(θ(0))>J(θ(0))]√Q(t)3 ∥∥∥∥ 2 ∥∥∥√Qg(θ(t))∥∥∥ 2
≤ γ ∥∥∥√Q∆t∥∥∥
2 + η ∥∥∥J(θ̃(t))>J(θ(t))− J(θ(0))>J(θ(0))∥∥∥ F √ 1 + 3 γt−t Bt R0
≤ γ ∥∥∥√Q∆t∥∥∥
2 + 2ηM2C0
√ 1 + 3 γt−t Bt R0d̃ −1/4
(73) This implies that
γ−(t+1) ∥∥∥√Q∆t+1∥∥∥
2 ≤ γ−t ∥∥∥√Q∆t∥∥∥ 2 + 2ηM2C0 √ 1 + 3 γ−1−t Bt R0d̃ −1/4 (74)
Combining with (70), it implies that for all t ≥ t ,∥∥∥√Q∆t∥∥∥ 2 ≤ 2γt−t ηM2C0Bt R0 [ t B −1 + √ 1 + 3 γ−1(t− t ) ] d̃−1/4 (75)
Next, we consider an arbitrary test point x such that ‖x‖2 ≤ 1. Denote δt = f (t) lin (x) − f (t)(x). Then we have { f
(t+1) lin (x)− f (t) lin (x) = −η∇θf(x; θ (0))>J(θ(0))Q(t)glin(θ (t))
f (t+1)(x)− f (t)(x) = −η∇θf(x; θ̃(t))>J(θ(t))Q(t)g(θ(t)) (76)
which yields
δt+1 − δt =η [ ∇θf(x; θ̃(t))>J(θ(t))−∇θf(x; θ(0))>J(θ(0)) ] Q(t)g(θ(t))
− η∇θf(x; θ(0))>J(θ(0))Q(t)∆t (77)
For t ≤ t , we have
‖δt‖2 ≤η t−1∑ s=0 ∥∥∥[∇θf(x; θ̃(s))>J(θ(s))−∇θf(x; θ(0))>J(θ(0))]Q(s)∥∥∥ 2 ∥∥∥g(θ(s))∥∥∥ 2
+ η t−1∑ s=0 ∥∥∥∇θf(x; θ(0))>J(θ(0))Q(s)∥∥∥ 2 ‖∆s‖2
≤η t−1∑ s=0 ∥∥∥∇θf(x; θ̃(s))>J(θ(s))−∇θf(x; θ(0))>J(θ(0))∥∥∥ F ∥∥∥g(θ(s))∥∥∥ 2
+ η t−1∑ s=0 ∥∥∥∇θf(x; θ(0))∥∥∥ 2 ∥∥∥J(θ(0))∥∥∥ F ‖∆s‖2
≤2ηM2C0d̃−1/4 t−1∑ s=0 BsR0 + ηM 2 t−1∑ s=0 (2sηM2C0B s−1R0d̃ −1/4)
(78)
So we can see that there exists a constant C1 such that ‖δt ‖2 ≤ C1d̃−1/4. Then, for t > t , we have
‖δt‖2 − ‖δt ‖2 ≤η t−1∑ s=t ∥∥∥∥[∇θf(x; θ̃(s))>J(θ(s))−∇θf(x; θ(0))>J(θ(0))]√Q(s)3 ∥∥∥∥ 2 ∥∥∥√Qg(θ(s))∥∥∥ 2
+ η t−1∑ s=t ∥∥∥∥∇θf(x; θ(0))>J(θ(0))√Q(s)3 ∥∥∥∥ 2 ∥∥∥√Q∆s∥∥∥ 2
≤2ηM2C0d̃−1/4 √ 1 + 3 t−1∑ s=t γs−t Bt R0
+ ηM2 √ 1 + 3 t−1∑ s=t ( 2γs−t ηM2C0B t R0 [ t B −1 + √ 1 + 3 γ−1(s− t ) ] d̃−1/4 ) (79)
Note that ∑∞ t=0 tγ
t is finite as long as γ ∈ (0, 1). Therefore, there is a constant C such that for any t, ‖δt‖2 ≤ Cd̃−1/4 with probability at least (1− δ) for any d̃ ≥ D̃.

heading: B.4 PROOF OF LEMMA 11

text:
We will use the following theorem regarding the eigenvalues of random Gaussian matrices: Theorem 12 (Corollary 5.35 in Vershynin (2010)). If A ∈ Rp×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(−t2/2),
√ p−√q − t ≤ λmin(A) ≤ λmax(A) ≤ √p+√q + t (80)
By this theorem, and also note that WL is a vector, we can see that for any δ, there exist D̃ > 0 and M1 > 0 such that if d̃ ≥ D̃, then with probability at least (1− δ), for all θ ∈ B(θ(0), C0), we have∥∥W l∥∥ 2 ≤ 3 √ d̃ (∀0 ≤ l ≤ L− 1) and ∥∥WL∥∥ 2 ≤ C0 ≤ 3 4 √ d̃ (81) as well as ∥∥βbl∥∥ 2 ≤M1 √ d̃ (∀l = 0, · · · , L) (82)
Now we assume that (81) and (82) are true. Then, for any x such that ‖x‖2 ≤ 1,∥∥h1∥∥ 2 = ∥∥∥∥ 1√d0W 0x+ βb0 ∥∥∥∥
2
≤ 1√ d0
∥∥W 0∥∥ 2 ‖x‖2 + ∥∥βb0∥∥ 2 ≤ ( 3√
d0 +M1)
√ d̃
∥∥hl+1∥∥ 2 = ∥∥∥∥∥ 1√d̃W lxl + βbl ∥∥∥∥∥
2
≤ 1√ d̃
∥∥W l∥∥ 2 ∥∥xl∥∥ 2 + ∥∥βbl∥∥ 2 (∀l ≥ 1)
∥∥xl∥∥ 2 = ∥∥σ(hl)− σ(0l) + σ(0l)∥∥ 2 ≤ L0 ∥∥hl∥∥ 2 + σ(0) √ d̃ (∀l ≥ 1)
(83)
where L0 is the Lipschitz constant of σ and σ(0l) = (σ(0), · · · , σ(0)) ∈ Rdl . By induction, there exists an M2 > 0 such that ∥∥xl∥∥ 2 ≤M2 √ d̃ and ∥∥hl∥∥ 2 ≤M2 √ d̃ for all l = 1, · · · , L.
Denote αl = ∇hlf(x) = ∇hlhL+1. For all l = 1, · · · , L, we have αl = diag(σ̇(hl))W l>√ d̃ αl+1 where σ̇(x) ≤ L0 for all x ∈ R since σ is L0-Lipschitz, αL+1 = 1 and ∥∥αL∥∥
2 =∥∥∥∥diag(σ̇(hL))WL>√
d̃ ∥∥∥∥ 2 ≤ 3 4 √ d̃ L0. Then, we can easily prove by induction that there exists an
M3 > 1 such that ∥∥αl∥∥
2 ≤ M3/
4 √ d̃ for all l = 1, · · · , L (note that this is not true for L + 1
because αL+1 = 1).
For l = 0, ∇W 0f(x) = 1√d0x 0α1>, so ‖∇W lf(x)‖2 ≤ 1√ d0 ∥∥x0∥∥ 2 ∥∥α1∥∥ 2 ≤ 1√ d0 M3/ 4 √ d̃. And for any l = 1, · · · , L, ∇W lf(x) = 1√ d̃ xlαl+1, so ‖∇W lf(x)‖2 ≤ 1√ d̃ ∥∥xl∥∥ 2 ∥∥αl+1∥∥ 2 ≤M2M3.
(Note that if M3 > 1, then ∥∥αL+1∥∥
2 ≤ M3; and since d̃ ≥ 1, there is ∥∥αl∥∥ 2 ≤ M3 for l ≤ L.)
Moreover, for l = 0, · · · , L,∇blf(x) = βαl+1, so ‖∇blf(x)‖2 ≤ βM3. Thus, if (81) and (82) are true, then there exists an M4 > 0, such that ‖∇θf(x)‖2 ≤M4/ √ n. And since ‖xi‖2 ≤ 1 for all i, so ‖J(θ)‖F ≤M4.
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
∥∥∥h1 − h̃1∥∥∥ 2 = ∥∥∥∥ 1√d0 (W 0 − W̃ 0)x+ β(b0 − b̃0) ∥∥∥∥
2
≤ 1√ d0 ∥∥∥W 0 − W̃ 0∥∥∥ 2 ‖x‖2 + β ∥∥∥b0 − b̃0∥∥∥ 2 ≤ ( 1√ d0 + β )∥∥∥θ − θ̃∥∥∥ 2∥∥∥hl+1 − h̃l+1∥∥∥
2 = ∥∥∥∥∥ 1√d̃W l(xl − x̃l) + 1√d̃ (W l − W̃ l)x̃l + β(bl − b̃l) ∥∥∥∥∥
2
≤ 1√ d̃
∥∥W l∥∥ 2 ∥∥xl − x̃l∥∥ 2
+ 1√ d̃ ∥∥∥W l − W̃ l∥∥∥ 2 ∥∥x̃l∥∥ 2 + β ∥∥∥bl − b̃l∥∥∥ 2
≤ 3 ∥∥xl − x̃l∥∥
2 + (M2 + β) ∥∥∥θ − θ̃∥∥∥ 2
(∀l ≥ 1)∥∥xl − x̃l∥∥ 2 = ∥∥∥σ(hl)− σ(h̃l)∥∥∥
2 ≤ L0 ∥∥∥hl − h̃l∥∥∥ 2
(∀l ≥ 1)
(84)
By induction, there exists an M5 > 0 such that ∥∥xl − x̃l∥∥
2 ≤M5 ∥∥∥θ − θ̃∥∥∥ 2 for all l.
For αl, we have αL+1 = α̃L+1 = 1, and for all l ≥ 1,
∥∥αl − α̃l∥∥ 2 = ∥∥∥∥∥diag(σ̇(hl))W l>√d̃ αl+1 − diag(σ̇(h̃l))W̃ l>√ d̃ α̃l+1 ∥∥∥∥∥ 2
≤ ∥∥∥∥∥diag(σ̇(hl))W l>√d̃ (αl+1 − α̃l+1) ∥∥∥∥∥
2
+ ∥∥∥∥∥diag(σ̇(hl)) (W l − W̃ l)>√d̃ α̃l+1 ∥∥∥∥∥
2
+ ∥∥∥∥∥diag((σ̇(hl)− σ̇(h̃l)))W̃ l>√d̃ α̃l+1 ∥∥∥∥∥
2 ≤ 3L0 ∥∥αl+1 − α̃l+1∥∥ 2 + ( M3L0d̃ −1/2 + 3M3M5L1d̃ −1/4 )∥∥∥θ − θ̃∥∥∥ 2
(85) where L1 is the Lipschitz constant of σ̇. Particularly, for l = L, though α̃L+1 = 1, since ∥∥∥W̃L∥∥∥
2 ≤ 3d̃1/4, (85) is still true. By induction, there exists anM6 > 0 such that ∥∥αl − α̃l∥∥
2 ≤ M6 4 √ d̃ ∥∥∥θ − θ̃∥∥∥ 2
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), C0), any x such that ‖x‖2 ≤ 1, we have∥∥∥∇W 0f(x)−∇W̃ 0 f̃(x)∥∥∥ 2 = 1√ d0 ∥∥xα1> − xα̃1>∥∥ 2
≤ 1√ d0
∥∥α1 − α̃1∥∥ 2
≤ 1√ d0 M6 4 √ d̃ ∥∥∥θ − θ̃∥∥∥ 2
(86)
and for l = 1, · · · , L, we have
∥∥∥∇W lf(x)−∇W̃ l f̃(x)∥∥∥ 2 = 1√ d̃ ∥∥xlαl+1> − x̃lα̃l+1>∥∥ 2
≤ 1√ d̃
(∥∥xl∥∥ 2 ∥∥αl+1 − α̃l+1∥∥ 2 + ∥∥xl − x̃l∥∥ 2 ∥∥α̃l+1∥∥ 2 ) ≤ ( M2M6
4 √ d̃ + M5M3√ d̃ )∥∥∥θ − θ̃∥∥∥ 2
(87)
Moreover, for any l = 0, · · · , L, there is∥∥∥∇blf(x)−∇b̃l f̃(x)∥∥∥ 2 = β ∥∥αl+1 − α̃l+1∥∥ 2 ≤ βM6 4 √ d̃ ∥∥∥θ − θ̃∥∥∥ 2
(88)
Overall, we can see that there exists a constant M7 > 0 such that ∥∥∥∇θf(x)−∇θ̃f̃(x)∥∥∥
2 ≤
M7 √ n· 4 √ d̃ ∥∥∥θ − θ̃∥∥∥ 2 , so that ∥∥∥J(θ)− J(θ̃)∥∥∥ F ≤ M7 4 √ d̃ ∥∥∥θ − θ̃∥∥∥ 2 .

heading: B.5 PROOF OF THEOREM 6

text:
Let η1 = min{η0, η∗}, where η0 is defined in Corollary 3 and η∗ is defined 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 sup t≥0 ∣∣∣f (t)lin (x)− f (t)(x)∣∣∣ ≤ Cd̃−1/4 sup t≥0
∣∣∣f (t)linERM(x)− f (t)ERM(x)∣∣∣ ≤ Cd̃−1/4 (89) By Corollary 3, we have
lim t→∞ ∣∣∣f (t)lin (x)− f (t)linERM(x)∣∣∣ = 0 (90) Summing the above yields
lim sup t→∞ ∣∣∣f (t)(x)− f (t)ERM(x)∣∣∣ ≤ 2Cd̃−1/4 (91) which is the result we want.

heading: B.6 PROOF OF THEOREM 7

text:
To minimize the regularized risk (12) with gradient descent, the update rule is
θ(t+1) = θ(t) − η n∑ i=1 q (t) i ∇θ`(f (t)(xi), yi)− ηµ(θ(t) − θ(0)) (92)
We can see that under the new rule, θ(t) − θ(0) ∈ span(∇θf (0)(x1), · · · ,∇θf (0)(xn)) is still true for all t. Let θ∗ be the interpolator in span(∇θf (0)(x1), · · · ,∇θf (0)(xn)), then the empirical risk of θ is 12n ∑n i=1〈θ − θ∗,∇θf (0)(xi)〉2 = 1 2n ∥∥∇θf (0)(X)>(θ − θ∗)∥∥22. Thus, there exists T > 0 such that for any t ≥ T , ∥∥∥∇θf (0)(X)>(θ(t) − θ∗)∥∥∥2
2 ≤ 2n (93)
Let the smallest singular value of 1√ n ∇θf (0)(X) be smin, and we have smin > 0. Note that the column space of ∇θf (0)(X) is exactly span(∇θf (0)(x1), · · · ,∇θf (0)(xn)). Define H ∈ Rp×n
such that its columns form an orthonormal basis of this subspace, then there existsG ∈ Rn×n such that ∇θf (0)(X) = HG, and the smallest singular value of 1√nG is also s min. Since θ(t) − θ(0) is also in this subspace, there exists v ∈ Rn such that θ(t) − θ∗ = Hv. Then we have √
2n ≥∥∥G>H>Hv∥∥ 2 = ∥∥G>v∥∥ 2 . Thus, ‖v‖2 ≤ √ 2
smin , which implies∥∥∥θ(t) − θ∗∥∥∥ 2 ≤ √ 2 smin (94)
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 ‖x‖2 ≤ 1,
|f (t)linreg(x)− f (t) linERM(x)| = |〈θ (t) − θ∗,∇θf (0)(x)〉| ≤ M0 √ 2
smin (95)
which implies (14).

heading: B.7 PROOF OF THEOREM 8

text:
First of all, with some simple linear algebra analysis, we can prove the following proposition: Proposition 13. For any positive definite symmetric matrix H ∈ Rn×n, denote its largest and smallest eigenvalues by λmax and λmin. Then, for any q ∈ Rn+ and Q = diag(q1, · · · , qn), HQ has n positive eigenvalues that are all in [mini qi · λmin,maxi qi · λmax].
Proof. H is a positive definite symmetric matrix, so there existsA ∈ Rn×n such thatH = A>A, and A is full-rank. First, any eigenvalue of AQA> is also an eigenvalue of A>AQ and vice versa, because for any eigenvalue λ of AQA> we have some v 6= 0 such that AQA>v = λv. Multiplying both sides by A> on the left yields A>AQ(A>v) = λ(A>v) which implies that λ is also an eigenvalue of A>AQ because A>v 6= 0 as λv 6= 0. We can prove the other direction similarly.
Second, by condition we know that the eigenvalues ofA>A are all in [λmin, λmax] where λmin > 0, which implies for any unit vector v, v>A>Av ∈ [λmin, λmax], which is equivalent to ‖Av‖2 ∈ [ √ λmin, √ λmax]. Thus, we have v>A>QAv ∈ [λmin mini qi, λmax maxi qi], which implies that the eigenvalues ofA>QA are all in [λmin mini qi, λmax maxi qi].
Thus, the eigenvalues ofHQ = A>AQ are all in [λmin mini qi, λmax maxi qi].
Now return back to the proof of Theorem 8. We still use the shorthand (40). With L2 penalty, the update rule of the reweighting algorithm for the neural network is:
θ(t+1) = θ(t) − ηJ(θ(t))Q(t)g(θ(t))− ηµ(θ(t) − θ(0)) (96)
And the update rule for the linearized neural network is:
θ (t+1) lin = θ (t) lin − ηJ(θ (0))Q(t)glin(θ (t))− ηµ(θ(t)lin − θ (0)) (97)
First, we need to prove that there exists D0 such that for all d̃ ≥ D0, supt≥0 ∥∥θ(t) − θ(0)∥∥ 2 is bounded with high probability. Denote at = θ(t) − θ(0). By (96) we have
at+1 =(1− ηµ)at − η[J(θ(t))− J(θ(0))]Q(t)g(θ(t)) − ηJ(θ(0))Q(t)[g(θ(t))− g(θ(0))]− ηJ(θ(0))Q(t)g(θ(0))
(98)
which implies ‖at+1‖2 ≤ ∥∥∥(1− ηµ)I − ηJ(θ(0))Q(t)J(θ̃(t))>∥∥∥
2 ‖at‖2 + η ∥∥∥J(θ(t))− J(θ(0))∥∥∥
F ∥∥∥g(θ(t))∥∥∥ 2 + η ∥∥∥J(θ(0))∥∥∥ F ∥∥∥g(θ(0))∥∥∥ 2
(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 R0 > 0 such that with
probability at least (1 − δ/3), ∥∥g(θ(0))∥∥
2 < R0. Let M be as defined in Lemma 11. Denote
A = ηMR0, and let C0 = 4Aηµ in Lemma 11 6. By Lemma 11, there exists D1 such that for all d̃ ≥ D1, with probability at least (1− δ/3), (47) is true. Now we prove by induction that ‖at‖2 < C0. It is true for t = 0, so we need to prove that if ‖at‖2 < C0, then ‖at+1‖2 < C0. For the first term on the right-hand side of (99), we have∥∥∥(1− ηµ)I − ηJ(θ(0))Q(t)J(θ̃(t))>∥∥∥
2 ≤(1− ηµ) ∥∥∥∥I − η1− ηµJ(θ(0))Q(t)J(θ(0))> ∥∥∥∥
2 + η ∥∥∥J(θ(0))∥∥∥
F ∥∥∥J(θ̃(t))− J(θ(0))∥∥∥ F
(100)
Like what we have done before, we can show that all non-zero eigenvalues of J(θ(0))Q(t)J(θ(0))> are eigenvalues of J(θ(0))>J(θ(0))Q(t). This is because for any λ 6= 0, if J(θ(0))Q(t)J(θ(0))>v = λv, then J(θ(0))>J(θ(0))Q(t)(J(θ(0))>v) = λ(J(θ(0))>v), and J(θ(0))>v 6= 0 since λv 6= 0, so λ is also an eigenvalue of J(θ(0))>J(θ(0))Q(t). On the other hand, by Theorem 4, J(θ(0))>J(θ(0))Q(t) converges in probability to ΘQ(t) whose eigenvalues are all in [0, λmax] by Proposition 13. So there exists D2 such that for all d̃ ≥ D2, with probability at least (1 − δ/3), the eigenvalues of J(θ(0))Q(t)J(θ(0))> are all in [0, λmax + λmin] for all t. Since η/(1− ηµ) ≤ (λmin + λmax)−1 by our choice of η, we have∥∥∥∥I − η1− ηµJ(θ(0))Q(t)J(θ(0))> ∥∥∥∥ 2 ≤ 1 (101) On the other hand, we can use (47) since ‖at‖2 < C0, so ∥∥J(θ(0))∥∥
F ∥∥∥J(θ̃(t))− J(θ(0))∥∥∥ F ≤
M2 4 √ d̃ C0. Therefore, there exists D3 such that for all d̃ ≥ D3,∥∥∥(1− ηµ)I − ηJ(θ(0))Q(t)J(θ̃(t))>∥∥∥
2 ≤ 1− ηµ 2 (102)
For the second term, we have∥∥∥g(θ(t))∥∥∥ 2 ≤ ∥∥∥g(θ(t))− g(θ(0))∥∥∥ 2 + ∥∥∥g(θ(0))∥∥∥ 2
≤ ∥∥∥J(θ̃(t))∥∥∥
2 ∥∥∥θ(t) − θ(0)∥∥∥ 2 +R0 ≤MC0 +R0 (103)
And for the third term, we have η ∥∥∥J(θ(0))∥∥∥
F ∥∥∥g(θ(0))∥∥∥ 2 ≤ ηMR0 = A (104)
Thus, we have
‖at+1‖2 ≤ (
1− ηµ 2
) ‖at‖2 +
ηM(MC0 +R0) 4 √ d̃ +A (105)
So there exists D4 such that for all d̃ ≥ D4, ‖at+1‖2 ≤ ( 1− ηµ2 ) ‖at‖2 + 2A. This shows that if ‖at‖2 < C0 is true, then ‖at+1‖2 < C0 will also be true. In conclusion, by union bound, we have proved that for any δ > 0, with probability at least (1− δ) for all d̃ ≥ D0 = max{D1, D2, D3, D4}, ∥∥θ(t) − θ(0)∥∥ 2 < C0 is true for all t. This also implies
that for C1 = MC0 +R0, we have ∥∥g(θ(t))∥∥
2 ≤ C1 for all t by (103).
6Note that Lemma 11 only depends on the network structure and does not depend on the update rule, so we can use this lemma here.
Second, let ∆t = θ (t) lin − θ(t). Then we have
∆t+1 −∆t = η(J(θ(t))Q(t)g(θ(t))− J(θ(0))Q(t)glin(θ(t))− µ∆t) (106) which implies
∆t+1 = [ (1− ηµ)I − ηJ(θ(0))Q(t)J(θ̃(t))> ] ∆t + η(J(θ (t))− J(θ(0)))Q(t)g(θ(t)) (107)
By (102), with probability at least (1− δ) for all d̃ ≥ D0, we have ‖∆t+1‖2 ≤ ∥∥∥(1− ηµ)I − ηJ(θ(0))Q(t)J(θ̃(t))>∥∥∥
2 ‖∆t‖2 + η ∥∥∥J(θ(t))− J(θ(0))∥∥∥ F ∥∥∥g(θ(t))∥∥∥ 2
≤ (
1− ηµ 2
) ‖∆t‖2 + η
M 4 √ d̃ C0C1
(108)
Again, as ∆0 = 0, we can prove by induction that for all t,
‖∆t‖2 < 2MC0C1 µ d̃−1/4 (109)
For any test point x such that ‖x‖2 ≤ 1, we have∣∣∣f (t)reg (x)− f (t)linreg(x)∣∣∣ = ∣∣∣f(x; θ(t))− flin(x; θ(t)lin )∣∣∣ ≤ ∣∣∣f(x; θ(t))− flin(x; θ(t))∣∣∣+ ∣∣∣flin(x; θ(t))− flin(x; θ(t)lin )∣∣∣
≤ ∣∣∣f(x; θ(t))− flin(x; θ(t))∣∣∣+ ∥∥∥∇θf(x; θ(0))∥∥∥
2 ∥∥∥θ(t) − θ(t)lin ∥∥∥ 2
≤ ∣∣∣f(x; θ(t))− flin(x; θ(t))∣∣∣+M ‖∆t‖2
(110)
For the first term, note that{ f(x; θ(t))− f(x; θ(0)) = ∇θf(x; θ̃(t))(θ(t) − θ(0))
flin(x; θ (t))− flin(x; θ(0)) = ∇θf(x; θ(0))(θ(t) − θ(0))
(111)
where θ̃(t) is some linear interpolation between θ(t) and θ(0). Since f(x; θ(0)) = flin(x; θ(0)),∣∣∣f(x; θ(t))− flin(x; θ(t))∣∣∣ ≤ ∥∥∥∇θf(x; θ̃(t))−∇θf(x; θ(0))∥∥∥ 2 ∥∥∥θ(t) − θ(0)∥∥∥ 2 ≤ M 4 √ d̃ C20 (112) Thus, we have shown that for all d̃ ≥ D0, with probability at least (1− δ) for all t and all x,∣∣∣f (t)reg (x)− f (t)linreg(x)∣∣∣ ≤ (MC20 + 2M2C0C1µ ) d̃−1/4 = O(d̃−1/4) (113)
Given that R̂(f (t)linreg) < for sufficiently large t, this also implies that∣∣∣R̂(f (t)linreg)− R̂(f (t)reg )∣∣∣ = O(d̃−1/4√ + d̃−1/2) (114) So for a fixed , there exists D > 0 such that for all d ≥ D, for sufficiently large t,
R̂(f (t)reg ) < ⇒ R̂(f (t) linreg) < 2 (115)
By Theorem 5, we have
sup t≥0 ∣∣∣f (t)linERM(x)− f (t)ERM(x)∣∣∣ = O(d̃−1/4) (116) Combining Theorem 7 with (113) and (116) derives
lim sup t→∞ ∣∣∣f (t)reg (x)− f (t)ERM(x)∣∣∣ = O(d̃−1/4 +√ ) (117) Letting d̃→∞ leads to the result we need.

heading: C A NOTE ON THE PROOFS IN LEE ET AL. (2019)

text:
We have mentioned that the proofs in Lee et al. (2019), particularly the proofs of their Theorem 2.1 and Lemma 1 in their Appendix G, are flawed. In order to fix 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 fix 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.

heading: C.1 THEIR PROBLEMS

text:
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 sufficiently wide fully-connected neural network and the output of its linearized neural network at any test point x can be uniformly bounded by
sup t≥0 ∣∣∣f (t)(x)− f (t)lin (x)∣∣∣ ≤ Cd̃−1/2 (claimed) (118) where they used the original NTK formulation and initialization in Jacot et al. (2018):h l+1 = W l√ dl xl + βbl
xl+1 = σ(hl+1)
and
{ W
l(0) i,j ∼ N (0, 1)
b l(0) i ∼ N (0, 1)
(∀l = 0, · · · , L) (119)
where x0 = x and f(x) = hL+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:{
hl+1 = W lxl + βbl xl+1 = σ(hl+1) and
W l(0) i,j ∼ N (0, 1 dl )
b l(0) i ∼ N (0, 1)
(∀l = 0, · · · , L) (120)
See their Appendix F for the definition 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 η′ = η/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 η′ to the standard formulation is equivalent to applying the following learning rates
ηlW = dl dmax η and ηlb = 1 dmax η (121)
to W l and bl of the NTK formulation, where dmax = max{d0, · · · , dL}. 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 first problem is that ηlb = 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 fixed by modifying their standard formulation ashl+1 = W lxl+β √ dlb l where bl(0)i ∼ N (0, d −1 l ). The real problem that is non-trivial to fix is that by (121), there is
η0W = d0 dmax η. However, note that d0 is a constant since it is the dimension of the input space, while dmax goes to infinity. With that being said, in (121) they were essentially using a very small learning rate for the first layer W 0 but a normal learning rate for the rest of the layers, which definitely 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 constantM > 0, for any two models with the parameters θ and θ̃ such that θ, θ̃ ∈ B(θ(0), C0) for some constant C0, there is∥∥∥J(θ)− J(θ̃)∥∥∥
F ≤ M√
d̃
∥∥∥θ − θ̃∥∥∥ 2
(claimed) (122)
Note that the original claim in their paper was ∥∥∥J(θ)− J(θ̃)∥∥∥
F ≤M
√ d̃ ∥∥∥θ − θ̃∥∥∥
2 . This is because
they were proving this result for their standard formulation. Compared to the standard formulation, in the NTK formulation θ is √ d̃ times larger, while the Jacobian J(θ) is √ d̃ times smaller. This is also why here we have θ, θ̃ ∈ B(θ(0), C0) instead of θ, θ̃ ∈ B(θ(0), C0d̃−1/2) for the NTK formulation. Therefore, equivalently they were claiming (122) for the NTK formulation.
However, their proof of (122) in incorrect. Specifically, the right-hand side of their inequality (S86) is incorrect. Using the notations in our Appendix B.4, their (S86) essentially claimed that∥∥αl − α̃l∥∥ 2 ≤ M√
d̃
∥∥∥θ − θ̃∥∥∥ 2
(claimed) (123)
for any θ, θ̃ ∈ B(θ(0), C0), where αl = ∇hlhL+1 and α̃l is the same gradient for the second model. Note that their (S86) does not have the √ d̃ in the denominator which appears in (123). This
is because for their standard formulation, θ is √ d̃ times smaller than the original NTK formulation,
while ∥∥αl∥∥
2 has the same order in the two formulations because all hl 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 bl1 for some 1 ≤ l < L. Then, ∥∥∥θ − θ̃∥∥∥ 2 = ∣∣∣bl1 − b̃l1∣∣∣. We can see that hl+1 and h̃l+1 only differ in the first element, and∣∣∣hl+11 − h̃l+11 ∣∣∣ = ∣∣∣β(bl1 − b̃l1)∣∣∣. Moreover, we have W l+1 = W̃ l+1, so there is
αl+1 − α̃l+1 =diag(σ̇(hl+1))W l+1>√ d̃ αl+2 − diag(σ̇(h̃l+1))W̃ l+1>√ d̃ α̃l+2
= [ diag(σ̇(hl+1))− diag(σ̇(h̃l+1)) ] W l+1>√ d̃ αl+2
+ diag(σ̇(h̃l+1)) W l+1>√
d̃ (αl+2 − α̃l+2)
(124)
Then we can lower bound ∥∥αl+1 − α̃l+1∥∥
2 by
∥∥αl+1 − α̃l+1∥∥ 2 ≥ ∥∥∥∥∥[diag(σ̇(hl+1))− diag(σ̇(h̃l+1))] W l+1>√d̃ αl+2 ∥∥∥∥∥
2
− ∥∥∥∥∥diag(σ̇(h̃l+1))W l+1>√d̃ (αl+2 − α̃l+2) ∥∥∥∥∥
2
(125)
The first term on the right-hand side is equal to ∣∣∣[σ̇(hl+11 )− σ̇(h̃l+11 )] 〈W l+11 /√d̃,αl+2〉∣∣∣ where
W l+11 is the first row of W l+1. We know that ∥∥W l+11 ∥∥2 = Θ(√d̃) with high probability as its
elements are sampled from N (0, 1), and in their (S85) they claimed that ∥∥αl+2∥∥
2 = O(1), which
is true. In addition, they assumed that σ̇ is Lipschitz. Hence, we can see that∥∥∥∥∥[diag(σ̇(hl+1))− diag(σ̇(h̃l+1))] W l+1>√d̃ αl+2 ∥∥∥∥∥
2
= O (∣∣∣hl+11 − h̃l+11 ∣∣∣) = O (∥∥∥θ − θ̃∥∥∥
2 ) (126)
On the other hand, suppose that claim (123) is true, then ∥∥αl+2 − α̃l+2∥∥
2 = O
( d̃−1/2 ∥∥∥θ − θ̃∥∥∥ 2 ) .
Then we can see that the second term on the right-hand side is O ( d̃−1/2 ∥∥∥θ − θ̃∥∥∥ 2 ) because∥∥W l+1∥∥
2 = O(
√ d̃) and σ̇(x) is bounded by a constant as σ is Lipschitz. Thus, for a very large d̃,
the second-term is an infinitely small term compared to the first term, so we can only prove that∥∥αl+1 − α̃l+1∥∥ 2 = O (∥∥∥θ − θ̃∥∥∥
2
) (127)
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 √ d̃ in the denominator on the right-hand side. As a result, their Lemma 1 and Theorem 2.1 cannot be proved without this critical d̃−1/2. Similarly, we can also construct a counterexample where θ and θ̃ only differ in the first row of some W l.

heading: C.2 OUR FIXES

text:
Regarding Problem 1, we can still use anO(1) learning rate for the first layer in the NTK formulation given that ‖x‖2 ≤ 1. This is because for the first layer, we have
∇W 0f(x) = 1√ d0 x0α1> = 1√ d0 xα1> (128)
For all l ≥ 1, we have ∥∥xl∥∥
2 = O(d̃1/2). However, for l = 0, we instead have ∥∥x0∥∥ 2
= O(1). Thus, we can prove that the norm of ∇W 0f(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 first 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 WL(0)i,j ∼ N (0, 1) to the zero initialization WL(0)i,j = 0. Now we show how this modification solves Problem 2.
The main consequence of changing the initialization of the last layer is that (81) becomes different: instead of ∥∥WL∥∥ 2 ≤ 3 √ d̃, we now have ∥∥WL∥∥ 2 ≤ C0 ≤ 3 4 √ d̃. In fact, for any r ∈ (0, 1/2), we
can prove that ∥∥WL∥∥
2 ≤ 3d̃r for sufficiently large d̃. In our proof we choose r = 1/4. Consequently, instead of ∥∥αl∥∥
2 ≤M3, we can now prove that ∥∥αl∥∥ 2 ≤M3d̃r−1/2 for all l ≤ L by
induction. So now we can prove ∥∥αl − α̃l∥∥
2 = O
( d̃r−1/2 ∥∥∥θ − θ̃∥∥∥ 2 ) instead of O (∥∥∥θ − θ̃∥∥∥ 2 ) ,
because • For l < L, we now have ∥∥αl+1∥∥
2 = 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 ∥∥αL+1∥∥ 2 = 1, note that ∥∥WL∥∥ 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
sup t≥0 ∣∣∣f (t)(x)− f (t)lin (x)∣∣∣ ≤ Cd̃r−1/2 (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′) still depends on x and x′), and under the overparameterized setting where dL n, chances are high that Θ is full-rank. Hence, we can still assume that λmin > 0.
As a final remark, one key reason why we need to initialize WL as zero is that the dimension of the output space (i.e. the dimension of hL+1) is finite, and in our case it is 1. Suppose we allow the dimension of hL+1 to be d̃which goes to infinity, then using the same proof techniques, for the NTK formulation we can prove that supt ∥∥∥hL+1(t) − hL+1(t)lin ∥∥∥ 2 ≤ C, i.e. the gap between two vectors of infinite dimension is always bounded by a finite constant. This is the approximation theorem we need for the infinite-dimensional output space. However, when the dimension of the output space is finite, supt ∥∥∥hL+1(t) − hL+1(t)lin ∥∥∥ 2 ≤ C no longer suffices, so we need to decrease the order of the norm of WL in order to obtain a smaller bound.

heading: D EXPERIMENT DETAILS AND ADDITIONAL EXPERIMENTS

heading: D.1 EXPERIMENT DETAILS

text: 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 ResNet18 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 definition of ν in (3)). The batch size used for Waterbirds is 128, and for CelebA it is 400. Data augmentation including random cropping, random horizontal flip and normalization is performed on both datasets.

heading: D.2 SAMPLE WEIGHTS CONVERGE IN GROUP DRO

text:
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 satisfied on Waterbirds and CelebA.
Recall that in Section 2.2 we empirically showed that reweighting algorithms could easily overfit without regularization. Here using the same experimental settings, we keep track of the weight of each group gk 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 /nk for all zi ∈ Dk, so the sample weights also converge.