Focus on the Common Good: Group Distributional Robustness Follows
Authors: Vihari Piratla
Praneeth Netrapalli
Sunita Sarawagi
Venue: ICLR 2022 Poster (2022)
Keywords: sub-population shift
robust optimization
domain generalization

Abstract: 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.

Published as a conference paper at ICLR 2022
FOCUS ON THE COMMON GOOD: GROUP DISTRIBU- TIONAL ROBUSTNESS FOLLOWS
Vihari Piratla*, Praneeth Netrapalli", and Sunita Sarawagi'
! Indian Institute of Technology, Bombay Google Research, India
ABSTRACT
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 minority groups and that group distributionally robust optimization (Group-DRO) objective a better alternative. The starting point of this paper is the observation that though Group-DRO per- forms 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 in- stead, on groups that enable better performance even 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 contempo- rary baselines including ERM and Group-DRO on standard benchmarks minority groups and across all groups Theoretically, we show that the proposed algorithm is descent method and finds first order stationary points of smooth nonconvex functions. Our code and datasets can be found at this URL.
both
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 small groups/subpopulations of the data, in order achieve 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 majority groups which are either nonexistent, or worse oppositely correlated, on the minority groups (Sagawa et al., 2020a). In such cases, ERM ex- ploits these spurious correlations achieve high accuracy on the majority groups, thereby suffering from poor performance on minority groups. Consequently, this is a form of unfairness, where accu- racy on minority subpopulation or group is being heavily sacrificed achieve near perfect accuracy on majority groups.
More concretely, we are given training examples, stratified non-uniformly in to multiple sub- populations, 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 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-DRe) DRO) method (Sagawa et al., 2020a), which at every update step focuses on the group with the highest regularized loss. While Group-DRO has
ratt 1.com

---

Published as a conference paper at ICLR 2022
been shown to obtain reduction in gorst-group error over ERM 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 ERM (Koh et al., 2020).
The key issue with Group-DRO ' strategy of just focusing on the group with the highest training loss that, it 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 inter- group interactions essential for properly addressing spurious correlations and coll-ooulation 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 naacoogroup-avvrrae loss, and consequently finds first order stationary points. CGD's emphasis all groups, makes it robust to the presence of spurious features, enabling it to achieve better gener- alization across all groups compared 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 Batasets-which include two text and five image tasks with a mix of sub-populle poppllatouoonnn and domain shifts, we demonstrate that CGD either matches or obtains im- proved performance over several existing algorithms for robustness that include: ERM, PGI (Ahmed et al., 2021), IRM (Arjovsky et al., 2019), and Group-DRO.
common gradients that lead to improvements across
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 E G include n; instances from a In addition to the label yj E V each training example 2- E X
probability distribution also annotated by the group'subpopulation i E G from which comes. The group annotations available only during training and not during testing time, hence the learned model is required to be group-agnostic, i.c., 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 n; as majority groups and those with small n, as minority groups We use the terms group and uub-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 E G and average accuracy denoting the average test accuracy across all examples belonging to all groups
We denote with b, the average loss over examples of group using a predictor fa. (,(0) ((x.)))-,ixx)) C(x, yi fe), for an appropriate classification loss L, a parameterized predictor fo with parameter A. We refer to the parameters at step 't' of training by or.
The ERM training algorithm learns parameters over the joint risk, which the population weighted sum of the losses: E, n,f4/ Σ, n, However, as described in Section 1 when there are spurious correlations of certain features in the majority groups, ERM's loss cannot 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 ERM- UW, is to reshape the risk function so that the minority groups are up- weighted, for some predefined group-specific weights a. ERM-UW while simple, however, could overfit to the up- weighted mi- nority 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 GRo,p-DRO, which at any update step trains on the group with the highest loss, shown
2

---

