PaperlayerView exact revision on arXiv

Source-linked AI summary

Improving Out-of-Distribution Robustness via Selective Augmentation

Huaxiu Yao, Yu Wang, Sai Li, Linjun Zhang, Weixin Liang, James Zou, Chelsea Finn

arXiv:2201.00299v3cs.LG

TL;DR

Distribution shifts can cause severe OOD failures, especially under subpopulation and domain shifts driven by spurious correlations. The paper proposes LISA, which selectively applies mixup to same-label cross-domain or same-domain different-label pairs, and reports consistent gains across nine benchmarks, more invariant predictors, and smaller worst-group error in theory.

  • Problem

    Distribution shifts between training and test data create OOD generalization challenges, including failures from spurious correlations under subpopulation and domain shifts.

  • Method

    LISA learns invariant predictors through selective mixup, interpolating either same-label samples from different domains or different-label samples from the same domain.

  • Results

    Across nine benchmarks, LISA consistently outperforms prior methods and produces more domain-invariant predictors; theory shows smaller worst-group error than ERM and vanilla interpolation.

  • Takeaways & Limitations

    Selective interpolation targets domain-specific information and spurious correlations without explicitly constraining model representations or predictors.

  • Takeaways & Limitations

    LISA is difficult to apply when examples with the same label are hard to obtain, such as in object detection or generative modeling.

Abstract

from arXiv · show

Machine learning algorithms typically assume that training and test examples are drawn from the same distribution. However, distribution shift is a common problem in real-world applications and can cause models to perform dramatically worse at test time. In this paper, we specifically consider the problems of subpopulation shifts (e.g., imbalanced data) and domain shifts. While prior works often seek to explicitly regularize internal representations or predictors of the model to be domain invariant, we instead aim to learn invariant predictors without restricting the model's internal representations or predictors. This leads to a simple mixup-based technique which learns invariant predictors via selective augmentation called LISA. LISA selectively interpolates samples either with the same labels but different domains or with the same domain but different labels. Empirically, we study the effectiveness of LISA on nine benchmarks ranging from subpopulation shifts to domain shifts, and we find that LISA consistently outperforms other state-of-the-art methods and leads to more invariant predictors. We further analyze a linear setting and theoretically show how LISA leads to a smaller worst-group error.

1. Introduction

The paper addresses OOD generalization under subpopulation and domain shifts, where spurious domain-label correlations can degrade performance. LISA learns invariant predictors through selective interpolation rather than explicit regularization.

  • Distribution shift occurs when test data differ from training data, substantially degrading model performance.
  • Subpopulation shifts involve seen but underrepresented test domains, while domain shifts involve changes across domains.
  • Prior invariant-learning regularizers can restrict model expressiveness or create difficult optimization problems, and designing broadly suitable regularizers is challenging.
  • LISA selectively interpolates same-label samples from different domains or same-domain samples with different labels.
  • LISA consistently outperforms seven prior methods across nine benchmarks and produces more domain-invariant predictors.
  • Theoretical analysis shows that LISA can mitigate spurious correlations and achieve smaller worst-domain error than ERM and vanilla mixup.

2. Preliminaries

The paper formalizes distribution shifts through training and test distributions over domains, distinguishing changed domain proportions from disjoint test domains. Evaluation emphasizes worst-group, worst-domain, and average test performance.

  • Traditional generalization assumes equal training and test distributions, whereas this paper studies the setting P_tr ≠ P_ts.
  • The overall distribution is modeled as a mixture of domains, each containing feature, label, and domain-associated samples.
  • Training and test domains are defined by domains with positive mixture probabilities in their respective distributions.
  • Subpopulation shifts retain training domains but change their subpopulation proportions, with D_ts ⊆ D_tr.
  • For group-based spurious correlations, each group combines a domain and label, and performance is measured by worst test-group 0-1 error.
  • For domain shifts, test domains are disjoint from training domains but may share common properties, and evaluation considers worst-domain or average performance.

