Focus on the Common Good: Group Distributional Robustness Follows

We consider the problem of training a classification model with group annotated training data . Recent work has established that , if there is distribution shift across different groups , models trained using the standard empirical risk minimization ( ERM ) objective suffer from poor performance on minority groups and that group distributionally robust optimization ( Group-DRO ) objective is a better alternative . The starting point of this paper is the observation that though Group-DRO performs better than ERM on minority groups for some benchmark datasets , there are several other datasets where it performs much worse than ERM . Inspired by ideas from the closely related problem of domain generalization , this paper proposes a new and simple algorithm that explicitly encourages learning of features that are shared across various groups . The key insight behind our proposed algorithm is that while Group-DRO focuses on groups with worst regularized loss , focusing instead , on groups that enable better performance even on other groups , could lead to learning of shared/common features , thereby enhancing minority performance beyond what is achieved by Group-DRO . Empirically , we show that our proposed algorithm matches or achieves better performance compared to strong contemporary baselines including ERM and Group-DRO on standard benchmarks on both minority groups and across all groups . Theoretically , we show that the proposed algorithm is a descent method and finds first order stationary points of smooth nonconvex functions . Our code and datasets can be found at this URL . 1 INTRODUCTION . It is by now well known ( Sagawa et al. , 2020a ) that , in several settings with subpopulation shift , large neural networks trained using empirical risk minimization ( ERM ) objective heavily sacrifice prediction accuracy on small groups/subpopulations of the data , in order to achieve a high overall prediction accuracy . One of the key reasons for this behavior of ERM is the existence of spurious correlations between labels and some features on majority groups which are either nonexistent , or worse oppositely correlated , on the minority groups ( Sagawa et al. , 2020a ) . In such cases , ERM exploits these spurious correlations to achieve high accuracy on the majority groups , thereby suffering from poor performance on minority groups . Consequently , this is a form of unfairness , where accuracy on minority subpopulation or group is being heavily sacrificed to achieve near perfect accuracy on majority groups . More concretely , we are given training examples , stratified non-uniformly in to multiple subpopulations , also called groups . Our goal is to train a model that generalizes to all training groups without severely sacrificing accuracy on the minority groups . This problem , which is also known as the sub-population shift problem , was popularized by Sagawa et al . ( 2020a ) , and is well-studied in the literature ( Sagawa et al. , 2020a ; b ; Menon et al. , 2021 ; Zhai et al. , 2021 ; Goel et al. , 2020 ; Bao et al. , 2021 ) . Among all algorithms proposed for this problem , one of the most influential ones is the Group Distributionally Robust Optimization ( Group-DRO ) method ( Sagawa et al. , 2020a ) , which at every update step focuses on the group with the highest regularized loss . While Group-DRO has been shown to obtain reduction in worst-group error over ERM on some benchmark datasets , it has obvious but practically relevant failure modes such as when different groups have different amounts of label noise . In such a case , Group-DRO ends up focusing on the group ( s ) with the highest amount of label noise , thereby obtaining worse performance on all groups . In fact , on several other datasets with subpopulation shift , Group-DRO performs poorly compared to ERM ( Koh et al. , 2020 ) . The key issue with Group-DRO ’ s strategy of just focusing on the group with the highest training loss is that , it is uninformed of the inter-group interactions . Consequently , the parameter update using the group with the highest training loss may increase the loss on other groups . In fact , this observation also reveals that Group-DRO is not properly addressing the issue of spurious correlations , which have been identified by prior work as a key reason behind the poor performance of ERM ( Sagawa et al. , 2020a ) . Inspired by ideas from the closely related problem of domain generalization ( Sun & Saenko , 2016 ; Shankar et al. , 2018 ; Arjovsky et al. , 2019 ) , we hypothesize that modeling intergroup interactions is essential for properly addressing spurious correlations and sub-population shift problem . More concretely , in each update step , we propose the following simple strategy to decide which group to train on : Train on that group whose gradient leads to largest decrease in average training loss over all groups . In other words , our strategy focuses on the group that contributes to the common good . We call the resulting algorithm Common Gradient Descent ( CGD ) . We show that CGD is a sound optimization algorithm as it monotonically decreases the macro/group-average loss , and consequently finds first order stationary points . CGD ’ s emphasis on common gradients that lead to improvements across all groups , makes it robust to the presence of spurious features , enabling it to achieve better generalization across all groups compared to usual gradient descent . We present insights on why CGD may perform better than Group-DRO by studying simple synthetic settings . Subsequently , through empirical evaluation on seven real-world datasets–which include two text and five image tasks with a mix of sub-population and domain shifts , we demonstrate that CGD either matches or obtains improved performance over several existing algorithms for robustness that include : ERM , PGI ( Ahmed et al. , 2021 ) , IRM ( Arjovsky et al. , 2019 ) , and Group-DRO . 2 METHOD . Problem statement Let X and Y denote the input and label spaces respectively . We assume that the training data comprises of k groups from a set G where each i ∈ G include ni instances from a probability distribution Pi ( X , Y ) . In addition to the label yj ∈ Y , each training example xj ∈ X is also annotated by the group/subpopulation i ∈ G from which it comes . The group annotations are available only during training and not during testing time , hence the learned model is required to be group-agnostic , i.e. , it may not use the group label at test time . The number of examples could vary arbitrarily between different groups in the train data . We refer to groups { i } with ( relatively ) large ni as majority groups and those with small ni as minority groups . We use the terms group and sub-population interchangeably . Our goal is to learn a model that performs well on all the groups in G. Following prior work ( Sagawa et al. , 2020a ) , we use two metrics to quantify performance : worst group accuracy denoting the minimum test accuracy across all groups i ∈ G and average accuracy denoting the average test accuracy across all examples belonging to all groups . We denote with ` i , the average loss over examples of group i using a predictor fθ , ` i ( θ ) = E ( x , y ) ∼Pi ( X , Y ) L ( x , y ; fθ ) , for an appropriate classification loss L , a parameterized predictor fθ with parameter θ . We refer to the parameters at step ‘ t ’ of training by θt . The ERM training algorithm learns parameters over the joint risk , which is the population weighted sum of the losses : ∑ i ni ` i/ ∑ i ni . However , as described in Section 1 when there are spurious correlations of certain features in the majority groups , ERM ’ s loss can not avoid learning of these spurious features . This leads to poor performance on minority groups if these features are either nonexistent or exhibit opposite behavior on the minority groups . A simple alternative , called ERMUW , is to reshape the risk function so that the minority groups are up-weighted , for some predefined group-specific weights α. ERM-UW while simple , however , could overfit to the up-weighted minority groups . Sagawa et al . ( 2020a ) studied an existing robust optimization algorithm ( Duchi & Namkoong , 2018 ) in the context of deep learning models . They proposed a risk reshaping method called Group-DRO , which at any update step trains on the group with the highest loss , as shown below . Group-DRO update step : j∗ = arg max j ` j ( θ ) and θt+1 = θt − η∇θ ` j∗ . Training on only high loss groups avoids overfitting on the minority group since we avoid zero training loss on any group while the average training loss is non-zero . However , this approach of focusing on the worst group is subject to failure when groups have heterogeneous levels of noise and transfer . In Section 4 , we will illustrate these failure modes . Instead , we consider an alternate strategy for picking the training group : the group when trained on , most minimizes the overall loss across all groups . Let gi = ∇θ ` i represent the gradient of the group i , we pick the desired group as : j∗ = arg min j ∑ i ` i ( θ t − η∇θ ` j ( θt ) ) ≈ arg max j ∑ i gTi gj [ First-order Taylor ] ( 1 ) The choice of the training group based on gradient inner-product can be noisy . To counteract , we smooth the choice of group with a weight vector αt ∈ ∆k−1 , at step t , and also regularize between consecutive time steps . The amount of regularization is controlled by a hyperparameter : ηα . Our parameter update at step t+ 1 takes the following form . αt+1 = arg max α∈∆k−1 ∑ i αi〈gi , ∑ j gj〉 − 1 ηα KL ( α , αt ) ( 2 ) θt+1 = θt − η ∑ j αt+1j gj ( θt ) . The update of α in equation 2 can be solved in closed form using KKT first order optimality conditions and rearranging to get : αt+1i = αti · exp ( ηα〈gi , ∑ j gj〉 ) ∑ s α t s · exp ( ηα〈gs , ∑ j gj〉 ) . ( 3 ) Pseudocode for the overall algorithm is shown in Algorithm 1 . Algorithm 1 CGD Algorithm 1 : Input : Number of groups : k , Training data : { ( xj , yj , i ) : i ∈ [ k ] , j ∈ [ ni ] } , Step sizes : ηα , η 2 : Initialize θ0 , α0 = ( 1 k , · · · , 1 k ) 3 : for t = 1 , 2 , · · · , do 4 : for i ∈ { 1 , · · · , k } do 5 : αt+1i ← αti exp ( ηα∇ ` i ( θt ) > ∑ s∈ [ k ] ∇ ` s ( θt ) ) 6 : end for 7 : αt+1i ← α t+1 i /‖αt+1‖1 ∀i ∈ [ 1 . . . k ] . Normalize 8 : θt+1 ← θt − η ∑ i∈ { 1 , ··· , k } α t+1 i ∇ ` i ( θt ) . Update parameters 9 : end for The scale of the gradients can vary widely depending on the task , architecture and the loss landscape . As a result , the α update through equation 3 can be unstable and tuning of the ηα hyperparam tedious . Since we are only interested in capturing if a group transfers positively or negatively to others , we retain the cosine-similarity of the gradient dot-product but control their scale through ` ( θ ) p for some p > 0 . That is , we set the gradient : ∇ ` i ( θt ) to ∇ ` i ( θ t ) ‖∇ ` i ( θt ) ‖ ` i ( θ ) p. In our implementation , we use p = 1/2 . In Appendix A , we discuss in more detail the range of gradient norm , why we scale the gradient by loss , and how we pick the value of p. Finally , the αt+1 update of equation 3 is replaced with equation 4 . When we assume the groups do not interact , i.e . 〈gi , gj〉 = 0 ∀i 6= j , then our update ( equation 4 ) matches Group-DRO . αt+1i = αti exp ( ηα ∑ j √ ` i ` j cos ( gi , gj ) ) ∑ s α t s exp ( ηα ∑ j √ ` s ` j cos ( gs , gj ) ) . ( 4 ) Group Adjustment The empirical train loss underestimates the true loss by an amount that is inversely proportional to the population size . Owing to the large group population differences , we expect varying generalization gap per group . Sagawa et al . ( 2020a ) adjusts the loss value of the groups to correct for these generalization gap differences as ` i = ` i + C/ √ ni , ∀i. , where C > 0 is a hyper-parameter . We apply similar group adjustments for CGD as well . The corrected loss values simply replace the ` i of Line 5 in Algorithm 1 .

Summary: The paper proposes a new method for robust ML under distribution shifts. Past work has looked at formulations that minimize the worst group error. This paper adds a new twist on it and instead argues for focusing on the group that leads to the greatest decrease in average training error for all the groups. This intuition is combined into an algorithm and the paper proves that though their proposed algorithm doesn't minimize a specific loss function, it still finds first-order-stationary points. The results are shown on several synthetic datasets as well as on the WILDS Robust ML benchmark that show the superior performance of the proposed algorithm over several baselines.    Main Contributions:   1). The paper proposes a new approach for robust ML under distribution shifts that performs gradient descent not on the group with worst error but on the group which decreases the average error of all other groups.  2). Results are shown on synthetic and real-world datasets which show the superior performance of the proposed method in achieving group robustness.