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

David Krueger; Ethan Caballero; Joern-Henrik Jacobsen; Amy Zhang; Jonathan Binas; Dinghuai Zhang; Remi Le Priol; Aaron Courville

Source-linked AI summary

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

David Krueger, Ethan Caballero, Joern-Henrik Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Remi Le Priol, Aaron Courville

arXiv:2003.00688v5Mar 2, 2020cs.LGcs.AIcs.NEstat.ML

TL;DR

Distributional shift limits deployment because models can be sensitive to changes beyond the training domains. The paper proposes Risk Extrapolation, a robust-optimization framework that equalizes training risks to extrapolate domain variation, and reports improved performance over IRM when covariate and causally induced shifts co-occur.

Problem
Distributional shift makes prediction systems sensitive to test domains that may differ from training domains, including shifts involving both causal and anti-causal input elements.
Method
Risk Extrapolation uses robust optimization over affine combinations of training risks, with equality-promoting risk penalties as a simpler variant.
Results
REx uncovers invariant input–target relationships, and experiments show it significantly outperforms IRM in settings combining covariate shift with invariant prediction.
Takeaways & Limitations
REx can provide out-of-distribution robustness when inputs contain both causes and effects of targets while also handling covariate shift more robustly than IRM in the reported scenarios.
Takeaways & Limitations
REx performs poorly on heteroskedastic structural-equation tasks because it does not distinguish underfitting from irreducible noise.

Abstract

from arXiv · show

Distributional shift is one of the major obstacles when transferring machine learning prediction systems from the lab to the real world. To tackle this problem, we assume that variation across training domains is representative of the variation we might encounter at test time, but also that shifts at test time may be more extreme in magnitude. In particular, we show that reducing differences in risk across training domains can reduce a model's sensitivity to a wide range of extreme distributional shifts, including the challenging setting where the input contains both causal and anti-causal elements. We motivate this approach, Risk Extrapolation (REx), as a form of robust optimization over a perturbation set of extrapolated domains (MM-REx), and propose a penalty on the variance of training risks (V-REx) as a simpler variant. We prove that variants of REx can recover the causal mechanisms of the targets, while also providing some robustness to changes in the input distribution ("covariate shift"). By appropriately trading-off robustness to causally induced distributional shifts and covariate shift, REx is able to outperform alternative methods such as Invariant Risk Minimization in situations where these types of shift co-occur.

1. Introduction

REx addresses out-of-distribution generalization by extrapolating variation observed across training domains, aiming to uncover invariant input–target relationships and improve robustness to larger shifts.

Motivation: Distributional shift can make models rely on spurious features, such as background scenery rather than the target animal.This sensitivity is a major obstacle to applying prediction systems outside their training distribution.
Problem setting: OOD generalization evaluates performance on a new test domain without labeled or unlabeled training data, requiring assumptions about possible test domains.The paper formulates this challenge as worst-case optimization over a perturbation set of possible test domains.
Risk Extrapolation: Unlike DRO, which considers convex combinations of training risks, MM-REx uses affine combinations to represent a larger set of extrapolated domains.Under specific circumstances, MM-REx can be viewed as DRO over extrapolated domains.
Risk Extrapolation: REx seeks invariant relationships between inputs X and targets Y, such as the animal–label relationship rather than the background–label relationship.The paper distinguishes invariant prediction from learning an invariant representation.
Comparison and contributions: REx differs from IRM by addressing covariate shift and significantly outperforming IRM when covariate shift co-occurs with invariant-prediction requirements.The reported settings include modified CMNIST and simulated robotics tasks from the DeepMind Control Suite.
Comparison and contributions: The paper contributes MM-REx, evidence that REx can solve tasks where IRM fails under covariate shift, and a proof that equal risks can suffice for discovering causal structure.The framework is motivated as robust optimization over extrapolated domains.

2. Background & Related work

The paper frames OOD generalization as robust optimization over plausible shifts and contrasts REx with invariant-prediction and representation approaches. REx uses training-risk variation to target robustness beyond interpolation, while its theory connects equalized risks to invariant causal prediction.

