Amortized Inference for Causal Structure Learning
Authors: Lars Lorch
Scott Sussex
Jonas Rothfuss
Andreas Krause
Bernhard Schölkopf
Venue: NeurIPS 2022 Accept (2022)
Keywords: causality
amortized inference
causal discovery
structure learning
Bayesian causal discovery

Abstract: Inferring causal structure poses a combinatorial search problem that typically involves evaluating structures with a score or independence test. The resulting search is costly, and designing suitable scores or tests that capture prior knowledge is difficult. In this work, we propose to amortize causal structure learning. Rather than searching over structures, we train a variational inference model to directly predict the causal structure from observational or interventional data. This allows our inference model to acquire domain-specific inductive biases for causal discovery solely from data generated by a simulator, bypassing both the hand-engineering of suitable score functions and the search over graphs. The architecture of our inference model emulates permutation invariances that are crucial for statistical efficiency in structure learning, which facilitates generalization to significantly larger problem instances than seen during training. On synthetic data and semisynthetic gene expression data, our models exhibit robust generalization capabilities when subject to substantial distribution shifts and significantly outperform existing algorithms, especially in the challenging genomics domain. Our code and models are publicly available at: https://github.com/larslorch/avici

Amortized Inference for Causal Structure Learning
Lars Lorch ETH Zurich Zurich, Switzerland llorch@ethz.ch
Scott Sussex ETH Zurich Zurich, Switzerland
Jonas Rothfuss
ETH Zurich Zurich, Switzerland ro waooonnstthzcccmmm ch
ssussex@ethz.co
Andreas Krause"
Bernhard Schölkopf"
ETH Zurich Zurich, Switzerland
MPI for Intelligent Systems
Tübingen, Germany
krausea@ethz.ch
bs@tuebingen. mpg.
Abstract
Inferring causal structure poses combinatorial search problem that typically in- volves evaluating structures with a score or independence test. The resulting search is costly, and designing suitable scores tests that capture prior knowledge is difficult. this work, we propose to amortize causal structure learning. Rather than searching over structures, we train a variational inference model directly predict the causal structure from observational or interventional data. This allows our inference model to acquire domain-specific inductive biases for causal discov- ery solely from data generated by simulator, bypassing both the hand-engingg -agiinnning of suitable score functions and the search over graphs. The architecture of our inference model emulates permutation invariances that are crucial for statistical efficiency in structure learning, which facilitates generalization to significantly larger problem instances than seen during training. On synthetic data and semisyn- thetic gene expression data, our models exhibit robust generalization capabilities when subject to substantial distribution shifts and significantly outperform existing algorithms,, especially in the challenging genomics domain. Our code and models are publicly available at: https: //gi thub. laassolchhaavii.....
Introduction
Learning the causal structure among set of variables is a fundamental task in various scientific disciplines (Spirtes et al., 2000; Pearl, 2009). However, inferring this causal structure from observa- tions of the variables is difficult inverse problem. The solution space of potential causal structures, usually modeled as directed graphs, grows superexponentially with the number of variables. To infer a causal structure, standard methods have to search over potential graphs, usually maximizing either graph scoring function or testing for conditional independences (Heinze-Deml al., 2018).
Specifying realistic inductive biases universally difficult for existing approaches causal discovery. Score- based methods use strong assumptions about the data tiatt--enerating process, such as linearity (Shimizu al., 2006), specific noise models (Hoyer et al., 2008; Peters and Bühlmann, 2014), and the absence of measurement егтоr (cf. Scheines and Ramsey 2016; Zhang et al. 2017), which are difficult to verify (Dawid, 2010; Reisach al., 2021). Conversely, constraint-basee based methods do not have enough domain-specift inductive bias. Even with an arbitrarily large dataset, they are limited identifying equivalence classes that may be exponentially large (He al., 2015b). Moreover, the search over directed graphs itself may introduce unwanted bias and artifacts (cf. Colombo et al. 2014).
36th Conference Neural Information Processing Systems (NeurIPS 2022).
"Equal supervision.

---

