Understanding Overfitting in Reweighting Algorithms for Worst-group Performance

1 INTRODUCTION . 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 can not 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 . 1.1 RELATED WORK . 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 . 2 PRELIMINARIES . 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 ) 2.1 REWEIGHTING ALGORITHMS . 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 .

Summary: 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).