Out-of-Distribution Generalization via Risk Extrapolation (REx)
Authors: David Krueger
Ethan Caballero
Joern-Henrik Jacobsen
Amy Zhang
Jonathan Binas
Rémi LE PRIOL
Dinghuai Zhang
Aaron Courville
Venue:  (2021)
Keywords: out of distribution
domain generalization
invariant risk minimization
robust optimization
invariant causal prediction
spurious features
generalization

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

Under review as a conference paper at ICLR 2021
OUT-OF-DISTRIBUTION GENERALIZATION VIA RISK EXTRAPOLATION
Anonymous authors Paper under double-blind review
ABSTRACT
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 magnitude. In particular, domains can reduce a model's 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 penalty on the variance of training risks (V-REx) as 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.
show that reducing differences in risk across training
INTRODUCTION
While neural networks often exhibit super-human generalization on the training distribution, they can be extremely sensitive to distributional shift, presenting a major roadblock for their practical application (Su al., [ |Engstrom et al. 2017 Recht et al.) 2019 Hendryeks Dietterich] 2019). This sensitivity is often caused by relying on "spurious" features unrelated to the core concept we are trying to learn (Geirhos et al. 2018). For instance, Beery et al. (2018) give the example of an image recognition model failing to correctly classify cows on the beach, since it has learned to make predictions based on the features of the background (e.g. a grassy field) instead of just the animal.
In this work, we consider out-of-distribution (OOD) generalization, also known as domain gener- alization, where a model must generalize appropriately to a new test domain for which it has neither labeled nor unlabeled training data. Following common practice (Ben-Tal et al. 2009). formulate this as optimizing the worst-case performance over a perturbation set of possible test domains, F:
ROOD
(0) max (!) eE
(I)
Since generalizing to arbitrary test domains is impossible, the choice of perturbation set encodes our assumptions about which test domains might be encountered. Instead of making such assumptions priori, we assume access to data from multiple training domains, which can inform our choice of perturbation set. A classic approach for this setting group distributionally robust optimization (DRO) (Sagawa et al., 2019). where F contains all mixtures of the training distributions. This is mathematically equivalent to considering convex combinations of the training risks.
However, we aim for more ambitious form of OOD generalization, over a larger perturbation set. Our method minimax Risk Extrapolation (MM-REx) is an extension of DRO where F instead contains affine combinations of training risks, see Figure []] Under specific circumstances, MM-REx can be thought of as DRO over a set of extrapolated domains]] allowing us to carry over machinery
We define "extrapolation" mean "ousidee the convex hull", see Appendix B for more.

---

Under review as a conference paper at ICLR 2021
R
R
RRES
Rgi
convex hull of training
extrapolation region
distributions
6
6
(2
c1
(2
e2
pd (X))
P))
Left: Robust optimization optimizes worst-case performance over the convex hull of training distributions. Right: By extrapolating risks, REx encourages robustness to larger shifts. Here e1, e2, and represent training distributions, and P* P2 represent some particular directions of variation in the affine space of quasiprobability distributions over (X,Y),
Figure 1:
developed for DRO, such as the O(I/VT) convergence rate proven by Sagawa et al. (211)). But MM-REx also unlocks fundamental new generalization capabilities unavailable to DRO.
In particular, focusing on supervised learning, we show that Risk Extrapolation can uncover invariant relationships between inputs X and targets Y. Intuitively, an invariant relationship is statistical relationship which maintained across all domains in F.. Returning the now-on-the-bead example, the relationship between the animal and the label is expected to be invariant, while the relationship between the background and the label is not. A model which bases its predictions on such an invariant relationship is said to perform invariant prediction
Many domain generalization methods assume P( is an invariant relationship. limiting distri- butional shift to changes in P(X), which are known covariate shift (David et al. 2010). This assumption can easily be violated, however. For instance, when causes X, a more sensible assump- tion is that P(X|Y) is fixed, with PYY) varying across domains (Schölkopf et al. 2012 Lipton et al. 2018). In general, invariant prediction may involve an aspect of causal discovery. Depending on the perturbation set. however, other, more predictive, invariant relationships may also exist (Koyama & Yamaguchi 2020).
The first method for invariant prediction be compatible with modern deep learning problems and techniques is Invariant Risk Minimization (IRM) (Arjovsky et al. 2019). making it a natural point of comparison. Our work focuses on explaining how REx addresses OOD generalization, and highlighting differences (especially advantages) compared with IRM and other domain generalization methods, see Table Broadly speaking, REx optimizes for robustness to the forms distributional shift that have been observed to have the largest impact on performance in training domains. This can prove a significant advantage over the more focused (but also limited) robustness that IRM targets. For instance, unlike IRM, REx can also encourage robustness to covariate shift (see Section 3)
And indeed, our experiments show that REx significantly outperforms IRM settings that in- volve covariate shift and require invariant prediction, including modified versions of CMNIST and simulated robotics tasks from the Deepmind control suite. On the other hand, because REx does not distinguish between underfitting and inherent noise, IRM has advantage in settings where some domains are intrinsically harder than others. We perform several other sets of ex- periments in order to better understand and compare REx and IRM. Our contributions include:
1) MM-REx, a novel domain generalization problem formulation suitable for invariant prediction.
2) Demonstrating that REx solves invariant prediction tasks where IRM fails due covariate shift.
3) Proving that equality of risks can be a sufficient criteria for discovering causal structure.
Note this is different from learning an invariant representation Ganin et al. 2016 see Section 2.2
2

---

Under review as a conference paper at ICLR 2021
Method
Invariant Prediction
Cov. Shift Robustness
Suitable for Deep Learning
DRO
(C-)ADA
X
ICP
X
IRM
/
X
REx
Table 1: A comparison of approaches for OOD generalization.
2 BACKGROUND & RELATED WORK
We consider multi-source domain generalization. where our goal is to find parameters 0 that perform well on unseen domains, given a set of m training domains, & = (ei.... Cn sometimes also called environments. We assume the loss function, E is fixed, and domains only differ in terms of their data distribution Po(X,Y) and dataset De. The risk function for a given domain/distribttion
is:
(0) (xx.)))
We refer to members of the set RR.e E E} as the training risks or simply risks. Changes in P. (x Y) can be categorized as either changes in P(X) (covariate shift), changes in P(Y]X) (concept shift), or a combination. The standard approach to learning problems Empirical Risk Minimization (ERM), which minimizes the average loss across all the training examples from all the domains:
Del RERM(() (( {(fo(z), 9) = Σ BL ((fe()),))
(3)
2.1 ROBUST OPTIMIZATION
An approach more taylored to OOD generalization is robust optimization (Ben- Tal et al. 2009), which aims optimize model's worst-case performance over some perturbation data distributions, JF (see Eqn T. When only single training domain is available, it is common to assume that XX) is fixed and let be all distributions within some f f-divergence ball of the training P(X) (Hu et al. 2016 Bagnell 2005). ) As another example, adversarial robustness can be seen as instead using a Wasserstein ball as perturbation set (Sinha et al.) 2017).
possible
In multi-source domain generalization, test distributions are often assumed to be mixtures (i.e. convex combinations) of the training distributions; this is equivalent to setting F
RuI (0) - Eev=
R.
(e) eEE
(4)
We call this objective Risk Interpolation (RI), (group) Domain Robust Optimization (DRO), following Sagawa et al. (2019). While single-source methods classically assume that the probability vary independently (Hu et al. 2016). DRO yields much lower dimensional perturbation set, with at most one direction variation per domain, regardless of the dimensionality of X and Y. training domains, whereas single-source methods typically assume only P(X) can change. Minimax REx is an extension of this approach affine combinations of training risks.
of each data-point
also provide robustness to any form of shift in P(X, Y) which oceurs across
2.2 INVARIANCE AND CAUSALITY
An invariant predictor, 4, is function of X with the property that P.(Y equal Ve E. In other words, the relationship between such a and Y is invariant to the choice of domain. Invariant relationships between X and Y are those that can be written as P4 (Y = x) = P(Y))(x)) with
This assumption often called "covariate shift", but we assume covariate/concept shift co-occur.

---

Under review as a conference paper at ICLR 2021
Algorithm 1 Variance Risk Extrapolation (V-REx)
Require: Di...., Dm: training sets from different domains
Require: 8: learning rate and penalty term hyperparameters
1: randomly initialize model parameters 0
2: while not done do
3: for all D, do
4: Estimate risk (0) using minibatch of K examples from D;
5: end for
6: Update 0 - + aVo (E, R,(0) + aV (R,(0())))))
7: end while
an invariant predictor (although REx does not use such an explicit decomposition), Koyama Yamaguchi (2020) prove that maximal invariant predictor, that is, one that maximizes mutual information with the targets, &* argmax, (1)) Y). solves the robust optimization problem (Eqn. ! under fairly general assumptions; when d- is unique, we call the features it ignores spurious.
The result of Koyama & Yamaguchi 2020) provides theoretical reason for favoring invariant prediction over the common approach of learning invariant representations (Pan et al., 2010).
which make (() or Pe (c equal Ve E E. Popular methods here include Adversarial domain adaptation (ADA) (Ganin al. 2016) and conditional ADA (C-ADA) (Long et al. (2018). Unlike invariant predictors, invariant representations can easily fail to generalize OOD: ADA forces the predictor to have the same marginal predictions P(Y). which is a mistake when P(Y) in fact changes
across domains (Zhao al. 2019); C-ADA suffers from more subtle issues (Arjovsky et al. 2019).
Causal relationships are a paradigmatic example of invariant relationships. In a Structural Causal Model (SCM), the mechanism for variable describes how it's value computed based on the values of its causes, which can be thought of as its parents in a Pytorch (Paszke et al., 2019) computational graph | Works that take a causal approach to domain generalization often assume that the mechanism for Y is fixed, while X may be subject to different interventions in different domains (Bühlmann) 2018ah. We call resulting changes in P(X, Y) interventional shift. Interventional shift can involve both covariate shift and/or concept shift. In their seminal work on Invariant Causal Prediction (ICP), Peters its nonlinear extension (Heinze- Deml et al., 2018) of a linear model are equal across domains. Our work differs from ICP in that:
Peters et al. (2016) leverage this invariance to learn which elements of X cause Y. ICP and statistical tests to detect whether the residuals
1. Our method is model agnostic and scales to deep networks.
2. Our goal is OOD generalization, not causal inference. These are not identical: invariant prediction can sometimes make use of non-causal relationships, but when deciding which interventions to perform, truly causal model is called for.
3. Our learning principle only requires invariance of risks, not residuals. Nonetheless, we prove that this can ensure invariant causal prediction.
more similar method to REx is Invariant Risk Minimization (IRM) (Arjovsky et al. 2019). which shares properties (1) and (2) of the list above. Like REx, IRM also uses weaker form of invariance than ICP; namely, they insist that the optimal linear classifier must match across domains Still, REx differs significantly from IRM. While IRM specifically aims for invariant prediction, REx seeks robustness to whatever forms of distributional shift are present. Thus, REx is more directly focused on the problem of OOD generalization, and can provide robustness to wider variety of distributional shifts, inluding covariate shift. Also, unlike REx, IRM seeks to match across domains, not the full P(Y (X)). This, combined with IRM's 's indifference to covariate shift, make it more effective in cases where different domains or examples are inherently more noisy.
(XX)
2.3 FAIRNESS
Equalizing risk across different groups (e.g. male vs. female) has been proposed as a definition of fairness (Donini et al. 2018), generalizing the equal opportunity definition of fairness (Hardt et al.)
*See Appendix for more in-depth technical overview of SCMs. In practice, IRMv1 replaces this hilevel optimization problem with gradient penalty classifier weights.
4

---

Under review as a conference paper at ICLR 2021
2016). Williamson & Menon (2019) propose using the variance or absolute difference of risks to measure deviation from this notion of fairness; these correspond to our V-REx and (in the case of only two domains) MM-REx. However, in the context of fairness, equalizing the risk of training groups is the goal. Our work goes beyond this by showing that it can serve as a method for OOD) generalization.
RISK EXTRAPOLATION
Before discussing algorithms for REx and theoretical results, we first expand on our high-level explanations of what REx does, what kind of OOD generalization it promotes, and how, The principle of Risk Extrapolation (REx) has two aims:
1. Reducing training risks
2. Increasing similarity of training risks
In general, these goals can be at odds with each other; decreasing the risk in the domain with the lowest risk also decreases the overall similarity of training risks. Thus methods for REx typically seeks to increase risk on the best performing domains. From geometric point of view, encouraging equality of risks flattens the "risk plane" (the affine span of the training risks, considered as function of the data distribution, see Figure []. While this can result in higher training risks, it also means that the risk changes less if the distributional shifts between training domains are magnified at test time.
1-10000 (=12033
1.70
0.4
0.3
Training accuracies (left) and risks (right) colored MNIST domains with varying P(Y = Olcolor = red) after 500 epochs. Dots represent training risks, lines represent test risks. Increasing the V-REx penalty (i.e. B) leads to a flatter "risk plane" and more consistent performance across domains, as the model learns to ignore color in favor of shape- based invariant prediction.
Figure 2:
Figure 2 illustrates how flattening the risk plane promote OOD generalization real data, using the Colored MNIST (CMNIST) task as an example (Arjovsky al. 2019). In the CMNIST training domains, the color of digit is more predictive of the label than the shape is. But because the correlation between color and label is not invariant, predictors that use the color feature achieve different risk on different domains. By enforcing equality of risks, REx prevents the model from using the color feature enabling successful generalization to the test domain where the correlation between color and label is reversed.
Probabilities vs. Risks. Figure ddepicts how the extrapolated risks considered in MM-REx can be translated into a corresponding change in P(X, Y). Training distributions can be thought as points in an affine space with dimension for every possible value of (X, Y); see Appendix C. for an example. Because the risk is linear w.r.t. P(z. y). a convex combination of risks from different domains is equivalent to the risk a domain given by the mixture of their distributions. The same holds for the affine combinations used in MM-REx, with the caveat that the negative coefficients may lead to negative probabilities, making the resulting P(X, Y) quasiprobability distribution, i.e. signed measure with intregral 1. We explore the thcoretical implications of this in Section
Covariate Shift. When only P(X) differs across domains, as in Figure 3 then (xx) x is an invariant predictor, and thus learning invariant predictor is not expected to improve OOD generalization (compared with ERM). Instead, what is needed robustness to covariate shift, which REx, but not IRM, can provide Robustness covariate shift can improve OOD generalization by
Arjovsky et al. 2019 recognize this limitation of [RM in what they call the "realizable" case.
5

---

Under review as a conference paper at ICLR 2021
Pr
PI-3
PIK
interpolation extrapolation
Figure 3: Extrapolation can yield distribution with negative P(x) for some æ. Left: P(x:) for domains 6j and eg. Right: Point-wise interpolatinn/ettappllttio of ped (x) and (x).
ensuring that low-capacity models spend sufficient capacity on low-density regions of the input space, as we demonstrate in Appendix C.2 But even for high capacity models, P(X) can have significant influence on what is learned; for instance Sagawa et al. (2019) show that DRO can significantly improves the performance on rare groups in their with a model that achieves 100% training accuracy in their Waterbirds dataset. Pursuing robustness to covariate shift also comes with drawbacks for REx, however: REx does not distinguish between underfitting and inherent noise in the data, and so can force the model make equally bad predictions everywhere, even some examples are less noisy than others.
3.1
METHODS OF RISK EXTRAPOLATION
We now formally describe the Minimax REx (MM-REx) and Variance-REx (V-REx) techniques for risk extrapolation. Minimax-RE, performs robust learning over perturbation of affine combinations of training risks with bounded coefficients:
X
RMM-REE
max
E.. (1 mdmiu) max (e) Amin Σ (O),
(5)
Ne>e=
Ae 2.nin
where is the number of domains, and the hyperparameter Amin controls how much extrapolate. For negative values of Amin. MM-REx places negative weights on the risk all but the worst-case domain, and Amin -o, this criterion enforces strict equality between training risks; Amin 0 recovers risk interpolation (RI). Thus, like RI, MM-REx aims to be robust the direction of variations in P(X, Y) between test domains. However, negative coefficients allow to extrapolate more extreme variations. Geometrically, larger values of Amin expand the perturbation set farther away from the convex hull of the training risks, encouraging flatter "risk-plane" (see Figure 2}
While MM-REx makes the relationship to RI/RO clear, we found using the variance of risks as a regularizer (V-REx) simpler, stabler, and more effective:
RVERES (0) = 3 Var (ar([Ri (R), (0), Rm (0))) Σ (o)
Here E [0, oo) controls the balance between reducing average risk and enforcing equality of risks, with 3 = 0 recovering ERM, and 3 - de leading V-REx to focus entirely making the risks equal.
See Appendix for the relationship between V-REx and MM-REx and their gradient vector fields.
EXPERIMENTS
We evaluate REx and compare with IRM on range of tasks requiring OOD generalization. REx provides generalization benefits and outperforms IRM on wide range of tasks, including: i) variants of the Colored MNIST (CMNIST) dataset (Arjovsky et al.) 2019) with covariate shift, ii) continuous control tasks with partial observability and spurious features, iii) domain generalization tasks from the BemainBed suite (Gulrajani Lopez-Paz) 2020). On the other hand, when the inherent noise in Y varies across environments, IRM succeeds and REx performs poorly.
4.1 COLORED MNIST
Arjovsky et al. 2019) construct a binary classification problem (with 0-4 and 5-9 each collapsed into single class) based the MNIST dataset, using color spurious feature. Specifically, digits

