Doubly robust identification of treatment effects from multiple environments
Authors: Piersilvio De Bartolomeis
Julia Kostin
Javier Abad
Yixin Wang
Fanny Yang
Venue: ICLR 2025 Poster (2025)
Keywords: treatment effect
confounding
heterogenous data
causality
causal inference
unobserved variables
post-treatment variables
collider bias
Primary area: other topics in machine learning (i.e., none of the above)

Abstract: Practical and ethical constraints often require the use of observational data for causal inference, particularly in medicine and social sciences. Yet, observational datasets are prone to confounding, potentially compromising the validity of causal conclusions. 
 While it is possible to correct for biases if the underlying causal graph is known, this is rarely a feasible ask in practical scenarios. A common strategy is to adjust for all available covariates, yet this approach can yield biased treatment effect estimates, especially when post-treatment or unobserved variables are present.
We propose RAMEN, an algorithm that produces unbiased treatment effect estimates
by leveraging the heterogeneity of multiple data sources without the need to know or learn the underlying causal graph. Notably, RAMEN achieves *doubly robust identification*: it can identify the treatment effect whenever 
the causal parents of the treatment or those of the outcome are observed, and the node whose parents are observed satisfies an invariance assumption. Empirical evaluations across synthetic, semi-synthetic, and real-world datasets show that our approach significantly outperforms existing methods.

Published as a conference paper at ICLR 2025
DOUBLY ROBUST IDENTIFICATION OF TREATMENT EFFECTS FROM MULTIPLE ENVIRONMENTS
Piersilvio De Bartolomeis', Julia Kostin', Javier Abad', Yixin Wang", Fanny Yang* ETH Zurich UUniversity of Michigan
ABSTRACT
Practical and ethical constraints often require the use of observational data for causal inference, particularly in medicine and social sciences. Yet, observational datasets are prone to confounding, potentially compromising the validity of causal conclusions. While it is possible to correct for biases if the underlying causal graph known, this is rarely a feasible ask in practical scenarios. A common strategy is to adjust for all available covariates, yet this approach can yield bi- ased treatment effect estimates, especially when post-treatment or unobserved variables are present. We propose RAMEN, an algorithm that produces unbi- ased treatment effect estimates by leveraging the heterogeneity of multiple data sources without the need know or learn the underlying causal graph. Notably, RAMEN achieves doubly robust identification: it can identify the treatment effect whenever the causal parents of the treatment or those the outcome are observed, and the node whose parents are observed satisfies an invariance assumption. Em- pirical evaluations across synthetic, semi-synthetic, and real-world datasets show that our approach significantly outperforms existing methods'.
INTRODUCTION
Treatment effects are key quantities of interest applied domains such as medicine and social sci- ences, they determine the impact of interventions like novel treatments policies outcomes of interest. To achieve this goal, researchers often rely on randomized trials since randomizing the treatment assignment guarantees unbiased treatment effect estimates under mild assumptions. How- ever, methods relying on randomized data face several issues, such as small sample sizes, sample populations that do not reflect those seen in the real world, and ethical or financial constraints. As result, there is growing interest in using observational data to estimate treatment effects.
fundamental challenge in using observational data is the selection of a valid adjustment set, i.e. set of covariates that can be used to identify and estimate the treatment effect. Although cri- teria for identifying valid adjustment sets are well-established, they rely on the knowledge of the underlying causal graph. When the graph is not known, practitioners often adjust for all available covariates (Austin, 2011). 1). Yet, this approach runs the risk of including bad control--covariant open backdoor paths between the treatment (7) and the outcome (Y), thereby introducing bias into the treatment effect estimate. For instance, consider the causal graphs illustrated in Figure 1, where {X1, X2) are the observed covariates. In Figure la, both X, and X2 are parents o' and adjusting for {X1,, X2) blocks all backdoor paths between T and Y, making it valid adjustment set. In contrast, Figure 1b, X1 is child of T, and conditioning on it opens backdoor path between T and Y, introducing bias in the effect estimate. In the latter case, X1 is referred to as bad control.
that
Bad controls pose a significant challenge when the causal ordering of the observed covariates is not clear (King, 2010; Montgomery et al., 2018). A prominent example for bad control is the birth-weight paradox (Wilcox, 2001) from which it was concluded that when estimating the effect of maternal smoking (I) on infant mortality (Y)), the birth weight would be bad control like X, in Figure 1b, as it is child node of T and likely leads to collider bias. Further, Acharya et al. (2016)
See GitHub repository: https: htttp::/github... com/ jaabmar wawaamarr/AAM///// /RAMNN/

---

Published as a conference paper at ICLR 2025
X1
X2
X1
X2
T
(a) (X1, X2) is a valid adjustment set
(b) {X1, X2) } is not a valid adjustment set
Figure 1: Two causal graphs illustrating when the set of all covariates is or is not a valid adjustment set: (a) (X1, X2) blocks all backdoor paths between T and Y, making it valid adjustment set; (b) x, opens a backdoor path between and Y, introducing bias in the treatment effect estimate adjusted for. Unobserved variables are dashed and colored in white, bad controls are colored in red, and good
that can be included in the adjustment set-are colored in green.
found that up two-thirds of empirical studies in political science that make causal claims inad- vertently include bad controls in their analysis, leading to biased treatment effect estimates. Several works have tried tackle the problem of bad controls using expert-driven structural knowledge, e.g. by leveraging anchor variables (Cheng et al., 2022a; Shah et al., 2022), but such domain expertise often unavailable in practice. Shi al. (2021) propose an alternative approach that leverages ac- cess to multiple heterogencous data sources-c.. observational studies from different countries to identify and estimate the treatment effect in the presence of bad controls. However, their approach can fail to identify the treatment effect when some variables in the causal graph are unobserved and their distribution shifts across environments, e.g. for the causal graph illustrated in Figure 1b.
In this work, we propose Robust ATE identification from Multiple ENvironments (RAMEN e), an algorithm that identifies and estimates the average treatment effect in the presence of bad controls without the need know learn the complete causal graph. Notably, RAMEN leverages the hetero- geneity of multiple data sources to achieve doubly robust identification: it can identify the average treatment effect whenever the causal parents of the treatment or those of the outcome are fully ob- served, and the node whose parents are observed satisfies an invariance assumption. In particular, our methodology relaxes the full observability requirements of Shi et al. (2021), requiring only partial observability of the causal graph. Our key contributions are outlined below.
We propose the first algorithm, data sources to identify and estimate the average treatment effect the presence of both bad controls and unobserved variables. Our algorithm is based novel double robustness property that offers two strategies to identify the treatment effect: either observe the causal parents of the treatment those of the outcome.
knowledge, that leverages multiple heterogeneous
We demonstrate that our algorithms significantly outperform existing approaches for treat- ment effect estimation in the presence bad controls on synthetic, semi-synthetic and real- wolld datasets. We further evaluate method on real-world example, showing that our results align with established epidemiological knowledge.
RELATED WORK
Various criteria and methods have been proposed for covariate selection, often in the form of neces- sary and sufficient conditions for a given causal graph, such as the backdoor criterion and its vari- ations (Pearl, 1995; Shpitser et al., 2010; Vander Weele & Shpitser, 201 Maathuis & Colombo, 2015; Perković et al., 2018). However, since the causal graph is rarely known in real-world appli- cations, the most common heuristic approach assumes that all observed covariates are mretrreament and includes all of them (Austin, 2011, p. 414). Yet, including all covariates has several drawbacks: certain pre-treatme treatment covariates can introduce M-bias (Entner et al., 2013; Gultchin et al., 2020; Cheng et al., 2022b; Shah et al., 2022), and even when bias is not issue, selecting smaller subset
2

---

Published as a conference paper at ICLR 2025
of covariates leads to more efficient estimates (Hahn, 2004; White & Lu, 2011; De Luna al., 2011; Rotnitzky& & Smucler, 2020; Witte et al., 2020; Henckel et al., 2022; Guo et al., 2023).
The problems described above are orthogonal to our focus in this paper, which on scenarios where bad controls are present and the underlying causal graph is unknown (see Appendix B for a complete literature review). In this context, previous works have achieved partial identification, albeit with significant computational costs (Hyttinen et al., 2015; Malinsky & Spirtes, 2017). More recently, several works have proposed methods for point identificatios using expert-driven structural knowledge. Cheng et al. (2022a) rely on known anchor variable, Shah et al. (2022) assume that a direct parent of the treatment variable is known, and Shah et al. (2024) assume all children of the treatment variable are observed and known. However, significant limitation of these approaches is their dependence on structural knowledge of the causal graph, which is often unavailable in practice. In contrast, our methodology achieves point identification by leveraging multiple heterogeneous data sources, effectively circumventing the need for partial knowledge of the underlying causal graph.
Finally, our notion of double robustness differs significantly from classical results in estimation (Robins et al., 1994; Vansteelandt et al., 2008; Chernozhukov et al., 2018) and identification (Arkhangelsky & Imbens, 2022), which focus on robustness to model misspecification. A more similar concept is robust identification in instrumental variable settings (Kang et al., 2016; Hartwig et al., 2017; Guo et al., 2018; Kuang et al., 2020; Hartford et al., 2021). which allows for a fraction of instruments be invalid. In contrast, our method guarantees identification that is robust to un- observed variables when either (i) the observed covariates include all parents of T and T satisfies invariance assumptions, or (ii) the same holds for Y hence yielding double robustness.
PROBLEM SETTING
For fixed directed acyclic graph (DAG) G. we denote the complete set of its nodes by Ž and the observed nodes by Z. We denote the index set of parents, ancestors, and descendants for any node Z; by Pa(22 Z,), An(Z,) and De(Z,), respectively. Additionally, for an index set S, Zg denotes the subvector of 2 corresponding to the indices in S. We assume the data is collected under different conditions, represented by environments e E E, with E| = Re.- For each environment have access to a dataset D {(X,, Ti, Y)) J== which contains n i.idd tuples sampled from the marginal induced by the joint distribution Pe' over (X, U, Y). Here, X E Idd are the observed covariates, U R are the unobserved covariates, 7 {0, }} is a treatment assignment variable and Y E R is the outcome. We denote by P |8 Eesg Pe the distribution of the pooled environments.
&, we
For each environment e E E, the distribution P is induced by structural causal model (SCM), defined as tuple M = (G. (G.ffi P2) on p d d+ k + 2 variables (Z,,..., .z)). where the observed covariates are X = Zidj. the unobserved covariates are U 2d+1A)- the treatment variable T = 7p-1. and the outcome variable is Y Zpp with p & An(T). The SCM defines the probability distribution 5et by setting for each E [p]
2, + f)(2pa(2)) 5),,
j= 1= 1=,,..,,7
(1)
where f P + R is measurable function and e E IR is an exogenous noise vector following the joint distribution g_ over p independent variables.
Further, along the lines of the existing methods in the literature (Shi et al., 2021; Wang al., 2023), we require the absence of observed mediators between T and Y in the structural causal model.
Assumption 3.1 (Absence of Mediators). We assume that no observed mediators exist between T and Y, i.e. it holds that
De(T) Dn(()
0
We remark that Assumption 3.1 is falsifiable using statistical tests to determine whether covari- ate is a mediator between T and Y: see, e.g. Baron & Kenny (1986); Preacher & Hayes (2004). Furthermore, even when this assumption is violated, the causal quantity identified by method corresponds to the natural direct effect (Pearl, 2022), which remains a quantity of interest in fields such epidemiology (Tchetgen & VanderWeele, 2014) and the social sciences (Imai al., 2011).
3

---

Published as a conference paper at ICLR 2025
3.1 TREATMENT EFFECT IDENTIFICATION
Our goal is to identify the treatment effects for different environments in the presence of unob- served and post-treatment variables. More specifically, we are interested in the average treatment effects (ATEs) for all environments E E, defined as
Ege ydo(f=1) - ydo(T=0)
A common approach for identifying the ATE is to find valid adjustment set (Shpitser al., 2010),
that is, subset 5 S [d] of the observed covariates that satisfies both the classic outcome and treatment identification formulae, i.e. for all environments e and E {0, 1] holds that
[ddo(T==)]
YNT== Ep= Xx. = Ep PT = xx))
Epc
Eps
(2)
Several criteria have been proposed in the literature to find valid adjustment sets, with the backdoor criterion being the most prominent-see Peters et al. (2017, Sec. 6.6) for a detailed discussion. However, these criteria crucially rely on knowledge of the underlying causal graph. Therefore, it is commonly assumed among practitioners that the set [d] of all observed covariates is valid adjustment set. This is a reasonable assumption only in settings where all the observed covariates pre-treatment and there are unobserved covariates.
In contrast, our work focuses on settings where both post-treatment and unobserved covariates are present. To identify the ATE such settings, we introduce key assumption. While each environ- ment may have different joint distribution  e over (X, U, Y, I'), we assume that at least one of T or Y is an invariant node with its parents fully observed and the conditional mean the node given its parents is invariant across environments.
Assumption 3.2 (Invariant node). We assume that one of the following holds for all E E: (a) All parents Pa(T) of T are observed and Ep [T] ZPa(T:] Ep [ ZPa(T)), I - a.s. (b) All purenty Pa(Y) Y are observed and Ey- [PZ)) = E5 [Y Zpacr)] pp a.s.
We denote the node for which the above holds as the invariant node Viav E (T.Y).
worth emphasizing that each of the above assumptions can provide identification of the ATE on its own. Here, we combine these two identification assumptions to obtain doubly robust iden- tification: we only require that either (a) or (b) in Assumption 3.2 holds. This is similar in spirit to the double machine learning literature (Robins & Rotnitzky, 1995; Chernozhukov et al., 2018), where only one of two assumptions about model specification needs to hold to obtain consistent ATE estimates. However, the key difference is that our assumption offers robustness against poten- tial unobserved variables in the underlying causal graph, whereas classic double robustness offers robustness against misspecification of the outcome and treatment functions.
Further, the invariance assumptions (a) and (b) are closely related to the conditions in the invariance- based domain generalization literature, such as Peters et (2016); Rojas-Carulla et (2018); Gu et al. (2024). While these settings are included in Assumption 3.2 (as we discuss in Appendix A.1), our setting does not require full independence the noise variable". unlike Peters et al. (2016), nor is it limited to the additive noise case, as in Gu et al. (2024), which does not hold in the case of binary treatment variables. Finally, we comment on the observability part of Assumption 3.2: assuming Pa(Vin)) are observed is strietly weaker than causal sufficiency, which would require the full causal graph to be observed. Notably, our framework allows for scenarios where Assumption 3.2 (a) is satisfied with T as the invariant node, whereas the setting proposed in Shi et al. (2021) is more restrictive, since it requires Y to always be the invariant node.
4 METHODOLOGY
In this section, we introduce RAMEN, our method to identify the ATE by leveraging the heterogene- ity in the observed data. First, we present doubly robust population-level estimator and discuss
Athhough we require independence of exogenous noise variables for the full graph, here, we refer to the graph limited to the observed nodes, where the noise variables ey can he dependent on Zpa(Z,)-
4