Published as a conference paper at ICLR 2022
below.
Group-DRO update step: j arg max (5) and 0'-1 -0 Volg
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 9i = Vol, represent the gradient of the group we pick the desired group as:
j" arg min Σ Eello ((00~~)oe,((0)))
arg max
Σ oh 95 [First-order Taylor]
(1)
The choice of the training group based on gradient inner-produce can be noisy. To counteract, we smooth the choice of group with a weight vector a' E A 1-11, 1 at step and also regularize between consecutive time steps. The amount of regularization controlled by hyperparameter: !/a- Our parameter update at step i - takes the following form.
a ++1 arg max
Σ Oi (gis Σgi)
1 _K((aaa)) 1a
(2)
aE k-1
ge+ 0¹ Σ aj
(+1 (01).
The update of a in equation 2 can be solved in closed form using KKT first order optimality condi- tions and rearranging to get:
af . exp Nh (gis E, 93))
0 1+1
(3)
E. at exp na (gx. E, 93)
Pseudocode for the overall algorithm is shown in Algorithm 1.
Algorithm 1 CGD Algorithm
1: Input: Number of groups: k. Training data: {(zj, Mj i)
[k). E [n:]], Step sizes: na, 77
2: Initialize 80, a (t....,
3: for // 1,2, ... do
4: for i
k) do 5: a; 441 at exp(1a (B.( EVE(s) (110))TEEE (e)(
6:
end for
7:
a .+1 t- a; ++1 /|a
Vi E [1...K .k] ++1
> Normalize
8:
gh+l t 8t Eiel,,.. (4
oz (e,(t)
Update parameters
9:
end for
The scale of the gradients can vary widely depending on the task, architecture and the loss landscape. As result, the a update through equation can be unstable and tuning of the Je hyperparam dious. 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 {(0) for
[V (,))), In our implementation, p = 1/2. In Appendix A, we discuss in more detail the range gradient norm, why we scale the gradient by loss, and how we pick the value of Finally, the "+1 update of equation 3 is replaced with equation 4. When we assume the groups do not interact, i.e. (9t,9,) update (equation 4) matches Group-DRO.
some p 0. That is, set the gradient: (e*) to ve,(o")
Vi + j. then our
exp((n E, Vell cos( 9;))
4+
Σ. of exp(no E, Vlll, cos 9;))
3

---

Published as a conference paper at ICLR 2022
Group Adjustment The empirical train loss underestimates the true loss by amount that is in- versely 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 f = lg + C/ n vi., where C > hyper-parameter. We apply similar group adjustments for CGD well. The corrected loss values simply replace the & of Line 5 in Algorithm 1.
3 CONVERGENCE ANALYSIS
In this section, we will show that CGD is a sound optimization algorithm by proving that it mono- tonically decreases the function value and finds first order stationary points (FOSP) for bounded, Lipschitz and smooth loss functions. We now define the notion of e-FOSP which is the most com- mon notion of optimality for general smooth nonconvex functions.
Definition 1. A point 0 is said to be an c-FOSP of a differentiable function f(-) if S.
In the context of our paper, we consider the following cumulative loss function: R(0) = E (0). This is called the macro/group-aveere loss. We are now ready to state the conver- gence guarantee for our algorithm.
Theorem Suppose that (i) each f((-) is G-Lipschitz, (ii) R(-) is L-smooth and (iii) R(-) is bounded between -B and B. Suppose further that Algorithm 1 is run with = =22 LG and
7a = V G.T Then, Algorithm I will find an c-FOSP of R(0) in O (BLC²) iterations.
In other words, the above result shows that Algorithm finds an e-FOSP of the cumulative loss in O (c-1) iterations. We present high level outline of the proof here and present the complete proof of Theorem in Appendix B.
Proof outline of Theorem 1. Considering the tt iteration where the iterate is l and mixing weights are a.. Using the notation in Section 2, let us denote R(0) = 1/ Σ, (0),, 9i = v(((0)) and
E. Fr. The update of our algorithm is given by:
a,t1 af exp (No (g6,g)
(5)
Z
@++1 = 0¹ - Σ
++ Di.
(6)
where Z -= exp (n<lo,,))) Let us fix a" (1/k. 1/k) E R² and use KL(p, g) E. P: log 2: to denote the KL-divergence between and g. Noting that the update equation 5 on corresponds to mirror descent steps the function (g. Ex ongg)) and using mirror descent analysis, we obtain:
=
gi
KL(a", < KLla", Cainokonn)) + 2
-nall
+ - < -111 KL a²) - KL(a", aa(a) + G4
na
Using monotonicity of the exp function, we further show that Σ, a a 4 (g.. Using smoothness of R and update equation 6, we then show:
R((++1) < R(00)) |11ا|| + KL(a", "a'""aa)= (o)--KLE -KLa"",""
1 L22
no 0+
no 11 R(0') R(0++1) KT(a", 00) (os(a",,''-K((0)) "Kl(a"aa't 7
G4 nT.2
n7
2
Summing the above inequality over timesteps t 1,..,, T, and using the parameter choices for and 7) proves the theorem.
4

---

Published as a conference paper at ICLR 2022
4 QUALITATIVE ANALYSIS
In this section, we look at simple multi-group training scenarios to derive insights on the difference between Group-DRO and CGD. We start with simple setup where one of the groups is noisy in Section 1. In Section 4.2, we look at a setting with uneven inter-group transfer. In Section 4.3, we study a spurious correlation setup. For all the settings, we train linear binary classifier model for 400 epochs with batch SGD optimizer and learning rate 0.1. Unless otherwise stated, the number of input dimensions two and the two features æ1; 2g are sampled from a standard normal, with the binary label y [x1 22 > O). and number of groups is three. For each setting we will inspect the training weight (aς) assigned by each algorithm per group i at epoch (Figure 1, 1,3, 5). The plots are then used to draw qualitative judgements. See Appendix C for descriptive plots.
We also present parallel settings with MNIST images and deep models (ResNet-18) to stress that observations more broadly applicable. Appendix D presents the details of this setup.
4.1 LABEL NOISE SETUP: Group-DRO focuses on noisy groups.
We induced label noise in the first group, only during training, by flipping labels for randomly picked 20% examples. The first and second group formed the majority with population of 450 examples each and the minority last group had 100 examples. The noisy first group and the two subsequent groups are referred to Noisy- Majority, Clcan-Majority, and Clean-Minority. Due noise in the first group, and since we cannot overfit in this setup, the loss on the first group was consistently higher than the other two. As result, Group-DRO trained only on the first group (Fig la), while CGD avoided overly training the noisy first group (Fig lb). This results in not only lower worst-case error but more robust training Simple column, Table 1). The variance of the Iearned parameters across six runs drops from 1.88 to 0.32 using CGD.
1.0
Nasy muuuty
majarty
0.5
0.5
0.0
0
100
200
300
0.0
0
100
200
300
(a) Group-DRO
(b) CGD
Figure Noise-Simple. Training weight (a) of groups number of epochs.
Table (Noise-MNIST column) shows the same trend under similar setup on the MNIST dataset with the ResNet-18 18 model. These experiments demonstrate that -RO's method of focusing on the group with the worst training loss is sub-optimal when groups have heterogeneous label noise training difficulty. CGD's weights that are informed by inter-group interaction provide greater in- built robustness to group-specific noise.
4.2 UNEVEN INTER- GROUP SIMILARITY: CGD focuses central groups
Here we simulated setting such that the opti- mal classifier for the first, third group is clos- est to the second. The label (y) was set to a group-specific deterministic function of the in- put (T1, x2); for the first group the mapping function was y [[[ > (which was rotated by 30 degrees for the two subsequent groups), for the second group was y [00.8721 + 0.5x² > 0), and for the third group it was y = T(0.5z. + 0.87.a > 0). The angle be- tween the classifiers is proxy for the distance between the groups. We refer to the groups in the order as Left, Center, Right, and their spective training population was: 499, 499, 2.
1.0
1.0
0.0
200
400
200
400
(a) Group-DRO
(b) CGD
Figure 3: Rot-Simple. Training weight (a) of
groups vs number of epochs.
The optimal classifier that generalizes equally well all the groups is the center classifier. Group- DRO, CGD, differed in how they arrive at the center classifier: Group-DRO assigned all the training mass equally the Left and Right groups (Figure 2a), while CGD trained, unsurprisingly, mostly
5

---

Published as a conference paper at ICLR 2022
on the Center group that has the highest transfer (Figure 2b). Group-DRO up- up-weighted the noisy minority group (2 examples) and is inferior when compared with CGD's ' strategy, which is reflected by the superior performance on the test set shown in -initiion-Smplee column of Table 1. On the MNIST dataset we repeated a similar experiment by rotating digits by O, 30, and 60 degrees, and found CGD to provide significant gains over Group-DRO (Rotation-MNIST column of Table 2). These experiments demonstrate the benefit of focusing on the central group even for maximizing worst case accuracy.
4.3 SPURIOUS CORRELATIONS SETUP: CGD focuses on groups without spurious correlations
Here we create spurious correlations by adding third feature that takes different values across the three different groups whose sizes are 490, 490, and 20 respectively. The third feature was set to the value of the label y on the first group, 1-y the third group, and set y 1-y w.p. 0.6 on the second group. We perturbed the first two features of the first group examples such that they predict label correctly only for 60% of the examples. Therefore, the net label cor- relation of the first two and the third feature is about the same: 0.8. We refer to the first and last groups Spurious-majone feature-and the second as the clean majority.
1.0
1.0
Sou meleray
Chin-nuuortyy
Sourrinor
0.5
0.5
0.0
400
0.0
200
400
Group-DRO
(b) CGD
Figure 5: Spurious-Simple. Training weight
(a) of groups vs number of epochs.
minority- they both contain an active spurious third
Group-DRO balances training between the two spurious groups (Figure 4a), while CGD focuses mostly on the clean group (Figure 4b). Both the training strategies could avoid the spurious third feature, however, the training strategy of Group-DRO DRO is sub-optimal since up-weight the small, hence noisy, spurious minority group. The better test loss of CGD in Table column) further verifies the drawback of Group-DRO's training strategy. On the MNIST dataset too we observed similar trends by creating spurious correlation using colors of the digit (Spurious- MNIST column of Table 2).
Sspurious-Simpp
More generally, whenever relatively clean group with low spurious correlation strength is present, and when the train groups represent both positive and negative spurious correlations, we expect CGD to focus on the group with the least spurious correlation since is likely to have greater inter-group similarity.
Task
Noise-Simple
Rotalion-Simple
Spurious-Simple
Algi
Variance
Worst Loss
Variance
Worst Loss
Variance
Worst Loss
Group-DRO 1.88
0.35 (0.03)
0.41
0.77 (0.14)
0.17
0.70 (0.16)
CGD
0.32
0.25 (0.02)
0.08
0.59 (0.05)
0.04
0.43 (0.06)
Table 1: Comparison of Group-DRO, CGD across the three simple settings. Worst loss is the worst group binary cross entropy loss on the test set, and averaged over six seeds, shown in parenthesis is the standard deviation. Variance column shows the variance of Loc normalized solution across the six runs. CGD has lower variance and test loss when compared with Group-DRO.
Alg. + DRO
Noise-MNIST
Rotation-MNIST
Spurious-MNIST
Group- trouppDROO CGD
(1.0), 85.9 (1.0)
90.5 (0.2), 80.6 (1.1)
95.4 (0.2), 95.0 (0.4)
88.9 (1.0), 88.8 (1.0)
92.1 (0.6), 85.0 (0.7)
96.7 (0.1), 96.4 (0.3)
Table 2: Average, worst group accuracy on the test split. Shown in parenthesis is the standard deviation. All numbers are aggregated over three runs. More details of the setup in Appendix

---

Published as a conference paper at ICLR 2022
5 EXPERIMENTS
We compare CGD with state of the art methods from group-distribttin shift and domain general- ization literature. We enumerate datasets in Section 5.1, discuss the baselines and present imple- mentation details Section 5.2, and finally present our results in Section 5.3.
5.1 DATASETS
We evaluated on eight datasets, which include two synthetic datasets with induced spurious cor- relations: ColSTedMNIST. WaterBirds; two real-world datasets with known spurious correlations: CelebA, MultiNLI; four WILDS (Koh al., 2020) datasets with a mix of sub- cobb-ppuuatiinn and do- main shift. Table 3 summarises all the eight datasets. We describe all the datasets in Appendix E. For more details on WaterBirds, CelebA, MultiNLI, we point the reader to the original paper (Sagawa et al., 2020a), and Koh et al. (2020) for details on the WILDS datasets. We discuss how some of the real-world datasets relate to our synthetic setting in Appendix H.
Dataset
labels
# groups
Population type
Type
Size
Worst ratio
Colored MNIST
2
2
Labelx
Image
50K
1000
WaterBirds
2
Label Label
Group
Image
4.8K
62.5
CelebA
2
2
Label LabelxGroup
Image
162K
51.6
MultiNLI
3
2
Label x Group
Text
200K
44.3
NovilCommeett-WILL.D
2
LabelxGroup
Text
270K
74.5
PovertyMap-WiIe -ILDS
real
13
Group
Image
10K
5.9
FMoW-WILDS Camelyon17-WIILDS
62
11
Group
Image
80K
16.0
2
3
Group
Image
3K
2.5
Table 3: Summary of the datasets we used for evaluation. The Size column shows the number of training instances. The Worst ratio column shows the ratio of the size of the largest train group to the smallest; worst ratio reflects loosely the strength of the spurious correlations.
5.2 EXPERIMENT DETAILS
Baselines
ERM: Simple descent on the joint risk. ERM-UW, Group-DRe DRO: risk reshaping baselines described in Section With ERM-UW, instead of simply up-weighting all the groups equally, set the group weighting parameter a; of a group as a function of its size. For some non-negative constant C, og x exp(C/ /n.).
exp(C///m.).
PGI (Ahmed et al., 2021), predictive group invariance, penalizes the divergence between predictive probability distributions among the training groups, and was shown to he effective for in-domain as well as out-domain generalization. We provide more details about the algorithm Appendix F.
On the WILDS dataset, we also compare with two algorithms from Domain Generalization litera- ture: CORAL (Sun & Saenko, 2016), IRM (Arjovsky al., 2019). Domain Generalization methods are relevant since they could generalize to any domain including the seen train domains.
Implementation Details and Evaluation Metric: We use the codestack' released with the WILDS (Koh et al., 2020) dataset as our base implementation. Our code can be found at this URL. We report the average and standard deviation of the evaluation metric from at least three runs. not otherwise stated explicitly, we report worst accuracy over the groups evaluated on the test split.
Hyperparameters: We search C: the group adjustment parameter for Group-DRO, CGD, and group weighting parameter for ERM-UW, over the range [0, 20). We tune C' only for WaterBirds, CelebA and MultiNLI datasets for fair comparison with Group-DRO. For other dataset, we set to 0 because group adjustment tuning did not yield large improvements size parameter of Group- ErouppDROO, CGD, and À penalty control parameter of PGI, is picked from {1, 0.1, lc-2, 1e-3). We follow the same learning procedure of Sagawa et al. (2020a) Colored- WaterBirds, CelebA, MultiNLI, datasets, we pick the best learning rate parameter from ( (le-3, le- 5). weight decay from [1e4,, .1, 1}. use SGD optimizer, and set the batch size to 1024, 128, 64, 32
many datasets. The step
'https:/ // thub. com/ / p- ambda /wi 1ds

---

Published as a conference paper at ICLR 2022
respectively, On WILDS datasets, we adopt the default optimization parameters?, and only tune the step-size parameter.
We report the test set's performance corresponding to the best hyperparameter(s) and epoch on the validation split. We use standard train-validation-tess splits for all the datasets when available.
Backbone Model
We followed closely the training setup of Sagawa et al. (2020a); Koh et al. (2020). On WaterBirds, CelebA dataset, we used ResNet50; on Colored-MNIST, PavertyMa,, used ResNet18; on Cem- lyon17, FMoW, we used DenseNet121. All the image models except Colored-MNIST (which trains on 28 x 28 images), Poverty.Map (which deals with multi-spectral version of ResNet18), 8), are pre- trained on ImageNet. MultiNLI, CivilComments, use pretrained uncased DistilBERT-bree base model.
5.3 RESULTS
Table 4 shows the worst group test split accuracy for the four standard (non-WILDDS) sub-population shift datasets. For all the tasks shown, ERM performs worse than a random baseline on the worst group, although the average accuracy is high. ERM-UW is a strong baseline, and improves the worst group accuracy of ERM on three of the four tasks, without hurting much the average accuracy. PGI is no better than ERM or ERM-UW on all the tasks except on Golored-MNIST. Group-DRe DRO improves worst accuracy on most datasets, however, CGD fares even better. CGD improves worst- group accuracy on all the datasets over ERM. Except on MultiNLI text task, the gains of CGD are significant over other methods, and on MultiNLI task the worst group accuracy of CGD is at least as good as Group-DRO. On Colored-MNIST, which has the highest ratio of majority to minority group size, the gains of CGD are particularly large.
We report comparisons using the four WILDS datasets Table 5. We show both in-domain (ID) and out-domain (OOD) generalization performance when appropriate, they are marked in the third row. All the results are averaged over multiple runs; FMoW numbers are averaged over three seeds, CivilComments over five seeds, Camelyon17 over ten seeds, and PovertyMap over the five data folds. Confirming with the WILDS standard, we report worst group accuracy for FMoW, Civil- Comments, worst region Pearson correlation for Pap, average out-of-domain accuracy for 1amelyon17. We make the following observations from the results. ERM is surprisingly strong on all the tasks except CivilComments. Strikingly, Group-DRO is worse than ERM on four of the five tasks shown in the table, including the in-domain cuu-poopulation shift) evaluation on PovertyMap task. CGD is the only algorithm that performs consistently well across all the tasks. The results suggest CGD is significantly robust to sub- uub-population shift, and performs no worse than ERM on domain shifts. Further, we study Colored-MNIST dataset under varying ratio of majority to minority group sizes in Appendix G and demonstrate that CGD is robust to sub-population shifts even under extreme training population disparity.
Alg. +
CMNIST
WaterBirds
CelebA
MultiNLI
ERM
50 (0.0)
85.2 (0.8), (1.4) 95.2 (0.2), 45.2 (0.9) 81.8 (0.4), 69.0 (1.2)
ERM-UW
53.6 (0.1), 7.2 (2.1)
91.8 (0.2), 86.9 (0.5) 92.8 (0.1), 84.3 (0.7)
81.2 (0.1), 64.8 (1.6)
PGI
52.5 (1.7), (1.5)
89.0 (1.8), 85.8 (1.8)
92.9 (0.2), 83.0 (1.3) 80.8 (0.8). 69.0 (3.3)
G-DRO
73.5 (5.3), 48.5 (11.5)
90.1 (0.5). 85.0 (1.5)
92.9 (0.3), 86.1 (0.9)
81.5 (0.1), 76.6 (0.5)
CGD
78.6 (0.5), (5.9)
91.3 (0.6), 88.9 (0.8) 92.5 (0.2), 90.0 (0.8) 81.3 (0.2), 76.1 (1.5)
Table 4: Average and worst-group accuracy in that order. All numbers are averaged over 3 seeds and standard deviation is shown in parenthesis. G-DRO is abbreviation for Group-DRO.
RELATED WORK
Distributionally robust optimization (DRO) methods Duchi & Namkoong (2018) seek to provide uniform performance across all examples, through focus on high loss groups As result, they are not robust to the presence of outliers (Hashimoto et al., 2018: Hu et al., 2018; Zhai et al., 2021). Zhai et al. (2021) handles the outlier problem of DRO by isolating examples with high train loss.
WWLLDS dataset parameter configuration Github URI..
8

---

Published as a conference paper at ICLR 2022
Algorithml
Camelyon17
PovertyMap
FMoW
CivilComments
Metric +
Avg. Acc.
Worst Pearson
Worst-region Acc
Worst-Group Acc.
Eval. type
OOD
OOD
0OD
ID
CORAL
59.5 (7.7)
0.59 (0.03)
0.44 (0.06)
31.7 (1.2)
65.6 (1.3)
IRM
64.2 (8.1)
0.57 (0.08)
0.43 (0.07)
30.0 (1.4)
66.3 (2.1)
ERM
70.3 (6.4)
0.57 (0.07)
0.45 (0.06)
32.3 (1.2)
56.0 (3.6)
Group-DRe DRO
68.4 (7.3)
0.54 (0.11)
0.39 (0.06)
30.8 (0.8)
70.0 (2.0)
CGD
69.4 (7.8)
0.58 (0.05)
0.43 (0.03)
32.0 (2.2)
69.1 (1.9)
Table 5: Evaluation on WILDS datasets: All numbers averaged over multiple runs, standard de-
viation is shown in parenthesis. Second row shows the evaluation metric, and the third shows the evaluation type: in-oomiin (ID) or out-of-domain (OOD). Two highest absolute performance num- bers are marked in bold in each column.
Sagawa al. (2020a) extend DRO the case where training data comprises of a set of groups like in our setting. Dagaev et al. (2021); Liu et al. (2021): Creager et al. (2021); I); Ahmed et al. (2021) extend the sub-population shift problem to the case when the group annotations are unknown. They proceed by first inferring the latent groups of examples that negatively interfere in learning followed by robust optimization with the identified groups. In the same vein, Zhou al. (2021); Bao et al. (2021) build on Group-DRO for the case when the supervised group annotations cannot recover the ideal distribution with no spurious features in the family of training distributions they represent. Zhou et al. (2021) maintains per-group and per-example learning weights, in order to handle noisy group annotations. Bao et al. (2021) uses environment specific classifiers to identify example groupings with varying spurious feature correlation, followed by Group-DRO's worst-case risk minimization.
Goel et al. (2020) augment the minority group with generated examples. However, generating rep- resentative examples may not be easy for all tasks. In contrast, Sagawa et al. (2020b) propose to sub-sample the majority group to match the minority group. In the general case, with more than two groups and when group skew large, such sub-sampling could lead severe data loss and poor performance for the majority group. Menon et al. (2021) partially get around this limitation by pre-training on the majority group first. Ahmed al. (2021) considered as a baseline modified version of Group-DRO such that all the label classes have equal training representation. However, these methods have only been applied on a restricted set of image datasets with two groups, and do not consistently outperform Group-DRe DRO. In contrast, we show consistent gains over DRO and ERM up weighing, and our experiments are over eight datasets spanning image and text modalities.
Domain Generalization (DG) algorithms (Sun & Saenko, 2016; Shankar et al., 2018; Arjovsky al., 2019; Ahmed et al., 2021) train such that they generalize to all domains including the seen domains that is of interest in the sub-population shift problem. However, due to the nature of DG benchmarks, the DG algorithms are not evaluated on cases when the domain sizes are heavily dis- proportionate such as in the case of sub-population shift benchmarks (Table 3). We compare with two popular DG methods: CORAL Sun & Saenko (2016) and IRM Arjovsky et al. (2019) on the WILDS benchmark, and obtain significantly better results than both.
7 CONCLUSION
Motivated by the need for modeling inter-group interactions for minority group generalization, we present a simple, new algorithm CGD for training with group annotations. We demonstrated the qualitative as well as empirical effectiveness of CGD over existing and relevant baselines through extensive evaluation using simple and real-world datasets. We also prove that CGD converges to FOSP ERM objective even though it is not performing gradient descent on ERM.
Limitations and Future work: CGD critically depends on group annotated training data to remove spurious correlations or reduce worst case error on sub-population. This limits the broad applica- bility of the algorithm. As part of future work, we plan to extend CGD to work in settings with noisy or unavailable group annotations following recent work that extends Group-DRO to such settings. The merits of CGD, rooted in modeling inter- inter-group interactions, could be more generally applied to example weighting through modeling of inter-example interactions, which could be of potential leverage in training robust models without requiring group annotations.

---

Published as a conference paper at ICLR 2022
8 REPRODUCIBILITY STATEMENT
We release the anonymized implementation of our algorithms publicly at this link, with instructions on running the code and pointers to datasets. We describe all our datasets in Section and Ap- pendix E. In Section 5, we also detail the hyperameter ranges, strategy for picking the best model. The detailed proof of our convergence theorem can be found in Appendix
ACKNOWLEDGEMENT
The first author is supported by Google PhD Fellowship. We acknowledge engaging and insightful discussions with Divyat Mahajan and Amit Sharma during the initial phases of the project. We gratefully acknowledge the Google's TPU research grant that accelerated our experiments.
REFERENCES
Faruk Ahmed, Yoshua Bengio, Harm van Seijen, and Aaron Courville. Systematic generalisation with group invariant predictions. In International Conference Learning Representations, 2021.
URL https: /ttttp:://peee /ttp//pppeeevveeeeeeeeeeeeeeeeee iew net / forum? id=b9Poims
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paze Invariant risk minimization.
arXiv preprint /9X7v:/907.000 002933, 2019.
Peter Bandi, Oscar Geessink, Quirine Manson, Marcory Van Dijk, Maschenka Balkenhol, Meyke Hermsen, Babak Ehteshami Bejnordi, Byungjae Lee, Kyunghyun Paeng, Aoxiao Zhong, et al. From detection of individual metastases to classification of lymph node status at the patient level: the camelyon17 challenge. IEEE Transactions on Medical Imaging, 2018.
Yujia Bao, Shiyu Chang, and Regina Barzilay. Predict then interpolate: A simple algorithm to learn stable classifiers. arXiv preprint arXiv: 105. 22288, 2021.
Daniel Borkan, Lucas Dixon, Jeffrey Sorensen, Nithum Thain, and Lucy Vasserman. Nuanced metrics for measuring unintended bias with real data for text classification. In Companion Pro- ceedings of The 2019 World Wide Web Conference, 2019.
Gordon Christie, Neil Fendley, James Wilson, and Ryan Mukherjee. Functional map of the world. In Proceedings of the IEEE Conference on Computer Vision and Puttern Recognition, 2018 8.
Elliot Creager, förn-Henria Jacobsen, and Richard Zemel. Environment inference for invariant learning. International Conference on Machine Learning, pp. 2189-2200. PMLR, 2021.
Nikolay Dagaev, Brett D Roads, Xiaoliang Luo, Daniel N Barry, Kaustubh R Patil, and Bradley C Love, too-good-to- so-od-tt-be--ruee prior to reduce shorteut reliance. arXiv preprint arXiv:2102.06-406 2021.
John Duchi and Hongseok Namkoong. Learning models with uniform performance via distribution- ally robust optimization. arXiv preprint arXiv:/810.088550, 2018.
Karan Goel, Albert Gu Yixuan Li, and Christopher Ré. Model patching: Closing the subgroup performance gap with data augmentation. arXiv preprint arXiv:2008.06777 2020.
Suchin Gururangan, Swabha Swayamdipta, Omer Levy, Roy Schwartz, Samuel R Bowman, and Noah A Smith. Annotation artifacts in natural language inference data. arXiv preprint arXiv: 1803.02234 2018.
Tatsunori Hashimoto, Megha Srivastava, Hongseok Namkoong, and Percy Liang, Fairness without demographics in repeated loss minimization. In International Conference on Machine Learning, pp. 1929-1938. PMLR, 2018.
Weihua Hu, Gang Niu, Issei Sato, and Masashi Sugiyama. Does distributionally robust supervised learning give robust classifiers? In International Conference on Machine Learning, pp. 2029- 2037. PMLR, 2018.
10

---

Published as a conference paper at ICLR 2022
Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Bal- subramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, et al. Wilds: A benchmark of in-the-wiln wild distribution shifts. arXiv preprint arXlv:2012.07411, 2020.
Yann LeCun and Corinna Cortes. MNIST handwritten digit database. 2010. URL http://yann.
lecun. com/ exdb /t.ist/.
Evan Z Liu, Behzad Haghgoo, Annie S Chen, Aditi Raghumathan, Pang Wei Koh, Shiori Sagawa, Percy Liang, and Chelsea Finn. Just train twice: Improving group robustness without training group information. In International Conference on Machine Learning, pp. 6781-6792. PMLR, 2021.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Large-scale celebfaces attributes (celeba) dataset. Retrieved August, 15(201 11: 2018.
Aditya Krishna Menon, Ankit Singh Rawat, and Sanjiv Kumar. Overparameterisatioo and worst- case generalisation: friend or foe? In International Conference on Learning Representations, 2021. URL https::t///........ openreview net/ forump:dd=jjhhnnoooeee
Shiori Sagawa, Pang Wei Koh, Tatsunori Hashimoto, and Percy Liang. Distributionally robust neural networks. In International Conference Learning Representations, 2020a. URL https:
/ //oppnreview..o////////// wett/formmiiddeeeeeeeee xyxuuuuU ott
Shiori Sagawa, Aditi Raghunathan, Pang Wei Koh, and Percy Liang. An investigation of why overparameterizattoon exacerbates spurious correlations. In International Conference Machine Learning, pp. 8346-8356. PMLR, 2020b.
Shiv Shankar, Vihari Piratla. Soumen Chakrabarti, Siddhartha Chaudhuri, Preethi Jyothi, and Sunita Sarawagi. Generalizing across domains via cross-gradient training. arXiv preprint a8Xiv: 10745, 2018.
Baochen Sun and Kate Saenko. Deep coral: Correlation alignment for deep domain adaptation. In European conference on computer vision, pp. 443 450. Springer, 2016.
Adina Williams, Nikita Nangia, and Samuel R Bowman. A broad-coverage challenge corpus for sentence understanding through inference. arXiv preprint 17XYvv1704.05426. 2017.
Christopher Yeh, Anthony Perez, Anne Driscoll, George Azzari, Zhongyi Tang, David Lobell, Ste- fano Ermon, and Marshall Burke. Using publicly available satellite imagery and deep learning to understand economic well-being africa. Nature Communications, 2020.
Runtian Zhai, Chen Dan, Zico Kolter, and Pradeep Ravikumar. Doro: Distributional and outlier robust optimization. arXiv preprint (ew:2106.061/ 2021.
Chunting Zhou, Xuezhe Ma, Paul Michel, and Graham Neubig. Examining and combating spurious features under distribution shift, arXiv preprint arXiv:2106.07171, 2021.

---

Published as a conference paper at ICLR 2022
Appendix
A SCALING RULE FOR GRADIENTS USING LOCAL QUADRATIC APPROXIMATION
The gradient norms can vary widely even for given task depending on where we are in the loss landscape. We will now illustrate this issue. Consider any group i. Since the training loss (, 2 0 for all 6, any 9 satisfying (0)) 0 is approximate global minimizer. For any such global minimizer with (0) 0 and (Q)]l
the local approximation of & given by the second
order Taylor expansion at 0 is:
(,()) (0 - 0, (0) (o 0) and Ve((0) ~ (0-(0))
Denoting = Omin (0)) 2 0 and σ max (e2 2 0 to be the smallest and the largest eigenvalue of (0), have that |V,((0) bounded between v (e,(o) and v (0) Consequently, the gradient norm can vary widely between smallest and largest cigenvalues. Besides, the gradient scale also depends the number of parameters, task, dataset and architecture.
To summarize, we show that gradients are of the order of the loss, but can take a large spectrum of values, which makes the updates unstable and Ma tuning difficult if used is in equation 3. Since we are only interested in capturing if a group transfers positively or negatively, we retain the cosine- similarity of the gradient dot-product but control their scale through {()), for some p > O. That is, we set the gradient (0) to [[
(8, (0)) leading to the final a update shown in equation 4.
When we set to very large value, the gradient inner products E,"E," cos (O), Ve, and hence the a value are largely influenced by the loss values. On the other hand, when we set p to O, we could get stuck in picking low loss train groups repeatedly without significant parameter update and hence not converge. In practice, neither the extremes are ideal. We tried p = {0.25, 0.5, 1,2] on synthetic setup, WaterBirds, CelebA, and found p 0.5 to perform well empirically.
B PROOF OF THEOREM
Proof of Theorem 1. Let consider the 1th iteration where the iterate is 0 and mixing weights are a.. Using the notation in Section 2, we denote R(0)- = & Σ, (,))). 9r ve,(0') and g = Σ, gi-
The update of our algorithm is given by:
a 1+1 a; exp (Na (gr. g))
(7)
Z
g++l_ 0r - 11 Σ a
(+]
gi,
(8)
where =E exp (n<(9).g)). g)). Let us fix a" (1/k, 1/k) R* and use (L E, Ps log Ei to denote the KL-divergence between p and q. We first note that the update equation on corresponds to mirror descent steps on the function (g. E; augi) and obtain:
KT(a0, 1000,,a41) (1
Σ
a log
a
a Σ af log a exp (na (9. 9))
4I
Σ
a ag
log
a
naa (gi.)) g) + a log
-Klaa (a)a" na 1 -lo Σ af exp (Na (ma (1)6
(9)
12

---

Published as a conference paper at ICLR 2022
Since each (.(.) is G-Lipschitz, we note that || G for every i. Consequently, |l - G and
Kgs, g)l I G Using the fact that exp(x) I 1 + x + x² for a < 1, and since 7a S G2 we see
that:
log Σ of exp (na (2. log Σ ag
na (ar + (na (1a)9)))
[ (Cotnderr)))
(vn)2)²
= =log
1
Σ cimceorn:)) (Ma.)))
where we used log(1 2) in the last step. Plugging this back into equation 9, we obtain:
KL(a", < KL(a", (a)) Σ Eaina(onn)) + (na - 2 nnallll
- + KI(a", ,a)) KL(a", aa - S =19112
naG..
(10)
na
We will now argue that E, (9 s E, (gi.)) g). To show this, it suffices show that:
Eal g) s
++1 (gb.9)
a
Σ aditonn0) (Σ at exp (na g/)
Σ
af 9) exp (na (9t,9)) g))
of ((),,) g) ip) (exp (9. 9)) - cxp) 2 0,
where Σ, at (gt. and erp Σ: at (gi, g) exp (Mt The last inequality follows from the fact that exp(-) is an increasing function and hence for any random variable X, covariance of X and (X) is greater than equal zero.
Now coming back to the update equation 8, we can argue its descent property as follows:
a 9t) L
R(0++1)= R(() n(g. Σ
Σ (+1 9ill2
< RR((') + 1 (-1012+ KI(a" a²) "Tla",a'tt)
122L
+ +na
+
1
a,
2+1
9il12
7a (1191)2 KL(a",a') - "Kla",atti + +4
2
112 LR
S R(00) + 17
Ta | R(0") R(0+1)) KL(a" a²) -KL(("",'') Ga + 1L
2
7 Summing the above inequality over timesteps t 1,... ,T. we obtain:
Ta
2
|R((0))2 R((0) R((0)
KL(a", a0) KI(a", + +1 + LLG
na
2B T
< + n). G²
2
where we used KL(a" +1 2 0, KL(a", .0) 0 and R(00) ((80)) S 2B in the last step.
Using the parameter choices 1 = 2 LG and 1/2- = v T. we obtain that:
T-1
BLG T
VR((')) ||2 < 33/
T
This proves the theorem.
13

---

Published as a conference paper at ICLR 2022
C
MORE DETAILS OF THE SYNTHETIC SETTINGS
Table 6: Plot of data from first two features for the three synthetic settings. The three columns show distribution examples from the three groups the order. The rows from top to bottom show the plots for Noisy (Section 4.1), Spurious (Section 4.3) and Rotation (Section 4.2) simple settings. The minority group for Rotation setting (right bottom) has only two examples. We omitted from the plot the third feature of the spurious setting (second row).
In Table 6, we show the data for each group and each simple setting considered in Section 4. In the noise setting of the first row, the noise is only introduced in one of the groups. For the spurious case of second row, the first features perform poorly on the first group while the third spurious feature (not shown in the figure) performs poorly on both the second and the third group. Shown in the last row, the classifier gradually rotated for the rotation case and the minority group is extremely sparse with only two examples.
D QUALITATIVE EXPERIMENTS WITH DEEP MODELS
In this section, we repeat the qualitative experiments of Section with MNIST examples and deep models. For all the setting below, we train a ResNet-18 18 model with standard SGD, pick the best hyperparameter based on the validation set, and report the performance on the test set corresponding to the checkpoint performing best on the validation set. We binarized the label space by grouping the first five digits to one class and the next five to another. of examples by randomly partitioning the training set into 4,900, 4,900, and 200 examples. We illustrate various aspects of the differences between Group-DRO and CGD by perturbing instances in group in three different ways below:
cases, we create three groups
D.1 LABEL NOISE SETUP
In this case, we randomly flipped all the labels the first group, meaning the first group is highly noisy in the training data while the test/validation set is noise-free. We trained a ResNet- 18 model
14

---

Published as a conference paper at ICLR 2022
on all the three groups with a parameter weight decay of 1. Since CGD focuses on high transfer groups instead of high loss groups, it is more robust to the presence of a noisy training group when compared with Group-DRO As a result, model trained with CGD has better worst group accuracy when compared with Group-DRO (Noise-MNIST column of Table 2).
D.2 UNEVEN INTER-GROUP SIMILARITY
In this case, we rotate digits in the second and third groups by 30, 60 degrees. The difference in degree of rotation quantifies the transferability between the groups. The second group transfers better to all the groups compared to the first and third. CGD, by design, focuses training on the high-transfer second group more than first and performs better on all the rotations when compared with Group-DRO (RoIa column of Table 2).
D.3 SPURIOUS CORRELATIONS SETUP
Here, three group setup, and last group, but not in the second group as follows: Digits in the first group are colored red for label 0 and blue for label 1, colored blue or red randomly in the second group, and colored blue for label and red for label 1 in the third group. CGD performs better than Group-DRO by focusing mostly on the clean and high transfer second group. As result, generalizes better than Group-DRO Spprriou-MNIST column of Table 2)
color feature to spuriously correlate with label in the first
E DATASETS
CoNOred-MNIST: Using MNIST (LeCun & Cortes, 2010) digits, we create a dataset where the foreground color is spuriously correlated with the label. We split 50,000 examples into majority, minority with 1000: group size ratio, and binarize the label space to predict digits 0-4, 5-9 as the two classes. Test data has equal proportion of each group. In the majority group, the foreground red for examples of label 0 and blue for label I; the minority group examples were colored in reverse. Color based prediction of the label cannot generalize, however, color is spuriously and strongly correlated with the label in the training data. Many previous work (Arjovsky et al., 2019; Ahmed et al., 2021) adopted Colored MNIST for probing generalization.
WaterBirds task is to predict an image contains water-bird or a land-bird and contains two groups: majority, minority in ratio 62:1 in training data. The minority group has reverse background- foreground coupling compared to the majority. Test/validation data has equal proportion of each group.
CelebA (Liu et al., 2018) task is to classify a portrait image of a celebrity as blonde/non-blonde, examples are grouped based on the gender of the portrait's subject. The training data has only 1,387 male blonde examples compared to 200K total examples. Test/validation data follows same group-class distribution as train.
MultiNLI (Williams et al., 2017) task is to classify pair of sentences as one of entailment, neutral, contradiction. Gururangan et (2018) identified that negation words in the second sentence are spuriously correlated with the contradiction labels. Accordingly, examples are grouped based on the second sentence contains any negation word. Test/validation data follows same group-class distribution as train.
CivilComments- WILDS (Borkan et al., 2019) task is to classify comments as toxic/non-toxis toxic. Ex- amples come with eight demographic annotations: white, black, gay, muslim, based on if the com- ment mentioned terms that are related to the demographic. Label distribution varies across demo- graphics. In the bechmark train examples grouped on black demographic, and testing is the worst group accuracy among all 16 combinations of the binary class label and eight demographics.
Poverty WILDS (Yeh et al., 2020) task is to classify satellite images in to real valued wealth index of the region. The rural and urban sub- aub-oopulation from different countries make the differ- ent train groups. Bechmark evaluates true wealth index in two settings: An in-domain setting that measures sub- populatton shift to seen regions, and out-domain setting that evaluates generalization new countries.
worst-region Pearson correlation between predicted and
15

---

Published as a conference paper at ICLR 2022
FMoW-WILDS (Christie al., 2018) task is to classify RGB satellite images to of 62 land use categories. Land usage differs across countries and evolves over years. Training examples are stratified in to cleven groups based the year of satellite image acquisition. We report out-domain evaluation using test data from regions of seen countries but from later years, and measures worst accuracy among the five geographical regions: Africa, Americas, Oceania, Asia, and Europe. The in-domain evaluation with worst accuracy on the eleven years is unavailable for other algorithms, so stick to only out-domain evaluation.
Camelyon177 1eeelyon177yy 17-WSD (Bandi et al., 2018) is a binary classification task of predicting from a micro- scope image a tissue it contains a tumour. The training data contains scans from three hospitals, and the test data contains scans from multiple unseen hospitals. The evaluation metric is average accuracy on the test set hospitals.
F MORE TECHNICAL DETAILS
Predictive Group Invariance (Ahmed et al., 2021) proposed an algorithm that was shown to gener- alize various kinds of shifts including in-distribution shifts. The algorithm penalizes discrepancy in predictive probability distribution across groups with the same label as shown below.
Eq- (y x)] Ep (p x)]
R(0) -
44 E, ΣΣ
Eq. (p x)] log
Where 'c' the class label and the rest of terms in the equation follow the notation of this paper. For given class label "c", the distributions: Q", Pc, represent the examples of the same label but belong to two different inferred groups. They used an existing work (Creager et al., 2021) to stratify the examples in to groups. Here, we used the available group annotations for creating the example splits: Q for a fair comparison. We picked the value of À from {1e-3, le-2, le-1, 1}.
The original paper noted that the second distribution of the KL divergence is an "easy" group such that its average predictive distribution is mostly correct. In our case, since all the groups contain bias, they all "casy", and hence the direction of divergence is unclear. Since ERM tends to learn the biases aligned with the majority group, we set the pe distribution from the majority and Q" from the minority.
We model inter-group transfer characteristics using inner products of first order gradients (equation 4). Since per-group gradient computation for all the parameters is computationally expen- sive, we use only a subset of parameters. We use only the last three layers for ResNet-50, Densenet- 101 and DistilBERT, and all the parameters for gradient computation for any other network.
CGD:
G PERFORMANCE AT VARYING LEVELS OF HETEROGENEITY
We study performance on the simple synthetic setting of Colored-MNIST under varying levels of heterogeneity, i.e. ratio majority to the mi- nority group. As discussed in the datasets tion, ColoredMNIST dataset has two groups: majority group where the label and the digit's color match and the minority group where they do not match. The examples are grouped in to majority and minority in the ratio of r to As the value of increases, the net (spu- rious) correlation of the color with the label increases, furthering the gap between the best and worst groups. The train split has 50,000 examples and validation, test split have 10,000 examples each. The validation, test split con- tain all the groups in equal proportions. We follow the same optimization procedure de-
100
G-DRO
CG
80
ERM
ERM-UW
40
20
10
50
100
500 1000
Figure 6: Worst group error rate with increasing ratio of majority to minority. Shaded region shows one standard error. All estimates are aggregated over three runs.
16

---

Published as a conference paper at ICLR 2022
scribed for the ColoredMNIST as described in the main content, with the exception that we set the weight decay to O.
We trained and evaluated different algorithms for the ratio (r) going from 2 to 1000. Figure 6 shows the worst-group error rate with the value of r for various algorithms. CGD performed well across all the values of I, and for the extreme value of r=1000: CGD's worst error rate of 34.4% far better than the 51.5% of Group-DRO.
H How DOES SYNTHETIC EXPERIMENTS RELATE TO REAL-WORLD DATASETS?
In this section we discuss the similarities between our synthetic settings of Section 4 and the real- world datasets 5.1.
The text benchmarks (MultiNLI, CivilCommens-WILDD and CelebA resemble our toy setup of Sec. 4.3. In MultiNLI, the examples with negation words have spurious correlations that negatively transfer between the contradiction and entailment labels. The examples from the group with no negation words do not contain such spurious features. Similarly, in CivilComments-WiLe WILDSS the examples from the black group contain spurious correlation (they contain tokens that identify the demographic, which can be easily exploited to classify examples from black group as mostly toxic), while such spurious features are absent in the non-black group. In the CelebA dataset too, male non- blonde (majority) negatively transfers to the male blonde (minority) since the classifier may learn to interpret short hair (male) to be non- non-beende. On the other hand, the female blond/non- badddoon--londd groups do not contain any known spurious correlation. In all the cases, CGD could avoid learning spurious features by focussing training groups with what was demonstrated in Sec 4.3, thereby learning a more robust solution with dampened strength of spurious features.
relatively low spurious correlation similar to
FMOW-WILDS, Mav WILDS are similar to our label noise simple setup of Sec. 4.1. FMOW-WILDS task is to classify a satellite image into one of 62 land-use categories. The dataset is annotated with human curators labeling if marked region in an image contains, say a "police station" (Yeh al., 2020). Depending on the demographic spread of the human curators, the label correctness is expected to vary from one region to another. Also, some land-use categories are far easier to classify than others (for eg. "police station" vs "helipad"). Similarly, PovertyMap-WIl. task is to map satellite image to its poverty index. Poverty index per region (urban/rural settle- ment) ground-truth was acquired through secondary sources such as asset index of the region from the national demographic surveys (Christie et al., 2018). The asset to wealth index per region was found to vary per country and hence the quality of the label. The non-uniform label noise of the two datasets is similar to setup in Sec. 4.1. CGD focuses only the difficult groups that transfer well to the rest, unlike Group-DRO that only pursue the maximum loss groups
17


Reviews:
[{"review_id": "rKfj51T4aAD", "reviewer": "Reviewer_UQv8", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The proposed method is easily understood and implemented. It also has good performance empirically. The writing is clear and the intuition behind the method is clearly conveyed in the simulation study. The intuition of inter-group interactions seems reasonable to me.\n\nWeakness:\nFor the detailed algorithm, There is somewhere that is unclear to me.\n1. Close to convergence, the statement that the gradient $\\nabla ||\\ell_i( \\theta_t)||$ can be much smaller than the loss value $\\ell_i(\\theta_t)$ is unclear to me. In fact, consider $\\ell(\\theta) = ||\\theta ||^2_2$. It should be $\\ell(\\theta_t) << ||\\nabla \\ell(\\theta_t)||$ when close to the convergence. In Appendix A, why ||nabla \\ell(\\theta_t)|| could be as small as $\\sqrt{\\underline{\\sigma} \\ell_i(\\theta)}$. And why is $\\underline{\\sigma}$ close to zero?\n\n2. When $<g_i,g_j> = 0$, it does not exactly match the group-DRO update steps mentioned at the bottom of Page 2, since there is regularization between consecutive time steps.\n3. Please formally define $cos(g_i,g_j)$.\n3. At the end of paper 3, the author mentioned: \" The corrected loss values simply replace the $\\ell_i$ of Line 5 in Algorithm 1\". However, the correction $\\ell_i = \\ell_i + C/\\sqrt{n}$ does affect $\\nabla \\ell_i(\\theta)$. I think the author should refer to updating step (4).\n\nFor the simulations and experiments:\n1. At Page 6, \"y, 1-y on the third group, and set to y or 1-y w.p. 0.6\" should be set to y w.p 0.6\n2. For the datasets, what are population types \"Label * Group\" or Group? Why does the worst ratio refect loosely the strength of the spurious correlation?", "summary_of_the_paper": "The paper provides an efficient method to generalize to all groups in the presence of sub-population shifts and domain adaptation. The paper conducts extensive simulations to derive insights and also numerical experiments on the benchmark dataset to demonstrate the performance. The proposed method is intuitive, easily implemented, and has good performance.", "summary_of_the_review": "This paper provides a practical algorithm to the broad field of distributional shifts, domain adaptation, and fairness.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "LkOgO0yq69Q", "reviewer": "Reviewer_cTFe", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "PROS:\n\n- Building an ERM-based method that performs well on the whole dataset without sacrificing the prediction accuracy on the group-level worse case really makes sense to me, especially when different groups have different amounts of label noise.\n- The paper is well written, with a clear presentation style, and well-supported contributions, and easy to follow with good details on the experimental setup.\n- The theoretical analyses and empirical understanding provide additional insight into this new algorithm. \n\nConcerns:\n\n- In some scenarios, the group annotations are available. Should we directly discard this information?\n- What is the computational complexity of the algorithm? \n- Is this method vulnerable to label noise?\n\n\n*** Post Rebuttal: After reading the authors' response and the updated components of the manuscript, I thank the authors for addressing nearly all of my concerns. The inclusion of a clearer motivation, more discussion w.r.t. CGD and datasets, all enhance my understanding of the contributions of the paper beyond my original review. I decide to vote for acceptance. ***", "summary_of_the_paper": "This paper proposes a novel ERM-based method for classification task with group annotated training data. The goal is to be group distributionally robust while enhancing the minority performance. The authors make an improvement to an existing method named Group-DRO by modifying the focus on the group with the highest regularized loss to focus on the group that leads to the largest decrease in average training loss. They analyze the convergence and present detailed comparisons with Group-DRO.", "summary_of_the_review": "Overall, I think this is an interesting and solid paper.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "L85QoS8Wbz", "reviewer": "Reviewer_4c9A", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros \n\n- I feel the idea of training on the group which leads to largest overall decrease in loss is natural and interesting. The use case of one group having larger noise and thus the issue of group DRO focusing on that too much is interesting. The idea is a natural solution to that issue.\n- The synthetic examples presented in the paper are interesting and clearly bring out the use cases of the method proposed and comparison with group DRO.\n- The empirical results presented lead to improved results on a variety of benchmark tasks.\n\nCons\n\n- I understand the synthetic examples presented but it is still hard for me to understand why this method also works better for the benchmark datasets. For example, in the simple spurious correlation setting with majority and minority groups, I would expect the spurious feature to also have a large reduction in the loss due to the imbalance between majority and minority groups. For group DRO, it makes sense the any large correlation with the spurious feature would hurt one of the groups and thus, the worst loss amongst groups can increase. But, the intuition is not so clear for the method proposed in this paper. Since, the gradient method is not minimizing any fixed loss function, it is hard to understand what is going on. For the real world datasets, it would have been nicer is the paper included some results on why the proposed method did better as it is hard to see why the benchmark datasets reflect the structure proposed in the toy setups.\n\nQuestions/Clarifications/Suggestions\n\n- It would be helpful for the readers if the synthetic examples were presented more clearly. For example, I am guessing that the groups are formed randomly in each of the synthetic examples by fixing their sizes. Some figures on how the data actually looks using figures would be helpful for the readers. \n- The authors mention their method does not optimize any particular loss function. Is it possible that there exists a loss function which their method optimizes but that is hard to find? What is the precise statement that the authors mean here? \n- The authors mention that they have to scale the gradient on page 3 because the norm of the gradient can get very small close to convergence. Can the authors comment more on why scaling by l_i(\\theta^t)^p is the right thing to do and how do they choose p=0.5? Why does the scaling need to depend on the loss value at all? Does the analysis in Theorem 1 take into account this scaling or is this an implementation detail? ", "summary_of_the_paper": "The paper gives a new algorithm for the setup where the test distribution is different from the train distribution. The setup includes multiple groups whose information is present during the training time but not during test time and the relative proportion of these groups change during test. The most commonly used method group-DRO does distributionally robust optimization or finds a classifier which performs well on the group with worst loss. This paper proposes to focus instead on the group which leads to maximum decrease in the loss while training instead of the group which has the maximum loss. The paper present several synthetic toy cases where their approach could be useful and concludes with experiments on a variety of benchmarks for this setup and shows improved results.", "summary_of_the_review": "I would make a recommendation for accepting this paper.\n\nOverall, I think the idea of the paper is interesting and also leads to improved results on benchmark datasets. The main concern is that there could have been more discussion  on why this method works for these real world datasets by connecting them to the synthetic setups presented in the paper.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "CQtXfw6W09", "reviewer": "Reviewer_4DaZ", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality & Quality: While I like the simplicity of the proposed approach, there are a couple of glaring weaknesses in the paper: \n\n\n\n1). The idea that we shouldn't minimize the worst-group error but rather minimize the error on the group that decreases all other groups' errors is intriguing! Definitely, the extant approach of minimizing the worst group error has generalization problems and hence doesn't perform best on all the benchmarks. However, the alternative proposed in the paper is not motivated, but rather comes across as an \"ad-hoc\" idea that \"just seems to work\". This is also corroborated by the fact that the proposed algorithm doesn't perform gradient descent on a particular loss function. That's the biggest weakness of this paper, especially in times when we want to infuse more *science* and *rigor* into deep learning. Sure, the paper tries to motivate the algorithm by saying that the \"worst group\" approach doesn't model inter-group similarity, but that intuition is nebulous at best. \n\n\n2). The results on the synthetic datasets seem contrived, but that's OK as one can design synthetic dataset to make their point. However, the standard benchmarks shown in Table 3 are also altered. For example, the CivilComments dataset is shown as a 2 group when it is originally a 8-group task (the groups being the demographics of the users) as shown in the WILDS dataset paper. Similarly, the MultiNLI dataset contains 2-groups that seem contrived. \n\n\nClarity: The paper is well written and puts itself nicely in context of previous work. The overall presentation of the paper is good.\n\nSignificance: The paper addresses an important problem of robust ML and proposes a new method for improving the \"group\" robustness.\n", "summary_of_the_paper": "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.\n\n\n\nMain Contributions: \n\n1). 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.\n\n2). Results are shown on synthetic and real-world datasets which show the superior performance of the proposed method in achieving group robustness.\n", "summary_of_the_review": "The paper provides an intriguing idea to improve group distributional robustness, but at its heart it is just an *ad-hoc algorithm* that *seems to work*. It is unclear what really is being optimized since there is no clear loss function that is minimized. The results also seem a little unconvincing since some of the real-world datasets are altered and not used in their original form, e.g., CivilComments, which raises questions regarding if the approach even works at all, irrespective of the issues regarding lack of theory behind the algorithm.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The manuscript proposes a method for addressing spurious correlations and sub-population (group) shift problem by modelling intergroup interactions. Past work (GroupDRO) focuses on the worst group which is subject to failure when groups have heterogeneous levels of noise and transfer. This work focuses on the group whose gradient leads to largest decrease in average training loss over all groups. The manuscript presents insights on why the proposed method called CGD may perform better than GroupDRO by studying simple synthetic settings. The manuscript also provides empirical evaluation on seven real-world datasets–which include two text and five image tasks with a mix of sub-population and domain shifts.

There are several positive aspects of the manuscript, including:
1. The idea of training on the group which leads to largest overall decrease in loss is natural and interesting;
2. The synthetic examples presented in the manuscript clearly bring out the use cases of the method proposed and comparison with GroupDRO;
3. The empirical results presented lead to improved results on a variety of benchmark tasks.

There are also several major concerns, including:
1. More discussion on why the proposed method works for the chosen real world datasets by connecting them to the synthetic setups presented in the manuscript;
2. The proposed algorithm does not minimize a specific loss function;
3. The standard benchmarks are altered. For example, the CivilComments dataset is shown as a 2 group when it is originally a 8-group task (the groups being the demographics of the users) as shown in the WILDS dataset paper.

Authors clarified, among others, that the proposed approach optimizes the macro-average loss function, and the standard benchmarks are not modified and the experiment setup is exactly like GroupDRO evaluation on the CivilComments-WILDS dataset. Reviewers noted that the generative model has not added anything new since it is essentially the synthetic example and it just shows what every robust machine learning method is supposed to do, i.e., don't rely on e_s (group-specific components) but on e_c (common components) while making predictions. It doesn't justify the procedure of choosing to focus on the group that minimizes the error for the group that decreases all other groups' errors. 

The revised manuscript includes a clearer motivation, and more discussion on how the synthetic examples connect to the real world datasets. Based on that, I put an accept recommendation.
Decision: Accept (Poster)