---

Under review as a conference paper at ICLR 2021
0.5
0.6
0.4 TottA00
0.5
0.4
0.4
Random Guessing
0.3 Ran dom Guessing
0.3
Random Sumwang
V-NEx
V-ItEx
V4KEX
IRNL
IRMV1
J8Mv1
03
0.0
(0122344))
Pishapabo 21224462
PJC1)Gemen)
Figure 4: REx outperforms IRM on Colored MNIST variants that include covariate shift. The X-axis indexes increasing amount of shift between training distributions, with 0 corresponding to disjoint supports. Left: class imbalance, Center: shape imbalance, Right: color imbalance.
are either colored red or green, and there is strong correlation between color and label, which is reversed test time. The goal is learn the causal "digit shape" feature and ignore the anti-causal "digit color" feature. The learner has access to three domains:
1. A training domain where green digits have 80% chance of belonging to class (digits 5-9).
2. A training domain where green digits have 90% chance of belonging to class 1.
3. A test domain where green digits have a 10% chance of belonging to class 1.
We use the exact same hyperparameters as Arjovsky et al. (2019). only replacing the IRMv1 penalty with MM-REx or V-REx penalty These methods achieve similar performance, see Table
Method
train acc
test ace
V-REx (ours)
71.5 + 1.0
68.7
IRM
70.8 1 0.9
2.5
MM-REx (ours)
72.4 1.8
66.1 1.5
CMNIST with covariate shift To test our hypothesis that REx should outperform IRM under covariate shift, we construct 3 variants of the CMNIST dataset. Each variant represents a different way inducing covariate shift to ensure differences across methods are consis- tent. These experiments combine covariate shift with interventional shift, since P(Green|X Y = 1) still differs across training domains as the original CMNIST.
RI
88.9 1 0.3
22.3 1 4.6
ERM
87.4 0.2
17.1 0.6
Grayscale oracle
73.5 +: 0.2
73.0 + 0.4
Optimum
75
75
Chance
50
50
Table 2:
Accuracy (percent) on Colored MNIST. REx and IRM learn to ignore the spurious color feature. Strikethrough results achieved via tuning on the test set.
1. Class imbalance: varying P(shape(x) E 2,,,,,,) as in Wu et al. 2020)).
2. Digit imbalance: varying p = P(shape(x) E (1,2) U (6, 7)); digits 0 and 5 are removed.
3. Color imbalance: We use 2 versions of each color, for total channels: R1, R2, G1, G2- We vary p = P(R Redd PrGilGreen).
While (1) also induces change in P.Y), (2) and (3) induce only covariate shift in the causal shape and anti-causal color features (respectively). We compare across several levels of imbalance, p E j0, 0.5), using the same hyperparameters from Arjovsky et al. (2019), and plot the mean and standard error over trials.
V-REx significantly outperforms IRM in every case, see Figure Al In order to verify that these results not due to bad hyperparameters for IRM, we perform random search that samples 340 unique hyperparameter combinations for each value of and compare the the number of times each method achieves better than chance-level (50% accuracy). Again, V-REx outperforms IRM; in particular, for small values of IRM never achieves better than random chance performance, while REx does better than random in 4, 4.4% 22..23..%..9/ of trials, respectively, scenarios for 0.1/0.1 /..2. This indicates that REx can achieve good OOD generalization in settings involving both covariate and interventional shift, whereas IRM struggles to.
the class/digit/color imbalance
4.2 ToY STRUCTURAL EQUATION MODELS (SEMs)
REx's sensitivity covariate shift can also be a weakness when reallocating capacity towards domains with higher risk does not help the model reduce their risk, e.g. due to irreducible noise. We illustrate
When there are only 2 domains, MM-REx is equivalent to penalty on the Mean Absolute Error (MAF), Appendix F2.2
7

---