The intractable search space ultimately imposes hard constraints on the causal structure, e.g., the node degree (Spirtes et al.. 2000), which limits the suitability of search in real-world domains.
In the present work, we propose to amortize causal structure learning. In other words, our goal is to optimize an inference model to directly predict causal structure from a provided dataset. We show that this approach allows inferring causal structure solely based on synthetic data generated by simulator of the data-generating process we are interested in. Much effort in the sciences, for example, goes into the development of realistic simulators for high-impact and yet challenging causal discovery domains, like gene regulatory networks (Schaffter et al., 201 I; Dibaeinia and Sinha, 2020), fMRI brain responses (Buxton, 2009; Bassett and Sporns, 2017). and chemical kinetics (Anderson and Kurtz, 2011; Wilkinson, 2018). Our approach based on amortized variational inference (AVICI) ultimately allows us to both specify domain-spectie inductive biases not easily represented by graph scoring functions and bypass the problems of structure search. Our model architecture is permutation in- and equivariant with respect to the observation and variable dimensions of the provided dataset, respectively, and generalizes to significantly larger problem instances than seen during training.
On synthetic data and semisynthetic gene expression data, our approach significantly outperforms existing algorithms for causal discovery, often by a large margin. Moreover, we demonstrate that our inference models induce calibrated uncertainties and robust behavior when subject to substantial distribution shifts of graphs, mechanisms, noise, and problem sizes. This suggests that our pretrained models are not only fast but also both reliable and versatile for future downstream use. In particular, AVICI was the only method to infer plausible causal structures from noisy gene expression data, advancing the frontiers of structure discovery in fields such as biology.
2 Background and Related Work
2.1 Causal Structure
In this work, we follow Mooij et al. (2016) and define the causal structure G of a set of d variables (#1...., 22) as the directed graph over whose edges represent all direct causal effects among the variables. A variable z; has direct causal effect on æy if intervening on æ; affects the outcome
of xy independent of the other variables X,6 x (x,, Jr i.e., there exists o a' such that
p(ay do(r; a, X c)) + p(r, do(z; d' Xب = c))
(I)
for some An intervention do(-) denotes any active manipulation of the generative process of x, like gene knockouts, in which the transcription rates of genes are externally set to zero. Other models such as causal Bayesian networks and structural causal models (Peters et al., 2017) are less well-suited for describing systems with feedback loops, which we consider practically relevant. However, we note that approach does not require any particular formalization of causal structure. In particular, we later show how to apply our approach when G is constrained to be acyclic. We assume causal sufficiency, i.e., that x contains all common causal parents of the variables 2; (Peters al., 2017).
2.2 Related Work
Classical methods for causal structure learning search over causal graphs and evaluate them using a likelihood or conditional independence test (Chickering, 2003; Kalisch and Bühlman, 2007: Hauser and Bühlmann, 2012: Zheng et al., 2018: Heinze-Demn Deml et al., 2018). Other methods combine constraint- and acore-based ideas (Tsamardinos et al., 2006) or use the noise properties of an SCM that is postulated to underlie the gantraenng process (Shimizu al., 2006: Hoyer al., 2008).
Deep learning has been used for causal inference, e.g., for estimating treatment effects (Shalit et al., 2017; Louizos et al., 2017; Yoon et al., 2018) and in instrumental variable analysis (Hartford et al., 2017; Bennett et al., 2019). In structure learning, neural networks have primarily been used to model nonlinear causal mechanisms (Goudet et al., 2018: Yu al., 2019; Lachapelle et al., 2020: Brouillard et al., 2020; Lorch et al., 2021) or to infer the structure of a single dataset (Zhu et al., 2020). Prior work applying amortized inference to causal discovery only studied narrowly defined subproblems such as the bivariate case (Lopez-Paz et al., 2015) and fixed causal mechanisms (Löwe et al., 2022) or used correlation coefficients for prediction (Li et al., 2020). In concurrent work, Ke et al. (2022) also frame causal discovery as supervised learning, but with significant differences. Most importantly, we optimize a variational objective under a model class that captures the symmetries of structure learning. Empirically, our models generalize to much larger problem sizes, even on realistic genomics data.

---

3 AVICI: Amortized Variational Inference for Causal Discovery
3.1 Variational Objective
To amortize causal structure learning, we define data- Gata-generating distribution p(D) that models the domain in which we infer causal structures. The observations 1) = {x* ,x,,,,,,"", x")
(D) are gener- ated by sampling from distribution over causal structures p(G) and then obtaining realizations from a data- generating mechanism p( D G). The data-generating process p(D G) characterizes all direct causal effects (1) in the system, but is not necessarily induced by ancestral sampling over a directed acyclic graph. Real-world systems are often more naturally modeled at different granularities or as dynamical systems (Mooij et al., 2013; Hoel et al., 2013; Rubenstein et al., 2017; Schölkopf, 2019).
Given a set of observations D, our goal is to approximate the posterior over causal structures p(G D) with a variational distribution g(; (). To amortize this inference task for the domain distribution D), we optimize an inference model fo to predict the variational parameters 0 by minimizing the expected forward KL divergence from the intractable posterior p(G D) to q(;; 0) for ~ p(D):
min Ep(D) DRL (p(G) D)]l((C:fa(D))
Since it is not tractable to compute the true posterior in (2), we make use of ideas by Barber and Agakov (2004) and rewrite the expected forward KL to obtain an equivalent, tractable objective:
ByDD DEL (p(G) ()
EEDD EyG D) log ((G D) log fo())
(3)
-Bp(G) ENDI (log (G; fo(D)) + const.
The constant does not depend o, so we can maximize L(Ø) Eptc D G) Glog
which allows us to perform amortized variational inference for causal discovery (AVICI). While the domain distribution p( D) Epp C) [p( D G)] can be arbitrarily complex, L is tractable whenever have access to the causal graph G underlying the generative process of D, i.e., to samples from the joint distribution p p(G, D). In practice, p(G) and p(D G) can thus be specified by simulator.
From an information- -hhoreetcc viewpoint, the objective (2) maximizes a variational lower bound on the mutual information /[;; DJ between the causal structure G and the observations D (Barber and Agakov, 2004). Starting from the definition of mutual information, we obtain
I[G; Dj H[G] H DJ [ + bl p(G D)]
(4)
2 [[] + Ey ((;;Jo(D)) HIG + L(0)
where the entropy H[G] is constant. The bound is tight if Ep(D) (p( ()) = 0.
3.2 Likelihood-Fres Inference using the Forward KL
The AVICI objective in (3) intentionally targets the forward KL (KL(P optimizing EppG.D) D) [log (G; 0)). This choice implies that we both model the density (G: 8) explic- itly and assume access to samples from the true data-generating distribution p(G. D). Minimizing the forward KL enables us to infer causal structures in arbitrarily complex domains-that is, even domains where it is difficult to specify
which requires
explicit likelihood p(D G). Moreover, the forward KL typically yields more reliable uncertainty estimates since it does not suffer from the variance underestimation problems common to the reverse KL (Bishop and Nasrabadi, 2006).
In contrast, variational inference usually optimizes the reverse KL DK [I p). which involves the reconstruction term E,(G.,) [lo GG] (Blei et al., 2017). This objective requires tractable marginal likelihood p(D |G). Unless inferring the mechanism parameters jointly (e.g. Brouillard et al. 2020; Lorch et al. 2021), this requirement limits inference to conjugate models with linear Gaussian or categorical mechanisms that assume zero measurement error (Geiger and Heckerman, 1994; Heckerman et al., 1995), which are not justified in practice (Friston et al., 2000; Schaffter et al., 2011: I; Runge et al., 2019; Dibaeinia and Sinha, 2020). Furthermore, unless the noise scale is learned jointly, likelihoods can be sensitive to the measurement scale of x (Reisach et al., 2021).
4 Inference Model
In the following section, we describe a choice for the variational distribution g(G; 0) and the inference model fa that predicts 0 given D. After that, we detail our training procedure for optimizing the model parameters o and for learning causal graphs with acyclicity constraints.
3

---

XL
(a)
(b)
(c)
Figure 1: Model architecture. (a) Our model maps the input three dimensional tensor of shape n x d x k and remains permutation in- and equivariant over axes n and d, respectively. (b) Each of the layers first self-attends over axis d and then over n, sharing parameters across the other axis. (c) The inner product of two variables' representations models the probability of a direct causal effect.
4.1 Variational Family
While any inference model that defines density is feasible for maximizing the objective in (3), we opt to use factorized variational family in this work.
GG; 8) = II p(oi5;01.) with Pi.j ~ Bern(O..) 2.5
(5)
The inference model fo maps a dataset D corresponding to n samples (o'...., to a d-by-d matrix 0 parameterizing the variational approximation of the causal graph posterior. In addition the joint observation x* = (zf,.... xyy), each sample of = (of....; ........., o²)) may contain interventional infor-
mation for each variable. When interventions or gene knockouts are performed, we set of (z).. u²) and un E {0, } indicating whether variable j was intervened upon in sample i. Other settings could be encoded analogously, e.g., when the intervention targets are unknown or measurements incomplete.
4.2 Model Architecture
To maximize statistical efficiency, fo should satisfy the symmetries inherent to the task of causal structure learning. Firstly, fo should be permutation invariant across the sample dimension (axis n). Shuffling the samples should not influence the prediction, i.c., for any permutation n, we have fa( = fa((o)). Moreover, fo should be permutation equivariant across the variable dimension (axis d). Reordering the variables should permute the predicted causal edge probabilities, (=1:)))(())x()))) (())) Lastly, f should apply to any d, 2 1.
i.e., fo( (f(x(11))) ((x( He, = fe(((o:d (for:d))
In the following, show how to parameterize fo as neural network that encodes these properties. After first mapping each 0; to real-valued vector using a position-wise linear layer, Ss operates over continuous, dinee-dimesiionl tensor of rows for the observations, columns for the variables, and feature size k. Figure illustrates the key components of the architecture.
Attending axes and The core of fo is composed of L = 8 identical layers. Each layer con- sists of four residual sublayers, where the first and third apply multi- head self-attention and the second and fourth position- wise feed-forward networks, similar the Transformer encoder (Vaswani et al., 2017). To enable information flow across all n x tokens of the representation, the model alternates in attending over the observation and the variable dimension (Kossen et al., 2021). Specifically, the first self-attention sublayer attends over axis d, treating axis n as a batch dimension; the second attends over axis n. treating axis batch dimension. Since modules are shared across non-attended axes, the representation is permutation equivariant over axes and d at all times (Lee et al., 2b).
Variational parameters After building up a representation tensor from the input using the attention layers, we max-pool over the observation axis n to obtain a representation (z)...., z²) consisting of one vector z' E R for each causal variable. Following Lorch al. (2021), we use two position
linear layers to map each z' to two embeddings u', v E Rt, which are & normalized. We then model the probability of each edge in the causal graph with an inner product:
0, 0 v b)
(6)
where a is the logistic function, a learned bias, and J positive scale that is learned in log space. Since max- max-pooling invariant to permutations and since (6) permutes with respect to axis d, fe satisfies the required permutation invariance over axis and permutation equivariance over axis d.

---

4.3 Acyclicity
Cyclic causal effects often occur, e.g., when modeling stationary distributions of dynamical systems, and thus loops in causal structure possible. However, certain domains may be more accurately modeled by acyclic structures (Rubenstein et al., 2017). While the variational family in (5) cannot en- force it, we can optimize for acyclicity through c. Whenever the acyclicity prior is justified, amend the optimization problem in (2) with the constraint that only models acyclic graphs in expectation:
F(Ø) = Ep(D) [h(f&(D)) =
(7)
The function h is zero if and only if the predicted edge probabilities induce an acyelic graph. We use the insight by Lee al. (2019a), who show that acyclicity equivalent to the spectral radius i.e., the largest absolute eigenvalue, of the predicted matrix being zero. We use power iteration to approximate and differentiate through the largest eigenvalue of (D) (Golub and Van der Vorst, 2000; Lee et al., 2019a):
a a Wb k(W) =- p(W)
+ TW ||aTaa W b - Wb [Wbl|2
where for steps:
(8)
a a
and a, b E Rd are initialized randomly. Since a few steps are sufficient in practice, (8) scales with Odd) and is significantly more efficient than O(d") constraints based on matrix powers (Zheng et al., 2018; Yu et al., 2019). We do not backpropagate gradients with respect to Φ through a, b.
4.4 Optimization
Combining the objective (3) with our inference model (5), we can directly use stochastic opti- mization to train the parameters & of the inference model. The expectations over p(G. D) inside L and FF are approximated using samples from the data- generating process of the domain. When enforcing acyclicity, causal discovery algorithms often use the augmented Lagrangian method for constrained optimization (e.g., Zheng et al. 2018; Brouillard et al. 2020). In this work, we optimize the parameters of neural network, so rely methods specifically tailored for deep learning and solve the constrained program max (()) s.t. F(Ø) = 0 through its dual formulation (Nandwani et al., 2019):
Algorithm 1 Training the inference model fo
Parameters: á variational, À dual, 7 step size while not converged do for / steps do
Ao x V (1())
x A (F))
min max L(Ø) XF(
Algorithm I summarizes the general optimization procedure for qa. which converges to a local optimum under regularity conditions on the learning rates (Jin et al., 2020). Without an acyelicity constraint, training reduces to the primal updates of 6 with X = 0.
5 Experimental Setup
Evaluating causal discovery algorithms is difficult since there are few interesting real-world datasets that come with ground-truth causal structure. Often, the believed ground truths may be incomplete or change as expert knowledge improves (Schaffter et al., 2011; Mooij al., 2020). Following prior work, we deal with this difficulty by evaluating our approach using simulated data with known causal structure and by controlling for various aspects of the task. In Appendix E, additionally report results on a real-world proteomics dataset (Sachs et al., 2005).
5.1 Domains and Simulated Components
We study three domains: two classes of structural causal models (SCMs) as well as semisynthetic single-cell expression data of gene regulatory networks (GRNs). To study the generalization of AVICI beyond the training distribution p( (D), we carefully construct a spectrum of test distributions ş D) that incur substantial shift from p(D) in terms of the causal structures, mechanisms, and noise, which we study in various combinations. Whenever we consider interventional data our experiments, half of the dataset consists of observational data and half of single-variable interventions.
Data-ginn processes p( ( We consider SCMs with linear functions (LINFAR) and with nonlinear functions of random Fourier features (RFF) that correspond to functions drawn
5

---

-
NCO, 33
-20
Wont Strogocz
Stochask a Block Model
Geameztic Random
2-20,0=5
C- 20.1-5
Eastere33, ol N
-20
-1
-0
20
-5
-12
(a)
Figure 2: Moving out-of-distribution in the RFF domain. Randomly sampled data-generating components of the nonlinear SCM domain during training p D) (top) and o.o.d. evaluation p D) (bot- tom). For visualization, the adjacency matrices (a) are topologically sorted, the causal mechanisms (b) have two parents, where and are output and length scales of the underlying GP, and the noise (c) is shown as function of one parent, where dashed lines indicate 0.66, 0.95, and 0.999 coverage.
from a Gaussian process with squared exponential kernel (Rahimi and Recht, 2007). In the out-of-distribution (o.o.d.) setting P(D). we sample the linear function and kernel parameters from the tails of p(D) and unseen value ranges. Moreover, we simulate homoscedastic Gaussian noise in the training distribution P(D) but test heteroscedastic Cauchy and Laplacian noise o.o.d. that is induced by randomly drawn, nonlinear functions hy. In LINEAR and RFF, interventions set variables to random values and are performed
subset of target variables containing half of the nodes.
In addition to SCMs, we consider the challenging domain of GRNs (GRN) using the simulator of Dibaeinia and Sinha (2020). Contrary to SCMs, gene expression samples correspond to draws from the steady state of a stochastic dynamical system that varies between cell types (Huynh- Thhu Thu and Sanguinetti, 2019). In the o.o.d. setting, the parameters sampled for the GRN simulator are drawn from significantly wider ranges. In addition, we use the noise levels of different single-cell RNA sequencing technologies, which were calibrated on real datasets. In GRN, interventions are performed on all nodes and correspond to gene knockouts, forcing the transcription rate of a variable to zero.
Causal structures p(G) Following prior work, we use random graph models and known biological networks to sample ground-truth causal structures. In three domains, the training data distribution DDD is induced by simple Erdös-Rényi and scale-free graphs (Erdōs and Rényi, 1959; Barabási and Albert, 1999). the o.o.d. setting, çi D) of the LINEAR and RFF domains are simulated using causal structures from the Watts-Strogatz model, capturing 1998); the stochastic block model, generalizing Erdōs-Rény Rényi to community structures (Holland et al., 1983); and geometric random graphs, modeling connectivity based on spatial distance (Gilbert, 1961). In the GRN domain, we use subgraphs of the known cerevisiae and E. coli GRNs and their effect signs whenever known. To extract these subgraphs, we use the procedure by Marbach et al. (2009) to maintain structural patterns like motifs and modularity (Ravasz et al., 2002; Shen- Orr et al., 2002).
phenomena (Watts and Strogatz,
To illustrate the distribution shift from p( D) to (D), Figure 2 shows set of graph, mechanism, and noise distribution samples in the RFF domain. Appendix A, we give the detailed parameter configurations and functions defining p( D) and p(D) in the three domains. We also provide details on the simulator by Dibacinia and Sinha (2020) and subgraph extraction (Marbach et al., 2009).
5.2 Evaluation Metrics
All experiments throughout this paper are conducted on datasets that AVICI has never seen during training, regardless of whether we evaluate the predictive performance in-distribution or o.o.d. To assess how well a predicted structure reflects the ground truth, we report the structural Hamming distance (SHD) and the structural intervention distance (SID) (Peters and Bühlmann, 2015). While the SHD simply reflects the graph edit distance, the SID quantifies the closeness of two graphs in terms of their interventional adjustment sets. For these metrics and for single-edge precision, recall, and F1 score, we convert the posterior probabilities predicted by AVICI to hard predictions using a threshold of 0.5. We evaluate the uncertainty estimates by computing the areas under the precision-recall curve (AUPRC) and the receiver operating characteristic (AUROC) (Friedman and Koller, 2003). How

---

Prec.
raD
Prec.
Prec.
saD
Rec.
0.0
SHD r0b0
Rec.
0.0
SHD rab0
Rec.
0.0
SHD 200
0.2
0.4
0.4
0.4
0.8
0.8
A
AUPRC
acyclic
AUPRC
% acyclic
AUPRC
acyclic
AUROC
AUROC
AUROC
In-dist
0.0.d.
o.o.d. G.
6
0.0.d. G.
0.0.0.
0.0.d G. sim. noise
10
10
10
20
20
20
100
0.4
1.0
0.0
0.6
0.4
AUPRC
AUPRC
AUPRE
A- 30
n-300
O- 30
n-300
30
n-300
(a) LINEAR
(b) REF
(c) GRN
Figure 3: Generalization properties of the inference model fa- Top row plots show performance metrics of AVICI under increasing distributional shift given n 1000 observations for = 30 variables. SIDratio is defined as SIDin-ist//ID and analogously for SHD. Thus, higher is better for all metrics. Bottom row shows the in-distribution AUPRC for various as we vary the number of observations provided to AVICI. The datasets contain interventional data (cf. Section 5.1). All values are the mean over fifteen random task instances. Error bars indicate the interquartile range.
well these uncertainty estimates are calibrated is quantified with the expected calibration error (ECH) (DeGroot and Fienberg, 1983). More details on the metrics are given in Appendix B.
5.3 Inference Model Configuration
We train three inference models overall, one for each domain, and perform all experiments these three trained models, both when predicting from only observational and from interventional data. During training, the datasets sampled from P(D) have d 2 to 50 variables and 200 samples. With probability 0.5, these training datasets contain 50 interventional samples. The inference models in the three domains share identical hyperparameters for the architecture and optimization, except for the dropout rate. We add the acyclicity constraint for the SCM domains LINEAR and RFF. Details on the optimization and architecture are given in Appendix C.
6 Experimental Results
6.1 Out-Or-Distribution Generalization
Sensitivity to distribution shift In our first set of experiments, we study the generalization capabilities of our inference models across the spectrum of test distributions described in Section 5.
We perform causal discovery from 1000 observations in systems of d 30 variables. Starting from the training distribution D)) we incrementally introduce the described distribution shifts in the causal structures, causal mechanisms, and finally noise, where fully o.o.d. corresponds to P(D). The top row of Figure 3 visualizes the results of an empirical sensitivity analysis. The radar plots disentangle how combinations of the three o.o.d. aspects, i.e., graphs, mechanisms, and noise, affect the empirical performance in the three domains LINEAR. RFF, and GRN. In addition to the metrics in Section 5.2, we also report the percentage of predicted graphs that are acyelic.
In the LINEAR domain, AVICI performs very well in all metrics and hardly suffers under distribution shift. In contrast, GRN is the most challenging problem domain and the performance degrades
7

---

Table 1: Benchmarking results (d = 30 variables). Mean SID (L) and Fl score (1) with standard error of all methods on 30 random task instances. Methods in the top section use only observational data, in the bottom section both observational and interventional data. We highlight the best result of each section and those within its 95% confidence interval according to an unequal variances t-test.
LINEAR
RFF
GRN
Algorithm
SID
F1
SID
F1
SID
F1
GES
215.6 (35.)
0.548 (0.0)
346.3 (44.4)
0.285 (0.03)
573.6 (29.2)
0.058 (0.00)
LiNGAM
413.4 (48.4)
0.369 (0.4)
410.3 (47.6)
0.238 (0.12)
617.5
(31.7)
0.044
PC
400.5 (53.7)
0.338 (0.0))
370.1 (51.2)
0.421
(0.1))
594.0
(30.0)
0.061
(0.01)
DAG-GNN
474.5 (50.8)
0.154 (0.1))
425.3
(50.2)
0.221
(0.13)
588.7
(36.6)
0.078 (0.02)
GraN-DAG
466.0 (54.3)
0,200
(0.0))
328.6
(4R.4)
0,476
(0.05)
582.4
(33.4)
0.073 (0.02)
AVICI (ours)
145.6 (21.5)
0.672
(6.04)
255.1
(48.2)
0.618
(0.06)
641.7
(M47)
0.000
GIES
120.8 (26.2)
0.736 (0.13)
304.8
(441))
0.338 (0.14)
545.5
(26.9)
0.092 00.01)
IGSP
244.0 (34.4)
0.559 (0.12)
374.1
(45.0)
0.407
(0.04)
597.4 (31.7)
0.057
(0.01)
DCDI
383.5 (45.1)
0.327
(0.03)
282.8
(46.3)
0.409
(0.04)
590.9
(30.6)
0.075
(0.02)
AVICI (ours)
110.9 (19.3)
0.819 (8.02)
192.7
(44.8)
0.707
(8.06)
416.9
(47.1)
0.338
(9.06)
more significantly for the o.o.d. scenarios. We observe that AVICI can perform better under certain distribution shifts than in-distribution, e-g., in GRN. This is because AVICI empirically performs better at predicting edges adjacent to large-degree nodes, a common feature of the E. coli and S. cerevisiae graphs not present in the Erdōs-Rényi training structures. We also find that acyclicity is perfectly satisfied for LINEAR and RFF and that AUPRC and AUROC do not suffer as much from distributional shift as the metrics based on thresholded point estimates.
In Appendix E. we additionally report results for generalization from LINEAR to REF and vice versa, i.e., to entirely unseen function classes of causal mechanisms in addition to the previous o.o.d. shifts.
Generalization to unseen problem sizes In addition to the sensitivity to distribution shift, we study the ability to generalize to unseen problem sizes. The bottom row of Figure 3 illustrates the AUPRC for the edge predictions of AVICI when varying d and n on unseen in-distribution data. The predictions improve with the number of data points n while exhibiting diminishing marginal improvement when seeing additional data, Moreover, the performance decreases smoothly as the number of variables d increases and the task becomes harder. Most importantly, this robust behavior can be observed well beyond the settings used during training (n = 200 and d < 50).
6.2 Benchmarking
Next, we benchmark AVICI against existing algorithms. Using only observational data, we compare with the PC algorithm (Spirtes al., 2000), GES (Chickering, 2003), LiNGAM (Shimizu et al., 2006), DAG-GNN (Yu et al., 2019). and GraN-DAG (Lachapelle et al., 2020). Mixed with interventional data, we compare with GIES (Hauser and Bühlmann, 2012), IGSP (Wang et al., 2017), and DCDI (Brouillard et al., 2020). We tune the important hyperparameters of each baseline on held-out task instances of each domain. When computing the evaluation metrics, we favor methods that only predict (interventional) Markov equivalence classes by orienting undirected edges correctly when present in the ground truth. Details on the baselines are given in Appendix D.
The benchmarking is performed on the fully o.o.d. domain distributions j(D). i.e., under distribution shifts on causal graphs, mechanisms, and noise distributions w.r.t. the training distribution of AVICI. Table shows the SID and Fl scores of all methods given = 1000 observations for d = 30 variables. We find that the AVICI model trained on LINEAR outperforms all baselines, both given observational interventional data, despite operating under significant distribution shift. Only GIES achieves comparable accuracy. The same holds for RFF, where GraN-DAG and DCDI perform well but ultimately do not reach the accuracy of AVICI.
In the GRN domain, where inductive biases are most difficult to specify, classical methods fail infer plausible graphs. However, provided interventional data, AVICI can use its learned inductive bias to infer plausible causal structures from the noisy gene expressions, even under distribution shift. This is a promising step towards reliable structure discovery in fields like molecular biology. Even without gene knockout data, AVICI achieves nontrival AUROC and AUPRC while classical methods predict close to randomly (Table 9 in Appendix E; see also Dibacinia and Sinha 2020; Chen and Mar 2018).
8

