Year: 2021

Abstract:
Importance weighting is a classic technique to handle distribution shifts. However, prior work has presented strong empirical and theoretical evidence demonstrating that importance weights can have little to no effect on overparameterized neural networks. Is importance weighting truly incompatible with the training of overparameterized neural networks? Our paper answers this in the negative. We show that importance weighting fails not because of the overparameterization, but instead, as a result of using exponentially-tailed losses like the logistic or crossentropy loss. As a remedy, we show that polynomially-tailed losses restore the effects of importance reweighting in correcting distribution shift in overparameterized models. We characterize the behavior of gradient descent on importance weighted polynomially-tailed losses with overparameterized linear models, and theoretically demonstrate the advantage of using polynomially-tailed losses in a label shift setting. Surprisingly, our theory shows that using weights that are obtained by exponentiating the classical unbiased importance weights can improve performance. Finally, we demonstrate the practical value of our analysis with neural network experiments on a subpopulation shift and a label shift dataset. Our polynomially-tailed loss consistently increases the test accuracy by 2-3%.

heading: 1 INTRODUCTION
text: Machine learning models are often evaluated on test data which differs from the data that they were trained on. A classic statistical technique to combat such distribution shift is to importance weight the loss function during training (Shimodaira, 2000). This procedure upweights points in the training data that are more likely to appear in the test data and downweights ones that are less likely. The reweighted training loss is an unbiased estimator of the test loss and can be minimized by standard algorithms, resulting in a simple and general procedure to address distribution shift.
Surprisingly, recent papers (Byrd & Lipton, 2019; Xu et al., 2020) have found that importance weighting is ineffective in the current deep learning paradigm, where overparameterized models interpolate the training data or have vanishingly small train loss. In particular, Byrd & Lipton (2019) empirically showed that when no regularization is used, overparameterized linear and nonlinear models trained with the importance weighted cross-entropy loss ignore the importance weights. Xu et al. (2020) followed up and provided a theoretical justification for this observation in overparameterized linear and non-linear models.
To build intuition about why importance weighting fails, consider linear classifiers as an example. Given linearly separable data (x1, y1), . . . , (xn, yn) ∈ Rd × {−1, 1}, Soudry et al. (2018) showed that if gradient descent is applied to minimize an exponentially-tailed classification loss ( ∑ i∈[n] `exp(yixi)) then the iterates converge in direction to the maximum margin classifier θ̂MM := arg min‖θ‖=1 {γ : yixi · θ ≥ γ, for all i ∈ [n]}. Xu et al. (2020) showed that in this same setting, minimizing the importance weighted loss ( ∑ i∈[n] wi`exp(yixi)) with gradient descent also results in convergence to the maximum margin classifier, regardless of the weights. To see why, consider the special case where the weights (w1, . . . , wn) are positive integers. This reweighting is equivalent to simply repeating each datapoint wi times, and the maximum margin classifier over this “new dataset” remains unchanged. Thus, invoking the original result by Soudry et al. (2018) proves that the importance weights has no effect in correcting the distribution shift. This result can be seen in Figure 1 where we demonstrate this phenomenon in a simple toy problem.
Such evidence has led some to wonder if importance weighting is fundamentally incompatible with overparameterized interpolating models. In this paper, we show that this is not the case. We find that the culprit behind the ineffectiveness of importance weighting is the exponential tail of popular losses such as the cross-entropy or the logistic. We propose altering the structure of the loss to have fatter, polynomially decaying tails instead. We theoretically and empirically demonstrate that importance weights do correct for distribution shift under such losses even for overparameterized classifiers.
Our first contribution is to characterize the limiting direction of the iterates of gradient descent (its implicit bias) when minimizing reweighted polynomially-tailed losses with linear classifiers. We show that this limiting direction is a function of both the datapoints as well as the importance weights, unlike the maximum margin classifier that only depends on the data (see the right half of Figure 1). Next, we analyze the generalization behavior of this classifier in a label shift setting. We prove that when the weights are an exponentiation of the unbiased importance weights, the test error decays to zero in the large sample limit, regardless of the level of imbalance in the data. In contrast, we prove that the test error of the maximum margin classifier in this same setting must be at least 1/8.
Finally, we demonstrate the practical benefits of our framework by applying this approach to experiments with neural networks. In both a label shift dataset (Imbalanced Binary CIFAR10), and a subpopulation shift dataset with spurious correlations (CelebA (Sagawa et al., 2019)), we find that using polynomially-tailed losses consistently leads to a gain of 2-3% in test accuracy over using the cross-entropy loss.

heading: 2 RELATED WORK
text: Early work (Shimodaira, 2000; Wen et al., 2014) already warned against the potential ineffectiveness of importance weights on interpolating overparameterized models. Shimodaira (2000) showed that when the model is well-specified, importance weights can fail to have an effect, and that the ordinary maximum likelihood estimate is asymptotically optimal. Wen et al. (2014) showed that when there is a zero-loss minimizer of an unweighted convex loss minimization problem, then it is also a minimizer of the (adversarially) reweighted loss as well. Recent work (Byrd & Lipton, 2019; Xu et al., 2020) has shown that importance weighting fails to have an effect on neural networks trained with gradient descent, though always in the setting of exponentially-tailed losses. Sagawa et al. (2019) demonstrated that reweighting can fail to have the desired effect when unregularized distributionally robust optimization (DRO) methods are used in conjunction with the cross-entropy loss. They empirically showed that regularization is necessary to reap the benefits of reweighting, also observed by Byrd & Lipton (2019).
Our work also connects to literature that has studied the implicit bias of gradient descent (Soudry et al., 2018; Ji & Telgarsky, 2019; Nacson et al., 2019). Especially relevant is the work by Ji et al. (2020) who relate the implicit bias of gradient descent with exponentially and polynomially-tailed losses for linear classifiers to a solution of a regularized loss minimization problem.
Finally, our generalization analysis draws from the growing literature focused on finite sample bounds on the test error of the maximum margin classifier in the overparameterized regime (Chatterji & Long, 2021; Muthukumar et al., 2020; Wang & Thrampoulidis, 2021; Cao et al., 2021).

heading: 3 SETTING
text: We consider a distribution shift setting where the training samples {(x1, y1), . . . , (xn, yn)} ∈ Rd × {−1, 1} are drawn i.i.d. from Ptrain, and the test samples are drawn from a different distribution Ptest that is absolutely continuous with respect to Ptrain. Let fθ denote a classifier parameterized by θ. Given a feature x, a classifier maps this feature to fθ(x) ∈ R. In this paper we shall consider cases where the classifier is either linear (for our theory) or a neural network (for our experiments).
Our goal is to find a classifier fθ that minimizes the 0-1 loss with respect to the test distribution:
TestError[fθ] = P(x,y)∼Ptest [sign(fθ(x)) 6= y] .
To handle the mismatch between Ptrain and Ptest, we shall study importance weighting algorithms. Given a datapoint (x, y) ∈ Rd×{−1, 1}, the classical unbiased importance weight at this datapoint is given by the ratio of densities between the test and the train distributions Ptest(x,y)Ptrain(x,y) . Using these unbiased importance weights ensure that the reweighted training loss is an unbiased estimate of the test loss.
However, as noted above, past work has shown that interpolating classifiers trained with gradient descent on importance weighted exponentially-tailed losses, such as the logistic loss `log(z) := log (1 + exp(−z)), the exponential loss `exp(z) := exp(−z), and the cross-entropy loss, ignore the importance weights. For example, consider the case when the classifier is linear fθ(x) = x · θ, the weights are w1, . . . , wn > 0, and the reweighted loss function is L̂(θ) = ∑n i=1 wi`log(yixi · θ). Xu et al. (2020) showed that if the data is linearly separable then the iterates of gradient descent converge in direction to the `2-maximum margin classifier,
θ̂MM := arg max θ:‖θ2‖=1
{γ : yixi · θ ≥ γ for all i ∈ [n]} . (1)
Observe that the maximum margin classifier does not depend on the set of importance weights (w1, . . . , wn) and hence may suffer large test error when there is distribution shift. Xu et al. (2020) further showed that when separability assumptions hold, non-linear classifiers (like multilayer neural networks) trained with gradient descent on exponentially-tailed losses are also unaffected by importance weights.
We initiate a study of polynomially-tailed losses in the distribution shift setting and show that they have improved behavior with respect to importance weighting even when the model is overparameterized. Given parameters α > 0 and β ∈ R define the polynomially-tailed loss as follows:
`α,β(z) :=
{ `left(z) if z < β
1 [z−(β−1)]α if z ≥ β,
where `left is any loss function such that the overall loss function `α,β is convex, differentiable and strictly decreasing. Several natural choices for `left include the scaled logistic (c1 log(1 + exp(−c2z))), exponential (c1 exp(−c2z)) or linear (−c1z + c2) losses. Given a training dataset {(x1, y1), . . . , (xn, yn)} and a set of weights w1, . . . , wn ≥ 0 we let
L̂α,β(fθ) := n∑ i=1 wi`α,β(yifθ(xi))
be the reweighted empirical loss on this dataset.
Notation. Given a vector v, let ‖v‖ denote its Euclidean norm. For any j ∈ N, we denote the set {1, . . . , j} by [j]. A random variable ξ is 1-sub-Gaussian if for any λ ∈ R, E [ eλξ ] ≤ eλ2/2.

heading: 4 THEORETICAL RESULTS
text: In this section, we present several theoretical results that justify the use of polynomially-tailed losses in conjunction with importance weights to handle distribution shifts. To let the analysis proceed, we restrict our theoretical study to linear classifiers, fθ(x) = x · θ, for some θ ∈ Rd. First, in Section 4.1 we shall characterize the limiting direction of gradient descent on reweighted polynomially-tailed losses and show that this direction depends on both the weights as well as the datapoints. Next, in Section 4.2, we upper bound the test error of this limiting solution in a label shift setting. We also show that choosing weights that are obtained by exponentiating the unbiased importance weights helps in reducing the test error. Finally, in this label shift setting, we show that the maximum margin classifier suffers an error that is at least 1/8.