Under review as a conference paper at ICLR 2021
Algorithm
ColoredMNNSS
VLCS
PACS
OfficeHome
ERM
52.0 + 0.1
77.4 + 0.3
85.7 = 0.5
67.5 = 0.5
IRM
51.8 + 0.1
78.1 + 0.0
84.4 + 1.1
66.6 1.0
V-REx
52.1 + 0.1
77.9 + 0.5 85.8 0.6
66.7 0.5
Table 3: REx, IRM, and ERM all perform comparably on set of domain generalization benchmarks.
this using the linsar-Gaussian structural equation model (SEM) tasks introduced by Arjovsky et al. (2019). Like CMNIST, these SEMs include spurious features by construction. They also introduce 1) heteroskedasticity, 2) hidden confounders, and/or 3) elements of X that contain a mixture of causes and effects of Y. These three properties highlight advantages of IRM over ICP (Peters et al. 2016), as demonstrated empirically by Arjovsky et al. (2019). REx is also able to handle (2) and (3), but it performs poorly in the heteroskedastic tasks. See Appendix G.2 for details and Table for results.
4.3 DOMAIN CENNEALLIATION IN THE DOMAINBED SUITE
Methodologically, it is inappropriate to assume access to the test environment in domain generalization settings, as the goal to find methods which generalize to unknown test distributions. Gulrajani & Loope7PPaz (2020) introduced the DomainBed evaluation suite to rigorously compare existing approaches to domain generalization, and found that no method reliably outperformed ERM. We evaluate V-REx DomainBed using the most commonly used triiiig-domaai validation method for model selection. Due to limited computational resources, we limited ourselves to the 4 cheapest datasets. Results of baseline are taken from Gulrajani & Lopez-Paz (2020). who compare with more methods. Results in Table give the average over 3 different train/valid splits.
4.4 REINFORCEMENT LEARNING WITH PARTIAL OBSERVABILITY AND SPURIOUS FEATURES
Finally, we turn to reinforcement learning, where covariate shift (potentially favoring REx) and heteroskedasticity (favoring IRM) both occur naturally as a result of randomness in the environment and policy. In order to show the benefits of invariant prediction, we
modify tasks from the Deepmind Control Suite (Tassa et al. (2018) to include spurious features in the observation, and train Soft Actor-Critic (Haarnoja et al.] (20T8) agent. REx outperforms both IRM and ERM, suggesting that REx's robustness to covariate shift outweighs the challenges We average over 10 runs finger_sph s_in and wa walke lerr wa k,, using hyperparameters tuned on cartpole_swi setee_sssiiee sww.pusissiguuppppppppppppp ngup (to avoid overfitting). See Appendix for details and further results.
faces with heteroskedasticity in this setting, see Figure [
370
800
CRM
250
100
320
230
300
260
100
100
1/5
Figure 5: Performance and standard error on walker_ walk (left), finger_spiin (right).
C CONCLUSION
We have demonstrated that REx, method for robust optimization, can provide robustness and hence out-of-distributio generalization in the challenging case where X contains both causes and effects of Y. In particular, like IRM, REx can perform causal identification, but REx can also perform more robustly in the presence of covariate shift. Covariate shift is known to be problematic when models misspecified. when training data is limited, or does not cover areas of the test distribution. As such situations are inevitable in practice, REx's ' ability to outperform IRM in scenarios involving a combination of covariate shift and interventional shift makes it powerful approach.
8

---

Under review as a conference paper at ICLR 2021
REFERENCES
Isabela Albuquerque, Nikhil Naik, Junnan Li, Nitish Keskar, and Richard Socher. Improving out-of-distribution generalization via multi-task self-supervised pretraining. 2020.
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint 1iXi:1907.0222 2019.
Philip Bachman, R Devon Hjelm, and William Buchwalter. Learning representations by maximizing mutual information across views, 2019.
J. Andrew Bagnell, Robust supervised learning, In Proceedings of the 20th National Conference on Artificial Intelligence Volume 2, AAAI055, pp. 714-719. AAAI Press, 2005. ISBN 157735236x.
Sara Beery, Grant Van Horn, and Pictro Perona. Recognition in terra incognita. Lecture Notes in Computer Science, pp. 472-489, 2018. ISSN 1611-3349.
Shai Ben-Benniiiie John Blitzer, Koby Crammer, Alex Kulesza, Fernando Pereira, and Jennifer Wortman Vaughan, A theory of learning from different domains. Machine learning, 19(1-2):551-175, 2010.
Aharon Ben-Tal, Laurent El Ghaoui, and Arkadi Nemirovski. Robust optimization, volume 28.
Princeton University Press, 2009.
Peter Bühlmann. Invariance, causality and robustness, 2018a.
Peter Bühlmann. Invariance, causality and robustness, 2018b.
Fabio Carlucci, Antonio D'Innocente, Silvia Bucci, Barbara Caputo, and Tatiana Tommasi. Domain generalization by solving jigsaw puzzles. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 2229-2238, 2019.
Ekin D. Cubuk, Barret Zoph, Dandelion Mane, Vijay Vasudevan, and Quoc V. Le. Autoaugment: Learning augmentation policies from data, 2018.
Shai Ben David, Tyler Lu, Teresa Luu, and Dávid Pál. Impossibility theorems for domain adaptation. In Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics, pp. 129-136, 2010.
Guillaume Desjardins, Karen Simonyan, Razvan Pascanu, et al. Natural neural networks. In Advances in Neural Information Processing Systems, pp. 2071-2079, 2015.
Michele Donini, Luca Oneto, Shai Ben-David, John Shawe-Taylor, and Massimiliano Pontil. Empiri- cal risk minimization under fairness constraints, 2018.
Logan Engstrom, Brandon Tran, Dimitris Tsipras, Ludwig Schmidt, and Aleksander Madry, Exploring the landscape of spatial robustness. arXiv preprint arXiv: 2017.
Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky. Domain-adversarial training of neural networks. The Journal of Machine Learning Research, 17(1):2096-0300 2016.
Robert Geirhos, Patricia Rubisch, Claudio Michaelis. Matthias Bethge, Felix A Wichmann, and Wieland Brendel. Imagenet-trained cnns are biased towards texture; increasing shape bias improves accuracy and robustness. arXiv preprint arXiv: 181. 181..1231,, 2018.
Ian J Goodfellow, Jonathon Shlens, and Christian Szegedy. Explaining and harnessing adversarial examples. arXiv preprint arXiv: 14fi:::41..65111 2014.
Sven Gowal, Chongli Qin, Po-Sen Huang, Taylan Cemgil, Krishnamurthy Dvijotham, Timothy Mann, and Pushmeet Kohli. Achieving robustness in the wild via adversarial mixing with disentangled representations. arXiv preprint arXiv: 2019.
Ishaan Gulrajani and David Lopez-Paz. In search of lost domain generalization, 2020.

---

Under review as a conference paper at ICLR 2021
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In Jennifer Dy and Andreas Krause (eds.). Proceedings of the 35th International Conference Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 1861-1870, Stockholmsmässsan Stockholm Sweden, 10-15 Jul 2018. PMLR.
Patrick Haffner. Escaping the convex hull with extrapolated vector machines. In G. Dietterich, S. Becker, and Z. Ghahramani (eds.), Advances in Neural Information Processing Systems 14, pp. 753-760.
MIT Press, 2002.
Moritz Hardt, Eric Price, and Nathan Srebro. Equality of opportunity in supervised learning, 2016.
Trevor Hastie, Robert Tibshirani, and Jerome Friedman. The elements of statistical learning: data mining, inference, and prediction. Springer Science & Business Media, 2009.
Yue He, Zheyan Shen, and Peng Cui. Towards non-i.idd, image classification: A dataset and baselines, 2019.
Christina Heinze-Deml, Jonas Peters, and Nicolai Meinshausen. Invariant causal prediction for nonlinear models. Journal of Causal Inference, 6(2), Sep 2018. ISSN 2193-3685. doi: 10.1515/ jci-2017-0016. URL/ttpp: //dx doi. ort. (tlp:////ci12111-0110
Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness to common corruptions and perturbations. arXiv preprint 1903. 2019.
Dan Hendrycks and Kevin Gimpel. baseline for detecting misclassified and ousto--distt examples in neural networks. arXiv preprint arXiv:1610.0211 2016.
Dan Hendrycks, Mantas Mazeika, and Thomas Dietterich. Deep anomaly detection with outlier exposure. arXiv preprint
12.04666, 2018.
Dan Hendryeks, Mantas Mazeika, Saurav Kadavath, and Dawn Song. Using self-supervised learning can improve model robustness and uncertainty, 2019a.
Dan Hendrycks, Norman Mu, Ekin D. Cubuk, Barret Zoph, Justin Gilmer, and Balaji Lakshmi- narayanan. Augmix: A simple data processing method to improve robustness and uncertainty, 2019b.
Devon Hjelm, Alex Fedorov, Samuel Trischler, and Yoshua Bengio. Learning deep representations by mutual information estimation and maximization, 2018
Karan Grewal, Phil Bachman, Adam
Weihua Hu, Gang Niu, Issei Sato, and Masashi Sugiyama. Does distributionally robust supervised learning give robust classifiers?, 2016.
Maximilian Ilse, Jakub M Tomezak, and Patrick Forré. Designing data augmentation for simulating interventions. arXiv preprint arXiv:2005.01S: 2020.
Fredrik D. Johansson, David Sontag, and Rajesh Ranganath. Support and invertibility in domain- invariant representations, 2019.
Daphne Koller and Nir Friedman. Probabilistic graphical models: principles and techniques. 2009.
Masanori Koyama and Shoichiro Yamaguchi. Out-of-distribution generalization with maximal invariant predictor. 2020.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convo- lutional neural networks. In Advances in neural information processing systems, pp. 1097-1105. 2012.
Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M Hospedales. Deeper, broader and artier domain generalization. In Proceedings of the IEEE international conference on computer vision, pp. 5542- 5550, 2017.
10

---

Under review as a conference paper at ICLR 2021
Ya Li, Xinmei Tian, Mingming Gong, Yajing Liu, Tongliang Liu, Kun Zhang, and Dacheng Tao. Deep domain generalization via conditional invariant adversarial networks. In Proceedings of the European Conference on Computer Vision (ECCV), pp. 624-639 2018.
Zachary C Lipton. Yu-Xiang Wang, and Alex Smola. Detecting and correcting for label shift with black box predictors. arXiv preprint arXiv: 10V::1802.0391. 2018.
Mingsheng Long, Zhangjie Cao, Jianmin Wang. and Michael Jordan. Conditional adversarial domain adaptation. In Advances in Neural Information Processing Systems, pp. 1640-1650, 2018.
Nicolai Meinshausen, Peter Bühhmann, et al. Maximin effects in inhomogeneous large-scale data. The Annals of Statistics, 43(4):1801-1830; 2015.
Jose G Moreno-Torres Troy Raeder, RocíO Alaiz-RodriCo Nitesh V Chawla, and Francisco Herrera, A unifying view on dataset shift in classification. Pattern recognition, 45(1):521-530, 2012.
Sinno Jialin Pan, Ivor W Tsang, James T Kwok, and Qiang Yang. Domain adaptation via transfer component analysis. IEEE Transactions Neural Networks, 22(2):199-210, 2010.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Vetito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. dAlché- Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp. 8024-8035, Curran Associates, Inc., 2019. URL /ttptttpp//ppeees.uuiusss...e 9015 5015-10 pytorch an imperat ive style high
-eee Secteep--eearinnn learning 1 library.
pdf
Judea Pearl. Causality: Models, Reasoning and Inference. Cambridge University Press, New York, NY, USA, 2nd edition, 2009. ISBN 055689, 9780521895606.
Jonas Peters, Peter Bühlmann, and Nicolai Meinshausen. Causal inference by using invariant prediction: identification and confidence intervals. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 1855 2016.
Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Elements of causal inference: foundations and learning algorithms. 2017.
Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do imagenet classifiers generalize to imagenet? arXiv preprint arXiv: 1900..01111111111 2019.
Dominik Rothenhäusler, Nicolai Meinshausen, Peter Bühlmann, and Jonas Peters. Anchor regression: heterogeneous data meets causality, 2018.
Shiori Sagawa, Pang Wei Koh, Tatsunori B. Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization, 2019.
Subham S. Sahoo, Christoph H. Lampert, and Georg Martius. Learning equations for extrapolation and control, 2018.
Bernhard Schölkopf, Dominik Janzing, Jonas Peters, Eleni Sgouritsa, Kun Zhang, and Joris Mooij. On causal and anticausal learning. In Proceedings of the 29th International Coference International Conference on Machine Learning, ICML' 12, pp. 459-466, Madison, WI, USA, 2012. Omnipress.
ISBN 978145031 12851.
Connor Shorten and Taghi M Khoshgoftnar. A survey image data augmentation for deep learning. Journal of Big Data, 6(1):60, 2019.
Aman Sinha, Hongseok Namkoong, Riccardo Volpi, and John Duchi. Certifying some distributional robustness with principled adversarial training, 2017.

---

Under review as a conference paper at ICLR 2021
Amos J Storkey. When training and test sets different: characterising learning transfer. In In Dataset Shift in Machine Learning. Citeseer, 2009.
Jiawei Su, Danilo Vasconcellos Vargas, and Kouichi Sakurai. One pixel attack for fooling deep neural networks. IEEE Transactions on Evolutionary Computation, 2385):828884,, 2019.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, Timothy Lillicrap, and Martin Riedmiller. DeepMind control suite. Technical report, DeepMind, January 2018.
Yonglong Tian, Dilip Krishnan, and Phillip Isola. Contrastive multiview coding, 2019.
Antonio Torralba and Alexei A Efros. Unbiased look at dataset bias. In CVPR 2011, pp. 1521-1528. IEEE, 2011.
Eric Tzeng, Judy Hoffman, Kate Saenko, and Trevor Darrell. Adversarial discriminative domain adaptation. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 7167-7176, 2017.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding, 2018
Haohan Wang, Zexue He, Zachary C Lipton, and Eric P Xing. Learning robust representations by projecting superficial statistics out. arXiv preprint arXiv: 19Xvv01903 2019.
Robert C. Williamson and Aditya Krishna Menon. Faimness risk measures, 2019.
Xi Wu, Yang Guo, Jiefeng Chen, Yingyu Liang. Somesh Jha, and Prasad Chalasani. Representation bayesian risk decompositions and multi-source domain adaptation, 2020.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. arXiv preprint arXiv:161 1.03530, 2016.
Hongyi Zhang, Moustapha Cisse, Yann N. Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization, 2017.
Han Zhao, Remi Tachet des Combes, Kun Zhang, and Geoffrey J. Gordon. On learning invariant representation for domain adaptation, 2019.
12

---

Under review as a conference paper at ICLR 2021
Appendices
A APPENDIX OVERVIEW
Our code available online at:
c75e2344a397/
The sections of our appendix are as follows:
A) Overview
B) Definition and discussion of extrapolation in machine learning
C) Illustrative examples of how REx works in toy settings
D) Background on causal models
E) Theory
F) The relationship between MM-REX V-REx, and the role each plays in work
G) Further results and details for experiments mentioned in main text
H) Experiments not mentioned in main text
1) Overview of other topics related to OOD generalization
B DEFINITION AND DISCUSSION OF EXTRAPOLATION IN MACHINE LEARNING
We define interpolation and extrapolation as follows: interpolation refers to making decisions or predictions about points within the convex hull the training examples and extrapolation refers to making decisions predictions about points outside their convex hull 8 This generalizes the familiar sense these terms for one-dimensional functions. An interesting consequence of this definition is: for data of high intrinsic dimension, generalization requires extrapolation (Hastie et al. 2009), even in the i.i.d. setting. This is because the volume of high-dimensional manifolds concentrates near their boundary; Figure 6
Extrapolation in the space of risk functions. The same geometric considerations apply to extrap- olating to new domains. Domains can be highly diverse, varying according to high dimensional attributes, and thus requiring extrapolation to generalize across. Thus Risk Extrapolation might often do better job of including possible test domains in its perturbation than Risk Interpolation does.
Training points
Test point
Figure 6: Illustration of the importance of extrapolation for generalizing in high dimensional space. In high dimensional spaces, mass concentrates near the boundary of objects. For instance, the uniform distribution over ball in N + 1-dimensional space can be approximated by the uniform distribution over the N -idimensional hypersphere. We illustrate this in 2 dimensions, using the 1-sphere (i.e. the unit circle). Dots represent finite training sample, and the shaded region represents the convex hull of all but one member of the sample. Even dimensions, we can see why any point from a finite sample from such distribution remains outside the convex hull of the other samples, with probability 1. The only exception would be if two points in the sample coincide exactly.
SSurprssingly we were not able to find existing definition of these terms in the machine learning literature. They have been used in this sense (Hastic et al. 2009 Haffner 2002) but also to refer to strong generalization capabilities generally (Sahon al. 2018
13

---

Under review as a conference paper at ICLR 2021
C ILLUSTRATIVE EXAMPLES OF HOW REX WORKS IN TOY SETTINGS
Here, we work through two examples to illustrate:
1. How to understand extrapolation in the space of probability density/mass functions (PDF/MMss)
2. How REx encourages robustness to covariate shift via distributing capacity more evenly across possible input distributions.
C.I 6D EXAMPLE OF REX
Here we provide a simple example illustrating how to understand extrapolations of probability distributions. Suppose X E {0, 1, and Y É {0,1]; so there are a total of 6 possible types of examples, and we represent their distributions in particular domain point in 6D space: (P,,,)), P(1, 0), P(1, 1), 0), P(2, 1)). Now, consider three domains c1, C2, eg given by
1. (a,,,e,,,e,f)
2. (a,b,e,d,e k. f.ff+h. k)
3. (2a, 2b, c(1 atb) d(1 e, f)
The difference between 61 and 62 corresponds to shift in P(Y 2), and suggests that Y cannot be reliably predicted across different domains when el and eg tells that the relative probability of X = want our model to be robust to these sorts of covariate shifts. Extrapolating risks across these domains effectively tells the model: "don't 't bother trying to predict Y when
2. Meanwhile, the difference between
X can change, and we might
2 (i.e. aim for and X cases". By way of comparison. IRM would also aim for P(Y = 1 = 2) = .5, whereas ERM would aim for 3/+k (assuming Di1 = D21 = Dzl). And unlike REx, both ERM and IRM P(YY = X = 2) 3e+37
P(Y = 1|X = 2) .5), and split your capacity equally across the
would split capacity between X x= 0/11/2 cases according to their empirical frequencies.
C.2 COVARIATE SHIFT EXAMPLE
Viewing REx as robust learning over the affine span of the training distributions reveals its potential to improve robustness to distribution shifts. Consider a situation in which model encounters two types of inputs: COSTLY inputs with probability and CHEAP inputs with probability 1 q. The model tries to predicts the input - it outputs COSTLY with probability p and CHEAP with probability 1 - p. If the model predicts right its risk is O. but if it predicts COSTLY instead of CHEAP it gets a risk u = 2, and if it predicts CHEAP instead of COSTLY gets a risk v = 4. The risk has expectation (p) (1 (1 q)u + pqv. We have access to two domains with different input probabilities < Q2- This is an example of pure covariate shift.
We want to guarantee the minimal risk over the set of all possible domains:
min max Rg(r) (1 g)u + pqv pe 4E(0.1)
as illustrated in Figure The saddle point solution of this problem is p = w = N/uto and that R, (p) Ryp (p) can only happen for p w, so
(p) wn/u+o, Vg. From the figure the risk extrapolation principle will return the minimax optimal solution.
If ERM to minimize the risk, will pool together the domains into new domain with COSTLY input probability q = (q1 + /2) ERM will return p if q > w and = otherwise. Risk interpolation (RI) min, maxgelg ga) R2 (p) will predict p 0 q 02 w, p = 1 if qı, Q2 < and p w if qı < w < q2. We see that only REx finds the minimax optimum for arbitrary values of Q1 and 92.
14

---

Under review as a conference paper at ICLR 2021
4.0
Ra(p)
3.5
Ru(p)
maxRy(p)
3.0
(w))(w))
2.5
2.0
1.5
1.0
0.5
0.0
0.2
0.4
0.6
1.0
P
Figure 7: Each grey line is a risk R (p) as functions of p for a specific value of q. The blue line is when g w. We highlight in red the curve max, Rg (p) whose minimum is the saddle point marked by a purple star in w.
D BACKGROUND ON CAUSAL MODELS
D.1 CAUSAL MODELS: DEFINITIONS AND NOTATION
causal graphical model (Klllr & Friedman 2009 Pearl 2009) is a directed acyclic graph, where
edges point from causes (including noise variables) effects. Specifying the way effect N is computed based on the values of its causes (the mechanism, structural equation, or structural assignment for Y. commonly denoted fN. although Causal Model (SCM), denoted by C. We refer the reader to Peterss et al. (2017). definition 6.2, for more technical details At high level, SCMs are to causal graphical models as joint probability distributions are to probabilistic graphical models (PGMs); rather than just specifying causal or independence assumption, they fill in the details what values each node can take on and how exactly the nodes influence depend on one another.
m in this work) yields Structural
To highlight the difference between causal and probabilistic models, helps to think of an SCM as analogous to computational graph in deep learning framework such as Pytorch TensorFlow, with root nodes corresponding to independent noise variables. Elaborating: both objects specify how to compute node, given the values of its parents [ And while the value of node may be statistically dependent on any node to which is path- connectee, its value only directly influenced by the programmer setting it particular value. In this analogy,
by the value of its parents, setting the value of the node corresponds to an intervention. More formally, an intervention, can introduce edges, so
be any modification the mechanisms of one long as it does not introduce cycle. Applying an intervention on defines a new SCM, C.
variables;
The mechanisms of describe generative process for the data: absent any conditioning or in- terventions (which are viewed as producing a new SCM, rather than modifying C), the data can be generated using ancestral sampling, defining a joint distribution over all the nodes. This sam- pling process where, N m (Pa(N), defines an entailed distribution. P (N) for every node
Examining this sampling process reveals that the entailed distribution of N de- pends only on its mechanism and the value of its parents. As simple consequence, we have that: (Y Q) = Pa(Y)) = P(Y P(o\ddoQQ = q)) for Q representing all the other nodes in Al-
9SCMs sometimes called Structural Equation Models (SEMs), although this term is often reserved for SCMs with linear mechanisms. We do make assumptions ahout SCM mechanisms (c.g. linearity), unless stated otherwise.
InIn an SCM, every node has its unique independent noise variable, N. as parent, which is independent of all other nodes. Any "inherent randomness" in N arises from eN.
15

---

Under review as a conference paper at ICLR 2021
though all the nodes of C will not necessarily be observed, assume that they (i.e. X for supervised learning) for the remainder of the appendix, and overload m m (z) to refer to the entailed distribution for Y, given the values of its parents (only).
D.2 CAUSAL APPROACHES TO DOMAIN GENERALIZATION
number of works draw connections between causality, domain generalization, and robustness. Storkey 2009); Moreno- Torres et al. (2012); Schölkopf et al. (2012) characterize different types of distributional shift in probabilistic and causal terms. common thread in these works that the statistical relationship P(effect cause) should typically be preserved. In other words, interventions are assumed to be rare or sparse; or perhaps only certain types of interventions are allowed. For instance, the assumption of constant P(Y[X) is appropriate when X causes Y, and X but not Y may be subject to interventions. But when the roles are reversed, and causes X, constant P(XJY) is a better assumption (Lipton al. 2018). Bühlmann (2018b) and Rothenhäusler al. (2018) describe causal learning as a special case of robust learning, with the perturbation set given by the set of possible interventions that do not change the mechanism of Y.
E
THEORY
E.
CAUSAL DISCOVERY WITH REX
We prove that Risk Extrapolation learns the causal mechanism of Y under the same assumptions as used by Peters et al. (2016) their work on Invariant Causal Prediction (ICP). These are somewhat restrictive FFF
We assume that the causes of Y are observed, i.c. Pa(Y) X.
We assume that the homoskedasticity (a slight generalization of the additive noise setting assumed by Peters et al. (2016)).
The contribution of our theory (vs. ICP) is to prove that equalizing risks is sufficient learn the causes of Y. In contrast, they insist that the entire distribution of error residuals (in predicting Y) be the same across domains.
Theorem demonstrates practical result: we only need 3 interventions on each dimension of the input in order to identify linear SEM model using REx. Theorem 2 on the other hand, meant to provide insight into how the REx principle relates to causal invariance, not a practical result; the perturbation set in these theorems is uncountably infinite. It might also be possible to prove that REx identifies underlying causal structure in the limit of countably many diverse interventions, given some assumptions about the underlying casual structure, but we leave this to future work.
E.1.1 THE LINEAR CASE
We begin with theorem based on the setting explored by Peters et al. 2016):
Theorem 1. Given Linear SEM, X; +- Ejjt B(s X; es with Y = Xo. suppose predictor
fo(X) Eyzso X2 Ej satisfies for the mean-squared error (MSE) and perturbation set of domains that contains distinct dol) interventions for each X; >> 0. Then B3 = Bo.j. for all
Proof. We adapt the proof of Theorem 4i from Peters et al. (0016) to show that REx will learn the correct model under similar assumptions. Let Y f- >X be the mechanism for Y. assumed to be fixed across all domains, and let Ý Bx be our predictor. Then the residual is R(3) ( 3) +e. Define a; = % - B, and consider an intervention do(X, æ) on the youngest node X, with a, # O. Then as eqn 36/37 of Peters al. (2016), we compare the residuals R of this intervention and of the observational distribution:
Although we believe they could be substantially weakened, we emphasize that the goal of REx is OOD generalization, not causal discovery, which we view merely means to that end. Thus, in cases where the causal model is not maximal invariant predictor (MIP), expect RFx learn the MIP, not the mechanism.
16