Robust optimization: Domain generalization seeks good performance on unseen domains by optimizing worst-case performance over an assumed perturbation set.The perturbation set encodes assumptions about possible test distributions.
Robust optimization: DRO and Risk Interpolation restrict the perturbation set to convex mixtures of training distributions, whereas MM-REx extends it to affine combinations of training risks.Affine combinations permit extrapolation beyond the convex hull of observed training risks.
Invariant prediction: REx promotes invariant prediction by enforcing equality of risks across training domains without explicitly learning an equipredictive representation.An equipredictive representation makes the relationship between representation and target fixed across domains.
Invariant prediction: Unlike invariant-representation methods, invariant predictors can preserve target-relevant relationships when P(Y) changes across domains.The paper notes that matching representations or conditional representations can fail when label distributions vary.
REx relative to prior work: REx is model agnostic, scales to deep networks, and requires risk invariance rather than residual invariance while still supporting invariant causal prediction.The paper distinguishes its goal of OOD generalization from causal inference, although causal mechanisms can be recovered under suitable conditions.

3. Risk Extrapolation

Risk Extrapolation balances lowering average training risk with making domain risks similar, using extrapolated-risk robustness or a variance penalty. The paper shows that this principle can improve OOD robustness and recover causal mechanisms under explicit assumptions.

REx principle: REx pursues two potentially competing goals: reducing training risks and increasing their similarity across domains.Achieving risk equality may require increasing risk on the best-performing domains.
REx principle: On Colored MNIST, enforcing equal risks discourages reliance on non-invariant color and promotes shape-based prediction for unseen domains.The training-domain color-label correlation differs across domains, making color-based predictors domain-sensitive.
Covariate shift: MM-REx can target covariate shifts through extrapolated distributions, although extrapolation may produce invalid negative pointwise probabilities.The paper describes this as robustness to changes in P(X).
Algorithms: MM-REx performs robust optimization over bounded affine combinations of training risks, with extrapolation controlled by λmin.λmin = 0 recovers risk interpolation, while increasingly negative coefficients permit more extreme shifts and approach strict risk equality.
Algorithms: V-REx regularizes the variance of training risks and was found simpler, more stable, and more effective than MM-REx.Its parameter β trades off average-risk reduction against equality of risks; β = 0 recovers ERM.
Causal discovery: Under assumptions including observed causes, interventions on X, and homoskedasticity, equalizing risks can recover the causal mechanism of Y.The theory states that a predictor satisfying REx over the intervention domains uses fY as its predictive distribution.
Causal discovery: For linear structural equation models, REx can identify the model using a number of domains linear in the dimensionality of X.The result assumes three distinct do() interventions for each non-target variable.
Causal discovery: The causal theorem is an insight result whose perturbation set is uncountably infinite, and ERM generally depends on the distribution over domains.The paper gives Colored MNIST as an example where ERM can learn color when those training domains dominate the data.

4. Experiments

Experiments compare REx with IRM and ERM across Colored MNIST, structural-equation, domain-generalization, and control tasks. REx is strongest when covariate shift co-occurs with invariant prediction, while IRM is stronger under heterogeneous inherent noise.

Overview: REx outperforms IRM across tasks involving covariate shift and invariant prediction, including modified Colored MNIST and simulated robotics.The broader evaluation also includes DomainBed benchmarks.
Structural Equation Models: REx performs poorly when domain-dependent risk differences reflect irreducible noise rather than underfitting, giving IRM an advantage in heteroskedastic SEM settings.The authors identify this as a weakness of REx’s sensitivity to covariate shift.
Colored MNIST: V-REx significantly outperforms IRM on every tested Colored MNIST covariate-shift variant.The variants induce class, digit, and color imbalance, with some settings combining covariate and interventional shift.
Colored MNIST: For small imbalance values, IRM never exceeds chance in the hyperparameter search, whereas REx exceeds chance in 4.4%/23.7%/2.0% of class/digit/color trials.The reported values correspond to p = 0.1/0.1/0.2, respectively.
Domain Generalization: REx, IRM, and ERM perform comparably on the evaluated domain-generalization benchmarks.The DomainBed evaluation used training-domain validation for model selection and was limited to four inexpensive datasets.
Reinforcement Learning: REx outperforms both IRM and ERM in modified DeepMind Control Suite tasks combining spurious observations, covariate shift, and heteroskedasticity.The result suggests that robustness to covariate shift outweighs REx’s heteroskedasticity challenges in this setting.

5. Conclusion

The conclusion presents REx as a robust-optimization method for OOD generalization when inputs contain both causes and effects of targets. It emphasizes REx’s ability to combine causal identification with greater robustness to covariate shift than IRM.

Conclusion: REx provides robustness and OOD generalization when X contains both causes and effects of Y.The conclusion frames this as the challenging setting addressed by the work.
Conclusion: REx can perform causal identification like IRM while behaving more robustly under covariate shift.The authors connect this advantage to settings involving both covariate and interventional shift.

Appendices

