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

By Shiori Sagawa, Pang Wei Koh, Tatsunori B. Hashimoto, Percy Liang. 2019

1,692 citations recorded by the source corpus.

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---a stronger-than-typical L2 penalty 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, with convergence guarantees, to efficiently train group DRO models.

Subfields: Artificial Intelligence, Safety, Robustness & Interpretability

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Safety, Robustness & Interpretability

Citations in the corpus

Adversarial Training for High-Stakes Reliability (2022, methodology): “[44] try to solve simpler vision tasks such as WaterBirds [44] or CelebA [45], which resembles our approach of picking a relatively simple task of “don’t exacerbate injuries”.”
Improving Task-free Continual Learning by Distributionally Robust Memory Evolution (2022, background): “There are various applications of DRO in machine learning problems, including tackling the group-shift (Sagawa et al., 2020), subpopulation shift (Zhai et al., 2021), and class imbalances (Xu et al., 2020).”
In Search of Lost Domain Generalization (2020, methodology): “Hu et al. [2016], Sagawa et al. [2019] weight the importance of the minibatches of the training distributions proportional to their error.”
WILDS: A Benchmark of in-the-Wild Distribution Shifts (2020, positive): “• Group DRO (Hu et al., 2018; Sagawa et al., 2020) uses distributionally robust optimization (DRO) to explicitly minimize the loss on the worst-case domain, RDRO(θ) := max d∈Dtrain ÊPd [`(θ; (x, y))] .”
Meta-learning via Language Model In-context Tuning (2021, background): “For example, the test distribution might constitute new input compositions (Lake and Baroni, 2018), rare subgroups (Sagawa et al., 2019), other types of distribution shifts (Hendrycks and Diet-terich, 2019), or even adversarial examples (Kang et al., 2019).”
An Empirical Study on Robustness to Spurious Correlations using Pre-trained Language Models (2020, background): “This issue has also been referred to as annotation artifacts (Guru-rangan et al., 2018), dataset bias (He et al., 2019; Clark et al., 2019), and group shift (Oren et al., 2019; Sagawa et al., 2020) in the literature.”
Beyond Leaderboards: A survey of methods for revealing weaknesses in Natural Language Inference data and models (2020, methodology): “it the dataset biases (Clark et al., 2019; He et al., 2019), re-weighting the loss function to incorporate the bias in the data (Schuster et al., 2019; Zhang et al., 2019b), parameter regularisation (Sagawa et al., 2020) and the use of external resources, such as linguistic knowledge (Zhou et al., 2”
Test-Time Classifier Adjustment Module for Model-Agnostic Domain Generalization (2021, methodology): “These results include the following algorithms: Empirical Risk Minimization (ERM), Group Distributionally Robust Optimization (GroupDRO) [41], Inter-domain Mixup (Mixup) [56, 58, 53], Meta-Learning for Domain Generalization (MLDG) [28], DomainAdversarial Neural Networks (DANN) [15],…”
Causally motivated shortcut removal using auxiliary labels (2021, methodology): “For example, a DNN trained for image classiﬁcation could exploit correlations between the foreground object and background of images in the training distribution, and use a representation of the background as a shortcut to predict the foreground object [Beery et al., 2018, Sagawa et al., 2019].”
Out-of-Distribution Generalization via Risk Extrapolation (REx) (2020, background): “A classic approach for this setting is group distributionally robust optimization (DRO) (Sagawa et al., 2019), where F contains all mixtures of the training distributions.”
An Investigation of Why Overparameterization Exacerbates Spurious Correlations (2020, methodology): “(2) However, in line with Sagawa et al. (2020), we find that models trained via ERM have poor worst-group test error regardless of whether they are under- or overparameterized (Appendix A.1).”
Meta-learning without data via Wasserstein distributionally-robust model fusion (2022, methodology): “…characterize the underlying uncertainty in various ways, including KL-divergence, Wasserstein ball, etc. DRO has been applied to many machine learning problems, including federated learning [Deng et al., 2021], group shift [Sagawa et al., 2020] and reinforcement learning [Smirnova et al., 2019].”
Just Train Twice: Improving Group Robustness without Training Group Information (2021, methodology): “01 from Sagawa et al. (2020a), without tuning.”
MEMO: Test Time Robustness via Adaptation and Augmentation (2021, background): “The general problem of distribution shift has been studied under a number of frameworks [29], including domain adaptation [33, 4, 37], domain generalization [3, 26, 7], and distributionally robust optimization [2, 14, 31], to name just a few.”
dcbench: a benchmark for data-centric AI systems (2022, background): “It can also help practitioners debug and improve models: after an underperforming slice is identified, we can improve model robustness by either updating the training dataset or using robust optimization techniques [17, 18].”
Influence Tuning: Demoting Spurious Correlations via Instance Attribution and Instance-Driven Updates (2021, methodology): “Spurious correlations can also be mitigated by data-based, optimization-based, and post-processing methods (Zmigrod et al., 2019; Kaushik et al., 2020; Sagawa et al., 2020a; Yaghoobzadeh et al., 2021; Clark et al., 2019).”