---

Under review as a conference paper at ICLR 2021
Rohn (B) = azX, + CaxX, +e
R R = a,z + Eaxi +E
(7)
We now compute the MSE risk for both domains, set them equal, and simplify to find a quadratic formula for
E
(ajX;
Ea;Xi
=B (a + + Caxi+
(8)
0 = 022²² 22= + 20, E
X; + 5/ E (a;X,)²- 22 Σ
X, +
Since there are at most two values of x that satisfy this equation, any other value leads to violation of REx, so that 0 needs to zero contradiction. particular having domains with different do-interventions on every X; guarantees that the risks are not equal across all domains.
E.1.2 THE GENERAL CASE
Given the assumption that a predictor satisfies REx over all interventions that do not change the mechanism of Y, we can prove a much more general result.
Preliminaries: Our setting and assumptions for this case are:
1. We consider an arbitrary SCM, C, generating Y and X.
2. We use m to denote the mechanism for Y. as well as the entailed distribution m (x) =
(Y) = P(Y Pa(Y)) = (do(X = =)).
3. An intervention represents arbitrary change to the mechanisms of C, similarly to Arjovsky et al. (011)). and we define the causal perturbation set of g as the set of domains whose intervention does not change m' similarly to Peters al. (016)).
4. We assume that l is (strictly) proper scoring rule.
5. We say the SEM is homoskedastic when the Bayes error rate for E(m (x), m (x)) the same for all x E X, otherwise, it is heteroskedastic [2]
Theorem 2. In the homoskedastic case, a predictor that satisfies REx over Ed uses m² (x) as its predictive distribution input x. for all X.
The intuition of the proof is as follows:
1. We first show that that the distribution of Y given its parents doesn't 't depend the domain, and so m can be used to make reliable point-wise predictions across domains.
2. This translates into equality risk across domains when the overall difficulty of the examples is held constant across domains, by assuming homoskedasticity [3
3. While different predictor might a better job on some domains, we can always find an domain where it does worse than mY, and so m is both unique and optimal.
We emphasize that the predictor is not restricted to any particular class of models, and a generic function f X + P(Y). where P(Y) is the set of distributions over Y. Hence, drop & from the below discussion and simply use f to represent the predictor, and R(f) ) its risk.
Note that our detinitions constructed in Arjovsky et al. (2019). Section 5.1. but rather are generalization of the definitions of these terms commonly used in statistics. Specifically, for us, heteroskedasticity means that the "predicatability" (e.g. inputs x, whereas for Arjovsky al. (2019), it means the predicatahility of Y given input varies across domains; we refer
aoamoskedattic/heeerssseee
do not correspond to the types of domains
variance) of Y differs
hemtshctrrokkcdassiittt for clarity. 13 Note could also assume no covariate shift in order to fix the difficulty, but this seems hard to motivate in the context of interventions X.
second type as
17

---

Under review as a conference paper at ICLR 2021
Proof. Let (f. 2:) be the loss of predictor f on point in domain e, and (f) [pp(() (f.
be the risk of f in e. Define ((x) as the domain given by the intervention do(X = x)), and note that (()=(()) = R( (J,
Let X1 = Pa(Y). For every I E X, (x) = P(Y]do(X = x)) = P(Y]do(X; = =1)) = P = x1) is invariant (meaning "independent of domain") by definition; P(Y ddo(X = x1) follows from the semantics of SEM/SCMs, and
r)) P(Y]do(X, the fact that we don't allow m² to change across domains. Specifically Y is always generated by the same ancestral sampling process that only depends on X1 and ex_ Thus the risk of the predictor m v (x) at point I, R'mm x)) = {(m² (z), m (x)) is also invariant, soit R(m z). Thus Rs (m depend on x (homoskedastic case).
xι)) P(Y]X
(p)(z) R'mm x) mm()) R(m x) is invariant whenever R(m Y ,2) does
Now, we show that any other 9 achieves higher risk than m for at least one domain. This demon- strates both that m² achieves minimal risk (thus satisfying REx), and that it is the unique pre- dictor which does so (and thus no other predictors satisfy REx). We suppose such a exists and construct an domain where it achieves higher risk than m.. Specifically, if g * m then let
E X' be point such that g(z) * m m²(x). (x). And since l is a strictly proper scoring rule, this
implics that (((()) m (x)) > 8(m² y (z), m (x)). But the is exactly the risk of on the domain
= u(do(X r)), and thus و achieves higher risk than m in 6.
Remark. Even given infinite data from distribution with full support over Et, the ERM principle
does not, in general, recover the causal mechanism for Y. Rather, the ERM solution depends
on the distribution over domains. As an example, consider the colored MNIST task with 1
mass on the training domain where P(color = RedlY = 1) = P(color = Green Y = 0) = .9, P(color RRd = 0) P(color Green|Y = 1) P(Y = 0) P(Y = 1) = .5.
and mass distributed across other domains using a distribution with full support. For e close to 1, the ERM solution will be dominated by this domain, and the model's predictions will approximate
P(Y = l/color = Red) x .9, P(Y l/color Green) R .1.
E.2 REX AS DRO
We note that MM-REx is also performing robust optimization over a convex hull, see Fig- ure [ The corners of this convex hull correspond to "extrapolated domains" with coefficients (Amin Amin:.... .... (1 (m 1)Amin)) (up to some permutation). However, these domains do not necessarily correspond to valid probability distributions: in general, they are quasidistributions, which can assign negative probabilities to some examples. This means that, even if the original risk functions were convex, the extrapolated risks need not be. However, in the case where they are convex, then existing theorems, such as the convergence rate result of (Sagawa et al. 2019). This raises several important questions:
1. When is the affine combination of risks convex?
2. What are the effects of negative probabilities on the optimization problem REx faces, and the solutions ultimately found?
We hypothesize that affine combinations of risks will remain convex in mean- -quured error regression problems under fairly weak assumptions.
Negative probabilities: Figure illustrates this for case where X' Z2, i.e. a binary vector of length 2. Suppose æ1, 22 are independent in our training domains, and represent the distribution for a particular domain by the point (P(X1 = 1), P(X, = 1)). And suppose our training distributions have (P(X, 1), P(x, = 1)) equal to { (44, (.6, .),, (.6, .9)} with P(Y fixed.
F THE RELATIONSHIP BETWEEN MM-REX VS. V-REx, AND THE ROLE EACH PLAYS IN OUR WORK
The MM-REx and V-REx methods play different roles in our work:
18

---

Under review as a conference paper at ICLR 2021
1.0
P(X2 = 1)
Training environments
0.5
Valid distributions
Perturbation set
0.0
0.0
0.5
1.0
P(X1 =
Figure 8: The perturbation set for MM-REx can include "distributions" which assign invalid (e.g. negative) probabilities to some data-points, The range of valid distributions P(X) is shown in grey, and P(X) for 4 different training domains are shown as red points. The interior of the dashed line shows the perturbation set for Amin = 1/22.
We use MM-REx to illustrate that REx be instantiated variant of robust optimization, specifically a generalization of the common Risk Interpolation approach. We also find MM- REx provides a useful geometric intuition, since we can visualize its perturbation set as an expansion of the convex hull of the training risks or distributions.
We expect V-REx to be the more practical algorithm. It is simple to implement. And it performed better in our CMNIST experiments; believe this may be due V-REx providing a smoother gradient vector field, and thus more stable optimization, see Figure
Either method recovers the REx principle limiting case, as we prove in Section F.I We also provide sequence of mathematical derivations that sheds light on the relationship between MM-REx and V-REx in Section F22 we can view these as a progression of steps for moving from the robust optimization formulation of MM-RES REx the penalty term of V-REx:
1. From minimax to closed form: We show how to arrive at the closed-form version of MM-REx provided in Eqn.5)
2. Closed form mean absolute error: The closed form of MM-REx is equivalent to mean absolute error (MAE) penalty term when there are only two training domains.
3. V-REx as mean squared error: V-REx exactly equivalent penalty term (always). Thus in the case of only two training domains, the difference between MM-REx and V-REx just different choice of norm.
mean squared error
F.1 V-REX AND MM-REX ENFORCE THE REx PRINCIPLE IN THE LIMIT
We prove that both MM-REx and V-REx recover the constraint of perfect equality between risks in the limit of Amin + -0 or 3 + co, respectively. For both proofs, we assume all training risks are finite.
The MM-REx risk of predictor fe. RMM-REX + Amin -0o unless Rd = RD for all training domains d,
Proposition 1.
Proof. Suppose the risk is not equal across domains, and let the largest difference between any two training risks be > O. Then RMM-RRB max, (0) m.min max, R. (0) + Amin Σm (0) 2 max, (e) Amine, with the inequality resulting from matching up the copies of Amio max. R2 with the terms in the sum and noticing that each pair has non-negative value (since R, max, R. is non-positive and Amin is negative), sends this lower bound on
(1 mAmin) max, (e) + Amin Σ (,)
and least one pair has the value -uunn E. Thus sending À RMM- REx to oc and hence -MM-REX + oc well.
Proposition 2. The V-REx risk of predictor fo. RV-REX (f) - oo as B oo unless Rd = R for all training domains d,
19

---

Under review as a conference paper at ICLR 2021
V
Restal
0
R2(0)
Figure 9: Vector fields of the gradient evaluated at different values of training risks (0), (2( We compare the gradients for RMM.RE (left) and RV.RE (right). Note that for RVRRS the gradient vectors curve smoothly towards the direction of the origin, as they approach the diagonal (where training risks are equal); this leads to a smoother optimization landscape.
Proof. Again, let e > O be the largest difference in training risks, and let μ be the mean of the training risks. Then there must exist an such that RR #4| c/2. And thus (ar;(R,(0)) (R, /1²² 2 (e/2)². since all other terms in the sum are non negatveee Since e > 0 by assumption. penalty term positive and thus RV-REx(( = Σ; (+) 3Var (V, goes infinity as 00.
F.2
CONNECTING MM-REX TO V-REX
F.2.1 CLOSED FORM SOLUTIONS TO RISK INTERPOLATION AND MINIMAX- REX
Here, we show that risk interpolation equivalent to the robust optimization objective of Eqn. 4 Without loss of generality, let R, be the largest risk, R. R1, for all Thus express
R. R₁ de for some non-negative de, with d1 = 0 2 de for all e. And thus we can write the weighted sum of Eqn. 5 as:
(M) = max Ne (e
(10)
Ain e=1
max Exd==1 de Z2mmm
(R - de)
(11)
(i)
(dd)
(12)
Amin
(13)
Now, since de are non-negative, -de is non-positive, and the maximal value of this sum is achieved when X Amin for all 2, which also implies that ê = 1 (m 1)Amin This yields the closed
form solution provided in Eqn. 5 The special case of Risk Interpolation, where Amin = 0, yields Eqn.
F.2.2 MINIMAX-REX AND MEAN ABSOLUTE ERROR REX
In the case of only two training risks, MM-REx equivalent to using a penalty on the mean absolute error (MAE) between training risks. However, penalizing the pairwise absolute errors is not equivalent when there are m > training risks, as we show below. Without loss of generality, assume that R1 R2 Rmn Then (1 (1/2 /2 of) the RMAE penalty term is:
20

---

Under review as a conference paper at ICLR 2021
ΣΣ(R, (R.- R₂) mR Σ
R + (m 1)
Σ R₂.
(14)
3mm
5mm-1
= IRR - ΣΣ
Ri
(15)
) iS.
ΣJR, - - (m j 1)
(16)
Σ(25 (2 1
(17)
For m 2, have 1/2RMAE (2 x - 1)R, + (2+ 2- 1)R2 R2 R1. Now, adding this penalty term with some coefficient SMAE to the ERM term yiclds:
RMAE = R1 + R2 + AMAE R1) (1 - (MAA) + (1 + SM
(18)
(19)
We wish to show that this is equal to RMM for an appropriate choice learning rate MMA and hyperparameter MAME. Still assuming that R1 R2, we have that:
RMM (1 Amia) + Amnin R1
(20)
Choosing "MAE = 1/2 is equivalent to multiplying RMM by 2, yielding:
2RMM 2(1 Amin RR₂ + 2Amin R1
(21)
Now, in order for RMAR 2RMM, we need that:
2 2Amin BMAE
(22)
2Amin 1 BMAE
(23)
(24)
And this holds whenever BMAΣ = 1 2Amin When m > 2, however, these are not equivalent, since
RMM puts equal weight on all but the highest risk, whereas RMAE assigns different weight each risk.
F.2.3 PENALIZING PAIRWISE MEAN SQUARED ERROR (MSE) YIELDS V-REx
The V-REx penalty (Eqn. 6 is equivalent to the average pairwise mean squared error between all training risks (up constant factor 2). Recall that R, denotes the risk domain i. We have:
1
2n ΣΣ (R, R₂)²
ΣΣ
(R. + R2 2R,R,)
(25)
2n
1
Σ
R₂
Σ
R3
ΣΣ
R.R,j
(26)
2n
2n
n2
2
1 Σ Ri
R2
(27)
"
Var(RR.
(28)
G FURTHER RESULTS AND DETAILS FOR EXPERIMENTS MENTIONED MAIN
TEXT
G.1 CMNIST WITH COVARIATE SHIFT
Here present the following additional results:
1. Figure of the main text with additional results using MM-REx, see G. used the "default" parameters from the code of Arjovsky et al. 2009
These results
21

---

Under review as a conference paper at ICLR 2021
2. A plot with results on these same tasks after performing random search over hyperparameter values similar to that performed by Arjovsky et al. (2019).
3. A plot with the percentage of the randomly sampled hyperparameter combinations that have satisfactory (>> 50%) accuracy, which we count as "success" since this is better than random chance performance.
These results show that REx is able handle greater covariate shift than IRM, given appropriate hyperparameters. Furthermore, when appropriately tuned, REx can outperform IRM in situations with covariate shift. The lower success rate of REx for high values of p is because produces degenerate results (where training accuracy is less than test accuracy) more often.
The hyperparameter search consisted of uniformly random search of 340 samples over the following intervals of the hyperparameters:
1. HiddenDim = 12**7,
2. L:2RegularizerWeiihtt [10**-2, 10**44]
3. =10*** 10**-4.3]
4. Penalty-Anneallterss [50, 250]
5. PenaltyWeight [10**2, 10**6]
6. Steps [201, 601]
0.6
0.6
0.5
Tea accuNTT
0.5
TeMe ACTUPMC
0.5
0.4
0.4
Fonvier Cleesdinn
Randres Bisescine
Randorn Guslarg
V
55
02
Figure 10: This Figure of main text with additional results using MM-REx. For each covariate shift variant (class imbalance, digit imbalance, and color imbalance from left to right as described in "CMNIST with covariate shift" subsubsection Section 4.1 in main text) of CMNIST, the standard error (the vertical bars in plots) is higher for MM-REx than for V-REx.
Test AECuNcy
Tee AcTurnY
0.5
Randum Guessing
Rancori Guessing
00.
Random Guessing
REX
INW
0.6
REx
IN
0.3
REx
IR2
Pahapelii (0.12.3.4))
Pahhss3d (1.2)
PIRI/Res) PSGl/Crea))
Figure 11: This is Figure of main text (class imbalance, digit imbalance, and color imbalance from left to right as described in "CMNIST with covariate shift" subsubsection of Section 4.1 in main text), but with hyperparameters REx and IRM each tuned perform as well as possible for each value of for each covariate shift type.
G.2 SEMS FROM "INVARIANT RISK MINIMIZATION"
Here we present experiments on the (linear) structural equation model (SEM) tasks introduced by Arjovsky et al. 2019). Arjovsky et al. (2019) construct several varieties of SEM where the task is to predict targets from inputs X1.X, where X1 are (non-anti-causal) causes of Y, and X₂ are (anti-causal) effects of Y. We refer the reader to Section 5.1 and Figure 3 of Arjovsky et al. (2019) for more details. We use the same experimental settings as Arjovsky et al. (2019) (except we only run 7 trials), and report results in Table
22

---

Under review as a conference paper at ICLR 2021
pitterta ofre N0,, an 2.1.22000000
30-
70-
20
20-
KEX I821
RIE
HEE
parcent of
1RM
0.2
Pichapeldd (0.22.3.11
Phstspejcl (1.2345,75)
Figure 12: This also corresponds to class imbalance, digit imbalance, and color imbalance from left to right as described in "CMNIST with covariate shift" subsubsection of Section 4.1 in main text; but now the y-axis refers to what percentage of the randomly sampled hyperparameter combinations deemed to to be satisfactory. We define satisfactory simultaneously being better than random guessing and having train accuracy greater than test accuracy. For less than .5, a larger percentage of hyperparameter combinations a larger percentage of hyperparameter combinations are often satisfactory for IRM than for REx because train accuracy is greater than test accuracy for more hyperparameter combinations for IRM. We stipulate that train accuracy must be greater than test accuracy because test accuracy being greater than train accuracy usually means the model has learned a degenerate prediction rule such as "not color".
often satisfactory for REx than for IRM; for p greater than .5,
These experiments include several variants of simple SEM, given by:
X1 N1
Y WI-Y + Ny
X₂ Wy + N2
Where Ny, N2 are all sampled i.i.d. from normal distributions. The variance of these distributions may vary across domains.
While REx achieves good performance the hoaaninn-oommkkdassiie case, performs poorly the domain-heteroskedastie case, where the amount of intrinsic noise, og in the target changes across domains Intuitively, this because the irreducible error varies across domains in these tasks, meaning that the risk will be larger on some domains than others, even if the model's predictions match the expectation E( ((Y|Pa(()) We tried using a "baseline" (see Eqn. 4 of re Var(Y,) MMeinshassen et al. 2015) to account for the different noise levels in Y, but this did not work.
include a mathematical analysis the simple SEM given above in order to better understand why REx succeeds in the domain- homoskedastiicc but not the domain-heterrokedasti case. Assuming that Y, X2, X2 are scalars, this SEM becomes
X1 N1
Y W1+t N1 + Ny
X2 Wy-+2011 -y Ni + Wy-22N N2
We consider learning a model y aX1 3X2. Then the residual is:
Y = (a + w 2 (3 1))N, + (Bwy-12 1)Ny + 3N22
Since all random variables have zero mean, the MSE loss is the variance of the residual. Using the fact that the noise N1, Ny, N2 are independent, this equals:
B( y> = (a + w w11-y 1)) + (Buy-12 1)203 + 83
Thus when (only) d2 changes, the only way to keep the loss unchanged is to set the coefficient in front of 02 to O, meaning O. By minimizing the loss, then recover homaann-homokkddastiic setting, the loss equality constraint of REx yields the causal model. On the other hand, if (only) σy changes, then RFx enforces 3 =1/ 1/wy-42 which then induces a recovering the anticausal model.
Wi ; i.e. in the
14See Footnote 12
23

---

Under review as a conference paper at ICLR 2021
HOU(e)
FOU(ne)
FOS(c)
FOS(nc)
IRM
00001000.000
0.001=00.000
0.001=0.000
0.000-20.000
REx, Pa
0:0.00000000
0.008:10.002
0.00710.002
0.000::0.000
REx,
V(Y.)
0.816=0.149
1.417#0.442
0.919=0.991
0.000=000000
POU())
POU(ne)
POS(e)
POS(ne)
IRM
0.004.1.0.01
0.006.10.003
0.002.10.000
0.000.10.000
REx,
0.004-10.001
0.004-10.001
0.002-000000
0.000:10.000
REx, r. v0G)
0.915=00.055
L.133:0.085
0.937=0.090
0.000=000000
FEU(G)
FEU(ne)
FES(c)
FES(nc)
IRM
0.0053-1::0015
0.1025-0.0073
0.0193-10.054
0.0000:10.000
REx, n. =0
0.0390=0.00899
19.1518+3.301
1.7646#1.186
0.0000+:000000
REx. r. V(G)
0.7713=t01402
1.0358+011244
0.8603=0.0233
0.0000+000000
PEU())
PEU(e))
PES(c)
PES(nc)
IRM
0.0102:1000099
0.09911000216
0.0510:10.0049
0.0000-10.0000
REx, 0
0.0784=000211
46.7235=11.4400
8. 8.364=226108 0.0000=:00000
REx, r. v(G))
1.099710.00299
0.9946.10.0877
1.025210.0819 0.000011..0000
Table 4: Average mean-squared error between true and estimated weights on causal (X1) and non-causal (X2) variables. Top 2: When the level of noise in the anti-causal features varies across domains, REx performs well (FOU, FOS, POU, POS). Bottom 2: When the level of noise in the targets varies instead, REx performs poorly (FEU, FES, PEU, PES). Using the baselines re V(Y) does not solve the problem, and indeed, hurts performance on the homoskedastic domains.
While REx (like ICP (Peters et al. 2016)) assumes the mechanism for Y is fixed across domains (meaning P(Y\Pa(Y ) is independent of the domain, e), IRM makes the somewhat weaker assump- tion that (Y (aY)) is independent of domain. While is plausible that an appropriately designed variant of REx could work under this weaker assumption, we believe forbidding interventions on Y not overly restrictive, and such an extension for future work.
G.3 REINFORCEMENT LEARNING EXPERIMENTS
Here provide details and further results the experiments in Section 4.4 We take tasks from the Deepmind Control Suite Tassa et al.] 2018) and modify the original state, s, to produce observation, (s e, vis') including noise and spurious features rys', where s' contains 1 or dimensions of 1/2/3 3 for the two training and test domains, respectively. The
8. The scaling factor takes values agent takes input and learns representation using Soft Actor-Critic (Haarnoja et al. 2018)) and an auxiliary reward predictor, which is trained to predict the next 3 rewards conditioned on the next actions. Since the spurious features are copied from the state before the noise is added, they are more informative for the reward prediction task, but they do not have an invariant relationship with the reward because the domain-dependentt, 17.
The hyperparameterss used for training Soft Actor-Critic can be found in Table 5
We used cartpole e_swi ngup as development task to tune the hyperparameters of penalty weight (cho- sen from 10.01, 0.1, 1, 10)) and number of iterations before the penalty is turned up (chosen from 5000, 10000, 20000)), both for REx and IRM. The plots with the hyperparameter sweep are in Figure
3
24

---

Under review as a conference paper at ICLR 2021
Hyperparameter sweep for IRM and REx on cartpole_swinneeeeeeeeeee Green, blue, and orange curves correspond to REx, ERM, and IRM, respectively. The subfigure titles state the penalty strength ("penalty") and after how many iterations the penalty strength was increased ('iters")) We chose penalty factor of and 10k iterations.
Figure 13:
Parameter name
Value
Replay buffer capacity
1000000
Batch size
1024
Discount
0.99
Optimizer
Adam
Critic learning rate
10
Critic target update frequency
2
Critic Q-function soft-update rate TO
0.005
Critic encoder soft-update rate Tene
0.005
Actor learning rate
10-5
Actor update frequency
2
Actor log stddev bounds
[-5,2]
Encoder learning rate
10
Decoder learning rate
10
Decoder weight decay
10
L1 regularization weight
10-5
Temperature learning rate
10
Temperature Adam's Init temperature
0.9
0.1
Table 5: A complete overview of hyperparameters used for reinforcement learning experiments.
H EXPERIMENTS NOT MENTIONED IN MAIN TEXT
We include several other experiments which do not contribute directly to the core message of our paper. Here is a summary of the take-aways from these experiments:
1. Our experiments in the CMNIST domain suggest that the IRM/V-REx penalty terms should amplified exactly when the model starts overfitting training distributions.
2. Our financial indicators experiments suggest that IRM and RFx often perform remarkably similarly in practice.
H.1 A POSSIBLE APPROACH TO SCHEDULING IRM/REX PENALTIES
We' ve found that REx and IRM are quite sensitive to the choice of hyperparameters. In particular, hyperparameters controlling the scheduling of the IRM/V-REx penalty terms of critical importance.
25

---

Under review as a conference paper at ICLR 2021
For the best performance, the penalty should be increased the relative weight of the penalty term after approximately 100 epochs of training (using so-called "waterfall" schedule (Desjardins et al. 2013)). See bb for comparison. We also tried an exponential decay schedule instead of the waterfall and found the results (not reported) were significantly worse, although still above 50% accuracy.
[4(b)
Given the methodological constraints of out-of- of-diiiributtion generalization mentioned in (Gulrajani & Lopez-Paz 2020), this could be a significant practical issue for applying these algorithms. We aim to address this limitation by providing a guideline for when to increase the penalty weight, based only on the training domains. We hypothesize that successful learning of causal features using REx or IRM should proceed in two stages:
1. In the first stage, predictive features are learned.
2. In the second stage, causal features are selected and/or predictive features are fine-tuned for stability.
This viewpoint suggests that we could use overfitting on the training tasks as indicator for when to apply (or increase) the IRM or REx penalty.
The experiments presented in this section provide observational evidence consistent with this hypoth- esis. However, since the hypothesis was developed by observing patterns in the CMNIST training runs, requires further experimental validation on a different task, which leave for future work.
H.1.1 RESULTS AND INTERPRETATION
In Figure I4| we demonstrate that the optimal point to apply the waterfall in the CMNIST task is after predictive features have been learned, but before the model starts memorize training examples. Before predictive features are available, the penalty terms push the model to learn a constant predictor, impeding further learning. And after the model starts memorize, it become difficult distinguish anti-causal and causal features. This second effect is because neural networks often have the capacity to memorize all training examples given sufficient training time, achieving and near-O loss (Zhang et al., (016). In the limits of this memorization regime, the differences between losses become small, and gradients of the loss typically do as well, and so the REx and IRMv1 penalties no longer provide strong or meaningful training signal, see Figure 15
H.2 DOMAIN GENERALIZATION, VLCS AND PACS
Here we provide earlier experiments on the VLCS and PACS dataset. We removed these experiments from the main text of our paper in favor of the more complete DomainBed results.
To test whether REx provides a benefit REx, IRM and ERM performance on the VLCS (orralba & Efros, [01I] and PACS (Li et al. 2017) image datasets. Both datasets are commonly-used for multi-source domain generalization. The task is to train on three domains and generalize to a fourth one at test time.
realistic domain generalization tasks, we compared
Since every domain in PACS is used test set when training on the other three domains, it is not possible to perform a methodologically sound evaluation on PACS after examining results on any of the data. Thus to avoid performing any tuning on test distributions, we use VLCS tune hyperparameters and then apply these exact same settings to PACS and report the final average over 10 runs on each domain.
We use the same architecture, training procedure and data augmentation strategy as the (formerly) state-of-the-art Jigsaw Puzzle approach (Carlucci et al. 2019) (except with IRM or V-REx intead of JigSaw auxilliary loss) for all three methods. As runs very noisy, we ran each experiment 10 times, and report average test accuracies extracted at the time of the highest validation accuracy on each run. Results PACS are in Table [ while detailed results and performance VLCS are left to the Appendix. On PACS we found that REx outperforms IRM and IRM outperforms ERM on average, while all are worse than the sfate-of oftthe-arttt Jigsaw method.
26

---

Under review as a conference paper at ICLR 2021
0.70
IRM
accuracy (%)
0.65
REx
0.60
0.55
0.50
0
100
200
300
400
0.06
Dtrain generalization gap (%)
accuracy (%)
overfitting
0.04
0.02
0.00
100
200
300
400
0.8
accuracy (%)
0.6
test
0.4
train
0.2
grayscale
100
200
300
400
epoch
Figure 14: Stability penalties should be applied around when traditional overfitting begins, to ensure that the model has learned predictive features, and that penalties still give meaningful training signals. Top: Test accuracy function of epoch at which penalty term weight is increased (learning rate is simultaneously decreased proportionally). Choosing this hyperparameter correctly essential for good performance. Middle: Generalization gap validation with 85% correlation between color and label (the same as the average training correlation). The best test accuracy is achieved by increasing the penalty when the generalization gap begins to increase. The increase clearly indicates memorization because color and shape are only 85%/75% correlated with the label, and so cannot be used make predictions with higher than 85% accuracy. Bottom: Accuracy on training/test sets, as well as an auxilliary grayscale set. Training/test performance reach 85%/15% after few epochs of training, but grayscale performance improves, showing that meaningful features are still being leamned.
We use all hyperparameters from the original Jigsaw codebase We use Imagenet pre-trained AlexNet features and chose batch-size, learning rate, as well as penalty weights based on performance on the VLCS dataset where test performance on the holdout domain was used for the set of parameters producing the highest validation accuracy. The best performing parameters on VLCS were then applied to the PACS dataset without further changes. We searched over batch-sizes in {128, 384]. over penalty strengths in {0. 0001001, 0.001, 0.01, 0.1, 1, 10). learning rates in {0.001, 0.01) } and used average performance over all 4 VLCS domains to pick the best performing hyperparameters, Table shows results on VLCS with the best performing hyperparameters.
LS
htttps://gthhub....../ hwwp://caaaamlcm geng
27

---

Under review as a conference paper at ICLR 2021
10a
train_nll
irmv1_pena rex_penalty
10
10-4
10-7
loss
10-10
10
16
0
5000
10000
15000
20000
25000
epoch
Figure 15: Given sufficient training time, empirical risk minimization (ERM) minimizes both REx and IRMv1 penalty terms Colored MNIST (without including either term in the loss function). This is because the model (a deep network) has sufficient capacity to fit the training sets almost perfectly. This prevents these penalties from having the intended effect, once the model has started to overfit. The y-axis is in log-scale.
28

---

Under review as a conference paper at ICLR 2021
The final parameters for all methods on PACS were a batch size of 384 with 30 epochs of training with Adam, using a learning rate of 0.001, and multiplying it by 0.1 after 24 epochs (this step schedule was taken from the Jigsaw repo). The penalty weight chosen for Jigsaw was 0.9; for IRM and REx was 0. We used the same auta-tugmentttion pipeline as the original Jigsaw code for ERM, IRM, Jigsaw and REx to allow for fair comparison.
The
VLCS
CALTECH
SUN
PASCAL
LABELME
Average
REx (ours)
96.72
63.68
72.41
60.40
73.30
IRM
95.99
62.85
74.7
59.61
72.54
ERM
94.76
61-92
69,03
60.55
74.56
Jigsaw (SOTA)
96.46
63.84
70.49
60.06
72.71
Table 6: Accuracy (percent) of different methods on the VLCS task. Results are test accuracy at the time of the highest validation accuracy, averaged over 10 runs. On VLCS REx outperforms all other methods. Numbers are shown in strike-thruugh because selected our hyperaarameters based on highest test set performance; the goal of this experiment was to find suitable hyperparameters for the PACS experiment.
PACS
Art Painting
Cartoon
Sketch
Photo
Average
REx (ours)
66.27
68.8
59.57=0.78
89.60=0.11
71.07
IRM
66.46110.31
68.60::0.00
58.6610.73
13.94
70.91
ERM
66.01
68.62++0.66
58.38++0.00
89.40++0.18
70.60
Jigsaw (SOTA)
66.96:10.39
66.67:0.41
61.27::0.73
89.54:10.19
71.11
Table 7: Accuracy (percent) of different methods on the PACS task. Results are test accuracy at the
time of the highest validation accuracy, averaged over 10 runs. REx outperforms ERM average, and performs similar to IRM and Jigsaw (the state-of-the-art).
H.3 FINANCIAL INDICATORS
We find that IRM and REx seem to perform similarly across different splits of the data in a prediction task using financial data. The dataset is split into five years, 2014-18, containing 37 publicly reported financial indicators of several thousand publicly listed companies each. The task is to predict if company's value will increase or decrease in the following year (see Appendix for dataset details.) We consider each year different domain, and create 20 different tasks by selecting all possible combinations of domains where three domains represent the training sets, one domain the validation set, and another the test set. We train MLP using the validation stopping point, with 3 = 10 The per-task results summarized in fig. 16 indicate substantial differences between ERM and IRM, and ERM and REx. The predictions produced by IRM and REx, however, only differ insignificantly, highlighting the similarity of IRM and REx. While performance on specific tasks differs significantly between ERM and IRM/REx, performance averaged over tasks is not significantly different.
determine early
H.3.1 EXPERIMENT DETAILS
We use v1 of the dataset published on 16 and prepare the data as described in 17 We further remove all the variables that are not shared across all years, leaving us with 37 features, and whiten the data through centering and normalizing by the standard deviation.
On each subtask, we train an MLP with two hidden layers of size 128 with tanh activations and dropout (p=0.5) after each layer. We optimize the binary cross-entropy loss using Adam (learning rate 0.001, 31 = 0.9, 32 = 0.999, = 10 10-8), 3), and an L2 penalty (weight 0.001). In the IRM/REx experiments, the respective penalty is added to the loss (19 1) and the original loss is scaled by
httpp::/// wwt::://wwwwww.. kaggl e. wtt.c/tccccccccccctcc wxw.://eepppeeeeeeee
29

---

Under review as a conference paper at ICLR 2021
Training envs.
Test IRM
Test REx
0.05
Task
0.00
Performance
-0.05
-0.10
2014
2015
2016
2017
2018
2014
2015
2016
2017
2018
2014
2015
2016
2017
2019
Figure 16: Financial indicators tasks. The left panel indicates the set of training domains; the middle and right panels show the accuracy the respective domains relative to ERM black dot corresponds to a training domain; a colored patch indicates the test accuracy on the respective domain.)
Overall accuracy
Min
Max acc.
ERM
54.6 4.6
47.6
66.2
IRM
55.3 + 5.9
45.9
67.5
REx
55.5 + 6.0
47.2
68.0
Table 8: Test accuracy of models trained on the financial domain dataset, averaged over all 20 tasks, as well as min./max. accuracy across the tasks.
factor 10 after 1000 iterations. Experiments are run for a maximum of 9000 training iterations with early stopping based on the validation performance. All results are averaged over trials. The overall performance of the different models, averaged over all tasks, is summarized in Tab. [ The difference in average performance between ERM, IRM, and REx not statistically significant, as the error bars very large.
I OVERVIEW OF OTHER TOPICS RELATED TO OOD GENERALIZATION
Domain adaptation (Ben-Daviee et al. 2010) shares the goal of generalizing to distributions test time, but allows some access to the test distribution. A common approach is to make different domains have similar distribution of features (Pan et al. (010). A popular deep learning method for doing so is Adversarial Domain Adaptation (ADA) (Ganin et al. 2016 Tzeng et al.) 2017 Long al. 2018 Li al. 2018), which seeks "invariant representation" of the inputs, distribution is domain-independent. Recent works have identified fundamental shortcomings with
one whose
this approach, however (Zhao al. 2019 Johansson et al., 2019 Arjovsky et al. 2019 Wu et 2020).
Complementary the goal of domain generalization out-of-distributiio detection (Hendrycks & Gimpell 2016; Hendrycks et al. 2018), where the goal is to recognize examples as belonging to new domain. Three common deep learning techniques that can improve OOD generalization adversarial training (Goodfellow et al. [2014] Hendrycks & Dietterich] 2019). self-supervised learning (van den Oord et al./ 2018 2018|Hjelm al. [2018 Henddycks et al.)
201 |019b] Albuqueqqu et al.
2020) and data augmentation (Krizhevsky
al. 2012 Zhang et al. 2017 Cubuk et al.) 2018
Shorten & Khoshgoftaar 2019 Hendrycks et al. 2019a] Carlucci et al,
2019). These methods can
also been combined effectively in various ways (Tian et al. 2019 Bachman et al.) 2019 Gowal
et al. 2019). Data augmentation and self-supervised learning methods typically use prior knowledge such 2D image structure. Several recent works also prior knowledge design augmentation strategies for invariance to superficial features that may be spuriously correlated with labels in object recognition tasks (He al. 20199; Wang al. 2019; Gowal al.) 2019 [Ise contrast, REx can discover which features have invariant relationships with the label without such prior knowledge.
2220)). In
30