The appendices define extrapolation geometrically, explain why high-dimensional and domain generalization require it, and illustrate how REx allocates capacity under covariate shift. They also include causal-model material and supplementary experiments.

Appendix contents: The appendices include toy examples, causal-model summaries, REx-versus-V-REx analysis, further experiment details, unreported experiments, and related OOD topics.These materials are listed as appendix sections B–I.
B. Definition and discussion of extrapolation in machine learning: Extrapolation predicts outside the convex hull of training examples, whereas interpolation predicts within it.The appendix notes that high intrinsic dimension can make extrapolation necessary even in i.i.d. settings.
B. Definition and discussion of extrapolation in machine learning: High-dimensional domain diversity can require extrapolation across risk functions, making REx’s perturbation set broader than risk interpolation’s.The comparison concerns extrapolating to new domains rather than merely interpolating among observed training domains.
C.2. Covariate shift example: A discrete probability example shows that REx can ignore an unreliable conditional relationship while distributing capacity across inputs subject to covariate shift.The example contrasts REx with ERM and IRM on the X = 0, 1, 2 cases.

C.2. Covariate shift example

The covariate-shift example models risks over domains with varying input probabilities and shows that risk extrapolation selects the minimax solution. The appendix situates this example within the paper’s broader causal-model terminology.

C.2. Covariate shift example: REx is motivated as robust learning over the affine span of training distributions when limited capacity or data makes covariate shift consequential.The example distinguishes COSTLY inputs with probability q from CHEAP inputs with probability 1 − q.
C.2. Covariate shift example: The saddle-point solution is p = ω = u/u+v with risk uv/u+v for every q, making the risks equal across domains.The appendix identifies this as the minimax-optimal solution.
C.2. Covariate shift example: ERM and risk interpolation select boundary solutions for many domain probabilities, whereas REx finds the minimax optimum for arbitrary q1 and q2.Risk interpolation returns p = ω only when q1 and q2 straddle ω.
Causal-model background: The appendix distinguishes causal graphs, causal graphical models, and structural causal models by their treatment of independences, distributions, interventions, and counterfactuals.The comparison is presented as a summary following standard causal-model definitions.

E.1. Proofs of theorems 1 and 2

The proofs show that, under causal and homoskedasticity assumptions, satisfying REx across interventions identifies the causal predictive mechanism. Exact risk equality is essential, while practical use may trade it against minimizing risk.

REx is defined through two goals: reducing training risks and increasing their similarity across domains.The analysis prioritizes exact equality of risks, then selects a minimal-risk predictor among those satisfying equality.
Under the stated assumptions, three distinct do-interventions for every input variable force the linear predictor to recover the causal coefficients.The assumptions include observed causes of Y, intervention-based domains, and homoskedasticity.
For an arbitrary structural causal model and strictly proper scoring rule, REx over interventions on X uniquely recovers the causal predictive distribution fY.The causal mechanism equalizes risk across intervention domains, while any other predictor has higher risk in at least one domain.
MM-REx extrapolates beyond ordinary mixtures of training risks, but its perturbation set can contain quasidistributions with negative probabilities.The resulting optimization questions concern convexity and the effects of negative probabilities on optimization and learned solutions.

F. The relationship between MM-REx vs. V-REx, and the role each plays in our work

MM-REx provides a robust-optimization formulation and geometric view of Risk Extrapolation, while V-REx is the simpler practical penalty. They coincide in enforcing equal risks in the infinite-penalty limit but differ in optimization behavior and multi-domain weighting.

MM-REx expands DRO from convex combinations to affine combinations of training risks, enabling robust optimization over extrapolated domains.MM-REx offers geometric intuition through an expanded convex hull, including extrapolated domains that may be quasidistributions.
V-REx is the more practical algorithm because it is simple to implement, performed better in CMNIST experiments, and may optimize more stably through smoother gradients.The smoother gradient field is illustrated by vectors curving toward equal-risk directions near the diagonal.
Both MM-REx and V-REx recover perfect equality of training risks as λmin → −∞ or β → ∞, respectively.The convergence results assume all training risks are finite.
With two domains, MM-REx is equivalent to an MAE penalty and V-REx to an MSE penalty; with more domains, MM-REx and pairwise MAE differ in their risk weighting.For two domains, the MM-REx and MAE formulations match when βMAE = 1 − 2λmin.
V-REx is equivalent, up to a constant factor of 2, to the average pairwise mean squared error between training risks.

G.1. CMNIST with covariate shift

The CMNIST covariate-shift experiments evaluate REx across class, digit, and color imbalance using default and tuned hyperparameters. With appropriate tuning, REx handles greater covariate shift and can outperform IRM, though success falls at high p because degenerate results occur more often.