---

1.0
1.0
LINEAR
RIT
GRN
empirica plg. 11 (I
0.8
GES
0.031
0.068 mun:
0.092 (000
LiNGAM
0.066
0.054 mm::
0.053 (000
0.1
frequency
PC
0.036 (0.0)
0.033 mun
0.065 (0.00
0.6
DAG-GNN
0.078 (001)
0.063 m.n)
0.063 (0.00
GraN-DAG
0.046 (0.0)
0.042 (un))
0.199 (0.06)
0.4
0.01
AVICI (ours)
0.013 (U0))
0.024 (Un))
0.018 (000)
0.2
GIES
0.027 (0:00)
0.074 (0:0)
0.094 (000)
IGSP
0.042
0.083 m.o1)
0.077 (0.00
0.001
DCDI
0.068
0,087 m.n2)
0.170 (0.00)
0.0
0.2
0.4
0.6
0.8
1.0
DiBS
0.056
0.035
0.093 (0.0)
predicted p(g.. D)
AVICI (ours)
0.011
0.022 (um)
0.024
Linear
GRN
Nonparametric DAG bootstrap (Friedman et al., 1999)
(a)
(b)
Figure 4: Uncertainty calibration (d = 30 variables). As previously, datasets are held-out and o.o.d. in terms of graph, parameters, and noise. (a) Calibration plots for AVICI aggregating the predictions for ten test cases of each domain. The histograms on the right y-axis show the frequency of predictions at each confidence level. (b) ECE (J) with standard error averaged over ten test cases. Methods in the top (bottom) section use observational (and interventional) data. We highlight the best result and those within its 95% confidence interval according to unequal variances f-test.
Results for in-distribution data and for larger graphs of d = 100 variables are given in Appendices E.2 and E.3. In Appendix E.4, we also report results for real proteomics dataset (Sachs et al., 2005).
Uncertainty quantification Using metrics of calibration, we can evaluate the degree to which predicted edge probabilities are consistent with empirical edge frequencies (DeGroot and Fienberg, 1983; Guo et al., 2017). We say that predicted probability p is calibrated if we empirically observe an event in (p- 100)% of the cases. When plotting the observed edge frequencies against their predicted probabilities, a calibrated algorithm induces a diagonal line. The expected calibration error (ECE) represents the weighted average deviation from this diagonal. For further details, see Appendix B.
Since the baseline algorithms only infer point estimates of the causal structure, we use the non- parametrie DAG bootstrap to estimate edge probabilitiess (Friedman et al. 1999, Appendix D). We additionally compare AVICI with DiBS, which infers Bayesian posterior edge probabilities like AVICI (Lorch et al., 2021). Figure 4 gives the calibration plots for AVICI and Table 4b the ECE for all methods. In each domain, the marginal edge probabilities predicted by AVICI are the most calibrated in terms of ECE. Moreover, Figure 4a shows that AVICI closely traces the perfect calibration line, which highlights its accurate uncertainty calibration across the probability spectrum.
In Appendix E.5, we additionally report AUROC and AUPRC metrics for all methods. We also provide calibration plots analogous to Figure for the baselines (Figure 6). which often show vastly overconfident predictions where the calibration line is far below the diagonal.
6.3 Ablations
Finally, analyze the importance of key architecture components of the inference network fe- Focusing on the REF domain, we train several additional models and ablate single architecture compo- nents. We vary the network depth L, the axes of attention, the representation of 8, and the number of training steps for ø. All other aspects of the model, training and data simulation remain unchanged.
Table 2 summarizes the results. Most noticeably, we find that the performance drops significantly when attending only over axis and aggregating information over axis n only once through pooling after the 2L self-attention layers. Attending only over axis is not sensible since variable interactions are not processed until the prediction of 0, but we still include the results for completeness.
We also test an alternative variational parameter model given by 0. = oJ tanh (Φ_ u + o vi) that uses an additional, learned vector Φe and matrices Φa, Φ.. This model has been used in related causal discovery work for searching over high-scoring causal DAGs (Zhu et al., 2020) and is a relational network (Santoro et al., 2017). This variant also satisfies permutation equivariance (cf. Section 4.2) since it applies the same MLP elementwise to each edge pair [u', vd]. Ultimately, we find no statistically significant difference in performance to our simpler model in Eq. (6), hence we opt for less parameters and lower memory requirement.

---

Table 2: Ablations of the architecture of Sa- Models are evaluated on 100 interventional datasets of d 30 variables in the RFF domain. Top row (x) corresponds to the model used the main experiments. In (a), we vary the number of blocks T,; in (b), the axes over which attention is performed; in (c), the generative model of the variational parameters; in (d), the number of update steps of Φ. We again highlight the best result and those within its 95% t-test confidence interval.
RFF (in-dist.)
RFF (o.o.d.)
d ax. n
model
steps
SID>
AUPRC
SID
AUPRC
(x)
Eq. (6)
300k
65.2 (8A)
0,972 (008)
221.5 (247)
0.650
267.2 (22.0)
0.635 (w01)
394.2 (284)
0.242 (062)
(a)
195.9 (I8.))
0.825 (Mn1)
343.1 (27.1)
0,400 (0.3)
4
116.6 (3S
0,937 (w01)
264.0 (248)
0,566 (0,3)
351.5 (27.9)
0.552 (w01)
414.2 (9
0.209 (0.2))
(b)
V
416.8 (29.0)
0.256 (n.1))
390.2 (27.))
0.078 (0.00)
(c)
(Santoro et al., 2017)
72.4 (9:3)
0.971 (108)
225.7 (25.2)
0.634 (ua3)
(d)
100k
96.9 (11.6)
0.955 (0.00)
259.3 (26.6)
0.589
Lastly, Table shows that the causal discovery performance of AVICI scales up monotonically with respect to network depth and training time. Even substantially smaller models of L = 4 shorter training times achieve an accuracy that is on par with most baselines (cf. Table 1). Our main models (x) have moderate size of 4.2 10 parameters, which amounts to only 17.0 MB at f32 precision. Performing causal discovery (computing a forward pass) given on a trained model takes only a few seconds on CPU.
7 Discussion
We proposed AVICI, a method for inferring causal structure by performing amortized variational infer- ence over an arbitrary data-generatie distribution. Our approach leverages the insight that inductive biases crucial for statistical efficiency in structure learning might be more easily encoded in a simula- tor than in an inference technique. This is reflected in our experiments, where AVICI solves structure learning problems in complex domains intractable for existing approaches (Dibaeinia and Sinha, 2020). Our method can likely be extended to other typically difficult domains, including settings where we cannot assume causal sufficiency (Bhattacharya et al., 2021). Our approach will continually benefit from ongoing efforts in developing (conditional) generative models and domain simulators.
Using AVICI still comes with several trade-offs. First, while optimizing the dual program empirically induces acyclicity, this constraint is not satisfied with certainty using the variational family considered here. Moreover, similar to most amortization techniques (Amos, 2022), AVICI gives no theoretical guarantees of performance. Some classical methods can do so in the infinite sample limit given specific assumptions on the data- gettrggnerating process (Peters et al., 2017). However, future work might obtain guarantees for AVICI that are similar to learning theory results for the bivariate causal discovery case (Lopez-Paz et al., 2015).
Our experiments demonstrate that our inference models are highly robust to distributional shift, suggesting that the trained models could be useful out-of-the-box in causal structure learning tasks outside the domains studied in this paper. In this context, fine-tuning pretrained AVICI model on labeled real-world datasets promising avenue for future work. To facilitate this, our code and models are publicly available at: https:/ //gtttub..tttt//g///gg thub. comm/arssorcchavr
Acknowledgments and Disclosure of Funding
We thank Alexander Neitz, Giambattista Parascandolo, and Frederik Träuble for their feedback and the reviewers for their helpful comments. This research was supported by the European Research Council (ERC) under the European Union's Horizon 2020 research and innovation program grant agreement no. 815943 and the Swiss National Science Foundation under NCCR Automation, grant agreement 5INF4 180545. Jonas Rothfuss was supported by Apple Scholars in AI/ML fellowship.
10