Reviews:
[{"review_id": "TZNTpU6k5-_", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This manuscript studies the problem of domain generalization and proposes a method, dubbed Rex, for this purpose. The main movitation of this work over the invariant risk minimization (IRM) [1] paradigm is that IRM is not robust to covariate shift, while the authors claim that Rex can deal with both covariate shift and concept shift together. Although the argument that IRM is not robust to covariate shift in the feature space is true, out-of-domain generalization is not the original goal of IRM either. My main concern is that it is not clear to me why Rex could deal with both covariate and concept shift together, as from the optimization formulation in Eq. (6), neither invariant predictor nor invariant representation is enforced. In particular, no theoretical analysis is given to justify this claim. \n\nPerhaps what adds more confusion to me is that, what's the meaning of negative probability since the authors allow the combination weight \\lambda of different domains to be negative? As the authors have already pointed out on page 5 (Probabilities vs Risks), the risk is a linear functional of the joint distribution over X and Y, hence using affine combination in the risk functions from different domains directly translate to allowing the use of negative coefficient for probabilities. This is quite strange, since the mixture distribution is only a convex combination, not affine combination.\n\nFurthermore, I also found some of the discussions in the related work section misleading: \n-   \"The first method for invariant prediction ... is IRM\". This is not correct. As the authors have already realized, ICP [2] has been proposed in 2015, and the definition of IRM is essentially the same as ICP. \n\n-   The discussions about invariant representations on page 4 are not accurate. Only P_e(\\phi) (but not P_e(\\phi | Y)) is called invariant representations, and only this one can fail domain adaption if the marginal label distributions differ. The C-ADA does not try to find invariant P_e(\\phi | Y). Instead, it tries to find invariant P_e(\\phi x \\hat{Y}), where \\hat{Y} is the classifier output, and this is precisely the reason why C-ADA also fails under different label distributions. See more discussions in [4]. In fact, it has recently been shown in [4] that matching P_e(\\phi | Y) provably works for domain adaptation, see Theorem 3.1 of [4]. \n\n-   \"Also, unlike Rex, IRM seeks to match E[Y | \\phi], not the full P(Y | \\phi)\". Again, this discussion is misleading. For the purpose of out of domain generalization, the learner does not need to match the full distributions P(Y | \\phi). Only the conditional mean matters. See Theorem 4.1 of [5] as well as Theorem 3.3 of [6] in the context of fairness.\n\n-   About the discussion on fairness of equalizing risk across groups. In fact a sufficient condition for this goal has been proved in Theorem 3.3 of [6]. Given the close relationship between these two problems, I feel it's necessary to cite and have a discussion of this result here as well. \n\n\nMore detailed comments:\n-   The reference of David et al., 2010 should be Ben-David et al. 2010 on page 2\n\n-   I think the naming of invariant prediction on page 2 is not very accurate. To be precise, as long as the same hypothesis (classifier) is used over different domains, this amounts to be an invariant prediction rule. Instead, what IRM enforces is the invariant OPTIMAL predictor, i.e., invariant conditional means.  \n\n-   Eq. (3) is not correct: the second term is a weighted combination of samples from different domains, where the larger the sample size from a domain the more weight it has in the combination. However, the right most term has a wrong weight for a domain. The correct one should be |D_e| / \\sum_e |D_e|. \n\n-   On page 4, the citation of Pan et al. 2010 is not accurate. Invariant representations for domain adaptation is first proposed by [3].\n\n\n\n[1]     Invariant Risk Minimization\n[2]     Causal inference using invariant prediction: identification and confidence intervals\n[3]     Unsupervised Domain Adaptation by Backpropagation\n[4]     Domain Adaptation with Conditional Distribution Matching and Generalized Label Shift\n[5]     On Learning Invariant Representation for Domain Adaptation\n[6]     Inherent Tradeoffs in Learning Fair Representations\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "RiTpWcDI1QY", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\nThis paper addresses the problem of distributional shift in transfer learning from multiple training domains. The authors propose Risk Extrapolation (REx), which is a novel approach for out-of-distribution generalization when the new test domain for which we do not even have the covariate matrix. Thorough empirical experiments show that REx significantly outperforms state-of-the-art.\n\nPros:\n- This is a highly quality paper with strong theoretical and empirical results. \n- The paper is clearly written and easy to understand.\n- Based on the thorough literature review, this idea of this work is original.\n- The results of this work are highly significant and of interest to the domain adaptation and transfer learning community.\n\nCons:\n- Although I understand the page limit, most of the major parts of the paper (especially the theoretical aspects) can be found in the appendix. As a person who is not very familiar with the literature on distributional shift from multiple domains, I did appreciate having this thorough overview; however, the detailed discussions of the contributions of the paper might be overlooked if (when) located in the appendix. Specifically, there is  only half a page on introducing the proposed methods for REx  in Sec. 3.1.\n\nMinor comment(s):\n- Reference “Peter Bühlmann. Invariance, causality and robustness, 2018a.” is duplicated.\n\n############################################################\n\nPost-Rebuttal:\n\nAfter reviewing the concerns raised by the other reviewers, and the responses provided by the authors, I have decided to adjust my scores.\n\nMoreover, I was disappointed that the authors did not use the extra one page to move some material from the appendix to the main text in order to elaborate on the proposed method.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "76OUF4kNz-1", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Overview:\n\nThe authors propose Risk Extrapolation (Rex) which is an invariance-based approach to domain generalization. The main idea is to go from worst-casing over a convex combination of domains to an affine set of domains (MM-Rex) or to penalize the variance (V-Rex). Evaluations compare to IRM and ERM and show uniform improvements.\n\nPositives:\n\nThe paper has a well written and motivated introduction, and there's substantial added expository material in the supplement. I would maybe tone down on the amount of bolded text. \n\nThe arguments are pretty well-thought-out, including some discussion of the differences between causal recovery, invariant prediction, and domain generalization. \n\nThe method itself seems simple (in a good way), though I have some questions below. I also wonder if the variance version of the objective can be tied to DRO based approaches via the fact that DRO on a chi-squared perturbation ball is equivalent to variance regularization of the risk. \n\nThe method seems to work well overall compared to IRM, and although it only matches ERM in the domain generalization benchmark, I think this is still a decent result given that most methods underperformed ERM.\n\nNegatives:\n\nThis is possibly a comment that refers to a paper that's too recent, so not addressing this comment won't affect my rating of the paper, but it seems worthwhile from a scientific perspective to address how recent negative theory results (Rosenfeld, Ravikumar, Risteski 2020) about IRM (and REx) affect the intro framing. In particular, I'd like to see the claim in the intro about how REx can extrapolate can be reconciled with the claims in the other paper that IRM and REx succeed under the same conditions as ERM. \n\nAs a note: I also don't think Williamson and Menon suggest the use of variance of risks. They explicitly state that the R_{sd} risk aggregator is not a fairness risk measure. The risks fulfilling the axioms in that paper are coherent risk measures, and I believe they will all fall under the category of RI methods through duality arguments.\n\nI can believe that minimizing the max risk is a reasonable thing to do, but it seems like an added leap of faith to want to actively increase the risk of the other domains since minimax risk would naturally equalize the different domains (if it's possible to achieve equal risk). It would be nice to see a clearer justification for this behavior.\n\nHaving looked at the supplement, there's substantial and good material there, and I would maybe suggest that the authors cut something like figure 3 and algorithm 1 to bring back some more intuition and motivation about REx, maybe some result from section E.\n\nHaving looked at the Thm 1 result, wouldn't 3 interventions on every variable also recover the true beta with ERM? This result also depends very heavily on fixed noise across environments. I do think this is a neat result though, and it might be useful to use this to give additional intuition about REx in the main text.\n\nDoes E.1.2 require pointwise homoskedasticity ? that seems wildly strong for a general result... In classification, I think this means that the map from x → y has to be deterministic and in general, for any log-probably loss, this says that the entropies have to be identical everywhere. \n\nMinor: \n\nMight be worth re-stating what the methods are in the experiment section (RI, for example, is defined pretty early on).\n\nDRO is usually expanded to distributionally robust optimization, not domain robust optimization.\n\nFigure 3 seems unnecessary.\n\nIs Epsilon_j in the statement of Theorem 1 a typo? is there a typo in the subscript of beta_j = beta_{0,j} for all j?", "extra_scores": {}, "extra_text": {}}, {"review_id": "5kN2X-tfv3q", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper studies the out-of-distribution (OOD) generalisation problem via risk extrapolation (REx). The authors propose two methods, MM-REx and V-REx, and empirically show that REx can recover the causal mechanisms on Colored MNIST, while also providing some robustness to covariate shift. The authors deal with the relationship between robustness, invariance and causality carefully, and provide experimental evidence beyond Colored MNIST. \n\nI vote for weakly rejecting. This paper reviews various previous work but does not provide a clear comparison from them.\n\nI have some comments and questions as follows:\n1. Contribution 1 and Table 1 state that REx is suitable for invariant prediction. In the experiments, the authors also take IRM as the competitor. So I think the main objective of REx is the invariant prediction, rather than covariate shift. However, the authors emphasize that REx can deal with covariate shift. Please explain more on how REx discovers the invariant prediction.\n2. When considering the invariant prediction, we aim to discover a stable conditional distribution $P(Y|X)$ or an invariant conditional mean $E[Y|X].$ However, the covariate shift refers to the changes in the distribution of X. What is the novelty of assuming covariate shift here? Please also figure out the importance of considering covariate shift under the invariant prediction.\n3. What is the expression (1) in Page 1? If (1) is the risk function of the OOD generalization problem, $\\mathcal{F}$ is unseen and should be the same throughout this paper. At the end of Page 1, you say: \"Our method minimax Risk Extrapolation (MM-REx) is an extension of DRO where $\\mathcal{F}$ instead contains affine combinations of training risks, see Figure 1.\" Does MM-REx solve a different OOD problem? Please figure out the definitions of $\\mathcal{F}$ of  DRO, MM-REx and IRM respectively.\n4. Please explain Contribution 3. In general, the equality of risks is not a sufficient condition of the causality. \n5. In the VLCS and PACS experiments, the evaluation is incorrect. In Section 4.3, the task is to train on three domains and generalize to the fourth one at test time. However, this test accuracy is not worst-case performance.  According to (1), the problem is to generalize to all four domains at test time and to find out the worst domain. Then Table 3 should report the average of the worst-domain accuracy.\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper is proposing Risk Extrapolation (REX) as a domain generalization algorithm. Authors extends the distributionally robust learning to affine mixture of distributions from convex mixture. Authors later uses variances instead of this extension and demonstrate various empirical and theoretical properties. The paper is reviewed by four expert reviewers and the reviewers did not reach to a consensus. Hence, I also read the paper in detailed and reviewed it. In summary, reviewers argue the following:

- R#2: Main argument is the lack of justification of the claim "Rex could deal with both covariate and concept shift together". Authors try to address this in their response. Moreover, reviewer also argues in the private discussion that manuscript is not updated and authors did not address any of the issues during the discussion period.
- R#3: Argues that (similar to R#2), dealing with covariate shift is not explained properly. Reviewer is not persuaded that REX results in invariant prediction.
- R#1 and R#4: Largely positive about the paper. In the mean time, argue that organization of the paper is lacking and some of the material in the supplement is relevant and should be moved to the main text. R#1 decreases their score due to the lack of re-organization during the discussion.

The value of the paper is clear to me, the joint treatment of minimax perspective, domain generalization and invariances is definitely interesting and valuable. Hence, the paper has merit to be published. However, the presentation is lacking significantly.  The main contribution of the paper lies in Table 1 but the invariant prediction property is not justified at all in the main text. Hence, Table 1 is not justified properly. Authors discuss Thm 1&2 in their response but they both are in the supplement. From reading only the main text, confusion of the reviewers are well justified. ICLR guidelines clearly states that "...Note that reviewers are encouraged, but not required to review supplementary material during the review process..." It is authors' responsibility to make the main paper self contained. Even more worrisome is the fact that authors dismiss this concern in their response to R#1 which eventually leads to R#1 decreasing their score. Hence, I decided to reject the paper since the presentation is subpar and authors did not persuaded reviewers that they can fix this presentation issue by the camera-ready deadline. On the other hand, I think the paper can be really influential if it was written clearly. I suggest authors to revise the claims more precisely, extended the discussion on the claims and move the theorems to the main paper.
Decision: Reject