BaCaDI: Bayesian Causal Discovery with Unknown Interventions
Abstract
Learning causal structures from observation and experimentation is a central task in many domains.
For example, in biology, recent advances allow us to obtain single-cell expression data under multiple interventions such as drugs or gene knockouts.
However, a key challenge is that often the targets of the interventions are uncertain or unknown .
Thus, standard causal discovery methods can no longer be used.
To fill this gap, we propose a
Bayesian framework (BaCaDI) for discovering the causal structure that underlies data generated under various unknown experimental/interventional conditions.
BaCaDI is fully differentiable and operates in the continuous space of latent probabilistic representations of both causal structures and interventions.
This enables us to approximate complex posteriors via gradient-based variational inference and to reason about the epistemic uncertainty in the predicted structure.
In experiments on synthetic causal discovery tasks and simulated gene-expression data, BaCaDI outperforms related methods in identifying causal structures and intervention targets.
Finally, we demonstrate that, thanks to its rigorous Bayesian approach, our method provides well-calibrated uncertainty estimates.
INTRODUCTION
Identifying causal dependencies by empirical observation and experimentation is a problem of fundamental scientific interest.
If we understand the causal mechanisms that govern a system of interest, we can predict its behavior when parts of system are actively manipulated from outside.
For instance, if we understand the causal structure of a gene regulatory network, we can predict the effect of drugs or gene knockouts more reliably [Meinshausen et al., 2016].
While identifying the true causal structure from observational data is impossible in many cases [Pearl, 2009, Peters et al., 2017], intervening on some variables in the system and observing the outcome can provide valuable information helping with identification
In some settings, we can collect a range of datasets from the same causal system under multiple different interventions.
For instance, recent advances in biology allow us to obtain single-cell gene expression data under interventions such as different drug candidates or gene knockouts [Srivatsan et al., 2020, McFarland et al., 2020].
Causal discovery methods that can work with interventional data typically assume knowledge of the targets, i.e., which variables have been intervened upon, or the statistical effect of the interventions [Hauser and Bühlmann, 2012, Wang et al., 2017, Yang et al., 2018].
However, this limits their applicability because the intervention targets and effect are often uncertain or unknown in practice, for instance, due to unknown off-target effects upon drug administration.
An additional challenge is that we typically only have a small number of data points per intervention due to high experimental costs.
This renders joint inference over causal structures and interventions brittle, and a rigorous treatment of uncertainty is paramount for reaching reliable conclusions.
Recent methods [Mooij et al., 2016, Ke et al., 2019, Brouillard et al., 2020, Squires et al., 2020] that are able to deal with imperfect or unknown interventions do not account for epistemic uncertainty during inference.
Thus, they frequently fail in realistic scenarios where interventional data is scarce.
Addressing these shortcomings, we introduce Bayesian
Causal Discovery with unknown Interventions (BaCaDI) , a fully Bayesian approach for inferring the complete set of causal structure, functional mechanisms and intervention targets given various unknown experimental/interventional conditions.
In particular, BaCaDI performs joint inference over causal Bayesian Networks (CBNs) as well as intervention targets and effects across multiple experimental contexts.
Our method is fully differentiable and operates in the continuous space of latent probabilistic representations of both BNs and interventions.
This makes our approach particularly scalable to causal systems of many variables, and it enables us to approximate complex posteriors via gradient-based particle variational inference.
In a range of experiments with synthetic causal BNs as well as a realistic gene-expression simulator, BaCaDI outperforms related methods in recovering the underlying causal structure and intervention targets.
Moreover, the principled treatment of uncertainty allows BaCaDI to provide well-calibrated uncertainty estimates for its predictions, even when the underlying Bayesian model is misspecified as in the case of gene-expression data.
RELATED WORK
Causal discovery from multiple contexts.
Learning a causal structure from datasets collected from different interventional contexts of the same causal system has been referred to as Joint Causal Inference (JCI)
[Mooij et al., 2016].
Multiple methods handle such combinations of observational and interventional data [Magliacane et al., 2016,
Zhang et al., 2017, Hauser and Bühlmann, 2012, Wang et al., 2017, Yang et al., 2018] but assume full knowledge of the interventions.
Other methods build on the notion of invariance [Schölkopf et al., 2012, Peters et al., 2016, Meinshausen et al., 2016, Ghassami et al., 2017, Heinze-Deml et al., 2018, Huang et al., 2020], but either do not generalize to full graphs or make restrictive assumptions about the local causal effects.
Causal discovery with unknown interventions.
Mooij et al.
[2016] and Squires et al.
[2020] reduce the unknown intervention setting to standard causal discovery tools such as conditional independence and invariance tests.
However, such hypothesis tests typically require large datasets, making these methods brittle for realistic datasets of small size.
Another line of recent work formulates continuous relaxations of the joint inference problem under unknown interventions and uses gradient-based optimization to find the graph and intervention targets [Ke et al., 2019, Brouillard et al., 2020, Faria et al., 2022].
Our approach uses similar ideas to relax the semi-discrete inference problem into a continuous one.
However, we adopt a fully Bayesian approach that enables principled uncertainty quantification and does not only give point estimates.
Eaton and Murphy [2007] is the only work we are aware of that does the same, though only for discrete variables and with a costly dynamic programming approach that does not scale well to larger graphs.
BACKGROUND: CAUSAL DISCOVERY
Causal Bayesian Networks.
A Bayesian network (BN) ( G , Θ ) uses a directed acyclic graph (DAG) to model the joint density p ( x ) of d variables x
= x 1: d via conditional probabilities.
The joint distribution p factorizes as p ( x 1 , . . .
, x d | Θ , G ) =
(cid:81) di
=1 p i ( x i | x pa G ( i ) , Θ ) where pa G ( i ) is the set of parents of node i in G and the parameters Θ describe the exact local conditional distributions.
In a causal BN (CBN), the edges also describe direct causal relations.
For causal structure learning, we assume that there are no hidden confounding variables (causal sufficiency)
[Pearl, 2009, Spirtes et al., 2000, Peters et al., 2017].
Interventions.
In the causal graph, an intervention on the variable i corresponds to changing the conditional distribution p i and replacing it by a new distribution p Ii .
An intervention is typically considered imperfect (soft) if the distribution is changed but the dependence on the causal parents remains, or perfect (hard, structural) if all dependencies to the causal parents are removed, resulting in a mutilated graph G I k
[e.g. Pearl, 2009, Peters et al., 2017].
In this work, we assume the setting of a collection of M interventions I :=
( I 1 , . . .
I M ) , where each intervention
I k :=
( I tar k , Θ I k ) acts on a set of targets I tar k ⊆ { 1 , . . .
, d } .
We use Θ I k to denote parameters that describe the conditional distributions p I
k
i ( x i | x pa G ( i ) ,
Θ I k ) induced by the intervention I k on the target variables { x i | i ∈ I tar k } .
To keep the notation and following exposition simple, we assume perfect interventions, i.e., p I k
i
( x i | x pa G ( i ) , Θ I k )
= p
I k i ( x i | Θ I k ) .
However, all the arguments made in the remainder of the paper also hold for soft interventions.
The full data distribution under intervention I k factorizes into the observational and interventional conditionals:
[Equation]
The local conditional distributions of the variables that are not intervened upon do not change with respect to the observational distribution, a principle often referred to as invariance [Peters et al., 2016] or modularity [Peters et al., 2017].
BAYESIAN CAUSAL DISCOVERY WITH MULTI-CONTEXT DATA
Problem Statement.
In this section, we develop a method for Bayesian inference of the causal Bayesian
Network (CBN) given multiple interventional datasets generated from the same underlying causal model.
This setting is also known as learning from multiple contexts
Mooij et al.
Formally, we are given a set of M datasets D =
{D 1 , ..., D M } with corresponding (unknown) interventions
I =
{ I 1 , ..., I M } .
Each D k is a set of independent samples D k
= { x ( k, 1) , . . .
, x ( k,n k ) } obtained from the interventional data distribution p ( x | Θ gt , G gt , I k ) based on the ground truth CBN ( G gt , Θ gt ) .
Observational data, if available, can be added to D as D 0 :
=
D with intervention targets I 0 = ∅ .
Our goal is to infer the ground truth CBN ( G gt , Θ gt ) and interventions I given the dataset D .
Compared to standard causal inference, the key difficulty is that in addition to the ground truth CBN, the intervention targets (i.e I tar k ) and their statistical effects (i.e. Θ I k ) are unknown.
Therefore, we need to perform joint inference over M mutilated graphs that are all closely connected to an unknown “prototype” graph.
Naturally, such inference is well-posed only if the structural changes implied by each intervention I k are sparse compared to the overall size of G gt , i.e. | I k | ≪ d .
In many relevant application domains, such as biology, the observed samples n k are noisy and small in number.
Hence, it is paramount to not only to predict a single prototype
CBN alongside one intervention hypothesis per dataset, but to also to reason about the epistemic uncertainty of our empirical inferences.
Such uncertainty estimates allow us to quantify the reliability of our predictions and can be used to actively design future experiments.
We thus approach the problem from a Bayesian perspective.
This renders the task of learning from D as a posterior inference problem, which we are going to gradually build up in the following.
Known interventions.
When the intervention targets I tar k and the parameters
Θ I k of the intervention effect are known, for k = 1 , . . .
, M , the posterior over CBNs includes (i.) the product of data likelihoods over all datasets in D , and (ii.) interventional instead of observation likelihoods for D k ≥ 1 ,
[Equation]
(1) where p ( D k | Θ , G , I k ) = (cid:81)
n k i =1 p ( x ( k,i )
| Θ , G , I k ) is the interventional likelihood for D k given I k
= ( I tar k , Θ I k ) .
Unknown interventions.
We include unknown interventions in our inference model by introducing additional priors p ( I tar k ) and p ( Θ I k | I tar k ) .
Accordingly, the modified posterior follows as
[Equation]
The prior p ( I tar k ) over intervention targets can incorporate prior beliefs about the structure of interventions, e.g. that only a sparse set of variables are subject to an intervention at the same time.
The parametrization of the interventional distributions p Ii ( x i | Θ I ) is informed by the application, and reflects the general nature of interventions, e.g., gene knockdowns in biology.
[Figure]
A DIFFERENTIABLE GENERATIVE MODEL OVER CBNS AND INTERVENTIONS
In the following, we represent G ∈ { 0 , 1 } d × d as the adjacency matrix
and I tar k =
[ I tar k, 1 , ..., I tar k,d ] ⊤ ∈ { 0 , 1 } d as the indicator vector where I tar k,l = 1 if the l -th variable is intervened upon
and I tar k,l = 0 otherwise.
We write I tar short for the stack [ I tar 1 , ..., I tar M ] of intervention target masks and Θ I :=
[ Θ I 1 , ..., Θ I M ] similarly for the intervention effect parameters.
The posterior in Eq.
2 is highly intractable as the number of possible DAGs G grows super-exponentially with the number of variables d
[Robinson, 1973].
Furthermore, the number of possible intervention targets I tar k grows in the order of O (2 d ) .
To facilitate approximate inference using Eq. 2, we harness recent advances in Bayesian structure learning proposed by Lorch et al.
[2021] that allows for a fully differentiable posterior over graphs.
By computing the score of the posterior, we can use approximate inference methods such as variational inference [Blei et al., 2017] or Stein Variational
Gradient Descent (SVGD)
[Liu and Wang, 2016].
The key idea is to transform the inference problem over discrete structures G
and I tar k into one over continuous parameters with respect to which we can differentiate.
To this end, we introduce continuous latent variables Z and Γ k for k = 1 , . . .
M that model the generative processes of G
and I tar k through p ( G | Z ) and p ( I tar k | Γ k ) .
This implies the following factorization of the generative model in Figure 1:
[Equation]
where we write Γ :
=
[ Γ 1 , ..., Γ M ] for brevity.
Under this generative model, it holds that
[Equation]
where H = p ( Θ | G ) p ( Θ I |I tar ) p ( D | G , I , Θ ) with p ( D | G , I , Θ ) = (cid:81) Mk =1 p ( D k | G , Θ , I tar k , Θ I k ) and p ( I| Γ ) =
(cid:81) Mk =1 p ( I tar k |
Γ k ) p
( Θ I k | I tar k ) .
A proof is provided in Appx.
A.1.
Eq.
4 allows us to compute Bayesian expectations under the posterior in Eq.
1 in terms of the posterior p ( Z , Θ , Γ , Θ I | D ) over parameters that are continuous.
Before we discuss how to perform approximate inference with p ( Z , Θ , Γ , Θ I | D ) , we further specify some of the conditional probabilities of our generative model and how to make them differentiable.
Generative model of DAGs G .
Following Lorch et al.
[2021], we define the latent variable Z =
[ U , V ] as a stack of two embedding matrices
U , V ∈ R d × d and the generative model for the adjacency matrix G by using the inner product between the latent variables in Z :
[Equation]
with p α ( g ij = 1 | u i , v j ) =
σ α ( u ⊤ i v j ) where σ α ( x ) = 1 / (1 + exp( − αx )) is the sigmoid function with inverse temperature α and u i , v j the i -th and j -th column vector of U and V respectively.
We denote the matrix of edge probabilities in G given Z by G α ( Z ) ∈ [0 , 1] d × d with G α ( Z ) ij :
=
σ α ( u ⊤ i v j ) .
The prior over Z utilizes (i.) factorized Gaussians with a variance of η 2 Z = 1 /d to ensure well-behaved gradients and (ii.)
an acyclicity prior using that penalizes the expected cyclicity of G given Z : p β ( Z ) = p ( U , V ) ∝ exp (cid:0)
[Equation]
( G | Z ) [
[Equation]
( G )
[Equation]
(cid:1)
[Equation]
Here, β is the inverse temperature parameter controlling how strongly the acyclicity is enforced, and h ( G ) = tr (cid:2) ( I + 1 d G ) d (cid:3) − d ≥ 0 .
We make use of Theorem 1 in
Yu et al.
[2019] states that G is acyclic
iff h ( G ) = 0 .
As β → ∞ , the support of p ( Z ) reduces to all Z that model
DAGs [cf.
Lorch et al., 2021].
Generative model of intervention targets I tar .
Similar to the generative model of G , we define the latent variables Γ ∈ R d as the logits of independent Bernoulli distributions that model the entries of the intervention target mask I tar =
[ I tar 1 , ..., I tar M ] ∈ { 0 , 1 } M × d : p ( I tar | [Equation] p
[Equation]
7) with p α
( I tar k,i = 1 | γ k,i ) =
σ α ( γ k,i )
We denote the matrix of intervention target probabilities as I tar α ( Γ ) ∈ [0 , 1] M × d with I tar α ( Γ ) k,i =
σ α ( γ k,i ) .
Similar to Z , the prior over Γ has three components: (i.)
A Gaussian term for numerical stability , (ii.)
a Beta-distribution prior that encourages σ α ( γ k,i ) to be close to 0 or 1, and (iii.)
a sparsity prior with the l 1 -norm of σ α ( Γ k ) and the inverse temperature parameter λ . p ( Γ ) ∝ M (cid:89) k
=
[Equation]
) (cid:124) (cid:123)(cid:122)
[Equation]
sparse masks · d
[Equation]
=
Beta (
[Equation]
(cid:124) (cid:123)(cid:122) (cid:125) sharp masks
N
[Equation]
) (cid:124) (cid:123)(cid:122) (cid:125) numerical stability (8)
We assume that interventions occur only infrequently, i.e., in expectation only on one variable.
Hence, we choose ζ 1 = 1 /d and ζ 2 = ( d − 1) /d .
The sparsity prior implies that, given an active intervention target i , it is a-priori less likely that a variable j ̸
= i is intervened upon in the same context k .
Interventional likelihood.
We obtain a differentiable version of the interventional likelihood by sampling masks
I tar ∼ Bernoulli ( σ α ( Γ )) with the Gumbel-Softmax trick [Jang et al., 2016, Maddison et al., 2017] and using them as a switch between observational and interventional likelihoods per variable:
[Equation]
Stein Variational Gradient Descent.
Having rewritten the inference problem to continuous latent variables Z and Γ , the final challenge is approximating the intractable posterior p ( Z , Θ , Γ , Θ I | D ) .
To this end, we employ the particle variational inference approach SVGD
[Liu and Wang, 2016], which minimizes the KL divergence to the intractable distribution of interest using a finite set of particles.
It relies on the score of the distribution to guide the particles towards regions of high probability while using a kernel function k ( · , · ) between them.
We employ a sum of RBF kernels as the kernel function over Ψ :=
( Z , Θ , Γ , Θ I ) .
The proposed algorithm is summarized in Alg.
1 of Appx.
B. There, we also give details on how to derive the scores of the log-likelihood and the chosen kernel.
Figure 4 in Appx. B illustrates an example of the returned posterior particles for G
and I tar alongside the ground truth for the case of a linear Gaussian 5 node graph.
EXPERIMENTS
We evaluate BaCaDI on different causal discovery tasks with data from multiple contexts.
Our aim is to empirically investigate how accurately BaCaDI predicts the causal structure as well as the intervention targets and how it compares to related state-of-the-art methods.
First, we focus on synthetic datasets.
Second, we use the SERGIO simulator [Dibaeinia and Sinha, 2020] to evaluate the methods on a more realistic task of simulated gene expression data.
Finally, we analyse how well the uncertainty estimates are calibrated.
[Figure]
AUPRC
[Figure]
AUPRC
[Figure]
6.1 EXPERIMENT METHODOLOGY
Datasets.
Following related work [Zheng et al., 2018, Yu et al., 2019, Zheng et al., 2020, Annadani et al., 2021, Scherrer et al., 2021, Lorch et al., 2021], we perform inference of random graphs with different local conditionals.
We consider Erd ˝ os-Rényi [Gilbert, 1959] and scale-free [Barabási and Albert, 1999] random graphs with d = 20 nodes and 2 d edges in expectation (ER-2 and SF-2).
We create datasets by randomly sampling parameters or simulating data with SERGIO.
We then split the data into training and held-out test datasets.
In all settings, we collect n 0 = 100 observational samples together with n k = 10 samples per intervention context for k ∈ { 1 , . . .
, M } .
We defer further details on data generation to Appx.
D.
Baselines.
We compare BaCaDI to existing algorithms that are capable of joint causal inference from multiple contexts with unknown interventions.
This includes the constraintbased methods UT-IGSP [Squires et al., 2020] and the JCI framework with the PC algorithm (JCI-PC)
[Mooij et al., 2016], and the score-based method DCDI [Brouillard et al., 2020] that can handle unknown interventions.
JCI-PC and
UT-IGSP are based on conditional independence or invariance tests.
For DCDI, we use a neural network to model the local conditionals with Gaussian additive noise (DCDIG).
This is the same model class and capacity as BaCaDI with a nonlinear likelihood model.
Since all of these methods arrive only at a single DAG estimate, we use the nonparametric DAG bootstrap approach [Friedman et al., 2013,
Agrawal et al., 2019] to obtain an approximate distribution over DAGs and intervention targets.
Throughout the experiments, 20 bootstrap samples are used.
BaCaDI instantiations.
We instantiate BaCaDI usingparticles for SVGD and run it for 2000 steps.
Unless specified otherwise, we model interventions using a Gaussian likelihood p
I k
i ( x i | Θ I k ) =
N ( x i | µ Ik,i , σ I ) with fixed variance σ I 2 = 0 .
We infer the means Θ I k
=
[ µ Ik, 1 , ..., µ Ik,d ] using a wide
Gaussian prior
[Equation]
This reflects an uninformative prior over a large effect range of the interventions.
The local conditionals for the observational likelihood are either modelled as linear or nonlinear models with additive Gaussian noise.
The latter uses 1 hidden-layer neural networks (NN) with 5 hidden units.
We use 20 particles to instantiate the SVGD algorithm.
More details about the models can be found in Appx.
D.1.
Metrics.
Our reported metrics focus on the essential aspects of our inference problem: causal graph prediction, intervention detection, and inference of the full CBN conditionals.
Causal Discovery:
We focus on the Structural Interventional Distance (SID)
[Peters and Bühlmann, 2015] that quantifies the closeness between two DAGs in terms of how well their interventional adjustment sets coincide.
Since we perform posterior inference, we consider the expected
SID: E -SID ( p, G gt ) :
= (cid:80) G p ( G |D ) · SID ( G , G gt ) .
Since
UT-IGSP and JCI-PC only return a CPDAGs of the Interventional Markov Equivalence Class (I-MEC), we calculate its lower and upper bound of the SID, and report their midpoint as the E -SID (see Appx.
D.5).
Furthermore, we compute the area under the precision recall curve (AUPRC) for pairwise edge prediction based on the posterior marginals p ( g ij = 1 | D ) .
Interventions : We report interventional AUPRC (INTVAUPRC) for the classification of targets.
Learning conditionals: We compute the average negative interventional log-likelihood (I-NLL) (see Appx.
D.5) on
M test = 10 heldout interventional datasets.
Since UT-IGSP and JCI-PC are not equipped with local conditional distributions, we use the linear Gaussian maximum-likelihood parameters (MLE) that are computed in closed-form [Hauser and Bühlmann, 2014] to compute the heldout I-NLL.
Result aggregation.
For all methods and settings, we perform a search over a specified range of hyperparameters with at least 20 settings.
For a specific choice of hyperparameters, we collect results over 30 different random graphs and pick the hyperparameters that resulted in the lowest
I-NLL across the 30 instances in the held-out interventional dataset.
We report the median of each metric together with its 90% confidence interval based on empirical percentiles.
[Figure] G
[Figure]
E -SID 3 .3 .4 .
I-NLL 0 .0 .
AUPRC 0 .0 .1 .
INTV-AUPRC
[Figure]
UT [Figure]
DCDI
[Figure]
G [Figure] (LinG) [Figure]
(NonlinG)
6.2 EXPERIMENT RESULTS ON SYNTHETIC DATA
First, we study the performance of BaCaDI and the baselines on joint causal inference from synthetic data.
We focus on hard interventions that set the value of the targeted variable to a Gaussian with a randomly chosen mean bounded away from zero.
Each interventional dataset is generated by intervening on a specific variable.
We target all variables in the graph, i.e., the number of interventional contexts is M = d .
In total, this results in 300 samples for synthetic datasets with d = 20 variables.
We first consider linear Gaussian
CBNs where each variable is a linear combination of its parents with additive Gaussian noise.
The corresponding results for ER-2 and SF-2 are shown in Fig.
Across the four synthetic evaluation settings, BaCaDI’s causal structure predictions are the closest to the groundtruth CBN a) in terms of their intervention implications (i.e., E -SID, I-NLL) as well as b) for predicting individual edges (i.e., AUPRC).
In most cases, our method outperforms the baselines by a significant margin.
Moreover, BaCaDI achieves strong AUPRC scores for the prediction of intervention targets.
Among the baselines, UT-IGSP is the best at detecting interventions, largely on par with BaCaDI.
Although UT-IGSP also achieves high AUPRC for the linear SF-2 setting, the method significantly falls behind
BaCaDI in other settings and metrics.
6.3 EXPERIMENTS ON SIMULATED GENE-REGULATORY NETWORKS
We evaluate all methods in a realistic application domain using SERGIO [Dibaeinia and Sinha, 2020], a simulator for single-cell expression data of gene regulatory networks.
Given a user-defined causal graph G , SERGIO utilizes stochastic differential equations to simulate the gene expression dynamics and generate realistic single-cell transcriptomic datasets, which correspond to samples from the steady state of this dynamical system.
Since real-world gene regulatory networks resemble scale-free structures [Albert, 2005, Ouma et al., 2018], we use randomly sampled SFgraphs with d = 20 .
In this domain, we perform M =gene knockout interventions on a single randomly selected target per context, resulting in a dataset size of 200 including the observational samples.
See Appx.
D.3 for more details.
Our Bayesian formulation allows us to embed prior knowledge into the inference process in a principled manner.
Since we perform knockout interventions, we expect intervention values to be close to zero and exhibit low variance.
To reflect this prior belief, we set the intervention noise to σ I 2 = 0 .and
use the prior p ( µ Ik,i | I tar k,i = 1)
=
N ( µ Ik,i | 0 , 1) .
Accuracy.
Fig.
3 shows the results for the SERGIO datasets.
As in the synthetic CBN domain, BaCaDI infers the ground-truth graph most accurately given the provided data.
Moreover, it is very accurate in predicting intervention targets as reflected by the intervention target AUPRC score, where it benefits from the prior.
Uncertainty.
We utilize the concept of calibration [Gneiting et al., 2007, Kuleshov et al., 2018] to quantify the reliability of our uncertainty estimates.
The results are shown in Appx.
E.3.
Notably, BaCaDI is the only method that takes into account the epistemic uncertainty when dealing with limited data and shows that its probabilistic predictions are reliable.
CONCLUSION
We introduce BaCaDI, a fully-differentiable Bayesian causal discovery framework for data generated under various unknown experimental/interventional conditions.
BaCaDI performs approximate inference jointly over the underlying graph, mechanisms, and unknown interventions.
A key feature is its principled treatment of epistemic uncertainty which allows it to work reliably even when data is scarce, providing well-calibrated uncertainty estimates alongside its structural predictions.
Our work is motivated by the challenging problem of inferring the causal mechanisms of gene regulatory networks from real single-cell gene expression data.
Our experimental results for the simulated gene expression data show that BaCaDI brings us one step closer to this ambitious goal.
To ultimately reach it, a range of further challenges, such as dealing with the experimental measurement noise incurred by single-cell sequencing techniques, will have to be overcome by future work. Acknowledgements
This research was supported by the European Research
Council (ERC) under the European Union’s Horizonresearch and innovation program grant agreement no.
815943 and the Swiss National Science Foundation under
NCCR Automation, grant agreement 51NF40 180545.
Jonas Rothfuss was supported by an Apple Scholars in AI/ML fellowship. We thank Parnian Kassraie for valuable feedback. A PROOFS A.1 LATENT POSTERIOR EXPECTATION
Recall that we have the generative model under which we factorize
[Equation]
where we write Γ :
=
[ Γ 1 , ..., Γ M ] .
For brevity, we can express this as p ( Z , G , Θ , Γ , I , D ) =
p
[Equation]
Z ) p
[Equation]
G | Z ) p ( Θ | G ) · p ( Γ ) p ( I tar | Γ ) p ( Θ I |I tar ) · p ( D | G , I , Θ )
where p ( D | G , I , Θ ) = p ( D | G , Θ , I tar , Θ I )
= (cid:81)
Mk =1 p ( D k | G , Θ , I tar k , Θ I k ) .
This gives us
E p ( G ,
[Equation]
I| D
[Equation] [Equation]
G ,
[Equation]
, I
[Equation] [Equation]
tar
[Equation]
( splitting I to I tar , Θ I ) = (cid:88)
G (cid:90) Θ (cid:88)
I tar (cid:90) Θ I p ( G , I tar , Θ , Θ I , D ) f
( G , Θ , I ) p ( D )
d Θ d Θ I = (cid:88)
G (cid:90) Θ (cid:88)
I tar (cid:90) Θ I (cid:90)
Z (cid:90) Γ p
( Z , G , Γ , I tar , Θ , Θ I , D )
f
( G , Θ , I ) p ( D )
d Z d Γ d Θ d Θ I ( extending by Z , Γ )
= (cid:90)
Z , Γ , Θ , Θ I (cid:88)
G , I tar p ( Z , G , Γ , I tar , Θ , Θ I , D ) f
( G , Θ , I ) p ( D )
d Z d Γ d Θ d Θ I ( rearranging ) =
(cid:90)
Z , Γ , Θ , Θ I (cid:88)
G , I tar p ( Z ) p ( G | Z ) p ( Θ | G ) p ( Γ ) p ( I tar | Γ ) p ( Θ I |I tar ) p ( D | G , I , Θ )
f ( G , Θ , I ) p ( D )
d Z d Γ d Θ d Θ I ( by the generative model ) =
E p ( Z , Θ , Γ , Θ I | D )   (cid:88)
G , I tar p ( G | Z ) p ( Θ | G ) p
( I tar | Γ ) p ( Θ I |I tar ) p ( D | G , I , Θ )
f ( G , Θ , I ) p
( Θ , Θ I , D | Z , Γ ) 
 ( since p ( Z , Θ , Γ , Θ I | D ) =
p ( Z ) p ( Γ ) p ( Θ , Θ I , D | Z , Γ ) p ( D ) )
=
E p ( Z , Θ , Γ , Θ I | D ) (cid:34)(cid:80) G , I tar p (
[Equation]
[Equation]
p
[Equation]
( rearranging ) =
E
[Equation]
(cid:34)(cid:80)
[Equation]
(cid:80) [Equation] ) (cid:35) ( law of total probability ) = [Equation]
(cid:34)(cid:80)
[Equation]
I )
[Equation]
G
[Equation]
I
[Equation]
D
[Equation]
Θ )
[Equation]
( by the generative model )
[Equation]
A.2 SCORES
We will derive the gradients of the unnormalized posterior since
[Equation]
and analogously for the gradients w.r
[Equation]
t. to other variables Z , Γ , Θ , Θ I .
By basic rules of probability theory and using the identity ∇ x log p ( x ) = ∇ x p ( x ) /p ( x )
[Equation]
we obtain
[Equation]
and analogously
[Equation]
B ALGORITHM DETAILS B.1 BACKGROUND: BAYESIAN INFERENCE OF BNS
Given i
[Equation]
i
[Equation]
d. observations D = { x (1) , . . .
, x ( N ) } , Bayesian inference over BNs means constructing a full posterior probability density over BNs that model the observations.
Following Friedman and Koller [2003]
[Equation]
given a prior distribution over DAGs p ( G ) and a prior over BN parameters p ( Θ | G ) , Bayes’ Theorem yields the joint and marginal posterior distributions
[Equation]
[Equation]
(23) where p ( D| G , Θ )
[Equation]
Θ ) the likelihood of the independent observations in D and p ( D| G ) =
(cid:82) p ( Θ | G ) p ( D| G , Θ )
d Θ is the marginal likelihood.
Thus, p ( D| G ) in Eq. 23 is only available in closed form for special conjugate cases.
The Bayesian formalism allows us to compute expectations of the form
[Equation]
[Equation]
for any function f of interest.
For instance, for Bayesian model averaging, we would use f ( G , Θ ) =
p ( x | G , Θ ) or f ( G ) =
p ( x | G ) , respectively [Madigan and Raftery, 1994, Madigan et al., 1995].
In active learning of CBNs, a commonly used f is the information gain in G from an intervention [Tong and Koller, 2001, Murphy, 2001, Cho et al., 2016, Agrawal et al., 2019].
Inferring the posterior is computationally challenging since there are O ( d !2( d 2 )) possible DAGs with d nodes [Robinson, 1973].
Hence, computing the normalization constant p ( D ) is also generally intractable.
B.2 STEIN VARIATIONAL GRADIENT DESCENT OVER CBNS AND INTERVENTIONS
Estimating Gradients of the Log-Posterior.
The score w.r.t.
auxiliary variables Z and Γ requires marginalization over the corresponding discrete structures G
and I tar ; in particular
[Equation] [Equation]
where p ( Θ , Θ I , D | G , I tar )
=
p ( Θ | G ) (cid:81) mk =1 p ( D k | G , Θ , I tar k ,
Θ I k ) p
( Θ I k | I tar k ) .
The expectations in Eq.
25 and Eq.
26 can be estimated by sampling the intervention masks I tar ∼ Bern ( σ α ( Γ )) and the adjacency matrices G ∼ Bern (cid:0) σ α ( UV ⊤ ) (cid:1) in a differentiable manner with the Gumbel-Softmax trick [Jang et al., 2016, Maddison et al., 2017].
For obtaining a differentiable observational log-likelihood, we mask individual log-likelihood summands based on G , as discussed in
Sec 4.3 of Lorch et al.
In a similar fashion, we use the I tar to switch between observational and interventional likelihoods per variable:
[Equation]
SVGD instantiation.
In Sec.
5 we have introduced a fully differentiable Bayesian model which translates the posterior inference over discrete graphs G and intervention targets I tar into an inference problem over the continuous latent variables
Z and Γ .
In this section, we discuss how to approach the final challenge of approximating the intractable posterior p ( Z , Θ , Γ , Θ I | D ) over our continuous latent variables.
To this end, we employ the particle variational inference approach SVGD
[Liu and Wang, 2016] which approximates the intractable distribution of interest by a finite set of particles.
SVGD uses the score of the distribution to guide the particles towards regions of high probability while using a kernel function k ( · , · ) between them.
The latter introduces repulsive forces which make the particles disperse well across the domain.
For a brief review of SVGD, please see Appx.
F. We employ a sum of RBF kernels (details in Appx.
B.4)
as the kernel function over Ψ :
=
( Z , Θ , Γ , Θ I ) .
We also considered a product of RBF kernels, though, found that the additive kernel composition performed better.
Starting with an initial set of L particles { Ψ ( l ) 0 } Ll =
= { ( Z ( l ) 0 , Γ ( l ) 0 , Θ ( l ) 0 , Θ ( l ) I , 0 ) }
Ll =1 , we perform T iterations of particle
SVGD updates.
Following Lorch et al.
[2021], we use annealing schedules α t → ∞ and β t → ∞ so that our continuous relaxations G α ( Z )
and I tar α ( Γ ) converge to DAGs and actual sets of intervention targets (details in Appx B.3).
We return { ( G ∞ ( Z ( l ) T ) , Θ ( l ) T , I tar ∞ ( Γ ( l ) T ) , Θ ( l ) I ,T ) } Ll =1 as the particle approximation of the posterior p ( G , Θ , I tar , Θ I | D ) with discrete
DAGs and interventions targets.
Figure 4 illustrates an example of the returned posterior particles for G
and I tar alongside the ground truth for the case of a linear Gaussian 5 node graph.
The proposed algorithm is summarized in Alg.
While
SVGD yields a set of particles of with equal weights, we weight each particle by its unnormalized posterior probability p ( G , Θ , Θ I , I tar , | D ) for performing approximate Bayesian model averaging.
We find that this improves the empirical performance.
Although we employ SVGD, our general framework from Sec.
5 can also be instantiated with other score-based sampling [e.g.
Welling and Teh, 2011, Chen et al., 2014] or variational inference methods [Blei et al., 2017].
Such alternatives may become favorable for systems of many variables, since the size of the adjacency matrix G grows quadratically in d and the respective kernel component between G ’s becomes useless in high-dimensions.
B.3 ANNEALING OF α AND β
The latent variables Z and Γ probabilistically model the causal graph G and intervention target masks I tar .
In a similar way, they can be viewed as as continuous relaxations of G and I tar respectively, where the α trades-off between smoothness
[Figure]
and accuracy of these relaxations.
If α → ∞ , the sigmoid σ α ( · ) converges to the unit step function so that the probability distributions p ( G | Z ) and p ( I tar | Γ ) become deterministic indicator functions, representing only single discrete graph
G ∞ ( Z ) ∈ { 0 , 1 } d × d and target mask I tar ∞ ( Γ ) ∈ { 0 , 1 } M × d .
To be able to make this simplification, as proposed in Lorch et al.
[2021], we anneal α t such that α t → ∞ and, at terminal iteration of the SVGD training loop, convert Z and Γ into their discrete counterparts G ∞ ( Z ) and I tar ∞ ( Γ ) .
Similarly, we set an annealing schedule β t → ∞ for inverse temperature parameter in the latent prior p β ( Z ) such that
and we only model DAGs as the training progresses.
Importantly, this annealing process simplifies the expectation in Eq.
4 to
[Equation] [Equation]
where I ∞ :
= ( I tar ∞ ( Γ ) , Θ I ) .
This holds since the inner expectations evaluate to a single point for α t , β t → ∞ and thus cancel out.
B.4 SVGD KERNEL
In our experiments, we employ an additive RBF kernel.
We also considered a product of RBF kernels, though, found that the additive kernel composition performed better.
This additive kernel is defined as
[Equation]
with lengthscales τ Z , τ γ , τ θ .
For brevity, we also write k ( Ψ , Ψ ′ ) for (28) where Ψ :=
( Z , Θ , Γ , Θ I ) .
For SVGD, the kernel introduces repulsive forces which make the particles disperse well across the domain (see F).
Importantly, the lengthscale hyperparameters provide the possibility to fine-tune the repulsion and hence calibrate our inference model.
We show this in
Sec.
E.3.
B.5 ALGORITHM OVERVIEW
We provide a pseudocode of our algorithm with its SVGD instantiation in Alg. 1.
C MARGINAL INFERENCE WITH THE BGE SCORE
In addition to joint posterior inference that includes the parameters Θ of the model, we also consider the marginal posterior p ( G | D ) .
We employ the commonly used Bayesian Gaussian Equivalent (BGe) marginal likelihood that scores Markov
Estimate ∇ log p ( Z
Θ , Γ
Θ I | D ) for
Ψ l t =
( Z ( l ) t , Θ ( l ) t , Γ ( l ) t , Θ ( l ) I ,t )
▷ see do 5:
Z ( l ) t +1 ← Z ( l ) t + η t ϕ Z t ( Ψ ( l ) t )
▷ where ϕ Z t ( · ) :
= 1 L
M (cid:88) l = L (cid:104) k (cid:16) Ψ ( l ) t , · (cid:17) ∇ Z ( l ) t log p ( Ψ ( l ) t | D ) + ∇ Z ( l ) t k (cid:16) Ψ ( l ) t , · (cid:17) (cid:105) 6: Θ ( l ) t +1 ← Θ ( l ) t + η t ϕ Θ t ( Ψ ( l ) t ) 7: Γ ( l ) t +1 ← Γ ( l ) t + η t ϕ Γ t ( Ψ ( l ) t ) 8: Θ ( l ) I ,t +1 ← Θ ( l ) I ,t + η t
ϕ
Θ
I t ( Ψ ( l ) t ) where ϕ Θ t , ϕ Γ t , ϕ Θ I t are analogous to ϕ Z t but use gradients ∇ Θ ( l ) t , ∇ Γ ( l ) t , ∇ Θ ( l ) I ,t + 9: return { ( G ∞ ( Z ( l ) T ) , Θ ( l ) T , I tar ∞ ( Γ ( l ) T ) , Θ ( l ) I ,T ) }
Ll =
▷ see Appx.
B. equivalent structures equally [Geiger and Heckerman, 1994, 2002].
This model factorises the marginal likelihood into components for each node given its parents.
Details on the computation of the BGe score are provided by Kuipers et al.
Other good explanations are given by Grzegorczyk [2010] and Kuipers and Moffa [2022].
Compared to the case of considering parameters Θ , the marginal likelihood does not yield factorization over different datasets.
That is,
[Equation]
Instead, we have to consider the fused data from all contexts by
[Equation]
for N = (cid:80) Mk =1 n k and variables c i ∈ { 0 , . . .
, M } that indicate from which context sample x i originates.
As before, c i =denotes the observational context.
Interventions.
When considering hard interventions, we cut off the connections of a variable to its parents.
This effectively means that the data at this variable only contributes a constant (wrt.
the current hypothesis graph) factor to the likelihood, and thus the scoring of a hypothesis graph should not be affected.
In other words, when computing the score for a certain variable j in the graph G , we drop all datapoints in X where j was the target of a hard intervention and then compute the
BGe score for the remaining datapoints ˆ X j .
In our implementation, we achieve this by masking out datapoints in X .
This enables us to still make use of the Gumbel-Softmax estimator for the Bernoulli intervention targets.
Priors.
Following the notation of Geiger and Heckerman [2002] and Kuipers et al.
[2014], we use the standard effective sample size hyperparameters α µ = 1 and α ω
= d + 2 .
Moreover, we use the diagonal form as the Wishart inverse scale matrix for the Normal-Wishart parameter prior underlying the BGe score (cf.
[Grzegorczyk, 2010]) with a prior mean of µ = [0 , . . .
, 0] T .
For interventions, we use the same approach but set the prior values µ to the real intervention mean from which the interventions where sampled.
We then compute a BGe score for the intervened datapoints given an empty graph in order to have comparable likelihood.
This helps when estimating the intervention targets that otherwise
Experiments.
We evaluate BaCaDI using the marginal BGe score to the baselines of JCI-PC and UT-IGSP.
Note that the algorithms of both JCI-PC and UT-IGSP are not affected by the different score; the difference is the scoring of the predicted
[Figure]
[Figure] [Figure]
I
[Figure]
NLL 0 .0 .0 .
AUPRC 0 .0 .
INTV-AUPRC
[Figure] [Figure]
graph structures, where we compute a likelihood score using the marginal BGe (replacing the MLE parameter estimate).
We do not compare to DCDI-G as it always performs joint inference with parameters.
Similar to Sec.
6, we focus on linear
Gaussian BNs with d = 20 nodes.
We report the results in Fig.
As done previously, we select the models with the lowest heldout I-NLL.
We again see how BaCaDI outperforms related methods in most metrics, particularly E -SID and I-NLL.
Only UT-IGSP performs competitively for AUPRC and INTV-AUPRC of SF graph structures, but fails to be on par with BaCaDI in other settings.
These results resonate with the other experimental evaluations, showing promising results for making causal structure predictions.
D EXPERIMENTAL SETUP
In the following, we describe the models, datasets and implementations that were used to perform experiments in detail.
D.1 MODEL DETAILS & BACADI
First, we discuss the models of graphs as well as the exact usage of BaCaDI 1 .
Graphs.
We consider DAGs that follow either Erd˝os-Rényi (ER)
[Gilbert, 1959] or scale-free (SF) [Barabási and Albert, 1999] distributions.
ER graphs follow a prior distribution
[Equation]
that describes that each edge exists independently w.p. q .
For SF graphs, we define a prior
[Equation]
analogous to their power law degree distribution p ( deg ) ∼ deg − 3 .
G i describes the i -th row of the adjacency matrix.
For all our experiments, we use priors that result in 2 d edges in expectation.
Building on Lorch et al.
[2021], BaCaDI can incorporate such a graph distribution in the prior p ( Z ) (see Sec.
4.2 in their paper).
Overview: Gaussian BNs.
As instantiations of BaCaDI ’s inference models that describe the local conditional distributions, we consider Bayesian networks with additive Gaussian noise.
This means that the variables x = ( x 1 , . . .
, x d ) follow a distribution
[Equation]
For the inference with BaCaDI , we assume a fixed observation noise σ 2 i =
σ 2 = 0 . 1 .
The function f can be modelled in different ways as follows.
Linear BNs.
Linear Gaussian BNs model the mean of a given variable as a linear function of its parents:
[Equation]
) where ◦ denotes the element-wise multiplication.
We use this parametrization as it allows for constant dimensionality of Θ ∈ R d × d and the elementwise multiplication only keeps the parents of each variable.
Moreover, this allows the use of the
Gumbel-Softmax estimator in Eq.
As a prior, we use a standard Gaussian p ( Θ i,j | G i,j = 1) =
N (0 , 1) .
Nonlinear BNs.
The local conditionals can be extended to nonlinear dependencies modelled by neural networks.
We consider feed-forward neural networks of the form
[Equation]
where
Θ =
(( Θ ( L ) , θ ( L ) b ) , . .
, ( Θ (1) , θ (1) b )) describe the parameters of the l -th layer for l ∈ { 1 , . . .
, L } .
The function σ is the element-wise nonlinear activation function and θ ( l ) b is the bias.
This then gives the distribution
[Equation]
Note that we thus have one neural network defined by Θ i for each of the local conditional distributions of variable i , that is, d networks in total.
For all experiments, we use one hidden layer with 5 units and the sigmoid activation function.
As a prior for the parameters, we analogously use a standard Gaussian with mean 0 and variance 1.
Interventions.
We model interventions using a Gaussian likelihood p
I k
[Equation]
( x
[Equation]
Θ
[Equation]
= N ( x
[Equation]
µ Ik,i , σ I ) with fixed variance σ I 2 = 0 .
We infer the means Θ I k
=
[ µ Ik, 1 , ..., µ Ik,d ] for which we set a wide Gaussian prior p ( Θ I k | I tar k )
= (cid:81)
i ∈
I tar k p ( µ Ik,i | I tar k,i = 1) =
(cid:81)
i ∈
I tar k N ( µ Ik,i | 0 , 10) .
This reflects an uninformative prior over a large effect range of the interventions.
Initialization.
For the linear Gaussians, the parameters are initialized closed to zero via Θ init ∼ N (0 , σ init I ) with σ init = 0 .
The closeness to zero is important to avoid inducing a bias at the start of the SVGD inference process.
Similarly, we sample Θ I , init ∼ N (0 , σ I , init I ) with σ 2 I , init = 0 . 1 .
For nonlinear Gaussians, we use the Glorot (sometimes called Xavier) normal to initialize the weights of the neural networks.
[Glorot and Bengio, 2010].
D.2 DATA GENERATION FOR THE SYNTHETIC CAUSAL INFERENCE TASKS
The data generation is done as follows: we first sample a random graph (either ER or SF) and then sample random parameters.
We collect data samples by iterating through the topological ordering of the graph and sample a variable given its local parents.
Since we consider additive Gaussian models, each variable is described by a distribution of the form x i | Θ , G ∼ N ( f ( x pa G ( i ) , Θ ) , σ i ) where f is either a linear function or a nonlinear feedforward neural network (cf.
Sec.
D.1).
The noise variables σ i are sampled per variable and fixed once from σ 2 i ∼ U [0 . 05 , 0 . 15] .
If the noise variables had the same variance across all variables in the graph, this would render identification possible [Peters and Bühlmann, 2014].
Parameters.
The parameters and generative models are initialized as follows:
• Linear BNs: We sample the parameters Θ uniformly and independently from U ([ − 2 , − 0 .
5] ∪ [0 . 5 , 2]) in order to bound the weights sufficiently away from zero.
• Nonlinear BNs: The NNs for each local conditional are the same model as used for BaCaDI as well as DCDI-G, that is, a fully connected NN with single hidden layer of size of 5 with biases.
The nonlinear activation function is sigmoid.
All weights and biases are drawn randomly and independently from a Gaussian N (0 , 1) .
Interventions.
As described in the main text, we perform hard interventions on every node for the 20 -node graphs.
We create random values by first sampling ˆ µ k,i ∼ N (0 , 2) and then setting µ Ik,i = sign (ˆ µ k,i ) · 5 + ˆ µ k,i .
This ensures that the interventions performed are bounded away from zero and out-of-distribution.
If a variable x i is the target of the intervention in context k , we then have p
I k
i ( x i | Θ I k ) =
p I
k i ( x i | µ Ik,i ) =
N ( x i ; µ Ik,i , σ I ) , where σ 2 I = 0 .
D.3 DATA GENERATION WITH THE SERGIO GENE-EXPRESSION SIMULATOR
SERGIO [Dibaeinia and Sinha, 2020] is a single-cell expression simulator guided by Gene Regulatory Networks (GRN).
The software tool can generate realistic single-cell transcriptomics datasets based on a user-defined graph input that describes the gene-regulatory network.
SERGIO uses a stochastic differential equation (SDE), the so-called chemical Langevin equation (CLE), to simulate a gene’s expression dynamics as a function of the changing levels of its regulators.
Of interest to us, the simulations resemble the data collected by modern high-throughput, single-cell RNA sequencing (scRNA-seq) technologies.
For more details, please consider the original paper by Dibaeinia and Sinha [2020].
We give a brief overview of how we simulate scRNA-seq data with SERGIO.
We use an implementation that is adapted from the open-surce code available under https://github.com/PayamDiba/SERGIO which is published under a GPL-3.0 license.
Simulation.
SERGIO generates synthetic scRNA-seq data D for a given causal graph with d genes in two stages: first, it collects clean gene expression snapshots; then they are altered by technical noise.
The N observations in D correspond to
N cell samples, i.e. one row in D describes the joint expression of the d genes in a single cell.
In the first stage, SERGIO samples clean gene expressions through snapshots at random timesteps from the steady state of a dynamical system.
In this regulatory process, the genes are expressed at rates influenced by other genes using the chemical
Langevin equation (CLE).
The source nodes in the causal graph G are denoted master regulators (MRs), whose expressions evolve at constant product and decay rates.
The expressions of all downstream genes evolve non-linearly under production rates caused by the expression of their causal parents in G .
In addition, cell types are defined by specifications of the MR production rates, which significantly influence the evolution of the system.
Thus, the data contains variation due to biological noise within collections of cells of the same type as well as across cell types.
Ultimately, we generate single-cell samples collected from ten cell types [Dibaeinia and Sinha, 2020].
In the second stage, the clean gene expressions sampled previously are corrupted with technical noise that resembles the noise phenomena found in real scRNA-seq data.
For simplicity, we do not make use of the technical noise for our experiments.
Interventions.
We consider knockout interventions that clip a specific gene to zero.
To that end, we extend SERGIO by forcing the production rate of knocked-out genes to be zero during simulation.
As described in the main text, we do this for
M = 10 gene targets and arrive at 11 contexts (including the observational setting).
Parameters.
Given a causal graph G , the data generation to simulate c cell types of d genes is governed by the following parameters:
• k ∈ R d × d : interaction strengths (only used if edge i → j exists in G ) – We sample each k
i,j ∼ U [1 , 5] .
• b ∈ R
d
×
c
+ : MR production rates (only used if gene j is a source node in G ) –
We sample each b i,j ∼ U [1 , 3] .
• γ ∈ R d × d
+ : Hill function coefficients controlling nonlinearity of interactions – We fix a single γ const = 2 . 0 and set γ i,j = γ const .
• λ ∈ R d : decay rates per gene – We fix each λ i = 0 .
• ζ ∈ R d + : scale of stochastic process noise in chemical Langevin equation simulating the dynamical system.
– We fix ζ = 1 . 0 .
We standardize the collected data for all methods
D.4 BASELINES & IMPLEMENTATION
In this section, we provide additional details on the baseline methods and their implementations used in our experiments.
Our source-code and instructions on how to reproduce results are available under https://www.dropbox.com/sh/ vtj4zp9h8sr9zq/AACkJ5naAxhXnpIByNEezI4Ia?dl=0 .
For all methods, we use implementations adapted from the source code of DCDI available under https://github.c om/slachapelle/dcdi (where the authors also included their baselines’ code).
The implementation of JCI-PC is a modified version of the R package using code from the JCI repository https://github.com/caus-am/jci/tree/ master/jci .
UT-IGSP has an implementation available from https://github.com/uhlerlab/causaldag .
DAG Bootstrap.
To compare all methods for Bayesian model averaging, we employ the nonparametric DAG bootstrap
[Friedman et al., 2013].
It performs model averaging by bootstrapping the observations D to yield a collection of synthetic data sets, each of which is used to learn a single graph.
We sample with replacement for each dataset D k individually.
The collection of unique single graphs approximates the posterior by weighting each graph by its unnormalized posterior probability.
All of the methods are evaluated with 20 bootstrap samples, i.e. the same number of samples as particles used for the SVGD instantiation of BaCaDI.
The only exception is DCDI-G, where we only use 5 bootstrap samples due to longer runtimes.
JCI-PC.
The Joint Causal Inference (JCI) framework [Mooij et al., 2016] introduces a general formulation to extend the initial causal graph by auxiliary nodes that describe the different contexts, effectively performing causal discovery over a graph of size d + M .
This can be instiantiated with different standard algorithms.
We use the PC algorithm [Spirtes et al., 2000] that relies on conditional independence tests (CI) to discover the Markov Equivalence Class (MEC), i.e. the skeleton of a graph with v-structures as well as possible identifiable edge directions.
We use the Gaussian CI tests that are best-suited for the Gaussian BNs that we consider with the threshold parameter α JCI .
For computing the graph metrics, we compute the E -SID between the CPDAGs of the GT graph and predictions of JCI-PC.
Additionally, we favor JCI-PC when computing AUPRC scores.
See Sec.
D.5 for more details.
To arrive at a DAG that we can use to estimate MLE parameters and compute log-likelihood metrics, we generate a consistent expansion of the CPDAG as defined by Chickering [2002].
That is, we generate a DAG s.t.
the CPDAG has the same skeleton and v-structures and every directed edge in the CPDAG has the same direction in the DAG.
To that end, we perform a random consistent expansion as described by Dor and Tarsi [1992].
UT-IGSP.
The interventional greedy sparsest permutation (IGSP) method [Wang et al., 2017, Yang et al., 2018] proposes an algorithm that learns causal structures via local scores based on CI relations and permutation search.
The work is extended by Squires et al.
[2020] to the case of unknown targets (UT-IGSP).
Analogous to JCI-PC, we make use of Gaussian CI and invariance tests with parameters α UT-IGSP and α UT-IGSPinv .
Since we consider a low sample setting, it is possible that the algorithm as provided by their open source implementation does not succeed when computing the correlation matrix.
This is because it computes a correlation matrix which becomes singular in case the number of samples is smaller than the number of variables considered, thus rendering an inversion of this matrix impossible.
In case this happens, we retry the inference with a halved α UT-IGSP confidence threshold.
The maximum number of restarts is set to 10 .
Should this maximum be reached, the bootstrap sample is simply dropped.
Similar to JCI-PC, for computing the graph metrics, the E -SID is computed as the midpoint of the lower and upper bound between the CPDAGs of the GT graph and predictions of UT-IGSP.
Additionally, we favor UT-IGSP when computing
AUPRC scores.
See Sec.
D.5 for more details.
Moreover, we obtain a DAG by the random consistent expansion of the CPDAG (as described above for JCI-PC) to compute the log-likelihood metrics.
DCDI-G.
The work of Brouillard et al.
[2020] introduced a model for Differentiable Causal Discovery from Interventional
Data (DCDI) that performs causal structure learning via the augmented Lagrangian method.
They do so by formulating a continuous-constrained optimization problem that relies on stochastic gradient descent and neural networks to fit the local conditionals.
For a fair comparison, we employ the exact same model as used for the nonlinear Gaussian BNs used in BaCaDI, that is, a feedforward neural network with one hidden layer of size 5 and Gaussian additive noise.
This model is also called
DCDI-G in [Brouillard et al., 2020], however, a smaller model than the default configuration that was used in their paper ( hidden layers with 16 hidden units); with that, we performed initial experiments, but saw strong overfitting and thus reduced model-capacity.
As elementwise activation function, we use the leaky ReLU with negative slope of 0 .
25 as suggested by their work.
Additionally, to have comparable log-likelihood metrics, we fix the noise variables to σ 2 = 0 .
1 just as done for
BaCaDI.
These noise variables could generally be learned by both models.
Since DCDI-G takes longer to compute, we restrict the number of bootstrap samples to 5.
D.5 METRICS
Here, we describe the evaluation metrics more in detail.
Our reported metrics focus on three essential aspects of our inference problem: causal graph prediction, intervention detection, and learning the local conditionals of individual variables/nodes.
We describe the metrics in the following.
• SID: The Structural Interventional Distance (SID) [Peters and Bühlmann, 2015] quantifies the closeness between two DAGs in terms of how well their interventional adjustment sets coincide.
Since we perform posterior inference and arrive at a distribution over graphs, we consider the expected SID:
[Equation]
We use the implementation provided by the Causal Discovery Toolbox [Kalainathan and Goudet, 2019].
Since UT-IGSP and JCI-PC only return a CPDAGs of the Interventional Markov Equivalence Class (I-MEC)
we calculate its lower and upper bound of SIDs in the I-MEC, and report their midpoint as the E -SID.
Note that the DAG bootstrap variants for all baselines as well as BaCaDI use the weighted mixture rather than the empirical distribution of samples; the weight is based on the achieved unnormalized log-likelihood on the bootstrap sample.
• SHD : Another commonly used metric is the structural hamming distance (SHD) that reflects the graph edit distance to the ground truth.
However, it often does not properly reflect the closeness of two DAGs in terms of their causal interpretation.
For instance, the trivial prediction of the empty graph achieves competitive SHD scores for the sparse graphs we consider in this paper.
In our main text, we thus focus on the E -SID as well as other metrics that together better assess the quality of
causal graph
predictions.
For completeness, we report the E -SHD results in Appx.
E.6.
• Threshold metrics : Treating the edge prediction as a classification task, we compute the area under the precision recall curve (AUPRC) for pairwise edge prediction based on the posterior marginals p ( g ij = 1 | D ) .
This marginal is simply defined as averaging the presence of edges in the posterior samples: p ( g ij = 1 | D ) =
E p ( G | D ) 1 [ g ij = 1] .
Note that this is a more suitable metric than, e.g., the AUROC, as sparse graphs translate to highly imbalanced edge classification tasks.
For both baselines JCI-PC and UT-IGSP, which both possibly return a CPDAG with edges that are undecided, we favor them when computing the AUPRC metric.
We orient a predicted undirected edge correctly when a ground truth edge exists and only count a falsely predicted undirected edge as a single mistake.
• Interventional AUPRC : Similarly, we report the interventional AUPRC (INTV-AUPRC) for the classification of targets.
This again captures how well an algorithm predicts which variables have been intervened on.
Since we are performing sparse interventions, this better captures a model’s performance for a highly imbalanced classification task.
• I-NLL: We compute the average negative interventional log-likelihood (I-LL) on M test = 10 heldout interventional datasets D test =
{D test 1 , ..., D test 10 } , where different interventions are performed compared to the training datasets.
Each interventional test dataset comprises 100 samples, and has known intervention targets I tartest ,k and effect distributions p ( x i | Θ I test ,k )
The I-NLL is computed via
[Equation]
Since UT-IGSP and JCI-PC are not equipped with local conditional distributions, we use the linear Gaussian maximumlikelihood parameters (MLE) that can be computed in closed-form [Hauser and Bühlmann, 2014] to compute the heldout
I-NLL.
Note that for the nonlinear datasets, this creates a model mismatch since the closed form is only available for the linear Gaussian BNs.
[Figure]
AUPRC
[Figure]
INTV
[Figure]
AUPRC
[Figure]
DCDI
[Figure]
BaCaDI (NonlinG)
D.6 HYPERPARAMETER SEARCH
To ensure a fair comparison, we perform a hyperparameter search for all baselines.
The search ranges can be found in Table 1.
Throughout all experiments, we use at least 20 hyperparameter samples and aggregate results over 30 different random seeds and graphs.
For the results of BaCaDI in the main text, we fix the hyperparameters α = 0 .
{+01 , β = 
+}
and λ = 1 , which are the linear scale for the temperature parameter of the sigmoids, the linear scale for the sparsity regularizer and the sparse intervention regularizer, respectively.
Method
Hyperparameter
Comment
Search range
BaCaDI τ Z
Kernel lengthscale log 10 U [ − 1 , 1 . 7]
τ γ
Kernel lengthscale log 10 U [ − 1 , 1 . 7] τ θ
Kernel lengthscale log 10 U [1 . 2 , 5]
JCI-PC α JCI
CI tests log 10 U [ − 5 , − 1]
UT-IGSP α UT-IGSP
CI tests log 10 U [ − 5 , − 1] α UT-IGSPinv
Invariance tests log 10 U [ − 5 , − 1]
DCDI-G batch sizeU ( { 16 , 32 , 64 } )
λ R
Sparsity coefficient for interventions log 10 U [ − 8 , − 1] λ
Sparsity
[Table]
E ADDITIONAL EXPERIMENTS E.1 SYNTHETIC NONLINEAR DATASETS
In addition to the results for synthetic linear Gaussian BNs in Sec. 6, we perform experiments for synthetic nonlinear
Gaussian BNs.
As described in Sec.
D.1 and Sec.
D.2, the local conditionals are modelled by feedforward neural networks and hence the same model as used DCDI-G and the nonlinear BaCaDI.
The corresponding results can be found in Fig.
Analogous to the previous evaluations, BaCaDI gives predictions that are the closest to the ground-truth CBN.
It strongly outperforms the baselines in terms of the E -SID and AUPRC while achieving low I-NLL and high INTV-AUPRC.
Among the baselines, UT-IGSP is the best at detecting interventions, largely on par with BaCaDI, but it fails behind in the other metrics.
E.2 LARGER DATASETS
We report additional results when doubling the dataset size, i.e. we collect n 0 = 200 observational samples and n k =samples per interventional context.
The results for synthetic linear and nonlinear Gaussian BNs can be found in Fig.
7 and
Fig. 8, respectively.
[Figure] [Figure]
E.3 CALIBRATION
Having compared BaCaDI with the baselines across the metrics for causal structure learning in the previous sections and the main text, we now provide additional experiments to show how BaCaDI incorporates epistemic uncertainty .
To quantify the reliability of our uncertainty estimates, we utilize the concept of calibration [Gneiting et al., 2007, Kuleshov et al., 2018] to quantify the reliability of our uncertainty estimates.
We show the results in Appx.
E.3.
Notably, BaCaDI is the only method that takes into account the epistemic uncertainty when dealing with limited data and shows that its probabilistic predictions are reliable.
Formally, we consider a probabilistic predictor to be well calibrated if in expectation, its α -% confidence intervals cover α -% of the true targets.
We visualize this in Fig.
9 for the task of predicting edges in CBNs.
The plot compares the confidence intervals for varying thresholds of α with the empirical coverage, that is, the average percentage of edges present in the ground-truth graphs.
Perfectly calibrated predictions exhibit a proportional one-to-one relationship between confidence level and empirical coverage, indicated by the black straight line in Fig.
All three baseline methods are severely underconfident in their predictions.
BaCaDI, in contrast, produces edge predictions that are close to the ideal calibration line.
This indicates that
[Figure]
BaCaDI properly takes into account the epistemic uncertainty when dealing with limited data and shows that its probabilistic predictions are reliable.
Additionally, BaCaDI offers the possibility to fine-tune the calibration of the predictor through the SVGD kernel parameters.
Calibration for synthetic datasets.
We show the calibration for all methods with the synthetic datasets used in the main text.
The calibration plots can be found in Fig.
Notably, BaCaDI is the only method that does not give grossly overconfident predictions.
Instead, it is conservative in predicting edges for the linear BNs, and gets closest to the ideal calibration line for nonlinear BNs.
Ablation study.
In addition to well capturing the causal mechanisms when performing posterior inference, we now demonstrate the flexibility of calibrating the inference of BaCaDI.
Through the use of SVGD and a kernel that can be designed specifically for the application, BaCaDI can be fine-tuned to a desired calibration.
We show different inference results in Fig.
11 for linear as well as nonlinear datasets.
Depending on the lengthscale hyperparameters for the kernel in 28, we see that the predictions of BaCaDI can be easily adapted to varying levels.
It thus properly takes into account the epistemic uncertainty when dealing with limited data.
E.4 OBSERVATIONAL VS. INTERVENTIONAL DATA
As an interesting use case, we now evaluate how interventional data helps predicting the causal structure with BaCaDI.
In general, intervening on variables in the system and observing the outcome provides information that helps discovering the causal mechanisms; in the perfect setting, when the interventions are fully known, they increase identifiability (reducing the interventional Markov equivalence class)
[Hauser and Bühlmann, 2012].
However, the information gain can be limited when the interventions are unknown [Squires et al., 2020].
We investigate how BaCaDI can leverage such unknown interventions compared to using just observational data.
To that end, we evaluate 3 different settings: i) when only observational data is available, ii) interventional data with full knowledge of the interventions, and iii) unknown interventions.
When the interventions are known, the posterior inference of BaCaDI can be reduced to Eq. 1.
We consider nonlinear Gaussian BNs for d = 20 node graphs.
All methods receive the exact same number of samples for inference.
That is, we collect n
= 300 samples from the ground-truth BN without interventions for the observational dataset; the interventional datasets, as before, have n 0 = 100 observations and n k = 10 samples per interventional context.
We report the results in Fig.
We can see how BaCaDI achieves much better results when interventional data is available.
The case of known interventions serves as a natural baseline, where BaCaDI achieves the best performance and gets closest in recovering the true causal structure.
This is most clearly shown by the E -SID and AUPRC scores. Notably, BaCaDI is able to get close this baseline even when interventions are unknown and outperforms the setting where just observational data is available. This shows the
.conﬁd.
level α 0 .0 .1 .e m p
i r i c .
c o v e r ag
e
[Figure]
[Figure]
.conﬁd.
level α 0 .0 .1 .e
m p
i r i c .
c o v e r ag
e
[Figure]
[Figure] [Figure] [Figure] conﬁd.
[Figure]
α 0 .0 .1 .e
m p
i r i c .
c o v
[Figure]
a [Figure] e
[Figure]
.conﬁd.
level α 0 .0 .1 .e
m p
i r i c .
[Figure]
v
[Figure]
a g e
[Figure]
benefit of performing and collecting interventions even when some of the effects may be unknown and is a promising result for future work.
E. 50 NODE GRAPHS
As additional experiments, we perform posterior inference for larger graphs with 50 nodes in total.
Analogous to the previous evaluations, we consider synthetic linear Gaussian BNs with hard interventions on every node.
We use a larger dataset of n 0 = 200 observational samples and n k = 20 samples per interventional context.
We show the results in Fig.
We see how BaCaDI scales to larger graphs, performing competitively across all metrics.
In particular, it outperforms the baselines by a large margin in terms of the E -SID and AUPRC for ER graphs.
E.6 STRUCTURAL HAMMING DISTANCE
For completeness, we include the Expected Structural Hamming Distance ( E -SHD) metric for all results discussed in the main text for d = 20 node graphs in Table 2.
Analogous to the E -SID, the E -SHD is defined as E -SHD ( p, G gt ) :
=
.conﬁd.
level α 0 .0 .0 .0 .1 .e
m p
i r i c .
c o v e r ag
e
[Figure]
.conﬁd.
level α 0 .0 .0 .0 .1 .e
m p
i r i c .
c o v e r ag
e
[Figure]
[Figure]
G p ( G |D ) · SHD ( G , G gt ) .
The SHD simply reflects the graph edit distance where wrongly inverted edges are counted as only one error.
However, while it captures closeness to the ground truth graph, trivial predictions like the empty graph achieve competitive results for sparse graphs.
For example, in the setting of d = 20 nodes and 2 d edges in expectation, the empty graph will achieve a E -SHD of 40 in expectation.
Similar conclusions hold if only a handful of edges are predicted.
In the main text [Table] we thus resort to the E -SID as well as additional metrics that together establish the quality of the causal predictions. [Table]
[Figure] AUPRC [Figure] [Figure]
AUPRC
[Figure]
[Figure]
INTV-AUPRC
[Figure]
E.7 RUNTIMES
We report the average runtimes (in minutes) for all methods and main results for d = 20 node graphs in Table 3.
E.8 COMPUTATIONAL RESOURCES
All experiments reported in this paper were performed in bulk and in parallel on 2-core CPU nodes of an internal cluster.
No
GPUs were used.
We estimate the total compute to roughly amount to 40000 hours of CPU time.
[Table]
Nonlinear Gaussian BNs
[Table]
We here give the most important points of the Stein Variational Gradient Descent (SVGD) introduced by Liu and Wang [2016].
Fundamentally, the work of Liu and Wang [2016] connects the mathematical notions of probability discrepancies with a variational inference method, which closely resembles the gradient descent algorithm.
For a comprehensive overview, we refer to the original paper.
Stein’s identity and discrepancy.
Formally, let p ( x ) be a continuously differentiable (i.e. smooth) density supported on
X ⊆ R d .
For a smooth vector function ϕ ( x ) , the Stein’s identity states that for sufficiently regular ϕ , we have
[Equation]
37) where A p ϕ ( x ) =
[Equation]
∇ x log p
[Equation]
+ ∇ x
[Equation]
38) where A p is the Stein operator acting on the function ϕ .
This equation can be subsequently used as a discrepancy measure: when considering the expectation over a smooth density q different than p , we obtain the so called Stein discrepancy by considering the "maximum violation of Stein’s identity".
That is, we have
[Equation]
39) for a choice of a set of functions F , for instance the reproducing kernel Hilbert space (RKHS) denoted by H d .
Variational Inference using smooth transforms.
The goal of variational inference is to approximate a target distribution p using a simpler distribution q ∗ which minimizes the KL-divergence
[Equation]
In order to minimize the KL-divergence, the authors in [Liu and Wang, 2016] consider incremental transforms formed by a small perturbation to the identity map that make up the set of distributions Q .
That is, the transform T : X → X is defined as
[Equation]
41) where ϕ is a smooth function that characterizes the perturbation direction.
As a key result, Liu and Wang [2016] connect these transforms to the Stein operator and the derivative of the KL divergence.
The authors show that if the function ϕ lies in the ball of the vector valued RKHS H d , the direction of the steepest descent on the KL divergence between a fixed q and the target p is given by
[Equation]
where one can easily identify the form of the Stein operator.
What is more, the value of the obtained gradient equals the (negative) kernelized Stein discrepancy measure − S ( q, p ) .
General Algorithm.
This mathematical result suggests an iterative method that transforms an initial reference distribution q 0 to the target distribution p .
Starting with a finite set of random particles { x ( m ) } Mm =1 , for some iteration t each particle is updated deterministically according to
[Equation]
where
[Equation]
These steps iteratively decrease the KL divergence between q t and p , ultimately converging.
Importantly, the advantage of SVGD is that it only depends on the gradient of the kernel k ( x, · ) that can be defined for the application, as well as the score function ∇ x log p ( x ) which can be computed without knowing the (intractable) normalization constant of p .
On an intuitive level, the first part of the perturbation direction expression drives the particles to high density regions close to the mode of the target distribution p , whereas the term ∇ x k ( x , · ) acts as a repulsive force that prevents the particles from collapsing together.