3. Learning Invariant Predictors with Selective Augmentation

LISA improves robustness by selectively applying mixup to pairs chosen by label and domain relationships. Its two strategies target spurious correlations through cross-domain same-label or within-domain different-label interpolation.

  • LISA encourages invariant predictors through selective data interpolation to alleviate domain-related spurious correlations.
  • Mixup forms interpolated features and labels as x_mix = λx_i + (1 − λ)x_j and y_mix = λy_i + (1 − λ)y_j.The interpolation ratio λ is sampled from Beta(α, β), and classification labels are one-hot vectors.
  • Vanilla mixup can improve in-distribution generalization but may fail to cancel spurious correlations and therefore provide weaker OOD generalization.
  • Intra-label LISA: Intra-label LISA interpolates same-label samples from different domains to eliminate domain-label spurious correlations.
  • Intra-domain LISA: Intra-domain LISA interpolates different-label samples within the same domain so predictions vary with the interpolation ratio rather than relying on domain information.
  • LISA randomly selects intra-label or intra-domain augmentation with probabilities p_sel and 1 − p_sel, respectively.The selection probability is treated as a hyperparameter determined by cross-validation.
  • Intra-label LISA helps more with more domains or weaker spurious correlations, whereas intra-domain LISA benefits settings with highly spurious domain-label correlations.

4. Experiments

The experiments evaluate LISA across subpopulation and domain shifts, comparing it with prior methods, substitute interpolation strategies, and measures of predictor invariance. LISA consistently improves robustness, with gains attributed to selective interpolation and stronger domain invariance.

  • Evaluation setup: The evaluation covers four subpopulation-shift datasets and five domain-shift datasets, using worst-group accuracy for subpopulation shifts.The subpopulation benchmarks are CMNIST, Waterbirds, CelebA, and CivilComments; domain-shift benchmarks include Camelyon17, FMoW, RxRx1, Amazon, and MetaShift.
  • Subpopulation shifts: LISA consistently outperforms invariant-learning baselines across subpopulation-shift datasets, achieving the best performance on CMNIST, CelebA, and CivilComments.On Waterbirds, LISA is slightly worse than GroupDRO but has comparable performance.
  • Subpopulation shifts: Both intra-label and intra-domain LISA contribute to performance, with intra-domain LISA performing slightly better on CMNIST, Waterbirds, and CelebA.The experiments use both strategies with psel = 0.5 on these three datasets.
  • Domain shifts: For natural domain shifts, experiments use only intra-label LISA, while same-label interpolation disregarding domain information performs best on Camelyon17, FMoW, and RxRx1.The authors associate this choice with weaker label-domain spurious correlations in those datasets.
  • Domain shifts: LISA consistently outperforms prior methods on all five domain-shift datasets across image and text data and different model architectures.For Camelyon17, results use 10 seeds; other domain-shift datasets use 3 seeds.
  • Ablations: Selective interpolation improves out-of-distribution robustness beyond vanilla and in-group mixup, including when upweighting enhances those substitute strategies.The comparisons support canceling spurious correlations and encouraging invariant predictors rather than providing only data augmentation.
  • Invariance and shift severity: LISA learns more domain-invariant predictors and representations, and its performance advantage increases as the distance between training and test domains grows.The degree-of-shift analysis varies MetaShift training-test distance using node similarity on a meta-graph.

5. Theoretical Analysis