---

Published as a conference paper at ICLR 2025
under which conditions it equals to the ATE. Then, we show how compute this estimator tractably by minimizing a novel invariance loss and propose two algorithms to do so: a combinatorial search over subsets and more scalable differentiable approach for high-dimensional covariate settings.
4.1 POPULATION-LEVEL ESTIMATOR
In what follows, we denote by Zr = [d] and Iy -d + 1] the index sets corresponding to the observed variables Z (X, T). For any node V E (T, Y) and any observed subset S Iv, define the conditional means over the pooled and individual environments
ms( = Epp IV Zs] and mz(Z:V) = Ex- [V Es].
1. Identify an invariant set We begin by observing that, by Assumption 3.2, there exists an invariant node Vinv (T, Y and a subset of covariates (given by, egg, Pa(Vin))) for which the following conditional moment constraint holds for all environments E E:
SS C IVi.. m2(Z; Vinv ms Z Vinw,, IPe - a.s.
The set S is not necessarily unique: besides the (observed) parents of Vinv for instance, the invari- ance could also hold for certain supersets of Pa(Vie). By observing that the conditional moment constraint above is equivalent the following infinite set of unconditional moment constraints
Ep [(Vinv mi (Z; Vinv [Z)] for all measurable h,
any set S that satisfies the invariance constraint Equation (3) is also contained in
argmin max SCIVin 3EE
Exe [(Viv ms Vinn))(((ss) (RIS)))
sup
= argmin Js(Z; Vine,, S SCIv
(4)
where (IRI)) denotes the space of measurable functions RR However, since the invariant node Vinv is not known beforehand, we search for a set of observed nodes that satisfy the invariance with respect to either or Y, that is, we want find
Se min argmin Js(Z:V). VE(TY) SCIv
(S)
Further, we can leverage the structural knowledge that always an ancestor of to simplify the optimization problem. Specifically, we know that T must be part of the invariant when V Y. Therefore, we can condition the expectation in Equation (4) on T = 0 and T'' 1 separately and take the maximum the resulting losses. This gives us slightly different loss function for Y:
T= (()))((()))) T +) ) 2
Js max (E(0.1) cEE
Eze - EplY Zs, = (RIS)
sup
(6)
2. Estimate the ATE For minimizer S,, we then define the corresponding population-level RAMEN, estimator for all environments e E as
(Y-Pi(Xss)) (Y- (1-
(Sx) = Egg (i((Xs))-nolf
a(xs²)
1 1 (xx
where we define the pooled conditional outcome and treatment functions as
(Xss) = E T T' and ((ss))=E<(TT]) Xs.l
In what follows, we show that under a condition on data heterogeneity. detailed in Assumption 4.1, population-level estimator 8 (Sa) is equivalent for all S& that satisfy Equation (5), and it is equal to the true treatment effect. In Section 4.3, we then discuss how we can construct a good finite-sample ATE estimate in computationally efficient way.
5

---

Published as a conference paper at ICLR 2025
4.2 DOUBLY ROBUST IDENTIFICATION GUARANTEES
Without further assumptions, finding the minimizer of Equation (5) is not sufficient for identifying the ATE: for instance, if there is no variability between distributions pe, our objective could be trivially minimized by any observed subset S. Only when there is sufficient heterogeneity in the observed environments will (S+) be equivalent to the ATE. We formalize this condition below.
Assumption 4.1 (Identification condition). For all V } and S C Iv, it holds that:
Ve E & VZ ms(Z:V), p
P (mms(Z:V) MP.(V) (P:(v)(ZVV)) 0.
Assumption 4.1 can be understood ensuring that the environments present sufficient heterogene- ity. This heterogeneity is crucial because it guarantees that conditioning on any set S with invariant outcome treatment functions across environments is equivalent conditioning the invariant node. Although our environment heterogeneity assumption is relatively strict, it is common requirement in the invariance literature (cf. Peters al. (2016); Arjovsky al. (2019)). For example, in the simultaneous noise intervention setting described in Peters et al. (2016, Section 4.2.3), Assumption 4.1 can be satisfied with few two environments, while in the case of single- node interventions it requires (p) environments, where p is the number of observed variables.
the parents of
We now present our formal identification result for the ATE
Theorem (Doubly robust identification). Let Se be any minimizer of the invariance loss in Equa- tion (5). Then, under Assumptions 3.1,3.2, 4.1, if positivity holds, that is
Ve E PT - ( Xs x) 0, VI {0, 13 and Vo E supp (PPX), we can identify the averuge trealment effect 0r 0 (Se), all environments e E E.
We remark that the positivity assumption is standard and widely used for identifying treatment effects in observational studies (Hernán & Robins, 2010, Sec. 3.2). Theorem states that any solution to our invariance loss valid adjustment in the sense that is sufficient identify the average treatment effect in all the environments.
4.3 AN EFFICIENT FINITE-SAMPLA ESTIMATOR
The population-Ievel estimator presented above involves two significant computational challenges. First, the invariance loss involves a supremum over an infinte-dimmnssioaal space of measurable functions, making it intractable to compute directly. Second, it requires searching over all possible subsets of covariates, which is computationally infeasible for high-dimensional settings. To address these issues, we introduce practical estimator based on a kernelized invariance loss and differen- tiable relaxation the subset selection problem.
Kernelized invariance loss A major problem of the loss function in Equation (4) is that it is computationally infeasible to search over the entire space of measurable functions. However, we can simplify the problem by restricting long as the reproducing kernel of the RKHS is universal (e.g. Gaussian kernel), the two formulations are equivalent (Gretton et al., 2012). More formally, for any subset S C Iv and environment e E:
be in a reproducing kernel Hilbert space (RKHS). As
1)
Ep- (R(S)))
sup
(V
ms (z(Z
sup Exe [ōs(Z,V)M(Zs)) |M H11
(2(V
Ep- [8 (8, V).(: ())(;
Ep- [8g (Z, V)k (Zs, (Zs,Z') óg (Z', VV)],
where k is a uniformly bounded reproducing kernel corresponding to a universal RKHS H (Stein- wart, 2001, Definition 4), and (V', Zz' an independent copy of (V, Z) following the same distri- bution. Hence, we can rewrite our invariance loss in closed form:
VZ = max Ep [6,(Z,V)) (Zs, Z'2) 65 V')].

---

Published as a conference paper at ICLR 2025
A long line of work has proposed methods estimate invariant predictors, especially when the optimal predictor is linear. These methods broadly fall into two categories: hypothesis test-based methods (Peters et al., 2016; Heinze-Deml et al., 2018; Pfister et al., 2019) and optimization-base methods (Arjovsky et al., 2019; Ghassami et al., 2017: Rothenhäusler et al., 2019; 2021; Pfister et al., 2021; Yin et al., 2024; Shen et al., 2023; Gu al., 2024; Wang et al., 2024). Our approach falls in the latter category, with a fundamental distinction. While all these works utilize the invariance principle to improve prediction in unseen environments and generalize to new settings, we aim to identify a treatment effect within the observed environments. This is reflected in our loss function, as it does not measure the quality of the predictor (e.g. using a least squares loss). Nonetheless, our invariance loss could also be of interest in the domain generalization literature as it retains the benefits of the invariance loss in Gu et al. (2024) while significantly simplifying their optimization procedure.
A fully differentiable loss When searching over all possible subsets of covariates is compu- tationally infeasible, we propose a continuous relaxation of the optimization problem in Equa- tion (5) that can be efficiently solved using gradient descent. Specifically, we select the nodes as Zu = B(w) 3 Z, where Z (X, and the j-th component of B(w) {0, 10]1+ sampled independently from a Bernoulli distribution with probability sigmoid (wz)). We parametrize the con- ditional mean using neural network fo and we aim to solve the following optimization problem:
_h E argmin wERd+1 BER (.VEYTY
min
max cEE
Epp, I(V (o (Z Zw) (V' fo((z)))
Since the weights are discrete, direct differentiation is not possible. To overcome this, we use a Gumbel approximation (Jang et al., 2017; Maddison et al., 2017: Gu et al., 2024), where the j-th component of B(w) approximated as:
(yw sigmoid W + G1. G2.)
T 0+,
T
with G1, and G2.5 being Gumbel(O, 1) random variables. This approximation makes B(w) dif- ferentiable (where it was previously discontinuous in wi). allowing us to optimize using gradient descent while gradually annealing the hyperparameter т. Finally, we construct the subset of covari- ates Sinsta-e by including Z, only if the weights are positive, that Sistt- = (i:w²² > 0}. We refer the reader to Appendix A.3 for the complete implementation details of our algorithms.
EXPERIMENTS
In this section, we evaluate our method through experiments on synthetic, semi-synthetic, and real- world datasets. We first present experiments on several known DAGs, where the invariances are known and satisfy our assumptions. In line with our theory, RAMEN correctly identifies the ATE, resulting in a low estimation error, whereas other methods tend to fail. We also test RAMEN on a more challenging benchmark by uniformly sampling DAGs using the Erdös- dard approach for testing causal methods across a wide variety of graph topologies (Huang et al., 2020). Finally, we validate our estimator beyond purely synthetic data: first in semi-synthetic set- ting with real-world covariates and then in a roal-world setting where we compare the conclusions from RAMEN with established epidemiological findings.
model-a a stan-
In our experiments, we focus on the statistical task of estimating the average treatment effect (ATE) 8 for each environment E E. To evaluate the performance of an estimator Ør, compute the mean absolute error (MAE) averaged across environments; 1 Eeeg | 0... We evaluate two (gl implementations of RAMEN: (i) òe- ((S)), based on combinatorial subset search (Section 4.3), and (ii) 0 = 0( ((Siusta-)). based on the Gumbel trick (Section 4.3); see Appendix A.3 for the complete implementation details of our algorithms. We compare RAMEN against three baselines: 8rm. the IRM approach for treatment effect estimation proposed by Shi et al. (2021); dau. which adjusts for all available covariates; and duil, which does not adjust for any covariates.

---

(Top row) Graphical models illustrating three scenarios in which the unobserved con- founder may break different invariances between T, Y, and the observed covariates X. In each graph, dashed arrows denote optional edges, dashed nodes indicate unobserved variables, green and red nodes indicate good and bad controls, respectively. Panel (a) shows the case where U does not break any invariance. Panel (b) shows the case where U breaks the invariance between X, and T. Panel (c) shows the case where U breaks the invariance between Xp and Y. (Bottom row) Mean 2500, d ob- Published as a conference paper at ICLR 2025
X
X
X
T
Y
T
T
Y
(a) YyY-invarie
(b) Y-invariance
(c) Tiivaariacce
1.2
1.2
1.2
0.8
0.8
MAE
MAE
MAE
0.6
0.6
0.2
A
collider descendant
noise
collider
descandant
noise
collier descendant
noise
0rrm
aaull
iista-
P
Figure 2:
absolute errors (MAE) of four estimators are shown across 5 environments served covariates), with error bars representing standard errors over 20 runs.
5.1 SYNTHETIC EXPERIMENTS WITH KNOWN DAGS
We start with data generated from distributions with simple underlying DAGs that satisfy our invari- ance assumptions, as illustrated Figure 2 (Row 2). Most importantly, we consider three distinct scenarios": (a) Y and T-invariances, i.e. both (a) and (b) in Assumption 3.2 hold; (b) Y-invariance, ic. only Assumption 3.2 (b) holds; (c) T-invariance, i.c. only Assumption 3.2 (a) holds. For each of the three different invariance scenarios, we further consider three variants: where X, is either descendant of Y, a collider between and Y, or independent noise. We describe the complete data- generating process in Appendix D.I. In the infinite sample limit, 0man should generally be biased since there is a confounder between 7' and Y; 6an should be biased only when X, is a collider or descendant; 0rm should be biased in the T-invariance case; Og and O should never be biased.
In Figure 2 (Row 1), we present the empirical MAE for all methods on finite-sample experiments that confirm the predictions from theory. First of all, both of our methods, Og and @ con- sistently achieve lower MAE compared to the baselines in all scenarios. In particular, we observe that the differentiable relaxation of our method does not significantly compromise statistical perfor- mance. Further, for T-invariance, the performance of 0.m deteriorates markedly as expected -e.g- in scenarios where the cost-treatment variable is a descendant of Y, it performs worse than simply approach remains robust even when of
adjusting for all available covariates. In contrast, the invariances compromised. Finally, we observe that relying on T-invariance increases the error across methods, possibly because the adjustment to statistically less efficient estimator, see e.g. Henckel et al. (2022, Corollary 3.4).
recover, the parents of the treatment, leads
In Appendix C.2, also present experiments for cases when of the invariances hold.
8

---

Published as a conference paper at ICLR 2025
5.2 SYNTHETIC EXPERIMENT WITH RANDOM HIGH DIMENSIONAL DAGS
We randomly draw a graph from the Erdös-Rényi random graph model with total number of nodes p 20. We do rejection sampling exclude graphs that either contain mediators- -s they violate Assumption 3.1 do not contain at least a confounder. We then assign Y and to nodes such that the invariance assumption satisfied for at least one of them', and sam- ple all variables from the resulting DAG via linear structural causal model except for the treat- ment variable T. We sample from a Bernoulli distribution with parameter equal to the sig- moid function applied to the function f in the structural equation of ' as in Equation (I). We further post-process the graph by adding a node X. = Y + T to make sure that there at least one post-treatment covariate. We assign the parents of T or Y (except common parents) to be the set of unobserved covariates U, depending on the invariance we want to preserve. In each environment, we apply a random uniform mean and variance shift to all the nodes in the graph except for T' and Y while preserving Assumption 4.1; see Ap- pendix D. for further details on the data generation. This process leads to distributions that are guaranteed to satisfy all our assumptions in Theorem 1.
Oumm
San
metaa
0.8
MAE
0.6
In our experiments we sample 100 DAGs and for each DAG vary the number of environments while keeping the sample size fixed. In Figure 3, we plot the empiri- cal MAE of fissta- (fg is computationally infeasible) and the baseline estimators averaged across the DAGs as a function of the number of environments. Notably, we observe that across all settings and numbers of available environments, @issta- significantly outperforms all the other baselines. Expectedly, Óirm fails to surpass all triv- ial baselines, even with many environments, lacks the robustness to unobserved parents of Y.
0.4
0.2
5
10
20
30
Number of environments
Figure 3: We plot the mean absolute er- ror averaged across environments when the T T-invariance is preserved. We sample n 2000 points for each environment; we report mean and standard crror over 100 runs.
5.3 SEMI-SYNTHETIC EXPERIMENTS: IHDP
The IHDP dataset contains covariates from 748 low- birth-weight, premature infants enrolled in home visita- tion program designed to improve their cognitive scores (Hill, 2011). Instead of using the commonly adopted syn- thetic functions from Dorie (2016), we simulate more challenging non-linear version of the outcome and treat- ment mechanisms inspired by Kang & Schafer (2007). Specifically, we retain the 6 continuous covariates (out of
0
Amni
2.5
MAE
1.5
1-
25 total covariates) from the original dataset and simu- late the outcome Y and treatment assignment T by ran- domly sampling complex functional forms, such as ex- ponentials and polynomials. In addition, we introduce 2-dimensional synthetic collider, X,, linear function of T and Y. We generate environments using Gaussian mean shifts in both pre-and post-treatment features, as well in either Y or T neither of them, and set the number environments to [E] = 5. Finally, at inference of the parents of the node
0.5
T,Y
Figure 4: Mean absolute emror averaged across environments the IHDP dataset when different invariances are preserved (T, Y. or both). We consider five environments with 748 points each; mean and stan- dard error reported over 20 runs.
time we do not observe among Y, T that is not invariant. See Appendix D.2 for the complete experimental details.
In the main text present the results for the settings where the parents Y, excluding the parents of T, are unobserved. Please refer Appendix C.3 for additional experiments where other invariances hold

---

Published as a conference paper at ICLR 2025
Figure depicts the MAE of all the methods when different invariance assumptions hold. Similar to the synthetic experiments in Sections 5.1 and 5.2, 6irm exhibits higher MAE when Y is not invariant across environments, and Öan, adjusting for all features, generally results in poor performance. Inter- estingly. mull performs competitively since the confounders have a limited impact on the outcome and treatment assignment in this dataset. Additional experiments where the post-treatment feature is either descendant of the outcome, independent noise, where neither nor Y remains invariant provided in Appendix C.2. Moreover, we present experiments including mediators between the treatment and the outcome in Appendix C.I.
5.4 REAL-WORLD EXPERIMENT: EFFECT MATERNAL SMOKING ON BIRTH WEIGHT
In our real-world experiment, we evaluate our method on the observational dataset from Cattaneo (2010) that studies the effect of maternal smoking (treatment 7') during pregnancy on birth weight (out- come Y) using the data from n 4642 patients. We split the original dataset into || = environments defined by the trimester of birth. We then use 20 other variables from the original dataset as observed covariates X. Given the nature of the treatment, we expect that some features are post-treatment, i.e. measured after the mother started smoking, as noted in Wilcox (2001). We provide complete experimen- tal details in Appendix D.3.
Table I: ATE estimates for the Cattanco2 dataset using different baselines. We report the mean and standard deviation over 100 initializations of the random seeds the algorithms.
Method ATE (mean + std)
Omull
275.25 + 10-5
0all
-157.55 + 10-5
0rm
1182.65 48.32
0 Oinsta
-214.60 25.20
Table presents the results the differentiable version of our method, alongside various baselines. While the ground truth ATE is unknown, the effect estimated by adjusting for the set selected by Oinsta-c aligns with existing epidemiological literature: both observational and interventional stud- ies (Meyer & Comstock, 1972; Sexton & Hebel, 1984) as well as statistical analyses (Almond et al., 2005; Cattaneo, 2010) estimate a decrease birth weight that ranges from 200 250 grams for infants born to smoking mothers compared to non-smoking mothers. contrast, 0 overestimates the ATE, whereas both dall and frm underestimate it.
DISCUSSION AND FUTURE WORK
In this work, we proposed Robust ATE identification from Multiple ENvironments (RAMEN), method that leverages multiple environments to identify the ATE in the presence of post-treatment and unobserved variables. To the best of our knowledge, we present the first ATE identification guarantees in this highly relevant, but previously unexplored setting. Further, we introduce new version of double robustness that concerns unobserved variables rather than model misspecification.
Nevertheless, our method faces several limitations. First, similar to other kernel-based methods, approach suffers from the curse of dimensionality and the computational complexity associated with computing kernel matrix. Additionally, the requirement for sufficient heterogeneity across environments may be too stringent in some practical cases. Finally, the combinatorial subset is computationally demanding, and the Gumbel trick remains heuristic solution. Addressing any of these shortcomings would constitute interesting avenues for future work.
ACKNOWLEDGMENTS
We thank Jonas Peters for the helpful discussions on an earlier version of this work. PDB was supported by the Hasler Foundation grant number 21050. JK was supported by the SNF grant number 204439. JA was supported by the ETH AI Center. YW was supported in part by the Office of Naval Research under grant number N00014-23-1-2590. the National Science Foundation under
10

---

Published as a conference paper at ICLR 2025
Grant No. 2231174, No. 2310831, No. 2428059, No. 2435696, No. 2440954, and Michigan Institute for Data Science Propelling Original Data Science (PODS) grant. This work was done in part while JK and FY were visiting the Simons Institute for the Theory of Computing.
REFERENCES
Avidit Acharya, Matthew Blackwell, and Maya Sen. Explaining causal findings without bias: de- tecting and assessing direct effects. American Political Science Review, 110(3):512-52, 2016.
Douglas Almond, Kenneth Chay, and David Lee. The costs of low birth weight. The Quarterly Journal of Economics, 120(3):1031-100) 2005.
Martin Arjovsky, Léon Bottou, Ishaan Guirajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint urXiv:1907.0289) 2019.
Dmitry Arkhangelsky and Guido Imbens. Doubly robust identification for causal panel data models. The Econometrics Journal, 25()):649-674, 2022.
Peter Austin. An introduction to propensity score methods for reducing the effects of confounding in observational studies. Multivariate Behavioral Research, 46(3):399-424, 2011.
Rina Foygel Barber, Mathias Drton, Nils Sturma, and Luca Weihs. Half-trek criterion for identifia- bility of latent variable models. The Annals of Statistics, 50(6):3174 3196, 2022.
Reuben Baron and David Kenny. The moderator-mediator variable distinction in social psycho- logical research: Conceptual, strategic, and statistical considerations. Journal of Personality and Social Psychology, 51(6): 173, 1986.
Matias Cattaneo. Efficient semiparametric estimation of multi-valued treatment effects under ignor- ability. Journal of Econometrics, 155(2):138 154, 2010.
Debo Cheng, Jiuyong Li, Lin Liu, Jixue Liu, Kui Yu, and Thuc Duy Le. Causal query in observa- tional data with hidden variables. European Conference on Artificial Intelligence, 2020.
Debo Cheng. Jiuyong Li, Lin Liu, Kui Yu, Thue Duy Le, and Jixue Liu. Toward unique and unbiased causal effect estimation from data with hidden variables. IEEE Transactions on Neural Networks and Learning Systems, (000))6601661111 2022a.
Debo Cheng. Jiuyong Li, Lin Liu, Jiji Zhang, Jixue Liu, and Thuc Duy Le. Local search for efficient causal effect estimation. IEEE Transactions on Knowledge and Data Engineering, 2022b.
Debo Cheng, Jiuyong Li, Lin Liu, Jixue Liu, and Thuc Duy Le. Data-driven causal effect estimation based on graphical causal modelling: survey, ACM Computing Surveys, 56(5):1-37, 2024.
Victor Chernozhukov, Denis Chetverikov, Mert Demirer, Esther Duflo, Christian Hansen, Whitney Newey, and James Robins. Double/debiased machine learning for treatment and structural pa- rameters. The Econometrics Journal, 21(1):C1-C6. 2018.
Piersilvio De Bartolomeis, Javier Abad, Konstantin Donhauser, and Fanny Yang. Hidden yet quan- tifiable: A lower bound for confounding strength using randomized trials. International Confer- ence on Artificial Intelligence and Statistics, 2024a.
Piersilvio De Bartolomeis, Javier Abad, Konstantin Donhauser, and Fanny Yang. Detecting critical treatment effect bias in small subgroups. Uncertainty in Artificial Intelligence, 2024b.
Xavier De Luna, Ingeborg Waembbau, and Thomas Richardson. Covariate selection for the non- parametric estimation of an average treatment effect. Biometrika, 9844):861-875, 2011.
I1ker Demirel, Edward De Brouwer, Zeshan Hussain, Michael Oberst, Anthony Philippakis, and
David Sontag. Benchmarking observational studies with experimental data under right- arXiv preprint arXiv:2402. /51337 2024.
11

---

Published as a conference paper at ICLR 2025
Vincent Dorie. Npci: Non Non-parametrics for causal inference. 2016. URL https://githab: ht////doore/ci....
Mathias Drton, Rina Foygel, and Seth Sullivant. Global identifiability of linear structural equation models. The Annals of Statistics, 39(2):865-886 2011.
Doris Entner, Patrik Hoyer, and Peter Spirtes. Data-driven covariate selection for nonparametric estimation of causal effects. International Conference on Artificial Intelligence and Statistics, 2013.
Zhuangyan Fang and Yangbo He. IDA with background knowledge. Uncertainty in Artificial Intel- ligence, 2020.
Rina Foygel, Jan Draisma, and Mathias Drton. Half-trek criterion for generic identifiability of linear structural equation models. The Annals of Statistics, pp. 1682-1713, 2012.
AmirEmad Ghassami, Saber Salehkaleybar, Negar Kiyavash, and Kun Zhang. Learning causal structures using regression invariance. Advances in Neural Information Processing Systems, 2017.
Arthur Gretton, Karsten Borgwardt, Malte Rasch, Berhard Schölkopf, and Alexander Smola. A kernel two-sample test. The Journal of Machine Learning Research, 13(1):723-1) 273223-77,,00 2012.
Yihong Gu, Cong Fang, Peter Bühlmann, and Jianqing Fan. Causality pursuit from heterogencous environments via neural adversarial invariance learning. arXiv preprint urXiv:2405.047155 2024.
Limor Gultchin, Matt Kusner, Varun Kanade, and Ricardo Silva. Differentiable causal backdoor discovery. International Conference on Artificial Intelligence and Statistics, 2020.
Richard Guo, Anton Rask Lundborg, and Qingyuan Zhao. Confounder selection: Objectives and
approaches. arXiv preprint arXiv:2208. 13871, 2022.
Richard Guo, Emilija Perković, and Andrea Rotnitzky. Variable elimination, graph reduction and the efficient g-formula. Biometrika, 110(3):739-761 2023.
Zijian Guo, Hyunseung Kang, Tony Cai, and Dylan Small. Confidence intervals for causal effects with invalid instruments by using two-stage hard thresholding with voting. Journal of the Royal Statistical Society Series B: Statistical Methodology, 80(4):793-815, 2018.
Jinyong Hahn. Functional restriction and efficiency in causal inference. The Review of Economics and Statistics, 86(1):73-76, 2004.
Jason Hartford, Victor Veitch, Dhanya Sridhar, and Kevin Leyton-Brown. Brown. Valid causal inference with (some) invalid instruments. International Conference on Machine Learning, 2021.
Fernando Pires Hartwig, George Davey Smith, and Jack Bowden. Robust inference in summary data mendelian randomization via the zero modal pleiotropy assumption. International Journal of Epidemiology, 4666):1985-1999 2017.
Christina Heinze-Deml, Jonas Peters, and Nicolai Meinshausen. Invariant causal prediction for nonlinear models. Journal of Causal Inference, 602):20170016, 2018.
Leonard Henckel, Emilija Perković, and Marloes Maathuis. Graphical criteria for efficient total effect estimation via adjustment in causal linear models. Journal of the Royal Statistical Society Series B: Statistical Methodology, 84(2):579-599, 2022.
Miguel Hernán and James Robins. Causal inference, 2010.
Jennifer Hill. Bayesian nonparametric modeling for causal inference. Journal of Computational and Graphical Statistics, 20(1):217-240, 201
Biwei Huang, Kun Zhang, Jiji Zhang, Joseph Ramsey, Ruben Sanchez-Romero, Clark Glymour, and Bernhard Schölkopf. Causal discovery from haterogameouushoosttttinnaay data. Journal of Machine Learning Research, 21 (89): 2020.
12

---

Published as a conference paper at ICLR 2025
Zeshan Hussain, Michael Oberst, Ming-Chine Chich Shih, and David Sontag. Falsification before trapolation in causal effect estimation. Advances in Neural Information Processing Systems, 35, 2022.
Zeshan Hussain, Ming-Chieh Shih, Michael Oberst, Ilker Demirel, and David Sontag. Falsifica- tion of internal and external validity in observational studies via conditional moment restrictions. International Conference Artificial Intelligence and Statistics, 2023.
Antti Hyttinen, Frederick Eberhardt, and Matti Järvisalo. Do-caleulus when the true graph is un- known. Uncertainty in Artificial Intelligence, 2015.
Kosuke Imai, Luke Keele, Dustin Tingley, and Teppei Yamamoto. Unpacking the black box of causality: Learning about causal mechanisms from experimental and observational studies, Amer- ican Political Science Review, 105(4))765-789, 2011.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmas In- fernational Conference on Learning Representations, 2017.
Hyunseung Kang, Anru Zhang, Tony Cai, and Dylan Small. Instrumental variables estimation with some invalid instruments and its application to mendelian randomization. Journal of the American Statistical Association, 11 (11(533):132-44 2016.
Joseph Kang and Joseph Schafer. Demystifying double robustness: comparison of alternative strategies for estimating a population mean from incomplete data. Statistical Science, pp. 523- 539, 2007.
Rickard Karlsson and Jesse Krijthe. Detecting hidden confounding in observational data using multiple environments. Advances in Neural Information Processing Systems, 37, 2023.
Rickard KA Karlsson and Jesse H Krijthe. Falsification of unconfoundedness by testing indepen- dence of causal mechanisms. arXiv preprint arXiv:2502.0623/ 2025.
Ilmun Kim and Aaditya Ramdas. Dimension-agnostic inference using cross U-statistics. Bernoulli, 30(1):683-711 2024.
Gary King. A hard unsolved problem? Post-treas cestttreatment bias in big social science questions. Hard Problems in Social Science Symposium, 2010.
Zhaobin Kuang, Frederic Sala, Nimit Sohoni, Sen Wu, Aldo Cardera-Palomera,, Jared Dunnmon, James Priest, and Christopher Ré. Ivy: Instrumental variable synthesis for causal inference. International Conference Artificial Intelligence and Statistics, 2020.
Marloes Maathuis and Diego Colombo. A generalized back-door criterion. The Annals of Statisties, 4363):1000-1088, 2015.
Marloes Maathuis, Markus Kalisch, and Peter Bühlmann. Estimating high-dimensional intervention effects from observational data. The Annals of Statistics, (766))3133-3164 2009.
Chris Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous laxation of discrete random variables. International Conference on Learning Representations,
2017.
Daniel Malinsky and Peter Spirtes. Estimating bounds causal effects in high-dimeesional and possibly confounded systems. International Journal of Approximate Reasoning, 88:371-384, 2017.
Sarah Mameche, Jilles Vreeken, and David Kaltenpoth. Identifying confounding from causal mech- anism shifts. International Conference on Artificial Intelligence and Statistics, 2024.
Mary Meyer and George Comstock. Maternal cigarette smoking and perinatal mortality. American Journal of Epidemiology, 96(1):1-10, 1972.
Jacob Montgomery, Brendan Nyhan, and Michelle Torres. How conditioning posttreatment vari- ables can ruin your experiment and what to do about it. American Journal of Political Science, 62 (3):7607775, 2018.
13

---

Published as a conference paper at ICLR 2025
Marco Morucci, Vittorio Orlandi, Harsh Parikh, Sudeepa Roy, Cynthia Rudin, and Alexander Vol- fovsky. double machine learning approach to combining experimental and observational data. arXiv preprint arXhv:2307.01449, 2023.
Judea Pearl. Causal diagrams for empirical research. Biometrika, 82(4):669-688, 1995.
Judea Pearl. Direct and indirect effects. Probabilistic and causal inference: the works of Judea Pearl, pp. 373-392, 2022.
Emilija Perković, Markus Kalisch, and Marloes Maathuis. Interpreting and using CPDAGs with background knowledge. Uncertainty in Artificial Intelligence, 2017.
Emilija Perković, Johannes Textor, Markus Kalisch, and Marloes Maathuis. Complete graphical characterization and construction of adjustment sets in Markov equivalence classes of ancestral graphs. Journal of Machine Learning Research, 18(22)):1-62, 2018.
Jonas Peters, Peter Bühlmann, and Nicolai Meinshausen. Causal inference by using invariant pre- diction: identification and confidence intervals. Journal of the Royal Statistical Society Series B: Statistical Methodology, 78(5):947-1012, 2016.
Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Elements causal inference: foundations and learning algorithms. The MIT Press, 2017.
Niklas Pfister, Peter Bühlmann, and Jonas Peters. Invariant causal prediction for sequential data. Journal of the American Statistical Association, 2019.
Niklas Pfister, Evan Williams, Jonas Peters, Ruedi Aebersold, and Peter Bühlmann. Stabilizing variable selection and regression. The Annals of Applied Statistics, 15(3):120-1226,, 2021.
Kristopher Preacher and Andrew Hayes. SPSS and SAS procedures for estimating indirect effects in simple mediation models. Behavior research methods, instruments, & &computers 36:717-731, 2004.
James Robins and Sander Greenland. Identifiability and exchangeability for direct and indirect effects. Epidemiology, 3(2):143-155, 1992.
James Robins and Andrea Rotnitzky. Semiparametric efficiency in multivariate regression models with missing data. Journal of the American Statistical Association, (2(429)::22-129) 1995.
James Robins, Andrea Rotnitzky, and Lue Ping Zhao. Estimation of regression coefficients when some regressors not always observed. Journal of the American Statistical Association, 89 (427):846-86) 866, 1994.
Mateo Rojas- Carulla, Bernhard Schölkopf, Richard Turner, and Jonas Peters. Invariant models for causal transfer leamning. Journal of Machine Learning Research, 19(36):1-34, 2018.
Elan Rosenfeld, Pradeep Kumar Ravikumar, and Andrej Risteski. The risks of invariant risk mini- mization. International Conference on Learning Representations, 2021.
Dominik Rothenhäusler, Peter Bühlmann, and Nicolai Meinshausen. Causal Dantzig. The Annals of Statistics, 47(3):1688-1777 2019.
Dominik Rothenhäusler, Nicolai Meinshausen, Peter Bühlmann, and Jonas Peters. Anchor regres- sion: Heterogeneous data meet causality. Journal of the Royal Stutistical Suciety Series B: Statis- fical Methodology, 83(2):215-246, 2021.
Andrea Rotnitzky and Ezequiel Smucler. Efficient adjustment sets for population average causal treatment effect estimation in graphical models. Journal of Machine Learning Research, 21(1): 7642-7727, 2020.
Mary Sexton and Richard Hebel. A clinical trial of change maternal smoking and its effect on birth weight. Journal of the American Medical Association, (7):911-915, 1984.
14

---

Published as a conference paper at ICLR 2025
Abhin Shah, Karthikeyan Shanmugam, and Kartik Ahuja. Finding valid adjustments under non- ignorability with minimal DAG knowledge. International Conference on Artificial Intelligence and Statistics, 2022.
Abhin Shah, Karthikeyan Shanmugam, and Murat Kocaoglu. Front-door adjustment beyond Markov equivalence with limited graph knowledge. Advances in Neural Information Processing Systems, 2024.
Uri Shalit, Fredrik Johansson, and David Sontag. Estimating individual treatment effect: general- ization bounds and algorithms. International Conference on Machine Learning, 2017.
Xinwei Shen, Peter Bühhlmann and Armeen Taeb. Causality-oriented robustness: exploiting general additive interventions. arXiv preprint arXiv:2307.10299, 2023.
Claudia Shi, Victor Veitch, and David Blei. Invariant representation learning for treatment effect estimation. Uncertainty in Artificial Intelligence. 2021.
Ilya Shpitser, Tyler VanderWeele, and James Robins. On the validity of covariate adjustment for estimating causal effects. Uncertainty Artificial Intelligence, 2010.
Arvid Sjölander. Propensity scores and M-structures. Statistics Medicine, 2869):1466-1420, 2009.
Ingo Steinwart. On the influence of the kernel on the consistency of support vector machines. Journal of Machine Learning Research, 2:67-93, 2001.
Ichiro Takeuchi, Quoc Le, Timothy Sears, Alexander Smola, and Chris Williams. Nonparametric quantile estimation. Journal of Machine Learning Research, 7(7), 2006.
Eric Tchetgen Tehetgen and Tyler VanderWeele. Identification of natural direct effects when confounder the mediator directly affected by exposure. Epidemiology, 25(2):282-291, 2014.
Tyler Vander Weele and Ilya Shpitser. new criterion for confounder selection. Biometrics, 67(4): 1413-1413, 2011.
Stijn Vansteelandt, Tyler VanderWeele, Eric Tchetgen Tchetgen, and James Robins. Multiply robust inference for statistical interactions. Journal of the American Statistical Association, 103(484): 1693-1704, 2008
Haotian Wang, Kun Kuang, Haoang Chi, Longqi Yang, Mingyang Geng, Wanrong Huang, and Wenjing Yang. Treatment effect estimation with adjustment feature selection. Conference on Knowledge Discovery and Data Mining, 2023.
Zhenyu Wang, Yifan Hu, Peter Bühlmann, and Zijian Guo. Causal invariance learning via efficient optimization of a nonconvex objective. arXiv preprint arXiv:2412.1/8883 2024.
Luca Weihs, Bill Robinson, Emilie Dufresne, Jennifer Kenkel, Kaic Kubjas Reginald McGee II, McGee II Reginald, Nhan Nguyen, Elina Robeva, and Mathias Drton. Determinantal generaliza- tions of instrumental variables. Journal of Causal Inference, 6(1):2117000, 2018.
Halbert White and Xun Lu. Causal diagrams for treatment effect estimation with application to efficient covariate selection. Review of Economics and Statistics, 153-11945,-, 201
Allen Wilcox. the importance-and and the unimportance-off birthweight. International Journal of Epidemiology, 30(6): 2001.
2001.
Janine Witte, Leonard Henckel, Marloes Maathuis, and Vanessa Didelez. On efficient adjustment in causal graphs. Journal of Machine Learning Research, 21(246):1-45, 2020.
Shu Yang, Chenyin Gao, Donglin Zeng, and Xiaofci Wang. Elastic integrative analysis of ran- domised trial and real-world data for treatment heterogencity estimation. Journal of the Royal Statistical Society Series B: Statistical Methodology, 85(3):575-596, (44 2023.
Mingzhang Yin, Yixin Wang, and David Blei. Optimization- causal estimation from heteroge- neous environments. Journal Machine Learning Research, 25:1-44 2024.
15

---

Published as a conference paper at ICLR 2025
APPENDICES
The following appendices provide deferred proofs, experiment details, and ablation studies.
TABLE OF CONTENTS
Methodology
17
A.1 Discussion of Assumption 3.2
17
A.2 Proof of Theorem
18
A.3 Implementation details
19
A.3.1 Algorithm Combinatorial search over subsets
20
A.3.2 Algorithm 2: Gumbel trick
20
B Extended related work
23
Additional experiments
24
C.1 Robustness to mediators
24
C.2 Robustness to lack of invariance
24
C,3 Additional random graphs experiments
25
C,4 Additional semi-synthetic experiments
25
C.5 Robustness to small sample size
26
C.6 Robustness to violations of Assumption 4.1
26
Experimental details
28
D.1 Synthetic experiments
28
D.2 Infant Health and Development Program (IHDP) Dataset
................................
29
D.3 Cattaneo2
31
16

---

Published as a conference paper at ICLR 2025
A METHODOLOGY
A.1 DISCUSSION OF ASSUMPTION 3.2
First, we observe here that Assumption 3.2 is not minimal "ob- servability" condition the parents of Y and T: in some cases, might still be possible to find a valid adjustment set via the ob- served parents of either T or Y (or both), although no full set of parents was observed (see c.g. Figure 5). However, in such cases, the valid adjustment set the corresponding regression function cannot be recovered via invariance methods, since neither 7 nor Y are invariant across environments. Thus, in a way, Assumption 3.2 is minimal assumption the DAG if one wants to recover the ATE via invariance of conditional expectations.
X1
U1
T
Y
Further, we remark that our assumption neither stronger nor weaker than the commonly used ignorability assumption with re- spect X (Robins Greenland, 1992). For example, parents of both T and Y are unobserved, Assumption 3.2 will not hold, but ignorability could still apply no common parent of and unobserved. Conversely, in graphs with M-bias structures and colliders, the ignorability assumption will not hold.
U2
X2
Figure 5: Although neither full set of parents is observed, one can still find a valid adjustment (X1, X23 (in green).
Below, we give some examples of settings in which Assumption 3.2 holds, including scenarios in the existing invariance literature (Pe- ters et al., 2016; Gu al., 2024):
given multi-enviromment data {IP E). where each distribution pe is induced by an SCM M = (9, (1 PC), all variables (X1,-., Xa, Y) observed, and for all Tpe E it is satisfied that
Fully observed DAG In Peters al. (2016),
y = =(XPa(y)) 5"), e"),"" 5 ~ F. and E Il P
In particular, the independence condition implies equality of conditional distributions Pe P (a())) across environments and thus Assumption 3.2(b).
General DAG with additive noise In Gu et al. (2024), the target variable follows the following data generating process for all e E:
y -xx. + +c; E[c| 0.
where 5* is the "true important variable set". This setting is, in a way, more general than Peters et al. (2016), since the noise variable is not required to be independent of the parent variables- instead, the only condition is on the first conditional moment of e.. Due to the additivity of the noise, Assumption 3.2(b) follows immediately.
General DAG with multiplicative noise We can define a similar setting for multiplicative noise, setting for all e E:
= B[ = c
where S* is, again, the true parent/importan variable set, and c is independent of the environment. We observe that Assumption 3.2(b) follows since holds that
E s = E (X ) Xs-]
g(Xs- E [cx-- cg(Xs.).
General DAG with polynomial noise From the above two examples, it becomes clear that for any Y ((Xss (c), where p a polynomial degree k. have that Assumption 3.2(b) holds if for all E E it holds that
E je
[s=] = qt for all < k.
where S* is the important variable/parent set. This condition is strictly weaker than the indepen- dence condition since l is finite.
17

---

Published as a conference paper at ICLR 2025
A.2 PROOF OF THEOREM
First, we establish that the loss function in Equation (4) attains value of zero any minimizer Se. By Assumption 3.2, there exists an invariant node E (T,Y} whose parents are observed. Since it holds that JPa(V) (Z; V)
and Pa(V) S Iv, conclude that there exists subset S S Iv such that (i; Js( Z; Y)} = 0. Additionally, since the loss function is non-negative, any global minimizer of Equation (4) must have corresponding loss value of zero.
Recall that the RAMEN estimator is given by
0 (S+) = Ep [x((s))++ + (Y (X
-Y- (o ()x (1
Ep
To(Xsa)+
((xsg)
1 ((Xss)
T1
T
where we have slightly rearranged the order of the terms. In the following, we only consider the treated term Ti and prove that Ti = Ex- [ydo[[=1)]] Analogous reasoning for the control term T2 shows that Tg = Ep- [ydd((T=0] and thus proves the claim.
We now consider two cases, depending on whether the minimum is attained for the node or T.
Case 1: Js, (Z;T) = 0. Since we assume that the kernel belongs to a universal RKHS (Steinwart, 2001, Def. 4), it follows from Gretton et al. (2012, Theorem S) that for all E E
msy (Z:T) Eg- [T]
Xsy] = ã (Xsg): P a.s.
(8)
Then, by Assumption 4.1, holds that
#(Xsg) = 7 #(ZpaTT)) P 8.s.
Further, since the measure P dominates pe, it holds that
=(Xs²) = क (PP(T)); Pe a.8.
(9)
Using the tower property, we compute
Egue [1-- T(Xss))
) jn (Xxx)
T
(1=
= Epe 11 [xs)
Epm
Exe
|Xse
a(Xsg)
Epe [itn (Xss))
Epe
(Xxa)] = =0.
Hence,
TY
TY म ((PP((T))
Ti = Ep-
Ez.
7 7(Xs²)
where the second equality follows by (9).
Now, under the positivity assumption and Assumption 3.2, since the parents of T satisfy the back- door criteria, we can identify the treatment effect, that is, holds that
TY
[yda(T=1)]
Epe
Ep
T T(ZPa(T))
Thus, conclude that
dd)(T=1)]
TY 7(Xsp)
1
T
in (i(Xse)
Epe
Epe
क (Xss)
The analogous argument can be used to show T3 = Ege [dd)(T=0)).
Combining the terms yields
0 (Sa)) = Epe [ddo(==1)
Eps [ydo(T=o)]
18

---

Published as a conference paper at ICLR 2025
Case 2: Jsy. (Z:Y) 0.
We recall that this term of the loss function is defined
Js(Z;Y) = max max (0.1))) AEf
sup (RII(
Epp [[Y - EE(y
=
Thus, again, by the universal property of the kernel lt and Gretton et (2012, Theorem 5), it follows
E Ep- Y T= t, jul (Xsy), P a.s., At {0,1}.
(10)
Further, by Assumption 4.1, have
(((X) = Toe (r(Zra()), P P a.s.
This also implies 1 a.s. equality of the conditional means for all & by domination of measure. We now observe that
Eg. i1 (Xss)) Eg- (Z((Zr))) Epe
since ZPa(Y) is valid adjustment set. Then, we simplify the second treatment term in our estimand using the tower property the invariance of the conditional expectation:
in (TT
Ep-
Ep-
#(Xss) Ep- I(Y - -ii(Xs))) T (T Xss)
7(Xsa)
The RHS is equal since Epe [Y T in (Xsx)) pe- the invariance
from Equation (10) holds.
With an analogous argument for the control term it follows that
(Y Ro
Epe
= 0,
1 т (Xs})
and thus, the RAMEN estimator is equivalent to the ATE:
or (S+)) = Ep-
[ydo(T=1)]
Ep ())
A.3 IMPLEMENTATION DETAILS
In this section, we describe all the implementation details for our methodology.
Estimation of the loss function We have several choices when comes to estimating our loss function, there is trade-off between statistical and computational efficiency. For instance, one can choose the linear time estimator proposed in Gretton et al. (2012, Section 6) or the efficient estimator proposed in Kim & Ramdas (2024) that runs in quadratic time. In this paper, estimate
(S) := Ez- 0s(Z; V)k Zss Zg
(V, _00)
using the cross U-statistic from Kim & Ramdas (2024), defined as
n/2
i = Σ
hs(Zi, Vi),
==1
with hs( Vi) =
Σ j=n/2+1
Moreover, we would like the two loss functions, i.e., Jy and Jr, to be on the same scale to avoid any finite sample issues. Therefore, standardize the cross U-statistic by dividing the empirical variance à ((T(())) i.e. the finite sample estimate of the variance term
2² (IT(())) = Ep. (hs((2) (2) - Eze [hs (2)))
19

---

Published as a conference paper at ICLR 2025
Choice of kernel An important issue in practice is the selection of the kernel parameters. We used a Gaussian kereel all of our experiments. We set the bandwidth of the kernel a to be the median distance between points X in the pooled sample- this remains heuristic similar those described in Takeuchi et al. (2006), and the optimum kernel choice is an ongoing area of research.
A.3.1 ALGORITHM 1: COMBINATORIAL SEARCH OVER SUBSETS
We now describe the concrete implementation of our first algorithm.
Since we know that I' is parent of Y, we can simplify our loss function to incorporate this knowl- edge. Let us define the quantity (X,, Y) Y- (Xs), where fir (Xs) == | Xs, 7 t].
We can rewrite the Y-invariance loss function as follows
min SCM
max (E,,1
Exp (Y,,e(Y, (s))k (Xs, X's) OY,, X's) T = -]-
Similarly, we define # (Xxs) [T Xs] and minimize the T-invariance loss function
min
max Epe [6,(T, Xa)k: (Xs. X's) 6 (X)))
SCId eEC
where 6c(Xs, = T (Xs).
We explain how to compute the adjustment set explicitly in Algorithm I, assuming oracle access to the nuisance functions. practice, nuisance functions can be estimated using the pooled data from all environments.
A.3.2 ALGORITHM 2: GUMBEL TRICK
To deal with the computational infeasibility of searching over all possible subsets of covariates, we propose continuous relaxation of the optimization problem that can be efficiently solved using gradient descent. The method involves using Gumbel sampling to create differentiable binary masks for covariate selection, which allows optimization via gradient descent. We present the continuous relaxation Algorithm for obtaining the invariance loss with respect to the node TV the algorithm can be extended analogously to minimize the invariance loss for Yi and Yo. In this algorithm, we replace the max operator over the environments with an average to obtain smoother loss function.
20

---

Published as a conference paper at ICLR 2025
Algorithm Combinatorial search over subsets (0z)
1: Input: Data {(Xe, Ye: Te) Nuisance functions: n, po. P1
2: for each subset S S [d] do
3:
for each enviromment e E E do Compute T-invariance loss using dataset Dc:
4:
n/2
;s(D";T) - 22 1=1 jmn/2+1
Σ
(T- - (XX)))(XX..XXX =(x)))
TT())
;s(D);T)+
Var (Js(De:T))
Compute Y-invariance loss using dataset D restricted samples with T 1:
f/
Js(D"; <
Σ (Y, (Xx)) X,s))& (x())))
2 1=1 j=n/2+1
Js(D";Yi)
Js(D);Yi)
Var (Js(De:Vi))
6: Compute Y-invariance loss using dataset D restricted to samples with = 0:
Js(D<;Yo) +
(Y, xx Mx())))
1=1 j=n/2+1
3s(D";Y) Js(DD; Yo)
Var (Js(Dr:Yo)
7:
end for
8: Compute the worst environment losses:
Js(T)) + max (D);T),, Js(Y)) + max JD(x; Js(Y)) + max JD(%; eEE
eEE
REF
9: end for
10: Return: Se t- argming min (Js(T), max((s((33),
21

---

Published as a conference paper at ICLR 2025
Algorithm 2 Gumbel trick for subset selection (Oinsta-
1: Input: Data Ye, Je)) temperatures Timi: Thinal, interval k, rate a < learning rates
Ngaat: Tonn epochs Mlpaahs
2: Initialize gate weights Wa wy and neural networks 4.: 2g 8 Set T +- Finit
3: for epoch = 1 to Rapochs do
4: if epoch mod & = 0 then
5: Update temperature:
max(Tinal, 7
6: end
7: for each environment E do
8:
Compute masks:
3 sigmoid((w,,) + G1. G2,)/T)
with G1,, G2.5 ~ Gumbel(O, 1) with G G2,j ~ Gumbel(O,) 1)
B sigmoid((wy. + G1,5 G25)/T)
9:
Compute invariance losses:
Jw (D:TT) -
Σ
(Ti (x. for (Xj,B-)
1= j=n/2+1
Juc (DF;T)
Jm., (D²;)) <
Var (Jw, (D:;T)
1
(D²; 6
Σ
(Yi - Jan (X, Xy,n) Son, (Xi,n)) (using Dr with =
n
j=n/2+1
Ju, (D);
(D%;)1 5-
Var Jue, (Do;
Ju, (D%;Y) x j=n/2+1
4
(Y, - fo (X,,,))((Xx,, Xi,Bv)(y -ooy (X,,Bv) (using D with T = 0)
n
Juh (D%;Yo
Ju, (D%;)) - Var (Jus (De; (o)
10:
end for Compute losses:
11:
Σ Jwy (D%:)) Jw, 2 Σ Su + El
(D); (D,
12:
Update parameters:
Wr WT Tigate V w- Swn O₂ +- 02 Thn Ve S
Wy f- wy - Tizate V Van Swy, Oy + Oy - TJn Vo,, Jun, Oy - Oyn Tnn V Juy
end for 14: Determine subsets: ST < fi: Wx > 0}. Sy + {i: Wg > 0}
Calculate exact losses (as defined in Algorithm 1) for each subset S {Sr, Sy
Js(T)
Σ (DD) Js(Y))
Σ Js(D);Y)) Js(Y)) 1
Js(D";Y)
[
16:
Return: Smsta-c argminss min(Js(T), (ax(s
22

---

Published as a conference paper at ICLR 2025
B EXTENDED RELATED WORK
We discuss here the different challenges associated with the problem of identifying and estimating treatment effects. Our focus is to highlight differences and similarities with our methodology-wo leave out the orthogonal problem of statistical efficiency for the sake of clarity, and we refer the reader to Guo et al. (2022); Cheng et al. (2024) for a complete survey of methods.
Covariate selection with pre-treatment covariates Several works have relaxed the causal suffi- ciency assumption, allowing for unobserved variables as long as they not confounders-whiit constraining all observed covariates to be pre- treetreatment. In this setting, the main challenge is M- bias (Sjölander, 2009), which makes adjusting for the full set of covariates instance, EHS (Entner et al., 2013) was one of the first methods to obtain partial identification of treatment effects in this setting, however, at the cost of computational inefficiency. Gultchin et al. (2020) propose a more efficient relaxation for EHS to circumvent the computational inefficiency. Further, several more recent works leverage anchor variables obtain point identification com- putationally efficient way (Cheng et al., 2020; 2022b; Shah et al., 2022). In contrast, our setting is different since we do not assume that all observed covariates are pre-treatment.
viable solution. For
Covariate selection under causal sufficiency When all the variables in the causal graph are observed, the only challenge towards identifiability is the presence of post-treatment covariates that can introduce collider bias. Several methods have been proposed to tackle this setting- e.g- IDA (Maathuis et al., 2009) and its variants (Perković et al., 2017; Fang & He, 2020) aim to learn a complete graph from data and then infer valid adjustment set from to achieve identifiability. However, they suffer from computational inefficiency since they must first learn the entire causal graph, and they only achieve partial identification. More recently, Shi et al. (2021) consider the setting where multiple environments are available and apply invariant risk minimization (IRM) (Ar- jovsky et al., 2019) for treatment effect estimation. However, it is widely known that IRM re- quires many environments- regime (Rosenfeld et al., 2021). Finally, Wang et al. (2023) recently proposed a reinforcement learn- ing approach to identify the treatmente effect. In contrast, our approach achieves point identification while being computationally efficient and not requiring causal sufficiency.
in the number of covariates generalize well even in the linear
Identifiability in linear Gaussian SCMs In linear Gaussian structural causal models (SCMs), the structure of the causal graph imposes algebraic relationships among the entries in the covariance matrix of the associated distribution. Many researchers have exploited these relationships to de- rive graphical criteria for the identifiability of causal effects, even when some confounders remain unobserved. Specifically, several graphical criteria have been identified for deciding whether, in given causal graph, a specific causal effect can be identified from the covariance matrix for almost all linear Gaussian SCMs compatible with the graph (Drton et al., 2011: Foygel et al., 2012; Weihs et al., 2018; Barber et al., 2022). In contrast, our approach does not assume linearity or Gaussianity, and instead, leverage access to multiple heterogeneous data sources for identification.
Combining data from multiple environments Given the challenges associated with estimating treatment effects using non-randomized data, several works propose detecting bias in the treatment effect estimated from observational data by leveraging randomized trials (Yang et al., 2023; Morucci et al., 2023; Hussain et al., 2022; 2023; Demirel et al., 2024; De Bartolomeis et al., 2024b;a), or mul- tiple observational studies (Karlsson & Krijthe, 2023; Mameche et al., 2024; Karlsson & Krijthe, 2025). In contrast, we leverage the heterogeneity across multiple data sources to identify and esti- mate treatment effects in settings without unobserved confounders.
23

---

Published as a conference paper at ICLR 2025
prm
Binata
0g
2.5
2.5
MAE
1.5
MAE
1.5
0.5
0.5
0
0
Y
T
(a) With mediator
(b) Without mediator
Figure 6: Mean absolute enror averaged across environments for the IHDP dataset with descendant of the outcome Y when different invariances are preserved (I' or Y). We consider the setting with (a) mediator between 7 and Y and (b) without mediator. We consider five environments with n 748 points each; mean and standard error are reported over 20 runs.
C ADDITIONAL EXPERIMENTS
C.I ROBUSTNESS TO MEDIATORS
We study here the robustness of our method to violations of Assumption 3.1. More concretely, show how the inclusion of a mediator between 7 and Y affects the ATE estimate for our method and the baselines in several settings. We consider the semi-synthetic experiment setup from Section 5.3, using 2-d22- 2-dimensional descendant of Y as the post-treatment variable. In Figure 6, we present results for two settings: when the treatment or the outcome is invariant across environments (complete experimental details in Appendix D.2). All baselines show slightly worse performance when mediator is included. When T is invariant, our method remains competitive and outperforms the baselines, as the parents of T still form a valid adjustment set despite the mediator. However, both 8 and Piusta- experience significant drop in performance in the Y-invariance setup. i.e. when T-invariance is violated. This expected, as this scenario, we recover the parents of Y, which unfortunately also includes the mediator. closer inspection of the selected subsets reveals that they usually include the mediator, thus failing to estimate the full effect of T on Y. Instead, our method recovers the natural direct effect of T on Y (Pearl, 2022).
0
A
0
Busta +
0g
C.2 ROBUSTNESS TO LACK OF INVARIANCE
Next, we examine the robustness of our method to vio- lations of the invariance in Assumption 3.2. Specifically, we consider again the semi-synthetic experiments of Sec- tion 5.3 in the scenario where neither 7'- Moorr-ini Yiivvariance holds and there are post- post-treatment variables (i.e. not the independent noise setting). We provide the results Fig- ure 7. The performance of our method significantly wors- ens in this setting, with performance close to the Ómull baseline, as often recovers the empty set when no in- variant node is present. Nonetheless, our method still out- performs drmm in all the settings considered.
MAE
1.5
0.5
collider descendant noise
Figure 7: Mcan absolute enror averaged across environments for the IHDP dataset when no invariance preserved. sider five environments with ne = 748 points each; mean and standard error are reported over 20 runs.
con-
24

---

Published as a conference paper at ICLR 2025
jrm
Binata
0g
0.8
0.8
MAE
0.6
MAE
0.6
MAE
0.6
0.4
0.4
0.4
0.2
0.2
Number of environments
Number of environments
Number of environments
(a) TyY-invariance
(b) Y-invariance
(c) T-invariance
2.5
2.5
2
MAE
MAE
MAE
0.5
0.5
0.5
collider descondant
descendant
collider
doscendant
(d) T.Y-invariance
(e) Y-invariance
(f) T-invariance
Figure 8: (Row 1) We plot the mean absolute error averaged across environments when: (a) no unobserved variables and both invariance w.r.t T and Y are preserved; (b) the parents of T are unobserved but the invariance w.r.t. Y is preserved; (c) the parents of Y are unobserved but the invariance w.r.t T is preserved. For all plots, we sample n = 2000 points for each environment; we report mean and standard error over 100 runs. (Row 2) Complete experimental results for the semi- synthetic setup described in Section 5.3 using the IHDP dataset. The plots show the mean absolute error averaged across environments for: (a) both T- and Y-invariance, (b) Y-invariance only, (c) T-invariance only.We We consider five environments with n = 748 points each; mean and standard error reported over 20 runs.
C.3 ADDITIONAL RANDOM GRAPHS EXPERIMENTS
In Figure (Row 1), we report the MAE averaged across environments for all the invariance set- tings. First, observe that across all settings and numbers of available environments, method significantly outperforms existing baselines. Most notably, 0unttas achieves relatively small erross even with a limited number of environments. In contrast, 0rm requires a much larger number of environments to outperform the trivial baselines d and 0an-- Further, when the parents Y are unobserved, 0rrm fails to surpass all trivial baselines, even with many environments-thh outcome is expected, as the Y -iinvariance is broken in this case and 0mm lacks the double robustness.
C.4 ADDITIONAL SEMI-SYNTHETIC EXPERIMENTS
We present the complete experimental results using the IHDP dataset (see Section 5.3) in Fig- ure (Row 2). Specifically, we evaluate our proposed method and the baselines under three con- ditions, where the two-dimensional variable X, acts collider (as described in the main text), descendant, or independent noise. For T-, Y-, and T,Y- invariance, the results align with those obtained in previous sections for linear synthetic experiments. Both 0g and its differentiable ap- proximation, 0intt-c. outperform the baselines in most settings. The sole exception is when the post-treatment variables are independent noise, where 6mi achieves the best performance. In the case of 7 -'invariance, both method and 0mm exhibit slightly worse performance. 0rm gener- ally underperforms, showing the highest eитоr under the independent noise setting. The &mull baseline demonstrates competitive performance overall, likely due to the relatively low influence of confounders in this setup.
25

---

Published as a conference paper at ICLR 2025
ppm
mmull
Binat-
0g
1.4
1.4
1.2
1.2
1
MAE
0.8
MAE
0.8
MAE
0.0
0.6
0.6
0.6
0.4
0.2
0.2
collider
descendant
noise
collider descendant
noise
collder descendant
noise
(a) TyY-invariance
(b) Y-invariance
(c) T-invariance
1.6
1.6
1.6
1.4
1.4
12
1.2
MAE
MAE
MAE
0.8
0.8
0.6
0.6
0.4
0.4
02
0.2
00
collider descendant
noise
collider descendant
noise
collider
descendant
noise
(d) 0.0
(c) 0.3
(f) 0.5
Figure 9: (Row 1) For all the plots: n = 250, d = 3, E8 5. We plot the mean absolute error averaged across environments when: (a) both invariances
Y is preserved: (c) the invariance w.r.t T is preserved. We report mean and standard eπrо over 20 runs. (Row 2) For all the plots: n 2500, We plot the mean absolute eπror averaged across environments for different levels of heterogeneity the data (higher corresponds more heterogeneous data).
preserved; (b) the invariance
3, | 5, and only the invariant w.r.t Y is preserved.
C.5 ROBUSTNESS TO SMALL SAMPLE SIZE
In Figure 9 (a-c), we report the MAE averaged across environments for the three graphical models introduced Figure 2 (Row 2). We can observe that our method remains competitive even in the small sample size regime: our method consistently outperforms all basclines when collider or descendant is present. However, its performance declines in the edge case where the post-treatment variable is independent noise.
C.6 ROBUSTNESS TO VIOLATIONS OF ASSUMPTION 4.1
We evaluate here the robustness our method against violations our identification condition (i.e. Assumption 4.1). To do so, slightly modify the synthetic experiments presented in Figure 2: We introduce a parameter c² to control environment heterogeneity. distribution across environments. Therefore, there is heterogeneity across environments, and Assumption 4.1 is violated. On the other hand, 2 > 0 the mean and variance of X will shift across environments, with larger shifts the parameter e² increases. Therefore, the heterogeneity across environments increases with e2, and Assumption 4.1 1 is more likely to be satisfied.
2 0, has the same
Example C.1 (Post-treatment variables). Let & be the collection of environment indices. For each environment e E E, we first sample U ~ NN0, e2 18+1). Then, the data is given by
Xp, NUU,, 0.5 U?), for i i= 1.....,, ~ Ber (a (3, X, + <)), with 3. NO, Ia-1) and e, ~ ([Ua.0.. 0.5 + U3): N(Ua,
Y = + 8. X, + Ee, with By ~ (N(0 and e, ~ N((,
X, = T b Y Cc. with c NU 0.5+ +U
26

---

Published as a conference paper at ICLR 2025
In Figure 9 (d-f), plot the MAE for different levels of heterogeneity in the data. We can observe that both 0g and frrm suffer significantly when there is no heterogeneity (c 0.0). Nevertheless, our method 0g consistently outperforms Örrm. even under strong violations of the identification con- dition. Moreover, Og remains competitive against all baselines when Assumption 4.1 is only weakly satisfied (< 0.3).
27

---

Published as a conference paper at ICLR 2025
D EXPERIMENTAL DETAILS
Given an adjustment set, we estimate the ATE for each environment & follows
( 1) (X, + (V: --- 1))7:
(Y; pg 0))(1 TT)
0g =
Σ
(X;)
(x()
where (2, Ex- T= X3 = and T5(2) Ep [ Xg z].
For Prm. since the algorithm only learns the outcome function, estimate the ATE
frm
Σ
PS(X,, 1) AS(X,,
D.1 SYNTHETIC EXPERIMENTS
We now describe the data generating process for synthetic experiments in Section 5.1.
Example D.1 (Post-treatment variables). Let & be the collection of environment indices. For each environment e E E, we first sample U ~ NO, Iat1). For a, E {0,1}, we then observe the following variables:
Xp, ~ NUU,, for i i= 1=....,,d .....-- 1;
(N(0,1) if is invariant else
T ~ Ber (o (B) Xp + +))), with Br NO Ia-1) and a (N(Ua,
(N(0,1)
if is invariant
Y = T + 3 Xp + Cy. with By N~(000 Ia-1 and cy (N(Ua,
else
X. T Y + Ecr with e U(Us+1.U2+1)
Further, observe that for each choice of invariance, the post-treatment variable X, can either be descendant of Y (a and = 1), a collider between 7 and Y = and 1), or independent noise (a 0, b 0). Finally, under this data-generating process, the average treatment effect is constant across the environments, and it given by 0= 1, for all E
Random graph data generating process (Section 5.2) We randomly draw graph from the Erdös-Rényi random graph model with a density equal to 0.5 and consider graphs with total num- ber of observed nodes p 20. We do rejection sampling to exclude graphs that either contain mediators (as they violate Assumption 3.1) or do not contain at least a confounder. We then sample data from the resulting DAG via linear structural causal model with Gaussian weights using the causa dag python library, with the only exception being the treatment variable T, which is gen- erated by additionally applying sigmoid function and then sampling from Berouuli distribution. We further post-process the graph, adding post-treatment variable X, = Y + T and removing at random some parents of T Y depending on which invariance we want preserve. Therefore, consider a challenging scenario with both collider and unobserved variables. To sample data from multiple environments
E, within each environment we apply random uniform mean and variance shift to all the nodes in the graph, except for T and Y.
Implementation details We implement our method, 0intta by performing hyperparameter search over the following parameters each iteration: learning rate in the range [0.001, 0.01, 0.1), initial temperature values of [0.5, 0.8, 1.0], and annealing rates of [0.9, 0.95, 0.99]. The optimal combination of these hyperparameters is selected based on minimizing both T-invariance and Y- invariance loss. The outcome functions for dan. 0g. 0 and 6 are estimated using a linear regression model. Logistic regression is used for propensity score estimation.
28

---

Published as a conference paper at ICLR 2025
D.2 INFANT HEALTH AND DEVELOPMENT PROGRAM (IHDP) DATASET
The Infant Health and Development Program (IHDP) dataset is randomized controlled trial fo- cusing on low-birth-weight, premature infants. For our analysis, we keep six continuous covariates from Dorie (2016), representing the child's birth weight, head circumference birth, number of weeks pre-term, birth order, neonatal health index, and mother's age at birth.
Instead of adopting the treatment outcome functions from Dorie (2016), we simulate a more challenging scenario inspired by Kang & Schafer (2007). In this setting, each covariate assigned to the treatment (I") or outcome (Y) undergoes transformation using predefined set of com- plex functions similar to those encountered in real-world applications. We introduce the following relationships:
Confounders:
Three of the six covariates randomly selected act as confounders, affecting both T and Y.
Other pre-treatment covariates: The remaining covariates are assigned to affect either T or but not both.
Post-treatment treatment covariates: We include a two-dimensional post-treatment covariate, denoted Z, whose generation is detailed below.
Environmental variation: To introduce variation across environments, (i) randomly make a parent of either I or Y unobserved (the same one across environments) and (ii) introduce environment-speciin shifts, as detailed below. We apply both the same node (T or Y) so that the other remains invariant.
We set ATE 2 for all environments.
Modeling of T' and Y For each covariate X, affecting we apply randomly chosen transfor- mation 97 (x) from the following set:
1 9T (1))
0.5 1o.5lo((/o+ + 1), (5)² 0.2, exp 2
We then compute the logits for the treatment assignment as:
T Ingits
Σ ST (Xi).
where 3 are coefficients sampled independently from a uniform distribution B 0.5).
The binary treatment T is obtained by applying a sigmoid function to T Tingits and sampling from Bernoulli distribution:
P(T = (Tiigis))
T Bernouli(P([T 1)),
where (2x) = is the sigmoid function.
Similarly, for each covariate X, affecting Y, we apply a randomly chosen transformation 96 (x) from the set:
(x) { 21 log(l=)),
2 x + 1, exp 2
2
The outcome Y then computed as:
Y= T= 0 (x))
with coefficients 8 sampled from 849 ~ ~2(222222)
29

---

Published as a conference paper at ICLR 2025
Incorporating environment-speccficc shifts To introduce environment--peccin variability, de- fine a hidden variable U that modifies the pre-treatment and post-treatment covariates, outcome, and treatment assignment across different environments. The environments indexed by 0, 1, 2, 4. For each environment, we introduce shifts dependent on u.
We first sample coefficients:
Pun (005,,
3x ~ ((..5 1.0).
For each environment u, the shifts are generated as:
inw u Pun + Gives Ax u fx cx, A = a 3x Cposts
where all Eins EX; Spost are independently sampled from (0.
Then, for each environment, the covariates are modified:
x x + Ax,
where X represents the original covariate values.
Either or T is also shifted, depending on the invariance we aim to preserve:
If invariance in T Y = y + Ame: else invariance Y Toogis logits + in:
while we add ()) to the invariant node.
Generation of post-treatment variables X. For each environment, generate two- dimensional post-treatment variable X. as follows:
Collider:
X- Y + T (posts
Cpost N( Iz).
Descendant:
X. Y + Epost:
Epost
N I2).
Independent Noise:
X. Eposts
Epost
N I2).
Inclusion of mediators In some settings, we introduce an additional mediator variable influenced
by T:
Mediator = 9mod T Emeds
Bmed ~ 4(-110,,,1100 1.0), <med ~ ~N(00, 1).
The outcome Y is then adjusted:
+ Mediator.
Summary of data generation process
For each environment:
1. Modify covariates: X x + Ax.
2. Compute treatment: Tiogis E, (44 .9% (X.) Tw Bernouli((
3. Compute outcome: Y E, a(1) ,9
(1) (X,).
4. Apply environmental shift to Y or T and hide parent of Y or T (we hide the same parent for all environments).
5. Include the ATE 2.0 the outcome
6. If applicable, generate mediator and adjust Y
7. Generate post-treatment variables Z.
30

---

Published as a conference paper at ICLR 2025
Implementation details We implement method, mmsta-e. by performing hyperparameter search over the following parameters at each iteration: learning rate in the range [0.001. 0.01, 0.1], initial temperature values of [0.5, 0.8, 1.0], and annealing rates of [0.9, 0.95, 0.99]. The optimal selected based on the minimization of both T-invariance
combination of these hyperparameters and Y-invariance loss. The outcome and treatment assignment functions for both 0an and Oissta- are estimated using XGBoost. For these models, we set the number of estimators learning rate to 0,01, and the maximum tree depth to 6. For the non- IRM baseline, we employ the TARNet architecture (Shalit et al., 2017), which consists of shared representation with single hidden layer of 200 neurons, followed by two hypothesis-ppecife hidden layers, each with 100 neurons. Logistic regression is used for propensity score estimation.
1,000, the
D.3 CATTANEO2
The Cattanco2 dataset (Cattaneo, 2010) studies the effect of maternal smoking on newborn birth weight. We consider 21 covariates, including maternal and paternal age and education, marital status, maternal foreign status, Hispanic origin, alcohol consumption, receipt of prenatal care and the number of prenatal visits, whether the mother had previous children who died, an indicator for low birth weight, months since last birth by the mother, birth month, indicator for whether the baby first-born, and other variables for which full details are unavailable. The treatment indicator of smoking status, with 864 mothers in the treatment group and 3.778 in the control group. The outcome is a continuous variable representing birth weight, which we normalize to the interval [0, 1]. We exclude the month of birth from the observed features and instead use it to define the environments, creating four environments corresponding to the four quarters of the year.
binary
Implementation details We implement method, @mstta using the following hyperparame- ters: the number of epochs is set to 700, patience to 100, learning rate to 0.1, initial temperature to 1.0, and annealing rate 0.9. This configuration was chosen because it provided robust and fa- vorable results across experiments, specifically in minimizing T- and Y-invariance losses. All other hyperparameters are kept from previous experiments. The outcome and treatment assignment func- tions for both aall and 0iusta- are estimated using XGBoost, with the number of estimators set to 1,000, learning rate to 0.01, and maximum depth to 6. For the non-linear IRM implementation, we use the TARNet architecture, as the IHDP experiments.
31


Reviews:
[{"review_id": "TuIRDrdpod", "reviewer": "Reviewer_r3F3", "summary": "This work addresses the bias arising from adjusting for bad controls in observational causal inference by leveraging invariance conditional properties of either the treatment or the outcome across multiple environments. The methodology includes two practical solutions and they are validated across synthetic, semi-synthetic, and real-world datasets.", "questions": "1. In figure 2a and 2c, if $X_c$ is only a descendant of $Y$, why does adjusting for both covariates lead to bias?\n2. How are the standard errors calculated and why are they significantly higher for the proposed algorithm compared to the baselines in the application?\n3. Can this approach be generalized to non-binary treatments and other estimands?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The issue of bad controls is important in observational causal inference. The proposed approach of excluding them using multi-environment data appears to be a novel idea.\n3. Comprehensive simulations are done to demonstrate the performance and robustness of the proposed algorithms.\n4. The paper is well-written and clear.", "weaknesses": "1. The identification assumptions seem strong and the real-world applicability might be constrained. Are any parts of the assumptions testable using observed data, or can any robustness checks or sensitivity analyses be performed? Have you tested how violations of Assumption 4.1 impact the results?\n2. To provide more convincing results regarding the method's usefulness in real-world applications, could you elaborate more on the selected controls by the algorithm in the birthweight dataset? Additionally, why is it difficult to exclude those potential bad controls or colliders based solely on domain knowledge?\n3. Assumption 4.1 can be renamed since it's one of the identification assumptions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bfrexadmHi", "reviewer": "Reviewer_oUQ4", "summary": "This paper considers a novel setting in which data is collected from heterogeneous environments, aiming to identify causal effects for each environment without prior knowledge of the causal graph. Under certain assumptions, the authors propose two algorithms to identify the target causal quantities. The effectiveness of this approach is demonstrated through extensive experiments.", "questions": "1- Could you discuss the point mentioned above?\n\n2- What does “Descendant” mean in Figure 2?\n\n3- Could you elaborate on Lines 264 and 278? They are not clear to me.\n\n4- Regarding Theorem 1, we understand that the quantity is identifiable under certain assumptions. However, if some assumptions are not satisfied, can you demonstrate that the causal effect is not identifiable? This would be similar to the concept of completeness in the causal effect identification literature.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1- The paper is well-written, and related work is thoroughly discussed. Additionally, the connection between the paper’s assumptions and previous work is clearly presented, for example, following Assumptions 3.3 and 4.1.\n\n2- Various experiments have been conducted, demonstrating the significance of RAMEN.", "weaknesses": "1- The focus of the paper is solely on the identification of treatment effect; therefore, there is no analysis of sample complexity for the proposed algorithm.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "KhXqiK9lUI", "reviewer": "Reviewer_BoRS", "summary": "The paper proposes RAMEN, a method that leverages multiple environments to achieve doubly robust identification of the ATE in the presence of post-treatment and unobserved variables. Empirical evaluations across synthetic, semi-synthetic, and real-world datasets show that the proposed method significantly outperforms existing methods.", "questions": "1. What are the advantages of the proposed method compared with methods using neural networks, such as the method in the literature [1][2].\n2. How is the number of samples in different environments determined in synthetic data experiments? To vary the number of samples per environment, it is recommended that sensitivity analysis experiments be added to synthetic datasets.\n\n[1]Shalit U, Johansson F D, Sontag D. Estimating individual treatment effect: generalization bounds and algorithms[C].ICML’2017.\n\n[2]Shi C, Blei D, Veitch V. Adapting neural networks for the estimation of treatment effects[C]. NIPS’2019.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. This paper estimates causal effects in the presence of post-treatment and unobserved variables.\n2. The paper introduces a novel double robustness property.\n3. The authors demonstrate their method's effectiveness through extensive experiments on synthetic, semi-synthetic, and real-world datasets.", "weaknesses": "1. In the introduction, the explanation of valid and invalid adjustment sets lacks specific examples(such as in advertising recommendations or in the healthcare field), and it is difficult to understand the corresponding scenarios based only on the cause graph.\n2. RAMEN should satisfy the positivity and ignoreability assumptions, which are not given in the problem setting of the paper.\n3. There are many symbols and formulas in the paper. It may be better to list a symbol table.\n4. The experimental evaluation metrics(such as PEHE[1] or ATE[1]) and comparison algorithms(such as[1]) are insufficient.\n\n[1]Shalit U, Johansson F D, Sontag D. Estimating individual treatment effect: generalization bounds and algorithms[C].ICML’2017.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "chEixRl8eS", "reviewer": "Reviewer_6ghW", "summary": "This paper provides a new doubly robust identification framework given multiple data sources, in the sense that, it is able to identify the average treatment effect if in a causal DAG, the parent node of treatment or outcome is fully observed and conditional distribution of either treatment or outcome given their parents are the same across all data sources, without knowing which.\n\nTo identify the adjustment set, the paper proposed two losses based on the minimax problem outlined from the moment condition in the assumption above. \n\nOn the sample level, the paper conduct simulations to examine the performance of their proposed RAMEN estimator.", "questions": "Needs clarification:\n1. You assumed no presence of observed mediators (Assumption 3.2) but keeps emphasizing that the paper allows post-treatment variables and unmeasured variables, so do you mean you allow either unmeasured confounders, unmeasured mediators, or colliders (can be either observed or unobserved);\n2. In Assumption 3.1, you said that \\eps is an exogeneous noise vector following the joint distribution P_\\eps^e over p independent variables. But on Page 4 line 202, you said \"our setting does not require independence of the noise variable\", is this a contradiction?\n3. Page 2 line 83: \"We then provide the first, to our knowledge, doubly robust identification guarantees for treatment effect in the presence of both post-treatment and unobserved variables.\" This contribution is misleading to readers. This approach is not the first approach to handle both post-treatment and unobserved variables, but rather the first doubly robust one (if I understood correctly). For example, for the \"valid adjustment set\" approach, as long as practitioners know this set, it also allows both post-treatment and unobserved variables in the DAG.\n\nAddress a limitation:\n1. In abstract \"Notably, RAMEN achieves doubly robust identification: we identify the treatment effect if either the causal parents of the treatment or those of the outcome are observed. \" This needs more clarification because the doubly robust assumption not only requires either parent is observed but homogeneity of condiitonal distributions across sources of bias.\n2. Solving a minimax problem can be difficult and slow. Can you add comments on the latency (speed) of running your estimator?\n a. Please provide specific runtime measurements for your method on the datasets used in the paper.\n b. Compare these runtimes to those of the baseline methods.\n c. Discuss how the runtime scales with dataset size and number of covariates.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The presentation of the paper is clear. The problem tackled seems interesting. I am not entirely familiar with the literature on this direction so if I believe the contributions of other literature that this paper listed, I think the idea is novel in the literature.", "weaknesses": "1. I think one of key ingredient of this paper is Assumption 4.1. I found this assumption somewhat questionable and hard to believe. \n a. Can you provide what it means under some concrete real-data examples? Maybe explain for your real-data application specifically?\n b. Can you comment the testability and falsifiability of the assumption? Can one touch slightly or comment on some sort of sensitivitiy analysis you can imagine?\n c. Intuitively, not only different data sources should be heterogeneous, but the magnitude matter, especially in estimation when you identified the S_opt. So in the simulation, maybe you can add a sensitivity parameter to represent the strength of heterogeneity of data sources , and then twist that parameter (from zero (Assumption 4.1 fails) to strong) and see what happens?\n\n\n2. Apart from 1c above examining assumption 4.1, I think there are multiple angles the simulation can be strengthened, so that readers can better judge the value and contribution of this work.\nFor example, to examine assumption 3.3, can you check a fourth setting where both when both (a) and (b) fails. This is a common practice when evaluation classical doubly robust estimators. We expect RAMEN will fail under this setting, but it can help me to justify the difficulty of your simulations setting. For example, if RAMEN even performs reasonably well under the 4th setting, it means the simulation is too easy and failure of (a) or (b) creates not enough difficulty. I think a reasonably setting would be to combine the scenario when one of Assumption 3.3 (a) and Assumption 3.3 (b) fails in your simulation setting (b) and (c) into a case that Assumption 3.3 fails.\n\n3. In Section 5.4, I found using 4 trimesters of birth as different environment doubtful. Are they just repeated measure of the same pregant women for 4 times? Can you comment on what Assumption 3.3 and 4.1 means in your real-world experiment?\n a. Clarify if these are indeed repeated measures or separate groups of women.\n b. Explain how Assumptions 3.3 and 4.1 are expected to hold in this specific context.\n c. Suggest alternative ways to define environments in this dataset if trimesters are not appropriate.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)