heading: 4.1 IMPLICIT BIAS OF GRADIENT DESCENT ON POLYNOMIALLY-TAILED LOSSES
text: We begin by presenting a result that characterizes the implicit bias of gradient descent on reweighted polynomially-tailed losses for linearly separable classification. Understanding this situation is a key first step, as gradient descent for separable linear classification is often used as a simplified model to theoretically characterize the behavior of overparameterized models such as neural networks (Soudry et al., 2018; Ji & Telgarsky, 2019; Nacson et al., 2019; Lyu & Li, 2019; Ji & Telgarsky, 2020).
Given a linearly separable dataset (x1, y1), . . . , (xn, yn), we let zi := yixi. We shall analyze the iterates of gradient descent with step-size η > 0 and initial iterate θ(0) ∈ Rd:
θ(t+1) = θ(t) − η∇L̂α,β(θ(t)), for t ∈ {0, 1, . . .}.
Define the direction
θ̂α := arg min θ:‖θ‖=1 { n∑ i=1 wi (zi · θ)α , s.t. zi · θ > 0, for all i ∈ [n] } . (2)
The following proposition characterizes the limiting direction of gradient descent iterates.
Proposition 4.1. Suppose that the data is linearly separable. For any α > 0, β ∈ R, any initial point θ(0) ∈ Rd, and for all small enough step-sizes η the direction of the gradient descent iterates satisfy the following:
lim t→∞
θ(t)
‖θ(t)‖ → θ̂α.
The proof is presented in Appendix A. This proof relies on the recent result by Ji et al. (2020) who relate the limiting direction of gradient descent on the unregularized loss to the limiting solution of a norm-constrained loss minimization problem, where the limit is taken with respect to the norm constraint.
Note that, unlike the maximum margin classifier, it is immediately clear that this limiting direction θ̂α depends on the weights w1, . . . , wn. As one would intuitively expect, the direction θ̂α tries to achieve a larger margin zi · θ on points with larger weights wi. This behavior is also apparent in the simulation in the rightmost panel in Figure 1, where upweighting points in the minority class helps to learn a classifier that is similar in direction to the Bayes optimal classifier for the problem.
Our limiting direction θ̂α has the interesting property that it does not depend on several quantities: the initial point θ(0), the properties of `left, and the “switchover” point β. Linear separability ensures that in the limit, the margin on each point is much larger than β, and so the loss of each point is in the polynomial tail of part of `α,β .

heading: 4.2 GENERALIZATION ANALYSIS
text: The result in the previous subsection shows that the asymptotic classifier learnt by gradient descent, θ̂α respects importance weights. However, this does not demonstrate that using polynomially-tailed losses leads to classifiers with lower test error compared to using exponentially-tailed losses.
To answer this more fine-grained question, we will need to perform a more refined analysis. To do this we will make sub-Gaussian cluster assumptions on the data, similar to ones considered in the generalization analysis of the overparameterized maximum margin classifiers (Chatterji & Long, 2021; Wang & Thrampoulidis, 2021; Liang & Recht, 2021; Cao et al., 2021). In our setting, the features associated with each label are drawn from two different sub-Gaussian clusters. We shall consider a label shift problem where the training data is such that the number of data points from the positive (majority) cluster will be much larger than the number of datapoints from the negative (minority) cluster. The test datapoints will be drawn uniformly from either cluster.
Under these assumptions, we derive upper bounds on the error incurred by θ̂1, the limiting direction of gradient descent of polynomially tailed losses with α = 1 (Section 4.2.2). We will find that its test error is small when the reweighting weights are set to cubic powers of the unbiased importance weights. (A similar analysis can also be conducted for other α.) In the same setting in Section 4.2.3, we will show that the maximum margin classifier must suffer large test errors.

heading: 4.2.1 SETTING FOR GENERALIZATION ANALYSIS
text: Here we formally describe the setting of our generalization analysis. We let C ≥ 1 and 0 < C̃ ≤ 1 denote positive absolute constants, whose value is fixed throughout the remainder of the paper. We will use c, c′, c1, . . . to denote “local” positive constants, which may take different values in different contexts.
The training dataset S := {(x1, y1), . . . , (xn, yn)} has n independently drawn samples. The conditional distribution of the features given the label is
x | {y = 1} = µ1 + Uq x | {y = −1} = µ2 + Uq,
where µ1 · µ2 = 0, ‖µ1‖ = ‖µ2‖, U is an arbitrary orthogonal matrix, and q is a random variable such that
• its entries are 1-sub-Gaussian and independent, and • E [ ‖q‖2 ] ≥ C̃d.
We note that in past work that studied this setting (Chatterji & Long, 2021; Wang & Thrampoulidis, 2021; Cao et al., 2021), the cluster centers were chosen to be opposite one another µ1 = −µ2. Here, since our goal here is to study a label shift setting we consider the cluster centers µ1 and µ2 to be orthogonal. This ensures that learning the direction of one of the centers reveals no information about the other center, which makes this problem more challenging in this setting.
Define P := {i ∈ [n] : yi = 1} to be the set of indices corresponding to the positive labels and N := {i ∈ [n] : yi = −1} to be the set of indices corresponding to the negative labels. As stated above, we will focus on the case where |P| |N | ≥ 1. Let τ := |P||N| ≥ 1 be the ratio between the number of positive and negative samples.
The test distribution Ptest is balanced. That is, if (x, y) ∼ Ptest, then P [y = 1] = P [y = −1] = 1/2 and x | y follows the distribution as described above. We shall study the case where negative examples (which are in the minority) are upweighted. Specifically, set the importance weights as follows
wi = { 1 if i ∈ P w if i ∈ N .
Assumptions. Given a failure probability 0 < δ < 1/C, we make the following assumptions on the parameters of the problem:
1. Number of samples n ≥ C log(1/δ). 2. Norm of the means ‖µ‖2 := ‖µ1‖2 = ‖µ2‖2 ≥ Cn2 log(n/δ). 3. Dimension d ≥ Cn‖µ‖2.
Our assumptions allow for large overparameterization, where the dimension d scales polynomially with the number of samples n.

heading: 4.2.2 UPPER BOUND FOR THE REWEIGHTED POLYNOMIALLY-TAILED LOSS CLASSIFIER
text: First, we shall prove an upper bound on the test error of the solution learnt by gradient descent on the reweighted polynomially-tailed loss. Under the choice of weights described above, by equation (2)
θ̂1 = arg min θ:‖θ‖=1 ∑ i∈P 1 zi · θ + ∑ i∈N w zi · θ , s.t. zi · θ > 0, for all i ∈ [n],
where zi = yixi as defined previously. The following theorem provides an upper bound on the test error of θ̂1 and specifies a range of values for the weight w.
Theorem 4.2. For any 0 < C̃ ≤ 1, there is a constant c such that, for all large enough C and for any 0 < δ < 1/C, under the assumptions of this subsection the following holds. If the weight
τ3
2 ≤ w ≤ 2τ3,
then with probability at least 1− δ, training on S produces a classifier θ̂1 satisfying:
TestError[θ̂1] = P(x,y)∼Ptest [ sign ( θ̂1 · x ) 6= y ] ≤ exp ( −c‖µ‖ 4
d
) .
This theorem is proved in Appendix B. To prove this theorem, we first show that the data is linearly separable under our assumptions and hence by the implicit bias results of Proposition 4.1 our upper bound is equivalent to bounding the test error of the limiting iterate of gradient descent on our reweighted loss. The next step is to show that the sum of the gradients across each of the two groups remains roughly balanced throughout training under our choice of w. This ensures that the classifier aligns well with both cluster centers µ1 and −µ2 which then proves our bound on the test error via a standard Hoeffding bound.
If we consider a regime where d and n are growing simultaneously, then the test error goes down to zero if the norm of the cluster means ‖µ‖ grows faster than d1/4, regardless of the level of imbalance in the training data. Briefly note that if n is a constant, then the bound on the test error here matches the bound obtained for the maximum margin classifier in (Chatterji & Long, 2021; Wang & Thrampoulidis, 2021; Cao et al., 2021) (although these bounds are not directly comparable as the setting here is slightly modified for a label shift). The impossibility result by Jin (2009) shows that these previous results guarantee that learning occurs right up until the information theoretic frontier (‖µ‖2 = ω( √ d)).
Finally, it is rather interesting to note that the theorem requires the weight w to scale with τ3, instead of τ , the unbiased importance weight, which would ensure that the reweighted training loss is an unbiased estimate of the test loss. In our proof, we find that in order to suffer low test error it is important to guarantee that the norm of the gradients across the two groups is roughly balanced throughout training. We show that at any training iteration, the ratio of the sum of the derivative of the weighted losses in the majority and minority clusters scales as τ
w1/3 . Thus choosing w to
scale with τ3 ensures that the gradients are balanced across both groups. We also verify this in our simulations in Figure 2, and find that choosing w = τ3 ensures that the test error is equal across both classes leading to low overall test error.

heading: 4.2.3 LOWER BOUND FOR THE MAXIMUM MARGIN CLASSIFIER
text: In the previous section we derived an upper bound to demonstrate that classifiers trained with polynomially-tailed losses achieve low test error. We will now show that classifiers trained with exponentially-tailed losses have their error lower bounded by 1/8 in the same setting.
Theorem 4.3. Let q ∼ N(0, Id×d). There exist constants c and c′ such that, for all large enough C and for any 0 < δ < 1/C, under the assumptions of this subsection the following holds. With probability at least 1− δ, training on S produces a maximum margin classifier θ̂MM satisfying:
TestError[θ̂MM] = P(x,y)∼Ptest [ sign ( θ̂MM · x ) 6= y ] ≥ 1 2 · Φ ( − c τ · √ n‖µ‖2√ d ) ,
where Φ is the Gaussian cdf. Furthermore, if the imbalance ratio τ ≥ c′ √ n‖µ‖2√ d
then with probability at least 1− δ
TestError[θ̂MM] ≥ 1
8 .
This theorem is proved in Appendix C. Intuitively, the lower bound holds because when the imbalance is severe, the max-margin classifier essentially ignores the samples in the minority class and overfits to the majority class. Ignoring these negative samples results in a relatively small alignment between the maximum margin classifier and the negative of the minority class center, which leads to high error on this class.
Together, these two theorems demonstrate that there is a strict gap between the performance of exponentially tailed and polynomially tailed losses under distribution shifts. As a concrete example, consider the scaling where ‖µ‖2 = d 12+ 140 , n = d 15 and τ = d 320 ≤ n. For all large enough d, all our assumptions are satisfied. Now since, ‖µ‖2 > √ d, Theorem 4.2 guarantees that the test error of θ̂1 → 0 as d → ∞. However, as τ = d 3 20 ≥ c′‖µ‖2 √ n/d = c′d 5 40 , for all large enough d, the test error of the maximum margin classifier is guaranteed to be at least 1/8. This degradation of the test error with τ is also apparent in our simulations in Figure 2.

heading: 5 EMPIRICAL EVALUATION ON DEEP INTERPOLATING CLASSIFIERS
text: Inspired by our theoretical results, we use a polynomially-tailed loss with importance weights to train interpolating deep neural networks under distribution shift. We use the same polynomiallytailed loss with α = 1, β = 1, and `left(z) = log(1 + e−z)/ log(1 + e−1), ensuring that `α,β is continuous and differentiable at transition z = β. We train models on two image classification datasets, one with label shift and one with subpopulation shift, and include additional experiment details in Appendix D. Though these nonlinear networks violate the assumptions of our theory, polynomially-tailed loss with importance weights consistently improves test accuracy for interpolating neural networks under distribution shift compared to importance-weighted cross-entropy loss.
Imbalanced Binary CIFAR10. We construct our label shift dataset from the full CIFAR10 dataset. Similar to Byrd & Lipton (2019), we create a binary classification dataset out of the “cat” and “dog” classes. We use the official test examples as our label-balanced test set of 1000 cats and 1000 dogs. To form the train and validation sets, we use all 5000 cat examples but only 500 dog examples from the official train set, corresponding to a 10:1 label imbalance. We then use 80% of those examples for training and the rest for validation. We use the same convolutional neural network architecture as Byrd & Lipton (2019) with random initializations for this dataset.
CelebA. For our subpopulation shift dataset, we use the CelebA with spurious correlations dataset constructed by Sagawa et al. (2019). This dataset has two class labels, “blonde hair” and “dark hair”. Distinguishing examples by the “male” versus “female” attribute results in four total subpopulations, or groups. The distribution shift in the dataset comes from the change in relative proportions among the groups between the train and test sets. To reduce computation, we train on 2% of the full CelebA training set, resulting in group sizes of 1446, 1308, 468, and 33. We construct our test set by downsampling the original test to get group-balanced representation in the test set, resulting in 180 examples for each group. Following Sagawa et al. (2019), we use a ResNet-50 with ImageNet initialization for this dataset.

heading: 5.1 INTERPOLATING REGIME
text: To understand the impact of polynomial losses in the interpolating regime, we train unregularized neural networks with SGD past 100% train accuracy and report their final accuracies on the test sets. We compare models trained with our polynomially-tailed loss against those trained with the crossentropy loss which has exponential tails. Let w correspond to the unbiased importance weights. We consider three weighting scenarios for both loss functions: no importance weighting at all, “No IW”, the classical unbiased importance weighting, “IW (w)”, and biased weighting where we exponentiate the weights, “IW (wc))”, increasing the ratio of different weights. The third setting is inspired by our theory in Section 4 that shows biased importance weights can improve performance for polynomially-tailed losses. We set the exponents to be the largest value that still allowed for stable optimization. For CelebA, we exponentiate by 2 and for CIFAR10 we exponentiate by 3/2. We run 10 seeds for each experiment setting. We compare the two losses via paired one-sided Welch’s t-tests, pairing runs with the same random seed. We report the exact numbers in Appendix D.
1 The mean test accuracy for “No IW” with cross-entropy on CIFAR10 appears to deviate from “IW”, but this is due to the high variance across multiple runs of the model. The high variance arises due to 10:1 imbalanced training data, combined with the fact that we train the models until they interpolate the training data. These results for “No IW” differ slightly from those in Byrd & Lipton (2019) which used balanced training data and imbalanced test data.
Figure 3 shows the mean accuracy and the standard error for each of the three settings. As indicated by our theory, we find that importance weighting with polynomially-tailed losses leads to statistically significant gains over cross-entropy in all cases. Further we find that exponentiating weights boosts the performance of polynomially tailed losses, confirming our claims in Theorem 4.2. However, exponentiating the weights leaves the performance of cross-entropy loss largely unchanged. In the case of CelebA, we find that exponentiated reweighting with poly-tailed losses outperforms all other methods. While in the case of Imbalanced Binary CIFAR10, we find a smaller gap between polynomially-tailed losses and cross-entropy, partially due to a substantially higher run-to-run variability in training. This variability does not affect our main conclusion: exponentiated reweighting with polynomially tailed losses still outperforms this at a significance level of p = 0.01 even with this noise level.