---

References
Amos, B. (2022). Tutorial on amortized optimization for learning optimize over continuous domains. arXiv preprint arXiv:2202. 00665.
Anderson, D. F. and Kurtz, T. G. (2011). Continuous time markov chain models for chemical reaction networks. In Design and analysis of biomolecular circuits, pages 3-42. Springer.
Barabási, A.L and Albert, R. (1999). Emergence of scaling in random networks. Science, 286(5439):509 512. Barber, D. and Agakov, F. (2004). The IM algorithm: a variational approach to information maximization. Advances neural information processing systems, 16(320):201.
Bassett, D. S. and Sporns, O. (2017). Network neuroscience. Nature neuroscience, 20(3):353-30
Bennett, A., Kallus, N., and Schnabel, (2019). Deep generalized method moments for instrumental variable analysis. Advances in neural information processing systems, 32.
Bhattacharya, R., Nagarajan, T., Malinsky, D., and Shpitser, I. (2021). Differentiable causal discovery under unmeasured confounding. In International Conference Artificial Intelligence and Statistics, pages 2314-
2322. PMLR.
Bishop, C. M. and Nasrabadi, N. M. (2006). Approximate inference. In Pattern recognition and machine learning, volume 4, chapter 10. Springer.
Blei, D. M., Kucukelbir, A., and McAuliffe, J. D. (2017). Variational inference: A review for statisticians.
Journal of the American statistical Association, 12/5118)859-877
Bradbury, J., Frostig, R., Hawkins, P., Johnson. M. J., Leary. C., Maclaurin, D., Necula, G., Paszke, A., Vander-
Plas, J., Wanderman-Millee S., and Zhang, Q. (2018). JAX: composable transformations of Python+NumPP programs. http://giihub..oo/googge////
Brouillard, P., Lachapelle, S., Lacoste, A.,, Lacoste-Julien. S., and Drouin, A. (2020). Differentiable causal discovery from interventional data. Advances in Neural Information Processing Systems, 33:21865-
Buxton, R. B. (2009). Introduction functional magnetic resonance imaging: principles and techniques. Cambridge university press.
Chen. S. and Mar, J. C. (2018), Evaluating methods of inferring gene regulatory networks highlights their lack of performance for single cell gene expression data. BMC bioinformatics. 19(1):1-2) 21.
Chickering, D. M. (2003). Optimal structure identification with greedy search. J. Mach. Learn. Res., 3:507-554.
Colombo, D., Maathuis, M. H., et al. (2014), Order-independent constraint-based causal structure leamning.. J. Mach. Learn Res., 15(1):37413378)
Dawid, A. P. (2010). Beware of the DAG! In Proceedings of Workshop on Causality: Objectives and Assessment NIPS 2008, pages 59 86.
DeGroot, M. H. and Fienherg, S. E. (1983). The comparison and evaluation of forecasters. Journal of the Royal
Statistical Sociery: Series D (The Statistician), 32(1-2):
Dibacinia, P. and Sinha, S. (2020). Sergio: a single-cell expression simulator guided by gene regulatory networks. Cell Systems, 113))222--71.
Erdōs, P. and Rényi, A. (1959). On random graphs. Publicationes Mathematicae, 6:290-297.
Fawcett, T. (2004). ROC graphs: Notes and practical considerations for rescarchers. Machine learning,
31(1):1-38.
Friedman, N., Goldszmidt, M., and Wyner, A. (1999). Data analysis with bayesian networks: bootstrap approach. In Proceedings of the Fifteenth Conference on Uncertainty in Artificial Intelligence, UAI'99, page 196-205, San Francisco, CA, USA. Morgan Kaufmann Publishers Inc.
Friedman, N. and Koller, D. (2003). Being Bayesian About Network Structure. Bayesian Approach Structure Discovery Bayesian Networks. Machine Learning. 50(1):95-125
Friston, K. J., Mechelli, A., Turner, R., and Price, J. (2000). Nonlinear responses in fmri: the balloon model, volterra kernels, and other hemodynamies. Neurolmage, 1264):466-477
Geiger, D. and Heckerman, D. (1994). Learning gaussian networks. In Proceedings of the Tenth International Conference on Uncertainty in Artificial Intelligence, 9AI'94, page 235 223, San Francisco, CA, USA. Gilbert, N. (1961). Random plane networks. Journal of the society for industrial and applied mathematics. 94):5333553.
Golub, G. H. and Van der Vorst, H. A. (2000). Eigenvalue computation in the 20th century. Journal of Computational and Applied Mathematics, 123(1-2):35 65.
Goudet, O., Kalainathan, D.,. Caillou, P., Guyon. I.. Lopez- Paz, D., and Sebag, M. (2018). Causal generative neural networks. arXiv preprint arxiv:771..08936.

