Causal Reasoning in the Presence of Latent Confounders via Neural ADMG Learning
Abstract
Latent confounding has been a long-standing obstacle for causal reasoning from observational data.
One popular approach is to model the data using acyclic directed mixed graphs (ADMGs), which describe ancestral relations between variables using directed and bidirected edges.
However, existing methods using ADMGs are based on either linear functional assumptions or a discrete search that is complicated to use and lacks computational tractability for large datasets.
In this work, we further extend the existing body of work and develop a novel gradientbased approach to learning an ADMG with non-linear functional relations from observational data.
We first show that the presence of latent confounding is identifiable under the assumptions of bow-free ADMGs with non-linear additive noise models.
With this insight, we propose a novel neural causal model based on autoregressive flows.
This not only enables us to model complex causal relationships behind the data, but also estimate their functional relationships (hence treatment effects) simultaneously.
We further validate our approach via experiments on both synthetic and real-world datasets, and demonstrate the competitive performance against relevant baselines.
Introduction
Latent confounding has been a long-standing obstacle for causal reasoning from observational data.
If not properly accounted for, the presence of latent confounding can lead to incorrect evaluation of causal quantities of interest (Pearl, 2009).
Traditional causal discovery methods that account for the presence of latent confoundings, such as the fast causal inference algorithm (FCI) (Spirtes et al., 2000) and its extensions (Colombo et al., 2012; Claassen et al., 2013; Chen et al., 2021), rely on uncovering an equivalence class of acyclic directed mixed graphs (ADMGs) that share the same conditional independencies.
Without additional assumptions, however, these methods might return uninformative results as they cannot distinguish between members of the same Markov equivalence class (Bellot & van der Schaar, 2021).
More recently, causal discovery methods based on structural causal models (SCMs) (Pearl, 1998) have been developed for latent confounding (Nowzohour et al., 2017; Wang & Drton, 2020; Maeda & Shimizu, 2020, 2021; Bhattacharya et al., 2021).
By assuming that the causal effects follow specific functional forms, they have the advantage of being able to distinguish between members of the same Markov equivalence class (Glymour et al., 2019).
Yet, existing approaches either rely on restrictive linear functional assumptions (Bhattacharya et al., 2021; Maeda & Shimizu, 2020; Bellot van der Schaar, 2021), or discrete search over the space of causal graphs (Maeda & Shimizu, 2021) that are computationally burdensome and unintuitive to use.
As a result, modeling non-linear causal relationships between variables in a scalable way remains an outstanding task.
In this work, we seek to utilize recent advances in differentiable causal discovery (Zheng et al., 2018; Bhattacharya et al., 2021) and neural causal models (Lachapelle et al., 2019; Geffner et al., 2022) to overcome these limitations.
Our core contribution is to extend the framework of differentiable ADMG discovery for linear models (Bhattacharya et al., 2021) to non-linear cases using neural causal models.
This enables us to build efficient and flexible methods capable of discovering non-linear, potentially confounded relationships between variables and perform subsequent causal inference.
Specifically, our contributions include:
Sufficient conditions for ADMG identifiability with non-linear SCMs.
We assume: i) the functional relationship follows non-linear additive noise SCM; ii) the effect of observed and latent variables do not modulate each other, and iii) all latent variables confound a pair of non-adjacent observed nodes.
Under these assumptions, the underlying ground truth ADMG causal graph is identifiable.
A novel gradient-based framework for learning ADMGs from observational data.
Based on our ADMG identification conditions, we further propose Neural ADMG Learning (NADMG), a neural autoregressive-flow-based model capable of learning complex non-linear causal relationships in the presence of latent confounding.
N-ADMG utilizes variational inference to approximate posteriors over causal graphs and latent variables, whilst simultaneously learning the model parameters via gradient-based optimization.
Empirical evaluation on synthetic and real-world datasets.
We evaluate N-ADMG on a variety of synthetic and real-world datasets, comparing performance with a number of existing causal discovery and inference algorithms.
We find that N-ADMG provides competitive or state-of-the-art results on a range of causal reasoning tasks.
Background: acyclic directed mixed graphs (ADMGs) and its magnification
One of the most widely-used graphical representations of causal relationships involving latent confounding is the so-called acyclic directed mixed graph (ADMGs).
ADMGs are an extension of
DAGs, that contain both directed edges ( → ) and bidirected edges ( ↔ ) between variables.
More concretely, the directed edge x i → x j indicates that x i is an ancestor of x j , and the bidirected edge x i ↔ x j indicates that x i and x j share a common, unobserved ancestor.
An ADMG G over a collection of D variables x = ( x 1 , . .
, x D ) can be described using two binary adjacency matrices:
G D ∈ R D × D , for which an entry of 1 in position ( i, j ) indicates the presence of the directed edge x i → x j , and G B ∈ R D × D , for which an entry of 1 in position ( i, j ) indicates the presence of the bidirected edge x i ↔ x j .
Similar to a DAG, an SCM can be specified to describe the causal relationships implied by an ADMG through the so-called magnification process.
As formulated in (Peña, 2016), whenever a bidirected edge x i ↔ x j is present according to G B in an ADMG, we will explicitly add a latent node u m to represent the confounder of x i ↔ x j .
Then, the SCM of x i can be written as x by x i =
f i ( x pa ( i ; G D ) , u pa ( i ; G B ) )
+ ϵ i , where u pa ( i ; G B ) denotes the latent parents of x i in the set of all latent nodes u
= ( u 1 , . .
. , u M ) added in the magnification process.
In the compact form, we can write:
[Equation]
The "magnified SCM" will serve as a practical device for learning ADMGs.
Similar to DAGs, magnified SCMs induce a observational distribution on x , denoted by p θ ( x n ; G ) .
Note that given an ADMG, the magnified SCM is not unique, since latent variables may be shared.
For example, x 1 ↔ x 2 , x 2 ↔ x 3 , x 3 ↔ x 1 can be magnified into both { x 1 ← u 1 → x 2 , x 2 ← u 2 → x 3 , x 3 ← u 3 → x 1 } , and { u 1 → x 1 , u 1 → x 2 , u 1 → x 3 } .
Therefore, the of ADMG identifiability does not imply the structural identifiability of the magnified SCM.
In this paper, we focus on ADMG identifiability.
Establishing ADMGs identifiability under non-linear SCMs
To build flexible methods that are capable of discovering causal relationships under the presence of latent confounding, we first need to establish the identifiability of ADMGs under non-linear SCMs.
The concept of structural identifiability of ADMGs is formalized in the following definition:
Definition 1 (ADMG structural identifiability) .
For a distribution p θ ( x ; G ) , the ADMG G = { G D , G B } is said to be structurally identifiable from p θ ( x ; G ) if there exists no other distribution p θ ′
( x ; G ′ )
such that G ̸ =
G ′ and p θ ( x ; G )
=
p θ ′ ( x ; G ′ ) .
Assuming that our model is correctly specified and p θ 0 ( x ; G 0 ) denotes the true data generating distribution, then ADMG structural identifiability guarantees that if we find some p θ ( x ; G ) =
p θ 0 ( x ; G 0 ) (by e.g., maximum likelihood learning), we can recover G = G 0 .
In this section, we seek to establish sufficient conditions under which the ADMG identifiability is satisfied.
Let x = ( x 1 , . . .
, x D ) be a collection of observed random variables, and u = ( u 1 , . .
. , u M ) be a collection of unobserved (latent) random variables.
Our first assumption is that data generating process can be expressed as a specific non-linear additive noise SCM, in which the effect of the observed and latent variables do not modulate each other :
Assumption 1.
We assume that the data generating process takes the form
[Equation]
where each element of ϵ is independent of all other variables in the model and θ denotes the parameters of the non-linear functions f G D , x and f G B , u .
This assumption is one of the elements that separate us from previous work of (Bhattacharya et al., 2021; Maeda & Shimizu, 2020) (linearity is assumed) and (Maeda & Shimizu, 2021) (effects are assumed to be fully decoupled, x i = (cid:80) m f im ( x m ∈ x pa ( i ; G D ) )
+ (cid:80)
k
g ik ( u k ∈ u pa ( i ; G B ) )
+ ϵ i ).
As discussed in Section 2, the discovery of an ADMG between observed variables amounts to the discovery of their ancestral relationships.
Therefore, the mapping between ADMGs and their magnified SCMs is not one-to-one, which might cause issues when designing causal discovery methods using SCM-based approaches.
To further simplify the underlying latent structure (without losing too much generality), our second assumption assumes that every latent variable is a parentless common cause of a pair of non-adjacent observed variables:
Assumption 2 (Latent variables confound pairs of non-adjacent observed variables) .
For each latent variable u k in the data generating process, there exists a non-adjacent pair x i and x j that is unique to u k , such that x i ← u k → x j .
Arguments for Assumption 2 are made by (Pearl & Verma, 1992), who show that this family of causal graphs is very flexible, and can produce the same conditional independencies amongst the observed variables in any given causal graph.
Therefore, it has been argued that without loss of generality, we can assume latent variables to be exogenous, and have exactly two non-adjacent children.
This assumption also implies that we can safely specify an magnified SCM to the ADMG as shown in Section 2, which allows us to evaluate and optimize the induced likelihood later on.
Given Assumptions 1 and 2, we can prove the ADMG identifiability result:
Proposition 1 (Identifiability of ADMGs under non-linear SCMs) .
Assume Assumptions 1 to 2 hold.
Then, the ADMG matrices G D and G B are identifiable for the data generating process specified in
Equation 2 .
Whilst theoretically we could test for these conditions stated in Lemmas 1 to 3 directly, a more efficient approach is to use differentiable maximum likelihood learning.
Assuming the model is correctly specified, then ADMG identifiability ensures the maximum likelihood estimate recovers the true graph in the limit of infinite data (Appendix C).
Hence, we can design the ADMG identification algorithms via ML learning.
Neural ADMG Learning
Whilst we have outlined sufficient conditions under which the underlying ADMG causal structure can be identified, this does not directly provide a framework through which causal discovery and inference can be performed.
In this section, we seek to formulate a practical framework for gradientbased ADMG identification.
Three challenges that remain are: 1.
How can we parameterize the magnified SCM models for ADMG to enable learning flexible causal relationships?
How can we optimize the causal graphs in the space of ADMGs as assumed Section 3? 3.
How do we learn the ADMG causal structure efficiently, whilst accounting for the missing data ( u ) and graph uncertainties in the finite data regime?
In this section, we present Neural ADMG Learning (NADMG), a novel framework that addresses all three challenges.
4.1 Neural auto-regressive flow parameterization
We assume that our model used to learn ADMGs from data is correctly specified.
That is, it can be written in the same magnified SCM form as in Equation 2:
[Equation]
Following Khemakhem et al.
(2020), we factorise the likelihood p θ ( x n , u n | G ) induced by Equation 3 in an autoregressive manner.
We can rearrange Equation 3 as
[Equation]
where v = ( x , u ) ∈ R D + M , and ˜ G is the magnified adjacency matrix on v , defined as ˜ G i,j ∈ { 0 , 1 } if and only if v i → v j .
This allows us to express the likelihood as
[Equation]
Note that we have omitted the Jacobian-determinant term as it is equal to one since ˜ G is acyclic (Mooij et al., 2011).
Following Geffner et al.
(2022), we adopt an efficient, flexible parameterization for the functions f i taking the form (consistent with Assumption 1)
[Equation]
where ξ 1 ,i , ξ 2 ,i and ℓ i ( i = 1 , ..., D + M ) are MLPs.
A naïve implementation would require training 3( D + M ) neural networks.
Instead, we construct these MLPs so that their weights are shared across nodes as ξ 1 ,i ( · ) = ξ 1 ,i ( e i , · ) , ξ 2 ,i ( · ) =
ξ 2 ,i ( e i , · ) and ℓ i ( · ) =
ℓ ( e i , · ) , with e i ∈
R D + M a trainable embedding that identifies the output and input nodes respectively.
4.2 ADMG learning via maximizing evidence lower bound
Our ADMG identifiability theory in Section 3 suggests that the ground truth ADMG graph can, in principle, be recovered via maximum likelihood learning of p θ ( x | G ) .
However, the aforementioned challenges remain: given a finite number of observations x 1 , . . .
, x N , how do we deal with the corresponding missing data u 1 , . . .
, u N while learning the ADMG?
How do we account for graph uncertainties and ambiguities in the finite data regime?
To address these issues, N-ADMG takes a Bayesian approach toward ADMG learning.
We jointly model the distribution over the ADMG graph G , the observations x 1 , . . .
, x N , and the corresponding latent variables u 1 , . . .
, u N , as
[Equation]
[Equation]
(7) where p θ ( x n , u n | G ) is the neural SCM model specified in Section 4.1, θ denotes the corresponding model parameters, and p ( G ) is some prior distribution over the graph.
Our goal is to learn both the model parameters θ and an approximation to the posterior q ϕ ( u 1 , . . .
, u N , G )
≈ p θ ( u 1 , . . .
, u N , G | x 1 , . . .
, x N ) .
This can be achieved jointly using the variational inference framework (Zhang et al., 2018; Kingma & Welling, 2013), in which we maximize the evidence lower bound (ELBO) of (cid:80) n log p θ ( x n ) given by
[Equation]
We parameterize approximate posterior as q ϕ ( u 1 , . . .
, u N , G ) = (cid:81) n q ϕ ( u n | x n ) q ϕ ( G ) .
For q ϕ ( G ) , we use mean-field Bernoulli distributions to parameterize the ADMG adjancency matrices for G D and G B .
For G D , edge existence and edge orientation are parameterized separately using the ENCO (Lippe et al., 2021).
For the posteriors on u n , we apply amortized Gaussian VI as in VAE literature (Kingma & Welling, 2013), where q ϕ ( u n | x n ) is parameterized via a neural inference net.
In Appendix C we demonstrate that maximizing L ELBO ( θ, ϕ ) recovers the true ADMG causal graph in the limit of infinite data.
4.3 Choice of prior over ADMG graphs
As discussed in Section 2, the ADMG G can be parameterized by two binary adjacency matrices, G D whose entries indicate the presence of a directed edge, and G B whose edges indicate the presence of a bidirected edge.
As discussed in Section 3, a necessary assumption for structural identifiability is that each latent variable is a parent-less confounder of a pair of non-adjacent observed variables.
This further implies that the underlying ADMG must be bow-free (both a directed and a bidirected edge cannot exist between the same pair of observed variables).
This can be enforced by leveraging the bow-free constrain penalty introduced by Bhattacharya et al.
[Equation]
(9) which is non-negative and zero only if ( G D , G B ) is a bow-free ADMG.
As suggested in Geffner et al.
, we implement the prior as
[Equation]
10) where the coefficients α and ρ are gradually increased when maximizing L ELBO ( θ, ϕ ) , following an augmented Lagrangian scheme (Nemirovski, 1999).
Prior knowledge about the sparseness of the graphs is introduced by penalizing the norms ∥ G D ∥ 2 F and ∥ G B ∥ 2 F with scaling coefficients λ s 1 and λ s 2 .
Experiments
We evaluate N-ADMG in performing both causal discovery and causal inference on a number of synthetic datasets.
Additional result on a real-world dataset can be found in appendix I Note that we run our model both with and without the bow-free constraint, identified as N-BF-ADMG (our full model) and N-ADMG (for ablation purpose) , respectively.
We compare the performance of our model against five baselines: DECI (Geffner et al., 2022) (which we refer to as N-DAG for consistency), FCI (Spirtes et al., 2000), RCD (Maeda & Shimizu, 2020), CAM-UV (Maeda & Shimizu, 2021), and DCD (Bhattacharya et al., 2021).
We evaluate the causal discovery performance using F scores for directed and and bidirected adjacency.
The expected values of these metrics are reported using the learned graph posterior (which is deterministic for RCD, CAM-UV, and DCD).
Causal inference is evaluated using the expected ATE as described in
Appendix E. We evaluate the causal inference performance of the causal discovery benchmarks by fixing q ( G ) to either deterministic or uniform categorical distributions on the learned causal graphs, then learning a non-linear flow-based
ANM by optimizing Equation 8 in an identical manner to N-ADMG.
A full list of results and details of the experimental set-up are included in Appendix G.
5.1 Synthetic fork-collider dataset
We construct a synthetic fork-collider dataset consisting of five nodes ( Figure 1a).
The datagenerating process is a non-linear ANM with Gaussian noise.
Variable pairs ( x 2 , x 3 ) , and ( x 3 , x 4 ) are latent-confounded, whereas ( x 4 , x 5 ) share a observed confounder.
We evaluate both causal discovery and inference performances.
For discovery, we evaluate F-score measure on both G D and
G B .
For causal inference, we choose x 4 as the treatment variable taking and x 2 , x 3 and x 5 as the response variables, and evaluate ATE RMSE to benchmark performance.
Figure 1b-1f visualizes the causal discovery results of different methods, and Table 1 summarizes both discovery and inference metrics.
our method N-BF-ADMG-G and CAM-UV achieved the best overall performance.
N-BF-ADMG-G and N-ADMG-G on average are able to recover all bidirected edges from the data, while CAM-UV can only recover half of the latent variables.
Without the bowfree constraint, N-ADMG-G discovers a directed edge from x 4 to x 3 , which results in poor ATE
RMSE performance.
On the other hand, DAG-based method (N-DAG-G) is not able to deal with latent confounders, resulting in the poor f-scores.
Linear ANM-based methods (RCD and DCD) perform significantly worse than other methods, resulting in 0 f-scores for directed matrices and largest ATE errors. xxxxx(a) Ground truth. xxxxx(b) N-BF-ADMG. xxxxx(c) N-ADMG (no constraint).
xxxxx(d)
RCD. xxxxx(e) DCD. xxxxx(f) CAM-UV. 5.2 Random confounded ER synthetic dataset
[Table]
We generate synthetic datasets from
ADMG extension of
Erd˝os-Rényi (ER) graph model (Lachapelle et al., 2019), see Appendix H.2 for details.
We consider the number of nodes, directed edges and latent confounders triplets ( d, e, m ) ∈ { (4 , 6 , 2) , (8 , 20 , 6) , (12 , 50 , 10) } .
The resulting datasets are identified as
ER ( d, e, m ) .
Figure 2 compares the performance of N-ADMG with the baselines.
All variants of N-ADMG outperform the baselines for most datasets, highlighting its effectiveness relative to other methods (even those that employ similar assumptions).
Similar to the fork-collider dataset, we see that methods operating under the assumption of linearity perform poorly when the data-generating process is non-linear.
It is worth noting that even when the exogenous noise of
N-ADMG is misspecified, the causal discovery performance still exceeds that of other methods in most cases.
This robustness is a desirable property as in many settings the form of exogenous noise is unknown.
Conclusion
In this work, we proposed Neural ADMG Learning (N-ADMG), a novel framework for gradientbased causal discovery in the presence of latent confounding for nonlinear SCMs.
Building upon existing work, we established identifiability theory for ADMGs under latent confounding, and proposed a practical ADMG learning algorithm that is both flexible and efficient.
In future work, we will further extend our framework on how to more general cases (e g., the effect from observed and latent variables can modulate in certain forms; latent variables can confound adjacent variables; etc), and how to extend to more complicated scenarios
APPENDIX A Related Work
In this section, we review related works in causal discovery/inference under unobserved confounders.
Causal discovery with latent confounding.
In the literature, constraint-based causal discovery methods in the presence of latent confounding have been well-studied (Spirtes et al., 2000; Zhang, 2008; Colombo et al., 2012; Claassen et al., 2013; Chen et al., 2021).
In those methods, conditional independence tests are called iteratively to identify the graph skeletons and orientations to the extent that is possible.
Without assumptions on the data generating process, these approaches can only identify a Markov equivalence class of causal structures that imply the same conditional independencies.
When certain assumptions are made on the data generating process in the form of
SCMs (Pearl, 1998), additional constraints can help identify the true causal structure.
In the most general case, non-parametric constraints have been identified which can refine causal structures beyond those identified by conditional independency restrictions (Verma & Pearl, 1990; Shpitser et al., 2014; Evans, 2016).
Yet, further refinement can be made through the assumption of stricter SCMs.
For example, in the linear Gaussian additive noise model (ANM) case, Nowzohour et al.
(2017) proposes a score-based approach for finding an equivalent class of bow-free acyclic path diagrams.
Both Maeda & Shimizu (2020) and Wang & Drton (2020) develop additional conditional independency tests for linear non-Gaussian ANM case, with Maeda & Shimizu (2021) extending this to more general cases.
Differentiable characterization of causal discovery.
All aforementioned approaches employ a search over a discrete space of causal structures, which often requires task-specific search procedures, and imposes a computational burden for large-scale problems.
More recently, (Zheng et al., 2018) proposed a differentiable constraint on directed acyclic graphs (DAG), and frames the graph structure learning problem as a differentiable optimization task in the absence of latent confounders.
This is further generalized to the latent confounding case (Bhattacharya et al., 2021) through differentiable algebraic constraints that characterize the space of acyclic directed mixed graphs (ADMGs).
Given these differentiable constraints, the authors define a score based on the Bayesian information criterion to characterize how well each graph fits the data, enabling the deployment of gradientbased optimization techniques to find the best fitting graph within a discrete space.
Nonetheless, this work is limited in that it only considers linear Gaussian ANMs.
Causal inference under latent confounding.
Attempting to perform causal inference in the presence of latent confounding can lead to biased estimates (Pearl, 2012).
Whilst the observed data distribution may still be identifiable, estimating causal effects are not (Spirtes et al., 2000).
A recent string of work has made progress in the case where the effects of multiple interventions are being estimated (Tran & Blei, 2017; Ranganath & Perotte, 2018; Wang & Blei, 2019; D’Amour, 2019).
An alternative approach is to simply assume identifiability of the joint distribution over both latent and observed variables given just the observations.
Louizos et al.
(2017) point out that there are many cases in which this is possible (Khemakhem et al., 2020; Kingma & Welling, 2013).
Nevertheless, all these methods assume the underlying causal graph is known.
More recently, Mohammad-Taheri et al.
(2021) argue that a DAG latent variable model trained on data can be used for down-stream causal inference tasks even if its parameters are non-identifiable, as long as the query can be identified from the observed variables according to the do-calculus.
B Proof of proposition
Given Assumptions 1 and 2, to prove proposition 1, we need the following three lemmas which allow us to identify the ADMG matrices, G D and G B .
These lemmas are extensions of the previous results in Maeda & Shimizu (2021) under our new assumptions above.
We adopt the notation x − i to denote all variables in x excluding x i .
Definition 2 (Residual faithfulness condition) .
We say that nonlinear functions
g i , g j satisfies the residual faithfulness condition, if: for any two arbitrary subset of x , denote by M and N , when both ( x i − g i ( M )) and ( x i − g j ( M )) have terms involving the same exogenous noise ϵ k , then ( x i − g i ( M )) and ( x i − g j ( M )) are mutually dependent.
Lemma 1 (Case 1) .
Assume the data generation process follows Assumptions 1 and 2.
Then, [ G B ] i,j = 1 and [ G D ] i,j = {+
+} if and only if
[Equation]
where g i and g j denote regression functions satisfying the residual faithfulness condition in Definition 2.
Lemma 2 (Case 2) .
Assume the data generation process follows Assumptions 1 and 2.
Then, [ G B ] i,j = 0 and [ G D ] i,j = {+
+} if and only if [Equation]
[Equation]
where g i and g j denote regression functions satisfying the residual faithfulness condition in Definition 2.
Lemma 3 (Case 3) .
Assume the data generation process follows Assumptions 1 and 2.
Then, [ G B ] i,j = 0 and [ G D ] i,j = 1 if and only if
[Equation]
where g i and g j denote regression functions satisfying the residual faithfulness condition in Definition 2.
B.1 Proof of Lemmas 1 to
Here, we provide the proof for Lemmas 1 to 3 in Section 3.
To prove those lemmas, we need the help of the following lemma, which extends Lemma A in (Maeda & Shimizu, 2021), except we don’t assume any form for g i other than non-linearity.
Lemma 4.
Let s ( x i ) denote an arbitrary function of x i .
The residual of s ( x i ) regressed onto x − i cannot be independent of ϵ i :
[Equation]
Proof.
Assume that [ s ( x i ) − g i ( x − i ) |
= ϵ i ] holds, then x − i must contain at least one descendent of x i as it must have dependence on the noise ϵ i to cancel effect of ϵ i in s ( x i ) .
We can express g i ( x − i ) as u i ( ϵ ) .
Since g i operates on variables defined by non-linear transformations of the exogenous noise terms, we cannot express u i as a i ( ϵ − i )
+ b i ( ϵ i ) .
x − i contains a descendent of x i , so ϵ − i includes at least one noise term ϵ k that satisfies x i | =
ϵ k (i.e. is not in x i ).
Thus, terms containing ϵ i cannot be fully removed from s ( x i ) − g i ( x − i ) and so [ s ( x i ) − g i ( x − i ) |
=
ϵ i ] does not hold.
B.2 Proof of Lemma
Proof.
Define g i and g j as g i ( x − i ) =
f i, x ( pa x ( i ))
+ g ′ i
( x − i ) and g j ( x − j ) =
f j, x ( pa x ( j ))
+ g ′ i ( x − j ) respectively.
Then, Equation 11 becomes equivalent to
[Equation]
Given Lemma 4 and following the same arguments as in Maeda & Shimizu (2021), this is equivalent to
[Equation]
[Equation]
Since ϵ i | =
ϵ j , we have
[Equation]
[Equation]
The first implies the existence of an unobserved mediator between x j and x i , the second implies the existence of an unobserved mediator between x i and x j , and the third implies the existence of an unobserved confounder.
Given the assumption of latent variables being confounders and minimality, this indicates the presence of a latent confounder and no direct cause between x i and x j .
B.3 Proof of Lemma
Proof.
When Equation 12 holds, Equation 11 does not.
Thus, there is no unobserved confounder between x i and x j .
Assume that x j is a direct cause of x i , and that Equation 12 is satisfied for g i and g j .
x i contains a nonlinear function of ϵ j that cannot be removed by g i ( x − ( i,j ) ) , thus ( x i − g i ( x − ( i,j ) ))
(cid:0)(cid:0) |
=
ϵ j .
Similarly, ( x j − g j ( x − ( i,j ) ))
(cid:0)(cid:0) |
=
ϵ i .
Thus, we have [( x i − g i ( x − ( i,j ) )
(cid:0)(cid:0)
| =
( x j − g j ( x − ( i,j ) ))] which contradicts our initial assumption.
The same arguments apply when x i is a direct cause of x j , implying that there can be no causal relationship between x i and x j .
B.4 Proof of Lemma
Proof.
When Equation 14 holds, Equation 11 does not and so there is no latent confounder between x i and x j .
When Equation 13 holds, Equation 12 does not hold and so there is a direct causal relationship between x i and x j .
Assume that x j is a direct cause of x i .
Define g 1 ( x − i )
=
f i, x ( pa x ( i )) and g 2 ( x − ( i,j ) ) =
f j, x ( pa x ( j )) , giving x i − g i ( x − i ) =
f i, u ( pa u ( i ))
+ ϵ i and x j − g j ( x − ( i,j ) ) =
f j, u ( pa u ( j ))
+ ϵ j .
When there is no latent confounder, f i, u ( pa u ( i ))
|
= f j, u ( pa u ( j )) .
Thus, ( f i, u ( pa u ( i )) + ϵ i ) | =
( f j, u ( pa u ( j ))
+ ϵ j ) hods and Equation 14 is satisfied.
Now, assume that x i is a direct cause of x j .
Using Lemma 4, we have
[Equation]
[Equation]
Similarly, since x j is a function of x i we also have
[Equation]
Collectively, this implies
[Equation]
which contradicts Equation 14.
Thus, if Equation 14 is satisfied then x j is a direct cause of x i .
B.5 Finishing the proof for proposition
Although Lemmas 1 to 3 cannot cover the scenario where [ G B ] i,j =
and [ G D ] i,j = 1 , note this case is already eliminated by Assumption 2 (confounded pairs must be non-adjacent).
Since each case leads to mutually exclusive conditional independence/dependence constraints, for any p θ ( x ; G ) specified under the assumptions above, there cannot exist some p θ ′ ( x ; G ′ )
such that G ̸ = G ′ and p θ ( x ; G ) =
p θ ′ ( x ; G ′ ) , thus structural identifiability is satisfied.
This completes the proof for our
ADMG identifiability result in proposition 1.
C Maximizing L ELBO ( θ, ϕ ) Recovers the Ground Truth ADMG
In Section 3, we have proved the structural identifiability of ADMGs.
In this section, we further show that under certain assumptions, maximizing L ELBO ( θ, ϕ ) recovers the true ADMG graph (denoted by G 0 ) in the infinite data limit.
This result is stated in the following proposition:
Proposition 2 (Maximizing L ELBO ( θ, ϕ ) recovers the ground truth ADMG) .
Assume that:
• Assumptions 1 and 2 (hence the identifiability of ADMGs) holds for the model p θ ( x ; G ) .
• The model is correctly specified (there exists θ ∗ such that p θ ∗ ( x ; G 0 ) recovers the data-generating process) • Regularity condition: for all θ and G we have E p ( x ; G 0 ) [ | log p θ ( x ; G ) | ] < ∞ .
• The variational family of q ϕ ( u | x , G ) is flexible enough,
i
e., it contains the true posterior p θ ( u | x , G ) .
Then, the solution ( θ ′ , q ′ ϕ ( G )) that maximizes L ELBO ( θ, ϕ ) satisfies q ′ ϕ
( G ) =
δ ( G = G ′ ) , where
G ′ =
G 0 .
The proof of Proposition 2 can be found in Appendix D, which justifies performing causal discovery by maximizing ELBO of the N-ADMG model.
Once the model has been trained and the ADMG has been recovered, we could use the trained N-ADMG to perform causal inference tasks, as detailed in appendix E.
D Proof of Proposition
To prove Proposition 2, we need the following lemma:
Lemma 5.
Assume a variational distribution q ϕ ( G ) over a space of graphs G ϕ , where each graph
G ∈ G ϕ has a non-zero associated weight w ϕ ( G ) .
With the soft prior p ( G ) defined as Equationand bounded λ 1 , λ 2 , ρ, α , we have
[Equation]
Proof.
This directly follows from Lemma 1 of (Geffner et al., 2022).
Now we can proceed to prove Proposition 2:
Proof.
For N-ADMG, in the infinite data limit L ELBO becomes lim
[Equation]
where the zeroing of the KL divergence follows from Lemma 5.
Given fixed θ [Equation] the optimal posterior q ∗ ( u n | x n , G ) satisfies q ∗ ( u n | x n , G )
=
p θ ( u n | x n , G ) due to the flexibility assumption.
Thus,
[Equation]
where p ( x ; G 0 ) denotes the data generation distribution with the ground truth ADMG, G 0 .
Let ( θ ∗ , G ∗ ) =
arg max (cid:82) p ( x ; G 0 ) log p θ ( x | G )
d x be the MLE solution.
Since (cid:80) G ∈G ϕ w ϕ ( G ) = 1 , w ϕ ( G ) > 0 , we have
[Equation]
with the optimal value of (cid:80) G ∈G ϕ w ϕ ( G ) E p ( x ; G 0 ) [log p θ ( x | G )] is achieved when every graph G ∈
G ϕ and associated parameter θ G satisfies
[Equation]
Since the model is correctly specified, the MLE solution ( θ ∗ , G ∗ )
satisfies
[Equation]
Therefore, condition Equation 25 implies for every graph G ′ ∈
G ϕ , G ′ =
G 0 under the regularity condition; or equivalently, G ϕ =
{ G ′ =
G 0 } .
This proves our statement that q ′ ϕ
( G ) =
δ ( G = G ′ ) , where G ′ =
G 0 .
E Estimating Treatment Effects
For all experiments we consider, the causal quantity of interest we wish to estimate is the expected average treatment effect (ATE), E q ϕ ( G ) [ ATE ( a , b | G )] , where the expectation is taken with respect to our learned posterior over causal graphs q ϕ ( G ) :
[Equation]
This requires samples from p ( x Y | do ( x T = b ) , G )
=
p ( x Y | x T = b , G do ( x T ) ) , where G do ( x T ) is the ‘mutilated’ graph obtained by removing incoming edges into x T .
We can achieve this by simulating the learnt SCM on G do ( x T ) whilst keeping x T = b fixed.
Note that q ϕ ( u | x ) is not used to estimate the ATE; it suffices to sample u from the prior distribution, p ( u ) .
In our setting, the inference net q ϕ ( u | x ) is only used as a means through which the likelihood of the data can be evaluated efficiently, and thus model parameters learned.
This is in a similar spirit to deep generative models (Kingma &
Welling, 2013).
F Additional Discussions on Structural Identifiability of Latent Variables
In Section 2, we argued that the identifiability of ADMG does not imply the structural identifiability of the magnified SCM.
In this section, we will present more discussions on certain identifiability of latent structures (of magnified SCMs).
In general, these examples demonstrate that for linear non-Gaussian ANMs the structure of latent variables can be refined beyond the assumption of latent confounders acting between pairs of non-adjacent observed variables, whilst the same techniques cannot achieve the same for non-linear ANMs.
F.1 Determining Causal Structure Amongst Latent Variables
Recently, Cai et al.
(2019) demonstrated that it is possible to discover the structure amongst latent variables using their so-called Triad constraints.
Their method is limited to the linear non-Gaussian
ANM case.
In this section, we demonstrate that an analogous constraint isn’t available for non-linear
ANMs.
[Equation]
Consider the causal graphs shown in Figure 3a and Figure 3b.
Lemma 6 (Linear non-Gaussian identifiability) .
In the linear non-Gaussian ANM case, Equationis satisfied only for the causal graph shown in Figure 3b.
In the non-linear ANM case, Equationis not satisfied for either Figure 3a or Figure 3b.
[Equation]
Sketch of Proof.
For the causal graph in Figure 3a, Equation 27 is equivalent to
[Equation]
On the right, we have some function of the latent variables u 1 and u 2 , f 2 k ( f 12 ( u 1 ) , u 2 ) .
On the left, we have some function of the same noise terms, f 1 i ( u 1 ) − g ( f 2 j ( f 12 ( u 1 ) , ϵ j ) .
To remove u 1 from both sides, we require g to be non-zero and so a term including u 2 is still present.
To remove u 2 , we again require g to be non-zero and so a term including u 1 is still present.
Thus, Equation 27 is not satisfied in either the linear non-Gaussian or non-linear ANM case.
For the causal graph in Figure 3b, Equation 27 is equivalent to
[Equation]
In the linear case, we can construct a linear g to remove u 2 fom the left side so that Equationholds (i.e. g = f
i f 21 f 2 j ).
In the non-linear case, u 1 and u 2 are coupled in the leftmost term and cannot be removed by a term involving u 2 and ϵ j .
Hence, Equation 27 is violated.
F.2 Determining the Number of Latent Confounders
Here, we consider whether the number of latent confounders can be determined in the non-
linear
ANM case.
Lemma 7 shows that in the linear non-Gaussian case, the number of latent confounders acting between a triplet of confounded observed variables can be determined (by verifying certain constraints on marginal distributions on observed variables), whilst the same approach cannot be used in the non-linear ANM case.
Consider the two causal graphs shown in Figure 4a and Figure 4b.
Lemma 7 (Linear non-Gaussian identifiability) .
In the linear non-Gaussian ANM case, Equationis satisfied only for the causal graph shown in Figure 3b.
In the non-linear ANM case, Equationis not satisfied for either Figure 3a or Figure 3b.
[Equation]
Sketch of Proof.
For Figure 4b, Equation 30 is equivalent to
[Equation]
In the linear non-Gaussian case, it is straightforward to set g = f i f j to remove the common noise term u from the left term and make the two sides independent.
In the non-linear case, when f i ̸ =
f j the common noise term u cannot be removed from the left as g must be non-linear, and thus produces a term that involves both u and ϵ j .
For Figure 4a, Equation 30 is equivalent to
[Equation]
u 2 cannot be removed from the left, and so Equation 30 does not hold in either the linear nonGaussian or non-linear case.
[Equation]
G Optimisation Details G.1 Optimisation Details for N-ADMG
As discussed in Section 4, we gradually increase the prior hyperparameters ρ and α throughout training.
This is done using the augmented Lagrangian procedure for optimisation ().
The optimisation process interleaves two steps: 1) optimise the objective for fixed values of ρ and α for a certain number of steps; and ii) update the values of ρ and α .
Steps i) and ii) and ran until convergence, or the maximum number of optimisation steps is reached.
We describe these two steps in more detail below.
Step i).
The objective is optimised for some fixed vales of ρ and α using Adam ().
We use a learning rate of 10 − 3 for the model parameters and 5 × 10 − 3 for the variational parameters.
We optimise the objective for a maxmimum of 5000 steps or until convergence (we stop early if the loss does not improve for 1500 optimisation steps, moving to step ii)).
During training, we reduce the learning rate by a factor of 10 if the training loss does not improve for 1000 steps a maximum of two times.
If we reach the condition a third time, we assume optimisation has converged and move to step ii).
We apply annealing to the KL-divergence between the approximate posterior and prior over the latent variables.
The annealing contant is fixed for each step i), and increased linearly over the first optimisation loops.
Step ii).
We initialise ρ = 1 and α = 0 .
At the beginning of step i) we measure the DAG / bow-free penalty P 1 =
E q ϕ ( G )
[ h ( G )] .
At the beginning of step ii), we measure this penalty again,
P 2 = E q ϕ ( G ) [ h ( G )] .
If P 2 < P 1 , we leave ρ unchanged and update α ← α + ρP 2 .
Otherwise, if
P 2 ≥ 0 .
65 P 1 , we leave α unchanged and update ρ ← 10 ρ .
We repeat the steps i) to ii) a maximum of 30 times or until convergence (measured as α or ρ reaching some max value which we set to 10for both).
G.2 Additional Hyperparameters
Prior Hyperparameters.
We use the sparsity inducing prior hyperparameters λ s, 1 = λ s, 2 =
ELBO approximation.
We construct an approximation to the ELBO in Equation 8 using a single sample from the approximate posteriors.
For evaluating the gradients of the ELBO we use the
Gumbel softmax method with a hard forward pass and soft backward pass with temperature of 0.25.
Neural network architectures.
The functions ξ 1 , ξ 2 and ℓ used in the likelihood and the inference network used to parameterise q ϕ ( u | x ) are all two hidden layer MLPs with 80 hidden units per hidden layer.
Non-Gaussian noise model.
For the non-Gaussian noise model
ATE estimation.
For ATE estimation we compute expectations by drawing 1000 graphs from the learnt posterior, and for each graph we draw two samples of x Y for a total of 2000 samples which we used to form a Monte Carlo estimate.
H Dataset Details H.1 Synthetic Fork-Collider Dataset
We constructed a 2000 sample synthetic dataset with the causal structure shown in Figure 1 by sampling from the following SEM:
[Equation]
Variables u 1 and u 2 are latent confounders acting on variable pairs x 2 and x 3 , and x 3 and x 4 , respectively.
H.2 Latent confounded ER Dataset
We generate synthetic datasets from ADMG extension of Erd˝os-Rényi (ER) graph model (Lachapelle et al., 2019; Zheng et al., 2020).
An ER ( d, e, m ) dataset is generated according the following procedures:
Generate a d × d directed adjacency matrix G D of a ADMG from Erd˝os-Rényi (ER) graph model, whose expected directed edges equal to e ;
Simulate d × d bidirected adjacency matrix G B via random Bernoulli sampling, whose expected number equals to m ; 3.
Simulate exogenous noises
ϵ i from a zero mean Gaussian distribution with standard deviation of 0.1; 4.
Simulate latent variables u from a zero mean Gaussian distribution with standard deviation of 0.1; 5.
Simulate each observed variable as x i =
f i ( x pa ( i ; G D ) )
+ g i ( u pa ( i ; G B ) )
+ ϵ i , where f i , g i are randomly sampled nonlinear functions of the form: y = w T e − x 2 . 6.
Remove u from the sampled dataset.
H.3 IHDP dataset details
This dataset contains measurements of both infants (birth weight, head circumference, etc.) and their mother (smoked cigarettes, drank alcohol, took drugs, etc) during real-life data collected in a randomised experiment.
The main task is to estimate the effect of home visits by specialists on future cognitive test scores of infants.
The outcomes of treatments are simulated artificially as in Hill (2011); hence the outcomes of both treatments (home visits or not) on each subject are known.
Note that for each subject, our models are only exposed to only one of the treatments; the outcomes of the other potential/counterfactual outcomes are hidden from the mode, and are only used for the purpose of ATE evaluation.
To make the task more challenging, additional confoundings are manually introduced by removing a subset (non-white mothers) of the treated children population.
In this way we can construct the IHDP dataset of 747 individuals with 6 continuous covariates and 19 binary covariates.
We use 10 replicates of different simulations based on setting B (log-linear response surfaces) of Hill (2011), which can downloaded from https://github.com/AMLab-Amsterdam/CEVAE .
We use a 70%/30% train-test split ratio.
Before training our models, all continuous covariates are normalised.
I Additional Results I.1 Causal inference on Infant Health and Development Program (IHDP) Dataset
For the real-world datasets, we evaluate treatment effect estimation performances on infant health and development program data (IHDP).
This dataset contains measurements of both infants and their mother during real-life data collected in a randomized experiment.
The main task is to estimate the effect of home visits by specialists on future cognitive test scores of infants, where the ground truth outcomes are simulated as in (Hill, 2011).
To make the task more challenging, additional confoundings are manually introduced by removing a subset (non-white mothers) of the treated children population.
More details can be found in Appendix H.3.
We first perform causal discovery to learn the underlying ADMG of the dataset, and then perform causal inference.
Since the true causal graph is unknown, we evaluate the causal inference performance of each method by estimating the ATE
RMSE.
Apart from the aforementioned baselines, here we introduce four more methods: PC-DWL (PC algorithm for discovery, DoWhy (Sharma et al., 2021) linear adjustment for inference); PC-DwNL (PC for discovery, DoWhy double machine learning for inference); N-DAG-G-DwL (N-DAG Gaussian for discovery, linear adjustment for inference); and N-DAG-S-DwL (N-DAG Spline for discovery, linear adjustment for inference).
Results are summarized in Figure 5a-Figure 5c.
Generally, models with non-Gaussian exogenous assumptions tend to have lower ATE estimation errors; while models with linear assumptions (RCD and DCD) have the worst ATE RMSE.
Interestingly, the DoWhybased plug-in estimators tend to worsen the performances of SCM models.
However, regardless of the assumptions made on exogenous noise, our method (N-BF-ADMG-G and N-BF-ADMG-S) consistently outperforms all other baselines with the same noise.
It is evident that for causal inference in real-world datasets, the ability of N-BF-ADMG to handle latent confoundings and nonlinear causal relationships becomes very effective.
[Figure]