Optimization-based Causal Estimation from Heterogenous Environments
Abstract
This paper presents an optimization approach to causal estimation.
In classical machine learning, the goal of optimization is to maximize predictive accuracy.
However, some covariates might exhibit non-causal association to the outcome.
Such spurious associations provide predictive power for classical ML, but prevent us from interpreting the result causally.
This paper proposes CoCo, an optimization algorithm that bridges the gap between pure prediction and causal inference.
CoCo leverages the recentlyproposed idea of environments.
Given datasets from multiple environments—and ones that exhibit enough heterogeneity—CoCo maximizes an objective for which the only solution is the causal solution.
We describe the theoretical foundations of this approach and demonstrate its effectiveness on simulated and real datasets.
Compared to classical ML and the recently-proposed IRMv1,
CoCo provides more accurate estimates of the causal model.
1. Introduction
Consider a classical machine learning (ML) problem.
We observe a dataset of x i , y i pairs; x contains p covariates and y is an outcome.
In classical ML, our goal is to be able to predict y from x .
We want to learn a model y =
f ( α (cid:62) x ) , inferring the unknown coefﬁcients α from the training data.
Suppose our goal is not purely predictive but is one of causal estimation .
Consider that the outcome y is drawn from a true data generating process (DGP) that involves direct causes , a subset of the p covariates.
For simplicity we assume linearity that the outcome is truly drawn from y
i = g ( β (cid:62) x ) + (cid:15) where β k (cid:54) =
for the direct causes, β k = 0 for the other covariates and (cid:15) is independent noise.
The causal goal is to estimate the true causal coefﬁcients β , including both its structure and its value.
Classical empirical risk minimization (ERM) cannot reliably solve this problem.
In its search for ﬁnding the best predictor, it will capitalize on spurious (non-causal) associations between the components of x and the outcome y , such as due to confounding or conditioning on a collider.
Consequently, ˆ α , the resulting estimate of the coefﬁcients, will be a biased estimate of the causal coefﬁcients β .
We develop CoCo, an optimization-based approach to estimate causal coefﬁcients.
We derive CoCo in two steps.
In the ﬁrst step, we posit a risk-based objective for which the causal coefﬁcients are one of several optima.
In the second step, we use the idea of invariant environments —multiple training datasets that leave the targeted causal coefﬁcients intact—to whittle down the number of optima and leave only the causal coefﬁcients.
The result is a practical algorithm that analyzes data from multiple environments to produce an estimate of the causal coefﬁcients.
On synthetic and realworld data, relative to ERM and invariant risk minimization (IRM) [1], CoCo better estimates the causal coefﬁcients and predicts more robustly on new data.
2. Causal estimation as optimization Assumptions.
Consider an observed dataset of n datapoints D =
( X , Y ) , where Y ∈ R n is the outcome variable, X
=
[ X 1 , · · · , X p ] ∈ R n × p are the covariates.
Each column X j ∈ R n , j ∈ { 1 , 2 , · · · , p } is the observations of the j -th covariate on the n units.
Assume the underlying DGP of the outcome variable follows a linear SEM [11] as [Equation]
[Equation]
where β ∈ R p is the causal coefﬁcients [12].
(We absorb the intercept term into x and β .)
Eq.
(1) assumes linearity; we extend to nonlinear causal models in Appendix A.
The causal coefﬁcient might contain some zeros, indicating covariates that are not causally connected to the outcome.
The indices of non-zero coefﬁcients are the support set of β , denoted S ⊂ { 1 , 2 , · · · , p } , where β j (cid:54)
=
for j ∈ S and β j = 0 for j / ∈ S .
The support set represents the set of direct causes of the outcome.
Assume the noise is zero-mean E [ (cid:15) ]
= 0 , the covariates and noise have ﬁnite variance Var [ x j ] , Var [ (cid:15) ] < ∞ for all j .
We do not specify the SEM for the covariates x and allow their joint distribution to be arbitrary.
Assume the causal covariates x S are independent of the noise, i.e. x S ⊥⊥ (cid:15) .
This independence assumption implies that there is no unmeasured confounding between the true causes and the outcome.
Note that we allow unmeasured confounding between noncauses x \ S (or spurious variables) and the outcome.
Fig.shows the examples of endogenous covariates.
Pure prediction is biased.
With data from Eq.
(1) our goal is to estimate β .
We minimize the L 2 risk as a function of α
[Equation]
with linear predictions
[Equation]
When the form of ˆ y is clear in context, we abbreviate the risk function notation as R ( α ) .
Under the potential existence of spurious variables, directly optimizing Eq.
(2) will produce biased estimates of the causal coefﬁcients.
Idealized causal optimization.
Suppose we do not know the coefﬁcients β , but we do know which covariates are direct causes, the set S .
A key observation, though a simple one, is that among the models that share the true causal structure, the causal model has the best predictive accuracy.
Lemma 1 (Causal Optimality) .
The causal model with α = β is the optimal solution of the following constrained optimization problem
[Equation]
The proof is in Appendix C.
Lemma 1 is conceptually straightforward, but it provides a direct connection between optimization and causal estimation.
Of course, in practice, we do not know which covariates are causal and which are not.
We will build on this idealized optimization problem to construct a tractable objective for causal estimation.
3. Relaxed optimization for causal estimation
In this section, we use directional derivative to create a new optimization objective for causal estimation.
Directional derivatives and feasible directions.
We ﬁrst review the ideas of directional derivatives and its application in constrained optimization [16].
Consider a unit-length vector v , where (cid:107) v (cid:107) 2 = 1 .
The directional derivative in the direction v is denoted as operator D v and is deﬁned to be the rate of change of a function in that direction.
The directional derivative, as a scalar, can be computed as the inner product of the gradient and the direction vector
[Equation]
where we denote an inner product as (cid:104)· , ·(cid:105) .
In Eq.
(4), denote the constraints as g j ( α ) =
α j = 0 for j / ∈ S .
(Recall S is the support set of β , the indices of the non-zero causal coefﬁcients.)
Given a parameter α and the optimization problem Eq.
(4), the directions that violate the constraints at the maximum rate are the gradient direction of the constraint function,
[Equation]
The feasible directions are deﬁned as the directions orthogonal to the gradient of the constraints [10].
The feasible directions for Eq.
(4) form a linear space U = span { e j : j ∈ S } .
The ﬁrst-order condition for a point α to be an optima is that the directional derivative in the feasible directions vanish [10; 16].
For problem Eq.
(4), this condition means
D v = 0 for each v ∈ U .
The ﬁrst order condition can be guaranteed by
[Equation]
because v ∈ U is a linear combination of the basis { e j } j ∈ S .
More compactly, these conditions can be written as
[Equation]
with ◦ as Hadamard product because β j (cid:54) = 0 for j ∈ S .
Relaxing the causal optimization.
Lemma 1 states that α = β is an optimal solution of the problem in Eq.
(4), which means that it satisﬁes the ﬁrst-order condition of
Eq.
Plugging α = β into this condition reveals that (cid:107)∇ R ( β ) ◦ β (cid:107) 2 = 0 .
This fact, in turn, means that the causal coefﬁcients β is an optimum of the following optimization problem,
[Equation]
Notice that Eq.
(8) is an optimization problem that is entirely a function of the observed data; it does not require knowing the set S of non-zero causal coefﬁcients.
Yet the true causal coefﬁcient β is an optimum of this problem.
We call the set of points that minimizes Eq.
(8) the plausible set F .
However, the plausible set does not only contain the causal coefﬁcients.
By Eq.
(8), the plausible set also contains the all-zero vector 0 , the OLS solution, and the points “in-between” these two solutions, those that zero a subset of the covariates and zero the gradient of the remainder.
Our next step is to whittle down the plausible set until the causal coefﬁcients become the unique solution.
4. Optimization with multiple environments
In this section, we describe the invariance property of causal coefﬁcients under interventions and how it helps restore the identiﬁability of the causal model via optimization.
Environments and invariance.
As discussed in § 3, the causal model in general is non-identiﬁable by optimization with i.i.d.
data.
To obtain identiﬁability, we turn to the settings where data from multiple environments are available.
Denote E as a set of environments.
An environment e ∈
E speciﬁes a DGP.
For linear model, data from multiple environments is generated by an SEM similar to Eq.
[Equation]
The independence assumptions and moments conditions in § 2 apply to all environments.
Across environments, the set of observed variables, the set of direct causes of the outcome, and the causal coefﬁcients remain the same.
The environments are a set of DGPs.
The environments are heterogeneous if the DGPs are different.
Heterogeneous environments can be constructed by (hard) interventions that ﬁx a variable at a speciﬁc value.
They can also be DGPs where the joint distributions of variables are different in nature, also known as soft interventions [6].
For example, when studying the effect of health measurements on the probability of cancer, the environments can be different hospitals where the data are collected from [17].
The key property that environments enjoy is the invariance [1; 12].
Invariance assumes conditional on the same value of direct causes, the expectation of the outcome is the same across environments, i.e.
[Equation]
for all e, e (cid:48) ∈ E .
Note that we ask that the distribution of covariates x e and noise (cid:15) e changes across e ∈ E , which makes spurious associations vary with environments.
Narrowing down the optima set by environments.
The invariance property motivates us to aggregate the optimization problems and get the CoCo objective:
[Equation]
Denote the solutions of each environment as F e := arg min α (cid:107)∇ R e ( α ) ◦ α (cid:107) 2 .
The solutions of the CoCo objective Eq.
(11) is the intersection of all F e s,
[Equation]
as long as the intersection is not empty; this fact is guaranteed by the invariance assumption with β ∈ F e for all e .
Because of the intersection, the size of F E shrinks with the increasing number of environments, i.e. |F E 1 | ≤ |F E 2 | if E 2 ⊂ E 1 .
The multiple environments and heterogeneity therein induce differences among the set of solutions and, as a result, narrow down the solution set of CoCo objective
Eq.
The plausible sets are visualized with examples in
Fig. 4 in Appendix F.
Yet, as a last step, we need to remove the all-zero vector from the solutions.
Removing non-informative solution from the optima set.
We propose two modiﬁcations to remove the zero vector from the objective solutions.
Suppose there is prior knowledge of the underlying causal graph.
Speciﬁcally, suppose a set C of covariates that are known to be independent of the unobserved noise (cid:15) of the outcome (e.g. some non-descendants of the outcome).
We modify the CoCo objective Eq.
(11) to be
[Equation]
where ˜ α = α ◦ ( 1 − 1 C )+ 1 C .
For the risk function Eq.
(2), ∇ R ( β ) j = E [ x j (cid:15) ] = 0 for j ∈ C .
So the causal coefﬁcient β remains as a solution to the modiﬁed objective Eq.
The solution set of Eq.
(13) is a subset of that of Eq.
(11) with the all-zero vector being removed.
The algorithm is summarized in Alg.
The theoretical results on causal identiﬁcation is presented in Appendix D. CoCo objective is related to the invariant risk minimization (IRM) [1]; we illustrate the mathematical connections in Appendix B.
If there is no prior knowledge of the graph, we can add the risk function as a regularization term to Eq.
(11) as
[Equation]
where λ r ≥ 0 controls the regularization strength.
The regularization encourages the causal solution β against the zero vector because the zero vector has lower predictive accuracy than the causal model by the DGP.
Solving Eq.
(14) input : Data D e = { Y e , X e } , X e ∈ R n e × p the function estimation Initialize α randomly while not converged do for e in E do Compute the gradient of the empirical risk: g e ( α ) = (1 /n e ) ∂∂ α (cid:80) n e i
=1 R e ( α ; y
[Equation]
[Equation]
[Equation]
[Equation]
) ˜
[Equation]
= α ◦ ( 1 − 1 C ) + 1 C
Compute the loss: L e ( α )
= (cid:107)
g
e ( α ) ◦ ˜
α
(cid:107)end
Update α ← α − η ∂∂ α (cid:80) e ∈E L e ( α ) with step size η end input regularizer coefﬁcients λ r , λ w .
output: Predictor f α ( · ) that is robust to interventions
Initialize α randomly while not converged do for e in E do Compute the gradient of the empirical risk: g e ( α ) =
(1 /n e ) ∂∂ α (cid:80) n e i =1 R e ( α | y ei , ˆ y ei ) , ˆ y ei = f
( x ei ; α )
Compute: L e ( α ) =
(cid:107)
g e
( α ) ◦ α (cid:107)(Optional step
) add weak condition L e ( α ) += λ w ( (cid:104)
g e ( α ) , α (cid:105) )
Add risk function as regularization:
L e ( α ) += λ r (1 /n e )( (cid:80) n e i =1 R e ( α | y ei , ˆ y ei )) end
Update α ← α − η ∂∂ α (cid:80) e ∈E L e ( α ) with step size η end produces a model that enjoys distributional robustness under interventions [13].
The algorithm is summarized in Alg. 2.
Generalizing to nonlinear model.
We extend CoCo to predictors that is a nonlinear mapping of linear combinations of covariates, i.e. ˆ y = f
( A x ) , in Appendix A. It includes fully connected neural network as a special case.
The key is to build a constrained optimization problem similar to
Eq.
(4) and show that it admits the causal coefﬁcient as an optimum.
The analysis presented in §§ 3 and 4 can then be applied to such nonlinear models.
5. Empirical Studies
We study CoCo on simulated and real data.
Across datasets, we ﬁnd that CoCo produces an unbiased estimate of the causal structure and coefﬁcients.
CoCo can generalize its predictive ability from observed environments to new environments.
The details about the data generation and algorithms implementation in this section are in Appendix E.
5.1. Linear Synthetic Data
We study causal inference with optimization-based methods.
The data are generated from 5 different graphs in Fig.
3 in
Appendix E, each including a spurious variable.
The ﬁve graphs test different scenarios: (1) independent causes; (2) observed mediator; (3) observed confounder and mediator; (4) observed confounder and unobserved mediator; (5) collider.
We evaluate with mean absolute error (MAE) between the estimation α and true coefﬁcients β .
The results are summarized in Fig.
It shows when the covariates have spurious associations with the outcome, the estimate of ERM is biased.
IRM with proper hyper-parameter performs well in cases 2,3 while it has a large error in other cases.
CoCo has small estimation error in all cases.
As an ablation study, we replace the strong penalty in CoCo objective Eq.
(13) with the weak penalty of Eq.
(20) and minimize (cid:80) e ∈E ( (
cid:104)∇ R e ( α ) , ˜ α (cid:105) ) 2 .
This method is labeled
Naive-CoCo.
The comparison between Naive-CoCo with
CoCo in cases 1-4 shows that it is crucial to design the objective based on strong condition Eq.
(8) instead of weak condition Eq.
5.2. Colored MNIST (CMNIST)
In this section, we study whether CoCo produce a nonlinear model that uses causal variables to make predictions.
The covariates are the colored digit image and the binary label y e corresponds to the digit being odd or even.
By the data construction, the relationship between the digit shape and y e is genuinely causative while the relationship between the color and y e is spurious.
A desired predictor is one that make prediction based on the digit shape rather than the color.
A predictor using color information can neither accurately predict the label y e at testing.
Empirical results.
The results are shown in Fig.
8 in
Appendix, and Table 1.
In Table 1, ERM has the lowest accuracy in both training and testing.
The reason might be it largely depends on the color information to predict rather than on the digit shape, whereas the label is generated from the digit shape.
The testing accuracy for IRM increases in the early stage of training but drops in the later stage.
We [Table] hypothesize that the model at ﬁrst improves the prediction by utilizing all information including the digit shape, but later it relies more on the color information to boost the predictive accuracy, which reduces the accuracy at the test time.
5.3. Natural Image Classiﬁcation
In this example, following Cloudera [4], we adapt the iWildCam 2019 dataset [2] that contains wildlife images taken in the wild.
The goal is to classify coyote and raccoon in images.
The image background, such as plants and rocks, might be predictive to the species but in a spurious way.
This spurious association changes across locations hence we consider the images taken at different locations as coming from different environments.
Based on the setting of
Cloudera [4], we use images from two locations as the training data, images from one location as the validation data, and images from another location as the test data,
The results are summarized in Table 1 and Fig.
8 in Appendix G. ERM has high accuracy in training but low accuracy at testing.
CoCo performs on par with or slightly better than IRM.
Comparing to ERM, both methods have a slight drop in training accuracy but signiﬁcantly higher testing accuracy.
CoCo has a small performance gap between training and testing, indicating that it is not predicting animal labels via information from image backgrounds, i.e., information that varies across environments.
6. Conclusion
This paper formulated causal estimation as an optimization problem.
Using directional derivatives, we proposed the
CoCo objective, a computationally tractable optimization method for estimating causal coefﬁcients with datasets from multiple environments.
We discussed the mathematical connection between CoCo and IRM.
In empirical studies, we found that CoCo produces accurate causal estimation and distributionally robust predictions.
CoCo is applicable to high dimensional data, and to linear and nonlinear models.
A. Extension to Non-linear Models
In the main paper §§ 2 to 4, we focus on linear SEMs and linear predictors.
Here we generalize these results to nonlinear models where the outcome is generated by a nonlinear function of a linear transformation of direct causes and the noise is additive.
The key is to build a constrained optimization problem similar to Eq.
(4) and show that it admits the causal coefﬁcient as an optimum.
The analysis presented in § 3 can then be applied to nonlinear models.
Suppose we have a collection of environments E , and for each e ∈ E , we observed i.i.d.
data for variables ( x e , y e ) , x e ∈ R p , y e ∈ R .
Suppose the underlying DGP is
[Equation]
(15) where S ⊂ { 1 , 2 , · · · , p } , (cid:15) e ⊥⊥ x eS and E [ (cid:15) e ] = 0 .
f :
R K → R is an arbitrary function mapping with parameters β
= ( A , γ ∗ ) where A ∈ R K ×| S | and γ ∗ ∈ R M .
When
K = 1 and f ( · ) is an identity mapping, Eq.
(15) reduces to the linear SEM.
Eq.
(15) can represent a process when the outcome is generated through a deep neural network (DNN), where K and A are the width and weights of the ﬁrst hidden layer respectively.
Assume the nonlinear predictor is
[Equation]
where B ∈ R K × p , γ ∈ R M and α =
( B , γ ) are the parameters to optimize.
We can re-write A x eS =
A Λ x e where Λ ∈ R | S |× p has the i-th row as e (cid:62) i if i ∈ S and as 0 (cid:62) p if i / ∈ S .
Let B ∗ =
A Λ where the j -th column of B ∗ is 0 K if j / ∈ S .
Then for square error R e ( α ) we have the following proposition.
Proposition 1 (Causal Optimality, nonlinear) .
The causal model B = B ∗ , γ = γ ∗ is an optima of the following problem
[Equation]
Proposition 1 greenlights the analysis in § 3.
It implies that the CoCo objective Eq.
(11) and extension Eq.
(14) can be used for nonlinear models.
When B = B ∗ , the multiplication B x e zeros out non-causal covariates x e \ S , which become independent of prediction ˆ y e .
In the nonlinear regime, due to high ﬂexibility, identiﬁcation can be difﬁcult.
Different parameterizations can represent similar mappings on the training data.
Hence we expect the identiﬁcation is up to an equivalent class [3; 8].
Another challenge is that a sufﬁciently ﬂexible model may memorize all outcome labels and become the ERM optimum for each environment, but this model is not causative.
To avoid such a solution, one approach is to collect overlapped environments, for example, those sharing the same domain of inputs [18].
B. Connection to invariant risk minimization
In this section, we discuss the connections and distinctions between CoCo and IRM.
Arjovsky et al.
[1] introduces
IRM that can learn robust representation in the presence of spurious associations between covariates and the outcome.
In particular, IRM considers a predictor f ( x ; α ) : R p (cid:55)→ R with parameter α .
In a setting similar to CoCo , it considers a set of heterogeneous environments E and for each e ∈ E , a risk function R e ( α ; y, ˆ y ) .
Based on the intuition that invariant predictor induces invariant features, IRM proposes the following objective to ﬁnd an invariant model
[Equation]
where w ( · ) is a mapping from the range of f ( · ) to ˆ y .
For tractable computation, Arjovsky et al.
[1] further introduces the IRMv1 objective:
[Equation]
where λ > 0 and w is simpliﬁed as a dummy scalar variable.
The IRMv1 objective consists of an empirical risk term and an invariant risk term.
We make the connection between IRM and the constrained optimization in Lemma 1.
In § 3, we obtain the ﬁrst-order optimality condition Eq.
(7) from the directional derivative in the feasible directions { e j } j ∈ S .
In fact, any vector in the space U
= span { e j : j ∈ S } is a feasible direction.
Specially, the causal parameter β ∈ U is a feasible direction which implies the optima should have zero directional derivative in this direction, i.e. (cid:104)∇ R ( α ) , β (cid:105) = 0 .
By Lemma 1, plugging α to β we get (cid:104)∇ R ( β ) , β (cid:105)
= 0 , yielding another objective
[Equation]
that β satisﬁes.
Similarly, any partition P of the set { 1 , 2 , · · · , p } gives a necessary condition that admits causal model as an extreme point
[Equation]
When the outcome model is Linear-Gaussian or LinearBernoulli, minimizing the invariant risk term in Eq.
(19) is equivalent to Eq.
Suppose the DGP and the predictor are linear as in Eq.
(9), and L 2 risk function R e , then
[Equation]
where the left side is the invariant risk term and the right side is objective in Eq.
Similarly, suppose the outcome is generated by y e ←
Bernoulli ( σ ( β (cid:62) x e )) , the predictor is ˆ y e
=
σ ( α (cid:62) x e ) where σ ( x ) = 1 / (1 + exp( − x )) is the sigmoid function, and the risk function is the cross entropy loss R e ( α ; y e , ˆ y e ) = − E
[ y e log(ˆ y e )
+ (1 − y e ) log(1 − ˆ y e )] , then[-(cid:13)(cid:13) ∇ w | w =1 . 
-]
[Equation]
y, [Equation] y ) , α (cid:105) ) 2 . [Equation]
The connections between Eqs.
(20), (22) and (23) explain the mechanism behind IRMv1 for linear Gaussian or Bernoulli models, as the causal coefﬁcient belongs to the optima set of the invariant risk term.
The connection also indicates the sub-optimality of the
IRMv1 objective.
The invariance risk term, rewritten as the inner product between the gradient and parameter vectors, only considers a single feasible direction for the constrained optimization problem Eq.
(4), among all feasible directions that form a ( p − | S | )
-dimensional linear space.
The spectrum between CoCo and invariant risk term, as shown in Eq.
(21), tells that the ﬁner the partition is, the smaller the optima set of Eq.
(21) becomes.
This means, among all conditions in the form of Eq.
(21), the one given by CoCo as Eq.
(8) is the strongest and the one given by
IRMv1 as Eq.
(20) is the weakest.
Since the ultimate goal is to identify the causal coefﬁcient, we prefer the strong condition that gives a small set of solutions in a single environment.
See Appendix F for a case study comparing
ERM, IRMv1 and CoCo analytically.
Because of an excessive number of solutions of the invariant risk term, IRMv1 puts high requirement on the number of environments and sufﬁciency of heterogeneity.
In practice, there can be multiple parameters that minimize the IRMv objective, including that of non-causal models.
By simulations in § 5, we will show that optimizing the IRMv1 objective can fail to produce robust predictions, especially when the outcome is generated neither from Linear-Gaussian nor
Linear-Bernoulli models.
Similar failure modes of IRM are studied in cases of a two-bit model [9] and a nonlinear classiﬁcation model [14].
Lastly, we notice that adding any general condition in
Eq.
(21) to the strong condition Eq.
(8) does not change the optima set while in practice it may improve the smoothness of the optimization landscape.
C. Proofs
In this section, we present proofs for the results in the main paper.
The following proof is for Lemma 1.
Proof.
Let the random vector x = ( x 1 , · · · , x p ) (cid:62) denote the covariates.
The expected mean square error is
[Equation]
Since supp ( α ) =
supp ( β ) , the ( α − β ) (cid:62) x is a linear combination of the true causes as (cid:80) j ∈ supp ( β ) ( α j − β j ) x j which is independent of (cid:15) by the SEM, thus E [( α − β ) (cid:62) x (cid:15) ]
= 0 .
Since E [ xx (cid:62) ] is assumed to be positive deﬁnite, the unique optima of the square error is α = β .
The following proof is for Proposition 1.
Proof.
By the construction of Λ , B ∗ =
A Λ is a matrix where the j -th column B ∗ j =
if j / ∈ S .
Similar to the proof of Lemma 1, we can compute the L 2 risk as
[Equation]
Due to the constraints, B j = B ∗ j = 0 , B x ⊥⊥ (cid:15) , B ∗ x ⊥⊥ (cid:15) , therefore the second term is zero.
Then the L 2 risk reaches its minimum as E [ (cid:15) 2 ] when B = B ∗ , γ = γ ∗ .
D. Identiﬁcation with Heterogeneous Environments
We establish causal identiﬁcation for CoCo.
Causal identiﬁcation involves writing the causal quantity of interest as a functional of the observed data distribution; this functional is also known as the causal identiﬁcation strategy.
In the context of CoCo, we consider the functional that maps the joint distributions p ( x e , y e ) over a set of environments to the risk function, then to the optima of CoCo objective
Eq.
Causal identiﬁcation for CoCo thus amounts to proving any optima of the CoCo objective must coincide with the causal coefﬁcient of interest.
Since all solutions of Eq.
(8) can be characterized analytically (see § 3), we are able to deﬁne what the general effective and ineffective interventions are.
To keep notation consistent, denote C as the set of known non-descendants of the outcome and S as the unknown set of direct causes.
For any set H with C ⊂ H ⊂ { 1 , 2 , · · · , p } , we ﬁt a regression model on X eH in each environment, and collect the regression coefﬁcients as { ˆ α eH } e ∈E .
We call the set H an invariant set , if the estimations
[Equation]
If H is an invariant set, we deﬁne a length p vector as an invariant vector by equating it to ˆ α H when restricting to the set H and padding it with zeros at other elements.
When there is more than one invariant vector, we call the interventions that construct the environments as ineffective interventions .
Back to the linear SEM and linear predictor.
Denote E as a set of environments, R e ( α ) =
E [(1 / 2)(ˆ y
e − y e )] 2 as the risk, and W e :=
E [ x e ( x e ) T ] ∈
R p × p as the Gram matrix which is assumed to be positive deﬁnite [15].
With all this in place, the causal relationship and causal effects can be identiﬁed by CoCo, as long as the interventions are valid and effective.
Theorem 1.
For the linear SEM in Eq.
(9) and predictor in
Lemma 1, assume W e (cid:31) 0 for all e ∈ E , and assume the following conditions hold:
(A1)
Validity: ∃ S ⊂ { 1 , 2 , · · · , p } , x eS =
Pa ( y e ) , and E [ y e | x eS =
c =
E [ y e (cid:48) | x e (cid:48) S = c ] for all c ∈ R | S | , e, e (cid:48) ∈ E .
(A2) Effectiveness: exploring all the sets H with C ⊂ H ⊂ { 1 , 2 , · · · , p } , there are no distinct invariant vectors (deﬁned in Eq.
Then the causal coefﬁcients β are identiﬁable, and are given by
[Equation]
which is the solution of Eq.
Proof.
Let s ej =
[Equation]
] = [Equation]
= ( s e 1 , · · · , s ep ) T .
By the data generating process, s ej =for j ∈ { 1 , · · · , K } .
Let
[Equation]
where f ( α ) is CoCo objective.
Direct computation shows
[Equation]
Notice f ( α ) ≥ 0 and by the structual equation model, due to independence of the exogenous noise (cid:15) and causes Pa ( Y ) , we have s e ◦ β = 0 .
Hence for α ∗ =
β , f ( α ∗ ) = 0 .
This guarantees the existence of a solution as causal coefﬁcient β .
To prove the identiﬁcation, it is sufﬁcient to prove that for all α (cid:54) =
α ∗ , f ( α ) > 0 .
We use proof by contradiction.
Let H =
supp
(˜ α ) and H c as its compoment set in { 1 , 2 , · · · , p } .
We assume f ( α ) =
and α (cid:54) =
β and deduce a contradiction.
Since f ( α ) = 0 , for all e , (cid:107) g e ( α ) (cid:107) = 0 .
Since g e ( α ) =
∇ R e
( α ) ◦
˜ α , it means ∇ R e
( α ) H = 0 , for all e .
However, by the characterization of the plausible set in Section 3, Assumption A2) implies that there does not exist α ∈ R p , α (cid:54)
=
β , such that ∇ R e ( α ) H = 0 , ∀ e ∈ E .
Otherwise, the set H is an invariant set and both α and β are invariant estimations, which violates Assumption A2).
Hence for α (cid:54) =
β , there exists an environment e (cid:48)
∈ E with ∇ R e (cid:48)
( α ) H (cid:54) = 0 .
This yields a contradiction.
E. Data generation and implementation details for §
Synthetic data.
we generate data for the case 1 according to Fig. 3.
The DGPs for case 1 is
[Equation]
Case 2-3 with confounder and mediator are generated by the following SEM with c 1 = 1 , c 2 = 0 for case 2, c 1 = [Equation] for
[Equation]
Case 4 with an unobserved direct cause x e 3 is generated with
[Equation]
Case 5 with a collider is generated with
[Equation]
To generate data from different environments, we set the parameter γ e in DGP (Eqs.
(28), (30) and (31)) by γ e ∈ { 0 . 5 , 2 . 0 } ; as required, this leaves the causal effect invariant.
For IRM, we minimize IRMv1 objective Eq.
and report the hyper-parameter λ ∈ { 2 , 20 , 200 } that gives the lowest
MAE.
The function mapping from the causes to the outcome is linear with additive noise.
We specify x 1 as a known preoutcome variable (for use of the method in Eq.
and run
CoCo , IRM, and ERM to estimate the causal coefﬁcients.
CMNIST data.
CMNIST is a semi-synthetic data set for binary classiﬁcation, ﬁrst introduced in Arjovsky et al.
Based on the MNIST data set, the image of hand-written digits 0-4 and 5-9 are labels as ˜ y
= 0 and ˜ y
= 1 respectively.
For each environment, the outcome y e is generated with 0.75 probability as ˜ y and with 0.25 probability as 1 − ˜ y .
We call ˜ y
the clean labels and y e the noised labels .
The digit is colored green with probability p e if y e = 1 and 1 − p e if y e = 0 , otherwise it is colored red.
The DGP across environments differs in p e .
Environments are constructed for training with p e ∈ { 0 . 1 , 0 . 2 } , for validation p e = 0 .and for testing p e = 0 .
The predictor is a fully connected neural network with two hidden layers.
For CoCo, we use objective Eq.
(14) to optimize the predictor.
The risk penalty weight λ r is chosen on the validation environment and is reduced by a factor of 10 when the parameters are sufﬁciently away from 0 .
For IRM, we use a learning rate as 10 − 4 to ensure stability over long iterations and use other hyper-parameters and annealing strategy provided by the author’s code.
Natural image data.
The predictor is a fully connected neural network with one hidden layer of size 10.
The inputs are 512-dimensional features extracted from ResNet18 [7], a pre-trained model on the ImageNet dataset [5].
In this example, we ﬁnd for CoCo objective (14), adding the weak penalty (20) with weight λ w improves convergence.
It is possibly due to the smoothed landscape as discussed in
Appendix B. We set the weight λ w = 10 4 .
For both CoCo and IRM, we reduce the weight of risk regularization by a factor of 10 5 after 100 epochs.
The parameters are selected on the validation environment.
Here we ﬁnd annealing the risk necessary for both methods otherwise minimizing the risk term often forces the predictor to use spurious associations after long iterations.
The need of risk term being small may be due to a limited number of training environments.
F. Illustrative Example
In this section, we study the connection and distinction between ERM, IRM and CoCo by studying a speciﬁc example, which is adapted from the “minimal coding implementation” in [1] Appendix Section D.
The DGP is:
[Equation] [Equation]
The predictive model is ˆ y e
= α (cid:62) x e , where input x e = ( x e 1 , x e 2 , z e 1 , z e 2 ) (cid:62) , parameter α ∈ R 4 .
The risk function for the environment e is the mean square error, i.e.
R e ( α ; y e , ˆ y e ) =
E [(1 / 2)(ˆ y
e − y e ) 2 ] .
Variable z
=
( z 1 , z 2 ) are associated with the outcome spuriously.
Assume the number of training environments is |E| =
K .
Direct computation gives
[Equation]
From Eq.
(33), the optimum for ERM in environment e is ˆ α ERM
[Equation]
[Equation]
and the optimum for the ERM over K environments is
[Equation]
[Figure]
[Equation] [Figure]
The outcome is generated from a Linear-Gaussian model, so the invariant risk term in IRMv1 objective Eq.
(19) equals to
Eq.
(20) as shown in Appendix B.
Therefore, we can solve ( (cid:104)∇ R e ( α ; y e , ˆ y e ) , α (cid:105) ) 2 = 0 and ||∇ R e ( α ) ◦ α || 2 =and get the optima set of the invariant risk term and CoCo respectively.
As shown in Fig.
4, the solutions for the invariant risk term in environment
e form an ellipse in the space of ( α 1 , α 3 ) and ( α 2 , α 4 ) , which are inﬁnite points.
The optimum of
CoCo is a strict subset of solutions of the invariant risk term, which has size of 2 4 .
For a single environment, CoCo solutions contain the ERM solution and the causal coefﬁcients.
With data from three environments, the causal coefﬁcient β = (1 , 1 , 0 , 0) (cid:62) belongs to the solutions of the aggregated invariant risk term and CoCo, but does not belong to the solution of the aggregated ERM.
G. Additional Simulation Study
This section contains experimental results in addition to the simulations in § 5 in the main paper.
G.1. Model misspeciﬁcation
We further study the implication of model misspeciﬁcation using data from § 5.1 case 5.
We compare two predictors, one is a linear model and the other is a nonlinear neural network.
The data is generated linearly, so the linear predictor correctly speciﬁes the model and the neural network misspeciﬁes it.
Both models are trained with ERM and CoCo.
In case 5, the variable x is the cause and z is a predictive but non-causal covariate.
We study how the two models, trained by the two methods, can generalize its predictive accuracy to new observation ( x, z ) respectively.
The results are shown in Appendix Fig.
When the model is correctly speciﬁed, ERM estimation cannot generalize to new values of z , while CoCo estimation can predict on any ( x, z ) accurately.
When the model is misspeciﬁed, the model trained by ERM can only interpolate between the training points, while the model trained by CoCo can generalize to new z though not to new x .
This study demonstrates that if the model is not misspeciﬁed,
CoCo can learn the causal model; if model is misspeciﬁed,
CoCo learns a model that can generalize to new environments where the spurious association differs from that in the observed environments.
In both scenarios, model learned by CoCo has better generalization performance than that learned by ERM.
4 X
Z
ERM, linear            4 X           5 CoCo, linear      4 X
Z
ERM, nonlinear            4 X           5 CoCo, nonlinear
G.2. Gaussian mixture example
In this section, we study a multi-class classiﬁcation problem when the inputs contain non-causal covariates.
We modify
GMM to simulate the data set.
The observed covariates are ( x e , z e ) and the outcome is y e , where e is the environment index.
For each environment e , the data are generated with
SEM
[Equation]
where p k = N ( x e ; µ k , I ) / (cid:80) Kk
[Equation] =1 N ( x e ; µ k [Equation] , I ) , k 1 ∼
Multinomial (1 /K, · · · , 1 /K ) .
Among the covariates, the mapping from x e to the label y e is invariant across all e while z e is predictive to y e due to spurious associations.
We aim to learn a model that makes predictions based on the causal covariates x e .
In Eq.
(38), x e are generated from GMM with the component centers µ k = √ 1 .
5 K e k ∈ R K .
To generate the non-causal covariates z e , we ﬁrst generate K random vectors { u ek } Kk =1 with u ek ∼ (cid:81)
[ k/ 2] i =1 U (0 , 1) for environment e .
Then for a data point in the component y e , z e equals u ey e with probability 1 − p e and equals a random vector from { u ek } Kk =1 otherwise.
By doing so, z e is associated with y e but the association varies across environments when u e 1: K change with e .
The DGPs that generate the environments is characterized by the values of u e 1: K and p e .
We set the training environments with K = 5 and p e ∈ { 0 . 01 , 0 .
by Eq.
For a validation/test environment f we generate a new set of { u fk } Kk =1 and set p f = 0 .
We evaluate the test performance by averaging the accuracy over 10 testing environments.
If the predictor learns to predict based on the causes x e instead of z e , it can accurately predict y e in both training and testing environments.
The predictor is a fully connected neural network with two hidden layers.
For CoCo, we use objective Eq.
(14) to optimize the predictor.
The penalty weight λ r is chosen on the validation environment and is reduced to 0 when the parameters are sufﬁciently away from 0 .
For IRM, we choose λ in Eq.
(19) and step size on the validation environment.
The results are shown in Figs.
6 to 8, and Table 2. Fig.(a) is the trace plot for the predictive accuracy in the testing environments.
The testing accuracy increases for all methods in the early stage of training but drops in the later stage for ERM and IRM.
We hypothesize that
ERM and IRM at ﬁrst improves the prediction by utilizing all covariates including the causal ones.
But in the later stage of training, it relies more heavily on the spurious associations to boost [Table] the predictive accuracy, which reduces the accuracy at the test time.
We provide evidence to this hypothesis in Fig.
6, by plotting the weight matrix that connects the input and the ﬁrst hidden layer.
The model trained by CoCo manages to set the weights associated with the non-causes z (the right block) close to zero, aligned with the analysis in Proposition 1.
In comparison, these weights obtained by IRM and ERM are mostly non-zero, passing information from non-causal z to the subsequent hidden layers and outputs.
In Fig.
7, we study how CoCo performs if the invariance is violated and how sensitive it is with different hyperparameters.
In panel (a), we construct the training environments by changing the cluster centers { µ k } Kk =1 in Eq.
(38) to { µ k + (cid:15) ek } Kk =1 , (cid:15) ek ∼ N ( 0 , σ 2 I K ) .
The noise (cid:15) ek changes the mapping from the covariates x e to the label y e across the environments, and the noise scale σ 2 reﬂects the magnitude of change.
The panel (a) shows the testing predictive accuracy increases as the invariance tends to hold.
In panel (b), we compute the test accuracy when using different number of environments M in training.
we construct training environment e by Eq.
(38) with vectors u ek ∼ (cid:81) [ k/ 2] i =1 U (0 , 1) for all k and p e ∈ { 0 . 01 , 0 .
01 M } .
We ﬁnd a growing number of environments reduces the testing error monotonically which might due to the increased heterogeneity in data.
In panel (c), we study how the testing error changes with the penalty weights λ r in CoCo objective Eq.
When λ r is large, the objective is close to the empirical risk and the test error is high; when λ r is small, the parameters often collapse to 0 .
Between the two extremes, CoCo with a wide range of λ r can learn a model that makes robust prediction in new environments.
[Figure]
F i r s t
h i dd e n l a y e r
[Figure] [Figure] i [Figure]
s t
[Figure]
i [Figure] [Figure] [Figure]
c
[Figure]