---

Guo. C., Pleiss, G., Sun, Y., and Weinberger, K. Q. (2017). On calibration of modern neural networks. In International conference on machine learning, pages 1321-1330. PMLR.
Hartford, J., Lewis, G., Leyton-Brown, K., and Taddy, M. (2017). Deep IV: A flexible approach for counterfactual prediction. In International Conference on Machine Learning, pages 1414-1423. PMLR.
Hauser, A. and Bühlmann, (2012). Characterization and greedy learning of interventional markov equivalence classes of directed acyclic graphs. The Journal of Machine Learning Research, 13(1):2409-244 2464.
He. K., Zhang. X., Ren, S., and Sun. (2015a). Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In Proceedings of the IEEE international conference on computer vision, pages 1026-1034.
He, Y., Jia, J., and Yu, B. (2015b). Counting and exploring sizes of markov equivalence classes of directed acyclic graphs. The Journal of Machine Learning Research, 16( ((()))99)-260)
Heckerman, D., Geiger, D., and Chickering, D. M. (1995). Learning Bayesian Networks: The Combination of Knowledge and Statistical Data. Machine Learning, 20(3):17-2433,
Heinze-Demn C., Maathuis, M. H., and Meinshausen, N. (2018). Causal structure learning. Annual Review of Statistics and Its Application, 5:371-391.
Hennigan, T., Cai, T., Norman, T., and Babuschkin, I. (2020).
Haiku:
Sonnet for JAX.
Hoel, E. P., Albantakis, L., and Tononi, G. (2013). Quantifying causal emergence shows that macro can beat micro. Proceedingy the National Academy of Sciences,
Holland, P. W., Laskey, K. B., and Leinhardt, S. (1983). Stochastic blockmodels: First steps. Social networks, 5(2):109-137.
Hoyer, P., Janzing, D., Mooij, J. M., Peters, J., and Schölkopf, B. (2008). Nonlinear causal discovery with additive noise models. Advances in neural information processing systems, 21.
Thunhh-Thu, V. and Sanguinetti, G. (2019). Gene regulatory network inference: introductory survey. In Gene Regulatory Networks, pages 1-23. Springer.
Jin, C., Netrapalli, P., and Jordan, M. (2020). What is local optimality in nonconvex-nouu nonconcave minimax optimization? In International conference machine learning, pages 4880 4889. PMLR.
Kalainathan, D... Goudet, O.. and Dutta, R. (2020). Causal discovery toolbox: Uncovering causal relationships in
python. J. Mach. Learn. Res., 21:37-1.
Kalisch, M. and Bühlman, P. (2007). Estimating high-dimensional directed acyclic graphs with the pe-algorithm. Journal of Machine Learning Research, 8(3).
Ke, N. R., Chiappa. S., Wang. J., Bornschein, Weber, T., Goyal, A., Botvinic, M., Mozer, M., and Rezende, J. (2022). Learning induce causal structure. arXiv:2204. 04875.
Kossen, J., Band, N., Lyle, C., Gomez, A. N., Rainforth, T., and Gal, Y. (2021). Self-attention between datapoints:
Going beyond individual input-output pairs in deep learning. Advances in Neural Information Processing Systems, 34:28742-288 287566.
Lachapelle, S., Brouillard, P., Deleu, T., and Lacoste-Julien, S. (2020). Gradient-based neural dag learning. In International Conference on Learning Representations.
Loe, H.-C., Danieletto, M., Miotto, R., Cherng, S. T., and Dudley, J. T. (2019a). Scaling structural learning with no-bears to infer causal transcriptome networks. In PACIFIC SYMPOSIUM ON BIOCOMPUTING 2020, pages 391-402 World Scientific.
Lee, J., Lee, Y., Kim, J., Kosiorek, A., Choi, S., and Teh, W. (2019b). Set transformer: framework for attention-based permutation-invariant neural networks. In International Conference on Machine Learning, pages 3744-3753. PMLR.
Li, H., Xiao, Q. and Tian, J. (2020). Supervised whole dag causal discovery. arXiv preprint arXiv:2006.044 04697.
Lopez-Paz, D., Muandet, K., Schölkopf, B., and Tolstikhin, I. (2015). Towards a learning theory of cause-effect inference. In International Conference on Machine Learning, pages 1452-1461. PMLR.
Lorch, L., Rothfuss, J., Schölkopf, B., and Krause, A. (2021). DiBS: Differentiable Bayesian structure learning. Advances in Neural Information Processing Systems, 34.
Louizos, C., Shalit, U.. Mooij. J. M., Sontag, D.. Zemel, R., and Welling, M. (2017). Causal effect inference with deep latint-variable models. Advances in neural information processing systems, 30.
Löwe, S., Madras, D., Zemel, R., and Welling, M. (2022). Amortized causal discovery: Leaming infer causal graphs from time-series data, arXiv preprint arXiv:2006. 10833, 140:1 1-4..
Marbach, D., Schaffter, T., Mattiussi, C., and Floreano, D. (2009). Generating realistic in silico gene networks for performance assessment of reverse engineering methods. Journal of computational biology. 16(2):229-22 239.
12

