Improving Out-of-Distribution Robustness via Selective Augmentation (2201.00299v3)

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

• The paper introduces LISA, a selective augmentation strategy that improves OOD performance by mitigating spurious correlations.
• It employs intra-label and intra-domain mixup-inspired techniques to encourage the model to learn domain-agnostic features.
• Empirical results across nine benchmarks reveal significant gains in both average and worst-group performance compared to conventional methods.

Improving Out-of-Distribution Robustness via Selective Augmentation

The discussed paper presents a novel approach to address the common problem of distribution shifts—which manifest as subpopulation shifts and domain shifts—in machine learning models. The authors propose a method termed Learning Invariant Predictors with Selective Augmentation (LISA), which enhances the robustness of machine learning models to out-of-distribution (OOD) data by focusing on learning invariant predictors without imposing constraints on internal model representations. This approach is grounded in recognizing the limitations of existing methods that employ regularization strategies to enforce domain invariance.

Methodology

LISA leverages a selective data interpolation strategy inspired by the mixup technique. Mixup involves linearly interpolating features and labels of data samples; LISA extends this by conditioning the interpolation on specific attributes—either labels or domains—to tackle OOD issues strategically. The paper introduces two specific augmentation strategies:

1. Intra-label LISA (LISA-L): This involves interpolating samples that share the same label but originate from different domains. By blending such samples, the method aims to cancel out spurious domain-label correlations.
2. Intra-domain LISA (LISA-D): Here, interpolation is conducted between samples from the same domain but with different labels, thereby compelling the model to disentangle domain information from label prediction, pushing it towards learning more domain-agnostic representations.

This dual strategy aims to mitigate the risk of models relying on incidental correlations in the training data which do not hold across different domains or underrepresented subsets.

Empirical Evaluation

The empirical evaluation of LISA spans nine benchmark datasets characterized by variability in both domain shifts and subpopulation shifts. Results indicate that LISA consistently enhances performance over traditional methods across these benchmarks, achieving improvements in both average and worst-group performance metrics.

For instance, in benchmarks like CelebA, which involves significant domain-based biases (e.g., gender correlating with hair color in image classification), LISA outperforms traditional methods such as Invariant Risk Minimization (IRM) and Distributionally Robust Optimization (DRO) techniques by a notable margin. Similarly, in datasets from the WILDS benchmark targeting natural distribution shifts, LISA shows superior predictive consistency across previously unseen domains.

Theoretical Insights

The authors provide a theoretical analysis justifying LISA's efficacy. They draw formal connections between the proposed selective augmentation method and existing mixup techniques, extending understanding of how selective augmentation influences the classifier's boundary in feature space. They demonstrate that LISA results in a lower worst-group error compared to standard empirical risk minimization and vanilla mixup approaches, especially in settings with pronounced spurious correlations.

Implications and Future Work

Practically, LISA offers a low-complexity, computation-friendly enhancement to existing ML pipelines aimed at deployment in real-world settings where distributional assumptions are frequently violated. This has profound implications for fields like medical imaging or environmental mapping, where models trained on limited datasets often need to generalize well to diverse real-world scenarios.

Theoretically, the paper suggests new vistas for research into selective data augmentation, urging exploration of generalized frameworks beyond the binary classification setting or formats dependent on labels. Future research may evolve towards integrating LISA into multitask or transfer learning paradigms, potentially benefiting applications requiring cross-domain generalization capabilities despite marked disparities in training versus test data conditions.

In summary, through sophisticated yet intuitive selective augmentation, LISA provides a notable contribution to the domain of OOD generalization, revealing avenues for further research and application across varied machine learning tasks.