Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization

Introduction

Machine learning models are typically trained to minimize the average loss on a training set, with the goal of achieving high accuracy on an independent and identically distributed (i.i.d.) test set. However, models that are highly accurate on average can still consistently fail on rare and atypical examples (Hovy & Søgaard, 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 always hold. For example, in natural language inference (NLI)—determining if two sentences agree or contradict—the presence of negation words like ‘never’ is strongly correlated with contradiction due to artifacts in crowdsourced training data (Gururangan et 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 no negation words).

To avoid learning models that rely on spurious correlations and therefore suffer high loss on some groups of data, we instead train models to 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 no negation words in the NLI example above. This training procedure is an instance of distributionally robust optimization (DRO), 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; Shafieezadeh-Abadeh et al., 2015; Namkoong & Duchi, 2017; Duchi & Namkoong, 2018; Hashimoto et al., 2018).

We study group DRO in the context of overparameterized 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 modified version of the CUB dataset (Wah et al., 2011). The problem with applying DRO to overparameterized models is that if a 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., 2014). In the vanishing-training-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 test accuracies (Section 3.1). In other words, the generalization gap is small on average but large for the worst group.

Finally, to carry out the experiments, we introduce a new 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 5).

Setup

where ${\hat{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 expected loss over an uncertainty set of distributions $\mathcal{Q}$:

The uncertainty set $\mathcal{Q}$ encodes the possible test distributions that we want our model to perform well on. Choosing a general family $\mathcal{Q}$, such as a divergence ball around the training distribution, confers robustness to a wide set of distributional shifts, but can also lead to overly pessimistic models which optimize for implausible worst-case distributions (Duchi et al., 2019).

To construct a realistic set of possible test distributions without being overly conservative, we leverage prior knowledge of spurious correlations to define groups over the training data and then define the uncertainty set $\mathcal{Q}$ in terms of these groups. Concretely, we adopt 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 $\mathcal{G}=\{1,2,\ldots,m\}$. In our main experiments, $m=4$ or $6$; we also use $m=64$ in our supplemental experiments. We define the uncertainty set $\mathcal{Q}$ as any mixture of these groups, i.e., $\mathcal{Q}:=\{\sum_{g=1}^{m}q_{g}P_{g}:q\in\Delta_{m}\}$, where $\Delta_{m}$ is the $(m-1)$-dimensional probability simplex; this choice of $\mathcal{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 (2) is equivalent to a maximum over the expected loss of each group,

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 learn a group DRO model minimizing the empirical worst-group risk $\hat{}\mathcal{R}(\theta)$:

where each group ${\hat{P}}_{g}$ is an empirical distribution over all training points $(x,y,g^{\prime})$ with $g^{\prime}=g$ (or equivalently, a subset of training examples drawn from $P_{g}$). 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 $\delta:=\mathcal{R}(\theta)-\hat{\mathcal{R}}(\theta)$. We will show that for overparameterized neural networks, $\delta$ is large unless we apply sufficient regularization.

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)\in\mathcal{A}$ that is spuriously correlated with the label $y$, and we use this prior knowledge to form $m=|\mathcal{A}|\times|\mathcal{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 the correlation does not hold and hence do worse on the worst-group loss $\mathcal{R}(\theta)$.

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 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 $\mathcal{Y}=\{\text{waterbird},\text{landbird}\}$ and place it on one of $\mathcal{A}=\{\text{water background},\text{land background}\}$, with waterbirds (landbirds) more frequently appearing against a 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 ($\mathcal{Y}=\{\text{blond},\text{dark}\}$) as the target and gender ($\mathcal{A}=\{\text{male},\text{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 a 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 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 $\mathcal{Y}=\{\text{entailed},\text{neutral},\text{contradictory}\}$ and spurious attributes $\mathcal{A}=\{\text{no negation},\text{negation}\}$. There are $n=206175$ examples in our training set, with $1521$ examples in the smallest group (entailment with negations); see Appendix C.1 for more details on 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 inference which achieve high average test accuracies on their respective tasks.

We train the ERM (1) and group DRO (4) models using standard (minibatch) stochastic gradient descent and (minibatch) stochastic algorithm introduced in Section 5, respectively. 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 worst-group accuracy on the worst-performing group.

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 regularization and hyperparameter settings (He et al., 2016; Devlin et al., 2019), regardless of whether they are trained with ERM or group DRO. 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., 2019; Radford et al., 2019) but fine-tuned for a fixed small number of epochs because average test accuracy levels off quickly; we verified that training to convergence gave equally high average test accuracy.

As expected, ERM models attain near-perfect worst-group training accuracies of at least $99.9\%$ on all three datasets and also obtain high average test accuracies ($97.3\%$, $94.8\%$, and $82.5\%$ on Waterbirds, CelebA, and MultiNLI). However, they perform poorly on the worst-case group at test time with worst-group accuracies of $60.0\%$, $41.1\%$, and $65.7\%$ respectively (Table 1, 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 therefore near-optimal for both the ERM (1) and DRO (4) objectives. Indeed, we find that group DRO models perform similarly to ERM models, attaining near-perfect training accuracies and high average test accuracies, but poor worst-group test accuracies (Table 1, Figure 2).

Discussion.

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 regularization is not critical for average performance: models can do well on average even when all regularization is removed (Zhang et al., 2017), and default regularization settings (like in the models trained above) still allow models to perfectly fit the training data. Here, we study if increasing regularization strength—until the models no longer perfectly fit the training data—can rescue worst-case performance. We find that departing from the vanishing-training-loss regime allows DRO models to significantly outperform ERM models on worst-group test accuracy while maintaining high average accuracy. We investigate two types of regularization:

Early stopping.

Discussion.

We conclude that regularization—preventing 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 question of good worst-group training accuracy. Since no 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 models 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 Accounting for generalization through group adjustments improves DRO

The scaling with $1/\sqrt{n_{g}}$ reflects how smaller groups are more prone to overfitting than larger groups, and is inspired by the general size dependence of model-complexity-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 of using a $1/\sqrt{n}$ 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/\sqrt{n}$ term would simply be a constant and not affect the optimization.

Comparison between DRO and importance weighting

Our results above show that strongly-regularized DRO models can be significantly 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\in\Delta_{m}$ to groups, an importance-weighted estimator would learn

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.

This negative result implies that in the non-convex setting, there may not be any choice of weights $w$ such that the resulting minimizer ${\hat{\theta}_{w}}$ is robust. Even if such weights did exist, they depend on $\theta^{*}$ and obtaining these weights requires that we solve a dual DRO problem, making reweighting no 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).

Algorithm

To train group DRO models efficiently, we introduce an online optimization algorithm with convergence 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.

Recall that we aim to solve the optimization problem (4), which can be rewritten as

Extending existing minimax algorithms for DRO (Namkoong & Duchi, 2016; Oren et al., 2019), we interleave gradient-based updates on $\theta$ 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 $\theta$ and exponentiated gradient ascent on $q$ (Algorithm 1). (In practice, we use minibatches and a momentum term for $\theta$; see Appendix C.2 for details.) 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 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.

We analyze the convergence rate by studying the error $\varepsilon_{T}$ of the average iterate $\bar{\theta}{{}^{(1:T)}}$:

where the expectation is taken over the randomness of the algorithm.

Related work

Existing approaches to addressing non-uniform 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 et al., 2016a; 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 defines the uncertainty set $\mathcal{Q}$ as a divergence ball around the training distribution over $(x,y)$ (Ben-Tal et al., 2013; Lam & Zhou, 2015; Duchi et al., 2016; Miyato et al., 2018; Esfahani & Kuhn, 2018; Bertsimas et al., 2018; Blanchet & Murthy, 2019). With small divergence balls of radii $O(1/n)$, DRO acts as a regularizer (Shafieezadeh-Abadeh et al., 2015; Namkoong & Duchi, 2017). However, when the radius is larger, the resulting $\mathcal{Q}$ can be too pessimistic. In contrast, group DRO considers $\mathcal{Q}$ that is of wider radius but with fewer degrees of freedom (shifts over groups instead of over $(x,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 (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 non-vanishing training loss; and Oren et al. (2019) study neural networks but for generative modeling where loss tradeoffs 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., 2016b; Szegedy et 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., 2019). In the DRO literature, generalization bounds on the DRO objective exist for particular uncertainty sets (e.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., 2016; Bühlmann & Meinshausen, 2016; Heinze-Deml & Meinshausen, 2017; Rothenhäusler et al., 2018; Yang et al., 2019; Arjovsky et al., 2019). For example, the maximin regression framework (Meinshausen & Bühlmann, 2015) also assumes group-based shifts, but focuses on settings without the generalization problems identified in our work.

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 spurious 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 handling 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 biology (Leek et al., 2010), or image artifacts (Oakden-Rayner et al., 2019) and patient demographics (Badgeley et al., 2019) in medicine.

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 provided 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 https://github.com/kohpangwei/group_DRO. The datasets used in this paper are also available at that link, as well as scripts to modify dataset generation (e.g., to choose different spurious attributes for CelebA and MultiNLI, or different object backgrounds or relative group sizes for Waterbirds).

References

Appendix A Proofs

Together with the compactness of $\Theta$ and $\mathcal{Q}$, the above conditions imply (by Weierstrass’ theorem, proposition 3.2.1, Bertsekas (2009)), that the optimal value of the DRO objective

is attained at some $\theta^{*}\in\Theta$.

A similar argument implies that the sup-inf objective

attains its optimum at some $Q^{*}\in\mathcal{Q}$.

Moreover, because $\Theta$ and $\mathcal{Q}$ are compact and $h$ is continuous, we have the max-min equality (see, e.g., Ex 5.25 in Boyd & Vandenberghe (2004))

Together, the above results imply that $(\theta^{*},Q^{*})$ form a saddle point (proposition 3.4.1, Bertsekas (2009)), that is,

A.2 Convergence rate of Algorithm 1

where the expectation is taken over the randomness of the algorithm.

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

Online mirror descent with some $c$-strongly convex norm $\|\cdot\|_{\theta}$, yielding iterates $\theta{{}^{(1)}},\ldots,\theta{{}^{(T)}}$ and $q{{}^{(1)}},\ldots,q{{}^{(T)}}$, has the following guarantee.

Suppose that Assumptions 1-3 hold. Then the pseudo-regret of the average iterates $\bar{q}_{g}{{}^{(1:T)}}$ and $\bar{q}_{g}{{}^{(1:T)}}$ can be bounded as

for online mirror descent with $c$-strongly convex norm $\|\cdot\|_{\theta}$.

It 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:

Let $q$ be a distribution over $\xi=(x,y,g)$ that is a uniform mixture of individual group distributions $P_{g}$:

We now check that Assumptions 1-3 hold under the original assumptions in the statement of Theorem 2:

The expected value of $F_{g}(\theta)$ over distribution $q$ is $f_{g}(\theta)$:

We can compute an unbiased stochastic subgradient $\nabla F_{g^{\prime}}(\theta;(x,y,g))$

Finally, we compute the constants required for the regret bound in Theorem 1. Recalling the original assumptions of Theorem 2,

Bounded parameter norm: $\|\theta\|_{2}\leq B_{\Theta}$ for all $\theta\in\Theta$,

Plugging in these constants into the regret bound from Theorem 1, we obtain

This implies Theorem 2 because the minimax game is convex-concave. ∎

Appendix B Supplementary experiments

Instead of the ground-truth spurious attribute Male, we provide a related attribute Wearing Lipstick, and

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 (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-truth spurious attribute ($86.7\%$) and significantly outperform the ERM model ($37.8\%$).

Appendix C Experimental details

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 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 create our own validation and test sets by combining the training set and development set and then randomly shuffling them into a $50-20-30$ train-val-test split. We chose to allocates more examples to the validation and test sets than the standard split to allow us to accurately estimate performance on rare groups in the validation and test sets.

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 is a seabird (albatross, auklet, cormorant, frigatebird, fulmar, gull, jaeger, kittiwake, pelican, puffin, or tern) or waterfowl (gadwall, grebe, mallard, merganser, guillemot, or Pacific loon). Otherwise, we label it as 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 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 $95\%$ of all waterbirds against a water background and the remaining $5\%$ against a 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 a 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 us 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 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.

In a 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, if 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 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 official 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.2 Models

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.

For the group adjustment experiments in Section 3.3, we use the same settings of $\lambda=1.0$ for Waterbirds and $\lambda=0.1$ for CelebA, with both datasets using a learning rate of $0.00001$. For both datasets, we use the value of $C\in\{0,1,2,3,4,5\}$ found in the benchmark grid search described below.

BERT.

For the standard training experiments in Section 3.1, we train for $20$ epochs.

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. (2019).

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