The theoretical analysis models subpopulation and domain shifts with Gaussian mixtures and compares LISA against ERM and vanilla mixup using worst-group misclassification error. Under strong domain-label spurious correlation and stated assumptions, LISA achieves smaller asymptotic errors and mitigates spurious correlations.

  • Theoretical setup: The analysis uses a Gaussian mixture model with subpopulation and domain shifts to compare LISA, ERM, and vanilla mixup.The model includes conditional means, covariances, imbalanced groups, and assumed cross-domain invariance.
  • Theoretical setup: The worst-group error evaluates a linear classifier’s maximum misclassification error across domains and classes.For method A, the classifier has slope b_A and intercept b_0,A, with ERM, vanilla mixup, and LISA defined as comparison methods.
  • Subpopulation shifts: Smaller ξ indicates stronger spurious correlation between domains and labels, and the theorem compares LISA’s error with ERM and vanilla mixup under this regime.ξ measures alignment between the domain-specific difference and the marginal difference with respect to Σ.
  • Subpopulation shifts: When ξ is small and p = o(αn), LISA has asymptotically smaller worst-group classification errors than ERM and vanilla mixup.The analysis also finds that LISA yields a classification rule closer to the invariant classification rule.
  • Domain shifts: Under domain shifts, strong domain-label spurious correlation downgrades ERM and vanilla mixup, whereas selective interpolation enables LISA to mitigate this issue.The domain-shift comparison considers a new unseen domain and is formalized in Theorem 2.

6. Related Work and Discussion

Related work addresses distribution shifts through invariant representations, invariant predictors, or group-robust methods. LISA differs by using selective interpolation to target spurious correlations without relying on the same regularization strategies.

  • Learning Invariant Representations: Invariant-representation methods align feature distributions across domains using divergences or adversarial losses.This category is motivated by unsupervised domain adaptation and seeks to eliminate domain dependency in representations.
  • Learning Invariant Predictors: Invariant-predictor methods seek predictors that maintain label-relevant correlations across domains, including IRM and its variants.These approaches go beyond domain alignment by directly enhancing the relationship between invariant representations and labels.
  • Group Robustness: Group-robust methods address spurious correlations through worst-group optimization, minority-group sampling, reweighting, or regularization.Some approaches target subpopulation shifts without annotated group labels.

7. Conclusion

The paper proposes LISA as a simple, efficient approach for improving out-of-distribution robustness under subpopulation and domain shifts. Experiments and theory support its effectiveness, while the authors identify challenges for settings lacking same-label examples.

  • Conclusion: LISA improves out-of-distribution robustness by eliminating domain-related spurious correlations through selective interpolation.The paper evaluates this approach under both subpopulation shifts and domain shifts.
  • Conclusion: Experiments on nine datasets demonstrate LISA’s promise, while analyses associate its gains with more invariant predictors and representations.Theoretical results report smaller worst-group misclassification error than ERM and vanilla data interpolation.
  • Limitations: LISA is limited when obtaining examples with the same label is difficult, including object detection and generative modeling.The authors propose exploring more general selective augmentation strategies in future work.

A.1.1. DATASET DETAILS

The experiments use image and text benchmarks representing spurious associations and imbalanced subpopulations. Training applies selective mixup within specified label-domain groupings using pretrained vision or language models.

  • Dataset details: Colored MNIST classifies two digit groups while color acts as a spurious attribute with opposing training correlations across classes.Validation uses balanced color proportions, whereas the test set changes the green-red proportion.
  • Dataset details: Waterbirds classifies waterbirds versus landbirds, with bird-background combinations creating minority groups.The training set contains 4,795 samples, including 56 waterbirds on land and 184 landbirds on water.
  • Dataset details: CelebA predicts blond versus not blond hair color, with gender serving as the spurious domain attribute.The minority groups are blond male and not-blond female examples.
  • Dataset details: CivilComments predicts toxic versus non-toxic comments while treating demographic features as spurious domain identifications.The dataset contains 450,000 comments from online articles.
  • Implementation: Experiments use pretrained ResNet-50 for image datasets and BERT for CivilComments, applying LISA within same-label/different-domain or same-domain/different-label batches.The full subpopulation-shift results are reported in Tables 10 and 11.

A.2.1. DATASET DETAILS

