Out-of-Distribution Generalization via Risk Extrapolation (REx)

By David Krueger, Ethan Caballero, Jörn-Henrik Jacobsen, Amy Zhang, Jonathan Binas, Rémi Le Priol, Aaron C. Courville. 2020

1,228 citations recorded by the source corpus.

Generalizing outside of the training distribution is an open challenge for current machine learning systems. A weak form of out-of-distribution (OoD) generalization is the ability to successfully interpolate between multiple observed distributions. One way to achieve this is through robust optimization, which seeks to minimize the worst-case risk over convex combinations of the training distributions. However, a much stronger form of OoD generalization is the ability of models to extrapolate beyond the distributions observed during training. In pursuit of strong OoD generalization, we introduce the principle of Risk Extrapolation (REx). REx can be viewed as encouraging robustness over affine combinations of training risks, by encouraging strict equality between training risks. We show conceptually how this principle enables extrapolation, and demonstrate the effectiveness and scalability of instantiations of REx on various OoD generalization tasks. Our code can be found at this https URL.

Subfields: Artificial Intelligence

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

Citations in the corpus

Learning explanations that are hard to vary (2020, background): “This motivates our notion of invariant mechanisms, and inspired related notions which have been proposed for robust regression (Rojas-Carulla et al., 2018; Heinze-Deml et al., 2018; Arjovsky et al., 2019; Hermann and Lampinen, 2020; Ahuja et al., 2020; Krueger et al., 2020).”
SWAD: Domain Generalization by Seeking Flat Minima (2021, methodology): “12], obtained well-transferable model parameters by the meta-learning framework [13–16], data augmentation [17–19], or capturing causal relation [20, 21].”
In Search of Lost Domain Generalization (2020, background): “Krueger et al. [2020] propose an approximation to the IRM problem consisting in reducing the variance of error averages across domains.”
Inductive biases for deep learning of higher-level cognition (2020, neutral): “, 2016), spurious correlations (Krueger et al., 2021; Beery et al., 2018; Arjovsky et al., 2019) etc.”
Interpretable and Generalizable Graph Learning via Stochastic Attention Mechanism (2022, methodology): “Although G∗S determines Y , in the training dataset the data G and Y may have some spurious correlation caused by the environment (Pearl et al., 2016; Arjovsky et al., 2019; Chang et al., 2020; Krueger et al., 2021).”
Characterizing Generalization under Out-Of-Distribution Shifts in Deep Metric Learning (2021, background): “However, in practice, achieving effective OOD generalization is more challenging than in-distribution [28, 12, 21, 49, 31, 55].”
Environment Inference for Invariant Learning (2020, methodology): “Model selection Tuning hyperparameters when train and test distributions differ is a difﬁcult open problem (Krueger et al., 2021; Gulrajani & Lopez-Paz, 2021).”
Domain Generalization using Causal Matching (2020, background): “Even 5) from (Krueger et al., 2020) where ERM can minimize IRM penalty on Colored MNIST.”
Last Layer Re-Training is Sufficient for Robustness to Spurious Correlations (2022, methodology): “Related methods have been developed in several areas of machine learning, such as ML Fairness [22, 30, 47, 77, 2, 46], domain adaptation [23, 24] and domain generalization [9, 65, 57, 29, 84] including works on Invariant Risk Minimization and causality [75, 4, 53, 5].”
Neural Unsupervised Domain Adaptation in NLP—A Survey (2020, background): “This is the broader need for out-of-distribution generalization (Bengio, 2019), and a more challenging setup targeted at handling unknown domains (Volpi et al., 2018; Krueger et al., 2020).”
Generalizing to Unseen Domains: A Survey on Domain Generalization (2021, background): “researchers reduced the variance of training domain risks by risk extrapolation (VRex) [106] or reducing class-conditioned Wasserstein DRO [157].”
Causally motivated shortcut removal using auxiliary labels (2021, methodology): “(6) Rex : a recently suggested method that imposes an invariance penalty without re-weighting [Krueger et al., 2021].”
Gradient Starvation: A Learning Proclivity in Neural Networks (2020, methodology): “In several studies [9, 59, 79, 3], methods are developed to aggregate information from multiple training environments in a way that favors the generalizable / domain-agnostic / invariant solution.”
The Risks of Invariant Risk Minimization (2020, methodology): “Krueger et al. (2020) suggest penalizing the variance of the risks, while Xie et al. (2020) give the same objective but taking the square root of the variance.”
Understanding Continual Learning Settings with Data Distribution Drift Analysis (2021, background): “For example, the ﬁeld of out-of-distribution (OOD) generalization aims at developing learning algorithms able to perform well on test distributions that are different from training distribution (Arjovsky et al., 2019; Ahuja et al., 2020; Krueger et al., 2021).”