heading: 5.2 EARLY-STOPPING BEFORE MODELS INTERPOLATE
text: Prior works (Sagawa et al., 2019; Byrd & Lipton, 2019) found that adding regularization such as strong L2 regularization and early stopping can partially restore the effects of importance weights at the cost of not interpolating the training set. As such, we compare the performance of our polynomially-tailed loss against cross-entropy when models are early-stopped. For each run in Section 5.1, we select the model checkpoint with the best weighted validation accuracy and evaluate the checkpoint on the test set. Figure 4 compares the test accuracies of the early-stopped models (IW-ES (w), IW-ES (wc)) against reweighted models trained past interpolation (IW (w)).
Consistent with prior works, we see that early stopping as a form of regularization improves the test accuracies of both loss functions when used with importance weights. Our polynomially-tailed loss gives test accuracies that are better than or similar to cross-entropy in all weighted loss scenarios. The gain over cross-entropy is statistically significant in the Imbalanced Binary CIFAR10 runs. On CelebA, the polynomially-tailed loss with squared weights attains the highest mean test accuracy out of all settings.

heading: 6 CONCLUSION
text: Contrary to popular belief, importance weights are in fact compatible with the current paradigm of overparameterized classifiers, provided that the loss function does not decay exponentially. The limiting behavior of an interpolating linear classifier trained with a weighted polynomially-tailed loss provably depends on the importance weights, unlike the maximum-margin classifier recovered by an exponentially decaying loss function. From bounding the generalization error of our loss function, we found a setting where increasing the ratio of the weights improves the test accuracy. We empirically corroborated these theoretical intuitions, finding that weighted polynomially-tailed losses are also able to address distribution shift for interpolating neural networks. Our work suggests that heavy-tailed losses together with importance weighting serve as a simple and general candidate for addressing distribution shift in deep learning.
ETHICS STATEMENT
We use publicly available datasets in our paper, and do not foresee any potential harm arising from the use of our methods.
REPRODUCIBILITY STATEMENT
The full proofs of our theoretical results are presented in the appendix. All experimental details and hyperparameters are presented in Appendix D.

heading: A PROOF OF PROPOSITION 4.1
text: First, we restate the statement of the proposition.
Proposition 4.1. Suppose that the data is linearly separable. For any α > 0, β ∈ R, any initial point θ(0) ∈ Rd, and for all small enough step-sizes η the direction of the gradient descent iterates satisfy the following:
lim t→∞
θ(t)
‖θ(t)‖ → θ̂α.
Proof Recall that the loss `α,β is assumed to be convex, strictly decreasing to zero and is differentiable. Define the minimizer of the loss over a ball of radius R as follows
θR = arg min θ:‖θ‖≤R L̂α,β(θ).
Under the assumptions of this proposition, we can invoke Theorem 1 by Ji et al. (2020) to get that
lim t→∞
θ(t)
‖θ(t)‖ = limR→∞ θR R .
We will therefore instead demonstrate that limR→∞ θRR = θ̂α to establish our claim.
First, note that the classifier θR is the minimizer of loss in a ball of radius R, so there exists a R0 with
L̂α,β(θR0) = ∑ i∈[n] wi`α,β(zi · θR0) ≤ (min i wi)`α,β(β)
such that for all R > R0, each example is classified correctly by θR and the margin (zi ·θR) on each example is at least β. Therefore,
L̂α,β(θR) = n∑ i=1 wi`α,β(zi · θR) = n∑ i=1 wi [zi · θR − (β − 1)]α .
Thus, for any radius R > R0 we have
θR R = 1 R arg min θ:‖θ‖≤R L̂α,β(θ)
= 1
R arg min θ:‖θ‖≤R { n∑ i=1 wi [zi · θ − (β − 1)]α , s.t., zi · θ ≥ β }
= 1
R arg min θ:‖θ‖≤R  1Rα n∑ i=1 wi[ zi · ( θ R ) − β−1R ]α , s.t., zi · ( θR) ≥ βR 
= 1
R × R · arg min θ:‖θ‖≤1  n∑ i=1 wi[ zi · θ − β−1R ]α , s.t., zi · θ ≥ βR   .
Now taking the limit R→∞ we get that
lim R→∞ θR R = lim R→∞ arg min θ:‖θ‖≤1  n∑ i=1 wi[ zi · θ − β−1R ]α , s.t., zi · θ ≥ βR 
(i) = arg min
θ:‖θ‖≤1  limR→∞ n∑ i=1 wi[ zi · θ − β−1R ]α , s.t., zi · θ ≥ βR 
= arg min θ:‖θ‖≤1 { n∑ i=1 wi (zi · θ)α , s.t., zi · θ ≥ 0 }
(ii) = arg min
θ:‖θ‖≤1 { n∑ i=1 wi (zi · θ)α , s.t., zi · θ > 0 }
= θ̂α.
where (i) follows since for every R the function n∑ i=1 wi[ zi · θ − β−1R
]α is convex and continuous, and therefore
arg min θ:‖θ‖≤1  n∑ i=1 wi[ zi · θ − β−1R ]α , s.t., zi · θ ≥ βR 
has a unique minimizer (since the objective function is strictly convex, and the constraint set is convex and compact), and is a continuous map by Berge’s maximum theorem (see, e.g., Beavis & Dobbs, 1990, Theorem 3.6). Thus, it is possible to switch the order of the limit and the arg min. Equation (ii) follows since the data is linearly separable, so there exists a θ such that zi · θ > 0 for all i which has finite objective value.
Therefore,
lim t→∞
θ(t)
‖θ(t)‖ = limR→∞ θR R = θ̂α,
which finishes our proof.