---

Mooij, J. M.. Janzing, D., and Schölkopf. B. (2013). From ordinary differential equations to structural causal models: The deterministic case. In Proceedings of the Twenty-Nimth Conference Uncertainty in Artificial Intelligence, 13 page 440-448, Arlington, Virginia, USA. AUAI Press.
Mooij, J. M. Magliacane. S., and Claassen, T. (2020). Joint causal inference from multiple contexts. Journal of Machine Learning Research, 21(99):1-108 108.
Mooij. J. M., Peters, J., Janzing, D., Zscheischler, J., and Schölkopt, B. (2016). Distinguishing cause from effect using observational data: methods and benchmarks. The Journal of Machine Learning Research, 17(1):1033-120
Nandwani, Y. Pathak, A., and Singla, P. (2019), A primal dual formulation for deep learning with constraints. Advances in Neural Information Processing Systems, 32.
Pascanu, R., Mikolov, T., and Bengio, Y. (2013). On the difficulty of training recurrent neural networks. In International conference on machine learning, pages 1310-1318, PMI.R.
Pearl, J. (2009). Causaliry. Cambridge university press.
Peters, J. and Bühlmann, P. (2014). Identifiability of gaussian structural equation models with equal error variances. Biometrika, (1:211219-
Peters, J. and Bühlmann, P. (2015). Structural intervention distance for evaluating causal graphs. Neural computation, 27(3)771-799.
Peters, J., Janzing, D., and Schölkopf, B. (2017). Elements of causal inference: foundations and learning algorithmss. The MIT Press.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I,, et al. (2019). Language models are unsupervised multitask learners. OpenA/ blog, 1(8)9
Rahimi, A. and Recht, B. (2007). Random features for large-scale kernel machines. Advances neural information processing systems, 20.
Ravasz, E., Somera, L., Mongru, D. A., Oltvai, Z. N., and Barabási, A.-L. (2002). Hierarchical organization of modularity in metabolic networks. science, 297(5586)
Reisach, A. G., Seiler, C., and Weichwald, S. (2021). Beware of the simulated DAG! varsortability in additive noise models. Advances in Neural Information Processing Systems.
Robinson, M. D., McCarthy, D. J.. and Smyth, G. K. (2010). edger: a bioconductor package for differential expression analysis of digital gene expression data. Bioinformatics, 26( 139-1444
Rubenstein, P. K., Weichwald. S., Bongers, S., Mooij. M., Janzing, D., Grosse Weenssu. M., and Schölkopf, B. (2017). Causal consistency of structural equation models. In Proceedings of the 33rd Conference on Uncertainty in Artificial Intelligence (UAI), page ID 11.
Runge, J. Bathiany, S., Bollt, E., Camps- Valls, G., Coumou, D., Deyle, E., Glymour, C., Kretschmer, M.. Mahecha, M. D., Muñoz-Mari, J., et al. (2019). Inferring causation from time series earth system sciences. Nature communccations, 1033)
Sachs, K., Perez, O., Pe'er, D., Lauffenburger, D. A., and Nolan, (2005). Causal pooeii-siggnalin networks derived from multiparameter single-cell data. Science, 308(5721):523-59,,
Santoro, A., Raposo, D., Barrett, D. G., Malinowski, M., Pascanu, R., Battaglia, P., and Lillicrap. T. (2017). A simple neural network module for relational reasoning. Advances neural information processing systems, 30.
Schaffter, T., Marbach, D., and Floreano, D. (2011). Genenetweaver: in silico benchmark generation and performance profiling of network inference methods. Bioinformaticx, 27(16)2263-22770
Scheines, R. and Ramsey, J. (2016). Measurement error and causal discovery. In CEUR workshop proceedings, volume 1792 page 1. NIH Public Access.
Schölkopf. B. (2019). Causality for machine learning. arXiv preprint arXiv:1911.10500.
Shalit, U., Johansson, F. D., and Sontag, D. (2017). Estimating individual treatment effect: generalization bounds and algorithms. In International Conference on Machine Learning, pages 3076-3085, PMLR.
Shen-Orr, S. S., Milo, R., Mangan, S., and Alon, U. (2002). Network motif's in the transcriptional regulation network of escherichia coli. Nature genetics, 31(1):64-66 68.
Shimizu, S., Hoyer, P. O., Hyvärinen, A., Kerminen, A., and Jordan, M. (2006). A linear non-gaussian acyelic model for causal discovery. Journal of Machine Learning Research, 7(10).
Spirtes, P., Glymour, C. N., and Scheines, R. (2000), Causation, prediction, and search. Adaptive computation and machine learning. MIT Press, Cambridge, Mass, 2nd ed edition.
Squires. C. Belyaeva, A., Karnik, S., Saeed, B., Jablonski, K. P., and Uhler, C. (2018). Causaldag. https://ghhubb.oomuueerralauuuss
13

---

Tsamardinos, I., Brown, L. E., and Aliferis, C. F. (2006). The max-min hill-climbing Bayesian network structure learning algorithm. Machine learning, 65(1):31-78.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. (2017). Attention is all you need. Advances in neural information processing systems, 30.
Wang. Y. Solus, L.. Yang. K., Uhler. (2017). Permutation- based causal inference algorithms with interventions. Advances Neural Information Processing Systems, 30.
Watts, J. and Strogatz, S. H. (1998), Collective dynamics of 'small-worde networks, nature, 363(6)4)) 442.
Wilkinson, D. J. (2018). Stechastic modelling for systems biology. Chapman and Hall/CRC.
Yoon, J., Jordon.
Schaar, M. (2018). Ganite: Estimation of individualized treatment effects using generative adversarial nets. International Conference on Learning Representations.
You, Y., Li, J. Reddi, S., Hseu, J., Kumar, S., Bhojanapalli, S., Song. X., Demmel, J., Keutzer, K., and Hsich, C.-. (2019). Large batch optimization for deep learning: Training bert in 76 minutes. arXiv preprint WwXiv:/94400090
Yu, Y., Chen, J. Gao, T., and Yu, M. (2019). DAG-GNN: DAG structure learning with graph neural networks. In Proceedings the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 7154-7163, PMLR.
Zhang, K., Gong, M., Ramsey, J., Batmanghelich, K., Spirtes, P., and Glymour, C. (2017). Causal discovery in Identifiability conditions. arXiv preprint arXiv:1706.03766
the presence of measurement
Zheng, X., Aragam, B., Ravikumar, P. K., and Xing, E. P. (2018). Dags with no tears: Continuous optimization for structure learning. Advances in Neural Information Processing Systems, 31.
Zhu, S., Ng. I., and Chen, (2020). Causal discovery with reinforcement learning. International Conference on Learning Representations.
14

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] See in particular the discussion section for the trade-offss with using our approach versus other causal structure learning
algorithms.
(c) Did you discuss any potential negative societal impacts of your work? [Yes] We see our contribution as purely methodological contribution to causal structure learning that is unlikely to have any direct negative societal impacts.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full of assumptions of all theoretical results? [N/A]
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)? [Yes] We provide a repository URL with code and instructions for reproducing the results.
(b) Did you specify all the training details (c.g.. data splits, hyperparameters, how they were chosen)? [Yes] We give an overview of the experimental setup in the main body of the paper. Full technical details for reproducing the results are in the appendix.
(c) Did you report error bars (c.g., with respect to the random seed after running experi- ments multiple times)? [Yes] We report error bars for all experiments. For method, results are reported for a fixed model so error bars capture only noise due to randomness across tasks, since retraining is too computationally expensive.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] See Appendix E.
4. If you are using existing assets (e.g., code, data, models) or curating/rellusing new assets...
(a) your work uses existing assets, did you cite the creators? [Yes] We use simulator by Dibaeinia and Sinha (2020) to train and evaluate our models (GNU General Public License v3.0). For approaches that we compare our method against, we use code assets that implement those approaches: Causal Discovery Toolbox (MIT Licence) (Kalainathan al., 2020), Brouillard et al. (2020) (MIT License), Lachapelle et al. (2020) (MIT License), Yu et al. (2019) (Apache License 2.0), the CausalDAG package (Squires et al., 2018) (3-Clause BSD license). The GRN graphs used in experiments from Schaffter et al. (201 1) (MIT License). We also use the data of Sachs et al. (2005).
(b) Did you mention the license of the assets? [Yes] Listed above.
(c) Did you include any new assets either in the supplemental material as URL? [Yes] The code to reproduce our experiments is given the provided repository.
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
15


