A Unified Causal View of Domain Invariant Representation Learning
Authors: Zihao Wang
Victor Veitch
Venue: Submitted to ICLR 2023 (2023)

Abstract: Machine learning methods can be unreliable when deployed in domains that differ from the domains on which they were trained. One intuitive approach for addressing this is to learn representations of data that are domain-invariant in the sense that they preserve data structure that is stable across domains, but throw out spuriously-varying parts. There are many approaches aimed at this kind of representation-learning, including methods based on data augmentation, distributional invariances, and risk invariance. Unfortunately, it is often unclear when a given method actually learns domain-invariant structure, and whether learning domain-invariant structure actually yields robust models. The key issue is that, in general, it's unclear how to formalize ``domain-invariant structure''. The purpose of this paper is to study these questions in the context of a particular natural domain shift notion that admits a natural formal notion of domain invariance.
This notion is a formalization of the idea that causal relationships are invariant, but non-causal relationships (e.g., due to confounding) may vary. We find that whether a given method learns domain-invariant structure, and whether this leads to robust prediction, both depend critically on the true underlying causal structure of the data.

Under review as a conference paper at ICLR 2023
A UNIFIED CAUSAL VIEW OF DOMAIN INVARIANT SU- PERVISED REPRESENTATION LEARNING
Anonymous authors Paper under double-blind review
ABSTRACT
Machine learning methods can be unreliable when deployed in domains that differ from the domains on which they were trained. One intuitive approach for addressing this is to learn representations of data that are domain-invariant in the sense that they preserve data structure that stable across domains, but throw out spuriously- varying parts. There are many approaches aimed at this kind of representation- learning, including methods based on data augmentation, distributional invariances, and risk invariance. Unfortunately, is often unclear when given method actually learns domain- invariant structure, and whether learning domain-invarian structure actually yields robust models. The key issue is that, general, it's unclear how to formalize "domain-invarian structure". The purpose of this paper is to study these questions in the context of particular natural domain shift notion that admits a natural formal notion of domain invariance. This notion is formalization of the idea that causal relationships are invariant, but non-causal relationships (e.g., due to confounding) may vary. We find that whether a given method learns -nvariint structure, and whether this leads to robust prediction, both depend critically on the true underlying causal structure of the data.
INTRODUCTION
Machine learning methods could have unreliable performance at the presence of domain shift (Shimodaira] 2000 Ciatnonero-Canddee et al.] 2008). a structural mismatch between the training domain(s) and the deployed domain(s). A variety of techniques have been proposed mitigate domain shift problems. One popular class of approach- which we'll focus on in this paper-is to try to learn representation function of the data that is in some sense "invariant" across domains. Informally, the aim of such methods is to find a representation that captures the part of the data structure that is the "same" in all domains while discarding the part that varies across domains. It then seems intuitive that predictor trained on top of such a representation would have stable performance even in new domains. Despite the intuitive appeal, there are fundamental open questions: when does a given method actually succeed at learning the part of the data that is invariant across domains? When does learning a domain invariant representation actually lead to robust out-of-domain predictions?
There are many methods aimed at domain -nnvariant representation learning. When applied to broad ranges of real-world domain-shift benchmarks, there is no single dominant approach, and an attack that works well in one context is often worse than vanilla empirical risk minimization in another (i.e.. just ignoring the domain-shift problem). We'll study three broad classes of method:
Data augmentation Each example is perturbed in some way and we learn a representation that is the same for all perturbed versions. E.g., if L(X) is a small rotation of an image X, then
(xX) X
(((()) (Krizhevsky et al. 2012 Hendrycks al. 2019 Cubuk et al.) 2019 Xie
ct al.
2020
Wei & Zoul 2019: Paschali al. 2019 Hariharan & Girshick] 2017 Sennrich
et al.
2015
Kobayashi
2018
Nie et al. 2020).
Distributional invariance We learn a representation so that some distribution involving ((X) is constant in all domains. There three such distributional invariances that can be required
to hold for all domains e,
marginal invariance: ()) pd (o(X) Muandet et al. 2013 Ganin et al. 2016
Albuquerque al. 2020 Li et al. 2018a Sun al. 2017 Sun & Saenko 2016: Matsuura

---

Under review as a conference paper at ICLR 2023
& Harada 2020):
conditional invariance: When Y is a label of interest, (e X
Y)
Pe'
(0) X
ct al.] 2018b Long et al. 2018 Tachet des Combes et al.
2020
Goel et al.
2020
sufficiency: Po(Y o ( (Y) ((X)) (Peters et al.
2016
Rojas-Carully et al. 2018
Wald et al. (021))
Risk minimizer invariance For supervised learning, we learn representation ((X) so that there
is a fixed (domain beamainnindeeendent predictor w² on top of 6(X) that minimizes risk in all
domains (Arjovsky et al., 2019 Lu et al. 2021 Ahuja et al. 2020 Krueger et al. 202 et 2021).
Bae
In each case, the aim is to learn representation that throws away information that varies 'spuriously' across domains while preserving information that is reliably useful for downstream tasks. However, the notion of what is thrown away is substantively different across all the approaches and it is unclear which, if any, is appropriate for any particular problem.
The principle challenge to answering our motivating questions is that it's unclear in general how to formalize the idea of "part of the data that is invariant across domains". To make progress, it is necessary to specify the manner in which different domains are related. In particular, we require an assumption that is both reasonable for real-world domain shifts and that precisely specifies what structure is invariant across domains. In many problems, it is natural assume that causal deennnsippsdddeemiieedddd by the unobserved real real-world dynamics underpinning the data-should be the same in all domains. We'll use an assumption of this form; the first task is to translate into a concrete notion of domain S
Specializing to supervised learning with label Y and covariates X, we proceed as follows. The covariates X are caused by some (unknown) factors of variation. These factors of variation are also dependent with Y. For some factors of variation, jointly denoted as Z, the relationship between Y and Z is spurious in the sense that Y and Z are dependent due to an unobserved common cause. The distribution of this unobserved common cause may change across environments, which in turns means the relationship between Y and Z can shift. However, the structural causal relationships between variables will be the same in all environments- e.g., P(X pa(X)) is invariant, where pa(X) denotes the (possibly unobserved) causal parents of X. We call a family of domains with this structure Causally Invariant with Spurious Associations (CISA).
Concretely, consider the problem of classifying images X as either Camel Cow Y. In training, the presence of sand in the image background Z is strongly associated with camels. But, we may deploy our new classifier in an island domain where cows are frequently on crahhiinchanningg the Z-Y association. Nevertheless, the causal relationships between the factors of variation- Z, and others such as camera type or time of day- and the image X remain invariant.
In this example, a natural formalization of "domain invariant part" of the image is the part that does not changes if grass Z is added or removed from the background; an invariant representation learning method should learn a ((X) that throws away such information. The aim of CISA is to be a reasonably broad notion of domain-shiff that also allows us to formalize this intuitive notion of comain-invariiun
Namely, under CISA, we define the domain-invarian part of X to be the part that is not causally affected by the spurious factors of variation Z. Accordingly. a representation is domain invariant if ((X) is not causally affected by Z. We say representation with this property is counterfactually invariant to spurious factors (CF-invariant for short).
We now return the motivating questions: when does given method actually succeed at learning the part of the data that is invariant across domains? And, when does learning a domain invariant representation actually lead more robust out-of-domain predictions? In the context of CISA- compatible domain shifts, we can answer the first question by determining the conditions under which each approach learns a CF-invariant representation, and the second question by studying the relationship between CF-invariance and domain shift performance.
Informally, the technical contributions of the paper are:
1. Formalization of CISA and Counterfactually Invariant Representation Learning.
'This kind of causal-invariance assumption is already used in the domain-shift literature, though the formal domain shift notion we'll Il here differs from previous approaches, see see Section 2
2

---

Under review as a conference paper at ICLR 2023
2. A characterization of the causal (real-world) conditions where each domain invariant repre- sentation learning method yields in-invariant representations. This turns out to yield both a hierarchy of methods, and guidance for which methods are appropriate for given problem.
3. Results relating invariant learning to robust learning. We'll find that, in general, there is not a simple relationship. For instance, learning a representation via risk minimizer invariance (an intuitive formalization of robustness) is impossible in general. However, we'll also find that there are (common) cases where there is a close connection. These results highlight the fundamental role of the underlying causal structure of the data.
RELATED WORK
Causal invariance in domain generalization Several works (Peters et al., 2016, Rojas-Carulla
ctt al., 2018 Arjovsky et al. 2019 Lu et al. 202I) specify domain shifts with causal models. In these pa(Y)). However, this rules out, c.g., cases where
frameworks, the invariant predictor learns P(Y
Y causes X (Schölkopf al., 2012 Lipton et al.
2018). CISA allows learning with more general structures, which we'l II see is helpful in understanding the representation learning methods.
Domain generalization methods There are many methods for domain generalization: we give a categorization in the introduction. There have been a number of empirically-orienn
surveys testing domain generalization methods in natural settings (Koh et al. [222] Wiles et al., 2021 Gulrajani & Lopez-Paz, 2020). These find that method consistently beats HRM, but many methods work well in at least some situations. Our aim here is to give theoretical insight into when each might work.
Robust methods The CISA framework is reasonable for many weal-world problems, but certainly not all. There are other notions of domain shifts and differently motivated methods that do not fit under this framework. For example, many works (e.g.. Sagawa et al. 2019 Liu et al., 2021]: David cet al. 2022; Eastwood et al. 2022) assume the testing domains are not too different from the training domains (e.g., test samples are drawn from the mixture of training distributions). Thesse methods are complementary to the invariant representation learning approaches we study here; e.g., they would all fail in the -oo-iit-nnviromments experiment in Appendix
CAUSAL INVARIANCE WITH SPURIOUS ASSOCIATIONS
Setup In domain generalization problems, we have training and test data from several related domains. The goal is to learn predictor using data from training domains, and apply it to unseen test domains. The data comes in the form of (X, Y, E) where E X' is the input, Y E y the label and E E is the domain index. We abstract each domain probability distribution Pe, where X,, Y, E, = ~
lid
Pe. At training time, we have access to data from a finite set of domains Erain C E.
that ((X) x is in some sense invariant across the domains. Many methods (e.g.. Arjovsky et al., (219:; Ganin et al. 2016 Goel et al.) (0200) frame their learning procedures this way, though there are many distinct formal notions of "iivariance", The first task to establish canonical notion of invariance.
Our goal is to learn representation
Causal Invariance with Spurious Associations All methods for handling domain shifts must specify the structure that is preserved across domains, and the ways in which they are allowed to PI vary Here, we'll rely on a particular variant of the assumption that causal relationships are held fixed, while non-causal relationships may vary. To formalize this, assume that X is caused by some unobserved factors of variation, and give a notion of what it means for these factors to be spuriously associated with Y:
We say a latent cause of X is a spurious factor of variation if it is not a cause of Y and there is some latent confounder that affects (directly or indirectly) both and Y. Call the set of all such causes the spurious factors of variations.
Definition 1.
Figure I shows examples of causal structures for prediction problems involving spurious factors of variation (Z). X is divided into two parts: X_1 denotes the part of X not affected by the spurious
Withoout such an assumption, the test domain could be chosen adversarially to have support only on points where the training-domain predictor makes mistakes.
3

---

Under review as a conference paper at ICLR 2023
Figure 1: Examples of CISA-compatible causal structures. In each case, all causal relationships (solid arrows) are invariant across domains. Only the distribution of the unobserved confounder U/ can vary, changing the induced association between Y and the unobserved spurious factors Z.
factors Z, and Xg the remaining part affected by Z. We do not assume the division of X is known priori (indeed, Z is not even known!). Also note that X may affect X,, so these parts need not be independent. Formally, we use the notion of counterfactual invariance from Veitch al. (202T).
Definition Let & be a function from X' to H and Z the spurious factors. We say ó is counteffuctually invariant spurious factors (abbreviated CF-invariant), if ((X(2)) (2'()) a.e., Vz, z' Z. Here, (z) is potential outcomes notation, denoting the we would have seen had Z been set to 2.
Then, we can formalize the part of X that's not affected by Z:
Definition 3. X- is a X-measurable variable such that, for all measurable functions f is CF- invariant iff f(X) X
Now, we characterize how the relationship between Z and Y changes across domains with a definition of unobserved confounder.
Definition 4. We say an unobserved variable an unobserved confounder if it is a common cause of Z and Y, does not confound the relationship between X and Y, and is caused by E alone. Call the set of all such confounders unobserved confounders.
We introduce unobserved confounders U that induces the spurious correlations between Xg and Y. Now put together all the ingredients and formally specify how the domains are related:
Definition 5. We say set of domains & are Causally Invariant with Spurious Associations (CISA), there are unobserved spurious factors of variation Z, unobserved confounders U, and only U is directly caused by domain E, that Pe (X, Y) Po(X, (, (& dU) Ve E E, where Po is a fixed distribution determined by the underlying dynamics the system and P.(U) is a domain-specific distribution of the unobserved confounder.
Assumptions in domain shifts should about the real-world processes that govern the domains from which the data is collected. This allows us to assess whether assumption is reasonable for given real-world situation. CISA is such an assumption; the spurious factors and unobserved confounders U are unknown and unobserved, but are assumed to correspond real-world properties. This assumption about the real-world structure implies an assumption about the distributions Pe (x, Y) in each domain. However, in general, the variables Z and U need not be identifiable from (i.e., uniquely determined by) the observational data. This point is only conceptually important, and does not affect the technical development in the following.
CF-invariant representation learning We now return to the question of what the right notion of domain-invann invariint representation is. In the case of CISA domains, there a canonical notion for the part of X that has Nomain-invarian relationship with Y. Namely. X the part of X that is not affected by the spurious factors of variation. Accordingly. the goal of learning in this context is to find a representation á such that &i X) is counterfactually -antuallll-ivviinntt to Z (or, equivalently, ((X) X= measurable). Of course, we could satisfy this condition by simply throwing away all of the information in X (e.g., ((X) = O everywhere). So, we further look for the counterfactuullly-in -nvvarraaly representation that preserves the most predictive power for the label Y. Let deriiw denote the set of CF-in CF-invariant representations for CISA domains E. Then, we take our representation learning objective to be:
representation
min O. H.W
Erecm [L(Y, (w (X))
subject to & E (E)
Such variable exists under weak conditions; c.g., Z discrete Veitch al. 2021
4

---

Under review as a conference paper at ICLR 2023
Here, the predictor and loss function L capture the sense in which (XX) should be predictive of Y.
The challenge here is that the spurious factors of variation are unknown and unobserved, so identifying the set of caunterfactuall--nvvriiin representations is difficult. Now turn to understanding various approaches to domain adaptation methods as achievable relaxations of this ideal objective.
CAUSAL VIEW ON DOMAIN INVARIANT REPRESENTATION LEARNING
4.1
DATA AUGMENTATION
Data augmentation standard technique in machine learning pipelines, and has been shown to (sometimes) help when faced with domain shifts (Wiles et al. 2021p. Our goal now is to understand when and why data augmentation might enable domain-invariant representation learning. The basic technique first applies pre-determined "label- aabel-preserving" transformations & to original features X to generate artificial data (X)). There are two ways this transformed data can be used. The first option simply add the transformed data as extra data to standard learning procedure. Alternatively, we might pass pairs (Xx) to our learning procedure, and directly enforce some condition that (XX) x ((t(X)) (Garg al. 2019 Von Kügelgen al. 2021). In both case, the natural questions are: What transformations are "label pressrrres ? And, when do these techniques help with domain generalization?
We first formalize a notion of of transformation t(X) of X being equivalent changing cause of X and then propagating this change through. For example, suppose a particular transformation rotates the input images by 30 degrees, and W is the factor of variation corresponding
for CISA domains. The key idea is that we can think
the angle away from vertical. Then, can understand the action of as t(X (w)) = X X( + 30), where we again use the potential outcomes notation for counterfactuals. With this idea hand, we see that transformation label-preseriing in CISA domains if it is equivalent to change that affects only spurious factors of variation. That is, label- preberiig transformations cannot affect X Otherwise, the transformation may change the invariant relationship with Y; replacing the background of cows with "sand" with "grass" doesn't t change animal type; but distorting the part corresponding to the "cow" object may.
Definition 6. We say a data transformation ::A' + X' is label-preserving for CISA domains & if, for each (2) there is 2/ that t(X(:)) X(2').a.c. a.c..
Label preserving transformations leave the CISA invariant relationships (between X and Y) alone, but can change the relationship between Y and the spurious factors of variation Intuitively. if we have a 'large enough' set of such transformations, they can destroy the relationship between Y and Z that exists the training data. So, we might expect that training with such data augmentation will automatically yield optimal counterfactualll-innaan
predictor.
This is nearly correct, with the caveat that things can go wrong (for the naive training approach) there is part of X causally related to both Z and Y. That is, if there is a part of X that relies on the interaction between Z and Y. We follow Veitch et al. (202I) h formalizing how to rule out this case:
Definition 7. The spurious factors of variation Z are purely spurious if Y 1 X|X 7
That is, without the unstable correlation (removed by conditioning on Z), X is sufficient for Y.We state the main result connecting data augmentation and domain-invariance:
Theorem 8. For a CISA domain, if the set of transformations T satisfies label-preserving and enumerates all potential outcomes of Z, then
1. If the model is trained to minimize risk on augmented data, and Z is purely spurious, or
2. If the model is trained minimize risk original data, with hard consistency regularization (i.e. enforcing ((X) ((((X)) Vt E T
Then we recover the CF-invariant predictor that minimizes risk on original data.
Thus for CISA domains, ideal data augmentation (i.e. enumerating label aabel-preserving transformations) will exactly leam CF- i-vinarrinnt representations. Moreover, this holds irrespective of what the true underlying causal structure is. Accordingly, such data augmentation would be the gold standard for Domain-invarian representation learning.
5

---

Under review as a conference paper at ICLR 2023
X
Y
(a) anti-causal
(b) confounded-outene outcome
(c) ceafounded-descenddee
Figure 2: Every CISA compatible set of domains obeys one of the causal DAGs defined as follows:
The DAG must match one of the three templates, where the black edges must be included, and the non-edges must he excluded. Orange edges may he included or excluded. In the case of Figure 2c at least one the two dashed black arrows must included. (Note that in Figure 2a and Figure 2b the edges between X and Y are typically included, as otherwise there is no part of X that has a non-trivial and stable relalationship with Y.)
However, practice, it can be hard to satisfy the idealized conditions. Applying predefined transfor- mations without thinking about the specific applications can lead to violations of the label preserving condition. For example in bird classification task, changing color may really change the bird species (this is called manifold intrusion in Guo et al. () Further, heuristic transformations often cannot enumerate all potential outcomes of (().. Indeed, deep models can sometimes memorize predeter- mined transformations and fail to generalize with previously unseen transformations (Vasiljevic et al.)
2016 Geirhos et al. 2018).
Considering the limitations of using handerafted transformations, natural idea is replace them with transformations learned from data. However, in practice, Z is unknown and we only observe the data domains E. Then, learning transformations must rely either on detailed structural knowledge of the problem (Robey et al., [2021 e.g..), or on some distributional relationship between E, Y and X and t(X) (e.g.. G.el al., 2020). Since t(X) ) is only used for the representation learning, this is equivalent to learning based on some distributional criteria involving E, Y, and ((X)- the subject of the next section.
4.2 DISTRIBUTIONALLY INVARIANT LEARNING
Many domain-invariant representation learning methods work by enforcing some form of distribu- tional Srariance[Sectioon 1 There are three possible notions of distributional invariance, each of which appears in the literature. Early work tries to learn representation è that o(X) 1 E while ((X) still predictive Y (Muandet et al.) 2013 Ganin et al. 2016). The intuition is that predictions must be based features that cannot discriminate between training and test domains. Later work recognized that that P(Y may be unstable across domains, and instead aims at @(X) IL EJY (Li et al.] 2018b Long et al. 2018). Finally, some motivvtedd approaches have considered learning representations so that Y 1 Elo(X) (Peters et al. 2016 Wald et al. 2021).
predictive of Y and Y is dependent with E, these three indepen- dence statements-o(X J. E, d(X) 1 E\Y,Y JL Elø(X) the only possible distributional invariance relationships, which are usually mutually exclusive. The question now when, if ever, are each of these distributional invariances the right approach for domain insiiinninvarantt representation learning?
When the representation o(
CISA compatible causal structures To answer this, the first step is to characterize the set of causal structures that are compatible with CISA.
Theorem 9. Suppose a set of domains E share the common causal structure underlying 1o(X,Y2)) Z)) Then E satisfies CISA if and only if the corresponding causal DAG is one of the set given in Figure In particular, there are three families of allowed DAGs: anti-causal, and confondedd-descente
The idea here is that for set CISA domains, the causal structure the same in each domain. There only three kinds of causal structure, distinguished by whether Y causes X or X causes Y (anti- causal vs wtwpaeddoootooee//eeeeeeeeeeeede affects some (potential) causal descendant of Y (confounded-outccome vs confounded-desscedamn)
and by whether the confounding affects Y directly or just

---

Under review as a conference paper at ICLR 2023
Anti-causal Image classification can be naturally viewed as anti-causal problem. Various factors of variations such as lighting, background, angles, etc, and the object class Y, generate the image X. Some of the factors of variation Z are confounded with Y e-c... background and Y may be associated due to evolutionary pressures. The Cow/Camel on Grass/Sand example fits here
Confounded outcome The goal is to predict the helpfulness of a review. Each review receives a number of "helpful" votes Y, produced by site users. We the review's is text content X covariates. The data is collected for different types of products E. The model's 's performance drops significantly when deployed in new product type. We think that the general sentiment Z of the review, and the helpfulness has unstable relationship across E: e.g. for books, customers write very positive reviews which are often voted favorably; for electronics this relationship is reversed.
Confounded descendant Consider predicting unemployment rate Y from a variety of economic factors X. It's not clear which factors cause Y directly, and which are descendants of Y. The relationships among X, Y change under certain big events, say financial crisis or pandemics. We denote those events domains E, and take U E. U might affect Y's descendants jointly with Y (through intermediate variable Z), but Y is not changed. Or U might affect pa(Y) directly, which changes Y. Notice that in this case, the CF-invariant representation is also the representation that uses only pa(Y). Thus, for this causal structure, the counterfactually invariant notion matches the traditional causally-invarial representation learning desiderata (Peters et al. 2016 Arjovsky et al. 2019).
Distributionally invariant learning (partially) enforces conttrrfctualll-invaart representations. It turns out that the answer relies critically on the true causal structure of the problem:
We now return to the question of when distributional invariance
Theorem 10. If ó is coumterfuctuall--invriirn representation,
1. if the underlying causal graph is anti-causal, (xX) Y
2. if the underlying causal graph is confounded-outcome o(X) () 1 E;
3. if the underlying causal graph is confounded dessedddan, Y L EX)(X).
Remark 11. This theorem looks similar to Veitch et al. 2021 Thm. 3.2). This deceptive; here we observe the environment E, whereas they assumed observations of the spurious factors Z.
In words: each of the distributional invariances arises particular implication of CF-invariance. This suggests a learning strategy that relaxes CF-invariance. Namely, for a given problem, determine the distributional invariance implied by CF-invariance and then enforce that distributional invariance. That is, we learn representation according to:
min dE (Evan) 10
EPe.. [L(Y, (w (X))),
where (E) is the set of representations matching the causal structure of the domains.
Notice that the right distributional invariance is necessary but not sufficient for CF-invariance. There are two reasons: first, the causal condition of CF-invariance not implied by its distributional signature in Theorem 10] Second, the distributional invariance can only be measured on limited number of training domains, which may provide a limited constraints. this sense, distributional invariance is relaxation of CF-invariance; i.c., (r (E') Ş (Evinn)
Also notice that in general enforcing the wrong distributional invariance will directly contradict CF -Fviiaarianee, and increase dependency the spurious factors of is only a relaxation when we actually get the underlying causal structure correct!
variation. Distributional invariance
4.3 INVARIANT RISK MINIMIZATION
The Invariant Risk Minimization (IRM) paradigm (Arjossyy et al. 2019) aims to find representations that admit single predictor that has the optimal risk across all domains. That is, the set of IRM
*There is some controversy here since the label Y is often due human annotators LL al. 2021
7

---

Under review as a conference paper at ICLR 2023
representation is
PIRM (8) = {(b w st w E argmin Er, [L.(, (w ()))) Ve E EE)
(4.1)
Here, the question is: when, ever, does the IRM procedure yield a CF-invariant predictor?
Again the answer will turn out to depend on the underlying causal structure. IRM is justified in the case where X includes both causes and descendants of Y, and the invariant predictor should use only information in the parents of Y. As explained above, this coincides with CF-invariance-aan with distributional invariance-in the confounded-descenee relaxation of Y (X)( Instead of asking the distribution P (Y across domains for every h E H. we only require the risk minimizer under (" = h) to be invariant. So, a partial answer is that IRM is a relaxation of distributional invariance (and CF-invariance) when the underlying causal structure is confounded descendant.
case. Indeed, we can view IRM as
A) to be invariant
However, the situation is harsher in the case of other causal structures, For anti-causal and confounded- outcome problems, the typical case is (IRM(E) 0. For the anti-causal case, this is consequence of prior shift (Zhang al. 2013). We can change the optimal predictor just by changing the prior distribution Y without affecting the causal structure. Thus, there is no X measurabll predictor with invariant risk. The confounded-outcome case is even harder. In general, the generating function for Y could involve interaction between U and X. Then P(Y p can be quite different because of the shifting U. Thus the risk minimizers may be very different across domains.
Fortunately, for anti-causal problems there is natural generalization of IRM that partially enforces the distributional invariance. Similarly, we can relax the invariance of P (o()) = y) to the invariance of re-weighted risk. We define generalized-IRM (g-IRM) follows, which is essentially importance sampling based on Y (this simple procedure has causal implications (section 5)):
Definition 12. For anti-causal domains, we define the set of representations satisfying g-IRM as:
PY(Y) P.(Y)
& (E, == {6 3 w st w E argmin Ep.
L(Y, (w v
where Po(.) is chosen baseline distribution for Y.
If we know the label marginal in the test domain, we could use that as Po. Then the optimal predictor (on top of d) on training domains is optimal on the test domain. Next we can show the (generalized) IRM is a relaxation of the corresponding distributional invariance:
Theorem 13. Let & satissy CISA, then
1. & confonnded-descendee then Dr((E) C QIRM(EE)
2. E is anti-causal, then Qrr(E) C Ig. (E, Po) any chosen Po
Although this simple generalization works for the anti-causal case, there is no such easy fix for the confounded-outcome case. For confounded outcome, the idea of having the same risk minimizer across domains does not make sense without further assumptions. The general case has Y +- J(X, U, .1) where 7) is noise independent of X,U. For example when U = E can write Y <- fe(X, 7)) so there could be arbitrarily different relationships between X and Y in every domain. It may be possible to circumvent this by making structural assumptions on the form of f: e.g. there is an invariant risk minimizer in the case where the effect of U and X is additive (Veitch et al. 2021.
4.4 RELATIONSHIP OF METHODS
Data augmentation training is the gold standard for CF-invariant representation learning it's possible to enumerate all preserviesg transformations (Theorem 8). This is impossible in general as it requires direct manipulation of the spurious factors Z. Still, using augmentation with labee--res transformations (but not exhaustively) enforees relaxation of CF-invariance.
Distributional invariance relaxes CF-invariance if and only chosen to match the underlying causal structures (Theorem 10). This can be good option when full augmentation is not possible.
(generalized) IRM further relaxes distributional invariance for anti-causal and confounded- descendant problem, when it's chosen match the causal structure of the problem Theoreem 13). It
8

---

Under review as a conference paper at ICLR 2023
weakens the full independence criteria to use just the implication for single natural test statistic: the loss of the model. This allows for more efficient algorithms (Arjovsky et al. 2019).
5 INSIGHTS FOR ROBUST PREDICTION
representations is an intermediate step towards learning robust predictors. Here, robust means that the predictor should have good performance when deployed in previously unseen domain. We now discuss the implications of our domain-invariant representation learning results for robust prediction.
Often, learning
whether robust prediction is even possible depends on the underlying causal structure For manfounded-descendee problems, robust prediction is straightforward. Enforcing the correct distribu- tional invariance on representation function é leads to invariant risk in test domains for every value of (X)). If we enforce the weaker IRM requirement, the optimal predictor (on top of () on training domains optimal on test data. For anti-causal problems, there is no invariant risk minimizer in general because Pa PY)
change across domains indeed, even in the most simple anti-causal problem in [Appendix A. IRM fails. However, we can generalize IRM to match the causal structure: adjusting for prior shift using importance sampling, gIRM successfully recovers the optimal invariant predictor. For the confounded-outconm case, there is no notion of robust predictor without making some further structural assumptions.
no method works for all domain generalization problems Heuristic data augmentation enforces CF-invariance directly regardless of causal structures, but requires truly aabelppreserving transforma- tions and only solves the invariant representation problem if the transformation set affects all spurious factors. Distributionally-nnve invariant methods can work, but each approach is only valid if it matches the underlying causal structure of the problem. Enforcing the wrong distributional invariance harm performance (Veitch et al. 2021). Sometimes even seemingly innocent technique, such as im- portance sampling, has causal implications and could destroy invariant relalationship (appendix A.2). Similarly, (generalized) invariant risk minimizer can work for some types of problems, but only when matches the causal structure. This is consistent with the findings from various benchmark studies
actually
that there is no single method that does well in all tasks (Wiles et al. 2021 Koh et al. 2021 Gulrajani Lopez- 2020).
data augmentation helps in most cases generalization and can often help. This is true no matter the underlying causal structure of the problem. This matches empirical benchmarks where data augmentations usually help domain generalization performance, sometimes dramatically (Wiles et al., 2021] Koh et al. 2021 Gulrajani & Lopez
data augmentation won't hurt domain
Paz 2020 For example, WWles al. (2021) finds that simple augmentations used in Krizhevsky et al. 2012) generally improves performance when "augmentations approximate the true underlying generative model".
pick a method matching the true causal structure Many papers apply distributional invariance approaches with no regard the underlying causal structure of the problem. In particular, many tasks in benchmarks have the anti-causer structures, but the methods evaluated do not include those enforcing (X) Y (Koh et al. 2021). Tachet des Combes et al. (2020) find that methods enforcing Φ( X) IL E consistently improve over methods that enforce Φ X Fe-retrospectivelyy this is because they benchmark on anti-causal problems. Wiles et al. (021) finds that Icarned data augmentation (Goel al. 2020) consistently improves performance in deployment- this method can be viewed as enforcing @ X) E|Y; again, the benchmarks mostly have anti-causal structure.
CONCLUSION
We have introduced the CISA notion of domain shift and used it to study invariant representation learning methods. The key property of CISA is that it has cannonical notion of "domman-invarian"" representation: CF-invariance. We have seen that of data augmentation, distributional invariance learning, and risk invariant learning can be (sometimes) understood aelaeaaa relaxations of CF-invariance. The validity of the methods- as well as the connection between fundamentally the underlying causal structure of the data.
and domain- robust

---

Under review as a conference paper at ICLR 2023
REFERENCES
Kartik Ahuja, Karthikeyan Shanmugam, Kush Varshney, and Amit Dhurandhar. Invariant risk minimization games. In International Conference on Machine Learning, pp. 145-155. PMLR, 2020.
Isabela Albuquerque, João Monteiro, Mohammad Darvishi, Tiago H Falk, and Ioannis Mitliagkas. Adversarial target-invariant representation learning for domain generalization. 2020.
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint /rXiv: 2019.
Jun-Hyun Bae, Inchul Choi, and Minho Lee. Meta-learned invariant risk minimization. arXiv preprint arXiv:2103. 2021.
Eyal Ben-David, Nadav Oved, and Roi Reichart. Pada; Example-base based prompt learning for on-the-fly adaptation to unseen domains. Transactions of the Association for Computational Linguistics, 10: 413-433, 2022.
Ekin D Cubuk, Barret Zoph, Dandelion Mane, Vijay Vasudevan, and Quoc V Le. Autoaugment: Learning augmentation strategies from data. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 113-123, 2019.
Cian Fastwood, Alexander Robey, Shashank Singh, Julius von Kügelgen, Hamed Hassani, George Pappas, and Berhard Schölkopf. Probable domain generalization via quantile risk minimization. arXiv preprint arXiv:2207.008 0099444 2022.
Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky. Nomain-adversaee training of neural networks. The journal of machine learning research, 17(1):2096 2016.
Sahaj Garg, Vincent Perot, Nicole Limtiaco, Ankur Taly, Ed H Chi, and Alex Beutel. Counterfactual fairness in text classification through robustness. In Proceedings of the 2019 AAAI/ACM Conference on AI, Ethics, and Society, pp. 219-226, 2019.
Robert Geirhos, Carlos RM Temme, Jonas Rauber, Heiko H Schütt, Matthias Bethge, and Felix A Wichmann. Generalisation in humans and deep neural networks. Advances in neural information processing systems, 31, 2018.
Karan Goel, Albert Gu, Yixuan Li, and Christopher Ré. Model patching: Closing the subgroup performance gap with data augmentation. arXiv preprint UrXiv:2008.06777: 2020.
Ishaan Gulrajani and David Lopez-Paz. In search of lost domain generalization. arXiv preprint arXiv:2007.01434 2020.
Hongyu Guo. Yongyi Mao, and Richong Zhang. Mixup as locally linear out-of-manifold regulariza- tion. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 3714-3722 2019.
Bharath Hariharan and Ross Girshick. Low-shot visual recognition by shrinking and hallucinating features. In Proceedings of the IEEE International Conference on Computer Vision, pp. 3018-3027, 2017.
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. arXiv preprint arXiv: arXiv: /9/2.02781, 2019.
Pritish Kamath, Akilesh Tangella, Danica Sutherland, and Nathan Srebro. Does invariant risk minimization capture invariance? In International Conference on Artificial Intelligence and Statistics, pp. 4069-4077. PMLR, 2021.
Sosuke Kobayashi. Contextual augmentation: Data augmentation by words with paradigmatic relations. arXiv preprint arXiv: 1805.00000002000 2018.
10

---

Under review as a conference paper at ICLR 2023
Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Bal- subramami, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, et al. Wilds: A benchmark of in-the-wild distribution shifts. In International Conference Machine Learning, pp. 5637-5664. PMLR, 2021.
Alex Krizhevsky, lya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolu- tional neural networks. Advances in neural information processing systems, 25, 2012.
David Krueger, Ethan Caballero, Joern-Henrka Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville. Out-of-diitributtion generalization via risk extrapola- tion (rex). In International Conference on Machine Learning, pp- 5815-5826. PMLR, 2021.
Haoliang Li, Sinno Jialin Pan, Shiqi Wang, and Alex Kot. Domain generalization with adversarial feature learning. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 5400-5409, 2018a.
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 2018b.
Zachary Lipton, Yu-Xiang Wang, and Alexander Smola. Detecting and correcting for label shift with black box predictors. In International conference on machine learning, pp. 3122-3130. PMLR, 2018.
Evan Z Liu, Behzad Haghgoo, Annie S Chen, Aditi Raghunathan, Pang Wei Koh, Shiori Sagawa, Percy Liang, and Chelsea Finn. Just train twice: Improving group robustness without training group information. In International Conference on Machine Learning, pp. 6781 6792.. PMLR, 2021.
Mingsheng Long. Zhangjie Cao, Jianmin Wang. and Michael Jordan. Conditional adversarial domain adaptation. Advances in neural information processing systems, 31, 2018.
Chaochao Lu, Yuhuai Wu, Jose Miguel Hernández- Lobato, and Bernhard Schölkopf. Nonlinear invariant risk minimization: A causal approach. arXiv preprint arXiv:2102.12353 2021.
Toshihiko Matsuura and Tatsuya Harada. Domain generalization using a mixture of multiple latent domains. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pp. 11749-11756, 2020.
Krikamol Muandet, David Balduzzi, and Bernhard Schölkopf. Domain generalization via invariant feature representation. International Conference on Machine Learning, pp. 10-18. PMLR, 2013.
Yuyang Nie, Yuanhe Tian, Xiang Wan, Yan Song, and Bo Dai. Named entity recognition for social media texts with semantic augmentation. arXiv preprint arXiv:2010.15553 2020.
Magdalini Paschali, Walter Simson, Abhijit Guha Roy, Muhammad Ferjad Naeem, Rüdiger Göbl, Christian Wachinger, and Nassir Navab. Data augmentation with manifold exploring geometric transformations for increased performance and robustness. arXiv preprint arXiv:1901.044 0400000004420.... 2019.
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), 7845):9471 2016.
Joaquin Cannonero-Candlle Masashi Sugiyama, Anton Schwaighofer, and Neil D Lawrence. Dataset shift in machine learning. Mit Press, 2008
Alexander Robey, George J Pappas, and Hamed Hassani. Model-based domain generalization. Advances in Neural Information Processing Systems, 34:20210-20229 2021.
Mateo Rojas-Carulla Carulla, Bernhard Schölkopf, Richard Turner, and Jonas Peters. Invariant models for causal transfer learning. The Journal of Machine Learning Research, 19(1):1009-134) 2018.

---

Under review as a conference paper at ICLR 2023
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. arXiv preprint arXiv:1911.0873/1 2019.
Bernhard Schölkopf, Dominik Janzing, Jonas Peters, Eleni Sgouritsa, Kun Zhang, and Joris Mooij. 2012.
On causal and anticausal learning. arXiv preprini
Rico Sennrich, Barry Haddow, and Alexandra Birch. Improving neural machine translation models with monolingual data. arXiv preprint 1rXiv:1511.066 2015.
Hidetoshi Shimodaira. Improving predictive inference under covariate shift by weighting the log-
likelihood function. Journal of statistical planning and inference, 90(2):22-2244, 2000.
Baochen Sun and Kate Saenko. Deep coral: Correlation alignment for deep domain adaptation. In European conference on computer vision, pp. 443-450. Springer, 2016.
Baochen Sun, Jiashi Feng, and Kate Saenko. Correlation alignment for unsupervised domain adaptation. In Domain Adaptation Computer Vision Applications, pp. 153-171. Springer, 2017.
Remi Tachet des Combes, Han Zhao, Yu-Xiang Wang, and Geoffrey Gordon. Domain adaptation with conditional distribution matching and generalized label shift. Advances in Neural Information Processing Systems, 33:19276- 2020.
Igor Vasiljevic, Ayan Chakrabarti, and Gregory Shakhnarovich. Examining the impact of blur on recognition by convolutional networks. arXiv preprint 161iv:1611.05766 2016.
Victor Veitch, Alexander D) Amour, Steve Yadlowsky, and Jacob Eisenstein. Counterfactual invariance to spurious correlations: Why and how to pass stress tests. arXiv preprint arXiv:2 2021.
Julius Von Kügelgen, Yash Sharma, Luigi Gresele, Wieland Brendel, Bernhard Schölkopf, Michel Besserve, and Francesco Locatello. Self-supervised learning with data augmentations provably isolates content from style. Advances neural information processing systems, 34: 14:16451-16461 2021.
Yoav Wald, Amir Feder, Daniel Greenfeld, and Uri Shalit. On calibration and out-of-domain generalization. Advances in Neural Information Processing Systems, 34, 2021.
Jason Wei and Kai Zou. Eda: Easy data augmentation techniques for boosting performance on text classification tasks. arXiv preprint arXiv: 19X1v:1901.11196 2019.
Olivia Wiles, Sven Gowal, Florian Stimberg. Sylvestre Weiise-Rebufff, Ira Ktena, Taylan Cemgil, et al. A fine-gra mme-inained analysis on distribution shift. arXiv preprint 1rXiv:2110.1/328. 2021.
Qizhe Xie, Zihang Dai, Eduard Hovy, Thang Luong, and Quoc Le. Unsupervised data augmentation for consistency training. Advances in Neural Information Processing Systems, 33:6256-6268, 2020.
Kun Zhang, Bernhard Schölkopf, Krikamol Muandet, and Zhikun Wang. Domain adaptation under target and conditional shift. In International Conference on Machine Learning, pp. 819-827, PMLR, 2013.
12

---

Under review as a conference paper at ICLR 2023
A EXPERIMENTAL DEMONSTRATION
A key point in the above analysis is that the invariant representation learning method we use must match the underlying causal structure of the data. In In[Section 5] we discuss this point in the context of existing large-scale, real-world experiments. However, such demonstrations rely on out of domain performance of various methods, which depends on both the match to the true causal structure but also on lower level implementation issues e.g., hyperparameter tuning, overfitting, or optimization. This can make it somewhat difficult to draw precise conclusions.
To demonstrate the role of causal structure clearly we now study it in simple case: two bit environments from Kamath et al. (2021). ). [ This is a toy setting mimicing the well-known colored MNIST example used to demonstrate Invariant Risk Minimization Arrovsky et al. (2019). We create two domain shifts. The first is anti-causal and the second is soffounded-descendee
As predicted
from Theorem 10 second.
we find that IRM fails but gIRM works in the first case, and vice versa in the
In short, find that even in the simplest possible case, causal structure plays key role. In partieular, vanilla IRM fails totally on an apparently innocuous modification of the data generating process (prior shift), and this is readily fixed by modifying the method to match the correct causal structure. Further, the experiments demonstrate that importance sampling based on Y-an apparently innocent technique-actuieeee has causal implications and can destroy invariant relaltionship under certain causal structures (e.g. confounded descendaan in the second example).
A.1 TWO-BIT-ENVS (ANTI-CAUSAL)
For each domain E E, the data generating process is as follows:
Y f- Rad(Ye)
X1 + Y (a)
X2 Y (a.(
where Rad(T) is random variable taking value- -1 with probability and +1 with probability 1 1-
This a simplification of the ColoredMNIST problem (Arjovsky et al. (019)): we know X1 (corresponds to the digit shape) and Y have invariant relationship: P(X; Y) = -1 a. The
correlation between X, (corresponds to color) and Y is spurious, as P(X, = Y) = 1- - 3, that varies across domains. The label imbalance across domains is due to prior-shift: P(Y -1) = % We observe 4 training domains and predict on test domain. We use a = 0.25, and set 3e = 0.1, 0.2, 0.15, 0.05 in the training domains respectively, so that using the spurious correlation could get better in-domain performance. However, in the test domain the spurious correlation is flipped: we use 3_- = 0.9 so the out-of-domain performance would be very bad if X2 is used, Finally, we use Te = 0.9, 0.1, 0.7, 0.3 in training domains to create shifr-shift. In the test domain the label is balanced (% = 0.5). These 4 domains constitute E. The goal to find predictor f E [f+1,-11 1 + R} (prediction is == [g:(((a:) for data (I., w)). The optimal 75% test accuracy is obtained when satisfies ((1,,) > 0 and f(-1,,) S.
The original two-bit-envs problem has % = 0.5 across domains, and IRM obtains one optimal predictor. However, it fails in this modified problem with prior-shift TTable I). To understand and fix its failure, we can find causal interpretation that both explains the data and is CISA-compatile: introduce U, Z and set Z + U: < E as shown in Figure 3a Then this domain shift problem falls under anti-causal subtype, and CF-invariant predictor should only rely on X,. IRM does not enforce the right invariance and fails to remove spurious X2 as a result. Instead, we should use gIRM to (partially) enforce CF-invariance. Indeed, gIRM successfully forces the model to discard X₂ and obtain the optimal test accuracy (Table T}. Note that in the original two-bit-envs problem without prior-shift, both X1= + Y (ad(a) and Y + X1 (a) can explain the data. Therefore we we
*Code would be made available upon accepted conference paper.
'Our experiment uses IRMv1 (and analogously gIRMv1). which is shown to fail with some choices of a Kamath et al. (202I) because of the relaxation from IRM to IRMv1. We avoid those choices of a so that we can focus on the legel question instead of heing distracted by the fragility of IRMv1. Below we use IRM and IRMv1 (also gIRM and M exchangebly.
13

---

Under review as a conference paper at ICLR 2023
X1
X2
Z
X2
E
(a) anti-causal
(b) confoundedddeescendamt
Figure 3: Causal graphs that are both CISA-compatible, and explain the data generating process (for (lor Appendix A. [.ndddppppddii respeecccl. From original data generating process, we introduced extra variables 0 and set Z U/ and E. Table I: In two-bit (anti-causal) experiment, TRMv1 predictor fails to discard spurious features because of prior-shift in Y. However, after modifying the method to match the underlying causal structure (using gIRMvl), we recovers CF-invariant predictor that obtains optimal test accuracy. The results are under cross-entropy loss (similar for for squared loss).
FRRM1l
ggRRMv1
X2
+1
X2
=
X2
+1
X2
=-=
X1 =
+1
2.53
-0.93
1.16
1.08
X1
1
-0.08
-3.19
11.11
-1.03
can interpret the data generating process as either anti-causal or confounded-descendeee So both IRM and glRM (partially) enforce CF-invariance in fact (ER =+ (g.IR((E) so the resulting predictor is the same.
A.2 TWO-BIT-ENVS (CONF-DESC)
Our implementation of enforcing glRM regularization is equivalent to: first, perform importance Po(y),
sampling with weight we (Ex1 y) = P. (v) next, enforce IRM regularization. Similarly, it's
common practice to perform importance sampling based on label Y when it's imbalanced. However, as we shall show in this example, importance sampling can remove invariant features under certain causal structures.
In this example, the data generating process as follows: for each domain e E E,
Xi Rad(Y.)
Y + X, Rad(a)
X₂ - Y (ad)
Compared to the previous example, we not change P(X: = Y) and P(X2 = Y). The only change is how the label imbalance is created: through covariate shift in X2. We use the same parameters for a, Se, Ye-
"There are two ways to implement glRM. The first way is apply importance sampling to the regularization term only; the second way is to apply it to both the loss term and regularization term. In these two examples, the implementations give the same result. Thus better illustrate point importance sampling, we study the second implementation.
Table 2: In two-bit leitouudddd-dsccnddaae experiment, directly applying IRMv1 recovers CF- invariant predictor that obtains optimal test accuracy. On the contrary, applying importance sampling destroy the invariant relationship in data result, gIRMvl learns only the trivial invariant predictor. The results are under cross-entropy loss (similar for for squared loss).
IRMvI
X,
X2
X2
X2
X1 -
1.1
1.1
0
0
X1
-1.1
14

---

Under review as a conference paper at ICLR 2023
Again (Y) is different across domains, but this time gIRM forces the model to always predict O (Table 2) To understand why, we find a CISA-compatibl DAG that explains the data as shown in Figure 3b (similarly introduce U, Z and set Z U;U < E). This is IRM enforces the right invariance whereas gIRM enforces the wrong one.
and
To understand how the importance sampling destroys the invariant relationship between X1 and Y, we look at the target distribution after reweighting (call it Qe). Since the weighting function is
(a 1/ P=(3)
Po(as) and that w. (#1, y) Q.(x1.9) we have (e y) = (2 v) Pu(y)
P.(y)" Observe
that the probability (X, Y) = g(tre) where the [0, suppooote function (treat constant and assume a > 0.5) satisfies the following:
1. 9(7) = g(1 g(1->) so g is symmetric around 0.5.
2. g strictly increases on [0, 0.5) and strictly decreases on [0.5, 1]
3. ((0) = (1) = 0.5, and g(0.5) 5) = a, so decreases as % deviates from 0.5
Thus 0.5 < Q Y) Q2 Y) < Q2(X, = =Y) Qx(X; Y) < Therefore, importance sampling not only weakens the relalationship between X₁ and Y, but also makese unstable! As result, enforcing IRM the resampled distribution finds no non-trivial invariant predictors.
B PROOFS
Theorem 8. For a CISA domain, if the set of transformations T satisfies label-preserving and enumerates all potential outcomes of Z. then
1. If the model is trained to minimize risk on augmented data, and Z is purely spurious, or
2. If the model trained minimize risk on original data, with hard consistency regularization (i.e. enforcing ((X) ((t(X)), vt E TT)
Then we recover the CF-invariant predictor that minimizes risk on original data.
Proof. First, for the convenience notation let's assume = X(z)) a.e. for some 2n E 2. Then by the labee-preserving TI we have: for each E 'T we have ((X) [00]=[(([[[pp ((X(=)))) (z) for some z E Z.
Consider consistency training. Let (T( denote the set of representation functions satisfying consistency requirement under transformation set T, i.e. (.(T) == (ó: x(X) = +(((X)) a.e. VI
T}. If ó (T) then for any 2, 2' Z, find T such that X(2) +(X(2)) since T enumerates all potential outcomes Z; therefore ((X(2') consistency requirement. Thus
((t(X()))) = (x)) a.e. by
(E). On the other hand if o E Parin then for any t E T, we have o(((X)) = ((X()) = o(X) for some E 2. Thus & E (T( Therefore (T) (E). Therefore. training the model to minimize risk on original data, with hard consistency regularization is equivalent to CF-invariant representation learning, which recovers the optimal CF-invariant predictor on training distribution.
Consider ERM training on augmented data with purely-spurious Z. Let P denote the original distribution, and P denote the distribution after the augmentation. Let T be the random variable for transformation operation. First, the generating process of the augmented data is: first sample (r(:); then sample XX,Y)T t from the distribution of (t(X), Y). Then have:
P(X,Y) / Y(dP(()
/ P(X((),
P(X(=), X(X()))
/ P(Xx
by the label-preserving of T, and the fact that Y a descendant of Z.
15

---

Under review as a conference paper at ICLR 2023
Y
E
E
(a) base graph
anti-causal
E
(c) confoumded-outtomme
(d) confoumded-descendaat
We put below Figure again and the base graph for convenience of inspection. The black arrows are included in graphs. orange arrows are optional. least one of the two dashed blue arrows in Figure 4d must exist.
Figure 4:
Next, observe that P(ylæ, do(z)) = ~ This is because: in original probability we have Y X Z: the opera-one removes the incoming edges of and set Z 2; as result P(y|z, do(=)) Plyl= do(z)) P)yl=
Put together:
P(X,Y) = / P(do(e))))
/
/ PCXIXX)P((XXO)))PPP
= P ([(x(a(())PPPe=P(((Y))) =
Therefore the objective is:
Ep(L(Y, PENN)BEYYO
Then for any input x, the the optimal predictor output (x( = argmin ((a)) S L(y, (()))))())z This is the same as directly restricting predictor to be CF-invariant.
Theorem 9. Suppose set of domains E share the common causal structure underlying P, Z.) Then E satisfies CISA if and only if the corresponding causal DAG one of the given in Figure? 2 In particular, there three families of allowed DAGs: anti-causal, confounded-oute -ottoomee and confounded-deseee descendant.
Proof. There are finite number of possible causal DAGs relating the variables U, Z, X Y, Moreover, for DAG to be compatible with CISA it must satisfy some conditions that narrows down the set. In particular, Z causes X, but not X or Y (by definition of Z); U should confound Z and Y. but cannot confound X and Y (by definition of U'); E only affects U (as only PU) changes across environments and P(X, Y, Z(U) is invariant). Below we use these conditions to enumerate all CISA-compatible DAGs.
Now we have Z + X,, + (optional) and + U (and it's the only edge from E). Note that define Z to not have any causal effect on Y. Accordingly, the path Z - X, Y is ruled out. Thus X, - Y is not allowed but Y + X, is optional. These edges form the base graph 4a build upon.
16

---

Under review as a conference paper at ICLR 2023
Figure 5: Enumerating CISA-compatible causal graphs
base graph 4a
x - y
anti-causal 4b
- X
conf-outome 4c
conf-descend 4d
Next, we divide into two cases: X - Y and X +- (when there is no edge between them, we can treat it as either case and the resulting graphs the same).
When X + Y: we require that U confounds Z, Y, so we need + Y (otherwise can't cause Y) and U + Z (otherwise U can't confound the relationship between U, 2). We do not allow U + X as otherwise the relationship between X and Y confounded. There are a few optional edges U + Xx, Y Z, X Z, as they do not violate CISA assumptions. Other edges cannot be allowed as they will violate CISA ussumptions. These constitute the anti-causal subtype as illustrated in 4b
When X_ + Y, we can again divide into two exclusive cases: U - Y and U + x Why? We need at least one of these two edges, otherwise U does not cause Y; the two edges cannot exist simultaneously as otherwise the relationship between X and Y is confounded.
So, when X Y and U Y: need U 7 otherwise U does not confound Z, Y. There a few optional edges U -> Xz, Y + Z, X - Z, as they do not violate CISA assumptions, Other edges cannot be allowed they will violate CISA assumptions. These constitute the confounded- outcome subtype as illustrated in 4c
Next, when _= + and the two edges (or both) X! + Z, + Z.There are few optional edges U - Xx, Y + Z, as they do not violate CISA assumptions. Other edges cannot be allowed as they will violate CISA assumptions. These constitute the confounded-ddscendaat subtype as illustrated in 4d
X to let cause Z and confound Z, Y. need least one of
Theorem 10. ó is counterfactuullly-nvaai invariant representation,
1. if the underlying causal graph is anti-causal, ((X) 1 YY; 1 E;
2. if the underlying causal graph is confounded- oi
3. if the underlying causal graph is comfounded--esceente Y L ElØ(X).
Proof. Reading d-separation from the corresponding DAGs, have X! 1 E for anti-causal problems; X 1 E for confounded-outcome problems; Y 1 E X for menfounddd--descennee problems. Since o is CF-invariant, that means (xX) ) is -measuura.ll. Thus the claim follows. n
17

---

Under review as a conference paper at ICLR 2023
Theorem 13. Let & satisfy CISA, then
1. E
then PD((E) c. (IRM(E)
2. if E is anti-causal,, then Dpi(E) c Dg-IRM (E, Po) for any chosen Po
Proof. Confounded-desscendaat case: let Φ E DDi(E), i.e. IL E (XX). To show the risk minimizer the same, it suffices show P (XX) be the same for all This immediate from the distributional invariance.
Anti-causal case: the representation & E PDI(E), i.c ((X) EYY,
Po(Y) L(Y, (w () -YYmP. Po(Y) [ (1
Po(Y)
Ep. Po(Y)
Ey (Xx (w (X))))
The second equality is because ((X) 1 Y.Y
Thus the objective function the same across domains, so the optimal w the same. Therefore (.-IR((E)
18


Reviews:
[{"review_id": "yw1nLW3wh61", "reviewer": "Reviewer_1f5D", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "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": "This paper proposes a framework unifying multiple approaches for out-of-distribution generalization, namely, (i) data-augmentation, (ii) enforcing distributional invariances and (iii) invariant risk minimization. The framework is centered around the assumption that both the training and testing domains are \"Causally Invariant with Spurious Associations\" (CISA), which basically postulates the existence of a specific kind of causal graphical model with latent variables in which the causal relationships are assumed to be stable from one environment to another and only the distribution of a confounder is allowed to change. Based on this assumption, the authors propose \"counterfactually invariant representation learning\" as an ideal but generally intractable procedure to learn an invariant representation. Then, it is showed how the three approaches mentioned above can be seen as either exactly solving this problem or a relaxation of it, as long as the ground-truth generating process satisfies some assumptions (anti-causal, confounded-outcome & confounded-descendant) specific to each method. One of the key takeaway is that one should choose a method matching the true causal structure.", "summary_of_the_review": "The paper was an interesting read, and I believe gaining a better theoretical understanding of precisely when these methods are expected to yield an invariant representation could provide guidance to practitioners. However, mainly because of the weak motivation for the CISA assumption (which is really central to this work), but also because of the lack of clarity at times, I believe this paper is only marginally above the acceptance threshold. I'm open to reconsider my evaluation after discussing with the authors as well as the other reviewers.", "strength_and_weaknesses": "Disclaimer: \n\nI am not an expert on this topic and didn't know about most of the papers cited. As a consequence, I couldn't confirmed the accuracy of the claims made about those papers and cannot really assess how novel this contribution is. My review addresses mainly points about readability, motivation and clarity.\n\nStrengths:\n- This paper is well structured and for the most part easy to follow. I enjoyed reading it.\n- The general goal of understanding theoretically under which settings previously proposed methods are expected to work is very important and should be valuable to the ICLR community.\n\nWeaknesses:\n- My main concern with this work is the weak motivation for exploring this specific class of data generating processes, i.e. CISA. A \"cow on the beach” example is provided in the introduction to motivate CISA, but I could not understand how this example satisfies this assumption. It might be useful to provide more explanations and/or a graphical model in the appendix to illustrate the point, given how important this assumption is to the contribution. Right now, the CISA assumption seems arbitrary (at least to a non-expert). \n- I found the introduction a bit hard to follow, I believe it could be shorten quite a bit to leave space for motivating CISA.\n- Some mathematical statements lack clarity. I expand on that later in my review.\n\n\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Sqz6FAr6ZGf", "reviewer": "Reviewer_eUgk", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "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": "This paper proposes from a unified perspective to investigate different approaches to learning domain-invariant representations. They first formulate CF-invariant representations and CISA domains, and then show different domain-invariant representation learning methods can be placed under the same framework.", "summary_of_the_review": "My main concern is if CF-invariance will lead to the optimal models in various scenarios? If not, under which conditions? This point seems unclear and not discussed in the paper, which might weaken the impact of this paper. Also, how flexible and general are the causal diagrams considered in the paper?", "strength_and_weaknesses": "**Strength:**\n\n+ The paper proposes to look into different approaches to domain-invariant representation learning from a unified perspective \n+ The idea seems straightforward and makes sense\n\n**Weaknesses:**\n\n- I am wondering if CF-invariance will lead to the optimal solution in terms of performance in different tasks. This is because some paper (e.g., https://openreview.net/forum?id=12RoR2o32T) suggests that CF-invariance be only sub-optimal in various scenarios. \n- The causal diagrams shown in Figures 1&2 seem not flexible enough to cover many cases, e.g.,\n   - Why do we only have $X_z^{\\perp} \\rightarrow X_z$, not $X_z^{\\perp} \\leftarrow X_z$?\n   - Why is there no hidden confounders between $X_z^{\\perp}$ and $X_z$ that is affected by $U$ or $E$?\n- Is it really necessary to introduce both $Z$ and $U$ to the causal diagrams? Is it possible to just combine $Z$ and $U$ and treat it as one hidden variable that is affected by $U$ or $E$?\n- I am unsure if it makes sense to simply divide $X$ into $X_z^{\\perp}$ and $X_z$ in the causal diagrams? It might be better to explicitly include $X$?\n- There is no empirical demonstrations for the theoretical results.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "PJRwzRx-iKt", "reviewer": "Reviewer_2pqk", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "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": "This paper aims to provide a causal view of the domain invariant learning. To achieve this goal, the authors have analyzed three types of methods. For all these methods, the authors believe that the invariant parts are different. By taking from the causal view, the authors have presented detail analysis on how to build a causality enhanced model to understand the domain invariant concept. In  the experiment, the authors have conducted extensive experiments to demonstrate the effectiveness of the model.\n", "summary_of_the_review": "See the above comments.", "strength_and_weaknesses": "\nIn general, the paper is well written. The contribution is interesting and clear. The proposed model is sound and may theoretical analysis have been provided. \nMy major concerns are as follows, can we have some experiments to demonstrate the claims of the paper. At least, we can have some simulation studies. How can we ensure that the causal graph in Figure 1 is accurate. Since these causal graph is important for the following analysis, we should better provide more evidences on its rationalities.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "1IsBgumzWb", "reviewer": "Reviewer_nMg7", "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": "The paper aims to present a unifying account of different methods for domain-invariant (supervised) representation learning, that is, given data from multiple joint distributions over inputs X and labels Y, learning structure that is stable across them and discarding \"spurious\" parts that vary. Specifically, three families of methods are considered: data augmentation (DA), distributional invariance (DI), and invariant risk (IR). To analyse these methods, the paper considers a setting where different domains or datasets are related through the CISA assumption, which postulates that causal relationships stay the same while only the distribution of a latent common cause U changes across domains. More formally, any $P_e(X,Y)$ is obtained from the same $P_0(X,Y,Z|U)$ via integration w.r.t.  a domain-varying $P_e(U)$ and marginalisation of the latent factors of variation (FoV) $Z$. The paper uses the notion of counterfactual invariance (CF-invariance) of Veitch et al. (2021) to formalise the desired notion of invariance. It then proceeds to study the implications of learning via DA, DI, and IR under different causal structures satisfying the CISA assumption. \n\nIt demonstrates that the causal structure is important in determining whether or not different methods (and subvariants thereof) succeed in learning domain invariant structure in the form of a CF-invariant representation of the input. Roughly, it is argued that DA is the gold standard but typically infeasible, depending on the causal structure, different types of DI are necessary (but not sufficient) for CF-invariance, and IR further relaxes DI by only enforcing invariance w.r.t. aspects of the distribution on which the risk depends. The paper concludes by discussing its results and insights in the context of robust prediction and related literature. ", "summary_of_the_review": "The paper is original and tackles an important problem from a new and interesting angle. I see great potential for this work. However, the current presentation is in many parts too vague and unclear and sometimes fails to convincingly deliver the main points---even if these may be correct. I view this submission as a borderline case. Since the paper could be greatly improved through careful revision, rewriting, and reformatting, I lean toward rejection.", "strength_and_weaknesses": "### Strengths\n- The paper is well-motivated and original in pursuing a unified perspective of the by-now large literature on supervised invariant representation learning. To the best of my knowledge, no such unified account exists thus far.\n- The categorisation of existing methods into three families (DA, DI, IR) is helpful.\n- The paper succeeds to some extent in proposing a common framework for analysing different methods. I would agree that causality is the right language to do so, and the considered CISA setting is non-trivial and reasonably flexible in containing different causal structures as special cases. \n- Some of the presented (theoretical) results are novel and interesting.\n- The paper makes some prescriptive statements that can be useful for practitioners aiming to learn invariant representation while having available knowledge to decide which causal structure may closest to the truth.\n\n### Weaknesses\n- The paper is a bit hard to follow in parts; the presentation of the work could be substantially improved (see below for more details).\n- Especially for being a theoretical/conceptual paper, attempting to provide a unifying account of existing literature, it often stays a bit too vague: certain key concepts such as \"causally related\", \"confounded\", \"sufficient for Y\" etc are not formally defined and, despite being intuitively familiar to someone with a background in causal inference, leave some ambiguity. A rigorous description of the entire considered data-generating process within a causal modelling framework such as SCMs or potential outcomes would be a helpful addition. Without specifying additional assumptions (e.g., faithfulness---is this assumed?), talking about causality solely based on the graphs in Figs. 1 and 2 is insufficient. \n- Related to the previous two points, several claims made in the normal text are not made sufficiently clear and/or are not sufficiently backed by evidence or formal argument to allow for proper evaluation of their correctness. To give some examples: Sec 4.2 mentions confounding, but it is not clear exactly which variable is supposed to confound the relationship between which others; in the context of the examples for \"anti-causal\" and \"confounded outcome\", the arrows $X^\\perp_z\\to Z$ and $Y \\to X_z$ seem strange; in Defn 11, should this hold for all P_0, or really just \"any\" P_0 (no additional requirements?), and why is P_0 not an argument to the LHS? much of the discussion in 4.3 and 5 is also not fully convincing/hard to follow. [I'm not claiming that these are necessarily errors or incorrect claims, just that it is hard to assess based on the presentation]\n- The paper misses several references which I believe should be added: In the context of classical, simpler assumptions for domain adaptation, for example [R1,R2,R3] come to mind; it would also be interesting to see a discussion of other recent work providing a causal perspective on learning with spurious features [R4]; in the context of the line of work on IR learning, [R5] could also be mentioned, which can be seen as a generalisation of Krueger et al. (VRex); in the context of 4.1, [R6] also discusses generative process and counterfactual view for data augmentation that aims to learn invariant (\"content\") features and discard spurious \"style\" features; the term \"anticausal\" was coined in [R7]\n\n[R1] Shimodaira, H. (2000). Improving predictive inference under covariate shift by weighting the log-likelihood function. Journal of statistical planning and inference, 90(2), 227-244.\n\n[R2] Quinonero-Candela, J., Sugiyama, M., Schwaighofer, A., & Lawrence, N. D. (Eds.). (2008). Dataset shift in machine learning. Mit Press.\n\n[R3] Lipton, Z., Wang, Y. X., & Smola, A. (2018, July). Detecting and correcting for label shift with black box predictors. In International conference on machine learning (pp. 3122-3130). PMLR.\n\n[R4] Makar, M., et al. (2022, May). Causally motivated shortcut removal using auxiliary labels. In International Conference on Artificial Intelligence and Statistics (pp. 739-766). PMLR.\n\n[R5] Eastwood, C., et al. \"Probable domain generalization via quantile risk minimization.\" arXiv preprint arXiv:2207.09944 (2022).\n\n[R6] Von Kügelgen, J., et al. (2021). Self-supervised learning with data augmentations provably isolates content from style. Advances in neural information processing systems, 34, 16451-16467.\n\n[R7] Schölkopf, B., et al. (2012, January). On causal and anticausal learning. In ICML.\n\n### Comments\n- The proofs of Thms. 9 and 12 seem rather trivial in that the result essentially follows from the previous assumptions and definitions; hence, it may be more appropriate to call these Lemma or Proposition.\n- for the more involved proofs, it would be helpful to include at least a sketch in the main paper\n- since representation learning is often associated with an unsupervised setting, I think it would be good to add \"supervised\" to the title (it is already stated in the introduction that this work only focuses on *supervised* domain invariant representation learning)\n- the description of the \"two-bit-envs experiment\" seems out of place in the conclusion and lacks context and details", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Decision: Reject