Distributionally Robust Neural Networks
Authors: Shiori Sagawa*
Pang Wei Koh*
Tatsunori B. Hashimoto
Percy Liang
Venue:  (2020)
Keywords: distributionally robust optimization
deep learning
robustness
generalization
regularization

Abstract: Overparameterized neural networks can be highly accurate on average on an i.i.d. test set, yet consistently fail on atypical groups of the data (e.g., by learning spurious correlations that hold on average but not in such groups). Distributionally robust optimization (DRO) allows us to learn models that instead minimize the worst-case training loss over a set of pre-defined groups. However, we find that naively applying group DRO to overparameterized neural networks fails: these models can perfectly fit the training data, and any model with vanishing average training loss also already has vanishing worst-case training loss. Instead, the poor worst-case performance arises from poor generalization on some groups. By coupling group DRO models with increased regularization---stronger-than-typical L2 regularization or early stopping---we achieve substantially higher worst-group accuracies, with 10-40 percentage point improvements on a natural language inference task and two image tasks, while maintaining high average accuracies. Our results suggest that regularization is important for worst-group generalization in the overparameterized regime, even if it is not needed for average generalization. Finally, we introduce a stochastic optimization algorithm for the group DRO setting and provide convergence guarantees for the new algorithm.


Published as a conference paper at ICLR 2020
DISTRIBUTIONALLY ROBUST NEURAL NETWORKS FOR GROUP SHIFTS: ON THE IMPORTANCE OF REGULARIZATION FOR WORST-CASE GENERALIZATION
Shiori Sagawa*
Pang Wei Koh"
Stanford University ssagawałcs stanford.edu
Stanford University
pangweieca. st anford.
Tatsunori B. Hashimoto
Percy Liang
Microsoft
Stanford University pliangêcs stanford.edu
sehashin@microsoff.: com
ABSTRACT
Overparameterized neural networks can highly accurate on average on i.i.d. test set yet consistently fail on atypical groups of the data (e.g.. by learning spu- rious correlations that hold on average robust optimization (DRO) allows to learn models that instead minimize the worst-case training loss over a set of pre-defined groups. However, we find that naively applying group DRO to overparameterized neural networks fails: these models can perfectly the training data, and any model with vanishing average training loss also already has vanishing worst-case training loss. Instead, the poor worst-case performance arises from poor generalization on some groups. By cou- pling group DRO models with increased regularization-2 stronger-than-typiaal l2 penalty or early stopping- we achieve substantially higher worst-group accu- racies, with 10-40 percentage point improvements on natural language inference task and two image tasks, while maintaining high average accuracies. Our results suggest that regularization important for worst-group generalization the over- parameterized regime, even if it is not needed for average generalization. Finally, we introduce a stochastic optimization algorithm, with convergence guarantees, to efficiently train group DRO models.
not in such groups). Distributionally
INTRODUCTION
Machine learning models typically trained minimize the average loss training set, with the goal of achieving high accuracy on an independent and identically distributed (i.i.d)) test set. However, models that highly accurate on average can still consistently fail on rare and atypical examples (Hovy & Sgaard| 2015) Blodgett et al., 2016 Tatman 2017 Hashimoto et al.] 2018 Duchi al. 2019). Such models are problematic when they violate equity considerations (Jurgens et al.) 2017: Buolamwini & Gebru 2018) or rely on spurious correlations: misleading heuristics that work for most training examples but do not always hold. For example, in natural language inference (NIJ)- (l..)-determining if two sentences agree contradict- the presence of negation words like 'never' strongly correlated with contradiction due to artifacts in crowdsourced training data (Gururangan al. 2018 McCoy et al.) 2019). A model that learns this spurious correlation would be accurate on average on an i.i.d. test set but suffer high error on groups of data where the correlation does not hold (e.g., the group of contradictory sentences with negation words).
To avoid learning models that rely on spurious correlations and therefore suffer high loss on some groups of data, instead train models minimize the worst-case loss over groups in the training data. The choice of how to group the training data allows us to use our prior knowledge of spurious correlations, e.g., by grouping together contradictory sentences with negation words in the NLJ example above. This training procedure is an instance of distributionally robust optimization (DRO),
"Equal contribution.

---

