I NVARIANT C AUSAL R EPRESENTATION L EARNING FOR O UT - OF -D ISTRIBUTION G ENERALIZATION
Anonymous authors Paper under double-blind review
A BSTRACT
Due to spurious correlations, machine learning systems often fail to generalize to environments whose distributions differ from the ones used at training time.
Prior work addressing this, either explicitly or implicitly, attempted to ﬁnd a data representation that has an invariant relationship with the target.
This is done by leveraging a diverse set of training environments to reduce the effect of spurious features and build an invariant predictor.
However, these methods have generalization guarantees only when both data representation and classiﬁers come from a linear model class.
We propose invariant Causal Representation Learning (iCaRL), an approach that enables out-of-distribution (OOD) generalization in the nonlinear setting (i.e., nonlinear representations and nonlinear classiﬁers).
It builds upon a practical and general assumption: the prior over the data representation (i.e., a set of latent variables encoding the data) given the target and the environment belongs to general exponential family distributions, i.e., a more ﬂexible conditionally non-factorized prior that can actually capture complicated dependences between the latent variables.
Based on this, we show that it is possible to identify the data representation up to simple transformations.
We also prove that all direct causes of the target can be fully discovered, which further enables us to obtain generalization guarantees in the nonlinear setting.
Extensive experiments on both synthetic and real-world datasets show that our approach outperforms a variety of baseline methods.
I NTRODUCTION
Modern machine learning algorithms still lack robustness, and may fail to generalize outside of a speciﬁc training distribution because they learn easy-to-ﬁt spurious correlations which are prone to change between training and testing environments.
We recall the widely used example of classifying images of camels and cows (Beery et al., 2018).
Here, the training dataset has a selection bias, i.e., many pictures of cows are taken on green pastures, while most pictures of camels happen to be in deserts.
After training, it is found that the model builds on spurious correlations, i.e., it relates green pastures with cows and deserts with camels, and fails to recognize images of cows on the beach.
To address this problem, a natural idea is to identify which features of the training data present domain-varying spurious correlations with labels and which features describe true correlations of interest that are stable across domains.
In the example above, the former are the features describing the context (e.g., pastures and deserts), whilst the latter are the features describing the animals (e.g., animal shape).
By exploiting the varying degrees of spurious correlation naturally present in training data collected from multiple environments, one can try to identify stable features and build invariant predictors.
Invariant risk minimization (IRM) seeks to ﬁnd data representations (Arjovsky et al., 2019) or features (Rojas-Carulla et al., 2018) for which the optimal predictor is invariant across all environments.
The general formulation of IRM is a challenging bi-leveled optimization problem, and theoretical guarantees require constraining both data representations and classiﬁers to be linear (Arjovsky et al., 2019, Theorem 9), or considering the special case of feature selection (Rojas-Carulla et al., 2018, Theorem 4).
Ahuja et al.
(2020a) study the problem from the perspective of game theory, with an approach termed invariant risk minimization games (IRMG).
They show that the set of Nash equilibria for a proposed game is equivalent to the set of invariant predictors for any ﬁnite number of environments, even with nonlinear data representations and nonlinear classiﬁers.
However, these theoretical results in the nonlinear setting only guarantee that one can learn invariant predictors from training environments, but do not guarantee that the learned invariant predictors can generalize well across all environments including unseen testing environments.
We propose invariant Causal Representation Learning (iCaRL), a novel approach that enables outof-distribution (OOD) generalization in the nonlinear setting (i.e., nonlinear representations and nonlinear classiﬁers 1 ).
We achieve this by extending and using methods from representation learning and graphical causal discovery.
In more detail, we ﬁrst introduce our main general assumption : when conditioning on the target (e.g., labels) and the environment (represented as an index), the prior over the data representation (i.e., a set of latent variables encoding the data) belongs to a general exponential family .
Unlike the conditionally factorized prior assumed in recent identiﬁable variational autoencoders (iVAE) (Khemakhem et al., 2020a), this is a more ﬂexible conditionally non-factorized prior, which can actually capture complicated dependences between the latent variables.
We then extend iVAE to the case in which the latent variable prior belongs to such a general exponential family .
The combination of this result and the previous general assumption allows us to a guarantee that the data representation can be identiﬁed up to simple transformations.
We then theoretically show that the direct causes of the target can be fully discovered by analyzing all possible graphs in a structural equation model setting.
Once they are discovered, the challenging bi-leveled optimization problem in IRM and IRMG can be reduced to two simpler independent optimization problems, that is, learning the data representation and learning the optimal classiﬁer can be performed separately.
This leads to a practical algorithm and enables us to obtain generalization guarantees in the nonlinear setting.
Overall, we make a number of key contributions: (1) We propose a general framework for out-ofdistribution generalization in the nonlinear setting with the theoretical guarantees on both identiﬁability and generalizability; (2) We propose a general assumption on the underlying causal diagram for prediction (Assumption 1 and Fig. 1c), which covers many real-world scenarios (Section 3.2); (3)
We propose a general assumption on the prior over the latent variables (Assumption 2), i.e., a more ﬂexible conditionally non-factorized prior; (4) We prove that an extended iVAE with this conditionally non-factorized prior is also identiﬁable (Theorems 1, 2&3); (5) We prove that our framework has the theoretical guarantees for OOD generalization in the nonlinear setting (Proposition 1).
P RELIMINARIES 2. I DENTIFIABLE V ARIATIONAL A UTOENCODERS
Variational autoencoders (VAEs, see Appendix B) (Kingma & Welling, 2013; Rezende et al., 2014) lack identiﬁability guarantees.
Consider a VAE model where O ∈ R d stands for the observed variables (data) and X ∈ R n for the latent variables.
Khemakhem et al.
(2020a) show that a VAE with an unconditional prior distribution p θ ( X ) over the latent variables is unidentiﬁable.
However, they also show that it is possible to obtain an identiﬁable model if one posits a conditionally factorized prior distribution over the latent variables, p θ ( X | U ) , where U ∈ R m is an additional observed variable (Hyvärinen et al., 2019).
Speciﬁcally, let θ =
( f , T , λ ) ∈
Θ be the parameters of the conditional generative model
[Equation]
[Equation]
where p f ( O | X ) = p (cid:15) ( O − f ( X )) in which (cid:15) is an independent noise variable with probability density function p (cid:15) ( (cid:15) ) .
Importantly, the prior p T , λ ( X | U ) is assumed to be conditionally factorial, where each element of X i ∈
X has a univariate exponential family distribution given U .
The conditioning on U is through an arbitrary function λ ( U )
(e.g., a neural net) that outputs the individual exponential family parameters
λ
i ( U ) for each X i .
The prior probability density thus takes the form
[Equation]
where Q i is the base measure, X i the i -th dimension of X ,
Z i ( U ) the normalizing constant,
T i =
( T i, 1 , . . .
, T i,k ) the sufﬁcient statistics, λ i ( U ) =
( λ i, 1 ( U ) , . . .
, λ i,k ( U ))
the corresponding natural parameters depending on U , and k the dimension of each sufﬁcient statistic that is ﬁxed in advance.
It is worth noting that this prior is restrictive as it is factorial and therefore cannot capture dependencies.
As in VAEs, the model parameters are estimated by maximizing the corresponding evidence lower bound,
[Equation]
where we denote by p D the empirical data distribution given by the dataset D = (cid:8)(cid:0)
O ( i ) , U ( i ) (cid:1)(cid:9) N i = and q φ ( X | O , U ) denotes an approximate conditional distribution for X given by a recognition network with parameters φ .
This approach is called identiﬁable VAE (iVAE).
Most importantly, it can be proved that under the conditions stated in Theorem 2 of (Khemakhem et al., 2020a), iVAE can identify the latent variables X up to a permutation and a simple componentwise transformation, see
Appendix F.
2. I NVARIANT R ISK M INIMIZATION
Arjovsky et al.
introduced invariant risk minimization (IRM), whose goal is to construct an invariant predictor f that performs well across all environments E all by exploiting data collected from multiple environments E tr , where E tr ⊆ E all .
Technically, they consider datasets D e :=
{ ( o ei , y ei ) } n e i =1 from multiple training environments e ∈ E tr , where o ei ∈ O ⊆
R d is the input observation and its corresponding label is
y ei ∈ Y ⊆
R s .
The dataset D e , collected from environment e , consists of examples identically and independently distributed according to some probability distribution P ( O e , Y e ) .
The goal of IRM is to use these multiple datasets to learn a predictor
Y =
f ( O ) that performs well for all the environments.
Here we deﬁne the risk reached by f in environment e as R e ( f ) =
E O e , Y e [ (cid:96) ( f ( O e ) , Y e )] , where (cid:96) ( · ) is a loss function.
Then, the invariant predictor can be formally deﬁned as follows:
Deﬁnition 1 (Invariant Predictor (Arjovsky et al., 2019)) .
We say that a data representation Φ ∈
H Φ : O → C elicits an invariant predictor w ◦ Φ across environments E if there is a classiﬁer w ∈ H w :
C → Y simultaneously optimal for all environments, that is, w ∈ arg min ¯ w ∈H w R e ( ¯ w ◦ Φ) for all e ∈ E , where ◦ means function composition.
Mathematically, IRM can be phrased as the following constrained optimization problem:
[Equation]
Since this is a generally infeasible bi-leveled optimization problem, Arjovsky et al.
(2019) rephrased it as a tractable penalized optimization problem by transfering the inner optimization routine to a penalty term.
The main generalization result (Theorem 9 in Arjovsky et al.
(2019)) states that if both Φ and w come from the class of linear models (i.e., H Φ
= R n × n and H w = R n × 1 ), under certain conditions on the diversity of training environments (Assumption 8 in Arjovsky et al.
and the data generation, the invariant predictor w ◦ Φ obtained by solving Eq.
(4) remains invariant in E all .
P ROBLEM S ETUP 3. A M OTIVATING E XAMPLE
In this section, we extend the example which was introduced by Wright (1921) and discussed by
Arjovsky et al.
(2019), and provide a further in-depth analysis.
Model 1.
Consider a structural equation model (SEM) with a discrete environment variable E that modulates the noises in the structural assignments connecting the other variables (cf.
Fig.
1a below):
[Equation]
where Gaussian (0 , σ ) denotes a Gaussian random variable with zero mean and standard deviation σ , and σ 1 , . .
. , σ 3 are functions of the value e ∈ E all taken by the environment variable E .
To ease exposition, here we consider the simple scenario in which E all only contains all modiﬁcations varying the noises of X 1 , X 2 and Y within a ﬁnite range, i.e., σ i ( e ) ∈ [0 , σ 2 max ] .
Then, to predict Y from ( X 1 , X 2 ) using a least-square predictor
ˆ Y e
= ˆ
α X e[-
-]
+ ˆ α 2 X e 2 for environment e ,
we can
• Case 1: regress from X e 1 , to obtain ˆ α 1 = 1 and ˆ α 2 = 0 , • Case 2: regress from X e 2 , to obtain ˆ α 1 = 0 and ˆ α 2 = σ 1 ( e )+ σ 2 ( e ) σ 1 ( e )+ σ 2 ( e )+ σ 3 ( e ) , • Case 3: regress from ( X e 1 , X e 2 ) , to obtain ˆ α 1 = σ 3 ( e ) σ 2 ( e )+ σ 3 ( e ) and ˆ α 2 = σ 2 ( e ) σ 2 ( e )+ σ 3 ( e ) .
In the generic scenario (i.e., σ 1 ( e ) (cid:54) = 0 , σ 2 ( e ) (cid:54) = 0 , and σ 3 ( e ) (cid:54) = 0 ), the regression using X 1 in
Case 1 is an invariant correlation: it is the only regression whose coefﬁcients do not vary with e .
By contrast, the regressions in both Case 2 and Case 3 have coefﬁcients that depend on e .
Therefore, only the invariant correlation in Case 1 will generalize well to new test environments.
X
Y
X
E (a)
X
Y
X
E
O (b)
X p 1 ··· X pr
Y
X c 1 ··· X ck
E
O (c)
Another way to understand Model 1 is through its graphical representation 3 , as shown in Fig.
1a.
We treat the environment as a random variable E , where E could be any information speciﬁc to the environment (Storkey, 2009; Peters et al., 2015; Zhang et al., 2017; Huang et al., 2020).
For simplicity, we let E be the environment index, i.e., E ∈ { 1 , . . .
, N } , where N is the number of training environments.
A more realistic version appearing in many settings is shown in Fig.
1b, where the true variables { X 1 , X 2 } are unobserved and we can only observe their transformation O .
In this case, Invariant Causal Prediction (ICP) (Peters et al., 2015) will fail when applied to O , even when Y is not affected by E (i.e., the edge E
→ Y is removed).
The reason is that each variable (i.e., each dimension) of O is jointly inﬂuenced by both X 1 and X 2 so that ICP is unable to ﬁnd the variables containing the information only about X 1 by searching for a subset of variables O .
By contrast, both
IRM and IRMG work, as long as the transformation is linear.
These ﬁndings are also empirically illustrated in Section 5.1.
We now go even further and consider a more general causal graph in which
Y can have more than one parent or child.
3. A SSUMPTIONS ON THE C AUSAL G RAPH
We extend the causal graph in Fig.
1b to a more general setting 4 , as encapsulated in Fig.
1c.
In particular, we now have O ∈ O ⊆ R d ,
Y ∈ Y ⊆
R s , X
=
( X p 1 , . . .
, X p r , X c 1 , . . .
, X c k ) ∈
X ⊆ R n , where n = r + k , and { X i } i ∈ I p
= { p 1 ,...,p r } and { X j } j ∈ I c
= { c 1 ,...,c k } are multiple scalar causal factors and non-causal factors 5 of Y , respectively.
We denote X p .
=
( X p 1 , . . .
, X p r ) and
X c .
=
( X c 1 , . . .
, X c k ) for the ease of clariﬁcation.
We also assume that X is of lower dimension than O , that is, n ≤ d .
We allow for arbitrary connections between the latent variables X as long as the resulting causal diagram including Y is a directed acyclic graph (DAG).
We use dashed lines to indicate the causal mechanisms which might vary across environments and even be absent in some scenarios, whilst solid lines indicate that they are invariant across all the environments.
To sum up, we assume that the underlying causal graph encapsulated in Fig.
1c satisﬁes the following assumption 6 :
Assumption 1.
(a) X i depends on one or both of Y and E for any i ; (b) The causal graph containing X and Y is a DAG; (c) O ⊥⊥ Y , E | X , implying that p ( O | X ) is invariant across all the environments; (d) Y ⊥⊥ E | X p , implying that p ( Y | X p ) is invariant across all the environments.
One may be wondering how practical Assumption 1 is in real world applications.
Let us explore this in more detail.
Assumption 1a rules out all the useless X i in the task of predicting Y .
This is because if Assumption 1a is violated, meaning that X i is independent of Y and E and has no inﬂuence in predicting Y , then such X i should be viewed as noise and thus eliminated during learning.
Assumption 1b is a common assumption in causal discovery (Spirtes et al., 2000; Pearl, 2009; Peters et al., 2015).
It also makes sense in Assumption 1c that the generative mechanism p ( O | X ) is invariant across all the environments.
Otherwise, it is impossible to infer X from O in any unseen environment.
Assumption 1d is a widely-used default assumption in OOD generalization (Peters et al., 2015; Arjovsky et al., 2019). In fact, Assumption 1d can be further relaxed to the more
[Table]
practical one that E [ Y | X p ] is invariant across all the environments.
That is, given X p , we allow
E to only affect the amount of noise in the distribution of Y , because that would not change the expected value of Y since the mean of the noise would be zero.
Apparently, Assumption 1, together with the causal graph in Fig. 1c, covers most scenarios (e.g., the ones of Zhang et al.
; von
Kügelgen et al.
(2020); Sun et al.
; Ahuja et al. (2021); von Kügelgen et al.
and is a very ﬂexible model for causal analysis when predicting Y from O .
3. A SSUMPTIONS ON THE P RIOR
When the underlying causal graph satisﬁes Assumption 1, our primary assumption leading to identiﬁability in this general setting is that the conditional prior p ( X | Y , E ) belongs to a general exponential family.
This is formalized as follows:
Assumption 2.
p ( X | Y , E ) belongs to a general exponential family with parameter vector given by an arbitrary function λ
( Y , E ) and sufﬁcient statistics T ( X ) =
[ T f ( X ) T , T NN ( X ) T ] T given by the concatenation of a) the sufﬁcient statistics T f
( X ) =
[ T 1 ( X 1 ) T , · · · , T n ( X n ) T ] T of a factorized exponential family, where all the T i ( X i ) have dimension larger or equal to 2, and b) the output
T NN ( X ) of a neural network with ReLU activations.
The resulting density function is thus given by
[Equation]
where Q is the base measure and Z the normalizing constant.
A neural network with ReLU activation has universal approximation power.
Therefore, the term
T NN ( X ) in the above prior distribution will allow us to capture arbitrary dependencies between the latent variables.
The distribution in Eq.
(5) is more ﬂexible than the conditionally factorized prior assumed by iVAEs.
However, surprisingly, the identiﬁability results of iVAEs also hold when using the more ﬂexible prior in Eq.
(5), as we will show in Section 4.1.
This motivates using an extended iVAE model with the above prior to model data generated by the ground truth model in Fig. 1c.
However, in this case, the data generating model and the learned model might have different priors.
For example, in the ground truth model, the prior for each X i ∈ I p might be p ( X i | E ) , when X i is only caused by E .
By contrast, in the extended iVAE model the prior is p ( X | Y , E ) .
Is this going to affect the identiﬁability of the latent variables?
Well, in practice not because the posterior distribution for X given O , Y and
E would be equivalent in both models (up to identiﬁability guarantees).
I NVARIANT C AUSAL R EPRESENTATION L EARNING
We now introduce our algorithm, invariant Causal Representation Learning (iCaRL), which consists of 3 phases as summarized in Algorithm 1.
The basic idea is that we ﬁrst identify X by using an extended iVAE model under Assumptions 1&
(Phase 4.1), then discover direct causes of Y among the identiﬁed X (Phase 4.2), and ﬁnally learn an invariant predictor for Y from the discovered causes (Phase 4.3).
4. P HASE 1: I DENTIFYING L ATENT V ARIABLES U SING NF- I VAE
In this section, we describe our identiﬁable model, namely NF-iVAE, which is an extended iVAE with a general non-factorized prior that is able to capture complex dependences between the latent variables.
Technically, in the general setting under Assumption 1, it is straightforward to obtain a corresponding generative model by directly substituting U with ( Y , E ) in Eq.
[Equation]
The corresponding lower bound is
[Equation]
To obtain an identiﬁability result, we assume that the prior p T , λ
( X | Y , E ) satisﬁes Assumption(i.e.
, Eq.
Since the prior is a general multivariate exponential family distribution with unknown normalization constant, we cannot learn its parameters ( T , λ ) by directly maximizing Eq.
Instead, we use score matching, a well-known method for training unnormalized probabilistic models (Hyvärinen, 2005; Vincent, 2011), and learn ( T , λ ) by minimizing
[Equation]
In practice, we can use a simple trick of partial integration to simplify the evaluation of Eq.
(9), see Appendix C. Furthermore, we can jointly learn ( θ , φ ) by combining Eq.
(8) and Eq.
(9) in the following objective:
[Equation]
where ˆ f , ˆ T , ˆ λ , ˆ φ are copies of f , T , λ , φ that are treated as constants and whose gradient is not calculated during learning.
More details can be found in Appendix M.
We now state our main theoretical results:
Theorem 1.
Assume that we observe data sampled from a generative model deﬁned according to Eqs.
(5-7), with parameters θ := ( f , T , λ ) , where p T , λ ( X | Y , E ) satisﬁes Assumption 2.
Furthermore, assume the following holds: (i) The set { O ∈ O| ϕ (cid:15) ( O ) = 0 } has measure zero, where ϕ (cid:15) is the characteristic function of the density p (cid:15) deﬁned in Eq.
(ii) Function f in Eq.
(7) is injective, and has all second-order cross derivatives.
(iii)
The sufﬁcient statistics in T f are all twice differentiable.
(iv) There exist k + 1 distinct points ( Y , E ) 0 , . . .
, ( Y , E ) k
such that the matrix
L = (cid:0) λ
(( Y , E ) 1 ) − λ (( Y , E ) 0 ) , . . .
, λ (( Y , E ) k ) − λ
( Y , E ) 0 ) (cid:1) of size k × k is invertible, where k is the dimension of T .
Then the parameters θ are identiﬁable up to a permutation and a “simple transformation" of the latent variables X , deﬁned as a componentwise nonlinearity making each recovered T
i
( X i ) in T f ( X ) equal to the original up to a linear operation.
Note that, this theorem is inspired by but beyond the main results of iVAEs in that the former is predicated on Assumption 2 which is more ﬂexible than the conditionally factorized prior assumed in iVAEs.
It results in several key changes in the proof, clariﬁed in Appendix H. Interestingly, from (iv) we can further see that E is unnecessary when there exist k + 1 distinct points Y 0 , . . .
, Y k such that the matrix L = (cid:0) λ
( Y 1 ) − λ
( Y 0 ) , . . .
, λ ( Y k )
− λ
( Y 0 ) (cid:1) of size k × k is invertible.
Not requiring
E would make our approach even more applicable.
We further have the following consistency result for the estimation.
Theorem 2.
Assume that the following holds: (i) The family of distributions q φ ( X | O , Y , E ) contains p θ ( X | O , Y , E ) , and q φ ( X | O , Y , E ) > 0 everywhere.
(ii)
We maximize L phase1 ( θ , φ ) with respect to both θ and φ .
Then in the limit of inﬁnite data, we learn the true parameters θ ∗ up to a permutation and simple transformation of the latent variables X .
As a consequence of Theorem 1&2, we have:
Theorem 3.
Assume the hypotheses of Theorem 1 and Theorem 2 hold, then in the limit of inﬁnite data, we identify the true latent variables
X ∗ up to a permutation and simple transformation.
Theorem 3 states that we can use NF-iVAE to infer the true X up to a permutation and simple transformation.
We use the mean of q φ ( X | O , Y , E ) for this task.
Note that the noise (cid:15) may introduce uncertainty in the estimation of X .
However, when O is high dimensional and X is low dimensional (as common in real world applications), q φ ( X | O , Y , E ) will be highly concentrated and we will still be able to estimate X with high accuracy.
The good results obtained by our method in various experiments seem to corroborate this.
All three theorems are proven in Appendix H.
4. P HASE 2: D ISCOVERING D IRECT C AUSES
After estimating X for each data point, the next step is to determine which components of X are a direct cause of Y .
We denote these components by Pa ( Y ) .
We ﬁrst conduct the PC algorithm
(Spirtes et al., 2000) to learn a Markov equivalence class of DAGs, which gives us the direct neighbors of Y , denoted by Ne ( Y ) .
Then, from Assumption 1, one observation is that for any two latent variables X i and X j from Ne ( Y ) , only when both are causes of Y do we have that the dependency between them increases after additionally conditioning on Y .
Thus, when there exist at least two causal latent variables in Ne ( Y ) , we can test all pairs of latent variables 7 with conditional independence testing 8 (Zhang et al., 2012) to discover Pa ( Y ) by comparing p -values from the two tests: IndTest ( X i , X j | E ) and IndTest ( X i , X j | Y , E ) , where IndTest denotes (conditional) independence test.
Conversely, if no such a pair is found, it implies that there is at most one causal latent variable.
This is a highly unlikely case in real world applications, which is left to Appendix G.
4. P HASE 3: L EARNING AN I NVARIANT P REDICTOR
After having obtained the causal latent variables Pa ( Y ) for Y across training environments, we can learn w by solving the following optimization problem:
[Equation]
where (cid:96) 1 ( · ) could be any loss.
Since we assume that E ( Y | Pa ( Y )) is invariant across E all (the relaxed version of Assumption 1d), the learned w is guaranteed to perform well across E all .
The remaining question is how to infer Pa ( Y ) (i.e., X p ) from O in a new environment.
This can be implemented by leveraging the learned p ( O | X ) .
The rationale behind is that p ( O | X ) is assumed to be invariant across E all (Assumption 1c).
In light of this idea, we follow Sun et al.
and infer
X p from O in any new testing environment by solving the following optimization problem:
[Equation]
where the hyperparameters λ 1 > 0 and λ 2 > 0 control the learned X p and X c have a reasonable scale.
For optimization, we follow Schott et al.
(2018) to ﬁrst use values of X sampled from the training set as initial points and then use Adam to optimize for several iterations.
Note that the noise (cid:15) will introduce uncertainty in the estimation of X p and X c from O .
However, as we mentioned before (below Theorem 3), this noise is not going to affect the estimation much because the likelihood will be highly concentrated around the ground truth values, as corroborated by our good empirical results.
A key question is if iCaRL performs well across E all even though it uses only data from E tr .
That is, does iCaRL enable OOD generalization, as deﬁned by Arjovsky et al.
The answer is positive since Theorem A.1 in (Arjovsky, 2021) indicates that i) any predictor w ◦
Φ with optimal OOD generalization uses only Pa ( Y ) to compute Φ and ii) the classiﬁer w in this optimal predictor can be estimated using data from any environment e for which the distribution of Pa ( Y e ) has full support, which will always be the case since the conditional prior in Eq.
(5) has full support.
Finally, Theorem
A.1 in (Arjovsky, 2021) also indicates that iii) the optimal predictor will be invariant across E all .
Key to these results is that Pa ( Y e ) are available when solving (12).
This requires ﬁrst to identify the latent variables X from O , Y and E and second, to discover the direct causes of Y .
The hypotheses of
Theorems 1 and 2 and Assumption 1 provide this guarantee.
We therefore have the following result whose proof is in Appendix H.
Proposition 1.
Under Assumption 1 and the assumptions of Theorems 1 and 2, the predictor learned by iCaRL across E tr in the limit of inﬁnite data has optimal OOD generalization across E all .
E XPERIMENTS
We compare our approach with a variety of methods on both synthetic and real-world datasets.
In all comparisons, unless stated otherwise, we average performance over ten runs.
Due to space constraints, we only highlight some key results while pointing to the extensive appendices for more information.
The supplement contains all the details of the experiments, e.g., datasets (Appendix J), implementation (Appendix M), hyperparameters and architectures (Appendix N), etc.
5. S YNTHETIC DATA
To verify the identiﬁability of NF-iVAE, we conduct a series of experiments on synthetic data generated according to the causal graph shown in Fig.
2e.
Details of the ground truth data generating [Figure] process are given in Appendix K.
The reason we choose this setting is that it is the simplest case satisfying our requirements: a) For ease of visualization, the latent space had better be 2-dimensional; b) To introduce the non-factorized prior given Y and E (i.e., X i (cid:54)⊥⊥ X j | Y , E ), Y has at least two causes.
We draw 1000 samples from each of the four environments E = { 0 . 2 , 2 , 3 , 5 } , and thus the whole synthetic dataset consists of 4000 samples.
The task is to recover the true latent variable
X =
( X 1 , X 2 ) using the samples of O , E , and Y .
We compare with two widely-used baselines: VAE (Kingma & Welling, 2013) (without identiﬁability guarantees) and iVAE (Khemakhem et al., 2020a)
(with a conditionally factorized prior for identiﬁability).
Through the aforementioned theoretical analysis, it is evident that our method has a more general assumption on the prior leading to identiﬁability.
This is demonstrated empirically in Figs.
2b-2d.
Our method NF-iVAE can recover the original data X up to a permutation and a simple componentwise transformation, whereas all the other methods fail because they are unable to handle the non-factorized case in which X i (cid:54)⊥⊥ X j | Y , E .
We also compute the mean correlation coefﬁcient (MCC) used in Khemakhem et al.
(2020a), which can be obtained by calculating the correlation coefﬁcient between all pairs of true and recovered latent factors and then solving a linear sum assignment problem by assigning each recovered latent factor to the true latent factor with which it best correlates.
By deﬁnition, higher MCC scores indicate stronger identiﬁability.
From Fig.
2f, we can see that the MCC score for NF-iVAE is signiﬁcantly greater than those of VAE and iVAE, indicating much stronger identiﬁability.
Note that, in Appendix K we additionally compare with more methods, whose differences are further summarized in a table.
5. C OLORED MNIST, C OLORED F ASHION MNIST, AND VLCS
In this section, we ﬁrst report experiments on two datasets used in IRM and IRMG: Colored MNIST (CMNIST) and Colored Fashion MNIST (CFMNIST).
We follow the same setting of Ahuja et al.
(2020a) to create these two datasets (see the details in Appendix J).
The task is to predict a binary label assigned to each image which is originally grayscale but artiﬁcially colored in a way that the color is correlated strongly but spuriously with the class label.
For all the experiments on these two datasets, we set the number of the latent variables to n = 10 .
Likewise, we investigate the identiﬁability of NF-iVAE on CMNIST by computing the MCC score between samples of the true latent variable and of the recovered latent variable.
Since the true latent variable on CMNIST is inaccessible to us, we follow Khemakhem et al.
(2020b)
and compute an average MCC score between samples of latent variables recovered by different models trained with different random initialization.
As shown in Fig.
2g, it is evident that the MCC score for NF-iVAE greatly outperforms the others, showing that the latent variables recovered by NF-iVAE have much better identiﬁability.
Furthermore, we demonstrate the ability of iCaRL to discover the causal latent variables (Phase 2) by visualizing the generated images through performing intervention upon a causal latent variable and a non-causal latent variable, respectively.
Fig.
2h shows how intervening upon each of them affects the image.
Obviously, intervening on a causal latent variable affects the shape of the digit but not its color
[Table] [Table]
(top plots), whilst intervening on a non-causal latent variable, which is an effect in Fig.
2h, affects the color of the digit only (bottom plots).
This visually veriﬁes the results of iCaRL in Phase 2.
In terms of the OOD generalization performance
[Table]
we compare iCaRL with 1) IRM, 2) two variants of IRMG: F-IRM Game (with Φ ﬁxed to the identity) and V-IRM Game (with a variable Φ ), 3) three variants of ERM: ERM (on entire training data), ERM e (on each environment e ), and ERM
GRAYSCALE (on data with no spurious correlations), and 4) ROBUST MIN MAX (minimizing the maximum loss across the multiple environments).
Table 1 shows that iCaRL outperforms all other baselines on CFMNIST.
It is worth emphasising that the train and test accuracies of iCaRL closely approach the ones of ERM GRAYSCALE and OPTIMAL, implying that iCaRL approximately learns the true invariant causal representations with almost no correlation with the spurious color feature.
We can draw similar conclusions from the results on CMNIST (Appendix L).
We also report the results on one of the widely used realistic datasets for OOD generalization: VLCS (Fang et al., 2013).
This dataset consists of 10 , 729 photographic images of dimension (3 , 224 , 224) and 5 classes from four domains: Caltech101, LabelMe, SUN09, and VOC2007.
We used the exact experimental setting that is described in Gulrajani & Lopez-Paz (2020).
We provide results averaged over all possible train and test environment combination for one of the commonly used hyper-parameter tuning procedure: train domain validation.
As shown in Table 2, iCaRL achieves state-of-the-art performance when compared to those most popular domain generalization alternatives.
We further include experimental evidence on another popular dataset: PACS (Li et al., 2017a), all the details of which are placed in Appendix L.
R ELATED W ORK
Invariant Causal Prediction (ICP), aims to ﬁnd the causal feature set (i.e., all direct causes of a target variable of interest)
(Peters et al., 2015) by exploiting the invariance property in causality which has been discussed under the term “autonomy”, “modularity”, and “stability” (Haavelmo, 1944;
Aldrich, 1989; Hoover, 1990; Pearl, 2009; Dawid et al., 2010; Schölkopf et al., 2012).
This invariance property assumed in ICP and its nonlinear extension (Heinze-Deml et al., 2018) is limited, because no intervention is allowed on the target variable Y .
Besides, ICP methods implicitly assume that the variables of interest X are given.
The works Magliacane et al.
and Subbaswamy et al.
(2019) attempt to ﬁnd invariant predictors that are maximally predictive using conditional independence tests and other graph-theoretic tools, both of which also assume that the X are given and further assume that additional information about the structure over X is known.
Mitrovic et al.
(2020) analyze data augmentations in self-supervised learning from the perspective of invariant causal mechanisms.
Arjovsky et al.
(2019) reformulate this invariance as an optimization-based problem, allowing us to learn an invariant data representation from O constrained to be a linear transformation of X .
The risks of this approach have been discussed in Rosenfeld et al.
; Kamath et al.
; Nagarajan et al.
and its sample complexity is analyzed in Ahuja et al.
(2020b).
Another line of related work is in the ﬁeld of domain generalization, which we discuss in Appendix A.
C ONCLUSION
We have proposed a novel framework to learn invariant predictors from a diverse set of training environments.
It is based on a practical and general assumption: the prior over the data representation belongs to a general exponential family when conditioning on the target and the environment.
This assumption leads to guarantees that the components in the representation can be identiﬁed up to a permutation and simple transformation.
This allows us to discover all the direct causes of the target, which enables generalization guarantees in the nonlinear setting.
We hope our framework will inspire new ways to address the OOD generalization problem through a causal lens.
R EPRODUCIBILITY S TATEMENT
Proofs of all our theoretical results are given in Appendix H with disclosure of all assumptions.
All datasets used are publicly available or instructions are provided on how to generate them in Appendix
J. Implementation details are given in Appendix M. A description of the hyperparameters and network architectures used is included in Appendix N. Source code for all our experiments will be made available on GitHub after deanonymisation, providing also a complete description of our experimental environment, conﬁguration ﬁles and instructions on the reproduction of our experiments
A D OMAIN G ENERALIZATION
The goal of domain generalization (DG) (Muandet et al., 2013) is OOD generalization: learning a predictor that performs well at unseen test domains.
Unlike domain adaptation (Pan & Yang, 2009;
Ben-David et al., 2007; 2010; Crammer et al., 2008; Patel et al., 2015; Zhao et al., 2019; Wilson & Cook, 2020;
Zhang et al., 2015)
DG assumes that the test domain data are not available during training.
One thread of DG is to explore techniques from kernel methods (Muandet et al., 2013;
Niu et al., 2015; Erfani et al., 2016;
Li et al., 2017b).
Muandet et al.
(2013) propose a kernel-based optimization algorithm that learns an invariant transformation by minimizing the discrepancy among domains and preventing the loss of relationship between input and output features.
Another line of DG work is using end-to-end methods from deep learning: (a) reducing the differences of representations across domains through adversarial or similar techniques (Ghifary et al., 2015; Wang et al., 2017;
Motiian et al., 2017; Li et al., 2018b;c); (b) projecting out superﬁcial domain-speciﬁc statistics to reduce sensitivity to the domain (Wang et al., 2019); (c) fusing representations from an ensemble of models across domains (Ding & Fu, 2017; Mancini et al., 2018).
Meta-learning can also be applied to domain generalization, by dividing source domains into meta-training and meta-test sets, and aiming for a low generalization error on meta-test sets after training on meta-training sets (Balaji et al., 2018;
Dou et al., 2019; Li et al., 2018a; 2019a;b).
Recently, an extensive empirical survey of many DG algorithm (Gulrajani & Lopez-Paz, 2020) suggested that with current models and data augmentation techniques, plain ERM may be competitive with the state-of-the-art.
It is worth noting that Sun et al.
(2020) also propose an approach to learning latent causal factors for prediction.
However, their assumptions over the underlying causal graph are restricted due to two reasons: 1) they only consider the scenarios when X and Y are generated concurrently, which excludes the cases in which some part of X could also be affected by Y in some manner; 2)
They assume that the causal latent factors X
p and the non-causal latent factors
X c are independent when conditioning on E , that is,
X p ⊥⊥ X c | E .
In practice, during model learning, they actually further assume that X i ⊥⊥ X j | E for any i (cid:54) =
j so that VAE could be leveraged to learn the model.
In this sense, their approach has the same issue as the one in iVAE, i.e., unable to deal with the non-factorized cases.
This point is also veriﬁed in Table 2, where our approach greatly outperforms theirs.
B V ARIATIONAL A UTOENCODERS
We brieﬂy describe the framework of variational autoencoders (VAEs), which allows us to efﬁciently learn deep latent-variable models and their corresponding inference models (Kingma & Welling, 2013; Rezende et al., 2014).
Consider a simple latent variable model where O ∈ R d stands for an observed variable and X ∈ R n for a latent variable.
A VAE method learns a full generative model p θ ( O , X )
=
p θ ( O | X )
p θ ( X ) and an inference model q φ ( X | O ) , typically a factorized
Gaussian distribution whose mean and variance parameters are given by the output of a neural network with input O .
This inference model approximates the posterior p θ ( X | O ) , where θ is a vector of parameters of the generative model, φ a vector of parameters of the inference model, and p θ ( X ) is a prior distribution over the latent variables.
Instead of maximizing the data log-likelihood, we maximize its lower bound L VAE ( θ , φ ) :
[Equation]
where we have used Jensen’s inequality, and KL ( ·||· ) denotes the Kullback-Leibler divergence between two distributions.
C D ERIVATION C. I VAE
In Section 2.1, the evidence lower bound of iVAE is deﬁned by
[Equation]
C. S CORE M ATCHING
In Section 4.1, we follow Hyvärinen (2005) and use a simple trick of partial integration to simplify the evaluation of the score matching objective L SMphase1 in Eq.
(10) of the main text:
[Equation]
const.
where the last equality is due to the Theorem 1 in Hyvärinen (2005).
D D EFINITIONS FOR SEM AND IRM
Deﬁnition 2.
A structural equation model (SEM) C :=
( S , N ) governing the random vector X =
( X 1 , . . .
, X d ) is a set of structural equations:
[Equation]
where Pa ( X i ) ⊆
{ X 1 , . . .
, X d } \ { X i } are called the parents of X i , and the N i are independent noise random variables.
We say that “ X i causes X j ” if X i ∈
Pa ( X j ) .
We call causal graph of
X to the graph obtained by drawing i) one node for each X i , and ii) one edge from X i to X j if
X i ∈
Pa ( X j ) .
We assume acyclic causal graphs.
Deﬁnition 3.
Consider a SEM C :
= ( S , N ) .
An intervention e on C consists of replacing one or several of its structural equations to obtain an intervened SEM C e :=
( S e , N e ) , with structural equations:
[Equation]
The variable X e is intervened if S i (cid:54)
=
S ei or N i (cid:54) =
N ei .
Deﬁnition 4.
Consider a structural equation model (SEM) S governing the random vector ( X 1 , . . .
, X n , Y ) , and the learning goal of predicting Y from X .
Then, the set of all environments E all ( S ) indexes all the interventional distributions P ( X e , Y e ) obtainable by valid interventions e .
An intervention e ∈ E all ( S ) is valid as long as (i) the causal graph remains acyclic,
(ii) E
[ Y e | Pa ( Y )] =
E [ Y | Pa ( Y )] , and (iii) V [ Y e | Pa ( Y )] remains within a ﬁnite range.
E D EFINITIONS AND L EMMAS FOR THE E XPONENTIAL F AMILIES
Deﬁnition 5.
( Exponential family )
A multivariate exponential family is a set of distributions whose probability density function can be written as
[Equation]
(13) where Q : X → R is the base measure, Z ( θ ) is the normalizing constant, T : X → R k is the sufﬁcient statistics, and θ ∈ R k is the natural parameter.
The size k ≥ n is the dimension of the sufﬁcient statistics T and depends on the latent space dimension n .
Note that k is ﬁxed given n .
Deﬁnition 6.
( Strongly exponential distributions )
A multivariate exponential family distribution
[Equation]
is strongly exponential, if
[Equation]
where l is the Lebesgue measure.
The density of a strongly exponential distribution has almost surely the exponential component and can only be reduced to the base measure on a set of measure zero.
Note that all common multivariate exponential family distributions (e.g. multivariate Gaussian) are strongly exponential.
F D EFINITIONS FOR I DENTIFIABILITY
Deﬁnition 7.
Let Θ be the domain of the parameters θ = { f , T , λ } .
Let ∼ be an equivalence relation on Θ .
A deep generative model is said to be ∼ –identiﬁable if
[Equation]
The elements in the quotient space Θ \ ∼ are called the identiﬁability classes.
Deﬁnition 8.
Let ∼ A be an equivalence relation on Θ deﬁned by:
[Equation]
where A ∈ R k × k is an invertible matrix and c ∈ R k is a vector.
Deﬁnition 9.
Let ∼ P be an equivalence relation on Θ deﬁned by:
[Equation]
where P ∈ R k × k is a block permutation matrix and c ∈ R k is a vector
. G P HASE 2: D ISCOVERING S INGLE C AUSE G. S OME S PECIAL C ASES IN M ULTI -C AUSE S ETTINGS
These conditional independence tests can be performed in parallel to largely accelerate the testing procedure.
Note that, in practice, it might occur that there exist some X i which is independent of any other X j when conditioning on Y and E .
It is probably because such X i is a deterministic
X i
Y
E
O     (a)
X i
Y
E
O (b)
X i
Y
E
O (c)
X i
Y
E
O (d)
X i
Y
E
O (e)
X i
Y
E
O (f)
X i
Y
E
O (g)
X i
Y
E
O (h)
X i
Y
E
O (i)
X i
Y
E
O (j)
X i
Y
E
O (k) transformation of Y .
In this special case, we can use IGCI (Daniusis et al., 2012;
Janzing et al., 2012) to determine whether or not X i is a cause of Y .
Also note that, in some scenarios in which the dependence signals between a pair of causal latent variables might be weak due to the data issue, we can test the conditional independence of such a latent variable with all the other causal latent variables.
If it is conditionally dependent on more than half of them or its average p-value is larger than the pre-speciﬁed threshold, we will select it as one cause of Y .
G. D ISCOVERING S INGLE C AUSE
In the single-cause case, by following Wang et al.
(2014), we leverage the MB-by-MB algorithm to ﬁrst construct a local structure around Y and then discover the single parent of Y according to the constructed local graph.
One obvious advantage of this approach is in efﬁciency, because there is no need to construct the whole causal graph containing all the latent variables and Y .
We could have an even more efﬁcient solution to the single-cause case in some special scenarios where we assume that X i ⊥⊥ X j | Y , E for any i (cid:54) =
j .
In fact, this assumption covers more scenarios than the common assumption that X i ⊥⊥ X j for i (cid:54) =
j in latent variable models, e.g., disentanglement (Bengio et al., 2013; Locatello et al., 2019), autoencoders (Kingma & Welling, 2013; Rezende et al., 2014), ICA (Comon, 1994; Hyvärinen & Oja, 2000), etc.
If Y is caused by at most one of X i and
X j , and no matter whether X i and X j are caused by E or not, then X i ⊥⊥ X j | Y , E holds, but we may well have X i (cid:54)⊥⊥ X j (e.g., if Y causes both X i and X j , or if there is a chain X
i → Y → X j ).
If Y causes or is caused by at most one of X i and X j , and at most one of X i and X j is caused by
E , then both X i ⊥⊥ X j and X i ⊥⊥
X j | Y , E hold.
If { Y , X i , X j } form a collider X i → Y ← X j , and no matter whether X i and X j are caused by E or not, then X i ⊥⊥ X j | E hold, but we may have
X i (cid:54)⊥⊥ X j (e.g., when both X i and X j are caused by E ).
Under this assumption, we are able to separately look into each X i given Y and E , without considering any other X j . Fig.
7a shows that there exist only ﬁve possible connections between X i ,
Y , E , and O .
Among them, only the arrow from X i to O must exist, because O is generated from
X .
The other four arrows might not be present, with the exception that there must be at least one connection between X i and Y or E (Assumption 1a).
This leaves ten possible structures, shown in
Figs.
3b-3k.
Given data { ˆ X i , Y , E , O } in which the ˆ X i are obtained using q φ ( X | O , Y , E ) (for example, as given by the mean of this distribution), we are able to distinguish all ten structures in Figs.
3b-3k by using causal discovery algorithms (Peters et al., 2017; Zhang et al., 2017; Huang et al., 2020) and performing conditional independence tests (Spirtes et al., 2000; Zhang et al., 2012).
This is summarized in Proposition 2 below.
Its proof can be found in Appendix H, which also describes the speciﬁc assumptions made.
In practice, we can assess in parallel whether or not each X i is a direct cause of Y , which accelerates this phase signiﬁcantly.
Proposition 2.
Under the assumptions stated in Appendix H, the ten structures shown in Figs.
3b-3k can be identiﬁed using causal discovery methods consistent in the inﬁnite sample limit.
Note that there are only four cases in which X i is a parent of Y (i.e., Figs.
3b, 3e, 3g, and 3j).
We can identify these by applying rules 1.2 , 1.6 , 2.1 , and 3.1 from Appendix H.5.
H P ROOFS H. P ROOF OF T HEOREM
The proof of this theorem consists of three parts.
In Part I, we prove that the parameters θ are ∼ A identiﬁable (Deﬁnition 8) by using assumption (i), the ﬁrst half of assumption (ii), and assumption (iv) of Theorem 1.
In Part II, based on the result in Part I, we further prove that the parameters θ are ∼ P identiﬁable (Deﬁnition 9) by additionally using Assumption 2, the second half of assumption (ii) and assumption (iii) of Theorem 1.
In Part III, we combine the results (Theorems 4&5) in both Part I and Part II into one theorem (Theorem 1), which completes the proof.
It is worth noting that, compared to the proof in iVAE, the main changes in our proof consist of
• Part I, In step III.
– It has been updated to account for vectors of sufﬁcient statistics whose entries can be arbitrary functions of all entries in the random variable vector, while in the previous proof the sufﬁcient statistics contained entries that are functions of individual entries in the random variable vector.
– The assumption of “The sufﬁcient statistics in T are all linearly independent.” is not required in our proof, but it is in the proof of iVAE.
• Part II.
– It has been updated to account for the part of the sufﬁcient statistics which is the output of a deep neural network with ReLU activation functions.
H.1. P ART I
For notational simplicity, in the proof of this part we denote ( Y , E ) by U .
Hence, our generative model deﬁned according to Eqs.
(6-8) in the main text now becomes:
[Equation]
Theorem 4.
Suppose that we observe data sampled from a deep generative model deﬁned according to Eqs.
(19-21) with parameters ( f , T , λ ) .
Assume that
(i)
The set { O ∈
O | ϕ ε ( O ) = 0 } has measure zero, where ϕ ε is the characteristic function of the density p ε deﬁned in Eq.
(20); (ii) The mixing function f in Eq.
(20) is injective; (iii) There exist k + 1 points U 0 , U 1 , · · · ,
U k ∈ U such that the matrix
[Equation]
is invertible.
Then the parameters { f , T , λ } are ∼ A identiﬁable.
Proof.
Deﬁne vol ( B ) =
(cid:112) det( B T B ) for any full column rank matrix B .
Suppose that we have two sets of parameters θ =
( f , T , λ ) and ˜ θ =
( ˜ f , ˜ T , ˜ λ ) such that p θ ( O | U ) =
p ˜
θ ( O | U ) , ∀ ( O , U ) ∈
O × U .
We want to show θ ∼ A ˜ θ .
Step I.
The proof of this step is similar to Step I in the proof of Theorem 1 in Khemakhem et al.
(2020a).
We transform the equality of the marginal distributions over observed data into the equality of noise-free distributions.
For all pairs ( O , U ) ∈
O × U , we have
[Equation]
where
• in Eq.
(26), J denotes the Jacobian, and we made the change of variable ¯ O = f
( X ) on the left hand side, and ¯ O = ˜ f
( X ) on the right hand side.
• in Eq.
, we introduced ˜ p f , T , λ , U [Equation] O ) (cid:44) p T , [Equation]
f
[Equation]
1 ( O )
[Equation]
U )
[Equation]
on the left hand side, and similarly on the right hand side.
• in Eq.
(28), we used ∗ for the convolution operator.
• in Eq.
(29), we used F [ . ] to designate the Fourier transform, and where ϕ ε = F [ p ε ] (by deﬁnition of the characteristic function).
• in Eq.
(30), we dropped ϕ ε ( ω ) from both sides as it is non-zero almost everywhere (by assumption (i)).
Step II.
In this step, we remove all terms that are either a function of O or U .
First, by replacing both sides of Eq.
(31) by their corresponding expressions from Eq.
(32), we have
[Equation]
Then, by taking logarithm on both sides of Eq.
(33) and replacing p T , λ by its expression from Eq.
(21), we obtain
[Equation]
Let U 0 , U 1 , ·
· · , U k ∈
U be the k + 1 points deﬁned in assumption (iii).
We evaluate the above equation at these points to obtain k + 1 equations, and subtract the ﬁrst equation from the remaining k equations to obtain:
[Equation]
Let L be deﬁned as in assumption (iii) and ˜ L deﬁned similarly for ˜ λ .
Note that L is invertible by assumption, but ˜ L is not necessarily invertible.
Letting b ∈ R k in which b l = log
˜ Z ( U 0 )
Z ( U l )
˜
Z ( U l ) Z ( U 0 ) , we have
[Equation]
Left multiplying both sides of the above equation by L − T gives
[Equation]
where A = L − T ˜ L ∈ R k × k and c =
L − T b ∈ R k .
Step III.
To complete the proof, we need to show that A is invertible.
Let X l ∈ X , O l = f
( X l ) , l = 0 , · · · , k .
We evaluate Eq. (37) at these k + 1 points to obtain k + 1 equations and subtract the ﬁrst equation from the remaining k equations to obtain
[Equation]
We need to show that for a given X 0 ∈ X , there exist k points X 1 , · · · , X k ∈ X such that the columns of R are linearly independent.
Suppose, for contradiction, that the columns of R would never be linearly independent for any choice of X 1 , · · · , X k ∈ X .
Then the function g ( X ) (cid:44)
T ( X ) − T ( X 0 ) would live in a k − 1 or lower dimensional subspace, and thus we could ﬁnd a nonzero vector λ ∈ R k orthogonal to that subspace.
This would imply that (cid:104) T ( X ) − T ( X 0 ) , λ (cid:105) =and thus (cid:104) T ( X ) , λ (cid:105) =
(cid:104) T ( X 0 ) , λ (cid:105) =
const, ∀ X ∈ X , which contradicts the assumption that the prior is strongly exponential (Deﬁnition 6).
Therefore, we have shown that there exist k + 1 points
X 0 , X 1 , · · · , X k ∈ X such that R is invertible.
Since R =
A ˜ R and A is not a function of X , A must be invertible.
This completes the proof.
H.1. P ART II
Theorem 5.
Suppose that all assumptions in Theorem 4 hold.
Let the sufﬁcient statistics
T ( X ) =
[ T f ( X ) T , T NN ( X ) T ] T given by the concatenation of a) the sufﬁcient statistics T f
( X ) =
[ T 1 ( X 1 ) T , · · · , T n ( X n ) T ] T of a factorized exponential family, where all the T i ( X i ) have dimension larger or equal to 2, and b)
the output T NN
( X ) of a neural network with ReLU activations.
(note that a neural network with ReLU activation has universal approximation power and should be able to capture any dependencies of interest).
Let k (cid:48) be the dimension of T f and thus that k (cid:48) ≥ 2 n .
Assume that
(i) the sufﬁcient statistics T f have all second-order own derivatives; (ii) the mixing function f has all second-order cross derivatives.
Then the parameters { f , T , λ } are ∼ P identiﬁable.
Proof.
Let v = ˜ f
− 1 ◦ f : X → X .
Since all assumptions in Theorem 4 hold, we have
[Equation]
) where A ∈ R k × k is invertible.
We want to show that A is a block permutation matrix.
Step I.
In this step, we show that v is a componentwise function.
First we differentiate both sides of Eq. (39) with respect to X s and X t ( s (cid:54) =
t ) to obtain ∂ T ( X ) ∂X s =
[Equation]
By construction, the second-order cross derivatives of T and ˜ T are all zero.
Therefore, we have
[Equation]
The above equation can be written in the matrix-vector form:
[Equation]
where
[Equation]
45) and
[Equation]
[Equation]
Now by concatenating ˜
[Equation]
and
[Equation] [Equation]
we obtain
[Equation]
Finally, we take the rows of ˜ T (cid:48)(cid:48)(cid:48) ( X ) that corresponds to the factorized strongly exponential family distribution part and denote them by ˜ T (cid:48)(cid:48)(cid:48)
f ( X ) ∈ R k (cid:48) × 2 n .
By Lemma 5 in the iVAE paper (Khemakhem et al., 2020a) and the assumption that k (cid:48) ≥ 2 n , we have that the rank of ˜ T (cid:48)(cid:48)(cid:48) f
( X ) is 2 n .
Since k ≥ k (cid:48) ≥ 2 n , the rank of ˜ T (cid:48)(cid:48)(cid:48) ( X ) is also 2 n .
Since the rank of A is k , the rank of A ˜ T (cid:48)(cid:48)(cid:48) ( X ) ∈
R k × 2 n is 2 n .
This implies that v (cid:48)(cid:48)(cid:48) s,t ( X ) must be a zero vector.
In particular, we have that v (cid:48) s,t ( X ) = 0 , ∀ s (cid:54) =
t .
Therefore, we have shown that v is a componentwise function.
Step II.
To complete the proof, we need to show that A is a block permutation matrix.
Without loss of generality, we assume that the permutation in v is the identity.
That is v ( X ) =
[ v 1 ( X 1 ) , · · · , v n ( X n )] T for some nonlinear univariate scalar functions v 1 , · · · , v n .
Since f and ˜ f are bijective, we have that v is also bijective and v − 1 ( X ) =
[ v − 1 1 ( X 1 ) , · · · , v − 1 n ( X n )] T .
We denote ¯ T ( v ( X ))
=
˜ T ( v ( X ))
+ A − 1 c and plug it into Eq. (39) to obtain T ( X ) =
A ¯ T ( v ( X )) .
Applying v − 1 to the variables X at both sides gives
[Equation]
Let t be the index of an entry in the sufﬁcient statistics T that corresponds to the the factorized strongly exponential family distribution part T f .
For all s (cid:54) =
t , we have
[Equation]
Since the entries of ˜ T are linearly independent (if they were not linearly independent, then ˜ T can be compressed into a smaller vector by removing the redundant entries), we have that a tj is zero for any j such that ∂ ¯ T ( X ) j ∂X s (cid:54) = 0 .
This includes the entries j in the sufﬁcient statistics ˜ T that corresponds to 1) the factorized strongly exponential family distribution part which do not depend on X t ; and 2) the neural network part.
Therefore, when t is the index of an entry in the sufﬁcient statistics T that corresponds to factor i in the factorized strongly exponential family distribution part T f , the only non-zero a tj are the ones that map between T i ( X i ) and ¯ T i ( v i ( X i )) , where T i are the factors in T f that only depend on X i and ¯ T i is deﬁned similarly.
Therefore, we can construct an invertible submatrix A (cid:48)
i with all non-zero elements a tj for all t that corresponds to factor i , such that
[Equation]
where ˜ T i are the factors in ˜ T
f that only depends on X i , and
c i are the corresponding elements of c .
This means that the matrix A is a block permutation matrix.
For each i = 1 , · · · , n , the block A (cid:48) i of
A afﬁnely transforms T i ( X i ) into ˜ T
i ( v i ( X i )) .
There is also an additional block A (cid:48) which afﬁnely transforms T NN ( X ) into ˜ T NN
( v ( X )) .
This completes the proof.
H.1. P ART III
Now we combine Theorem 4 in Part I and Theorem 5 in Part II into one theorem, which completes the proof of Theorem 1.
H. P ROOF OF T HEOREM
We recall that the loss function in Phase 1 is as follows:
[Equation]
where
[Equation]
Proof.
If the family of q φ ( X | O , Y , E ) is ﬂexible enough to contain p θ ( X | O , Y , E ) , then by optimizing the loss over its parameter φ , we will minimize the score matching term L SMphase1 in
Eq.
(55), which will eventually reach zero.
If we assume that the model is not degenerate and that q φ > 0 everywhere, then having that L SMphase1 = 0 implies that ∇ X log q φ ( X | O , Y , E ) and ∇ X log p T , λ
( X | Y , E ) are equal.
This implies that log q φ ( X | O , Y , E ) =
log p T , λ
( X | Y , E ) + c for some constant c .
But c is necessarily 0 because both q φ ( X | O , Y , E ) and p T , λ
( X | Y , E ) are pdf’s.
Therefore, the VAE term L VAEphase1 in Eq.
(54) will be equal to the log-likelihood, meaning that the loss L phase1 in Eq.
(53) will be equal to the log-likelihood.
Under this circumstance, the estimation in Eq.
(53) inherits all the properties of maximum likelihood estimation (MLE).
In this particular case, since our identiﬁability is guaranteed up to a permutation and componentwise transformation, the consistency of MLE means that we converge to the true parameter θ ∗ up to a permutation and componentwise transformation in the limit of inﬁnite data.
Because true identiﬁability is one of the assumptions for MLE consistency, replacing it by identiﬁability up to a permutation and componentwise transformation does not change the proof but only the conclusion.
H. P ROOF OF T HEOREM
Proof.
Theorem 1 and Theorem 2 guarantee that in the limit of inﬁnite data, NF-iVAE can learn the true parameters θ ∗ :=
( f ∗ , T ∗ , λ ∗ ) up to a permutation and componentwise transformation of the latent variables.
Let ( ˆ f , ˆ T , ˆ λ ) be the parameters obtained by NF-iVAE.
We, therefore, have ( ˆ f , ˆ T , ˆ λ ) ∼ P ( f ∗ , T ∗ , λ ∗ ) , where ∼ P denotes the equivalence up to a permutation and componentwise transformation.
If there were no noise, this would mean that the learned ˆ f transforms
O into ˆ X
=
ˆ f − 1 ( O ) that are equal to X ∗ =
( f ∗ ) − 1 ( O ) up to a permutation and componentwise transformation (Deﬁnition 9).
If with noise, we obtain the posterior distribution of the latent variables up to an analogous indeterminacy.
H. P ROOF OF P ROPOSITION
Proof.
Theorem A.1 in (Arjovsky, 2021) has showed that i) any predictor w ◦
Φ with optimal OOD generalization uses only Pa ( Y ) to compute Φ ; ii) the classiﬁer w in this optimal predictor can be estimated using data from any environment e for which the distribution of Pa ( Y e ) has full support; iii) the optimal predictor will be invariant across E all .
In iCaRL, the hypotheses of Theorems 1 andand Assumption 1 guarantee that Pa ( Y ) can be recovered by ﬁrst identifying the latent variables X from O , Y and E and then discovering the direct causes of Y through solving Eq.
Furthermore, since the conditional prior in Eq.
(5) of the main text has full support, the distribution of Pa ( Y e ) always has full support.
Also, under Assumption 1 we have that p ( Y | Pa ( Y )) is invariant across E all .
Hence, the classiﬁer w in this optimal predictor can be estimated using data from any environment e .
We therefore have that the resulting optimal predictor will be invariant across E all .
This completes the proof.
H. P ROOF OF P ROPOSITION
Proof.
The following rules can be independently performed to distinguish all the 10 structures shown in Figs.
3b-3k.
For clarity, we divide them into three groups.
Note that, since these rules rely on different algorithms of causal discovery and conditional independence tests, unless stated otherwise, we assume that the assumptions of these algorithms are satisﬁed during the proof process.
Group 1 All the six structures in this group can be discovered only by performing conditional independence tests.
• Rule 1.1 If X i ⊥⊥ Y , X i (cid:54)⊥⊥ E , and E ⊥⊥ Y , then Fig.
3d is discovered.
• Rule 1.2 If X i (cid:54)⊥⊥ Y , X i ⊥⊥ E , and E (cid:54)⊥⊥ Y , then Fig.
3g is discovered.
• Rule 1.3 If X i (cid:54)⊥⊥ Y , X i (cid:54)⊥⊥ E , and E ⊥⊥ Y , then Fig.
3f is discovered.
• Rule 1.4 If X i (cid:54)⊥⊥ Y , X i (cid:54)⊥⊥ E , E
cid:54)⊥⊥ Y , and X i ⊥⊥ Y | E , then Fig.
3i is discovered.
Rule 1.5 If X i (cid:54)⊥⊥ Y , X i (cid:54)⊥⊥ E , E (cid:54)⊥⊥ Y , and X i ⊥⊥ E | Y , then Fig.
3h is discovered.
• Rule 1.6 If X i (cid:54)⊥⊥ Y , X i (cid:54)⊥⊥ E , E (cid:54)⊥⊥ Y , and Y ⊥⊥ E | X i , then Fig.
3e is discovered.
Group
If X i (cid:54)⊥⊥ Y , X i ⊥⊥ E , and E ⊥⊥ Y , then we can discover both Fig.
3b and Fig.
3c.
These two structures cannot be further distinguished only by conditional independence tests, because they come from the same Markov equivalence class.
Fortunately, we can further distinguish them by running binary causal discovery algorithms (Peters et al., 2017), e.g., ANM (Hoyer et al., 2009) for continuous data and CDS (Fonollosa,
2019) for continuous and discrete data.
• Rule 2.
If X i (cid:54)⊥⊥ Y , X i ⊥⊥ E , and E ⊥⊥ Y , and a chosen binary causal discovery algorithm prefers X i
→ Y to X i ← Y , then Fig.
3b is discovered.
• Rule 2.2 If X i (cid:54)⊥⊥ Y , X i ⊥⊥ E , and E ⊥⊥ Y , and a chosen binary causal discovery algorithm prefers
X
i ← Y to X i
→ Y , then Fig.
3c is discovered.
Group
If X i (cid:54)⊥⊥ Y , X i (cid:54)⊥⊥ E , E (cid:54)⊥⊥ Y , X i (cid:54)⊥⊥ Y | E , X i (cid:54)⊥⊥ E | Y , and Y (cid:54)⊥⊥ E | X i , then we can discover both Fig.
3j and Fig.
3k.
These two structures cannot be further distinguished only by conditional independence tests, because they come from the same Markov equivalence class.
They also cannot be distinguished by any binary causal discovery algorithm, since both X i and Y are affected by E .
Fortunately, Zhang et al.
(2017) provided a heuristic solution to this case based on the invariance of causal mechanisms, i.e., P ( cause ) and P ( effect | cause ) change independently.
The detailed description of their method is given in Section 4.2 of Zhang et al.
For convenience, here we directly borrow their ﬁnal result.
Zhang et al.
(2017) states that determining the causal direction between X i and Y in Fig.
3j and Fig.
3k is ﬁnally reduced to calculating the following term:
[Equation]
(56) where (cid:104)·(cid:105) denotes the sample average, ¯ P ( Y
|
X i ) is the empirical estimate of P ( Y | X i ) on all data points, and (cid:104) ˆ P ( Y | X i ) (cid:105) denotes the sample average of ˆ P ( Y | X i ) , which is the estimate of P ( Y | X i ) in each environment.
We take the direction for which ∆ is smaller to be the causal direction.
• Rule 3.1 If X i (cid:54)⊥⊥ Y , X i (cid:54)⊥⊥ E , E
(cid:54)⊥⊥ Y , X i (cid:54)⊥⊥ Y | E , X i (cid:54)⊥⊥ E | Y , Y (cid:54)⊥⊥ E | X i , and ∆ X i
→ Y is smaller than ∆ Y → X i , then Fig.
3j is discovered.
• Rule 3.2 If X i (cid:54)⊥⊥ Y , X i (cid:54)⊥⊥ E , E (cid:54)⊥⊥ Y , X i (cid:54)⊥⊥ Y | E , X i (cid:54)⊥⊥ E | Y , Y (cid:54)⊥⊥ E | X i , and ∆ Y → X i is smaller than ∆ X
i → Y , then Fig.
3k is discovered.
I I LLUSTRATIONS FOR M ODEL L EARNING
As described in Section 3.3, in the ground truth model, the prior for each X i ∈
I p is either p ( X i | E ) or p ( X i ) , depending on whether X i is caused by E or not.
By contrast, in the NF-iVAE model the prior is p ( X | Y , E ) .
Is this going to affect the identiﬁability of the latent variables?
Well, in practice not because the posterior distribution for X given O , Y and E would be equivalent in both models (up to identiﬁability guarantees).
J D ATASETS
For convenience and completeness, we provide descriptions of Colored MNIST Digits and Colored
Fashion MNIST here.
Please refer to the original papers (Arjovsky et al., 2019; Ahuja et al., 2020a;
Gulrajani & Lopez-Paz, 2020; Venkateswara et al., 2017) for more details.
J. S YNTHETIC D ATA
For the nonlinear transformation, we use the MLP:
• Input layer: Input batch (batch size, input dimension)
• Layer 1: Fully connected layer, output size = 6, activation = ReLU • Output layer: Fully connected layer, output size =
J. C OLORED MNIST D IGITS
We use the exact same environment as in Arjovsky et al.
Arjovsky et al.
(2019) propose to create an environment for training to classify digits in MNIST data 9 , where the images in MNIST are now colored in such a way that the colors spuriously correlate with the labels.
The task is to classify whether the digit is less than 5 (not including 5).
There are three environments (two training containing 30,000 points each, one test containing 10,000 points)
We add noise to the preliminary label ( ˜ y =
if the digit is between 0-4 and ˜ y
= 1 if the digit is between 5-9) by ﬂipping it withpercent probability to construct the ﬁnal labels.
We sample the color id z by ﬂipping the ﬁnal labels with probability p e , where p e is 0 . 2 in the ﬁrst environment, 0 . 1 in the second environment, and 0 .in
the third environment.
The third environment is the testing environment.
We color the digit red if z = 1 or green if z = 0 .
J. C OLORED F ASHION MNIST
We modify the fashion MNIST dataset 10 in a manner similar to the MNIST digits dataset.
Fashion
MNIST data has images from different categories: “t-shirt”, “trouser”, “pullover”, “dress”, “coat”, “sandal”, “shirt”, “sneaker”, “bag”, “ankle boots”.
We add colors to the images in such a way that the colors correlate with the labels.
The task is to classify whether the image is that of foot wear or a clothing item.
There are three environments (two training, one test)
We add noise to the preliminary label ( ˜ y = 0 : “t-shirt”, “trouser”, “pullover”, “dress”, “coat”, “shirt” and ˜ y = 1 : “sandal”, “sneaker”, “ankle boots”) by ﬂipping it with 25 percent probability to construct the ﬁnal label.
We sample the color id z by ﬂipping the noisy label with probability p e , where p e is 0 . 2 in the ﬁrst environment, 0 .in
the second environment, and 0 .
9 in the third environment, which is the test environment.
We color the object red if z = 1 or green if z = 0 .
K I N - DEPTH A NALYSIS ON S YNTHETIC D ATA K. C OMPARISONS WITH S TATE - OF - THE -A RT
We ﬁrst conduct a series of experiments on synthetic data generated according to an extension of the SEM in Model 1.
The extension is to map the variables X :=
( X 1 , X 2 ) into a 10 dimensional
[Table]
observation O through a linear or nonlinear transformation.
Our goal is to predict Y from O
[Table]
where
O = g ( X ) .
We consider three transformations: (a) Identity : g ( · ) is the identity matrix
I ∈
R 2 × 2 ,
i
[Table]
e., O = g ( X ) =
X .
(b) Linear : g ( · ) is a random matrix S ∈ R 2 × 10 , i.e.,
O = g ( X ) =
X · S .
(c)
Nonlinear : g ( · ) is given by a neural network with 2-dimensional input and 10-dimensional output, whose parameters are randomly set in advance.
Since this is a regression task, we use the mean squared error (MSE) as a metric of performance.
Note that, in this problem, there is only one causal latent variable X 1 , meaning that the conditional prior in Eq.
( ?? ) will not exhibit dependencies.
Because of this, in this case we do not include a T NN ( X ) term in our NF-iVAE conditional prior.
In the following section we do consider settings with many potential causal latent variables and, in that case, we do include T NN ( X ) in the NF-iVAE prior.
We consider a simple scenario in which we ﬁx σ 1 = 1 and σ 2 = 0 for all environments and only allow σ 3 to vary across environments.
In this case, σ 3 controls the spurious correlations between X 2 and Y .
Each experiment draws 1000 samples from each of the three environments σ 3 = { 0 . 2 , 2 , 100 } , where the ﬁrst two are for training and the third for testing.
We compare with the following baselines:
ERM, and two variants of IRMG: F-IRM Game (with Φ ﬁxed to the identity) and V-IRM Game (with variable Φ ).
As shown in Table 3, in the cases of Identity and Linear , our approach is better than IRMG but only comparable with ERM and IRM.
This might be because the identiﬁability result up to a simple nonlinear transformation renders the problem more difﬁcult by converting the original identity or linear problem into a nonlinear problem.
In the Nonlinear case, the gains of iCaRL are very clear.
We perform additional experiments to analyze the three phases of our approach, including a) an analysis of the importance of Assumptions 1&2 and of the necessity of NF-iVAE in Phase 1, b) how accurately the direct causes can be recovered in Phase 2, and c) how well the two optimization problems can be addressed in Phase 3, all of which can be found in Appendix K.
[Table]
K. V ERIFYING P HASE
Theorem 3 tells us that we can leverage NF-iVAE to learn the true latent variables up to a permutation and componentwise transformation.
We empirically verify this point by comparing the raw data X with the corresponding ˆ X inferred through the learned inference model in NF-iVAE.
Fig.
4 clearly shows that the inferred ˆ X is equal to X up to a permutation and componentwise transformation.
To show the importance of NF-iVAE, we conduct an experiment in which we replace NF-iVAE with the original VAE in Phase 1.
As shown in Table 4, the performance of iCaRL based on NF-iVAE signiﬁcantly outperforms the one based on VAE (aka., iCaRL-VAE).
It is worth noting that when
VAE is instead used in Phase 1, it usually occurs in Phase 2 that either all the dimensions or no dimension of ˆ X are identiﬁed as the parents of Y .
This is because all components of ˆ X are mixed together and will inﬂuence one another even when conditioning on Y and E .
K. V ERIFYING P HASE
To show how well our method can identify the direct causes of Y in Phase 2, we compare the ﬁnal performance when the identiﬁed direct cause (i.e., X 1 ) and the identiﬁed non-cause (i.e., X 2 ) are respectively used in Phase 3 to learn the predictor.
More speciﬁcally, we respectively treat X 1 and Xas Pa ( Y ) to solve Eq.
(12) and Eq.
(11) to learn the data representation Φ and the invariant classiﬁer w .
Note that, there might exist a permutation between the inferred { ˆ X 1 , ˆ X 2 } and the true { X 1 , X 2 } .
Fig.
5a and Fig.
5b show the results on the regression task, from which we can obviously see that the predictor elicited by the identiﬁed cause has a much better generalization performance.
The classiﬁcation result (i.e., Y is binarized according to whether it is positive or negative) in Fig.
5c further demonstrates this
[Table]
[Table]
K. V ISUALIZATION OF I DENTIFIABILITY OF NF- I VAE
To further verify identiﬁability of NF-iVAE
[Table]
we conduct a series of experiments on synthetic data generated according to the data generating process (Fig. 6b) corresponding to the causal graph shown in Fig.
6a.
The reason we choose this setting is that it is the simplest case satisfying our
[Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Equation]
[Figure]
, where U{·} denotes the discrete uniform distribution, N ( · ) the Gaussian distribution, and g ( · ) is given by a neural network with 2-dimensional input and 10-dimensional output, whose parameters are randomly set in advance.
requirements: a) For ease of visualization, the latent space had better be 2-dimensional; b) To introduce the non-factorized prior given Y and E (i.e., X i (cid:54)⊥⊥ X j | Y , E ), Y has at least two causes.
We draw 1000 samples from each of the four environments E = { 0 . 2 , 2 , 3 , 5 } , and thus the whole synthetic dataset consists of 4000 samples.
We compare with the following baselines: VAE Kingma & Welling (2013) (without identiﬁability guarantees), iVAE Khemakhem et al.
(2020a)
(with a conditionally factorized prior for identiﬁability), and ICE-BeeM Khemakhem et al.
(2020b).
It is worth noting that in ICE-BeeM the primary assumption leading to identiﬁability is similar to that in iVAE, where the base measure Q ( X ) could be arbitrary to capture the dependences between latent
[Figure]
[Table]
variables but the exponential term still has to factorize across components (dimensions).
All these are summarized in Table 5.
Clearly
[Table]
from the table we can see that our method has a more general assumption on the prior leading to identiﬁability.
This is also demonstrated empirically in Fig.
Our method iCaRL can recover the original data X up to a permutation and a simple componentwise transformation, whereas all the other methods fail because they are unable to handle the case in which
X i (cid:54)⊥⊥ X j | Y , E .
L I N -D EPTH A NALYSIS ON M ORE R EALISTIC D ATA L. C OLORED MNIST
We compare iCaRL with 1) IRM, 2) two variants of IRMG: F-IRM Game (with Φ ﬁxed to the identity) and V-IRM Game (with a variable Φ ), 3) three variants of ERM: ERM (on entire training data),
ERM e (on each environment e ), and ERM GRAYSCALE (on data with no spurious correlations), and 4) ROBUST MIN MAX (minimizing the maximum loss across the multiple environments).
Table 1 shows that iCaRL outperforms all other baselines on Colored MNIST.
However, this dataset seems more difﬁcult because even ERM GRAYSCALE, where the spurious correlation with color is removed, falls well short of the optimum. [Table]
L. PACS
We report the results on another one of the widely used realistic datasets for OOD generalization:
PACS (Li et al., 2017a). This dataset consists of 9 , 991 images of dimension (3 , 224 , 224) andclasses from four domains: art, cartoons, photos, and sketches.
We used the exact experimental setting that is described in Gulrajani & Lopez-Paz (2020). We provide results averaged over all possible train and test environment combination for one of the commonly used hyper-parameter tuning procedure: train domain validation. As shown in Table 7, iCaRL achieves state-of-the-art performance when compared to those most popular domain generalization alternatives. M I MPLEMENTATION D ETAILS M. J OINT T RAINING
As described in Section 4.1, we can jointly learn ( θ , φ ) by optimizing the following objective:
[Equation]
where the last equality is due to the equation in Appendix C.2, and ˆ f , ˆ T , ˆ λ , ˆ φ are copies of f , T , λ , φ that are treated as constants and whose gradient is not calculated during learning.
In practice, ˆ f , ˆ T , ˆ λ , ˆ φ can be easily implemented through either “ detach ” in PyTorch Paszke et al.
(2019) or “ stop_gradient ” in TensorFlow Abadi et al.
M. T HE G ENERAL N ON - FACTORIZED P RIOR
In the experiments, the general non-factorized prior in Assumption 2 is implemented as follows:
[Equation]
where (cid:104)· , ·(cid:105) is the dot product of two vectors, and concat ( · , · ) means the concatenation of two vectors.
Now let us explain each term in details.
Firstly, concat ( X , X 2 ) is a vector of the latent variables and their squared values, and NN ( Y , E ; param3 ) is a deep neural network parameterized by param3 that computes a vector of natural parameters as a function of Y and E .
Hence, the term (cid:10) concat ( X , X 2 ) , NN ( Y , E ; param3 ) (cid:11) is equivalent to the factorized exponential family, which also satisﬁes that each T i ( X i ) has dimension larger or equal to 2.
Secondly, NN ( X ; param1 ) is a neural network that receives as input a vector of latent variables and outputs another vector representing complicated nonlinear transformations of those variables.
NN ( Y , E ; param2 ) is another neural network that generates a corresponding vector of natural parameters.
Hence, the term (cid:10) NN ( X ; param1 ) , NN ( Y , E ; param2 ) (cid:11) will allow this prior to capture the dependencies between the latent variables X .
N H YPERPARAMETERS AND A RCHITECTURES
In this section, we describe the hyperparameters and architectures of different models used in different experiments.
N. S YNTHETIC D ATA
We used Adam optimizer for training with learning rate set to 1e-3 and batch size set to 128 .
N.1. ERM
Linear ERM
• Input layer: Input batch (batch size, input dimension)
• Output layer: Fully connected layer, output size =
Nonlinear ERM
• Input layer: Input batch (batch size, input dimension)
• Layer 1: Fully connected layer, output size = 6, activation = ReLU • Output layer: Fully connected layer, output size =
N.1. IRM
Linear Data Representation Φ
• Input layer: Input batch (batch size, input dimension)
• Output layer: Fully connected layer, output size =
Nonlinear Data Representation Φ • Input layer: Input batch (batch size, input dimension) • Layer 1: Fully connected layer, output size = 6, activation = ReLU • Output layer: Fully connected layer, output size =
N
F-IRM GAME
Linear Classiﬁer w
• Input layer: Input batch (batch size, input dimension)
• Output layer: Fully connected layer, output size =
Nonlinear Classiﬁer w
• Input layer: Input batch (batch size, input dimension)
• Layer 1: Fully connected layer, output size = 6, activation = ReLU • Output layer: Fully connected layer, output size =
N.1.
V
IRM GAME
Linear Data Representation Φ
• Input layer: Input batch (batch size, input dimension)
• Output layer: Fully connected layer, output size =
Nonlinear Data Representation Φ
• Input layer: Input batch (batch size, input dimension)
• Layer 1: Fully connected layer, output size = 6, activation = ReLU • Output layer: Fully connected layer, output size =
Linear Classiﬁer w
• Input layer: Input batch (batch size, 2) • Output layer: Fully connected layer, output size =
Nonlinear Classiﬁer
w • Input layer: Input batch (batch size, 2) • Layer 1: Fully connected layer, output size = 6, activation = ReLU • Output layer: Fully connected layer, output size = N.1.
I C A RL
NF-iVAE λ f -Linear Prior
• Input layer: Input batch (batch size, input dimension)
• Output layer: Fully connected layer, output size =
NF-iVAE λ f -Nonlinear Prior
• Input layer: Input batch (batch size, input dimension)
• Layer 1: Fully connected layer, output size = 6, activation = ReLU • Output layer: Fully connected layer, output size =
NF-iVAE Linear
Encoder • Input layer: Input batch (batch size, input dimension) • Mean Output layer: Fully connected layer, output size =• Log Variance Output layer: Fully connected layer, output size =
NF-
iVAE
Nonlinear
Encoder • Input layer: Input batch (batch size, input dimension)
• Layer 1: Fully connected layer, output size = 6, activation = ReLU • Mean Output layer: Fully connected layer, output size =• Log Variance Output layer: Fully connected layer, output size =
NF-iVAE Linear Decoder
• Input layer: Input batch (batch size, 2) • Mean Output layer: Fully connected layer, output size = output dimension • Variance Output layer: 0 .
01 × 1 , where 1 is a vector full of 1 with the length of output dimension
NF-iVAE Nonlinear Decoder
• Input layer: Input batch (batch size, 2) • Layer 1: Fully connected layer, output size = 6, activation = ReLU • Mean Output layer: Fully connected layer, output size = output dimension • Variance Output layer: 0 .
01 × 1 , where 1 is a vector full of 1 with the length of output dimension
Linear Classiﬁer w
• Input layer: Input batch (batch size, 1) • Output layer: Fully connected layer, output size =
Nonlinear
Classiﬁer w • Input layer: Input batch (batch size, 1) • Layer 1: Fully connected layer, output size = 6, activation = ReLU • Output layer: Fully connected layer, output size =
N. CMNIST AND CFMNIST
Considering that the results of most baselines come from IRMG (Ahuja et al., 2020a), for a fair comparison, we follow the same setting of IRMG in terms of hyper-parameters and validation considerations.
For example, the batch size is set to 256, and the learning rate is 10 − 4 .
We also did not use the test environment data for validation.
Please ﬁnd
more details in Ahuja et al.
(2020a).
NF-iVAE T NN -Prior
• Input layer: Input batch (batch size, input dimension)
• Layer 1: Fully connected layer, output size = 50,
activation = ReLU • Output layer: Fully connected layer, output size =NF-iVAE λ NN -Prior • Input layer: Input batch (batch size, input dimension) • Layer 1: Fully connected layer, output size = 50,
activation = ReLU • Output layer: Fully connected layer, output size =NF-iVAE λ f -Prior • Input layer: Input batch (batch size, input dimension) • Layer 1: Fully connected layer, output size = 50,
activation = ReLU • Output layer: Fully connected layer, output size =
NF-iVAE O -Encoder
• Input layer: Input batch (batch size, 2, 28, 28) • Layer 1: Convolutional layer, output channels = 32, kernel size = 3, stride = 2, padding = 1, activation = ReLU • Layer 2: Convolutional layer, output channels = 32, kernel size = 3, stride = 2, padding = 1, activation = ReLU • Layer 3: Convolutional layer, output channels = 32, kernel size = 3, stride = 2, padding = 1,
activation = ReLU • Output layer: Flatten NF-iVAE ( Y , E ) -Encoder
• Input layer: Input batch (batch size, input dimension)
• Output layer: Fully connected layer, output size = 100, activation = ReLU
NF-iVAE ( O , Y , E ) -Merger/
Encoder • Input layer: Input batch (batch size, input dimension)
• Layer 1: Fully connected layer, output size = 100, activation = ReLU • Mean Output layer: Fully connected layer, output size =• Log Variance Output layer: Fully connected layer, output size =
NF-iVAE
Decoder • Input layer: Input batch (batch size, 10) • Layer 1: Fully connected layer, output size = 32 × 4 × 4 , activation = ReLU • Layer 2: Reshape to (batch size, 32, 4, 4)
• Layer 3: Deconvolutional layer, output channels = 32, kernel size = 3, stride = 2, padding = 1, outpadding = 0, activation = ReLU • Layer 4: Deconvolutional layer, output channels = 32, kernel size = 3, stride = 2, padding = 1, outpadding = 1, activation = ReLU • Layer 5: Deconvolutional layer, output channels = 2, kernel size = 3, stride = 2, padding = 1, outpadding =• Mean Output layer: activation =
Sigmoid • Variance Output layer: 0 .
01 × 1 , where 1 is a matrix full of 1 with the size of 2 × 28 × 28 .
Classiﬁer w
• Input layer: Input batch (batch size, 50) • Layer 1: Fully connected layer, output size = 100, activation = ReLU • Output layer: Fully connected layer, output size = 1, activation =
Sigmoid
N.
VLCS AND PACS
We used the exact experimental setting that is described in Gulrajani & Lopez-Paz (2020).
Speciﬁcally, we trained our model over all possible train and test environment combination for one of the commonly used hyper-parameter tuning procedure: train domain validation.
We use ResNet-50 as an encoder and reverse the architecture of ResNet-50 as a decoder.
We set the number of the latent variables to n = 50 .
We do the hyperparameter search by exactly following the guides given in Gulrajani &
Lopez-Paz (2020).