The experiments cover subpopulation and domain-shift benchmarks, use specified pretrained backbones and interpolation settings, and evaluate spurious correlations and predictor invariance. LISA’s gains are strongest when spurious correlations are present and are accompanied by more invariant predictors.

  • Dataset details: The experiments evaluate LISA across nine diverse benchmarks spanning subpopulation and domain shifts.The supplied dataset passages describe Camelyon17, FMoW, RxRx1, Amazon, and MetaShift, alongside subpopulation-shift evaluations.
  • Dataset details: Camelyon17, FMoW, and RxRx1 use domain definitions based on hospitals, time and geography, and experimental batches, respectively.Their out-of-distribution splits use later hospitals, later time periods, or held-out experimental batches.
  • Spurious correlation: Cramér’s V measures domain-label spurious correlation, with higher values indicating stronger correlation.The analysis applies this measure across the experimental datasets.
  • Spurious correlation: LISA performs similarly to ERM without spurious correlations, while outperforming ERM when spurious correlations exist.Vanilla mixup is strongest without spurious correlations, whereas LISA is strongest when those correlations are present.
  • Predictor invariance: Risk Variance and Gradient Norm provide predictor-invariance metrics, and their results confirm that LISA improves predictor invariance.The table caption states that smaller values indicate stronger invariance.

A.5.2. ANALYSIS OF LEARNED INVARIANT REPRESENTATIONS

The representation analysis measures label-conditioned domain invariance with KL divergence and supplements it with Waterbirds visualizations. LISA learns representations that cluster same-label samples more closely across domains than competing methods.

  • Representation-level invariance: IRkl measures representation-level invariance using pairwise KL divergence between label-conditioned representation distributions across domains.Smaller IRkl values indicate more invariant representations with respect to labels.
  • Representation-level invariance: LISA learns stronger representation-level invariance than ERM and regularization-based methods including IRM, IB-IRM, and V-REx.The analysis reports results on CMNIST, Waterbirds, Camelyon17, and MetaShift.
  • Waterbirds visualization: On Waterbirds, LISA places same-label samples closer together regardless of domain information in hidden-representation visualizations.The visualization also includes learned decision boundaries for comparison across methods.

A.6. Full Results of WILDS data

Additional WILDS results show LISA’s robustness across validation and test settings, including reduced Camelyon17 validation-test gaps and stronger lower-tail Amazon accuracy despite lower average accuracy.

  • Full WILDS results: LISA outperforms other baseline approaches across the reported WILDS scenarios.The additional tables include validation performance and other metrics for Camelyon17, FMoW, RxRx1, and Amazon.
  • Camelyon17: Approximately 10% is the ERM validation-test performance gap on Camelyon17, which LISA significantly reduces.The passage attributes the gap to test data being more visually distinctive than validation data.
  • Amazon: LISA achieves the best Amazon 10th-percentile accuracy despite performing worse than ERM in average accuracy.The 10th percentile is presented as a metric for consistency across users.
  • Proof strategy: The theoretical proof analyzes population OLS errors, empirical convergence, and separate bounds for intra-label and intra-domain LISA.Together, these steps establish empirical misclassification-error bounds for the selective strategies.
  • Theoretical comparison: Under the stated linear-model conditions, LISA-L and LISA-D have errors no larger than ERM.The supplied results state ELL < EERM for LISA-L and ELD ≤ EERM for LISA-D under sufficient conditions.

B.5. Finite sample analysis

The finite-sample analysis decomposes empirical estimation and misclassification error, then derives method-specific concentration and comparison bounds. It analyzes vanilla mixup, LISA-L, and LISA-D under linear-model assumptions.

  • Error analysis: The analysis bounds misclassification error through the learned slope and estimates the deviation between empirical and population quantities.Larger bL(ˆb) corresponds to smaller misclassification error, while the deviation includes estimation and n^-1/2 terms.
  • Selective augmentation: The finite-sample proof separately bounds covariance, mean, and sample-covariance terms for LISA-L and LISA-D.These bounds use sub-Gaussian or sub-exponential concentration for the relevant empirical quantities.
Loading 2201.00299v3…