Published as a conference paper at ICLR 2020
Common training examples
Test examples
yt waterhird at water
y: landhied a: land background
y: waterbird ac land
Waterbirds background
background
yr blond havir az
y: dark hair at malls
y: blond nc male
CelebA
yI coutradiction
y: entailmmntt
y: eutailment
at has negation (P: coull still better. (11 The been betber.
no negaation
az hus negation
MulINLI
Read Slate's
'heee sllence
Jackson's faolings.
State had opiitio
(11) There shoet perfod whene spoles
Jackson's finding.
Figure 1:
Representative training and test examples for the datasets we consider. The correlation between the label y and the spurious attribute at training time does not hold at test time.
which optimizes for the worst-case loss over potential test distributions (Ben-Tal et al., 2013 Duchi
et al. 2016). Existing work on DRO has focused on models that cannot approach zero training loss, such as generative models (Oren et al. 2019) or convex predictive models with limited capacity
(Maurer & Pontil 2009 Shatheezadeh-ADeeeee Abadeh et al. 2015 Namkoong & Duchi 2017 Duchi &
Namkoong 2018
Hashimoto
al. 2018).
We study group DRO in the context of overparameterized neural networks in three applications (Fig- natural language inference with the MultiNLI dataset (Williams al. 2018), facial attribute recognition with CelebA (Liu et al. 2015), and bird photograph recognition with our modified ver- sion of the CUB dataset (Wah et al. 201T). The problem with applying DRO to overparameterized models is that if model achieves zero training loss then it is optimal on both the worst-case (DRO) and the average training objectives (Zhang et al.) 2017 Wen et al.) (0144). In the vanishing- -raisiing-trinnng-. loss regime, we indeed find that group DRO models do no better than standard models trained to minimize average loss via empirical risk minimization (ERM): both models have high average test accuracies and worst-group training accuracies, but low worst-group fest accuracies (Section 3.1
In other words, the generalization gap is small on average but large for the worst group.
In contrast, we show that gtrongl-rregularze group DRO models that do attain vanishing train- ing loss can significantly outperform both regularized and unregularized ERM models. We con- sider & penalties, early stopping (Section 3.2 and group adjustments that minimize risk measure which accounts for the differences in generalization gaps between groups (Section 3.3). Across the three applications, regularized group DRO improves worst-case test accuracies by 10-40 percentage points while maintaining high average test accuracies. These results give new perspective on gen- cralization in neural networks: regularization might not be important for good average performance (e.g., models can "train longer and generalize better" on average (Hoffer et al. 2017) but it appears important for good worst-case performance.
Finally, to carry out the experiments, introduce stochastic optimizer for group DRO that is stable and scales to large models and datasets. We derive convergence guarantees for our algorithm in the convex case and empirically show that it behaves well in our non-convex models (Section
2 SETUP
Consider predicting labels E y from input features & E X. Given a model family θ loss E: e x (X' x )) + R,, and training data drawn from some distribution P. the standard goal is to find a model 0 E θ that minimizes the expected loss Ep [e(;; (z, y)] under the same distribution P. The standard training procedure for this goal is empirical risk minimization (ERM):
ÒERM == arg min E P1000) (æ, v))],
(I)
dee
where P the empirical distribution over the training data.
In distributionally robust optimization (DRO) (Ben-Tal al. 2013 Duchi al. 2016 aim instead to minimize the worst-case expected loss over an uncertainty set of distributions Q
min R((0) = sup (awwr (((0) [(( (z.0(((,)))) (r.)))) QEo
(2)
eee
2

---

Published as a conference paper at ICLR 2020
The uncertainty set Q encodes the possible test distributions that want our model to perform well on. Choosing a general family Q, such as a divergence ball around the training distribution, confers robustness to wide set of distributional shifts, but also lead to overly pessimistie models which optimize for implausible worst-case distributions (Duchi et al. 2019).
To construct realistic set of possible test distributions without being overly conservative, lever- age prior knowledge of spurious correlations to define groups over the training data and then define the uncertainty set @ in terms of these groups. Concretely, we adopt the group DRO setting (Hu) iet al. 2018 Oren et al. 2019) where the training distribution P is assumed to be a mixture of m groups Pg indexed by groups, i.e., Q == {-1=11 4 y g Am where Am is the (m -idimnnnionaal probability sim- plex; this choice of Q allows us to learn models that are robust to group shifts. Because the optimum of a linear program attained at vertex, the worst-case risk (2) is equivalent a maximum over the expected loss of each group,
2 2.................. m²]
We define the uncertainty set as any mixture of these
R max B(z,))),,, [e(0; (z. y))l. gEG
(3)
We assume that we know which group each training point comes from-ic.. the training data com- prises (2r, V. g) hiplets-thougle we do not assume we observe g test time, the model cannot
g directly. Instead, we learn a group DRO model minimizing the empirical worst-group risk R(0):
ODRO := arg min { R(()= max E (a.,), [e(0; (z. v))l) AEG
(4)
gEg
where each group Pg is an empirical distribution over all training points (z, v.g') with g' (or equivalently, subset of training examples drawn from Pg). Group DRO learns models with good worst-group training loss across groups. This need not imply good worst-group test loss because of the worst-group generalization gap 8 := R(0) R(0). We will show that for overparameterized neural networks, is large unless we apply sufficient regularization.
2.1 APPLICATIONS
In the rest of this paper, study three applications that share similar structure (Figure []: each data point (x, y) has some input attribute a(x) E A that is spuriously correlated with the label y, and we use this prior knowledge to form m |A| | groups, one for each value of (a, v). We expect that models that learn the correlation between and y in the training data would do poorly on groups for which the correlation does not hold and hence do worse on the worst-group loss R()..
Object recognition with correlated backgrounds (Waterbirds dataset). Object recognition models can spuriously rely on the image background instead of learning to recognize the actual object (Ribeiro et al. 2016). We study this by constructing new dataset, Waterbirds, which com- bines bird photographs from the Caltech-UCSD Birds-200 200-2011 (CUB) dataset (Wah et al., 201T) with image backgrounds from the Places dataset (Zhou al. 2017). We label each bird as one of y = (waterbird, landbird] and place it on one of A = (water background, land background), with waterbirds (landbirds) more frequently appearing against water (land) background (Appendix C.1 There are n = 4795 training examples and 56 in the smallest group (waterbirds on land).
Object recognition with correlated demographics (CelebA dataset). Object recognition models (and other ML models more generally) can also learn spurious associations between the label and demographic information like gender and ethnicity (Buolamwini & Gebru 2018). We examine this on the CelebA celebrity face dataset (Liu et al. 2015). using hair color () = (blond. dark)) } as the target and gender (AA {male, female)) as the spurious attribute. There are n 162770 training examples in the CelebA dataset, with 1387 in the smallest group (blond-haired males).
Natural language inference (MultiNLI dataset). In natural language inference, the task is to determine if a given hypothesis is entailed by, neutral with, or contradicts given premise. Prior work has shown that crowdsourced training datasets for this task have significant annotation artifacts, such as the spurious correlation between contradictions and the presence of the negation words
In main experiments,
6; we also use 64 in supplemental experiments.
3

---

Published as a conference paper at ICLR 2020
nobody, no, never, and nothing (Gururangan et al. 2018). We divide the MultiNLI dataset (Williams
et al. 2018) into m = groups, one for each pair of labels y = fentailed, neutral, contradictory) and spurious attributes A {no negation, negation). There are n = 206175 examples in training set, with 1521 examples in the smallest group (entailment with negations); see Appendix C. for more details dataset construction and the training/test split.
COMPARISON BETWEEN GROUP DRO AND ERM
To study the behavior of group DRO vs. ERM in the overparametrized setting, we fine-tuned ResNet50 models (He et al., 2016 on Waterbirds and CelebA and a BERT model (Devlin et al.) 2019) on MultiNLI. These are standard models for image classification and natural language infer- ence which achieve high average test accuracies their respective tasks.
We train the ERM (1) and group DRO models using standard (minibatch) stochastic gradient de- scent and (minibatch) stochastic algorithm introduced in Section 5 respectively. We tune the learn- ing rate for ERM and use the same setting for DRO (Appendix C.2 For each model, we measure its average (in-distribution) accuracy over training and test sets drawn from the same distribution, as well as its worst- worst-group accuracy on the worst-performing group.
3.1 ERM AND DRO HAVE POOR WORST-GROUP ACCURACY THE OVERPARAMETERIZED REGIME
Overparameterized neural networks can perfectly fit the training data and still generalize well on average (Zhang et al. 2017). We start by showing that these overparameterized models do not generalize well on the worst-case group when they are trained to convergence using standard regu- larization and hyperparameter settings (He et al. 2016 Devlin et al. 2019), regardless of whether they are trained with ERM or group DRO
ERM. As expected, ERM models attain near-perfect worst-group training accuracies of at least 99.9% all three datasets and also obtain high average test accuracies (97. .%% 94.8%, and 82.5% on Waterbirds, CelebA, and MultiNLI). However, they perform poorly on the worst-case group at test time with porst-group accuracies of 60.0%, 41.1%, and 65.7% respectively (Table Figure 2 Their low worst-group accuracies imply that these models are brittle under group shifts.
DRO.
The ERM models trained above nearly perfectly classify every training point, and are there- fore near-optimal for both the ERM ① and DRO objectives. Indeed, we find that group DRO models perform similarly to ERM models, attaining near- near-perfect training accuracies and high aver- age test accuracies, but poor worst-group test accuracies (Table I Figure 2)
Discussion. The ERM and DRO models attain near-perfect training accuracy and vanishing train- ing loss even in the presence of default regularization (batch normalization and standard lg penalties for ResNet50, and dropout for BERT). However, despite generalizing well average, they do not generalize well on the worst-case group, and consequently suffer from low worst-group accuracies. This gap between average and worst-group test accuracies arises not from poor worst-group train- ing performance-the models are near-perfect at training time, even on the worst-case groups-but from variations in the generalization gaps across groups. Even though DRO is designed improve worst-group performance, we find no improvements on worst-group test accuracies since the models already achieve vanishing worst-group losses on the training data.
3.2 DRO IMPROVES WORST-GROUP ACCURACY UNDER APPROPRIATE REGULARIZATION
Classically, we can control the generalization gap with regularization techniques that constrain the model family's capacity to fit the training data. In the modern overparameterized regime, explicit
Training to convergence a widespread practice for image models (Zhang et al. 2017 Hoffer et al.
2017). Pre-trained language models are typically pretrained until convergence (Devlin et al. 2019 Radford al. 2019) but fine-tuned for fixed small number of epochs hecause average test accuracy levels off quickly; we verified that training to convergence gave equally high average test accuracy.
4

---

Published as a conference paper at ICLR 2020
Average Accuracy
Worst-Group Accuracy
ERM
DRO
ERM
DRO
Train
100.0
100.0
100.0
100.0
Standard Regularization
Waterbirds
Test
97.3
97.4
60.0
76.9
Train
100.0
100.0
99.9
100.0
CelebA
Test
94.8
94.7
41.1
41.1
Train
99.9
99.3
99.9
99.0
MultiNLI
Test
82.5
82.0
65.7
66.4
Strong E Penalty
Train
97.6
99.1
35.7
97.5
Waterbirds
Test
95.7
96.6
21.3
84.6
Train
95.7
95.0
40.4
93.4
CclebA
Test
95.8
93.5
37.8
86.7
Train
86.2
80.1
7.1
74.2
Early Stopping
Waterbirds
Test
93.8
93.2
6.7
86.0
Train
91.3
87.5
14.2
85.1
CelebA
Test
94.6
91.8
25.0
88.3
Train
91.5
86.1
78.6
83.3
MultiNII
Test
82,8
81.4
66.0
77.7
Table Average and worst-group accuracies for each training method. Both ERM and DRO models perform poorly on the worst-case group in the absence of regularization (top). With strong regu- larization (middle, bottom), DRO achieves high worst-group performance, significantly improving from ERM. Cells are colored by accuracy, from low (red) to medium (white) to high (blue) accuracy.
ERM
DRO Standard Regularization
ERM Strang Penalty
DRO Strong Penalty
Standard Regularization
Accuracy
Training Time
Training Time
Training Time
Training Time
Figure 2: Training (light) and validation (dark) accuracy CelebA throughout training. With de- fault hyperparameters and training to convergence, ERM and DRO models achieve perfect training accuracy across groups, but generalize badly on the worst-case group (red line in the left panels). With strong &a penalties, ERM models get high average train and test accuracies at the cost of the rare group (panel 3). DRO models achieve high train and test accuracies across groups (panel 4).
regularization is not critical for average performance; models can do well on average even when all regularization is removed (Zhang et al. (0177). and default regularization settings (like the models trained above) still allow models to perfectly fit the training data. Here, study if increasing regu- larization strength-und the models no longer perfectly the training data- can rescue worst-case performance. We find that departing from the beisshing-raiiinn significantly outperform ERM models on worst-group test accuracy while maintaining high average accuracy. We investigate two types of regularization:
regime allows DRO models to
l2 penalties. The default coefficient of the f2-norm penalty \||0
ResNet50 is 0.0001 (He
al. 2016). We find that increasing À by several orders of magnitude- À 1.0 for Waterbirds and = 0.1 for CelebA doee two things: 1) it prevents both ERM and DRO models from achieving perfect training accuracy, and 2) substantially reduces the generalization gap for each group.
With strong lz penalties, both ERM and DRO models still achieve high average test accuracies. However, because model can achieve perfect training accuracy in this regime, ERM models sacri- fice worst- worst-group training accuracy (35.7 (35.7% and 40.4% on Waterbirds and CelebA: Table Figure 2 and consequently obtain poor worst-group test accuracies (21.33 33% and 37.8%, respectively).
5

---

Published as a conference paper at ICLR 2020
In contrast, DRO models attain high worst-group training aceuracy (97.5% and 93.4% on Waterbirds and CelebA). The small generalization gap in the strong-f2-1 f2-penllty regime means that high worst- group training accuracy translates to high worst-group test accuracy. which improves over ERM from 21.3% to 84.6% on Waterbirds and from 37.8% to 86.7% on CelebA.
While these results show that strong l penalties have striking impact on ResNet50 models for Waterbirds and CelebA, we found that increasing the & penalty on the BERT model for MultiNLI resulted in similar or worse robust accuracies than the default BERT model with no lg penalty.
Early stopping. A different, implicit form of regularization is early stopping (Hardt et al. 2016b). We use the same settings in Section 3.1 but only train each model for fixed (small) number of epochs (Section (.2). As with strong l penalties, curtailing training reduces the generalization gap and prevents models from fitting the data perfectly. In this setting, DRO also does substantially better than ERM on worst-group test accuracy, improving from 6.7% to 86.0% on Waterbirds, 25.0% to 88.3% on CelebA, and 66.0% to 77.7% on MultiNLI. Average test accuracies are comparably high in both ERM and DRO models, though there is a small drop of 1 3% for DRO (Table Figure
Discussion. We conclude that eegulaizztioon-prvveninnn the model from perfectly fitting the training data-does matter for worst-group accuracy. Specifically, it controls the generalization gap for each group, even on the worst-case group. Good worst-group test accuracy then becomes a ques- tion of good worst-group training accuracy. Since regularized model can perfectly fit the training data, ERM and DRO models make different training trade-offs: ERM models sacrifice worst-group for average training accuracy and therefore have poor worst-group test accuracies, while DRO mod- els maintain high worst-group training accuracy and therefore do well at test time. Our findings raise questions about the nature of generalization in neural networks, which has been predominantly studied only in the context of average accuracy (Zhang et al. 2017 Hoffer et al. 2017).
3.3 ACCOUNTING FOR GENERALIZATION THROUGH GROUP ADJUSTMENTS IMPROVES DRO
In the previous section, we optimized for the worst-group training loss via DRO relying on regularization to control the worst-group generalization gap and translate good worst group training loss to good worst-group test loss. However, even with regularization, the generalization gap can vary significantly across groups: in the Waterbirds DRO model with strong &g penalty, the smallest group has a train-test accuracy gap of 15.4% compared to just 1.0% for the largest group. This suggests that we can obtain better gorst- test loss if at training time, prioritize obtaining lower training loss on the groups that we expect to have a larger generalization gap.
We make this approach concrete by directly minimizing an estimated upper bound the worst- group test loss, inspired by ideas from structural risk minimization (Vapnik 1992). The key con- sideration is that each group 9 has its own generalization gap 6, Ex [e(;; (ax (e,8)]] ( [e(0: (x, w))] To approximate optimizing for the worst-group test loss R(0)) maxges E (, [e(0; (z, + we propose using the simple,
heuristic
by = C.. where ng is the group size for g and C is a model capacity constant which we treat as hyperparameter. This gives the group-adjusted DRO estimator
(E [e(0; (z. v))l + Vng C)-
C
Payj == arg min max lge
(5)
gEg
The scaling with 1/ ng reflects how smaller groups are more prone to overfitting than larger groups, and is inspired by the general size dependence of model-complexity-bee based generalization bounds (see, e.g., Cao et al. (2019)).
By incorporating group adjustments in (5. we encourage the model to focus more on fitting the smaller groups. We note that this method using a 1/ l/vn surrogate for the generalization gap only works in the group DRO setting. where we consider the worst-group loss over groups of different sizes. It does not apply in the ERM setting; if we were minimizing average training loss, the 1/ Vn term would simply be constant and not affect the optimization.
Results. We evaluate group adjustments using group DRO models with strong & penalties (as in Section 3.2 In Waterbirds (A 1.0), worst-group test accuracy improves by 5.9%. cutting the

---

Published as a conference paper at ICLR 2020
Average Accuracy
Worst-Grroup Accuracy
Naive
Adjusted
Naive
Adjusted
Waterbirds
96.6
93.7
84.6
90.5
CelebA
93.5
93.4
86.7
87.8
Table 2: Average and worst-group test accuracies with and without group adjustments. Group ad- justments improve worst-group accuracy, though average accuracy drops for Waterbirds.
C=
C=1
C=2
C=4
1.0
Accuracy
Training Time
Training Time Landbird, water
Training Time
Training Time
Landbird, land
Waterbird, land Waterbird, water
Figure 3: Training (light) and validation (dark) accuracies for each group over time, for different adjustments C. When C = 0, the generalization gap for waterbirds on land (green line) is large, dragging down worst-group accuracy. At = 2, which has the best worst-group validation accu- racy, the accuracies are balanced. At C' 4, overcompensate for group sizes, so smaller groups (e.g., waterbirds on land) do better at the expense of larger groups (e.g., landbirds on land).
error rate by more than a third (Table 2 and Figure 3). The improvements CelebA (A 0.1) are more modest, with worst-group accuracy increasing by 1.1%; f2 penalties are more effective in CelebA and there is not as much variation the generalization gaps by group at À 0.1. We did not evaluate group adjustments on MultiNLI as it did not benefit from stronger le penalties.
Empirically, group adjustments also help in the early stopping setting of Section 3.2 (in the next section, we evaluate models with group adjustments and early stopping across grid of & penalty strengths). However, it is difficult to rigorously study the effects of early stopping (e.g., because the group losses have not converged to stable value), leave more thorough investigation of the interaction between early stopping and group adjustments to future work.
4 COMPARISON BETWEEN DRO AND IMPORTANCE WEIGHTING
DRO models can be significantly more robust than ERM models. Here, we show theoretically and empirically that DRO also outperforms strong importance weighting baseline that is commonly used in machine learning tasks where the train and test distributions differ (Shimodaira] 2000; Byrd Lipton 2019). Recall that in our setting, the test distribution can be any mixture of the group distributions. For some assignment of weights w E A.. to groups, an emportance-weighted estimator would learn
Our results above show that
i arg min E ((...g)) E(O; (a. (x)) dee
(6)
Empirical comparison. We consider an importance-weighted baseline with weights set to the in- verse training frequency of each group, Wg 1/ p/ [T('' g)]- This optimizes for a test distribution with uniform group frequencies and is analogous the common upweighting technique for label shifts (Cui et al. 2019: Cao et al.) 2019); intuitively, this attempts to equalize average and worst-group error by upweighting the minority groups. Concretely, we train our weighted model by sampling from each group with equal probability (Shen et al. 2016), since recent study found this to be more effective than similar reweighting/resmmping methods (Buda et al. 2018).
Unlike group DRO, upweighting the minority groups does not necessarily yield uniformly low train- ing losses across groups in practice, as some groups might be easier to than others. To compare upweighting (UW) with ERM and DRO, train models across the same grid of l, penalty strengths and early stopping at the epoch with best worst-group validation accuracy (Table In CelebA and
"To avoid advantaging the DRO models by allowing them to tune additional hyperparameters, we restrict search for group adjustments the one & penalty strength used in Scction 3.3 See Appendix
7

---

Published as a conference paper at ICLR 2020
Waterbirds, upweighting performs much better than ERM but is slightly outperformed by DRO. However, upweighting fails on MultiNLI, achieving lower average and worst-group accuracies than ERM. With upweighting, it appears that the group is overemphasized and extremely low training accuracy achieved for that group at the cost of others.
Average Accuracy
Worst-Group Accuracy
ERM
UW
DRO
ERM
UW
DRO
Waterbids
97.0 (0.2)
95.1 (0.3)
93.5 (0.3)
63.7 (1.9)
88.0 (1.3)
91.4 (1.1)
CelebA
94.9 (0.2)
92.9 (0.2)
92.9 (0.2)
47.8 (3.7)
83.3
(2.8)
88.9 (2.3)
MultiNLI
82.8 (0.1)
81.2 (0.1)
81.4 (0.1)
66.4 (1.6)
64.8 (1.6)
77.7 (1.4)
Table 3: Comparison of ERM, upweighting (UW), and group DRO models, with binomial standard deviation in parenthesis. For each objective, we grid search over & penalty strength, number of epochs, and group adjustments and report the model with highest validation accuracy. These numbers differ from the previous tables because the larger grid search.
10:22
0.6
118 z2)
82
Figure 4: Toy example illustrating that DRO and importance weighting equivalent. The DRO solution is 0*, while any importance weighting would result in solutions at 0, or 0z.
Theoretical comparison. Should expect importance weighting to learn models with good worst- case loss? We show that importance weighting and DRO can learn equivalent models in the convex setting under some importance weights, but not necessarily when the models are non-convex.
We analyze the general framework of having weights w(z) for each data point 2, which is more powerful than the specific choice above assigning weights by groups. By minimizing the weighted loss E--p(wt (v)))()())))) 2)] over some source distribution P. we can equivalently minimize the expected loss Eg~q [E(B; =)] over target distribution Q where Q(2) x P(2) However, we want good worst-case performance over a family of Q E Q, instead of a single Q. Are there weights such that the resulting model i., achieves optimal worst-group risk? In the convex regime, standard duality arguments show that this is the case (see Appendix for the proof):
Proposition 1. Suppose that the loss ((:;z) is continuous and convex for all in Z, and let the uncertainty set Q be a set of distributions supported on Z. Assume that Q and the model family O C Rdd are convex and compact, and let 6* E θ be minimizer of the worst-group objective R()).
Then there exists distribution Q E Q such that 0* E arg ming Ex-or [e(0;
However, this equivalence breaks down when the loss E is non-convex:
Counterexample 1. Consider a uniform data distribution P supported on two points = {z1, 22) and let {(0; 2) be as in Figure with θ [0, 1]. The DRO solution 0* achieves a worst-case loss of R(0*)= = 0.6. Now consider any weights (Wi, W2) E 42 and w.l.o.g. let w1 2 W2. The minimizer of the weighted loss wi wr[(f; 21) wzz((0; 22) is 61, which only attains a worst-case loss of R((""))) 1.0. Remark. Under regularity conditions, there exists a distribution Q such that 0* is first-order stationary point of T (f(f; 2)) (see e.g., Arjovsky et al. (2019). However, as the counterexample demonstrates, in the non-corvex setting this does not imply that 0+ actually minimizes Ez~o [e=Q[((0; 2)}.
This negative result implies that in the non-convex setting, there may not be any choice of weights w such that the resulting minimizer 9.. is robust. Even such weights did exist, they depend on 6* and obtaining these weights requires that solve a dual DRO problem, making reweighting easier to implement than DRO. Common choices of weights, such as inverse group size, are heuristics that may not yield robust solutions (as observed for MultiNLI in Table 3)
8

---

Published as a conference paper at ICLR 2020
5 ALGORITHM
To train group DRO models efficiently, we introduce an online optimization algorithm with conver- gence guarantees. Prior work on group DRO has either used batch optimization algorithms, which do not scale to large datasets, or stochastic optimization algorithms without convergence guarantees.
In the convex and batch case, there is rich literature on distributionally robust optimization which treats the problem as a standard convex conic program (Ben-Tal Tal et al. 2013: Duchi et al. 2016 Bertsimas al. 2018: Lam & Zhou 2015). For general non-convex DRO problems. two types of stochastic optimization methods have been proposed: (i) stochastic gradient descent (SGD) on the
Lagrangian dual of the objective (Duchi & Namkoong [0088] Haahimoto al. 2018), and (ii) direct minimax optimization (Namkoong & 2016)] The first approach fails for group DRO because the gradient of the dual objective is difficult to estimate in stochastic and unbiased manner An al- gorithm of the second type has been proposed for group DRO (Oren et al. 2011),, but this work does not provide convergence guarantees, and we observed instability in practice under some settings.
Recall that we aim to solve the optimization problem
4 which can be rewritten as
m
min sup EE ge
Σ Eg [e((; (z, [))]
Extending existing minimax algorithms for DRO (Namkoong & Duchi 2016 Oren et al. 2019),, interleave gradient-based updates on and q. Intuitively, we maintain a distribution over groups, with high masses on high-loss groups, and update each example proportionally the mass on its group. Concretely, we interleave SGD on 0 and exponentiated gradient ascent on (Algorithm [ (In practice, we use minibatches and momentum term for 0: see Appendix C.2 for details.) The key improvement from the existing group DRO algorithm (Oren et al. 2019) is that is updated using gradients instead of picking the group with worst average loss at each iteration, which is important for stability and obtaining convergence guarantees. The run time of the algorithm is similar to that of SGD for a given number of epochs (less than a 5% difference), as run time is dominated by the computation of the loss and its gradient.
Algorithm 1: Online optimization algorithm for group DRO
Input: Step sizes 11. 110; Pg for each g E G Initialize ((0) and q
for t 1..... T do
g ~ Uniform( (1...,,mm)) m) x, Pg
// Choose a group 9 at random
// Sample x,y from group g
o' + 9 1111); d/g t l (xp(),
// Update weights for group g // Renormalize
q(t) + -g//E, L//Eo' d'g
(t) Ø10) + 0((-1) 1oo9
ve(0l(-1)) (x,1))
// Use q to update 0
end
We analyze the convergence rate by studying the error eT of the average iterate ((:T))
max L(Q(1:T), ,)) min max L(0.g), g.am.
(8)
lee0 gE
where L(0, == Σn qy (r.)), [6: (z. y))] is the expected worst-case loss. Applying results from Nemirovski et al. (2009), we can show that Algorithm [ has a standard convergence rate of O(1/VT) in the convex setting (proof in Section A.2):
Proposition 2. Suppose that the loss E(:; (z. y)) is non-negative, convex, -ipLipcchitz continuous, and bounded by Be for all (I. y) in X' x y. and [|0]|2 S Be for all E θ with convex θ C Rd. Then, the average iterate of Algorithm Sanadhmmmachieeeee Machieves an expected error the rate
10(B2 B2 + Bi log m)
E[ET] < 2m1
(9)
T
where the expectation is taken over the randomness of the algorithm.
4 "The dual optimization problem for group DRO mine, x E, [mnax(., I P, [e(0; (a,w))
- +B
for constant o. The max over expected loss makes difficult obtain unbiased, stochastic gradient estimate.

---

Published as a conference paper at ICLR 2020
6 RELATED WORK
The problem of non-uniform accuracy. Existing approaches to addressing non-uniforn accuracy over the data distribution include domain adaptation techniques for known target distributions (Ben- David et al. 2006; Ganin & Lempitsky 2015) and work in ML fairness (Dwork et al. 2012 Hardt
2016a Kleinberg et al. 2017). As we discuss in Section 4 importance weighting a classic et al.
example of the former (Shimodaira, 2000). [yrd & Lipton (2019) empirically study importance weighting in neural networks and demonstrate that has little effect unless regularization is applied.
This is consistent with the theoretical analysis in Wen et al. 2014), which points out that weighting has little impact in the zero-loss regime, and with our own observations in the context of DRO.
Distributionally robust optimization. Prior work in DRO typically defines the uncertainty set Q as a divergence ball around the training distribution over (x, y) (Ben-Tal et al., 2013 Lam &
Zhou 2015 Duchi al., 2016 MMyato al., 2018 Esfahani & Kuhn 2018 Bertsimas et al.
2018 Blanchet & Murthy (2099). With small divergence halls of radii O(1/n), DRO acts as a regularizer (Shaticezadeh-Abaee Abadeh et al., 2015 Namkoong & Duchi 2017). However, when the radius is larger, the resulting Q can be too pessimistic. In contrast, group DRO considers Q that of wider radius but with fewer degrees of freedom (shifts over groups instead of over (a., y)). Prior work proposed group DRO in the context of label shifts (Hu et al., (2018) and shifts in data sources (Oren et al.| 2019). Our work studies group DRO in the overparameterized regime with vanishing training loss and poor worst-case generalization. In contrast, most DRO work has focused on the classic (underparameterizee) model setting (Namkoong & Duchi 2017 Hu et al. 2018 Duchi et al., 2019). Sinha et al. (2018) study neural networks but with more conservative Wasserstein uncertainty set that leads to non non-vanishing training loss; and Oren et al. (2019) study neural networks but for generative modeling where loss tradeoff' arise naturally.
Generalization of robust models. There is extensive work investigating generalization of neural networks in terms of average loss, theoretically and empirically (Hardt et al.) [016bb Szegedy al. 2016 Hoffer et al. 2017). However, analysis on robust losses is limited. For label shifts, prior work has observed overfitting on rare labels and proposed algorithms to mitigate it (Buda et al.) 2018 Cui et al.) 2019 Cao et al. (019). In the DRO literature, generalization bounds on the DRO objective exist for particular uncertainty sets (c.g.. Duchi & Namkoong (2018)). but those works do not study overparameterized models. Invariant prediction models, mostly from the causal inference literature, similarly aim to achieve high performance on a range of test distributions (Peters et al, Bühlmann & Meinshausen [016 Heinze-Deml & Meinshausen 2017] RRothenhäuleer al. 2018 Yang et al. 2019; Arjovsky et al. (009)). For example, the maximin regression framework (Meinshausen & Bühlmann, 2015) also assumes group-based shifts, but focuses settings without the generalization problems identified in our work.
2016
7 DISCUSSION
In this paper, we analyzed group DRO in overparameterized neural networks and highlighted the importance of regularization for worst-case group generalization. When strongly regularized, group DRO significantly improves worst-group accuracy at a small cost in average accuracy.
As an application, we showed that group DRO can prevent models from learning pre-specified spu- rious correlations. Our supplemental experiments also suggest that group DRO models can maintain high worst-group accuracy even when groups are imperfectly specified (Appendix B). While han- dling shifts beyond pre-specified group shifts is important future work, existing work has identified many distributional shifts that can be expressed with pre-specified groups, e.g., batch effects in bi- ology (Leek et al. 22010), or image artifacts BaadennRRanne et al. 2019) and patient demographics (Badgeley et al. 2019) in medicine.
More generally, our observations call for deeper analysis of average vs. worst-case generalization the overparameterized regime. Such analysis may shed light on the failure modes of deep neural networks as well as provide additional tools (beyond strong le penalties early stopping) to counter poor carst-case generalization while maintaining high average accuracy.
10

---

Published as a conference paper at ICLR 2020
ACKNOWLEDGMENTS
We are grateful to Shyamal Buch, Yair Carmon, Zhenghao Chen, John Duchi, Jean Feng, Christina Heinze-Deml, Robin Jia, Daphne Koller, Ananya Kumar, Tengyu Ma, Jesse Mu, Hongseok Namkoong, Emma Pierson, and Fanny Yang for helpful discussions and suggestions. This work was funded by an Open Philanthropy Project Award. Toyota Research Institute ("TRI") also pro- vided funds to assist the authors with their research but this article solely reflects the opinions and conclusions of its authors and not TRI or any other Toyota entity. SS was supported by a Stanford Graduate Fellowship and PWK was supported by the Facebook Fellowship Program.
REPRODUCIBILITY
Code for training group DRO models is available at lhttps://githii com/ wwtpkoopaogwwee///// @roup_DRC) The datasets used in this paper also available that link, well as scripts to modify dataset generation (e.g., to choose different spurious attributes for CelebA and MultiNLI, or different object backgrounds relative group sizes for Waterbirds).
REFERENCES
M. Arjovsky, L. Bottou, I. Gulrajani, and D. Lopez-Paz. Invariant risk minimization. arXiv preprint vrXiv:1907.002 2019.
M. A. Badgeley, J. R. Zech, L. Oakden-Rayner, B. S. Glicksberg, M. Liu, W. Gale, M. V. McConnell, B. Percha, T. M. Snyder, and Dudley. Deep learning predicts hip fracture using confounding patient and healthcare variables. npj Digital Medicine, 2, 2019.
S. Ben-David, J. Blitzer, K. Crammer, and F. Pereira. Analysis of representations for domain adap- tation. In Advances in Neural Information Processing Systems (NeurIPS), pp. 137-144, 2006.
A. Ben-Tal Tal, D. den Hertog, A. D. Waegenaere, B. Melenberg, and G. Rennen. Robust solutions of optimization problems affected by uncertain probabilities. Management Science, 59:341-357, 2013.
D. P. Bertsekas. Convex Optimization Theory. Athena Scientific Belmont, 2009.
D. Bertsimas, Gupta, and Kallus. Data-driven robust optimization. Mathematical Program- ming Series A, 167, 2018.
J. Blanchet and K. Murthy. Quantifying distributional model risk via optimal transport. Mathematics of Operations Research, 44(2):565-600, 2019.
S. L. Blodgett, I. Green, and B. O'Connor. Demographic dialectal variation in social media: A case study of African-American English. In Empirical Methods in Natural Language Processing (EMNLP), pp. 1119-1130, 2016.
S. Boyd and L. Vandenberghe. Convex Optimization. Cambridge University Press, 2004.
M. Buda, A. Maki, and M. A. Mazurowski. A systematic study of the class imbalance problem in convolutional neural networks. Neural Networks, 106:249-259, 2018.
Bühlmann and N. Meinshausen. Magging: maximin aggregation for inhomogeneous large-scale data. In IEEE, 2016.
J. Buolamwini and T. Gebru. Gender shades: Intersectional accuracy disparities in commercial gender classification. In Conference on Fairness, Accountability and Transparency, pp. 77-91, 2018.
Byrd and Z. Lipton. What is the effect of importance weighting in deep learning? In International Conference Machine Learning (ICMI), pp. 872-881, 2019.

---

Published as a conference paper at ICLR 2020
K. Cao, C. Wei, A. Gaidon, N. Arechiga, and T. Ma. Learning imbalanced datasets with aabel-distribution-aweee aware margin loss. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
Y. Cui, M. Jia, Lin, Y. Song, and Belongie. Class-balanced loss based on effective number of samples. In Computer Vision and Puttern Recognition (CVPR), pp. 9268-9277, 2019.
Devlin, M. Chang, K. Lee, and Toutanova. BERT: Pre-training of deep bidirectional trans- formers for language understanding, In Association for Computational Linguistics (ACL), pp- 4171-4186, 2019.
J. Duchi and H. Namkoong. Learning models with uniform performance via distributionally robust optimization. arXiv preprint arXiv:/810.08750, 2018.
J. Duchi, P. Glynn, and H. Namkoong. Statistics of robust optimization: A generalized empirical likelihood approach. arXiv, 2016.
Duchi, Т. Hashimoto, and H. Namkoong. Distributionally robust losses against mixture covariate shifts. http:////...soocccd////////////s condrisk. pdd, 2019.
C. Dwork, M. Hardt, T. Pitassi, O. Reingold, and R. Zemel. Fairness through awareness. In Innova- tions in Theoretical Computer Science (ITCS), pp. 214-226, 2012.
P. M. Esfahani and D. Kuhn. Data-driven distributionally robust optimization using the wasserstein metric: Performance guarantees and tractable reformulations. Mathematical Programming, 171 (1):115-166, 2018.
Y. Ganin and V. Lempitsky. Unsupervised domain adaptation by backpropagation. In International Conference on Machine Learning (ICML), pp. 1180-1189, 2015.
S. Gururangan, S. Swayamdipta, O. Levy, R. Schwartz, S. Bowman, and N. A. Smith. Annotation artifacts in natural language inference data. In Association for Computational Linguistics (ACL), pp. 107-112, 2018.
M. Hardt, E. Price, and N. Srebo. Equality of opportunity in supervised learning. In Advances in Neural Information Processing Systems (NeurlPSS, pp. 3315-3323, 2016a.
M. Hardt, Recht, and Y. Singer. Train faster, generalize better: Stability of stochastic gradient descent. In International Conference on Machine Learning (ICML), pp. 1225-1234, 2016b.
Hashimoto, M. Srivastava, H. Namkoong, and P. Liang. Fairness without demographics in repeated loss minimization. In International Conference on Machine Learning (ICML), 2018.
K. He, X. Zhang. S. Ren, and J. Sun. Deep residual learning for image recognition. In Computer Vision and Pattern Recognition (CVPR), 2016.
C. Heinze-Denn and N. Meinshausen. Conditional variance penalties and domain shift robustness. arXiv preprint arXiv: 11469, 2017.
E. Hoffer, I. Hubara, and D. Soudry. Train longer, generalize better: closing the generalization gap in large batch training of neural networks. In Advances Neural Information Processing Systems (NeurIPS), pp. 1731-1741, 2017.
D. Hovy and A. Sgaard. Tagging performance correlates with age. In Association for Computational Linguistics (ACL), pp. 483-488, 2015.
W. Hu, G. Niu, I. Sato, and M. Sugiyama. Does distributionally robust supervised learning give robust classifiers? In International Conference on Machine Learning (ICML), 2018.
S. Ioffe and C. Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International Conference on Machine Learning (ICML), pp. 448 448-456 456, 2015.
12

---

Published as a conference paper at ICLR 2020
D. Jurgens, Y. Tsvetkov, and D. Jurafsky. Incorporating dialectal variability for socially equitable language identification. In Association for Computational Linguistics (ACL), pp. 51-57, 2017.
J. Kleinberg, S. Mullainathan, and M. Raghavan. Inherent trade-off's in the fair determination of risk scores, In Innovations Theoretical Computer Science (ITCS), 2017.
H. Lam and E. Zhou. Quantifying input uncertainty in stochastic optimization. In 2015 Winter Simulation Conference, 2015.
J. T. Leek, R. B. Scharpf. H. C. Bravo, D. Simcha, Langmead, W. E. Johnson, D. Geman, K. Baggerly, and R. A. Irizarry. Tackling the widespread and critical impact of batch effects in high-throughput data. Nature Reviews Genetics, 1 11(10), 2010.
Z. Liu, P. Luo, X. Wang, and X. Tang. Deep learning face attributes in the wild. In Proceedings of the IEEE International Conference on Computer Vision, pp. 3730-3738, 2015.
A. Maurer and M. Pontil. Empirical bernstein bounds and sample variance penalization. In Confer- ence on Learning Theory (COLT), 2009.
R. T. McCoy, E. Pavlick, and T. Linzen. Right for the wrong reasons: Diagnosing syntactic heuristics in natural language inference. Association for Computational Linguistics (ACL), 2019.
N. Meinshausen and P. Bühlmann. Maximin effects in inhomogeneous large-scale data. Annals of Statistics, 43, 2015.
Miyato, S. Maeda, Ishii, and M. Koyama. Virtual adversarial training: regularization method for supervised and semi-supervised learning. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2018
H. Namkoong and J. Duchi. Stochastic gradient methods for distributionally robust optimization with E-divergences. In Advances in Neural Information Processing Systems (NeurIPS), 2016.
H. Namkoong and J. Duchi, Variance regularization with convex objectives. In Advances Neural Information Processing Systems (NeurIPS), 2017.
A. Nemirovski, A. Juditsky, G. Lan, and A. Shapiro. Robust stochastic approximation approach to stochastic programming. SIAM Journal on optimization, 19(4):1574-1609 2009.
L. Oakden-Rayna J. Dunnmon, G. Carneiro, and C. Ré. Hidden stratification causes clinically meaningful failures in machine learning for medical imaging. arXiv preprint arXiv: 1909. 12475, 2019.
Y. Oren, S. Sagawa, T. Hashimoto, and P. Liang. Distributionally robust language modeling. In Empirical Methods in Natural Language Processing (EMNLP), 2019.
J. Peters, P. Bühlmann, and N. Meinshausen. Causal inference by using invariant prediction: identi- fication and confidence intervals. Journal of the Royal Statistical Society. Series B (Methodologi- cul), 78, 2016.
A. Radford, Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever. Language models unsuper- vised multitask learners. OpenAl Blog, 1(8), 2019.
M. T. Ribeiro, S. Singh, and C. Guestrin. "why should I trust you?": Explaining the predictions of any classifier. In International Conference on Knowledge Discovery and Data Mining (KDD), 2016.
D. Rothenhäusler, P. Bühlmann, N. Meinshausen, and J. Peters. Anchor regression: heterogeneous data meets causality. arXiv preprint 1201.01801.0622 2018.
S. Shafieezadeh-Ahadeh P. M. Esfahani, and D. Kuhn. Distributionally robust logistic regression. In Advances Neural Information Processing Systems (NeurIPS), 2015.
L. Shen, Z. Lin, and Q. Huang. Relay backpropagation for effective learning of deep convolutional neural networks. In European Conference Computer Vision, pp. 467-482, 2016.
13

---

Published as a conference paper at ICLR 2020
H. Shimodaira. Improving predictive inference under covariate shift by weighting the log- oog-likeihhod function. Journal of Statistical Planning and Inference, 90:227-244, 2000.
A. Sinha, H. Namkoong, and J. Duchi. Certifiable distributional robustness with principled adver- sarial training. In International Conference on Learning Representations (ICLR), 2018.
N. Srivastava, G. Hinton, A. Krizhevsky, I. Sutskever, and R. Salakhutdinov. Dropout: simple way to prevent neural networks from overfitting. Journal of Machine Learning Research (JMLR), (911):1922-195,, 2014.
C. Szegedy, V. Vanhoucke, S. loffe, J. Shlens, and Z. Wojna. Rethinking the Inception architecture for computer vision. In Computer Vision and Pattern Recognition (CVPR), pp. 2818 2826, 2016.
R. Tatman. Gender and dialect bias in youtubes automatic captions. In Workshop on Ethics in Natural Langauge Processing, volume pp. 53-59, 2017.
V. Vapnik. Principles of risk minimization for learning theory. In Advances in Neural Information Processing Systems, pp. 831-838, 1992.
C. Wah, S. Branson, P. Welinder, P. Perona, and S. Belongie. The Caltech-UCSD Birds-200-2001 200-2001 dataset. Technical report, California Institute of Technology, 2011.
J. Wen, C. Yu, and R. Greiner. Robust learning under uncertain test distributions: Relating covariate shift model misspecification. International Conference on Machine Learning (ICML), pp. 631-639, 2014.
A. Williams, N. Nangia, and S. Bowman. A broad-coverage challenge corpus for sentence under- standing through inference. In Association for Computational Linguistics (ACL), pp. 1112-1122, 2018.
Yang, Z. Wang, and C. Heinze-Deml. Invariance-inducing regularization using worst-case trans- formations suffices to boost accuracy and spatial robustness. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
C. Zhang, S. Bengio, M. Hardt, B. Recht, and O. Vinyals. Understanding deep learning requires re- thinking generalization. In International Conference on Learning Representations (ICLR), 2017.
B. Zhou, A. Lapedriza, A. Khosla, Oliva, and A. Torralba. Places: 10 million image database for scene recognition. IEEE Transactions Pattern Analysis and Machine Intelligence, 40(6); 1452-1464, 2017.
A PROOFS
A.1 EOUIVALENCE DRO AND IMPORTANCE WEIGHTING IN THE CONVEX SETTING
Proposition I Suppose that the loss e(-;2) is continuous and convex for all a in Z, and let the uncertainty sei Q be a set of distributions supported on Z. Assume that Q and the model family θ C _ are convex and compact, and let 0* E θ be minimizer of the worst-group objective R())). Then there exists a distribution Q E Q such that 0* E argmiing Ez-qr [e(0;2) z)).
Proof. Let (68, :E::=) 2)). Since the loss E(O: 2) is continuous and convex in 0 for all in 2, we have that 1((0, Q) is continuous, convex in 0, and concave (linear) in Q. Moreover, since con- vexity and lower cemi-continuity are preserved under arbitrary pointwise suprema, supgeo h(0, Q) also convex and lower semi-continuous (therefore proper).
Together with the compactness of and Q, the above conditions imply (by Weierstrass' theorem, proposition 3.2.1, Bertsekas (2009)), that the optimal value of the DRO objective
inf R(O) = inf sup h(0, Q).
(10)
d dee QEo
is attained some 0* E O
14

---

Published as a conference paper at ICLR 2020
A similar argument implies that the sup-inf objective
sup inf h(0, () QE age
(11)
attains its optimum at some Q" E Q.
Moreover, because Q and Q are compact and is continuous, we have the max-min equality (see, e.g., Ex 5.25 in [oyd & Vandenberghe (2004))
sup inf h(0, Q) inf sup h(0, Q)- QEQ eEe
(12)
e5o QEQ
Together, the above results imply that (0", Q*) form saddle point (proposition 3.4.1, Bertsekas 2009)), that is,
sup h(0", Q) = h(0", Q*) inf h(0, Q*). QEQ
(13)
Age
In particular, the second equality indicates that the optimal DRO model 0" also minimizes the weighted risk h(O, Q*) = (Bz-q-(((6) Z)}, as desired.
A.2 CONVERGENCE RATE OF ALGORITIMM
Proposition 2. Suppose that the loss E( (z. y)) is non-negative, convex, Bi-Lipschitz continuous, and bounded by Bt for all (I. y) in X' x y. and ||01|2 Be for all f θ with convex θ S R Then, the average iterate of Algorithm Machieves an expected error at the rate
10 B& + Bi logm]
Bierl s 2m
(14)
T
where the expectation taken over the randomness of the algorithm.
Proof. Our proof is an application of the regret bound for online mirror descent on saddle point optimization from Nemirovski et al. (2009).
We first introduce the existing theorem. Consider the saddle-point optimization problem
min #Ee g
(0)
(15)
under the following assumptions: Assumption 1. fy is convex on e.
Assumption 2. (0) (e=o(P,(O:E)) for some function Fg.
Assumption 3. We generate i.i.d. examples & ~ For given 0 E θ and & E, we can compute (9, () and unbiased stochastic subgradient (0; that is, Egrg (VF (0).
Online mirror descent with some c-strongly convex norm || has the following guarantee.
lle. yielding iterates ((1) ((T) and
Theorem 1 (Nemirovski al. (209 Eq 3.23). Suppose that Assumptions 1-3 hold. Then the pseulo-regret of the average iterates bg
(1:1") and 0g (1-1) can be bounded
10 10 M2 20 + M2 log ml
max GEA
f (8(:T)
min 8E0
(1:T) (g(
S
2
(16)
T
1=1
g=1
15

---

Published as a conference paper at ICLR 2020
where
2
E
0
q F (0:E)
M.
(17)
0,0_
E
2 F (0: E)
S M
(18)
RØ = max 10113 min
||0
(19)
for online mirror descent with c-strongly convex norm
lle-
remains to formulate our algorithm as an instance of online mirror descent applied to the saddle- point problem above. We start by defining the following:
Definition 1. Let q be distribution over ξ = (x, y.g) g) that is a uniform mixture of individual group distributions Pg:
(α,8.9) q: =
Ps
(20)
g'=1
Definition Let Fg (0; (x,))))))) mllg Ep, [e(0; (z. v))].
g'}e(0; (x,,)),
Correspondingly, let fo
We now check that Assumptions 1-3 hold under the original assumptions in the statement Theo- rem 2
We assume that the loss E(: (I, y)) is non-negative, continuous, and convex for all (x, y) in X' x V. As result, (0) is non-negative, continuous, and convex on e.
2. The expected value of (0) over distribution q is (0):
Er.y.grglFgg (0; (I. v.9))]
Ep [ (0: (z,
Ep, [Fg (0: (z, ()) -- g'l
m
Ep, (0:(00,,1) g'l
m Ep [e(0,,,)) 9 g'1 (0).
3. We can compute unbiased stochastic subgradient V V (0: (z. y. g)) Ex Eg..g~g (0; (z. v.9))] = Ex... Vmllg = (() (z, y))]
Ep, Vmllg = g ((00)(00=,,))
m
E,,, Ve(0: (z. y))
(f).
Finally, compute the constants required for the regret bound in Theorem | Recalling the original assumptions of Theorem 2
1. Bounded losses: E(O; (z, w)) B3 for all x, 0
2. Bounded gradients: Ve(0; (v,, 1)) < By for all x,y
16

---

Published as a conference paper at ICLR 2020
3. Bounded parameter norm:
||0|| S Be for all 0 e e,
obtain:
E
V
4g Fg. g0 (0; (2,
S m B MM.0
(21)
9 =1
V q
4g Fgr (0: (2, (x)
S m B7 = .M..g
(22)
9 =1
R3 = max 110133 min |012 Bg.
(23)
Plugging in these constants into the regret bound from Theorem obtain
773
112
10 B& + B² log ml
E max 7EAm..
min #EB
4g (1:2) f2
fg(
2m
(24)
T
9=1
9=1
This implies Theorem
Deemmm@eeeeee bbccase the minimax game is convex-concave.
B SUPPLEMENTARY EXPERIMENTS
Group DRO can maintain high robust accuracy even when spurious attributes are not perfectly spec- ified. We repeat the CelebA experiment on models with strong l penalties (Section 3.2) but with inexact group specifications:
1. Instead of the ground-truth spurious attribute Male, we provide a related attribute Wearing Lipstick, and
2. We also specify four distractor/non-sprrious attributes (Eyeglasses, Smiling, Double Chin, and Oval Face).
Optimizing for worst-case performance over all 2 64 groups (for all combinations of 5 attributes and I label), the DRO model attains 78.9% robust accuracy across the 4 original groups (dark- haired males and females, and blond males and females). These robust accuracies are not far off from the original DRO model with just the ground-trute spurious attribute (86. 7%) and significantly outperform the ERM model (37.8%).
C EXPERIMENTAL DETAILS
C.1 DATASETS
The standard MultiNLI train-test split allocates most examples (approximately 90%) the training set, with another 5% as a publicly-availably development set and the last 5% held- out test set that is only accessible through online competition leaderboards (Williams et al.) 2018). Because we are unable to assess model accuracy on each group through the online leaderboards, we validation and test sets by combining the training set and development set and then randomly shuffling them into a 50 - 20 100-30 30 train-vale muin-val-test. split. We chose to allocates more examples to the validation and test sets than the standard split to allow accurately estimate performance on rare groups in the validation and test sets.
MultiNLI.
create
We the provided gold labels as the target. removing examples with consensus gold label (as is standard procedure). We annotate an example as having a negation word any of the words nobody, no, never, and nothing appear in the hypothesis (Gururangan al. 2018).
17

---

Published as a conference paper at ICLR 2020
Waterbirds. The CUB dataset (Wah et al.) 2011) contains photographs of birds annotated by species as well as and pixel-level segmentation masks of each bird. To construct the Waterbirds dataset, we label each bird waterbird if it is a seabird (albatross, auklet, cormorant, frigate- bird, fulmar, gull, jaeger, kittiwake, pelican, puffin, or tern) or waterfowl (gadwall, grebe, mallard, merganser, guillemot, Pacific loon). Otherwise, label it a landbird.
To control the image background, we use the provided pixel-level segmentation masks to crop each bird out from its original background and onto a water background (categories: ocean natural lake) or land background (categories: bamboo forest or broadleaf forest) obtained from the Places dataset (Zhou et al., 2017). In the training set, we place 95% of all waterbirds against water background and the remaining 5% against land background. Similarly, 95% of all landbirds are placed against a land background with the remaining 5% against water.
We refer to this combined CUB-Places dataset as the Waterbirds dataset to avoid confusion with the original fine-grained species classification task in the CUB dataset.
We use the official train-test split of the CUB dataset, randomly choosing 20% of the training data to serve as validation set. For the validation and test sets, we allocate distribute landbirds and waterbirds equally to land and water backgrounds (i.e., there are the same number of landbirds on land vs. water backgrounds, and separately, the same number of waterbirds on land vs. water backgrounds). This allows to more accurately measure the performance of the rare groups, and it is particularly important for the Waterbirds dataset because of its relatively small size; otherwise, the smaller groups (waterbirds on land and landbirds accurately estimate performance on. We note that we can only do this for the Waterbirds dataset because we control the generation process; for the other datasets, we cannot generate more samples from the rare groups.
water) would have too few samples to
In typical application, the validation set might be constructed by randomly dividing up the available training data. We emphasize that this is not the case here: the training set is skewed, whereas the validation set is more balanced. We followed this construction so that we could better compare ERM vs. reweighting vs. group DRO techniques using a stable set of hyperparameters. In practice, the validation set were also skewed, we might expect hyperparameter tuning based on worst-group accuracy to be more challenging and noisy.
Due to the above procedure, when reporting average test accuracy in our experiments, we calcu- late the average test accuracy over each group and then report a weighted average, with weights corresponding to the relative proportion of each group in the (skewed) training dataset.
CelebA. We use the official train-val-test split that accompanies the CelebA celebrity face dataset (Liu et al. 2015). We use the Blond_Hai attribute as the target label and the Male attribute as the asssocittd variable.
C.2 MODELS
ResNet50. We use the Pytorch torchvision implementation of the ResNet50 model, starting from pretrained weights.
We train the ResNet50 models using stochastic gradient descent with momentum term of 0.9 and a batch size of 128; the original paper used batch sizes of 128 or 256 depending on the dataset (He) jet al. 2016). As in the original paper, we used batch normalization (Toffe & Szegedy (2015) and no dropout (Srivastava al.) 2014). For simplicity, we train all models without data augmentation.
We use a fixed learning rate instead of the standard adaptive learning rate schedule to make our dif- ferent model types casier to directly compare, since we expected the scheduler interact differently with different model types (c.g., due to the different definition loss). The interaction between batch norm and & penalties means that we had to adjust learning rates for each different & penalty strength (and each dataset). The learning rates below were chosen to be the highest learning rates that still resulted in stable optimization.
For the standard training experiments in Section 3.1 we use a &g penalty of À = 0.0001 (as in He al. (2016)) for both Waterbirds and CelebA, with learning rate of 0.001 for Waterbirds and
18

---

Published as a conference paper at ICLR 2020
0.0001 for CelebA. We train the CelebA models for 50 epochs and the Waterbirds models for 300 epochs.
For the early stopping experiments in Section 3.2 train each ResNet50 model for epoch. For the strong l penalty experiments that section, we use = 1.0 for Waterbirds and X = 0.1 for CelebA, with both datasets using learning rate of 0.00001. These settings of À differ because found that the lower value was sufficient for controlling overfitting on CelebA but not on Waterbirds.
For the group adjustment experiments in Section 3.3
the same settings of À 1.0 for Waterbirds and A = 0.1 for CelebA, with both datasets using a learning rate of 0.00001. For both datasets, below.
the value of C (0, 2,3,,,,,55 found in the benchmark grid search described
For the benchmark Section (Table 3). we grid search over l2 penalties of E (0.0001, 0.1, 1.0} for Waterbirds and À E {0.0001, 0.01,0.13 for CelebA, using the corresponding learning rates for each and dataset listed above. (Waterbirds and CelebA at À 0.1. which not listed above, both use learning rate of 0.0001.) To avoid advantaging DRO by allowing it to try many more hyperparameters, we only test group adjustments (searching over C E 2,,1, 4, 4,5)) on the & penalties used in Section 3.3 i.e., X = 1.0 for Waterbirds and = 0.1 for CelebA. All benchmark models were evaluated the best early stopping epoch (as measured by robust validation accuracy).
implementation of the BERT bert-base loss--nnassed model, starting from pretrained weights (Devlin et al. 2019) 5 We use the default tokenizer and model settings from that implementation, including fixed linearly-decaying learning rate starting at 0.00002, Adam optimizer, dropout, and no Eg penalty (A = 0), except that we use a batch size of 32 (as in Devlin al. (2019) instead of 8. We found that this slightly improved robust accuracy across all models and made the optimization less noisy, especially on the ERM model.
BERT.
We the Hugging Face artorch-transfie
For the standard training experiments in Section 3.1 we train for 20 epochs.
For the bg penalty experiments in Section 3.2 we tried penalties of (0.01, 0.03, 0.3, 1.0, 3.0, 10.0). However, these models had similar or worse robust ac- curacies compared the default BERT model with no &2 penalty.
For the early stopping experiments in Section 3.2 we train for 3 epochs, which is the suggested early-stopping time in Devlin al. (2019).
For the benchmark in Section (Tblee we similarly trained for 3 epochs. All benchmark models were evaluated the best early stopping epoch (as measured by robust validation accuracy).
19


Reviews:
[{"review_id": "rJgFdjSYFr", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes a novel training method based on group DRO that makes it possible for overparameterized neural networks to achieve uniform accuracy over different groups of the data, where the data is required to obey a mixture distribution composed of these groups. In addition, the authors identify that increased regularization plays a critical role for worst-group performance in the overparameterized regime by a series of empirical studies. The authors also compare  their method with the traditional importance weighting method. Finally, they introduce a stochastic optimizer for their group DRO method.\n\nPros: \n\n* The experiments are well designed and the results show the effectiveness of the approach.\n* Most parts of the paper are well written and easy to follow.\n\n Cons:\n* The title is a little bit vague and may overstate the paper's contribution. In fact, the problem addressed in the paper is not as general as the title suggests and has little to do with the generic DRO framework.\n* In the group DRO setting the data is distributed as a mixture of different groups. However, the groups need to be chosen using prior knowledge of \"spurious associations\", which may arouse doubts about the algorithm's actual effectiveness and undermines the value of the algorithm on the application side.\n* Some terms used in the paper are ambiguous. For example the concept of  \"group\" is not defined properly. Sometimes it refers to a subset of the dataset while sometimes it denotes a sub-distribution in the mixture distribution. The authors should use more precise terminology.\n* There is a typo in Proposition 1, where the uncertainty set Q should be a set of probability distributions,  not a subset of R^m.", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper."}, "extra_text": {}}, {"review_id": "HJg4BjDy5B", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper describes a method of training neural networks to be robust to a worse case mixture of a set of predefined example attributes. This is done with a loss in accuracy in the average case but improvements in the worse case. The proposed algorithm is relatively simple and convergence rates are also given for this new algorithm.\n\nBuilding neural networks that perform well in the face of group-level worse case test-set distributions is a very important problem particularly in areas such as health and safety-critical applications as past work points out. This paper shows good results in the worse case and additionally shows that the common technique of importance reweighting cannot arrive at the same solution. The convergence analyses also yield additional insight into this new algorithm. The paper is well written and relatively easy to understand with good details on the experimental setup. The algorithm has a downside in that the groups must be known a priori, is it possible for these groups to be learnt? Also, can using a hinge loss also improve robustness to the worse case examples?\n\nHowever, there are some unanswered questions in the paper. What is the effect on the training time of this algorithm? Is it just the time for an additional forward prop? What is the effect on the worse-case examples of weight decay on the Bert model? Even though it hurts the average performance does it improve the worse case at all? In the last line of algorithm 1 why is q_G used instead of q_g?\n\nOther comments:\nAt the bottom of P5: The ordering of 93.4% and 97.1% seem to be reversed.\nAbove eq. 5 , \\delta_g seems to be overloaded. In the paragraph, it first refers to the generalization gap and then later to a heuristic.\nTable 2: The drop in average accuracy for waterbirds does not seem 'small'.\nBottom of P8: 'background is more unique', it seems this is supposed to mean the background appears less often?\n\n======================================================================================================\nUpdate after rebuttal:\n\nThanks for the detailed answers to my comments and the additional experiments done with a hinge loss. I will keep my rating.", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "BJexEXB09r", "reviewer": "AnonReviewer5", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "To the best of my knowledge, this is the first paper to carefully address and propose an algorithm (with guarantees) for distributionally robust learning in the overparametrized regime, which is typical of modern large deep neural networks. Following other work, the paper formalizes distributionally robust learning as the minimization of a worst-case loss over a set of possible distributions. The main message of the paper is that for distributionally robust learning, regularization plays an important role by avoiding perfect fitting to the training data, at the cost of poor generalization (thus lack of robustness) in some of the possible distributions. Furthermore, the paper proposes a new stochastic optimization algorithm to minimize the loss that corresponds to distributionally robust learning and gives convergence guarantees for the proposed algorithm.\n\nI think this is an interesting and solid paper, with a clear presentation style, and well-supported contributions. To the best of my knowledge, this is novel work and, in my opinion, it is relevant work, both in terms of applicability as well as in terms of contribution to the understanding of the generalization behavior of overparameterized deep neural networks.\n", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I did not assess the derivations or theory.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes distributionally robust optimization (DRO) to learn robust models that minimize worst-case training loss over a set of pre-defined groups. They find that increased regularization is necessary for worst-group performance in the overparametrized regime (something that is not needed for non-robust average performance).

This is an interesting paper and I recommend acceptance. The discussion phase suggested a change in the title which slightly overstated the paper's contributions (a comment which I agree with). The authors agreed to change the title in the final version. 

Decision: Accept (Poster)