L EARNING TO I NDUCE C AUSAL S TRUCTURE
Anonymous authors
Paper under double-blind review
A BSTRACT
One of the fundamental challenges in causal induction is to infer the underlying graph structure given observational and/or interventional data.
Most existing causal induction algorithms operate by generating candidate graphs and evaluating them using either score-based methods (including continuous optimization) or independence tests.
In our work, we instead treat the inference process as a black box and design a neural network architecture that learns the mapping from both observational and interventional data to graph structures via supervised training on synthetic graphs.
The learned model generalizes to new synthetic graphs, is robust to train-test distribution shifts, and achieves state-of-the-art performance on naturalistic graphs for low sample complexity.
I NTRODUCTION
The problem of discovering the causal relationships that govern a system through observing its behavior, either passively ( observational data ) or by manipulating some of its variables ( interventional data ), lies at the core of many important challenges spanning scientific disciplines, including medicine, biology, and economics.
By using the graphical formalism of causal Bayesian networks (CBNs) (Koller & Friedman, 2009; Pearl, 2009), this problem can be framed as inducing the graph structure that best represents the relationships.
Most approaches to causal structure induction are based on an unsupervised learning paradigm in which the structure is directly inferred from the system observations, either by ranking different structures according to some metrics (score-based approaches) or by determining the presence of an edge between pairs of variables using conditional independence tests (constraint-based approaches)
(Drton & Maathuis, 2017; Glymour et al., 2019;
Heinze-Deml et al., 2018a;b) (see Fig. 1(a)).
The unsupervised paradigm poses however some challenges: score-based approaches are burdened with the high computational cost of having to explicitly consider all possible structures and with the difficulty of devising metrics that can balance goodness of fit with constraints for differentiating causal from purely statistical relationships (e.g. sparsity of the structure or simplicity of the generation mechanism); constraint-based methods are sensitive to failure of independence tests and require faithfulness, a property that does not hold in many real-world scenarios (Koski & Noble, 2012; Mabrouk et al., 2014).
Recently, supervised learning methods based on observational data have been introduced as an alternative to unsupervised approaches (Lopez-Paz et al., 2015a;b; Li et al., 2020).
In this work, we extend the supervised learning paradigm to also use interventional data, enabling greater flexibility.
We propose a model that is first trained on synthetic data generated using different CBNs to learn a mapping from data to graph structures and then used to induce the structures underlying datasets of interest (see Fig. 1(b)).
The model is a novel variant of a transformer neural network that receives as input a dataset consisting of observational and interventional samples corresponding to the same
CBN and outputs a prediction of the CBN graph structure.
The mapping from the dataset to the underlying structure is achieved through an attention mechanism which alternates between attending to different variables in the graph and to different samples from a variable.
The output is produced by a decoder mechanism that operates as an autoregressive generative model on the inferred structure.
Our approach can be viewed as a form of meta-learning, whereby the relationship between datasets and causal structures underlying them are learned rather than built-in. [Figure] A [Figure]
C
[Figure]
B
[Figure] Intvn [Figure] A
B
C
[Figure]
Our supervised approach to causal structure induction (via attention; CSIvA): A model is presented with data and structures as training pairs and learns a mapping between them.
A requirement of a supervised approach would seem to be that the distributions of the training and test data match or highly overlap.
Obtaining real-world training data with a known causal structure that matches test data from multiple domains is extremely challenging.
We show that meta-learning enables the model to generalize well to data from naturalistic CBNs even if trained on synthetic data with relatively few assumptions, where naturalistic CBNs are graphs that correspond to causal relationships that exist in nature, such as graphs from the bnlearn repository ( www.bnlearn.com/ bnrepository ).
We show that our model can learn a mapping from datasets to structures and achieves state-of-the-art performance on classic benchmarks such as the Sachs, Asia and Child datasets (Lauritzen & Spiegelhalter, 1988; Sachs et al., 2005; Spiegelhalter & Cowell, 1992), despite never directly being trained on such data.
Our contributions can be summarized as follows:
• We tackle causal structure induction with a supervised approach (CSIvA) that maps datasets composed of both observational and interventional samples to structures.
• We introduce a variant of a transformer architecture whose attention mechanism is structured to discover relationships among variables across samples.
• We show that CSIvA generalizes to novel structures, whether or not training and test distributions match.
Most importantly, training on synthetic data transfers effectively to naturalistic CBNs.
• We show that CSIvA significantly outperforms state-of-the-art causal discovery methods such as DCDI (Brouillard et al., 2020) and ENCO (Lippe et al., 2021) both on various types of synthetic CBNs, as well as on naturalistic CBNs.
B ACKGROUND
In this section we give some background on causal Bayesian networks and transformer neural networks, which form the main ingredients of our approach (see Appendix A.1 for more details).
Causal Bayesian networks (CBNs).
A causal Bayesian network (Koller & Friedman, 2009;
Pearl, 2009) is a pair M = ⟨G , p ⟩ , where G is a directed acyclic graph (DAG) whose nodes
X 1 , . . .
X N represent random variables and edges express casual dependencies among them, and where p is a joint distribution over all nodes such that p ( X 1 , . .
X N )
=
(cid:81)
Nn =1 p ( X n | pa ( X n )) , where pa ( X n ) are the parents of X n , i.e. the nodes with an edge onto X n (direct causes of
X n ).
An input to the transformer neural network is formed by a dataset D
= { x s } Ss =1 , where x s :=
( x s 1 , . . .
, x sN )
T is either an observational data sample or an interventional data sample obtained by performing an intervention on a randomly selected node in G .
Observational data samples are samples from p ( X 1 , . . .
, X N ) .
Except where otherwise noted, for all experimental settings, we considered hard interventions on a node X n that consist in replacing the conditional probability distribution (CPD) p ( X n | pa ( X n )) with a delta function δ X n ′ =
x
forcing X n ′ to take on value x .
Additional experiments were also performed using soft interventions , which consisted of replacing p ( X n | pa ( X n )) with a different CPD p ′ ( X n | pa ( X n )) .
An interventional data sample is a sample from δ X n ′ =
x (cid:81)
Nn =1 ,n ̸ = n ′ p
( X n | pa ( X n )) in the first case, and a sample from p ′
( X n | pa ( X n )) (cid:81)
Nn =1 ,n ̸ = n ′ p ( X n | pa ( X n )) in the second case.
The structure of G can be repre- sented by an adjacency matrix A , defined by setting the ( k, l ) entry, A k,l , to 1 if there is an edge from
X l to X k and to 0 otherwise.
Therefore, the n -th row of A , denoted by A n, : , indicates the parents of
X n while the n -th column, denoted by A : ,n , indicates the children of X n .
Transformer neural network.
A transformer (Devlin et al., 2018; Vaswani et al., 2017) is a neural network equipped with layers of self-attention that make them suited to modeling structured data.
In traditional applications, attention is used to account for the sequentially ordered nature of the data, e.g. modeling a sentence as a stream of words.
In our case, each input of the transformer is a dataset of observational or interventional samples corresponding to the same CBN.
Attention is thus used to account for the structure induced by the CBN graph structure and by having different samples from the same node.
Transformers are permutation invariant with respect to the positions of the input elements, ensuring that the graph structure prediction does not depend on node and sample position.
C AUSAL S TRUCTURE I NDUCTION VIA A TTENTION (CSI V A)
Our approach is to treat causal structure induction as a supervised learning problem, by training a neural network to learn to map observational and interventional data to the graph structure of the underlying CBN.
Obtaining diverse, real-world data with known causal relationships in amounts sufficient for supervised training is not feasible.
The key contribution of this work is to introduce a method that uses synthetic data generated from
CBNs with different graph structures and associated CPDs that is robust to shifts between the training and test data distributions.
3. S UPERVISED APPROACH
We learn a distribution of graphs conditioned on observational and interventional data as follows.
We generate training data from a joint distribution t ( G , D ) , between a graph G and a dataset D comprising of S observational and interventional samples from a CBN associated to G as follows.
We first sample a set of graphs {G i } Ii =1 with nodes X i 1 , . . .
, X iN from a common distribution t ( G ) as described in Section 4 (to simplify notation, in the remainder of the paper we omit the graph index i when referring to nodes), and then associate random CPDs to the graphs as described in Section 4.
This results in a set of CBNs {M i } Ii =1 .
For each CBN M i , we then create a dataset D i
= { x s } Ss =1 , where each element x s :=
( x s 1 , . . .
, x sN )
T is either an observational data sample or an interventional data sample obtained by performing an intervention on a randomly selected node in G i .
Our model defines a distribution ˆ t ( G | D ; Θ) over graphs conditioned on observational and interventional data and parametrized by Θ .
Specifically, ˆ t ( A | D ; Θ) has the following auto-regressive form: ˆ t ( A | D ; Θ) =
(cid:81) N 2 l =1 σ ( A l ; ˆ A l = f Θ ( A 1 ,..., ( l − 1) , D )) , where σ ( · ; ρ ) is the Bernoulli distribution with parameter ρ , which is a function f Θ built from an encoder-decoder architecture explained in Section 3.2 taking as input previous elements of the adjacency matrix A (represented here as an array of Nelements) and D .
It is trained via maximum likelihood estimation (MLE), i.e Θ ∗ = argmin Θ L (Θ) , where L (Θ) = − E
( G , D ) ∼ t
[ln ˆ t ( G | D ; Θ)] , which corresponds to the usual cross-entropy (CE) loss for the Bernoulli distribution.
Training is achieved using a stochastic gradient descent (SGD) approach in which each gradient update is performed using a pair ( D i , A i ) .
The data-sampling distribution t ( G , D ) and the MLE objective uniquely determine the target distribution learned by the model.
In the infinite capacity case, ˆ t ( · | D ; Θ ∗ ) =
t ( · | D ) .
To see this, it suffices to note that the MLE objective L (Θ) can be written as L (Θ) =
E D∼ t [ KL (ˆ t ( · | D ; Θ); t ( · | D ))]
+ c , where KL is the Kullback-Leibler divergence and c is a constant.
In the finite-capacity case, the distribution defined by the model ˆ t ( · | D ; Θ ∗ ) is only an approximation of t ( · | D ) .
3. M ODEL ARCHITECTURE
The function f Θ defining the model’s probabilities is built using two transformer networks.
It is formed by an encoder transformer and by a decoder transformer (which we refer to as “encoder” and “decoder” for short).
At training time, the encoder receives as input dataset D i and outputs a representation that summarizes the relationship between nodes in G i .
The decoder then recursively outputs predictions of the elements of the adjacency matrix
A
i using as input the elements previously predicted and the encoder output.
This is shown in Fig.
2 (where we omitted index i , as in the remainder of the section).
At test time we obtain deterministic predictions of the adjacency matrix elements by taking the argmax of the Bernoulli distribution for use as inputs to the decoder.
3.2. E NCODER
Our encoder is structured as an ( N + 1) × ( S + 1) lattice.
The N × S part of the lattice formed by the first N rows and first S columns receives a dataset D
= { ( x s 1 , . . .
, x sN ) T } Ss =1 .
This is unlike standard transformers which typically receive as input a single data sample (e.g., a sequence of words in neural machine translation applications) rather than a set of data samples.
Row N + 1 of the lattice is used to specify whether each data sample is observational , through value − 1 , or interventional , through integer value in { 1 , . .
. , N } to indicate the intervened node.
The goal of the encoder is to infer causal relationships between nodes by examining the set of samples.
The transformer performs this inference in multiple stages, each represented by one transformer layer, such that each layer yields a ( N +1)
×
( S +1) lattice of representations.
The transformer is designed to deposit its summary representation of the causal structure in column S + 1 .
Embedding of the input.
Each data-sample element x sn is embedded into a vector of dimensionality
H .
Half of this vector is allocated to embed the value x sn itself, while the other half is allocated to embed the unique identity for the node X n .
The value embedding is obtained by passing x sn , whether discrete or continuous, through an MLPUsing an MLP for a discrete variable is a slightly inefficient implementation of a node value embedding, but it ensures that the architecture is general.
encoder specific to node X n .
We use a node-specific embedding because the values of each node may have very different interpretations and meanings.
The node identity embedding is obtained using a standard 1D transformer positional embedding over node indices.
For column S + 1 of the input, the value embedding is a vector of zeros.
Alternating attention.
Traditional transformers discover relationships among the elements of a data sample arranged in a one-dimensional sequence.
With our two-dimensional lattice, the transformer could operate over the entire lattice at once to discover relationships among both nodes and samples.
Given an encoding that indicates position n, s in the lattice, the model can in principle discover stable relationships among nodes over samples.
However, the inductive bias to encourage the model to leverage the lattice structure is weak.
Additionally, the model is invariant to sample ordering, which is desirable because the samples are iid .
Therefore, we arrange our transformer in alternating layers.
In the first layer of the pair, attention operates across all nodes of a single sample ( x s 1 , . . .
, x sN ) T to encode the relationships among two or more nodes.
In the second layer of the pair, attention operates across all samples for a given node ( x 1 n , . . .
, x Sn ) to encode information about the distribution of node values.
Alternating attention in transformers was also done in Kossen et al.
Encoder summary.
The encoder produces a summary vector e sum n with H elements for each node
X n , which captures essential information about the node’s behavior and its interactions with other nodes.
The summary representation is formed independently for each node and involves combining information across the S samples (the columns of the lattice).
This is achieved with a method often used with transformers that involves a weighted average based on how informative each sample is.
The weighting is obtained using the embeddings in column S + 1 to form queries, and embeddings in columns 1 , . .
. , S to provide keys and values, and then using standard key-value attention.
3.2. D ECODER
The decoder uses the summary information from the encoder to generate a prediction of the adjacency matrix A of the underlying G .
It operates sequentially, at each step producing a binary output indicating the prediction ˆ A k,l of A k,l , proceeding row by row.
The decoder is an autoregressive transformer, meaning that each prediction ˆ A kl is obtained based on all elements of A previously predicted, as well as the summary produced by the encoder.
Our method does not enforce acyclicity.
Although this could in principle yield cycles in the graph, in practice we observe strong performance regardless (Section 6.4), likely due to the fact that training and evaluation graphs typically studied (e.g., ER-1 and ER-2) are very sparse.
Nevertheless, one could likely improve the results e.g. by using post-processing (Lippe et al., 2021) or by extending the method with an accept-reject algorithm (Castelletti & Mascaro, 2022; Li et al., 2022).
Auxiliary loss.
Autoregressive decoding of the flattened N × N adjacency matrix can be difficult for the decoder to learn alone.
To provide additional inductive bias to facilitate learning, we added auxiliary task of predicting the parents A n, : and children A : ,n of node X n from the encoder summary, e sum n .
This is achieved using an MLP to learn a mapping f n , such that f n ( e sum n ) =
( ˆ A n, : , ˆ A T : ,n ) .
While this prediction is redundant with the operation of the decoder, it short-circuits the autoregressive decoder and provides a strong training signal to support proper training.
S YNTHETIC DATA
In this section, we discuss identifiability and describe how synthetic data were generated.
Identifiability.
Dataset D i associated to CBN M i is given by D i
= { x s } Ss =1 , where x s :=
( x s 1 , . . .
, x sN )
T is either an observational or interventional data sample obtained by performing a hard intervention on a randomly selected node in G i .
As discussed in Eberhardt et al.
(2006), in the limit of an infinite amount of such single-node interventional data samples, G i is identifiable.
Note, that identifiability here means the ability to recover the exact graph given the data.
As our model defines a distribution over graphs, its predictions are meaningful even when the amount of data is insufficient for identifiability: in this case, the model would sample graphs that are compatible with the given data.
Empirically, we found that our model can make reasonable predictions even with a small amount of samples per intervention and improves as more samples are observed (Section 6.4).
Graph distribution.
We specified a distribution over G in terms of the number of nodes N (graph size) and number of edges (graph density) present in G .
As shown in Zheng et al.
Yu et al.
(2019); Ke et al.
(2020a), larger and denser graphs are more challenging to learn.
We varied N from 5 to 80 .
We used the Erd ˝ os–R ´ enyi (ER) and the scale-free (SF) metrics to vary density and evaluated our model on ER-1 and ER-2 graphs, as in Yu et al.
; Brouillard et al.
; Scherrer et al.
We generated an adjacency matrix A by first sampling a lower-triangular matrix to ensure that it represents a DAG, and by then permuting the order of the nodes to ensure random ordering.
Conditional probability distributions.
We performed ancestral sampling on the underlying CBN.
We considered both continuous and discrete nodes.
For continuous nodes, we generated continuous data using three methods following similar setups in previous works: (i) linear models ( linear data ) (Zheng et al., 2018; Yu et al., 2019), (ii) nonlinear with additive noise models ( ANM data ) (Brouillard et al., 2020;
Lippe et al., 2021), and (iii) nonlinear with non-additive noise models using neural networks (NN data) (Brouillard et al., 2020; Lippe et al., 2021).
For discrete nodes, we generated discrete data using two different methods: MLP ( MLP data ) and Dirichlet ( Dirichlet data ) conditionalprobability table generators.
Following past work (Ke et al., 2020a; Scherrer et al., 2021), we used a randomly initialized network.
The Dirichlet generator filled in the rows of a conditional probability table by sampling a categorical distribution from a Dirichlet prior with symmetric parameters α .
(We remind the reader that this generative procedure is performed prior to node ordering being randomized for presentation to the learning model.)
Refer to Appendix A.2 for details.
R ELATED WORK
Methods for inferring causal graphs can broadly be categorized into score-based, constraint-based, and asymmetry-based methods.
Score-based methods search through the space of possible candidate graphs, and ranks them based on some scoring function (Chickering, 2002; Cooper & Yoo, 1999;
Goudet et al., 2017; Hauser & B¨uhlmann, 2012; Heckerman et al., 1995; Tsamardinos et al., 2006;
Huang et al., 2018; Zhu et al., 2019).
Recently, Zheng et al.
Yu et al.
; Lachapelle et al.
(2019) framed the structure search as a continuous optimization problem.
There exist score-based methods that use a mix of continuous and discrete optimization (Bengio et al., 2019; Zhu et al., 2019; Ke et al., 2020a; Lippe et al., 2021;
Scherrer et al., 2021).
Constraint-based methods (Monti et al., 2019; Spirtes et al., 2000; Sun et al., 2007; Zhang et al., 2012; Zhu et al., 2019) infer the DAG by analyzing conditional independencies in the data.
Eaton & Murphy (2007) use dynamic programming techniques.
Asymmetry-based methods (Shimizu et al., 2006; Hoyer et al., 2009;
Peters et al., 2011; Daniusis et al., 2012; Budhathoki & Vreeken, 2017; Mitrovic et al., 2018) uses asymmetry between cause and effect to estimate the causal structure.
Peters et al.
; Ghassami et al.
; Rojas-Carulla et al.
(2018); Heinze-Deml et al.
(2018a) exploit invariance across environments.
Mooij et al.
(2016) propose a modeling framework that leverages existing methods.
[Table]
✓ ✓ ✓
[Table]
X ✓
[Table]
✓ ✓ ✓
[Table]
✓
X ✓
Learning-based methods have been proposed (Bengio et al., 2019; Goudet et al., 2018; Guyon, 2013; 2014; Kalainathan et al., 2018; Ke et al., 2020a;b; Lachapelle et al., 2022;
Lopez-Paz et al., 2015b; Wang et al., 2021b; Zhu et al., 2019).
In particular, Zhu et al.
; Wang et al.
(2021b) use transformers.
These works are concerned with learning only part of the causal induction pipeline, such as the scoring function, and hence are significantly different from our work, which uses an end-to-end supervised learning approach to learn to map from datasets to graphs.
Neural network methods equipped with learned masks exist (Douglas et al., 2017; Goyal et al., 2021; Ivanov et al., 2018;
Li et al., 2019; Yoon et al., 2018), but only a few have been adapted to causal inference.
Several transformer models (Goyal et al., 2022; Kossen et al., 2021; M ¨ uller et al., 2021) have been proposed for learning to map from datasets to targets.
However, none have been applied to causal discovery, although L ¨ owe et al.
(2022) proposes a neural-network based approach for causal discovery on time-series data.
A few supervised learning approaches have been proposed either framing the task as a kernel mean embedding classification problem (Lopez-Paz et al., 2015a;b) or operating directly on covariance matrices (Li et al., 2020) or binary classification problem of identifying v-structures (Dai et al., 2021).
These models accept observational data only (see Table 1), and because causal identifiability requires both observational and interventional data, our model is in principle more powerful.
E XPERIMENTS
We report on a series of experiments of increasing challenge to our supervised approach to causal structure induction.
First, we examined whether CSIvA generalizes well on synthetic data for which the training and test distributions are identical (Section 6.1).
This experiment tests whether the model can learn to map from a dataset to a structure.
Second, we examined generalization to an out-of-distribution (OOD) test distribution, and we determined hyperparameters of the synthetic data generating process that are most robust to OOD testing (Section 6.2).
Third, we trained CSIvA using the hyperparameters from our second experiment and evaluated it on a different type of OOD test distribution from several naturalistic CBNs (Section 6.3).
This experiment is the most important test of our hypothesis that causal structure of synthetic datasets can be a useful proxy for discovering causal structure in realistic settings.
Lastly, we performed a set of ablation studies to analyze the performance of CSIvA under different settings (Section 6.4).
All models are trained for 500 k iterations, please refer to Appendix A.3 for details on the hyperparameters.
Comparisons to baselines.
We compared CSIvA to a range of methods considered to be state-of-theart in the literature, ranging from classic to neural-network based causal discovery baselines.
For both in-distribution and OOD experiments, we compare to 4 very strong baselines: DAG-GNN (Yu et al., 2019), non-linear ICP (Heinze-Deml et al., 2018b), DCDI (Brouillard et al., 2020), and ENCO (Lippe et al., 2021).
For OOD experiments to naturalistic graphs, we compare to 5 additional baselines (Chickering, 2002; Hauser & B ¨ uhlmann, 2012; Zheng et al., 2018; Gamella & Heinze-Deml, 2020;
[Figure] [Figure]
Li et al., 2020).
Because some methods use only observational data and others do not scale well to large N, we could not compare to all methods in all experiments.
See Appendix A.4 for further discussion of alternative methods and conditions under which they can be used.
6. I N - DISTRIBUTION EXPERIMENTS
We begin by investigating whether CSIvA can learn to map from data to structures in the case in which the training and test distributions are identical.
In this setting, our supervised approach has an advantage over unsupervised ones, as it can learn about the training distribution and leverage this knowledge during testing.
We evaluate CSIvA for N ≤ 80 graphs.
Hamming distance H between predicted and ground-truth adjacency matrices, averaged over 128 sampled graphs.
CSIvA significantly outperforms DCDI and ENCO, both of which are very strong baselines.
The difference in performance increases with N .
We examined the performance on data with increasing order of difficulty, starting with linear (continuous data), before moving to non-linear cases (ANM, NN, MLP and Dirichlet data).
See Fig.
4 for comparisons between CSIvA and strong baselines on all data types ( N = 30 ), showing that
CSIvA significantly outperforms baselines for a wide range of data types.
For all but the Dirichlet data, we summarize the high-level results here, but detailed results and plots for all experiments can be found in Appendix A.5.
Results on continuous linear data are presented in Tables 4, 5 and results on continuous ANM data are presented in
Table 6, results on continuous NN data are reported in Table 7 in Appendix.
Results on MLP data are shown in
Fig.
5(b) (Appendix).
For all data types and all graphs, CSIvA significantly outperforms non-linear ICP,
DAG-GNN, DCDI and ENCO.
Differences become more apparent with larger graph sizes ( N ≥ 10 ) and denser graphs (ER-2 vs ER-1).
The Dirichlet data requires setting the values of the parameter α .
Hence, we run two sets of experiments on this data.
In the first set of experiments, we investigated how different values of α impact learning in CSIvA. As shown in Table 10 in the appendix, CSIvA performs well on all data with α ≤ 0 . 5 , achieving H < 2 . 5 in all cases.
CSIvA still performs well when α = 1 . 0 , achieving
H < 5 on size 10 graphs.
Learning with α > 1 is more challenging.
This is not surprising, as α >tends to generate more uniform distributions, which are not informative of the causal relationship between nodes.
In the second set of experiments, we compared CSIvA to DCDI and ENCO, as they are the strongest performing baselines.
We run the comparisons on graphs of size up to N ≤ 80 .
To limit the number of experiments to run, we set α = 0 .
1 , as this allows the conditional probabilities to concentrate on non-uniform distributions, which is more likely to be true in realistic settings.
As shown in Fig.
3, our model significantly outperforms DCDI and ENCO on all graphs (up to N ≤ 80 ).
As with other data types, the difference becomes more apparent as the size of the graphs grows ( N ≥ 50 ), as larger graphs are more challenging to learn.
Soft interventions and unknown interventions.
We additionally evaluated CSIvA on soft interventions, and on hard interventions for which the model did not know the intervention node ( unknown interventions ).
In both cases, we focus on Dirichlet data with α = 0 .
1 and α = 0 .
25 for ER-1 and
ER-2 graphs of size N ≤ 20 .
Results for the performance of CSIvA on soft interventions are shown in Table 13 in the Appendix.
CSIvA performed well on all graphs, achieving Hamming distance
H < 5 on all graphs.
For details about the data generation process, refer to Appendix Section A.5.7.
Results for the performance of CSIvA on unknown interventions are shown in Appendix Table 16.
CSIvA achieves strong performance on all graphs, with Hamming distance H < 7 on all graphs.
Note that CSIvA loses almost no performance despite not knowing the intervention; the biggest difference in Hamming distance between known and unknown interventions on all graphs is H < 1 . 0 .
For more details, refer to Section A.5.8 in the Appendix.
6. O UT - OF - DISTRIBUTION EXPERIMENTS
In this set of experiments, we evaluated CSIvA’s ability to generalize with respect to aspects of the data generating distribution that are often unknown, namely graph density and parameters of the
CPDs, such as the α values of the Dirichlet distribution.
Varying graph density.
We evaluated how well our model performs when trained and tested on
CBNs with varying graph density on MLP and α = 1 Dirichlet data and the number of nodes to
N = 7 , with variables able to take on discrete values in { 1 , 2 , 3 } .
The graphs in training and test datasets can take ER degree ∈ { 1 , 2 , 3 } .
Results are shown in Table 17 (Appendix) for the MLP data and Table 18 (Appendix) for the Dirichlet data.
For the MLP data, models trained on ERgraph generalizes the best.
For Dirichlet data, there isn’t one value of graph density that consistently generalizes best across graphs with different densities.
Nevertheless, ER- 2 graphs give a balanced trade-off and generalizes well across graphs with different sparsity.
Varying α .
We next trained and evaluated on data generated from Dirichlet distributions with α ∈ { 0 . 1 , 0 . 25 , 0 . 5 } .
Results for ER- 1 graphs with N = 7 are found in Table 19.
There isn’t a value of α that performs consistently well across different values of α for the test data.
Nevertheless, α = 0 .
25 is a balanced trade-off and generalizes well across test data with 0 .
1 ≤ α ≤ 0 . 5 .
S IM - TO - REAL EXPERIMENTS
In this final set of experiments, we evaluated CSIvA’s ability to generalize from being trained on
MLP and Dirichlet data to being evaluated on the widely used Sachs (Sachs et al., 2005), Asia (Lauritzen & Spiegelhalter, 1988) and Child (Spiegelhalter & Cowell, 1992) CBNs from the bnlearn repository, which have N = 11 , N = 8 and N = 20 nodes respectively.
We followed the established protocol from Ke et al.
(2020a)
; Lippe et al.
; Scherrer et al.
(2021) of sampling observational and interventional data from the CBNs provided by the repository.
These experiments are the most important test of our hypothesis that causal structure of synthetic datasets can be a useful proxy for discovering causal structure in realistic settings.
We emphasize that all hyperparameters for the MLP and Dirichlet data and for the learning procedure were chosen a priori ; only after the architecture and parameters were finalized was the model tested on these benchmarks.
Furthermore, to keep the setup simple, we trained on data sampled from a single set of hyperparameters instead of a broader mixture.
Findings in Section 6.2 suggest that ERgraphs with α = 0 .
25 work well overall and hence were chosen.
Results are reported in Table 2, comparing to a range of baselines from Squires et al.
(2020); Heinze-Deml et al.
(2018b);
Yu et al.
(2019); Gamella & Heinze-Deml (2020); Brouillard et al.
; Lippe et al.
(2021) and others.
[Table]
CSIvA and data significantly outperforms all other methods on both the Asia and Sachs datasets.
This serves as strong evidence that our model can learn to induce causal structures in the more realistic real-world
CBNs, while only trained on synthetic data.
6. D ETAILED ANALYSES
In order to better understand model performance, and which model components were most critical, we performed further analyses.
We describe here a subset of the studies performed; please refer to Appendix A.7 for the full results, including ablations.
Acyclicity.
We analyzed the generated (sampled) graphs for acyclicity under several settings.
We found that none of the generated graphs contains any cycles.
For details, refer to Appendix A.7.3.
Visualization of generated graphs.
We visualized some of the generated graphs in Figure 7 and
Figure 8 in the Appendix.
Note that samples are randomly drawn, and all are acyclic.
Identifiability upon seeing intervention data.
Here, we investigate how varying the proportion of interventional data impacts the performance of the proposed model.
As shown in Figure(Appendix), the model’s performance improves (Hamming distance decreases) almost monotonically as the amount of interventional data increases, from 0 − 100% .
This is a clear indication that our model is able to extract information from interventional data for identifying the graph structure.
For more details on the results and the how the experiments are conducted, please refer to Section A.7.4.
Amount of training cases.
We conducted experiments to better understand how much synthetic data is needed for different graph sizes.
For smaller graphs ( N ≤ 25 ), there is a small improvement for using more than 10 k training cases.
For larger graphs ( N > 25 ) the performance consistently improves as the model uses more training data (more details in Appendix A.7.6).
Ablation studies.
We investigate the role of various components in CSIvA. We found that all the different components play an important role in our model.
For details, see Appendix A.7.
Computation time.
Our models are trained for 500 k iterations.
Training can take up to 8 hours, and then inference takes minutes.
Training time of CSIvA can be amortized for multiple test cases; training time for baseline models are proportional to the number of test cases.
For more details, please refer to Appendix A.7.9.)
D ISCUSSION
In this paper, we have presented a novel approach towards causal graph structure inference.
Our method is based on learning from synthetic data in order to obtain a strong learning signal (in the form of explicit supervision), using a novel transformer-based architecture which directly analyzes the data and computes a distribution of candidate graphs.
Through a comprehensive and detailed set of experiments, we demonstrated that even though only trained on synthetic data, our model generalizes to out-of-distribution datasets, and robustly outperforms comparison methods under a wide range of conditions.
A direction of future work would be to use the proposed framework for learning causal structure from raw visual data.
This could be useful, e.g. in an RL setting in which an
RL agent interacts with the environment via observing low level pixel data (Ahmed et al., 2020; Ke et al.,
Wang et al., 2021a).
A A PPENDIX C ONTENTS
A.1 Transformer Neural Networks . . . . . . . . . . . . . . . . . . . . . . . . . . . .
A.2 Synthetic Data . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
A.3 Hyperparameters . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
A.
[Table]
Comparisons to
[Table]
A
[Table]
results for Section 6.1 . . . . . . . . . . . . . . . . . . . . . . . . . . . .
A.5.1 Results on
[Table]
data . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
A.5.2 Results on ANM data . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
A.5.
[Table]
Results on NN data . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
A.5.4 Results on MLP data . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
A.5.5 Results on
[Table]
A.5.6 Results
[Table]
A.
[Table] [Table]
interventions . .
A.5.8 Unknown Interventions . . . . . . . . . . . . . . . . . . . . . . . . . . . .
A.6 Detailed results for
-of- experiments . . . . . . . . . . . . . . . . .
A.6.1 Varying graph density . . . . . . . . . . . . . . . . . . . . . . . . . . . .
A.6.2 Varying conditional distribution . . . . . . . . . . . . . . . . . . . . . . .
A.6.3 Varying graph models . .
A.7 Ablation studies and analyses . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
A.7.
Visualization
of sampled graphs . . . . . . . . . . . . . . . . . . . . . . .
A.7.
Additional evaluations
A.7.
3 Acyclicity of generated
graphs . . . . . . . . . . . . . . . . . . . . . . . .
A.7.
Identifiability upon seeing intervention
data. . . . . . . . . . . . . . . . . .
A.
Varying number of samples
A.
7.6 Varying amount of training data
A.7.
7 Interventions on a few variables
A.
7.8 Ablation studies
A.7.9 Computation time . .
A.8 Pseudocode for CSIvA Model . . . . . . . . . . . . . . . . . . . . . . . . . . . .
A. T RANSFORMER N EURAL N ETWORKS
The transformer architecture, introduced in Vaswani et al.
(2017), is a multi-layer neural network architecture using stacked self-attention and point-wise, fully connected, layers.
The classic transformer architecture has an encoder and a decoder, but the encoder and decoder do not necessarily have to be used together.
Scaled dot-product attention.
The attention mechanism lies at the core of the transformer architecture.
The transformer architecture uses a special form of attention, called the scaled dot-product attention.
The attention mechanism allows the model to flexibility learn to weigh the inputs depending on the context.
The input to the QKV attention consists of a set of queries, keys and value vectors.
The queries and keys have the same dimensionality of d k , and values often have a different dimensionality of d v .
Transformers compute the dot products of the query with all keys, divide each by √ d k , and apply a softmax function to obtain the weights on the values.
In practice, transformers compute the attention function on a set of queries simultaneously, packed together into a matrix Q .
The keys and values are also packed together into matrices K and V .
The matrix of outputs is computed as:
Attention
[Equation]
Encoder.
The encoder is responsible for processing and summarizing the information in the inputs.
The encoder is composed of a stack of N identical layers, where each layer has two sub-layers.
The first sub-layer consists of a multi-head self-attention mechanism, and the second is a simple, position-wise fully connected feed-forward network.
Transformers employ a residual connection (He et al., 2016) around each of the two sub-layers, followed by layer normalization (Ba et al., 2016).
That is, the output of each sub-layer is LayerNorm( x + Sublayer( x )) , where Sublayer( x ) is the function implemented by the sub-layer itself.
Decoder.
The decoder is responsible for transforming the information summarized by the encoder into the outputs.
The decoder also composes of a stack of N identical layers, with a small difference in the decoder transformer layer.
In addition to the two sub-layers in each encoder layer, a decoder layer consists of a third sub-layer.
The third sub-layer performs a multi-head attention over the output of the encoder stack.
Similar to the encoder, transformers employ residual connections around each of the sub-layers, followed by layer normalization.
Transformers also modify the self-attention sub-layer in the decoder stack to prevent positions from attending to subsequent positions.
This masking, combined with fact that the output embeddings are offset by one position, ensures that the predictions for position i can depend only on the known outputs at positions less than i .
A. S YNTHETIC D ATA
In Section 4, we introduced the types of methods used to generate the conditional probability distribution.
In this section, we discuss the details of these methods.
We evaluate our model on datasets generated from 5 different methods, which covers both continuous and discrete valued data generated with a varying degree of difficulty.
For continuous data, we generated the data using three different methods: linear , non-linear additive noise model (ANM) and non-linear non-additive noise neural network model (NN).
For discrete data, we generated data using two different methods: MLP and Dirichlet .
Let X be a N × S matrix representing S samples of a
CBN with N nodes and weighted adjacency matrix A , and let Z be a random matrix of elements in
N (0 , 0 . 1) .
We describe each type of how different types of data is generated.
• For linear data , we follow the setup in Zheng et al.
(2018) and Yu et al.
Specifically, we generated data as X n, : =
A n, : X + Z n, : .
The biases were initialized using U [ − 0 .
5 , 0 . 5] , and the individual weights were initialized using a truncated normal distribution with standard deviation of 1 . 5 .
For nodes with interventions, values are sampled from the uniform distribution U [ − 1 , 1] .
• For additive noise models (ANM) , we follow the setup in (Brouillard et al., 2020; Lippe et al.,
We generated the data as X n, : = F n, : ( X ) + 0 .
· Z n, : , where F is fully connected neural network, the weights are randomly initialized from N (0 , 1) .
The neural network has one hidden layer with 10 hidden units, the activation function is a leaky relu
with a slop of 0 . 25 .
The noise variables are sampled from N (0 , σ 2 ) , where σ 2 ∼ U [1 , 2] .
For nodes with interventions, values are sampled from the uniform distribution N (2 , 1) .
• For non-additive noise neural network (NN) models , we also follow the setup in (Brouillard et al., 2020; Lippe et al., 2021).
We generate the data as X n, : = F n, : ( X, Z n, : ) , such that F is fully connected neural network, the weights are randomly initialized from N (0 , 1) .
The neural network has one hidden layer with 20 hidden units, the activation function is a tanh function.
The noise variables are sampled from N (0 , σ 2 ) , where σ 2 ∼ U [1 , 2] .
For nodes with interventions, values are sampled from the uniform distribution N (2 , 1) .
• For MLP data , the neural network has two fully connected layers of hidden dimensionality 32 .
Following past work (Ke et al., 2020a;
Scherrer et al., 2021; Lippe et al., 2021), we used a randomly initialized network.
For nodes with interventions, values are randomly and independently sampled from U { 1 , . .
, K } where K indicates the number of categories of the discrete variable.
• For Dirichlet data, the generator filled in the rows of a conditional probability table by sampling a categorical distribution from a Dirichlet prior with symmetric parameters α .
Values of α smaller than 1 encourage lower entropy distributions; values of α greater thanprovide less information about the causal relationships among variables.
Similar to MLP data, for nodes with interventions, values are randomly and independently sampled from U { 1 , . .
, K } where K indicates the number of categories of the discrete variable.
A. H YPERPARAMETERS
[Table]
For all of our experiments (unless otherwise stated) our model was trained on I = 15 , 000 (for graphs
N ≤ 20 ) and on I
= 20 , 000 (for graphs N > 20 ) pairs { ( D i , A i ) } Ii =1 , where each dataset D i contained S = 1500 observational and interventional samples.
For experiments on discrete data, a data-sample element x s could take values in { 1 , 2 , 3 } .
Details of the data generating process can be found in Section 4.
For evaluation in Sections 6.1 and 6.2, our model was tested on I ′
=
128 (different for the training) pairs { ( D i ′ ,
A i ′ )
} I ′ i ′
=1 , where each dataset D i ′ contained S = 1500 observational and interventional samples.
For the Asia, Sachs and Child benchmarks, our model was still tested on I ′
=
128 (different for the training) pairs { ( D i ′ ,
A i ′ )
} I ′ i ′
=1 , however, A i ′ =
A j ′ since there is only a single adjacency matrix in each one of the benchmarks.
We present test results averaging performance over the 128 datasets and 3 random seeds and up to size N = 80 graphs.
The model was trained for 500 , 000 iterations using the Adam optimizer (Kingma & Ba, 2014) with a learning rate of 1e − 4 .
Also, refer to Table 3 for the list of hyperparameters presented in a table.
We parameterized our architecture such that inputs to the encoder were embedded into 128dimensional vectors.
The encoder transformer has 10 layers and 8 attention-heads per layer.
The final attention step for summarization has 8 attention heads.
The decoder was a smaller transformer with only 2 layers and 8 attention heads per layer.
Discrete inputs were encoded using an embedding layer before passing into our model.
A. C OMPARISONS TO B ASELINES
In Section 6.1, we compare CSIvA to four strong baselines in the literature, ranging from classic causal discovery baselines to neural-network based causal discovery baselines.
These baselines are
DAG-GNN (Yu et al., 2019), non-linear ICP (Heinze-Deml et al., 2018b), DCDI (Brouillard et al., 2020) and ENCO (Lippe et al., 2021).
Non-linear ICP, DCDI and ENCO can handle both observational and interventional data, while
DAG-GNN can only use observational data.
Non-linear ICP could not scale to graphs larger than 20, therefore we compare to DCDI (Brouillard et al., 2020) and ENCO (Lippe et al., 2021) on larger graphs.
All baselines are unsupervised methods, i.e. they are not tuned to a particular training dataset but instead rely on a general-purpose algorithm.
We also compared to an all-absent model corresponding to a zero adjacency matrix, which acts as a sanity check baseline.
We also considered other methods (Chickering, 2002; Hauser & B¨uhlmann, 2012; Zhang et al., 2012; Gamella & HeinzeDeml, 2020), but only presented a comparison with DCDI, ENCO, non-linear ICP and DAG-GNN as these have shown to be strong performing models in other works (Ke et al., 2020a;
Lippe et al., 2021;
Scherrer et al., 2021).
For Section 6.3, we also compared to additional baselines from Chickering (2002);
Hauser &
B¨uhlmann (2012); Zheng et al.
; Gamella & Heinze-Deml (2020);
Li et al.
Note that methods from Chickering (2002); Zheng et al.
Yu et al.
; Li et al.
(2020) take observational data only.
DAG-GNN outputs several candidate graphs based on different scores, such as evidence lower bound or negative log likelihood, DCDI can also be run in two different settings (DCDI-G and DCDI-DSF), we chose the best result to compare to our model.
Note that non-linear
ICP does not work on discrete data, i.e. on the MLP and Dirichlet data, therefore a small amount of
Gaussian noise N (0 , 0 . 1) was added to this data in order for the method to run.
A. D ETAILED RESULTS FOR S ECTION 6.
We present detailed results for experiments in Section 6.1 are described in the tables below.
A.5. R ESULTS ON LINEAR DATA
[Table]
et al., 2018b), the method required nodes to be causally ordered, and 20 repeated samples taken per intervention, as interventions were continuously valued.
”Abs” baselines are All-Absent baselines, which is an baseline model that outputs all zero edges for the adjacency matrix.
DCDI and ENCO are the strongest performing baselines.
CSIvA outperforms all baselines (including DCDI and ENCO).
Results for comparions between our model CSIvA and baselines non-linear ICP (Heinze-Deml et al., 2018b), DAG-GNN (Yu et al., 2019), DCDI (Brouillard et al., 2020) and ENCO (Lippe et al., 2021) are shown in Table 4 for smaller graphs with N ≤ 20 and Table 5 for larger graphs with 20 < N ≤ 80 .
For smaller graphs with N ≤ 20 , all models that takes interventional data perform significantly better compared to DAG-GNN (Yu et al., 2019), which only takes observational data.
CSIvA achieves
Hamming distance H < 7 on evaluated graphs up to size 20 .
Similar to previous findings (Yu et al.,
2019; Ke et al., 2020a), larger and denser graphs are more challenging to learn.
Non-linear ICP achieves fairly good performance on smaller graphs ( N ≤ 10 ), however, the performance drops quickly as size of graphs increases ( N > 10 ).
Also, note that Non-linear ICP can not scale to graphs larger than N > 20 .
It also required a modificationWithout this modification, the method achieved near chance performance.
to the dataset wherein multiple samples were collected from the same modified graph after a point intervention (20 samples per intervention), while other methods only sampled once per intervention.
For larger graphs of 20 < N ≤ 280 , we compare to strongest baselines: DCDI and ENCO.
Results are found in Table 5.
CSIvA significantly outperforms both DCDI and ENCO.
The difference because more apparent as the size of the graphs grow ( N ≥ 50 ).
[Table]
A.5. R ESULTS ON ANM DATA
For additive noise non-linear model (ANM) data, we compare to the strongest baseline models: DCDI and ENCO on N ≤ 80 graphs.
Results are found in Table 6.
CSIvA achives hamming distance
H < 11 on all graphs of size up to 80 .
Therefore, CSIvA significantly outperforms strong baseline models DCDI and ENCO on all graphs.
This difference becomes more apparent on larger graphs of
N ≥ 50 .
[Table]
Hamming distance H for learned and ground-truth edges on synthetic graphs, compared to other methods.
The number of variables varies from 30 to 80, expected degree = 1 or 2.
We compare to DCDI (Brouillard et al., 2020) and ENCO (Lippe et al., 2021) as they are the best performing baselines and they are able to handle larger graphs.
A.5. R ESULTS ON NN DATA
Results for comparisons between CSIvA and strongest baselines DCDI and ENCO on non-additive noise non-linear (NN) data are found in Table 7.
CSIvA achieves hamming distance H < 11 on all graphs.
Hence, CSIvA significantly outperforms DCDI and ENCO on all graphs.
The differences grow larger as the size of the graph grows.
Additionally, we generated non-additive noise non-linear (NN) data with scale-free (SF) graphs with degree of 4 and 6 for N ≤ 50 .
We compare our model to strongest baselines DCDI and ENCO on these data.
The results are found in Table 15.
[Table]
[Figure]
A.5. R ESULTS ON MLP DATA
Results for comparisons between our model CSIvA and baselines non-linear ICP (Heinze-Deml et al., 2018b) and DAG-GNN (Yu et al., 2019) on MLP data are shown in Table 9.
MLP data is non-linear and hence more challenging compared to the continuous linear data.
Our model CSIvA significantly outperforms non-linear ICP and DAG-GNN.
The difference becomes more apparent as the graph size grows larger and more dense.
[Table] [Table]
A.5. R ESULTS ON D IRICHLET DATA .
We run two sets of experiments on Dirichlet data.
The first set of experiments is aimed at understanding how different α values impact the performance of our model.
In the second set of experiments, we compare the performance of our model to four strong baselines: DCDI (Brouillard et al., 2020),
ENCO (Lippe et al., 2021), non-linear ICP (Heinze-Deml et al., 2018b) and DAG-GNN (Yu et al., 2019).
[Table]
truth edges on synthetic graphs, averaged over 128 sampled graphs.
Our model accomplished a hamming distance of less than 2 . 5 for Dirichlet data with α ≤ 0 . 5 . ”
Abs” baselines are All-Absent baselines, which is an baseline model that outputs all zero edges for the adjacency matrix.
The results for the first set of experiments are shown in Table 10.
Our model performs well on all graphs where α ≤ 0 . 5 , and the performance starts to degrading as α = 1 . 0 .
When α = 5 . 0 , our model is almost performing similarly to the All-absent model (outputting all zero edges).
This is to be expected, as larger alpha values is less informative of the causal relationships between variables.
For the second set of experiments, we compared
CSIvA to non-linear ICP and DAG-GNN on graphs of size up to 20.
To limit the number of experiments to run, we set α = 1 . 0 , as this gives the least amount of prior information to
CSIvA. As shown in Figure 6, our model significantly outperforms non-linear ICP and DAGGNN.
Our model achieves H < 5 on sizegraphs, almost half of the error rate compared
[Table]
to non-linear ICP and DAG-GNN, both achieving a significantly higher Hamming distance ( H = 9 .and
H = 9 . 5 respectively) on larger and denser graphs.
Refer to Table 11 for complete sets of results.
For larger graphs ( N > 20 ), we compare CSIvA to DCDI (Brouillard et al., 2020) and ENCO (Lippe et al., 2021) as they could scale to larger graphs and are the strongest baselines.
The results are illustrated in Figure 3 and the detailed results are found in Table 12.
As the graphs get larger, the performance of baselines DCDI and ENCO drops significantly, for dense and largr graphs (ER-2, N = 80 ), the baseline models achieve almost near all-absent performance, while our model performs significantly better (achieving almost 30% performance gain in terms of structured hamming distance).
[Table]
A.5. R ESULTS ON SCALE - FREE GRAPHS
We further evaluated our model’s performance on on scale-free (SF) graphs.
We run experiments for both discrete and continuous variables.
For discrete variables.
We train and test our model on
Dirichlet data with α = 0 .
1 on scale-free (SF) graphs.
We run them on SF graphs with degree 4 and 6 similar to the setup in Zheng et al.
Results on Dirichlet data is reported in Table 14 and results on NN data is found in Table 15.
Our model significantly outperforms strong baseline models
DCDI (Brouillard et al., 2020) and ENCO (Lippe et al., 2021) in all cases, regardless of the type of data, size of the graph and the sparsity of the graphs.
[Table] [Table]
et al.
(2020) and Lippe et al.
(2021) as they are the strongest baselines and they are able to handle larger graphs.
A.5. S OFT INTERVENTIONS
In addition to hard interventions, we also consider soft interventions.
Hence, we further evaluate our model on imperfect interventions as discussed in Section 6.1.
To limit the number of experiments to run, we focus on Dirichlet data with α = 0 .
1 and α = 0 .
Soft interventions on Dirichlet data is generated as follows.
For an intervention on variable X
i , we first sample the α i for the Dirichlet distribution from a Uniform distribution U { 0 . 1 , 0 . 3 , 0 . 5 , 0 .
then using the sampled α i , we sample the conditional probabilities from the Dirichlet distribution with the new α i .
The results for CSIvA on soft interventions are shown in Table 13.
CSIvA was able to achieve a hamming distance H < 5 on all graphs of size H ≤ 20 .
These results are strong indications that that our model still works well on the imperfect interventions.
A.5. U NKNOWN I NTERVENTIONS
Until now, we have considered cases where the target of the intervention is known, we now consider the case when the target of the intervention is unknown, this is also referred to as unknown interventions.
Again, to limit the number of experiments to run, we focus on Dirichlet data with α = 0 .
1 and α = 0 .
The model does not know the target of the intervention and all other training procedures remains exactly the same.
The results are shown in Table 16.
We compare how well CSIvAperforms on known and unknown interventions.
We can see that the performance for the known and unknown interventions are almost the same for sparser graphs (ER-1).
The differences increases slightly for denser graphs with higher values of α , for example ER − 2 graphs with α = 0 .
This is to be expected, as denser graphs with higher α values are more challenging to learn.
The biggest differences for the performances of known and unknown interventions is less than H < 1 . 0 in hamming distance.
These results shown clear indications that our model performs well even when the intervention target is unknown.
[Table] [Table]
and ground-truth edges on synthetic graphs, for known vs unknown interventions, averaged over 128 sampled graphs ( ± standard deviation).
The number of variables varies from 15 to 20, expected degree = 1 or 2, α = 0 .
1 or 0 . 25 , and the dimensionality of the variables are fixed to 3.
A. D ETAILED RESULTS FOR O UT - OF - DISTRIBUTION EXPERIMENTS
This section contains detailed results for the out-of-distribution experiments experiments in Section 6.2.
A.6. V ARYING GRAPH DENSITY
Results for the experiments on varying graph density in Section 6.2 are shown in Table 17 and Table 18.
Our model generalizes well to OOD test distribution where the graphs vary in terms of density.
A.6. V ARYING CONDITIONAL DISTRIBUTION
The results on varying α values in Section 6.2 are found in Table 19.
Our model generalizes well to OOD test distributions, where the conditional probability can vary in terms of α values for the
Dirichlet distributions.
[Table] [Table]
[Table] [Table]
[Table]
A.6. V ARYING GRAPH MODELS
We also evaluate how well our model generalizes when the model used for generating the graphs during training and test differs.
We trained and tested our model on Erd ˝ os–R ´ enyi (ER) and Scale-Free (SF) graph models.
To limit the number of experiments to run, we evaluate our model on Dirichlet data with α = 0 .
We run experiments both on N = 10 and N = 20 graphs.
The models are trained on Erd˝os–R´enyi (ER) graphs and tested on Scale-Free (SF) graphs.
The results are shown in Tableand Table 21.
A. A BLATION STUDIES AND ANALYSES
In order to better understand the performance of our model, we performed further analysis and ablations.
This section contains full results for Section 6.4.
In this section, we aim to answer the following questions:
• What does visualizations of the generated graphs look like?
• Additional evaluation metrics to measure structural accuracy.
• Are these generated graphs from our model acyclic?
• Does intervention data improve identifiability in our model?
• How does varying the number of samples impact the performance of our model?
• How does varying the amount of training datasets impact the peformance of our model?
• How does different components of our model (such as sample-level attention and auxiliary loss) impact the performance of our model?
A.7. V ISUALIZATION OF SAMPLED GRAPHS
We visualized samples that our model generated on the test data.
The samples are shown in Figureand Figure 7.
The samples are randomly chosen, each subplot is a sample from a distinct test data.
The edges in the graph are shown in 3 colors, they each represent the following: (a) Green edges indicate that our model has generated the correct edge.
(b) A red edge indicates a missing edge, that is our model did not generate the edge, which exist in the groundtruth graph.
(c)
A blue edge indicates a redundant edge, such that our model generated an edge that does not exist in the groundtruth graph.
As shown in Figure 8 and 7, our model is able to generate the correct graph almost all of the times.
A.7. A DDITIONAL EVALUATIONS
We evaluated our model on additional evaluation metrics to better understand the performance of the model.
In particular, we evaluated our model on the balanced scoring function (BSF), F1 score, precision, recall and area under roc (AUROC).
We evaluated our model both on discrete and continous data.
For discrete data, we used Dirichlet data with α = 0 .
1 and for continous data, we used nonlinear non-additive noise (NN) data.
Results on Dirichlet data is in Table 22, results on NN data is in
Table 23.
For Dirichlet data, our model achieved strong results.
Again, the area under ROC (AUROC) is greater than 0.95 for all cases.
The F1 score is also always above 0.90.
For continuous data, our model achieved even stronger results.
Our model scored between 0.98 and 0.99 for all scores.
This is an indication that the edges were almost always correct and the correct edges also have strong probability.
[Table]
[Table]
A.7. A CYCLICITY OF GENERATED GRAPHS
worst
As discussed in Section 6.
we analyzed the generated graphs for acyclicity on Dirichlet data with α = 0 .
1 and α = 0 .
25 for graphs of size up to N ≤ 20 .
We evaluated the model ontest datasets per setting, and found that none of the generated graphs contains cycles, all graphs are acyclic.
This is a clear indication showing that although we do not have an acyclicity regularizer, our model was still able to learn to generated DAGs.
A.7. I DENTIFIABILITY UPON SEEING INTERVENTION DATA .
Our previous experiments were all performed using a fixed amount of interventions ( 80% ) in the training and test sets.
To investigate how changing the proportion of interventional data impacts the performance of our model, we train the model with varying amounts of interventions in the training set and evaluate it using different amount of interventions during test time.
To be specific, during training, the model is trained on data with varying amount of interventions, which is randomly sampled from the uniform distribution U [0 , 0 . 1 , 0 . 3 , 0 . 5 , 0 . 7 , 0 . 9 , 1 . 0] .
During test time, we evaluate the model on different amount of interventions, and we report the performance in hamming distance.
We trained the model on Dirichlet data with 15 nodes and α = 0 .
The model is trained and tested on 1000 samples per dateset.
The results are found in Figure 9.
As shown in Figure 9, our model’s performance is worst if it only receives observational data ( 0% interventions), and the performance of our model improves as the amount of interventional data increases.
This is a clear indication that our model is able to extract information from interventional data for predicting the graph structure.
A.7. V ARYING NUMBER OF SAMPLES
[Figure]
We evaluated CSIvA on different amount of samples ( 100 , 200 , 500 , 1000 , 1500 ) per CBNs.
To limit the number of experiments to run, we focus on Dirichlet data sampled from N = 10 graphs.
During the ith iteration in training, the model takes in as input a dataset with l i samples, where l i is sampled from
U [100 , 200 , 500 , 1000 , 1500] .
During test time, we run separate evaluations on the model for test datasets with different number of samples and the results are reported in Figure 10.
We can see that the model performance improves as it observes up to 1000 samples for ER- 1 graphs, whereas having 1500 samples gives slightly better results compared to 1000 samples for ERgraphs.
A.7. V ARYING AMOUNT OF TRAINING DATA
In this section
[Table]
we evaluate how the amount of training data impacts the performance of our model (CSIvA ).
To limit the number of experiments to run, we focus on Dirichlet data with α = 0 .
The model receives M amount of training datasets, where M is sampled from
U [5000 , 10000 , 15000 , 20000] .
[Table]
A.7. I NTERVENTIONS ON A FEW VARIABLES
We performed additional experiments to understand the performance of our model when only a few variables are intervened on.
We varied the percentage of variables to intervene on from 10% to 100% .
Results are reported in Table 25.
We used Dirichlet data with alpha = 0 .
1 for both training and test data.
Our model’s performance improves consistently as more nodes are intervened on.
[Table]
Hamming distance H for learned and ground-truth edges on synthetic graphs, compared to other methods.
The percentage of node that the model can intervene on varies from 10% 100% .
The performance of the model improves consistently as it more nodes are intervened on.
A.7. A BLATION STUDIES
We investigate the role that different components play in our model, for example, what are the effects of having auxillary loss and sample-level attention for our model.
The details for each experiment is described below.
Sample-level attention We first evaluated the performance of our with and without sample-level attention.
In the case without sample-level attention, our model will only use node-level attention for all 10 layers in the encoder transformer.
The results are reported in Table 26.
The models with sample-attention is a vanilla CSIvA model; and the model without sample-attention is CSIvA without sample-level attention.
As one could see that the performance of the model drops significantly if sample-attention is not used.
Hereby making sample-level attention a crucial component in our model.
ER =
[Table]
Var = Var =
Var = Var =
α = 0 .α = 0 .α
= 0 .α = 0 .α
= 0 .α = 0 .α
= 0 .α
= 0 .
Sampleattention 1.
No sample attention 15.
[Table]
Auxiliary loss We also conducted ablation studies to understand the effect of the auxiliary loss in the objective function, results are reported in Table 27.
The model with auxiliary loss is the vanilla CSIvA model, where as the one without is CSIvA trained without the auxiliary loss objective.
Experiments are conducted on Dirichlet data with α = 0 .
1 , the model does gain a small amount of performance ( H < 1 . 0 ) by having the auxiliary loss, the difference becomes more apparent as the size of the graph grows, indicating that auxiliary loss plays an important role as a part of the objective function.
A.7. C OMPUTATION TIME
All the models with the proposed method are trained for 500 k iterations and all models have converged by then (in fact they have all converged by 400 k ).
Once the model is trained, evaluation/ inference is fast.
All models are evaluated on 128 test datasets (graphs).
The focus of causal discovery is on high-stakes test cases where accuracy matters more than compute.
All models (baselines and our model) are trained on GPUs.
For baseline methods like DCDI, each test
[Table]
case can to
(More details are provided below.)
The training time of our model can be amortized if the model can be used on multiple test cases; in contrast, training time for ENCO and DCDI is proportional to the number of test cases.
For our model CSIvA , training on smaller graphs ( N ≤ 20 ) converges within 50 k , which takes up to 45 minutes.
Medium sized graphs ( 20 < N ≤ 50 ) converge within 250 k and take 4 hours.
Large graphs ( 50 < N < 80 ) converge within 400 k and take 8 hours.
Note that once the model is trained, it is very fast to evaluate on different test cases.
Our model is evaluated on 128 test cases.
For the baseline model DCDI takes up to 4 hours on a single graph of size up to 20 ( N < 20 ).
The model takes up to 12 hours on a medium size graph ( N < 50 ) and 48 hours for larger graphs ( 50 < N < 80 ).
For the baseline model ENCO.
Smaller graphs ( N < 20 ) take 20 minutes per graph and larger graphs ( N < 80 ) take 30 minutes per graph.
To compute 128 test cases, ENCO takes 2.
days of computation time.
P SEUDOCODE FOR CSI V A M ODEL
First, we present a high-level description of our training algorithm in Algorithm 1.
# a
Transformers from Transformers import TransformerBlock # a generic helper function for Multi Head Dot Product Attention.
MultiheadAttention import functools import haiku as hk import jax import jax.numpy as jnp class MultiMLP(hk.
Module): def __init__(self, hidden_sizes, num_var, name=
None): super ().__init__(name=name) self._mlps =
[hk.nets.
MLP(hidden_sizes) for _ in range (num_var)] def __call__(self, inputs): results = [] for i, mlp in enumerate (self._mlps): results.append(mlp(inputs[i])) return jnp.stack(results) class Model(hk.
Module): def _init_(self, num_var=3, var_dim=2, num_layers=8, num_heads=8, encode_num_hidden=64, use_embedding=
True, emb_dim=32): self.enc_num_hidden
= encode_num_hidden self.emb_dim = emb_dim self.var_dim = var_dim self.num_heads = num_heads self.action_encoders = hk.nets.
MLP( [self.enc_num_hidden self.use_embedding = use_embedding
# a generic transformer block self.transformer_block =
TransformerBlock(causal=
False, self.num_heads)
# a generic multi-head attention block.
self.attention = MHA(num_heads=4, qk_channels=128, v_channels=128) self.embds = hk.get_parameter( ’embs’, [self._var_dim, self._emb_dim], init=
init.
TruncatedNormal(stddev=1.0))
self.input_encoder = MultiMLP([self._encode_num_hidden, self._emb_dim], self._num_var)
# A helper function to create positional encodings def positional_encodings(self, inputs): ..... return positional_encodings # a helper function to calculate auxiliary loss and return the loss and the sampled graph.
def compute_aux_loss(self, graph, cond): ....
return loss, sampled_graph # Step 1: Process the input.
def process_input(self, inputs, intvn_indices): # Args # intvn_indices: multi-hot mask describing interventions, eg,means there are interventions on node 1 and 4.
if self.use_embedding: inputs = self.embs[inp.astype( int )] else : inputs =
self.input_encoder(inp[..., np.newaxis])
action_embed = self.action_encoders(inputs)
action_embed =
jnp.concatenate((action_embed, pad[-1][jnp.newaxis, ...]),axis=-2) inputs =
jnp.concatenate((inputs, action_embed), axis=0) return inputs # Step 2: Alternate attention between different graphs and within graph.
def transformer_pass(self, inputs, intvn_mask, graph,
# preprocess input.
inputs = self.process_input(inputs, intvn_mask)
# create positional encodings.
pos_encodings = self.positional_encodings(inputs, dimension=0) batch_size, seq_len, _ = inputs.shape pos_encodings =
jnp.tile( pos_encodings.reshape(batch_size, 1, -1), (1, seq_len, 1)) inputs =
jnp.concatenate((inputs, pos_encodings), axis=-1) h = inputs if mask is not None: atn_mask = mask
[:, None, None, :] else : atn_mask =
None # Go through all possible layers in the forward pass.
for layer_i in range (num_layers): if layer_i % 2 == 0: h = h.transpose(1, 0, 2) block = self.transformer_block block = functools.partial(block, is_training=is_training)
h
= block(h, mask=atn_mask) if layer_i % 2 ==
h = h.transpose(1, 0, 2) self_attention_out =
None # h dimension: (num_var + 1) * (timestep + 1) * hidden_dim prev_h =
h[:, :-1, :] # prev_h: (num_var + 1)
* time_step * hidden_dim last_step =
h[:, -1, :] #
last_step: (num_var + 1 ) * hidden_dim all_out =
[] for i in range (self._num_var):
att = self.attention kv_inputs = prev_h[i][jnp.newaxis, ...]
# context: 1 * time_step * hidden_dim q_inputs = last_step[i][jnp.newaxis, jnp.newaxis, ...] out = att