Reviews:
[{"review_id": "mzG6eFj84Va", "reviewer": "Reviewer_imor", "summary": "In this paper, the authors aim to address the problem of causal discovery. Considering that the existing methods for causal discovery are hard to specify realistic inductive biases and require different assumptions, the authors propose an amortized variational inference model for causal discovery, which amortizes the process of causal structure learning and approximates the posterior over causal graphs given a data-generating distribution. The authors evaluate the performance of the proposed method on the synthetic data 13 and semi-synthetic gene expression data. However, there are some concerns as follows:", "questions": "Please refer 'Strengths And Weaknesses'", "limitations": "Please refer 'Strengths And Weaknesses'", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "\n1.\tOne of the main concerns is that the authors do not provide sufficient theoretical validation, such as in the form of identifiability or consistency proofs. Identifiability is an important problem for causal discovery from observational data, and the authors seem to ignore this problem. \n2.\tThe main contribution of this paper is to leverage the variational inference to discover the causal structure. However, similar ideas have been proposed by other researchers like [1]. The main difference between the proposed method and [2] is that [2] is used for time-series data.  Therefore, the contributions of these are limited. \n3.\tThe authors assume causal sufficiency, so all the data should be extracted from the same distribution with one causal mechanism. And the domain variables that result in the circumstance of domain-shift or out-of-distribution [2], however, the authors consider the OOD generalization and cannot observe the domain variables. Hence, these experiments seem to be not reasonable. \n4.\tTechnologically, since the training data are extracted from one distribution, so only one causal structure can be learned. However, since the authors employ mini-batch optimization to learn the structure. What if different structures are extracted from different samples? \n5.\tFinally, it is suggested that the authors should consider more latest baselines. \n\n\n[1] Amortized Causal Discovery: Learning to Infer Causal Graphs from Time-Series Data \n[2] domain adaptation as a problem of inference on graphical model \n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "O6j38FqYIxo", "reviewer": "Reviewer_5Lgh", "summary": "The paper introduces an amortized inference approach to learn causal structures from observational/interventional data, by leveraging variational inference method and permutation invariant/equivarient models. Experiment results on synthetic and semi-synthetic data are provided to show that the proposed approach exhibits generalization capabilities under distribution shifts.", "questions": "The suggestions and questions for the authors are outlined in my comment to the previous question.", "limitations": "The authors have adequately addressed the limitations and potential negative societal impact of their work.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Pros:\n- The proposed method is sound and novel.\n- The problem studied is interesting and relevant to the field of causal discovery, which helps bridge the gap between causality and machine learning.\n- The empirical studies are conducted across different data generating processes, number of variables, graph types, etc.\n\nCons:\n- The proposed idea of amortized inference for multivariate causal discovery is not completely novel. It has been explored by existing work \"Supervised Whole DAG Causal Discovery\", which should be discussed in the introduction and compared in the experiments.\n- The model architecture in Sec. 4.2 appears to be very similar to the encoder-decoder architecture for causal graph generation in paper \"Causal Discovery with Reinforcement Learning\", and should be discussed/compared in the paper. The latter also applies a transformer encoder that attends over axis $d$ and a similar decoder to generate graph. The only difference appears to be that the latter does not apply attention over axis $n$.\n- It would be great to provide some ablation studies on the attention over axes $n$ and $d$ to study the contribution of each module.\n- The current OOD experiments change only the linear function weights, kernel scales, and noise distribution (L224-233). I believe that the ability to generalize to a completely different data-generating process is also important, e.g., train distribution contains RFF dataset while test distribution contains MLP dataset. Another aspect is to generalize from synthetic data to real data, e.g., train distribution contains RFF dataset while test dataset is a real dataset (such as the one by \"Causal Protein-Signaling Networks Derived from Multiparameter Single-Cell Data\").\n- Could the authors explain how users should apply the learned amortized model in practice? Will there be a separate learned model for different data-generating processes like linear, RFF, MLP, and GRN? That means users have to manually train/pick an amortized model for a dataset of interest based on their understanding of the dataset. Could we combine all different data-generating processes into a large dataset and train a single, more powerful, amortized model? This is what I might expect from a amortized/supervised causal discovery method.\n- Stronger baselines:\n    - DAG-GNN is only applicable for a limited type of nonlinear data-generating process (and not RFF). I suggest the authors to include GRAN-DAG (\"Gradient-Based Neural DAG Learning\") with preliminary neighborhood selection and pruning as a baseline.\n    - The authors should also include non-parametric scoring criterion and test for GES and PC.\n\nMinor:\n- It was surprising to see that LiNGAM does not perform well on the linear non-Gaussian synthetic data. Have the authors considered the improved method in \"DirectLiNGAM: A Direct Method for Learning a Linear Non-Gaussian Structural Equation Model\"?\n- The proposed method cannot ensure convergence to acyclic solution, which might be important in some cases. Could the authors explain their choice of dual ascent to solve the acyclicity-constrained problem, instead of penalty/augmented Lagrangian method (which is more widely adopted by existing works on differentiable causal discovery)?\n\nPost-rebuttal:\n\nI have read the rebuttals and further responses from the authors. I think the paper still lacks a number of experiments:\n- Systematic ablation study (this is particularly important in my opinion)\n- Comparison with existing decoders in Zhu et al., 2020\n- Benchmarks for in-distribution setting, i.e., homogeneous noise (the proposed method should ideally perform well as compared to the baselines for both in-distribution and out-distribution settings)\n- Stronger baselines like GRAN-DAG and GES using nonparametric scoring criterion\n\nGiven the limited time of rebuttal period, it is understandable that it is difficult to finish conducting all experiments, and the authors have promised to include them in the final version. Therefore, my concerns have been mostly addressed and I have increased my score to 6.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "MOdGs7oPwLj", "reviewer": "Reviewer_My1M", "summary": "This paper proposes a method for learning a causal graph from data, by training a neural network on similar data from a simulator that generates both random graphs and data from those graphs, so that the model learns to predict graph structures given data. Data points can be marked as observational or interventional in arbitrary ways.", "questions": "1. (Section 2.1) I wasn't sure at first what definition of causal structure you meant here. But having read the paper, I suppose it doesn't really matter: whatever the relation between graph and data might be, if you have a simulator that captures this relation, AVICI will learn how to recover graphs given datasets. Or are there assumptions I'm overlooking here?\n\n2. (Section 2.2) The sentence on lines 80-82, comparing AVICI to work leveraging multi-domain data, came across as odd. The methods just have different (orthogonal?) goals, but here it sounds like those other methods are weakened by some assumption. (\"Orthogonal\", because maybe if AVICI or a method like it is fed multi-domain datasets, it could learn to do the same task.)\n\n3. (Figure 3) The SID- and SHD-ratios have a counterintuitive definition, and furthermore are apparently usually larger than 1, often by quite a bit it seems from the direction of the red and green lines at F1. Is there a better way to visualise these results?\n\nMinor points:\n* below (1): \"for any\" -> \"for some\" ($\\exists$, not $\\forall$)\n", "limitations": "Yes", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The method proposed seems extremely powerful and general, both judging by its design and by its empirical performance. I expect it will have a very high impact.\n\nThe paper is also very clearly written.\n\nI am not familiar enough with related work on amortized variational inference to assess the originality of this approach.\n\nUPDATE: I learned from the other reviewers that this submission is more closely related to existing work than I thought. I have reduced my score (as well as my confidence) accordingly.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept