PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:1911.08731v2cs.LGstat.ML

TL;DR

Overparameterized neural networks can perform well on average yet fail on atypical groups, and naive group DRO does not resolve this worst-group generalization problem. The paper studies strongly regularized group DRO and finds 10–40 percentage point worst-group accuracy improvements across three applications while preserving high average accuracy.

  • Problem

    Overparameterized neural networks can achieve high average accuracy yet fail on atypical groups, while naive group DRO offers no advantage when training loss vanishes.

  • Method

    The paper trains group DRO models with stronger L2 regularization or early stopping and introduces a stochastic optimizer with convex-case convergence guarantees.

  • Results

    10–40 percentage points: regularized group DRO improves worst-group test accuracy across natural language inference and two image applications while maintaining high average accuracy.

  • Takeaways & Limitations

    Strong regularization appears important for worst-group generalization in overparameterized neural networks, even when it is unnecessary for average generalization.

  • Takeaways & Limitations

    The interaction between early stopping and group adjustments is difficult to study rigorously and is left for future work.

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper studies group DRO for overparameterized neural networks, motivated by models that perform well on average yet fail on atypical groups. It shows that strong regularization substantially improves worst-group generalization and introduces a scalable stochastic optimizer with convergence guarantees.

  • Motivation: Average accuracy can coexist with consistent failures on rare and atypical examples, motivating methods that address group-specific performance.The introduction frames this as a limitation of standard average-loss training aimed at i.i.d. test accuracy.
  • Method: Group DRO minimizes worst-case loss over predefined training groups, using prior knowledge about spurious correlations to define those groups.The paper describes grouping contradictory sentences with no negation words in the NLI example as one use of such prior knowledge.
  • Problem setting: In overparameterized neural networks, zero training loss makes a model optimal for both average and worst-case training loss, so poor worst-group performance arises from generalization rather than training loss.The study evaluates this issue in MultiNLI, CelebA, and a modified CUB bird-recognition dataset.
  • Results: 10–40 percentage points: regularized group DRO improves worst-case test accuracies across natural language inference and two image applications while maintaining high average accuracies.The paper considers ℓ2 penalties, early stopping, and group adjustments that account for differences in group generalization gaps.
  • Optimization: The paper introduces a stable stochastic group-DRO optimizer that scales to large models and datasets, with convex-case convergence guarantees and favorable empirical behavior in non-convex models.The optimizer is introduced to carry out the experiments and is evaluated in the paper’s neural-network settings.

2 SETUP

The setup contrasts ERM, which minimizes expected loss under the training distribution, with group DRO, which minimizes worst-case loss over mixtures of predefined groups. Group DRO assumes training group labels are known but highlights that low worst-group training loss need not ensure low worst-group test loss when generalization gaps are large.

  • Standard and robust objectives: ERM minimizes expected loss under the training distribution using the empirical training distribution.The model family, loss, and training distribution define the standard prediction objective.
  • Standard and robust objectives: DRO minimizes worst-case expected loss over an uncertainty set Q representing possible test distributions.Broad uncertainty sets can improve robustness to distribution shifts but may produce overly pessimistic models.
  • Group DRO: Group DRO defines Q as mixtures of m training groups, making worst-case risk equivalent to the maximum expected loss across groups.The group mixture weights lie in the (m −1)-dimensional probability simplex.
  • Group DRO: Group DRO uses known training group labels but not test-time group labels, and its worst-group test loss can exceed training loss by the generalization gap δ.For overparameterized neural networks, the gap is large unless sufficient regularization is applied.
  • Applications: The applications form m = |A| × |Y| groups from spuriously correlated attributes and labels across Waterbirds, CelebA, and MultiNLI.MultiNLI uses m = 6 groups and has 206175 training examples, including 1521 in its smallest group.

3 COMPARISON BETWEEN GROUP DRO AND ERM

At convergence, ERM and group DRO achieve near-perfect training accuracy and high average test accuracy but generalize poorly on worst-case groups. Stronger L2 regularization or early stopping prevents perfect fitting, reduces groupwise generalization gaps, and enables DRO to achieve substantially higher worst-group test accuracy than ERM.

  • ERM: ERM reaches at least 99.9% worst-group training accuracy and average test accuracies of 97.3%, 94.8%, and 82.5%, yet worst-group test accuracies are 60.0%, 41.1%, and 65.7%.These values correspond respectively to Waterbirds, CelebA, and MultiNLI.
  • DRO: Group DRO performs similarly to ERM, with near-perfect training accuracy and high average test accuracy but poor worst-group test accuracy.Because ERM nearly perfectly classifies every training point, it is near-optimal for both the ERM and group DRO objectives.
  • Discussion: Default regularization permits vanishing training loss, so the average-to-worst-group test gap reflects poor worst-group generalization rather than poor worst-group training performance.Both ERM and DRO generalize well on average while failing on the worst-case group.
  • Regularization: Increasing λ from the standard 0.0001 to 1.0 for Waterbirds and 0.1 for CelebA prevents perfect training accuracy and substantially reduces each group’s generalization gap.The stronger penalties still allow both ERM and DRO to achieve high average test accuracies.
  • Regularization: 84.6% and 86.7% worst-group test accuracy are achieved by DRO on Waterbirds and CelebA, improving over ERM’s 21.3% and 37.8% under strong L2 penalties.DRO’s corresponding worst-group training accuracies are 97.5% and 93.4%, and the small generalization gap transfers this performance to test data.
  • Early Stopping: Early stopping improves DRO over ERM from 6.7% to 86.0% on Waterbirds and from 25.0% to 88% on the reported worst-group test accuracies.Like strong L2 penalties, early stopping reduces the generalization gap and prevents perfect fitting.

4 COMPARISON BETWEEN DRO AND IMPORTANCE WEIGHTING

Strongly regularized group DRO outperforms importance weighting for worst-group robustness, both empirically and theoretically in non-convex settings. Inverse-frequency upweighting improves over ERM on some image tasks but fails on MultiNLI, while convex equivalence does not extend generally to neural-network objectives.

  • Empirical comparison: Importance weighting improves substantially over ERM on CelebA and Waterbirds but is slightly outperformed by DRO.The comparison uses inverse training-frequency weights and models selected across matched ℓ2 penalties and early stopping criteria.
  • Empirical comparison: On MultiNLI, upweighting achieves lower average and worst-group accuracies than ERM because the rare group is overemphasized, producing extremely low training accuracy for that group.This failure illustrates that inverse-frequency weighting need not equalize group performance in practice.
  • Theoretical comparison: In convex settings, importance weighting and DRO can learn equivalent models under some weights, but this equivalence need not hold for non-convex models.Weighted loss corresponds to expected loss under a single target distribution, whereas DRO seeks good performance across a family of target distributions.
  • Theoretical comparison: No weighting choice is guaranteed to produce a robust non-convex minimizer, and finding suitable weights may require solving a dual DRO problem.Common heuristics such as inverse group size therefore may fail and are not necessarily easier to implement than DRO.

5 ALGORITHM

The paper introduces an online group DRO optimization algorithm with convergence guarantees, addressing scalability and theoretical limitations of prior approaches. It alternates model updates with adaptive group reweighting to emphasize high-loss groups.

  • Algorithm: The proposed online optimization algorithm trains group DRO models efficiently while providing convergence guarantees.Prior work used batch methods that do not scale to large datasets or stochastic methods without convergence guarantees.
  • Algorithm: The method interleaves stochastic gradient descent on θ with exponentiated gradient ascent on q, a distribution over groups.Groups with higher loss receive greater mass, and examples are weighted proportionally to their group’s mass.
  • Algorithm: In practice, the algorithm uses minibatches and a momentum term for θ.These implementation details are specified as practical modifications to the alternating update procedure.
  • Convergence analysis: In the convex setting, the average iterate achieves a standard convergence-rate guarantee under boundedness, convexity, and Lipschitz assumptions.The guarantee is stated in expectation over the algorithm’s randomness.
  • Optimization challenge: The dual group DRO formulation is difficult to optimize stochastically because its maximum over expected losses obstructs unbiased stochastic gradient estimates.This motivates the need for a specialized stochastic optimization approach.

6 RELATED WORK

Related work addresses non-uniform accuracy through domain adaptation and fairness, while distributionally robust optimization commonly uses divergence-based uncertainty sets. Prior research also studies average-loss generalization extensively, but robust-loss generalization remains comparatively limited.

  • The problem of non-uniform accuracy: Approaches to non-uniform accuracy include domain adaptation for known target distributions, machine-learning fairness, and importance weighting.The cited work spans domain adaptation, fairness, and importance weighting as strategies for addressing uneven accuracy across the data distribution.
  • Distributionally robust optimization: Prior distributionally robust optimization typically defines Q as a divergence ball around the training distribution over (x, y).The passage lists divergence-based uncertainty sets as the standard formulation in prior DRO work.
  • Distributionally robust optimization: With small divergence-ball radii of O(1/n), distributionally robust optimization acts as a regularizer.The regularization interpretation is specifically associated with small uncertainty-set radii.
  • Generalization of robust models: Generalization of neural networks has been extensively studied for average loss, whereas analysis of robust losses remains limited.Related work also examines overfitting on rare labels under label shifts and proposes mitigation algorithms.

7 DISCUSSION

The paper finds that strong regularization makes group DRO substantially improve worst-group accuracy at a small average-accuracy cost, while preventing pre-specified spurious correlations. It also motivates deeper study of worst-case generalization and methods beyond strong ℓ2 penalties or early stopping.

  • 7 DISCUSSION: Strong regularization lets group DRO substantially improve worst-group accuracy while incurring only a small cost in average accuracy.The paper identifies regularization as important for worst-case group generalization in overparameterized neural networks.
  • 7 DISCUSSION: Group DRO can prevent models from learning pre-specified spurious correlations and may retain high worst-group accuracy when groups are imperfectly specified.The latter result comes from supplemental experiments; shifts beyond pre-specified groups remain future work.
  • 7 DISCUSSION: The findings call for deeper analysis of average versus worst-case generalization and tools beyond strong ℓ2 penalties or early stopping.Such analysis could clarify deep-network failure modes while maintaining high average accuracy.

A PROOFS · A.1 EQUIVALENCE OF DRO AND IMPORTANCE WEIGHTING IN THE CONVEX SETTING

Under convexity, compactness, and continuity assumptions, every optimal DRO model corresponds to minimizing expected loss under some worst-case distribution. The proof establishes this through attainment and max-min equality, yielding a saddle point between the model and distribution.

  • A.1 EQUIVALENCE OF DRO AND IMPORTANCE WEIGHTING IN THE CONVEX SETTING: Proposition 1 states that an optimal worst-group model θ∗ minimizes expected loss under some distribution Q∗ in the uncertainty set.This requires continuous convex loss, convex compact model and distribution sets, and θ∗ minimizing the worst-group objective.
  • A.1 EQUIVALENCE OF DRO AND IMPORTANCE WEIGHTING IN THE CONVEX SETTING: The expected-loss function h(θ, Q) is continuous, convex in θ, and linear—and therefore concave—in Q.These properties follow from continuity and convexity of the loss in θ.
  • A.1 EQUIVALENCE OF DRO AND IMPORTANCE WEIGHTING IN THE CONVEX SETTING: The worst-case objective remains convex and lower semi-continuous because these properties are preserved under pointwise suprema.The resulting function is also proper under the stated assumptions.
  • A.1 EQUIVALENCE OF DRO AND IMPORTANCE WEIGHTING IN THE CONVEX SETTING: Compactness of Θ and Q, together with the preceding regularity conditions, ensures that the DRO objective attains its optimal value.The proof invokes Weierstrass’ theorem for this attainment result.
  • A.1 EQUIVALENCE OF DRO AND IMPORTANCE WEIGHTING IN THE CONVEX SETTING: A similar argument shows that the sup-inf objective attains its optimum at some Q∗ in Q.Thus, an optimizing distribution exists within the uncertainty set.
  • A.1 EQUIVALENCE OF DRO AND IMPORTANCE WEIGHTING IN THE CONVEX SETTING: Continuity of h and compactness of Θ and Q establish the max-min equality needed to connect the two optimization orders.The proof cites a standard max-min result for this step.
  • A.1 EQUIVALENCE OF DRO AND IMPORTANCE WEIGHTING IN THE CONVEX SETTING: The optimal pair (θ∗, Q∗) forms a saddle point of the DRO objective.This saddle-point characterization follows by combining objective attainment with max-min equality.
  • A.1 EQUIVALENCE OF DRO AND IMPORTANCE WEIGHTING IN THE CONVEX SETTING: Consequently, the optimal DRO model θ∗ also minimizes the weighted risk h(θ, Q∗).This is the proof’s importance-weighting equivalence: DRO optimization is represented by expected loss under Q∗.

A.2 CONVERGENCE RATE OF ALGORITHM 1

Proposition 2 establishes an expected-error convergence rate for the average iterate of Algorithm 1 under bounded, convex, Lipschitz loss and bounded parameters. The proof derives this result by casting the method as online mirror descent for a convex-concave saddle-point problem and applying a standard pseudo-regret guarantee.

  • Convergence guarantee: Proposition 2 gives an expected-error convergence guarantee for Algorithm 1’s average iterate under nonnegative, convex, B∇-Lipschitz, Bℓ-bounded loss and ∥θ∥2 ≤ BΘ.The expectation is over the algorithm’s randomness.
  • Proof strategy: The proof applies online mirror descent regret bounds to the corresponding saddle-point optimization problem.The argument relies on the theorem of Nemirovski et al. (2009).
  • Proof strategy: The algorithm is represented using a uniform mixture of group distributions and group-indicator losses, enabling verification of the required convexity and unbiased stochastic-subgradient assumptions.For each group g′, the construction uses Fg′(θ; (x, y, g)) := mI[g = g′]ℓ(θ; (x, y)).
  • Proof conclusion: Substituting the problem-specific constants into the regret bound yields Theorem 2 because the minimax game is convex-concave.The proof first computes the constants required by the online mirror descent guarantee.

B SUPPLEMENTARY EXPERIMENTS

Supplementary CelebA experiments show that Group DRO remains robust when the spurious attribute is misspecified or distractor attributes are included. Across 64 attribute-label groups, the model achieves 78.9% robust accuracy, near 86.7% for the original DRO model and above 37.8% for ERM.

  • Inexact group specifications: The experiments replace the ground-truth Male attribute with Wearing Lipstick and add four distractor attributes: Eyeglasses, Smiling, Double Chin, and Oval Face.The model optimizes worst-case performance over all 64 groups formed by five attributes and one label.
  • Robust accuracy: 78.9% robust accuracy is achieved across the four original CelebA groups under inexact group specifications.This is compared with 86.7% for original DRO using the ground-truth spurious attribute and 37.8% for ERM.

C EXPERIMENTAL DETAILS · C.1 DATASETS · C.2 MODELS

The experiments use deliberately constructed group-balanced evaluation sets for MultiNLI and Waterbirds, alongside CelebA, to assess worst-group performance under spurious correlations. They evaluate pretrained BERT and ResNet50 models under specified optimization, regularization, and early-stopping settings.

  • C.1 DATASETS: MultiNLI combines the training and development sets to create validation and test sets because leaderboard evaluation cannot report accuracy by group.Examples without consensus gold labels are removed, and negation is defined by the hypothesis containing nobody, no, never, or nothing.
  • C.1 DATASETS: Waterbirds labels CUB birds as waterbirds or landbirds and composites them with water or land backgrounds from Places using segmentation masks.The training set places 95% of each bird category in its corresponding background and 5% in the opposite background.
  • C.1 DATASETS: Waterbirds validation and test sets balance landbirds and waterbirds across land and water backgrounds to measure rare-group performance more accurately.This differs from the skewed training distribution and supports stable comparisons among ERM, reweighting, and group DRO.
  • C.1 DATASETS: CelebA uses its official train-validation-test split, with Blond Hair as the target and Male as the spuriously associated variable.Average test accuracy is computed across groups and weighted by their relative proportions in the skewed training data.
  • C.2 MODELS: ResNet50 experiments use pretrained torchvision weights, stochastic gradient descent with momentum 0.9, batch size 128, batch normalization, no dropout, and no data augmentation.A fixed learning rate is used for comparability, with adjustments across L2 penalty strengths and datasets because of batch normalization interactions.
  • C.2 MODELS: Standard ResNet50 training uses λ = 0.0001, with 50 CelebA epochs and 300 Waterbirds epochs.Early stopping trains for 1 epoch, while strong L2 uses λ = 1.0 for Waterbirds and λ = 0.1 for CelebA, both at learning rate 0.00001.
  • C.2 MODELS: Group-adjustment experiments use the same strong-L2 settings and search C over {0, 1, 2, 3, 4, 5}, while the benchmark grid searches dataset-specific λ values.The benchmark restricts group-adjustment searches to the L2 penalties used in the comparison.
  • C.2 MODELS: BERT uses pretrained bert-base-uncased with the default tokenizer, AdamW, dropout, λ = 0, batch size 32, and a fixed linearly decaying learning rate starting at 0.00002.Standard training lasts 20 epochs; early stopping and the benchmark use 3 epochs, while tested L2 penalties produced similar or worse robust accuracy than λ = 0.
Loading 1911.08731v2…