The reported success rate counts randomly sampled hyperparameter combinations with accuracy above 50%, exceeding random guessing.The success analysis uses a random search over hyperparameter combinations.
REx handles greater covariate shift than IRM and can outperform IRM when hyperparameters are appropriately tuned.The experiments use randomly searched hyperparameters and assess multiple CMNIST covariate-shift variants.
At high p, REx has a lower success rate because degenerate results with training accuracy below test accuracy occur more often.
The experiments cover class imbalance, digit imbalance, and color imbalance in CMNIST.These variants are arranged from left to right in the corresponding figures.
MM-REx has higher standard error than V-REx across all three CMNIST covariate-shift variants.The comparison uses the vertical error bars in the plots.

G.2. SEMs from “Invariant Risk Minimization”

The SEM experiments test REx when inputs contain causal and anti-causal variables under homoskedastic and heteroskedastic noise. REx performs well when anti-causal noise varies, but target-noise variation causes poor performance and can recover the anticausal model instead.

The SEM tasks predict Y from causal inputs X1 and anti-causal effects X2, using variants of the structural equation model from prior IRM work.The experiments retain the prior experimental settings except for the number of runs described in the source.
REx performs well when the noise level in anti-causal features varies across domains but poorly when target noise varies instead.Table 5 reports this contrast across the corresponding task groups.
When only anti-causal-feature noise varies, equal-risk constraints recover the causal model by setting the anti-causal coefficient β to zero.Loss minimization then recovers the causal coefficient α = w1→y in the homoskedastic setting.
When only target noise varies, REx enforces β = 1/wy→2 and induces α = 0, recovering the anticausal model.The differing irreducible errors make equal risks incompatible with the desired causal predictor in this setting.
REx assumes the mechanism for Y is fixed across domains, whereas IRM assumes only that E(Y | Pa(Y)) is domain-independent.Allowing interventions on Y under the weaker IRM-style assumption is identified as future work.

G.3. Reinforcement Learning Experiments

The experiments examine REx and IRM across reinforcement learning, domain generalization, and financial prediction settings, with emphasis on penalty scheduling and comparative robustness. Results indicate that performance depends strongly on scheduling and task characteristics.

Reinforcement learning: REx outperforms IRM and ERM on modified DeepMind Control Suite tasks containing spurious observations, while covariate shift and heteroskedasticity favor different methods.The observations combine task state with noise and scaled spurious features, creating both covariate shift and intrinsic difficulty variation.
Penalty scheduling: REx and IRM require careful penalty scheduling; the penalty should increase after approximately 100 epochs, and exponential decay performed worse than the waterfall schedule.The authors identify scheduling as critical and report that exponential decay remained above 50% accuracy but was significantly worse.
Penalty scheduling: In CMNIST, the penalty should be applied after predictive features are learned but before the model memorizes training examples.Before predictive features emerge, penalties can push the model toward a constant predictor; after memorization, penalty gradients become weak.
Domain generalization: On PACS, REx outperforms IRM and IRM outperforms ERM on average, although all three methods perform worse than Jigsaw.The evaluation trains on three domains and tests on a fourth, using settings selected on VLCS and transferred unchanged to PACS.
Financial prediction: Financial-indicator experiments find that IRM and REx perform similarly across different data splits and prediction tasks.The dataset treats each of five years as a separate domain and forms 20 prediction tasks.

H.3.1. EXPERIMENT DETAILS

The experiment details describe preprocessing and optimization for financial-domain prediction, while the surrounding material situates domain generalization among related distribution-shift methods.

Data preparation: The financial dataset retains 37 features shared across all five years and standardizes them by centering and dividing by their standard deviations.The data use version 1 of the dataset and remove variables unavailable in every year.
Optimization: Financial prediction uses a two-hidden-layer MLP trained with binary cross-entropy, Adam, dropout, and an L2 penalty.The hidden layers each contain 128 tanh units, with dropout after each layer.
Optimization: IRM and REx add their respective penalties with β = 1, scale the original loss by 10^-4 after 1000 iterations, and train for up to 9000 iterations with early stopping.These settings are specific to the reported IRM/REx financial experiments.
Related methods: Domain adaptation differs from domain generalization because it permits some access to the test distribution.A common domain-adaptation strategy is learning an invariant representation of the inputs.
Related methods: Out-of-distribution detection instead aims to recognize examples from a new domain, complementing the goal of generalizing predictions to new distributions.The passage also lists adversarial training, self-supervised learning, and data augmentation as techniques that can improve OOD generalization.