D ISTRIBUTIONALLY R OBUST N EURAL N ETWORKS FOR G ROUP S
HIFTS
: O N THE I MPORTANCE OF R EGULARIZATION FOR W ORST -C ASE G ENERALIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Overparameterized neural networks trained to minimize average loss 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) provides an approach for learning models that instead minimize worst-case training loss over a set of pre-deﬁned groups.
We ﬁnd that naively applying group DRO to overparameterized neural networks fails: these models can perfectly ﬁt the training data, and any model with vanishing average training loss will also already have vanishing worst-case training loss.
Instead, their poor worst-case performance arises from poor generalization on some groups.
By coupling group DRO models with increased regularization—stronger-than-typical (cid:96) 2 regularization or early stopping—we achieve substantially higher worst-group accuracies, with 10percentage point improvements over standard models on a natural language inference task and two image tasks, while maintaining high average accuracies.
Our results suggest that regularization is critical for worst-group generalization in the overparameterized regime, even if it is not needed for average generalization.
Finally, we introduce and provide convergence guarantees for a stochastic optimizer for this group DRO setting, underpinning the empirical study above.
I NTRODUCTION
Machine learning models are typically trained to maximize average accuracy on a test set that is independent and identically distributed (i.i.d.) to the training set.
However, even models that perform well on average can consistently fail on rare and atypical examples (Hovy & Sgaard, 2015; Blodgett et al., 2016; Tatman, 2017; Hashimoto et al., 2018; Duchi et 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 hold on different test distributions.
For example, in natural language inference (NLI)—determining if two sentences agree or contradict—negation words like ‘never’ are strongly correlated with contradiction due to artifacts in crowdsourced training data (Gururangan et al., 2018).
A model could achieve high average accuracy on an i.i.d.
test set by learning this spurious correlation, but perform poorly on test sets where the correlation does not hold (McCoy et al., 2019).
In this paper, we show that in the modern regime of training overparameterized models—i.e., large neural networks that can perfectly ﬁt the training data—this problem is one of generalization: these models can generalize well on average , as in Zhang et al.
but do not generalize well on some (sub)groups of data (e.g., the group of contradictions with no negations in the NLI example above).
To avoid the pitfalls of optimizing for average loss, we instead optimize for the worst-case loss over various groups in the data, using prior knowledge of spurious associations (e.g., between negation and contradiction) to choose these groups.
This is an instance of distributionally robust optimization (DRO), which studies worst-case performance over potential test distributions (Ben-Tal et al., 2013;
Duchi et al., 2016).
In this group DRO setting, the data is distributed as a mixture of different groups with mixture weights that can vary at test time (Hu et al., 2018; Oren et al., 2019).
Optimizing for the worst-case group shift —i.e., shift in group mixture weights at test time—is equivalent to minimizing the worst-case loss over each group.
Common training examples
Test examples
[Figure]
Prior work has applied DRO as a data-dependent regularizer (Maurer & Pontil, 2009; ShaﬁeezadehAbadeh et al., 2015; Duchi & Namkoong, 2016) and to learn models that are robust over different data subsets (Duchi & Namkoong, 2018; Hashimoto et al., 2018).
However, a major challenge with DRO in the overparameterized regime is that models can achieve vanishing training loss (Zhang et al., 2017)); any such model would be (near) optimal for both the average and worst-case training losses, so there is little incentive for their minimizers to differ (Wen et al., 2014).
Indeed, existing work on DRO has focused on either generative models or convex predictive models with limited capacities, both of which cannot get vanishing training loss.
In this paper, we study group DRO in the context of large, high-performing neural networks in three applications (Figure 1)—natural language inference with the MultiNLI dataset (Williams et al., 2018), facial attribute recognition with CelebA (Liu et al., 2015), and bird photograph recognition with our modiﬁed version of the CUB dataset (Wah et al., 2011).
In the vanishing-training-loss regime, we ﬁnd that group DRO models do no better than ERM models: both models achieve nearly perfect training accuracies across groups and continue to perform well on average at test time, but suffer low test accuracies on the worst-case group (Section 3.1).
In other words, the generalization gap is small on average, but the worst-case generalization gap over groups is large.
However, group DRO models can signiﬁcantly outperform ERM models when we account for these generalization gaps—through appropriate regularization, such as strong (cid:96) 2 regularization or early stopping (Section 3.2), and through group adjustments that explicitly handle the differences in generalization gaps between groups (Section 3.3).
Across the three applications, group DRO improves worst-case test accuracies by 10 to 40 percentage points while maintaining high average test accuracies.
These results give a new perspective on generalization 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 critical for good worst-case performance.
Finally, we introduce a stochastic optimizer for group DRO that underpins the experiments above 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 5).
S ETUP
Consider predicting labels y ∈ Y from input features x ∈ X .
Given a model family Θ , loss (cid:96) :
Θ × ( X × Y ) → R + , and training data drawn from some distribution P , the standard goal is to ﬁnd a model θ ∈ Θ that minimizes the expected loss E P [ (cid:96) ( θ ; ( x, y )] under the same distribution P .
The standard training procedure for this goal is empirical risk minimization (ERM):
[Equation]
(1) where ˆ P is the empirical distribution over the training data.
In distributionally robust optimization (DRO) (Ben-Tal et al., 2013; Duchi et al., 2016), we aim instead to minimize the worst-case loss over an uncertainty set of distributions Q ,
[Equation]
The uncertainty set Q encodes the possible test distributions that we want our model to perform well on. Choosing a general family Q , such as a divergence ball around the training distribution
P , confers robustness to a wide set of distributional shifts but can be too pessimistic because it optimizes for implausible worst-case scenarios (Duchi et al., 2019).
In other words, without any prior knowledge or structural assumptions on the distributional shifts, it is difﬁcult to specify an uncertainty set Q that is not too conservative but can still capture realistic shifts.
We thus consider the group DRO setting (Hu et al., 2018; Oren et al., 2019) where the training distribution P is assumed to be a mixture of m groups P g indexed by G = { 1 , 2 , . .
, m } . 1 We use the term group to refer to a sub-distribution of either the population or empirical training distribution ( P or ˆ P ); in the latter, each group ˆ P g corresponds to a subset of training examples drawn from P g .
Further, we assume the mixture weights can differ arbitrarily at train and test time, i.e., Q := { (cid:80) mg
=1 q g P g : q ∈ ∆ m } , where ∆ m is the ( m − 1) -dimensional probability simplex; this choice of
Q allows us to learn models that are robust to group shifts. Because the optimum of a linear program is attained at a vertex, the worst-case risk over this choice of Q is equivalent to a maximum over the expected loss of each group,
[Equation]
We assume that we know which group each training point comes from—i.e., the training data comprises ( x, y, g ) triplets—though we do not assume we observe g at test time, so the model cannot use g directly.
Instead, we use the training data to learn the group DRO model that minimizes the empirical robust risk
ˆ R ( θ ) :
[Equation]
where ˆ P g is the empirical distribution over all training points ( x, y, g (cid:48) ) with g (cid:48)
=
g .
Group DRO learns models with good robust training loss across groups.
This need not imply good robust test loss because of the generalization gap δ :=
R (ˆ θ DRO ) − ˆ R
(ˆ θ DRO ) .
We will show that for overparameterized neural networks, δ is large unless we apply sufﬁcient regularization.
2. A PPLICATIONS
In the rest of this paper, we study three applications that share a similar structure (Figure 1): each data point ( x, y ) has some input attribute a ( x ) ∈ A that is spuriously correlated with the label y , and we use this prior knowledge to form m = |A| × |Y| groups, one for each value of ( a, y ) .
We expect that models that learn the correlation between a and y in the training data would do poorly on groups for which a (cid:54) =
y , and hence do worse on the robust loss R ( θ ) .
Object recognition with correlated backgrounds (Waterbirds dataset).
Object recognition models can learn to make predictions from image backgrounds instead of the actual object (Ribeiro et al., 2016).
We investigate this by constructing a new dataset, Waterbirds, which combines bird photographs from the Caltech-UCSD Birds-200-2011 (CUB) dataset
(Wah et al., 2011) with image backgrounds from the Places dataset (Zhou et 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 a water (land) background (Appendix C.1).
There are n = 4795 training examples and 89 in the smallest group (waterbirds against land).
Object recognition with correlated demographics (CelebA dataset).
Object recognition models (and other ML models more generally) have also been shown to learn spurious associations between the target 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 ( Y = { blond , dark } ) as the target and gender ( A = { male , female } ) as the spurious attribute.
There are n = 162770 training examples, 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 a given premise.
Prior work has shown that crowdsourced training datasets for this task have signiﬁcant annotation artifacts, such as the spurious correlation between contradictions and the presence of the negation words nobody , no , never , and nothing (Gururangan et al., 2018).
We divide the MultiNLI dataset (Williams et al., 2018) into m = 6 groups, one for each pair of labels Y = { entailed , neutral , contradictory } and spurious attributes A =
{ no negation , negation } .
There are n = 206175 examples in the training set, with 1521 examples in the smallest group (entailment with negations); we modify the standard MultiNLI split to better estimate accuracy on small groups (Appendix C.1).
C OMPARISON BETWEEN GROUP DRO AND ERM
To study the behavior of group DRO vs. ERM in a modern setting, we ﬁne-tuned ResNet50 models (He et al., 2016) on Waterbirds and CelebA, and a BERT model (Devlin et al., 2018) on MultiNLI; these are popular and high-performing models for image and natural language tasks, respectively.
We train the ERM models using standard (minibatch) stochastic gradient descent.
For the DRO models, we introduce a stochastic minibatch algorithm based on online mirror descent that optimizes for an adaptively-weighted training distribution at each iteration; we defer discussion of this to
Section 5.
We tune the learning 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 robust accuracy on the worst-performing group.
3. ERM AND DRO HAVE POOR ROBUST ACCURACY IN THE OVERPARAMETERIZED REGIME
Overparameterized neural networks can perfectly ﬁt the training data and still generalize well on average (Zhang et al., 2017). We start by investigating if these overparameterized models also generalize well on the worst-case group—i.e., if they attain high robust test accuracy on top of high average test accuracy—after they are trained to convergence using standard hyperparameter settings (He et al., 2016; Devlin et al., 2018).
ERM.
As expected, ERM models reach near-perfect training accuracies of at least 99 .
8% on all three datasets, even on the worst-case group, and also obtain high average test accuracies ( 82 .
5% , 97 . 0% , and 94 .
8% on MultiNLI, Waterbirds, and CelebA).
However, they perform poorly on the worst-case group at test time ( 66 .
4% , and 41 .
1% ; Table 1, Figure 2).
Their low robust accuracies imply that these models are brittle under group shifts.
DRO.
In this vanishing-training-loss regime, there exist models that do almost perfectly on both the ERM (1) and DRO (4) objectives.
We train group DRO models in the same way as above, and indeed ﬁnd that they perform similarly to ERM models, attaining near-perfect training accuracies and high average test accuracies, but poor robust test accuracies (Table 1, Figure 2).
Discussion.
The ERM and DRO models attain near-perfect training accuracy and vanishing training loss even in the presence of default regularization (batch normalization and (cid:96) 2 regularization for ResNet50, and dropout for BERT).
However, they do not generalize well on the worst-case group, and consequently suffer from low robust accuracies.
This gap between average and robust test accuracies arises not from poor robust training performance—the models are near-perfect at training time, even on the worst-case groups—but from variations in the generalization gaps across groups.
3. G ROUP DRO IMPROVES ROBUST ACCURACY UNDER APPROPRIATE REGULARIZATION
Classically, we can control the generalization gap with regularization techniques that constrain the model family’s capacity to ﬁt the training data.
In the modern overparameterized regime, explicit 2 Training to convergence is 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., 2018;
[-The default weight decay strength (i.e., the coefﬁcient of the L 2 -norm penalty λ (cid:107) θ (cid:107) 22 ) in ResNet50 is λ = 0 . 
 (He-]Radford et al., 2019) but ﬁne-tuned for a ﬁxed small number of epochs because average test accuracy levels off quickly; veriﬁed training to equally high average accuracy.
[Table]
[Figure]
[Table]
[Figure]
regularization is not critical for average performance models can do well average regularization removed regularization perfectly the training if increasing regularization strength—until the no longer perfectly ﬁt the training data—can rescue worst- that departing from the vanishing training-loss regime allows group
DRO to outperform high (cid:96) 2 regularization. The default (cid:96) 2 regularization strength (i.e., the coefﬁcient of the (cid:96) 2 -norm penalty λ (cid:107) θ (cid:107) 22 ) in ResNet50 is λ = 0 .
(He et al., 2016). We ﬁnd that increasing λ by several orders of magnitude—to λ = 1 . 0 for Waterbirds and λ = 0 .
1 for CelebA—does 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.
Both ERM and DRO models still achieve high average test accuracies that are comparable to each other and to the standard models with weaker regularization.
However, because no model can achieve perfect training accuracy in this regime, ERM models sacriﬁce training accuracy on the worst-case group (robust accuracies of 37 . 1% and 40 .
4% for Waterbirds and CelebA; Table 1, Figure 2), and consequently obtain poor robust test accuracy.
In contrast, DRO models optimize for high robust training accuracy ( 97 .
1% and 93 .
4% on Waterbirds and CelebA).
The small generalization gap in the strong (cid:96) 2 regularization regime means that high robust training accuracy translates to high robust test accuracy: with group DRO, robust test accuracy improves from 21 .
5% to 86 .
1% on Waterbirds and from 37 .
8% to 86 .
7% on CelebA.
While these results show that strong (cid:96) 2 regularization has a striking impact on ResNet50 models for Waterbirds and CelebA, our initial experiments did not show a similar effect on BERT model for MultiNLI.
(cid:96) 2 regularization is not typically tuned for this task (Devlin et al., 2018)—our default hyperparameters set {+λ = 
+}
(Appendix C.2)—so we turn to other forms of regularization.
Early stopping.
A different, implicit form of regularization is early stopping (Hardt et al., 2015).
We use the same settings in Section 3.1, but only train each model for a ﬁxed (small) number of epochs (Section C.2).
As with strong (cid:96) 2 regularization, curtailing training reduces the generalization gap and prevents models from ﬁtting the data perfectly.
In this setting, DRO also does substantially better than ERM on robust test accuracy, improving from 66 .
0% to 77 .
7% on MultiNLI, 5 . 5% to 84 .
9% on Waterbirds, and 25 . 0% to 88 .
3% on CelebA. Average test accuracies are comparably high in both ERM and DRO models, though there is a small drop of 1 - 2% for DRO (Table 1, Figure 2).
Discussion.
We conclude that regularization—preventing the model from perfectly ﬁtting the training data—does matter for robust accuracy.
Speciﬁcally, regularization can control the generalization gap across each group, even on the worst-case group; good robust test accuracy then becomes a question of good robust training accuracy.
ERM and DRO models make different training trade-offs in the regularized regime, since no model can perfectly ﬁt the training data; ERM models sacriﬁce robust for average training accuracy and therefore do poorly at test time, while
DRO models maintain high robust training accuracy and therefore do well at test time.
Our ﬁndings raise additional questions about the nature of generalization in neural networks, which has been predominantly studied in the context of average accuracy (Zhang et al., 2017; Hoffer et al., 2017).
3. A CCOUNTING FOR GENERALIZATION THROUGH GROUP ADJUSTMENTS IMPROVES DRO
We have optimized thus far for the robust training loss via DRO (4), relying on regularization to translate good training loss to good test loss by controlling the generalization gap δ .
Here, we show that we can improve performance by directly optimizing for an estimated upper bound on the robust test loss, using ideas from structural risk minimization (Vapnik, 1992).
The key consideration is that each group corresponding to g has its own generalization gap δ
[Equation] X, Y )) [Equation]
X [Equation] Y ) [Equation]
[Equation] θ [Equation] ( X, Y ))] .
To approximate optimizing for the robust test loss R ( θ )
= max g [Equation] E
[Equation]
X
[Equation]
∼ [Equation] [Equation] θ ; ( X, Y ))]
+ δ g , we propose using the simple, parameter-independent heuristic ˆ δ g = C/ √ n g , where n g is the group size for g and C is a model capacity constant which we treat as a hyperparameter.
This gives the group-adjusted DRO estimator
[Equation]
The scaling with 1 / √ n g reﬂects how smaller groups are more prone to overﬁtting than larger groups, and is inspired by the general size dependence of model-complexity-based generalization bounds (see, e.g., Cao et al.
Indeed, even in the regularized setting of Section 3.2, different groups g can have signiﬁcant differences in their generalization gaps δ g : for example, in the strong- (cid:96) 2regularization DRO model, the smallest group in Waterbirds has a train-test robust accuracy gap of 13 .
9% compared to just 0 .
6% for the largest group.
By incorporating group adjustments in (5), we encourage the model to focus more on ﬁtting the smaller groups.
Results.
We evaluate the group adjustments in the strong- (cid:96) 2 -regularization setting of Section 3.2.
In Waterbirds ( λ = 1 . 0 ), group adjustments improve robust test accuracy by 6 . 3% , cutting the error rate almost in half (Table 2 and Figure 3).
The improvements in CelebA ( λ = 0 . 1 ) are more modest, with robust accuracy increasing by 1 . 6% ; (cid:96) 2 regularization is more effective in CelebA and there is not as much variation in the generalization gaps by group at λ = 0 .
Naïve
Adjust
Naïve
Adjust 92.8 96.1 92.4 86.
Waterbirds 96.
CelebA 93.
Average Accuracy Robust Accuracy
[Figure]
land is large, dragging down robust accuracy; at C = 3 (which has the best robust accuracy), the accuracies are balanced; and at = 5 , we overcompensate for group sizes, so smaller groups (e.g., waterbirds on land) do better at the expense of larger groups (e g , landbirds on Empirically, we ﬁnd that group adjustments also help in the early stopping setting of Section 3.2 (in the benchmark in the next section, we consider models with group adjustments and early stopping across a grid of (cid:96) 2 regularization strengths, and report on the one with highest validation accuracy).
However, it is difﬁcult to rigorously study the effects of early stopping (e.g., because the group losses have not converged to a stable value), so we leave a more thorough investigation of the interaction between early stopping and group adjustments to future work.
C OMPARISON BETWEEN GROUP DRO AND IMPORTANCE WEIGHTING
Our results above show that appropriately-regularized group DRO models can be signiﬁcantly more robust than ERM models.
Here, we show theoretically and empirically that DRO also outperforms a 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 ∈ ∆ m to groups, an importance-weighted estimator would learn
[Equation]
Empirical benchmark.
We consider an importance-weighted baseline with weights set to the inverse training frequency of each group, w g = 1 / E g (cid:48) ∼ ˆ P
[ I ( g (cid:48) = g )] .
This optimizes for a test distribution with uniform group frequencies and is analogous to the common reweighting technique for label shifts (Cui et al., 2019; Cao et al., 2019).
Concretely, we train our weighted model by sampling with equal probability from each group for each minibatch (Shen et al., 2016), since a recent study found this to be more effective than similar methods (Buda et al., 2018).
Unlike group DRO, resampling for uniform group size does not necessarily yield uniformly low training losses across groups in practice, as some groups are easier to ﬁt than others.
To compare resampling (RS) with ERM and group DRO, we benchmark each objective, training models across the same grid of (cid:96) 2 regularization strengths and early stopping at the epoch with best robust validation accuracy (Table 3).
In CelebA and Waterbirds, resampling performs much better than ERM but is slightly outperformed by group DRO.
However, resampling fails on MultiNLI, achieving lower clean and robust accuracies than even ERM.
With resampling, it appears that the rare group is overemphasized and extremely low training loss is achieved for that group at the cost of others. [Table]
[Figure]
Theoretical comparison.
Should we 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.
For this analysis, we shift to a more general framework of importance weighting, in which we minimize the expected loss E z ∼ P [ w ( z ) (cid:96) ( θ ; z )] over some source distribution P .
Here, we assign weights to each data point z , abstracting away from our particular choice of assigning the weights according to its group.
Minimizing the weighted objective is equivalent to minimizing the expected unweighted loss E z ∼ Q [ (cid:96) ( θ ; z )] over a target distribution Q such that Q ( z ) ∝ w ( z ) P ( z ) .
However, we want good worst-case performance over a family of Q ∈ Q instead of optimizing for a single Q .
Are there weights w such that the resulting model ˆ θ w achieves optimal robust risk?
In the convex regime, standard duality arguments show that this is true (see Appendix A.1 for the proof):
Proposition 1.
Suppose that the loss (cid:96) ( · ; z ) is continuous and convex for all z in Z , and let the uncertainty set Q be a set of distributions supported on Z . Assume that Q and the model family Θ ⊆
R d are convex and compact, and let θ ∗ ∈
Θ be a minimizer of the robust objective R ( θ ) .
Then there exists a distribution Q ∗ ∈ Q such that θ ∗ ∈ arg min E z ∼ Q ∗ [ (cid:96) ( θ ; z )] .
However, this equivalence breaks down when the loss (cid:96) is non-convex (as in neural networks).
Counterexample 1. Consider a uniform data distribution P supported on two points
Z = { z 1 , z 2 } , and let (cid:96) ( θ ; z ) be as in Figure 4, with Θ = [0 , 1] .
The robust solution θ ∗ achieves a worst-case loss of R ( θ ∗ ) = 0 .
Now consider any weights ( w 1 , w 2 ) ∈ ∆ 2 and w.l.o.g. let w 1 ≥ w 2 . The minimizer of the weighted loss w 1 (cid:96) ( θ ; z 1 )+ w 2 (cid:96) ( θ ; z 2 ) is θ 1 , which only attains a worst-case loss of R ( θ ∗ )
= 1 . 0 .
This negative result implies that in the non-convex setting, there may not be any choice of weights w such that the resulting minimizer ˆ θ w is robust.
Even if such weights did exist, their values are unknown a priori.
In particular, the weights depend on θ ∗ , highlighting the beneﬁt of algorithms like group DRO that simultaneously learn θ and Q .
Common choices of weights, such as inverse group size, are heuristics that may not yield robust solutions (as observed for MultiNLI in Table 3).
A LGORITHM
To train group DRO models efﬁciently, we introduce an online optimization algorithm; there are no existing stochastic optimization algorithms for group DRO with convergence guarantees.
In the convex and batch case, there is a rich literature on distributionally robust optimization which treats the problem as a standard convex conic program (Ben-Tal et al., 2013; Duchi et al., 2016;
Bertsimas et 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, 2018;
Hashimoto et al., 2018), and (ii) direct minimax optimization (Namkoong & Duchi, 2016).
The ﬁrst approach fails for group DRO because the gradient of the dual objective is difﬁcult to estimate in a stochastic and unbiased manner.
An algorithm of the second type has been proposed for group DRO (Oren et al., 2019), but it lacks convergence guarantees and we observed instability in practice under some settings.
Recall that we aim to solve the optimization problem
[Equation]
Extending existing minimax algorithms for DRO (Namkoong & Duchi, 2016; Oren et al., 2019), we interleave gradient-based updates on θ and q .
Intuitively, we maintain a distribution q over groups, with high masses on high-loss groups, and update on each example proportionally to the mass on its group.
Concretely, we interleave SGD on θ and exponentiated gradient ascent on q (Algorithm 1).
The key improvement from the existing group DRO algorithm (Oren et al., 2019) is that q is updated using gradients instead of picking the group with worst average loss at each iteration, which is important for algorithmic stability and obtaining convergence guarantees.
In practice, we implement the algorithm with minibatching and evaluate model performance with the last iterate θ ( T ) ; we observe that the algorithm reliably converges in loss.
The run time of the algorithm is similar to that of SGD for a given number of epochs (less than 5% difference), as run time is dominated by the computation of the loss and its gradient.
Input: Step sizes η q , η θ ; P g for each g ∈ G
Initialize θ (0) and q (0) for t=1,...,T do g ∼ Uniform (1 , . . .
, m ) x, y ∼ P g q (cid:48) ← q ( t − 1) q (cid:48) g ← q (cid:48) g exp( η q (cid:96) ( θ ( t − 1) ; ( x, y ))) q ( t ) ← q (cid:48) / (cid:80) g q (cid:48)
g θ ( t ) ← θ ( t − 1)
− η θ q g ( t ) ∇ (cid:96) ( θ ( t − 1) ; ( x, y )) end
We analyze the convergence rate by studying the error ε T of the average iterate ¯ θ (1: T ) :
[Equation]
(8) where L ( θ, q ) :
= (cid:80) mg
=1 q g E ( x,y ) ∼ P g [ (cid:96) ( θ ; ( x, y ))] is the expected worst-case loss.
Applying results from Nemirovski et al. (2009), we can show that Algorithm 1 has a standard convergence rate of O (cid:0) 1 / √ T (cid:1) in the convex setting (proof in Section A.2):
Theorem 1.
Suppose that the loss (cid:96) ( · ; ( x, y )) is non-negative, convex, B ∇ -Lipschitz continuous, and bounded by B (cid:96) for all ( x, y ) in X × Y , and (cid:107) θ (cid:107) 2 ≤ B θ for all θ ∈ Θ with convex Θ ⊆ R d .
Then, the average iterate of Algorithm 1 achieves an expected error at the rate
[Equation]
where the expectation is taken over the randomness of the algorithm.
R ELATED WORK
The problem of non-uniform accuracy.
Other approaches to addressing non-uniform accuracy over the data distribution include domain adaptation techniques for known target distributions (BenDavid et al., 2006;
Ganin & Lempitsky, 2015) and work in ML fairness (Dwork et al., 2012; Hardt et al., 2016; Kleinberg et al., 2017).
As we discuss in Section 4, importance weighting is a classic example of the former (Shimodaira, 2000).
Byrd & Lipton (2019) empirically study importance weighting in neural networks and demonstrate that it 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 deﬁnes the uncertainty set Q as a divergence ball around the training distribution over ( x, y )
(Ben-Tal et al., 2013; Miyato et al., 2015; Lam & Zhou, 2015; Blanchet & Murthy, 2016; Duchi et al., 2016; Esfahani & Kuhn, 2018; Bertsimas et al., 2018).
With small divergence balls of radii O (1 /n ) , DRO acts as a regularizer (Shaﬁeezadeh-Abadeh et al., 2015; Duchi & Namkoong, 2016).
However, when the radius is larger, the resulting Q can be too pessimistic. To better control Q , 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, prior work on DRO has operated in settings with training loss trade-offs and reasonable generalization. Most DRO work has focused on the classic (underparameterized) model setting (Namkoong & Duchi, 2017; Hu et al., 2018; Duchi et al., 2019); Sinha et al. (2018) study neural networks but with a more conservative Wasserstein uncertainty set that leads to nonvanishing training loss; and Oren et al. (2019) study neural networks but for generative modeling.
Generalization of robust models.
There is extensive work investigating generalization of neural networks in terms of average loss, theoretically and empirically (Hardt et al., 2015; Szegedy et al., 2016; Hoffer et al., 2017).
However, analysis on robust losses is limited.
For label shifts, prior work has observed overﬁtting on rare labels and proposed algorithms to mitigate it (Buda et al., 2018; Cui et al., 2019; Cao et al., 2019).
In the DRO literature, generalization bounds on the DRO objective exist for particular uncertainty sets (Duchi & Namkoong, 2018), but these bounds do not apply directly to group DRO.
Invariant prediction models from the causal inference literature similarly aim to achieve high performance on a range of test distributions (Peters et al., 2016; B ¨ uhlmann & Meinshausen, 2016; Heinze-Deml & Meinshausen, 2017; Rothenh ¨ ausler et al., 2018; Yang et al., 2019).
The maximin regression framework (Meinshausen & B ¨ uhlmann, 2015) also assumes groupbased shifts, but focuses on settings without the generalization problems identiﬁed in our work.
D ISCUSSION
In this paper, we analyzed group DRO in the context of overparameterized neural networks and highlighted the importance of regularization for worst-case group generalization.
When appropriately regularized, group DRO signiﬁcantly improves robust accuracy at a small cost in average accuracy.
As an application of group DRO, we showed that it can prevent models from learning pre-speciﬁed spurious correlations.
While handling shifts beyond pre-speciﬁed group shifts is important future work, existing work has identiﬁed many distributional shifts that can be expressed with pre-speciﬁed groups: e.g., batch effects in biology (Leek et al., 2010), or image artifacts (Oakden-Rayner et al., 2019) and patient demographics (Badgeley et al., 2019) in medicine. Our supplemental experiments also suggest that group DRO models can maintain high robust accuracy even when groups are imperfectly speciﬁed (Appendix B).
More generally, our observations call for a deeper analysis of average vs. worst-case generalization in the overparameterized regime. Such analysis may shed light on the failure modes of deep neural networks as well as provide more precise tools (beyond strong (cid:96) 2 regularization or early stopping) to counter poor worst-case generalization while maintaining high average accuracy.
R EPRODUCIBILITY
We will provide code for training group DRO models and scripts that replicate the experiments above.
The constructed Waterbirds dataset will be publically available for download, with a script to adjust its generation (e g., to choose different object backgrounds or proportions)
A P ROOFS A. E QUIVALENCE OF DRO AND IMPORTANCE WEIGHTING IN THE CONVEX SETTING
Proposition 1.
Suppose that the loss (cid:96) ( · ; z ) is continuous and convex for all z in Z and let the uncertainty set Q be a set of distributions supported on Z . Assume that Q and the model family Θ ⊆
R d are convex and compact and let θ ∗ ∈
Θ be a minimizer of the robust objective R ( θ ) .
Then there exists a distribution Q ∗ ∈ Q such that θ ∗ ∈ arg min E z ∼ Q ∗ [ (cid:96) ( θ ; z )] .
Proof.
Let h ( θ, Q ) :=
E Z ∼ Q
[ (cid:96) ( θ ; Z )] .
Since the loss (cid:96) ( θ ; Z ) is continuous and convex in θ for all Z in Z , we have that h ( θ, Q ) is continuous, convex in θ , and concave (linear) in Q .
Moreover, since convexity and lower semi-continuity are preserved under arbitrary pointwise suprema, sup Q ∈Q h ( θ, Q ) is 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
[Equation]
is attained at some θ ∗ ∈ Θ .
A similar argument implies that the sup-inf objective
[Equation]
attains its optimum at some Q ∗ ∈ Q .
Moreover, because Θ and Q are compact and h is continuous, we have the max-min equality (see, e.g., Ex 5.25 in Boyd & Vandenberghe (2004))
[Equation]
Together, the above results imply that ( θ ∗ , Q ∗ ) form a saddle point (proposition 3.4.1, Bertsekas (2009)), that is,
[Equation]
In particular, the second equality indicates that the optimal DRO model θ ∗ also minimizes the weighted risk h ( θ, Q ∗ ) =
E Z ∼ Q [ (cid:96) ( θ ; Z )] , as desired.
A. C ONVERGENCE RATE OF A LGORITHM
Theorem 1.
Suppose that the loss (cid:96) ( · ; ( x, y )) is non-negative, convex, B ∇ -Lipschitz continuous, and bounded by B (cid:96) for all ( x, y ) in X × Y , and (cid:107) θ (cid:107) 2 ≤ B θ for all θ ∈ Θ with convex Θ ⊆ R d .
Then, the average iterate of Algorithm 1 achieves an expected error at the rate
[Equation]
where the expectation is 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.
Consider the saddle-point optimization problem
[Equation]
under the following assumptions:
Assumption 1.
f g ( θ ) is convex on Θ .
Assumption 2.
f g ( θ ) =
E ξ ∼ q [ F g ( θ ; ξ )] for some function F g .
Assumption 3.
We can generate i.i.d.
examples ξ ∼ q .
For a given θ ∈ Θ and ξ ∈ Ξ , we can compute F g ( θ, ξ ) and unbiased stochastic subgradient ∇ F g ( θ ; ξ ) , that is, E ξ ∼ q [ ∇ F g ( θ ; ξ )]
=
∇ f g ( θ ) .
Assume that we apply online mirror descent with some c -strongly convex norm (cid:107) · (cid:107) θ to obtain the iterates θ (1) , . .
, θ ( T ) and q (1) , . . .
, q ( T ) .
From Nemirovski et al. (2009), we have the following:
Theorem 2 (Nemirovski et al. (2009), Eq 3.23) .
Suppose that Assumptions 1-3 hold.
Then the pseudo-regret of the average iterates ¯ q g (1: T ) and ¯ q g (1: T ) can be bounded as
[Equation]
where
[Equation]
for c -strongly convex norm (cid:107) · (cid:107) θ .
It remains to formulate our algorithm as an instance of online mirror descent applied to the saddlepoint problem above.
We start by deﬁning the following:
Deﬁnition 1.
Let q be a distribution over ξ =
( x, y, g )
that is a uniform mixture of individual group distributions Q g :
[Equation]
Deﬁnition 2.
Let F g (cid:48)
( θ ; ( x, y, g ))) :
= m I [ g = g (cid:48) ] (cid:96) ( θ ; ( x, y )) .
Correspondingly, let f g (cid:48) :
=
E
Q g (cid:48)
[ (cid:96) ( θ ; ( x, y ))] .
We now check that Assumptions 1-3 hold under the original assumptions in the statement of Theorem 1:
We assume that the loss (cid:96) ( · ; ( x, y )) is non-negative, continuous, and convex for all ( x, y ) in X × Y .
As a result, f g ( θ ) is non-negative, continuous, and convex on Θ . 2.
The expected value of F g ( θ ) over distribution q is f g ( θ ) :
[Equation]
We can compute an unbiased stochastic subgradient ∇ F
[Equation]
θ
[Equation]
x
[Equation]
y
[Equation]
Finally
[Equation]
we compute the constants required for the regret bound in Theorem 2.
Recalling the original assumptions of Theorem 1,
Bounded losses: (cid:96) ( θ ; ( x, y ))
≤ B (cid:96) for all x, y, θ 2.
Bounded gradients: (cid:107)∇ (cid:96) ( θ ; ( x, y ))
(cid:107) 2 ≤ B ∇ for all θ, x, y 3.
Bounded parameter norm: (cid:107) θ (cid:107) 2 ≤ B θ for all θ ∈ Θ ,
we obtain:
[Equation]
[Equation]
Plugging in these constants into the regret bound from Theorem 2, we obtain E
[Equation]
This implies Theorem 1 because the minimax game is convex-concave.
B S UPPLEMENTARY EXPERIMENTS
Group DRO can maintain high robust accuracy even when spurious attributes are not perfectly speciﬁed.
We repeat the CelebA experiment on strong (cid:96) 2 regularization (Section 3.2) but with inexact group speciﬁcations:
Instead of the ground-truth spurious attribute Male , we provide a related attribute Wearing Lipstick , and 2.
We also specify four distractor/non-spurious attributes ( Eyeglasses , Smiling , Double Chin , and Oval Face ).
Optimizing for worst-case performance over all 2 6 = 64 groups (for all combinations of 5 attributes and 1 label), the DRO model attains 78 .
9% robust accuracy across the 4 original groups (darkhaired males and females, and blond males and females).
These robust accuracies are not far off from the original DRO model with just the ground-truth spurious attribute ( 86 . 7% ) and signiﬁcantly outperform the ERM model ( 37 . 8% ).
C E XPERIMENTAL DETAILS C. D ATASETS
MultiNLI.
The standard MultiNLI train-test split allocates most examples (approximately 90% ) to the training set, with another 5% as a publicly-available development set and the last 5% as a heldout test set that is only accessible through online competition leaderboards (Williams et al., 2018).
To accurately estimate performance on rare groups in the validation and test sets, we combine the training set and development set and then randomly resplit it to a 50 − 20 − 30 train-val-test split that allocates more examples to the validation and test sets than the standard split.
We use the provided gold labels as the target, removing examples with no consensus gold label (as is standard procedure).
We annotate an example as having a negation word if any of the words nobody , no , never , and nothing appear in the hypothesis (Gururangan et al., 2018).
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 as a waterbird if it belongs to one of the following (groups of) species: albatross, auklet, frigatebird, fulmar, gull, jaeger, pelican, tern, fulmar, gadwall, grebe, mallard, or guillemot; and landbird otherwise.
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 or 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 90% of all waterbirds against a water background and the remaining 10% against a land background.
Similarly, 90% of all landbirds are placed against a land background with the remaining 10% against water.
We refer to this combined CUB-Places dataset as the Waterbirds dataset to avoid confusion with the original ﬁne-grained species classiﬁcation task in the CUB dataset.
We use the ofﬁcial train-test split of the CUB dataset, randomly choosing 20% of the training data to serve as a validation set.
For the validation and test sets, we allocate equal numbers of birds to each of the four groups (i.e., landbirds and waterbirds are uniformly distributed on land and water backgrounds) so as to be able to more accurately measure the performance of the rare groups.
This is particularly important for the Waterbirds dataset because of its relatively small size; otherwise, the smaller groups (waterbirds on land and landbirds on water) would have too few samples to 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.
Due to the above procedure, when reporting average test accuracy in our experiments, we calculate 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 ofﬁcial train-val-test split that accompanies the CelebA celebrity face dataset (Liu et al., 2015).
We use the Blond Hair attribute as the target label and the Male attribute as the spuriously-associated variable.
C. M ODELS
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 a 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 et al., 2016).
As in the original paper, we used batch normalization (Ioffe & Szegedy, 2015) and no dropout (Srivastava et al., 2014).
For simplicity, we train all models without data augmentation.
We use a ﬁxed learning rate instead of the standard adaptive learning rate schedule to make our different model types easier to directly compare, since we expected the scheduler to interact differently with different model types (e.g., due to the different deﬁnition of loss).
The interaction between batch norm and (cid:96) 2 regularization means that we had to adjust learning rates for each different setting of (cid:96) 2 regularization (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 (cid:96) 2 regularization strength of λ = 0 .
(as in He et al.
for both Waterbirds and CelebA, with a learning rate of 0 .
001 for Waterbirds and 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, we train each ResNet50 model for 1 epoch.
For the strong (cid:96) 2 regularization experiments in that section, we use λ = 1 . 0 for Waterbirds and λ = 0 .for
CelebA, with both datasets using a learning rate of 0 .
These settings of λ differ because we found that the lower value was sufﬁcient for controlling overﬁtting on CelebA but not on Waterbirds.
For the group adjustment experiments in Section 3.3, we use the same settings of λ = 1 . 0 for
Waterbirds and λ = 0 .
1 for CelebA, with both datasets using a learning rate of 0 .
For both datasets, we search over group adjustments of C ∈ { 0 , 1 , 2 , 3 , 4 , 5 } and pick the model with the best robust validation accuracy.
For the benchmark in Section 4 (Table 3), we grid search over (cid:96) 2 regularization strengths of λ ∈ { 0 . 0001 , 0 . 1 , 1 . 0 } for Waterbirds and λ ∈ { 0 . 0001 , 0 . 01 , 0 . 1 } for CelebA, using the corresponding learning rates for each λ and dataset listed above.
(Waterbirds and CelebA at λ = 0 . 1 , which is not listed above, both use a learning rate of 0 . 0001 .)
To avoid advantaging DRO by allowing it to try many more hyperparameters, we only test group adjustments on the (cid:96) 2 regularization strengths used in Section 3.3, i.e., λ = 1 . 0 for Waterbirds and λ = 0 .
1 for CelebA.
All benchmark models were evaluated at the best early stopping epoch (as measured by robust validation accuracy).
BERT.
We use the Hugging Face pytorch-transformers implementation of the BERT bert-base-uncased model, starting from pretrained weights (Devlin et al., 2018).
We use the default tokenizer and model settings from that implementation, including a ﬁxed linearly-decaying learning rate starting at 0 . 00002 , AdamW optimizer, dropout, and no (cid:96) 2 regularization ( λ = 0 ), except that we use a batch size of 32 (as in Devlin et al.
(2018)) 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.
For the standard training experiments in Section 3.1, we train for 20 epochs.
For the (cid:96) 2 regularization experiments in Section 3.2, we tried regularization strengths of λ ∈ { 0 . 01 , 0 .
However, these models had similar or worse robust accuracies compared to the default model with no (cid:96) 2 regularization.
For the early stopping experiments in Section 3.2, we train for 3 epochs, which is the suggested early-stopping time in Devlin et al.
For the benchmark in Section 4 (Table 3), we similarly trained for 3 epochs.
All benchmark models were evaluated at the best early stopping epoch (as measured by robust validation accuracy).