heading: B PROOF OF THEOREM 4.2
text: In this section, we will prove an upper bound on the test error θ̂1, which is the classifier learnt by gradient descent with the importance weighted Polynomially-tailed loss that decays as 1/z. By Proposition 4.1, we know the choice of `left and β does not affect the implicit bias of gradient descent, hence here, for the sake of convenience, we shall use the specific loss function
`1,β(z) :=
{ −(z − β) + 1 if z < β
1 z−(β−1) if z ≥ β,
with β = 1 − wn < 0. It can be easily checked that this loss function is convex, differentiable and monotonically decreasing.
We want to bound the test error of θ̂1, however, in light of Proposition 4.1 we will instead bound the test error of the limiting iterate of gradient descent starting from θ(0) = 8w2/3n ( µ1 − w1/3µ2 ) with step-size η > 0 (again by the implicit bias of gradient descent is not affected by the choice of the initial point):
θ(t+1) = θ(t) − η∇L̂(θ(t)),
where
L̂(θ) = ∑ i∈P `1,β(zi · θ) + w ∑ i∈N `1,β(zi · θ).
Recall that if the step-size is small enough then Proposition 4.1 guarantees that limt→∞ θ (t)
‖θ(t)‖ = θ̂1.
In this section we will simply let ` denote `1,β and also use the shorthands ` (t) i := `(zi · θ(t)) and ` ′(t) i := `
′(zi · θ(t)). With this setup in place let us begin the proof. All of the assumptions made in Section 4.2.1 are in scope here. First we have a lemma that upper bounds the test error using Hoeffding’s inequality.
Lemma B.1. There is a positive absolute constant c such that
P(x,y)∼Ptest [sign (θ · x) 6= y] ≤ 1
2
[ exp ( −c (θ · µ1) 2 ‖θ‖2 ) + exp ( −c (θ · µ2) 2 ‖θ‖2 )] .
This lemma follows by mirroring the proof of (Chatterji & Long, 2021, Lemma 9) which in turn follows by a simple application of Hoeffding’s inequality (Vershynin, 2018, Theorem 2.6.3).
Next we have a lemma that proves bounds on the norms of the samples, bounds the inner products between the samples and also shows that with high probability the data is linearly separable.
Recall that the constants C ≥ 1 and 0 < C̃ ≤ 1 were defined above in Section 4.2. Lemma B.2. For all 0 < C̃ ≤ 1, there is a c ≥ 1 such that, for all large enough C, with probability 1− δ over the draw of the samples the following events simultaneously occur:
for all k ∈ [n], d c ≤ ‖zk‖2 ≤ cd; (3)
for all i 6= j ∈ P, |zi · zj | ≤ c ( ‖µ‖2 + √ d log(n/δ) ) ; (4)
for all i 6= j ∈ N , |zi · zj | ≤ c ( ‖µ‖2 + √ d log(n/δ) ) ; (5)
for all i ∈ P and j ∈ N , |zi · zj | ≤ c √ d log(n/δ); (6)
for all i 6= j ∈ [n], |zi · zj | < c(‖µ‖2 + √ d log(n/δ)); (7)
for all k ∈ P, |µ1 · zk − ‖µ‖2| < ‖µ‖2/2; (8) for all k ∈ N , |(−µ2) · zk − ‖µ‖2| < ‖µ‖2/2; (9) for all k ∈ N , |µ1 · zk| < c‖µ‖ √ log(n/δ); (10)
for all k ∈ P, |µ2 · zk| < c‖µ‖ √
log(n/δ); (11) the samples are linearly separable. (12)
Proof We shall prove that each of the ten parts hold with probability at least 1 − δ/10 and take a union bound to prove our lemma. Under the assumptions of this lemma, Parts (3)-(5), (7)-(9) and (12) can be shown to hold with the required probability by invoking (Chatterji & Long, 2021, Lemma 13). We will show that Parts (6), (10) and (11) also hold with the required probability to complete the proof.
Proof of Part (6): Note that for all i ∈ P , zi = µ1 + Uqi, and for all i ∈ N , zi = µ2 + Uq2. Therefore, for any i ∈ P and j ∈ N , since µ1 · µ2 = 0,
|zi · zj | = µ1 · (Uqj) + µ2 · (Uqi) + qi · qj ≤ |µ1 · (Uqj)|+ |µ2 · (Uqi)|+ |qi · qj | = |(U>µ1) · qj |+ |(U>µ2) · qi|+ |qi · qj |. (13)
We will show that each of these terms is small with high probability for all pairs i ∈ P and j ∈ N .
For the first two terms, note that by Hoeffding’s inequality (Vershynin, 2018, Theorem 2.6.3)
P [ |(U>µ1) · qj | > c‖µ‖ √ log(n/δ)
3
] < 2 exp ( −c1c 3
9 ‖µ‖2 log(n/δ) ‖µ‖2 ) = 2 exp ( −c1c 3
9 log(n/δ) ) = 2 ( δ
n )c1c3/9 ≤ δ
30n ,
where the last inequality follows since the constant c ≥ 1 can be chosen to be large enough and because δ < 1/C for a large enough constant C. Therefore, by taking a union bound over all j ∈ N
P [ ∃j ∈ N , |(U>µ1) · qj | > c‖µ‖ √ log(n/δ)
3
] < δ
30 . (14)
Using an analogous argument we can also show that
P [ ∃i ∈ P, |(U>µ2) · qi| > c‖µ‖ √ log(n/δ)
3
] < δ
30 . (15)
Next, using (Chatterji & Long, 2021, inequality (15)) we get that
P [ ∃i 6= j ∈ [n], |qi · qj | > c √ d log(n/δ)
3
] ≤ δ
30 . (16)
Combining inequalities (13)-(16) we get that for all i ∈ P and j ∈ N with probability at least 1− δ/10
|zi · zj | ≤ 2c‖µ‖
√ log(n/δ) 3 + c √ d log(n/δ) 3 ≤ c √ d log(n/δ),
where the last inequality follows since by assumption d ≥ Cn‖µ‖2 and because C is large enough. This completes the proof of this part.
Proof of Part (10): For any k ∈ N , zk = µ2 + Uqk. Recall that µ1 · µ2 = 0, thus |µ1 · zk| = |µ1 · (Uqk)|. Invoking inequality (14) proves that this part holds with probability at least 1− δ/10. Proof of Part (11): This follows by an analogous argument as in the previous part.
We continue by defining a good event that we will work under for the rest of the proof. Definition B.3. If the training dataset S satisfies all the conditions specified in Lemma B.2 then we call it a good run.
Going forward in this section we shall assume that a good run occurs.
The following lemma provides some useful bound on the loss ratio at initialization and guarantees that the loss of example remains in the Polynomially-tailed part throughout training. Lemma B.4. Recall that θ(0) = 8w2/3n(µ1 − w1/3µ2), and that wi = 1 if i ∈ P and wi = w if i ∈ N . On a good run, for all i 6= j ∈ [n]
` (0) i ` (0) j ≤ 8 ( wj wi )1/3 .
Furthermore, if the step-size η is sufficiently small then on a good run, for all t ∈ {0, 1, . . .} and for all i ∈ [n]
` (t) i =
1
zi · θ(t) − (β − 1) .
Proof Consider an i ∈ P then
zi · θ(0) = 8w2/3n ( zi · µ1 − w1/3zi · µ2 ) (i) ≤ 8w2/3n ( 3‖µ‖2
2 + c1w
1/3‖µ‖ √ log(n/δ) ) ≤ 12w2/3n‖µ‖2 [ 1 + 2c1w 1/3 √ log(n/δ)
3‖µ‖ ] (ii)
≤ 16w2/3n‖µ‖2,
where (i) follows by Lemma B.2 and (ii) follows since ‖µ‖2 ≥ Cn2 log(n/δ) ≥ Cτ2 log(n/δ) ≥ C 22/3 w2/3 log(n/δ) for a large enough constant C. Similarly, we also have the lower bound
zi · θ(0) = 8w2/3n ( zi · µ1 − w1/3zi · µ2 ) ≥ 8w2/3n (‖µ‖2 2 − c1w1/3‖µ‖ √ log(n/δ) ) ≥ 4w2/3n‖µ‖2 [ 1− 2c1w 1/3 √ log(n/δ) ‖µ‖ ] .
Recall from the definition of the loss that we set β = 1 − nw. Now again since ‖µ‖2 ≥ C 22/3 w2/3 log(n/δ), and because we choose β = 1− nw, we infer that
zi · θ(0) − (β − 1) ≥ 8w2/3n‖µ‖2
3 − nw ≥ 2w2/3n‖µ‖2.
Since the margin of the point is larger than (β − 1) therefore by the definition of the loss function above we have that
` (0) i =
1
zi · θ(0) − (β − 1)
and hence
1
16w2/3n‖µ‖2 ≤ ` (0) i ≤
1
2w2/3n‖µ‖2 . (17)
By mirroring the logic we can also prove that for all j ∈ N 1
16wn‖µ‖2 ≤ ` (0) j ≤
1
2wn‖µ‖2 . (18)
By combining equations (17) and (18) we immediately get that
` (0) i ` (0) j ≤ 8 ( wj wi )1/3 .
This proves the first part of the lemma.
To prove the second part we shall prove that the loss on each example remains smaller than 1/2, which ensures `(t)i is equal to the polynomial loss function. Note that
L̂(θ(0)) = ∑ i∈P ` (0) i + w ∑ i∈N ` (0) i ≤ |P| 2w2/3n‖µ‖2 + w|N | 2wn‖µ‖2 ≤
wn
2wn‖µ‖2 ≤ 1/2. (19)
This proves that `(0)i ≤ 1/2 for all i ∈ [n]. By (Ji et al., 2020, Lemma 2) we know that if the stepsize is small enough then the sequence {L̂(θ(t))} is non-increasing. Hence, we have that `(t)i ≤ 1/2 for all t ∈ {0, 1, . . .} and for all samples. This wraps up our proof. The next lemma lower bounds the inner product between the normalized gradient descent iterates, and µ1 and µ2.
Lemma B.5. Let c be an absolute constant. Then, on a good run, for any t ∈ {0, 1, . . .}
µ1 · θ(t+1) ‖θ(t+1)‖ ≥ µ1 · θ(0) ‖θ(t+1)‖ + ‖µ‖2 2c √ d  1 1 + ‖θ
(0)‖ cη √ d ∑t s=0 ∑ i∈[n]−wi` ′(s) i × ∑t s=0 [∑ i∈P −wi` ′(s) i − c √ log(n/δ) ‖µ‖ ∑ i∈N −wi` ′(s) i
] ∑t s=0 ∑ i∈[n]−wi` ′(s) i
and
µ2 · θ(t+1) ‖θ(t+1)‖ ≥ µ2 · θ(0) ‖θ(t+1)‖ + ‖µ‖2 2c √ d  1 1 + ‖θ
(0)‖ cη √ d ∑t s=0 ∑ i∈[n]−wi` ′(s) i × ∑t s=0 [∑ i∈N −wi` ′(s) i − c √ log(n/δ) ‖µ‖ ∑ i∈P −wi` ′(s) i
] ∑t s=0 ∑ i∈[n]−wi` ′(s) i .
Proof Let us prove the first claim for the inner product with µ1. The second claim regarding µ2 shall follow analogously. For any t ∈ {0, 1, . . .}, by the definition of the gradient descent step, we have that
µ1 · θ(t+1) = µ1 · θ(t) + η ∑ i∈P wiµ1 · zi ( −`′(t)i ) + η ∑ i∈N wiµ1 · zi ( −`′(t)i ) .
Note that by the definition of the loss function we have that −`′(t)i ≥ 0 for all t and all i. Thus by Parts (3) and (10) of Lemma B.2 we get that
µ1 · θ(t+1) ≥ µ1 · θ(t) + η‖µ‖2
2 ∑ i∈P −wi`′(t)i − c1η‖µ‖ √ log(n/δ) ∑ i∈N −wi`′(t)i
= µ1 · θ(t) + η‖µ‖2
2 ∑ i∈[n] −wi`′(t)i − η‖µ‖2 2
( 1 + c1 √ log(n/δ)
‖µ‖
)∑ i∈N −wi`′(t)i
= µ1 · θ(t) + η‖µ‖2
2 [∑ i∈P −wi`′(t)i − c1 √ log(n/δ) ‖µ‖ ∑ i∈N −wi`′(t)i ] ,
where c1 is the constant from Lemma B.2. Unrolling this inequality over t steps we have that
µ1 · θ(t+1) = µ1 · θ(0) + η‖µ‖2
2 t∑ s=0 [∑ i∈P −wi`′(s)i − c1 √ log(n/δ) ‖µ‖ ∑ i∈N −wi`′(s)i ] . (20)
On the other hand, by the triangle inequality we know that
‖θ(t+1)‖ ≤ ‖θ(t)‖+ η‖∇L̂(θ(t))‖
≤ ‖θ(t)‖+ η ∥∥∥∥∥∥ ∑ i∈[n] zi(−wi`′(t)i ) ∥∥∥∥∥∥ ≤ ‖θ(t)‖+ η
∑ i∈[n] ‖zi‖ (−wi`′(t)i ) (since −wi` ′(t) i ≥ 0)
≤ ‖θ(t)‖+ c1η √ d ∑ i∈[n] −wi`′(t)i , (21)
where the last inequality again uses Part (3) of Lemma B.2. Therefore,
‖θ(t+1)‖ ≤ ‖θ(0)‖+ c1η √ d t∑ s=0 ∑ i∈[n] −wi`′(s)i
≤ c1η √ d t∑ s=0 ∑ i∈[n] −wi`′(s)i
[ 1 +
‖θ(0)‖ c1η √ d ∑t s=0 ∑ i∈[n]−wi` ′(s) i
] .
Thus, combined with inequality (20) we get that,
µ1 · θ(t+1) ‖θ(t+1)‖ ≥ µ1 · θ(0) ‖θ(t+1)‖ + ‖µ‖2 2c1 √ d  1 1 + ‖θ
(0)‖ c1η √ d ∑t s=0 ∑ i∈[n]−wi` ′(s) i × ∑t s=0 [∑ i∈P −wi` ′(s) i − c1 √ log(n/δ) ‖µ‖ ∑ i∈N −wi` ′(s) i
] ∑t s=0 ∑ i∈[n]−wi` ′(s) i ,
which completes the proof of the first part of the lemma. The second part follows by an identical argument.
Next, we prove a lemma that shows that throughout training the ratio between the losses between any two samples remains bounded. Lemma B.6. There is a positive absolute constant c ≥ 1 such that the following holds for all large enough C, and all small enough step-sizes η and for any τ 3
2 ≤ w ≤ 2τ3. On a good run, for all t ∈ {1, 2, . . .} and all i 6= j ∈ [n]
` (t) i ` (t) j ≤ c ( wj wi )1/3 .
Proof Let c1 ≥ 1 be the constant c from Lemma B.2 above. We shall show that the choice c = max { 8, 2(4c21) 1/3 } suffices.
We shall prove this via an inductive argument. For the base case, at step t = 0, we know that by Lemma B.4 that the ratio between the losses of sample i and j is upper bounded by 8 (wj/wi)
1/3. Now, we shall assume that the inductive hypothesis holds at an arbitrary step t > 0 and prove that it holds at step t+ 1.
Without loss of generality, we shall analyze the ratio between the losses of the samples with indices i = 1 and j = 2. A similar analysis shall hold for any other pair. Define Gt := ` (t) 1 , Ht := ` (t) 2 , At := Ht/Gt. Note that,
Gt+1 = 1
θ(t+1) · z1 − (β − 1)
= 1 θ(t) · z1 + η ∑ k∈[n] wk(−` ′(t) k )(zk · z1)− (β − 1)
(i) =
1 θ(t) · z1 − (β − 1) + η ∑ k∈[n] wk(zk · z1) ( ` (t) k )2 = ` (t) 1
1 + η` (t) 1 ∑ k∈[n] wk(zk · z1) ( ` (t) k )2 = Gt · 1
1 + η` (t) 1 ∑ k∈[n] wk(zk · z1) ( ` (t) k )2
where (i) follows since −`′(t)k = (` (t) k )
2 as the loss of each example is always in the polynomial tail of the loss by Lemma B.4. Therefore, we have that
At+1 = Ht+1 Gt+1 = Ht Gt ·
1 + η` (t) 1 ∑ k∈[n] wk(zk · z1) ( ` (t) k )2 1 + η`
(t) 2 ∑ k∈[n] wk(zk · z2) ( ` (t) k )2 = At · 1 + η` (t) 1 ∑ k∈[n] wk(zk · z1) ( ` (t) k
)2 1 + η`
(t) 2 ∑ k∈[n] wk(zk · z2) ( ` (t) k )2 . Now since the step-size η is chosen to be small enough, |zi · zj | ≤ c1d (by Part (4) of Lemma B.2 and by the assumption on d) and because the losses are all smaller than a constant by Lemma B.4, the following approximations hold
1 + η` (t) 1 ∑ k∈[n] wk(zk · z1) ( ` (t) k )2 ≤ exp η`(t)1 ∑ k∈[n] wk(zk · z1) ( ` (t) k )2 1 + η`
(t) 2 ∑ k∈[n] wk(zk · z2) ( ` (t) k )2 ≥ exp η`(t)2 2 ∑ k∈[n] wk(zk · z2) ( ` (t) k )2 , and thus,
At+1 ≤ At exp η`(t)1 ∑ k∈[n] wk(zk · z1) ( ` (t) k )2 − η` (t) 2 2 ∑ k∈[n] wk(zk · z2) ( ` (t) k )2 . (22) Let us further upper bound the RHS as follows
At exp η`(t)1 ∑ k∈[n] wk(zk · z1) ( ` (t) k )2 − η` (t) 2 2 ∑ k∈[n] wk(zk · z2) ( ` (t) k )2 = At exp ( ηw1 ( ` (t) 1 )3 ‖z1‖2 − 1
2 ηw2
( ` (t) 2 )3 ‖z2‖2 )
× exp η`(t)1 ∑ k 6=1 wk(zk · z1) ( ` (t) k )2 − 1 2 η` (t) 2 ∑ k 6=2 wk(zk · z2) ( ` (t) k )2 (i)
≤ At exp ( c1ηdw1 ( ` (t) 1 )3 − 1
2c1 ηdw2
( ` (t) 2 )3)
× exp c2η(`(t)1 + `(t)22 )( ‖µ‖2 + √ d log(n/δ) ) ∑ k∈[n] wk ( ` (t) k )2 ≤ At exp ( c1ηdw1G 3 t − 1
2c1 ηdw2H
3 t
)
× exp c2η (Gt +Ht)(‖µ‖2 +√d log(n/δ)) ∑ k∈[n] wk ( ` (t) k )2 = At exp ( − 1
2c1 ηdw2G
3 t ( A3t −
2c21w1 w2
))
× exp c2η (Gt +Ht)(‖µ‖2 +√d log(n/δ)) ∑ k∈[n] wk ( ` (t) k )2 , (23) where (i) follows since for all i ∈ [n], d/c1 ≤ ‖zi‖2 ≤ c1d and for any j 6= k, |zj · zk| ≤ c1 ( ‖µ‖2 + √ d log(n/δ) ) by Lemma B.2.
With this upper bound in place, consider two cases.
Case 1 (A3t ≤ 4c 2 1w1 w2 ): Using inequality (23) we know that
At+1 ≤ At exp ( − 1
2c1 ηdw2G
3 t ( A3t −
2c21w1 w2
))
× exp c2η (Gt +Ht)(‖µ‖2 +√d log(n/δ)) ∑ k∈[n] wk ( ` (t) k )2 ≤ At exp ( c1ηdw1G 3 t ) exp c2η (Gt +Ht)(‖µ‖2 +√d log(n/δ)) ∑ k∈[n] wk ( ` (t) k )2 .
Now the loss on each example is less than the total initial loss 1/2 (see equation (19)) and therefore,
At+1 ≤ At exp (c3ηdw) exp ( c4η ( ‖µ‖2 + √ d log(n/δ) ) wn )
(i) ≤ (
4c21w1 w2
)1/3 exp(1/8) ≤ 2 ( 4c21w1 w2 )1/3 (ii) ≤ c ( w1 w2 )1/3 ,
where (i) follows by choosing the step-size η to be small enough and because At ≤ ( 4c21w1 w2 )1/3 in this case, and (ii) follows by the choice of constant c ≥ 2(4c21)1/3 from above.
Case 2 ( 4c 2 1w1 w2 < A3t ≤ c 3w1 w2 ): In this case again by inequality (23)
At+1 ≤ At exp ( − 1
2c1 ηdw2G
3 t ( A3t −
2c21w1 w2
))
× exp c2η (Gt +Ht)(‖µ‖2 +√d log(n/δ)) ∑ k∈[n] wk ( ` (t) k )2 = At exp ( − 1
2c1 ηdw2G
3 t ( A3t −
2c21w1 w2
))
× exp c2ηG3t (At + 1)(‖µ‖2 +√d log(n/δ)) ∑ k∈[n] wk ( ` (t) k )2 G2t  (i)
≤ At exp ( − 1
2c1 ηdw2G
3 t ( A3t −
2c21w1 w2
))
× exp c6ηG3t max {(
w1 w2
)1/3 , 1 }( ‖µ‖2 + √ d log(n/δ) ) ∑ k∈[n] wk ( w1 wk )2/3 = At exp ( − 1
2c1 ηdw2G
3 t ( A3t −
2c21w1 w2
))
× exp c6ηw1G3t (‖µ‖2 +√d log(n/δ)) ∑ k∈[n] ( wk min{w1, w2} )1/3 (ii)
≤ At exp ( − 1
2c1 ηdw2G
3 t ( A3t −
2c21w1 w2 )) × exp ( c6ηw1G 3 t ( ‖µ‖2 + √ d log(n/δ) )( |P|+ w1/3|N |
)) = At exp ( − 1
2c1 ηdw2G
3 t ( A3t −
2c21w1 w2 )) × exp ( c6ηw1G 3 t ( ‖µ‖2 + √ d log(n/δ) ) |P| ( 1 + w1/3
τ )) (iii)
≤ At exp ( − 1
2c1 ηdw2G
3 t ( A3t −
2c21w1 w2
)) exp ( 2c6ηw1G 3 t ( ‖µ‖2 + √ d log(n/δ) ) n )
(iv) ≤ At exp ( − 1
2c1 ηdw2G
3 t · 2c21w1 w2
) exp ( 2c6ηw1G 3 t ( ‖µ‖2 + √ d log(n/δ) ) n )
= At exp [ −c1w1G3t ( d− c7n ( ‖µ‖2 + √ d log(n/δ) ))] ,
where (i) follows by the inductive hypothesis which guarantees that ` (t) k
` (t) 1
≤ c(w1wk ) 1/3 ≤ cw1/3, (ii)
follows since wk = 1 if k ∈ P and wk = w if k ∈ N . Inequality (iii) follows since w < 2τ3 and since |P| ≤ n, and finally (iv) follows since in this case A3t > 4c 2 1w1 w2
. Now since by assumption the dimension
d ≥ Cn‖µ‖2 ≥ C2n3 log(n/δ)
for a large enough constant C. Hence, we have that At+1 ≤ At ≤ c ( w1 w2 )1/3 in this case.
Recall that since we assumed the inductive hypothesis to hold, the two cases analyzed above are exhaustive. This completes our proof.
The next lemma uses the loss ratio bound that we established to show that the difference between the gradient of the losses over the positive cluster and the negative cluster is small at any iteration.
Lemma B.7. For any positive c̃, there exists a 0 < c < 1 such that, for all large enough C, if the step-size η is sufficiently small and τ 3
2 ≤ w ≤ 2τ3 then on a good run, for any t ∈ {0, 1, . . .}
∑ i∈P −wi`′(t)i −
c̃ √ log(n/δ)
‖µ‖ ∑ i∈N −wi`′(t)i ≥ c ∑ i∈[n] −wi`′(t)i
and
∑ i∈N −wi`′(t)i −
c̃ √ log(n/δ)
‖µ‖ ∑ i∈P −wi`′(t)i ≥ c ∑ i∈[n] −wi`′(t)i .
Proof We begin by proving the first part of the lemma. Note that since [n] = P ∪ N to prove the first part it suffices to instead show that
(1− c) ∑ i∈P −wi`′(t)i ≥
( c+ c̃ √ log(n/δ)
‖µ‖
)∑ i∈N −wi`′(t)i
⇐⇒ (1− c) ∑ i∈P −`′(t)i ≥
( c+ c̃ √ log(n/δ)
‖µ‖
) w ∑ i∈N −`′(t)i
⇐⇒ (1− c) ∑ i∈P (` (t) i )
2 ≥ ( c+ c̃ √ log(n/δ)
‖µ‖
) w ∑ i∈N (` (t) i ) 2
(since −`′(t)i = (` (t) i ) 2 for the polynomial loss)
⇐ (1− c)|P|min i∈P (` (t) i )
2 ≥ ( c+ c̃ √ log(n/δ)
‖µ‖
) w|N |max
i∈N (`
(t) i ) 2
⇐ (1− c)τ ≥ ( c+ c̃ √ log(n/δ)
‖µ‖
) w maxi∈N (` (t) i ) 2
mini∈P(` (t) i ) 2 ⇐ (1− c)τ ≥ ( c+ c̃ √ log(n/δ)
‖µ‖
) wc1 · 1
w2/3
(by invoking Lemma B.6; note that c1 ≥ 1)
⇐ (1− c) c1 ( c+ c̃√
C )τ ≥ w1/3 (since ‖µ‖ ≥ √ Cn2 log(n/δ), where C is a large enough constant).
Since C is large enough, we now choose the constant 0 < c < 1 to be such that (1 − c)/(c1(c + c̃/ √ C)) is at least (2)1/3. This proves the first part of the lemma.
To prove the second part of the lemma, note that again since [n] = P ∪N it suffices to show that
(1− c) ∑ i∈N −wi`′(t)i ≥
( c+ c̃ √ log(n/δ)
‖µ‖
)∑ i∈P −wi`′(t)i
⇐ (1− c)|N |wmin i∈N (` (t) i )
2 ≥ ( c+ c̃ √ log(n/δ)
‖µ‖
) |P|max
i∈P (`
(t) i ) 2
⇐ w ≥
( c+ c̃ √ log(n/δ)
‖µ‖ ) (1− c) maxi∈P(` (t) i ) 2
mini∈N (` (t) i )
2 · τ
⇐ w ≥
( c+ c̃ √ log(n/δ)
‖µ‖ ) (1− c) c1w 2/3τ
⇐ w1/3 ≥ c1
( c+ c̃√
C ) (1− c) · τ
⇐ w1/3 ≥ τ 21/3
where the last implication follows by the choice of c from above. This completes the proof.
We now have all the pieces required to prove our main theorem. Recall its statement.
Theorem 4.2. For any 0 < C̃ ≤ 1, there is a constant c such that, for all large enough C and for any 0 < δ < 1/C, under the assumptions of this subsection the following holds. If the weight
τ3
2 ≤ w ≤ 2τ3,
then with probability at least 1− δ, training on S produces a classifier θ̂1 satisfying:
TestError[θ̂1] = P(x,y)∼Ptest [ sign ( θ̂1 · x ) 6= y ] ≤ exp ( −c‖µ‖ 4
d
) .
Proof First, by Part (12) of Lemma B.2 we know that the data is linearly separable. Thus, by Proposition 4.1 we know that
θ̂1 = lim t→∞
θ(t)
‖θ(t)‖ .
Given this equivalence, by Lemma B.1 we know that
P(x,y)∼Ptest [ sign ( θ̂1 · x ) 6= y ]
= P(x,y)∼Ptest
[ sign ( lim t→∞ θ(t) ‖θ(t)‖ · x ) 6= y ]
≤ 1 2
[ exp ( −c′ ( lim t→∞ θ(t) · µ1 ‖θ(t)‖ )2) + exp ( −c′ ( lim t→∞ θ(t) · µ2 ‖θ(t)‖ )2)] . (24)
We shall now establish lower bounds on limt→∞ θ (t)·µ1 ‖θ(t)‖ and limt→∞ θ(t)·µ2 ‖θ(t)‖ to obtain the desired bound on the test error. Let us lower bound limt→∞ θ (t)·µ1 ‖θ(t)‖ . The bound on limt→∞ θ(t)·µ2 ‖θ(t)‖ shall follow by exactly the same logic.
By Lemma B.5 we know that
lim t→∞ µ1 · θ(t+1) ‖θ(t+1)‖
≥ lim t→∞ µ1 · θ(0) ‖θ(t+1)‖ + limt→∞ ‖µ‖2 2c1 √ d  1 1 + ‖θ
(0)‖ c1η √ d ∑t s=0 ∑ i∈[n]−wi` ′(s) i × t∑
s=0 [∑ i∈P −wi`′(s)i − c1 √ log(n/δ) ‖µ‖ ∑ i∈N −wi`′(s)i ]
(i) ≥ ‖µ‖ 2
2c1 √ d · lim t→∞  1 1 + ‖θ
(0)‖ c1η √ d ∑t s=0 ∑ i∈[n]−wi` ′(s) i × t∑
s=0 [∑ i∈P −wi`′(s)i − c1 √ log(n/δ) ‖µ‖ ∑ i∈N −wi`′(s)i ]
= ‖µ‖2
2c1 √ d · lim t→∞  1 1 + ‖θ
(0)‖ c1η √ d ∑t s=0 ∑ i∈[n]−wi` ′(s) i × lim t→∞ t∑ s=0 [∑ i∈P −wi`′(s)i − c1 √ log(n/δ) ‖µ‖ ∑ i∈N −wi`′(s)i ] , (25)
where (i) follows since µ1 · θ(0) is bounded and ‖θ(t+1)‖ → ∞ by (Ji et al., 2020, Lemma 2). We will now show that the first limit in RHS equals 1. To do this, first note that
1 ≥ lim t→∞  1 1 + ‖θ
(0)‖ c1η √ d ∑t s=0 ∑ i∈[n]−wi` ′(s) i
 =  1
1 + limt→∞ ‖θ(0)‖ c1η √ d ∑t s=0 ∑ i∈[n]−wi` ′(s) i  . Now we will show that
lim t→∞ ‖θ(0)‖ c1η √ d ∑t s=0 ∑ i∈[n]−wi` ′(s) i = 0.
First note that ‖θ (0)‖ c1η √ d ∑t s=0 ∑ i∈[n]−wi` ′(s) i ≥ 0, since −wi`′(s)i ≥ 0, so to show that this limit equals
0, it suffices to show that c1η √ d ∑t s=0 ∑ i∈[n]−wi` ′(s) i grows unboundedly as t → ∞. Using inequality (21) from above we get that,
‖θ(t+1)‖ ≤ ‖θ(t)‖+ c1η √ d ∑ i∈[n] −wi`′(t)i ≤ ‖θ(0)‖+ c1η √ d t∑ s=0 ∑ i∈[n] −wi`′(s)i .
The norm ‖θ(0)‖ is finite and we know that ‖θ(t)‖ → ∞ by (Ji et al., 2020, Lemma 2), therefore c1η √ d ∑t s=0 ∑ i∈[n]−wi` ′(s) i must grow unboundedly. This proves that
lim t→∞  1 1 + ‖θ
(0)‖ c1η √ d ∑t s=0 ∑ i∈[n]−wi` ′(s) i
 = 1
as claimed. This combined with inequality (25) yields the bound
lim t→∞ µ1 · θ(t+1) ‖θ(t+1)‖ ≥ ‖µ‖2 2c1 √ d lim t→∞
∑t s=0 [∑ i∈P −wi` ′(s) i − c1 √ log(n/δ) ‖µ‖ ∑ i∈N −wi` ′(s) i ] ∑t s=0 ∑ i∈[n]−wi` ′(s) i
(i) ≥ c2‖µ‖ 2
√ d lim t→∞
∑t s=0 ∑ i∈[n]−wi`
′(s) i∑t
s=0 ∑ i∈[n]−wi` ′(s) i
= c2‖µ‖2√
d ,
where (i) follows by invoking Lemma B.7. As stated above, using a similar argument we can also show that limt→∞ µ1·θ (t+1) ‖θ(t+1)‖ ≥ c2‖µ‖2√ d . Plugging these lower bounds into inequality (24) completes our proof.

heading: C PROOF OF THEOREM 4.3
text: In this section we will prove a lower bound on the test error of the maximum margin linear classifier θ̂MM. Here we will work with the exponential loss
`(z) = exp(−z). Define the loss L̂(θ) = ∑ i∈[n] `(yix > i θ) = ∑ i∈[n] `(z > i θ). For a step-size η and initial iterate θ(0) = 0 (this choice of initial point is for the sake of convenience, it does not affect the implicit bias of gradient descent) for any t ∈ {0, 1, . . .}
θ(t+1) = θ(t) − η∇L̂(θ(t)) be the iterates of gradient descent. We let `(t)i be shorthand for `(zi · θ(t)) and therefore,
θ(t+1) = θ(t) − η∇L̂(θ(t)) = θ(t) + η ∑ i∈[n] zi` (t) i .
The results by Soudry et al. (2018) guarantee that for a small enough step-size η the direction of the iterates of gradient descent limt→∞ θ (t)
‖θ(t)‖ = θ̂MM. Therefore, we will instead prove a lower bound on the asymptotic iterates of gradient descent.
As we did in the proof of Theorem 4.2, going forward we will assume that a good run occurs (see Definition B.3), which guarantees that all of the conditions specified in Lemma B.2 are satisfied by the training dataset S. With the setup in place, we shall now prove this theorem in stages. Throughout this section the assumptions stated in Section 4.2.1 shall remain in force. We begin with a lemma that shows that the margin of the maximum margin classifier scales with √ d/n.
Lemma C.1. There is an absolute constant c such that, on a good run, for all large enough C, for all i ∈ [n]
θ̂MM · zi ≥ c √ d
n .
Proof We will prove this result by constructing a unit vector φwith a margin that scales with √ d/n. This immediately implies that the maximum margin classifier must also attain this margin on all of the points.
Define φ to be as follows
φ :=
∑ i∈[n] zi
‖∑j∈[n] zj‖ .
We will first bound the norm of the denominator as follows∥∥∥∥∥∥ ∑ j∈[n] zj ∥∥∥∥∥∥ 2 = ∑ j∈[n] ‖zj‖2 + ∑ j 6=k zj · zk
≤ ∑ j∈[n] ‖zj‖2 + ∑ j 6=k |zj · zk|
(i) ≤ c1nd+ c1n2 ( ‖µ‖2 + √ d log(n/δ) ) = c1nd ( 1 +
c1‖µ‖2 d
+ n √
log(n/δ)√ d ) (ii)
≤ 2c1nd (26) where (i) follows by Lemma B.2 and (ii) follows since d ≥ Cn‖µ‖2 ≥ C2n3 log(n/δ) where recall C is sufficiently large.
Now we lower bound the margin between numerator of v and zk for any k ∈ [n]∑ i∈[n] zi  · zk = ‖zi‖2 +∑ i 6=k zi · zk ≥ ‖zi‖2 − ∑ k 6=i |zi · zk|
(i) ≥ d c1 − c1n(‖µ‖2 +
√ d log(n/δ))
= d
c1
( 1− c
2 1n‖µ‖2 d
− c 2 1n √
log(n/δ)√ d
)
≥ d 2c1 , (27)
where (i) again follows by invoking Lemma B.2 and by the assumption on d,
Combining inequalities (26) and (27) yields that for any k ∈ [n]
φ · zk ≥ d
2c1 √ 2c1nd = c
√ d
n .
This proves the result.
The next lemma provides control over the rate at which the norm of iterates θ(t) grows late in training. Lemma C.2. There is an absolute constant c such that, for all large enough C, if the step-size η is sufficiently small then on a good run, there exists a t0 such that for all t ≥ t0
‖θ(t+1)‖ ≥ ‖θ(t)‖+ cη √ d
n ∑ i∈[n] ` (t) i .
Proof By the definition of gradient descent ‖θ(t+1)‖2 = ‖θ(t) + η ∑ i∈[n] ` (t) i zi‖2
= ‖θ(t)‖2 + 2η ∑ k∈[n] ` (t) i zi · θ(t) + η2‖ ∑ i∈[n] ` (t) i zi‖2
≥ ‖θ(t)‖2 + 2η ∑ i∈[n] ` (t) i zi · θ(t). (28)
Now we know that θ (t) ‖θ(t)‖ → θ̂MM. Also note that in the previous lemma (Lemma C.1) we showed that for all i ∈ [n]
θ̂MM · zi ≥ c1 √ d
n .
Therefore, there exists a iteration t1 such that for all t ≥ t1 and all i ∈ [n]
θ(t) · zi ‖θ(t)‖ ≥ c1 2
√ d
n .
Continuing from (28), for any t ≥ t1 ‖θ(t+1)‖2 ≥ ‖θ(t)‖2 + 2η ∑ i∈[n] ` (t) i zi · θ(t)
≥ ‖θ(t)‖2 + c1η‖θ(t)‖ √ d
n ∑ i∈[n] ` (t) i
= ‖θ(t)‖2 1 + c1η √ d n ∑ i∈[n] ` (t) i
‖θ(t)‖  . Taking square roots we get that for any t ≥ t1
‖θ(t+1)‖ ≥ ‖θ(t)‖ √√√√ 1 + c1η √ d n ∑ i∈[n] ` (t) i
‖θ(t)‖ . (29)
Further by (Ji et al., 2020, Lemma 2) we know that ‖θ(t)‖ → ∞, so there exists a t2 such that for all t ≥ t2, ‖θ(t)‖ ≥ c1η √ nd 8 . Thus, for t ≥ t2, we have
1 ‖θ(t)‖ · c1η √ d n ∑ i∈[n] ` (t) i ≤ 8 · ∑ i∈[n] ` (t) i n ≤ 8, (30)
where the last inequality follows since the initial loss is equal to n, as θ(0) = 0, and the total loss is decreasing again by (Ji et al., 2020, Lemma 2) if the step-size is small enough. It is easy to check that for any 0 ≤ x ≤ 8
√ 1 + x ≥ 1 + x
4 .
Thus, combining inequalities (29) and (30) we get that for all t ≥ t0 = max{t1, t2}
‖θ(t+1)‖ ≥ ‖θ(t)‖ √√√√ 1 + c1η √ d n ∑ i∈[n] ` (t) i
‖θ(t)‖ ≥ ‖θ (t)‖
1 + c1η √ d n ∑ i∈[n] ` (t) i
4‖θ(t)‖  = ‖θ(t)‖+ c1
4 η
√ d
n ∑ i∈[n] ` (t) i ,
which completes our proof.
Continuing we will show that throughout training the ratio of the losses between the different examples are bounded by a constant. This ensures that each example roughly “influences” the gradient update by the same amount in each step. However, since the number of points from the positive cluster is larger, the gradient update shall overall be more highly correlated with the mean of the majority positive center µ1 than the mean of the minority negative center µ2.
The proof is identical to the proof of Lemma 11 by Chatterji & Long (2021). However, since our setting is slightly different to the setting studied in that paper we reprove the result here.
Lemma C.3. There is an absolute constant c such that, for all large enoughC, and all small enough step sizes η, on a good run, for all iterations t ∈ {0, 1, . . .} and all i, j ∈ [n]
` (t) i ` (t) j ≤ c.
Proof First note that L̂(θ(0)) = ∑ i∈[n] ` (0) i = n and since step-size η is small enough training loss is non-increasing by (Ji et al., 2020, Lemma 2).
Let c1 be the constant c ≥ 1 from Lemma B.2. We will show that c = 4c21 suffices. We shall prove this via an inductive argument. For the base case, at step t = 0, we know that θ(0) = 0, therefore the loss on all of the samples is equal to 1. Now, we shall assume that the inductive hypothesis holds at an arbitrary step t > 0 and prove that it holds at step t+ 1.
Without loss of generality, we shall analyze the ratio between the losses of the samples with indices i = 1 and j = 2. A similar analysis shall hold for any other pair. Define Gt := ` (t) 1 , Ht := ` (t) 2 , At := Ht/Gt. The ratio between these losses at step t+ 1 is
At+1 = exp(−θ(t+1) · z2) exp(−θ(t+1) · z1)
= exp(−(θ(t) + η∑k∈[n] `(t)k zk) · z2) exp(−(θ(t) + η∑k∈[n] `(t)k zk) · z1)
= At · exp −η ∑ k∈[n] ` (t) k (zk · z2 − zk · z1)  = At · exp −η `(t)2 ‖z2‖2 − `(t)1 ‖z1‖2 −∑
k 6=2
` (t) k zk · z2 + ∑ k 6=1 ` (t) k zk · z1  = At · exp −η Ht‖z2‖2 −Gt‖z1‖2 −∑
k 6=2
` (t) k zk · z2 + ∑ k 6=1 ` (t) k zk · z1  ≤ At · exp −η Ht‖z2‖2 −Gt‖z1‖2 −∑
k 6=2
` (t) k |zk · z2| − ∑ k 6=1 ` (t) k |zk · z1|  . Now note that by Lemma B.2, for all i 6= j ∈ [n], d/c1 ≤ ‖zi‖2 ≤ c1d and |zi · zj | ≤ c1 ( ‖µ‖2 + √ d log(n/δ) ) , and therefore,
At+1 ≤ At · exp −η Htd
c1 −Gtc1d− 2c1
( ‖µ‖2 + √ d log(n/δ) ) ∑ k∈[n] ` (t) k  = At · exp −η Gtd
c1
( At − c21 ) − 2c1 ( ‖µ‖2 + √ d log(n/δ) ) Gt ∑ k∈[n] ` (t) k Gt  (i)
≤ At · exp ( −η ( Gtd
c1
( At − c21 ) − 8c31 ( ‖µ‖2 + √ d log(n/δ) ) Gtn )) = At · exp ( −ηGtd
c1
( At − c21 − 8c41 ( ‖µ‖2n d + n √ log(n/δ)√ d ))) (ii)
≤ At · exp ( −ηGtd
c1
( At − c21 − 8c41 ( 1
C +
1
C ))) (iii)
≤ At · exp ( −ηGtd
c1
( At − 2c21 )) , (31)
where (i) follows since by the inductive hypothesis `(t)k /Gt ≤ c = 4c21, (ii) follows since by assumption d ≥ Cn‖µ‖2 ≥ C2n3 log(n/δ), and (iii) follows since C is sufficiently large. Now consider two cases.
Case 1 (At ≤ 2c21): By inequality (31) At+1 ≤ At · exp ( −ηGtd
c1
( At − 2c21 )) = At exp ( ηGtd
c1
( 2c21 −At )) ≤ 2c21 exp (2ηc1Gtd)
≤ 2c21 exp (2ηc1nd) ≤ 4c21,
where the last inequality follows if the step-size η is sufficiently small.
Case 2 (2c21 ≤ At ≤ 4c21 = c): Again by inequality (31) At+1 ≤ At · exp ( −ηGtd
c1
( At − 2c21 )) ≤ At ≤ 4c21.
Thus, we have shown thatAt+1 ≤ 4c21 = c in both cases. Since we assumed the induction hypothesis to hold at step t, these two cases are exhaustive, and hence the induction is complete.
The next lemma proves an upper bound on the difference between the inner product between θ(t+1) · µ2 and the corresponding inner product at iteration t. Since we start at θ(0), unrolling this over t steps gives us an upper bound on inner product θ(t) · µ2 for any t ≥ 0. Lemma C.4. There is an absolute constant c such that, for all large enough C, if the step-size η is sufficiently small then on a good run, for all t ∈ {0, 1, . . .}(
θ(t+1) − θ(t) ) · (−µ2) ≤
cη‖µ‖2 τ ∑ i∈[n] ` (t) i .
Proof By the definition of a gradient descent step (θ(t+1) − θ(t)) · (−µ2) = η ∑ i∈[n] ` (t) i zi · (−µ2)
(i) ≤ 3η‖µ‖ 2
2
∑ i∈N ` (t) i + c1η‖µ‖ √ log(n/δ) ∑ i∈P ` (t) i
(ii) ≤ 3η‖µ‖ 2 2 · c2|N | n ∑ i∈[n] ` (t) i + c1η‖µ‖ √ log(n/δ) · c2|P| n ∑ i∈[n] ` (t) i
= 3c2η‖µ‖2 2 · |P| n ( |N | |P| + √ log(n/δ) ‖µ‖ )∑ i∈[n] ` (t) i
≤ 3c2η‖µ‖2 max { |N | |P| , √ log(n/δ) ‖µ‖ }∑ i∈[n] ` (t) i ,
where (i) follows by Lemma B.2 and (ii) follows by the loss ratio bound in Lemma C.3. Since by assumption ‖µ‖2 ≥ Cn2 log(n/δ), we infer that√
log(n/δ) ‖µ‖ ≤ 1 n ≤ |N ||P| = 1 τ .
Thus, (θ(t+1) − θ(t)) · (−µ2) ≤ 3c2η‖µ‖2 max { |N | |P| , √ log(n/δ) ‖µ‖ }∑ i∈[n] ` (t) i = cη‖µ‖2 τ ∑ i∈[n] ` (t) i ,
wrapping up our proof.
With these lemmas in place we are now ready to prove our result. Let us restate it here.
Theorem 4.3. Let q ∼ N(0, Id×d). There exist constants c and c′ such that, for all large enough C and for any 0 < δ < 1/C, under the assumptions of this subsection the following holds. With probability at least 1− δ, training on S produces a maximum margin classifier θ̂MM satisfying:
TestError[θ̂MM] = P(x,y)∼Ptest [ sign ( θ̂MM · x ) 6= y ] ≥ 1 2 · Φ ( − c τ · √ n‖µ‖2√ d ) ,
where Φ is the Gaussian cdf. Furthermore, if the imbalance ratio τ ≥ c′ √ n‖µ‖2√ d
then with probability at least 1− δ
TestError[θ̂MM] ≥ 1
8 .
Proof The test error for θ̂MM is
TestError[θ̂MM] = P(x,y)∼Ptest [ sign ( θ̂MM · x ) 6= y ]
= 1
2 Pq∼N(0, 12 Ip×p)
[ sign ( θ̂MM · (µ1 + q) ) 6= 1 ] + 1
2 Pq∼N(0, 12 Ip×p)
[ sign ( θ̂MM · (µ2 + q) ) 6= −1 ] ≥ 1
2 Pq∼N(0, 12 Ip×p)
[ sign ( θ̂MM · (µ2 + q) ) 6= −1 ] = 1
2 Pq∼N(0, 12 Ip×p)
[ −θ̂MM · µ2 − θ̂MM · q < 0 ] = 1
2 Pq∼N(0, 12 Ip×p)
[ −θ̂MM · q < −θ̂MM · (−µ2) ] = 1
2 Pξ∼N(0,1)
[ ξ < −θ̂MM · (−µ2) ] (since −θ̂MM · q ∼ N(0, 1))
= Φ(−θ̂MM · (−µ2))
2 , (32)
where Φ is the Gaussian cumulative distribution function. With this inequality in place, we want to prove an upper bound on ( θ̂MM · (−µ2) )2 . Now, since
θ(t) ‖θ(t)‖ → θ̂MM, it suffices to prove an upper bound on ( limt→∞ θ(t)·(−µ2) ‖θ(t)‖ )2 instead.
To do this, going forward let us assume that a good run (see Definition B.3) occurs. Lemma B.2 guarantees that this happens with probability at least 1− δ. Let t0 be the constant from Lemma C.2, then by Lemma C.4 we have that for any t > t0
θ(t) · (−µ2) = θ(0) · (−µ2) + t0∑ s=1 (θ(s) − θ(s−1)) · (−µ2) + t∑
s=t0+1
( θ(s) − θ(s−1) ) · (−µ2)
≤ ψ + c1η‖µ‖ 2
τ
t−1∑ s=t0 ∑ i∈[n] ` (s) i , (33)
where we define ψ := θ(0) · (−µ2) + ∑t0 s=1 ( θ(s) − θ(s−1) ) · (−µ2).
On the other hand, by repeatedly applying Lemma C.2 we get that
‖θ(t)‖ ≥ ‖θ(t0)‖+ c2η √ d
n t−1∑ s=t0 ∑ i∈[n] ` (s) i . (34)
Furthermore, since ‖θ(t)‖ → ∞ (by Ji et al., 2020, Lemma 2) and
‖θ(t+1)‖ = ∥∥∥∥∥∥θ(0) + η t∑
s=0 ∑ i∈[n] ` (s) i zi ∥∥∥∥∥∥ ≤ η t∑ s=0 ∑ i∈[n] ` (s) i ‖zi‖
≤ c3ηd t∑
s=0 ∑ i∈[n] ` (s) i (by Part 3 of Lemma B.2)
= c3ηd t0−1∑ s=0 ∑ i∈[n] ` (s) i + c3ηd t∑ s=t0 ∑ i∈[n] ` (s) i
we can conclude that ∑t s=t0 ∑ i∈[n] ` (s) i →∞.
Thus combining inequalities (33) and (34) we get that
lim t→∞ θ(t) · (−µ2) ‖θ(t)‖ ≤ limt→∞
ψ + c1η‖µ‖ 2
τ
∑t−1 s=t0 ∑ i∈[n] ` (s) i
‖θ(t0)‖+ c2η √ d n ∑t−1 s=t0 ∑ i∈[n] ` (s) i = c4 τ
√ n
d ‖µ‖2.
Plugging this upper bound into inequality (32) completes the proof.

heading: D EXPERIMENTAL DETAILS
text: D.1 HYPERPARAMETERS
Imbalanced binary CIFAR10 experiments. We use the same convolutional neural network architecture with random initialization as (Byrd & Lipton, 2019). We train for 400 epochs with SGD with a batch size of 64. We chose hyperparameters that resulted in stable training for each setting. We use a constant 0.001 learning rate with 0.9 momentum for “No IW”, “IW (w)”, and “IW-ES (w)”. We use a constant 0.008 learning rate with no momentum for “IW (w3/2)”, and “IW-ES (w3/2)”.
CelebA experiments. We use a ResNet50 architecture with ImageNet initialization as done in (Sagawa et al., 2019). Due to the large dataset size, we use only 2% of the full training dataset as mentioned in the main text. Reducing the computation requirements allows us to perform statistical evaluations of the results over sufficiently many seeds. We train for 100 epochs with SGD with a batch size of 64. We chose hyperparameters that resulted in stable training for each setting. We use a constant 0.0004 learning rate with 0.9 momentum for all settings.
D.2 IMPORTANCE WEIGHT COMPUTATION.
Here we explain in greater detail how we compute the importance weights. Note that the scale of the importance weights is important in ensuring stable training. Let there beG groups of subpopulations in the training and testing set. G = 2 for binary CIFAR10 and G = 4 for CelebA. Let n1, . . . , nG be the counts of each group in the training set. Let n be the total number of training examples. In our case, the test set is balanced across the groups. Our procedure for computing weights is:
1. For each example i that belongs to group g, compute the unbiased weight: wi ← 1/ng . 2. Exponentiate the weights by c if necessary: wi ← wci . No importance weighting corre-
sponds to c = 0. Unbiased importance weighting corresponds to c = 1.
3. To adjust for the scale of weights, normalize by the average exponentiated weight across the full training set: wi ← wi/ ∑n j=1 wi. Only normalizing across each minibatch can
result in unstable training, especially when a minibatch contains no representatives from a group.
D.3 EARLY-STOPPING METRIC.
We use the checkpoint with the highest importance weighted validation accuracy when early stopping. We compute separate importance weights of the validation set with respect to the test set, and reweight the validation accuracy using unbiased weights, even when the training weights are biased. Our procedure allows for the situation where the validation set is not exactly the same distribution as the training set.
D.4 EXACT NUMERICAL RESULTS OF OUR EXPERIMENTS

Review summaries:
The paper shows both theoretically and experimentally that importance weighting is not incompatible  with the training of overparameterized models provided that the training loss is appropriately modified so that it does not have an exponential tail. Specifically, for binary classification the authors propose a new loss function with polynomial tail decay. Theoretically, the new loss is shown to outperform weighted cross entropy for linear models and an imbalanced mixtures of gaussian model. Empirically, the new loss outperforms weighted cross entropy on imbalanced CIFAR 10 and the CelebA dataset. 

Previous works pointed out an interesting incompatibility between importance sampling and training overparameterized neural nets. In the view of margin, Soudry et al. (2018) showed that GD on exponentially-tailed loss converges to the max-margin classifier, and Xu et al. (2020) pointed out that this margin maximization is unaffected by importance sampling. This paper then turns to studying polynomially-tailed loss.  * This paper establishes the implicit bias by formulating an optimization problem, which can change according to the importance weights. The convergent classifier is shown to generalize better than the exponentially-tailed loss in a simple linear classification setting, where the dataset contains two sub-Gaussian clusters. * Experiments on Imbalanced Binary CIFAR10 and CelebA show that the polynomially-tailed loss performs better on imbalanced datasets.

In this paper, the author investigates whether importance weighting is incompatible with the training of overparameterized neural networks. In contrast to the recent observation that  importance weighting is ineffective in current deep learning paradigm (Byrd & Lipton 2019), the authors show that it could actually be helpful using polynomially-tailed losses. Both theoretical justifications and empirical evidences are provided to support the claim.