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

Shiori Sagawa, Pang Wei Koh+2 more|arXiv (Cornell University)|Nov 20, 2019

Domain Adaptation and Few-Shot LearningComputer Science58 references364 citations

TL;DR

The paper shows that for overparameterized neural networks, group DRO improves worst-group generalization when combined with stronger regularization (e.g., strong L2 or early stopping), and introduces a scalable online algorithm with convergence guarantees for training group DRO models.

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

Motivation & Objective

• Motivate the problem of spurious correlations causing poor performance on atypical groups under i.i.d. training.
• Investigate why naively applying group DRO to overparameterized networks can fail to improve worst-group generalization.
• Demonstrate that stronger regularization enables group DRO to achieve substantial gains in worst-group accuracy while maintaining average accuracy.
• Propose an online optimization algorithm for group DRO with convergence guarantees and analyze its performance.

Proposed method

• Formulate group DRO with groups defined from known spurious correlations and derive that the worst-case risk equals the maximum group risk.
• Show that in overparameterized regimes, zero training loss for both ERM and group DRO yields poor worst-group test performance.
• Investigate regularization strategies (strong L2, early stopping) to prevent perfect training fit and reduce worst-group generalization gaps.
• Introduce group-adjusted DRO that adds a group-wise generalization-gap term C/√ng to prioritize smaller groups during training.
• Develop an online alternating gradient algorithm updating θ via SGD and the group distribution q via exponentiated gradient ascent, with convergence guarantees in the convex setting.

Experimental results

Research questions

• RQ1 Can group DRO improve worst-group generalization for overparameterized neural networks, and under what regularization conditions?
• RQ2 How do different regularization strategies (strong L2, early stopping) affect worst-group versus average performance in group DRO?
• RQ3 Does incorporating group-size based adjustments to account for group-specific generalization gaps further boost worst-group accuracy?
• RQ4 Is the proposed online training algorithm for group DRO stable and convergent, and what are its theoretical guarantees?
• RQ5 How does group DRO compare to importance weighting as a baseline for worst-case robustness under group shifts?

Key findings

• Overparameterized models with standard regularization exhibit near-perfect training loss but poor worst-group test performance, with worst-group accuracies like 60.0%, 41.1%, and 65.7% on Waterbirds, CelebA, and MultiNLI respectively.
• Strong regularization (large L2 penalty or early stopping) enables group DRO to achieve substantially higher worst-group accuracies while maintaining high average accuracies (e.g., up to 84.6% on Waterbirds and 86.7% on CelebA under strong regularization).
• Group DRO with regularization substantially improves worst-case performance by 10–40 percentage points across tasks (natural language inference and two image tasks).
• Incorporating group-adjusted terms that account for group-specific generalization gaps (proportional to 1/√ng) further improves worst-group test accuracy in some settings (e.g., Waterbirds improved by 5.9 percentage points).
• An online optimization algorithm interleaving SGD on θ with exponentiated gradient updates on a group weight distribution q yields convergence guarantees in the convex setting and scales to large models/datasets.