Amortized Inference for Causal Structure Learning

Introduction

Learning the causal structure among a set of variables is a fundamental task in various scientific disciplines (Spirtes et al., , 2000; Pearl, , 2009). However, inferring this causal structure from observations of the variables is a 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 a graph scoring function or testing for conditional independences (Heinze-Deml et al., , 2018).

Specifying realistic inductive biases is universally difficult for existing approaches to causal discovery. Score-based methods use strong assumptions about the data-generating process, such as linearity (Shimizu et al., , 2006), specific noise models (Hoyer et al., , 2008; Peters and Bühlmann, , 2014), and the absence of measurement error (cf. Scheines and Ramsey, 2016; Zhang et al., 2017), which are difficult to verify (Dawid, , 2010; Reisach et al., , 2021). Conversely, constraint-based methods do not have enough domain-specific inductive bias. Even with an arbitrarily large dataset, they are limited to identifying equivalence classes that may be exponentially large (He et al., 2015b, ). Moreover, the search over directed graphs itself may introduce unwanted bias and artifacts (cf. Colombo et al., 2014). 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 a causal structure from a provided dataset. We show that this approach allows inferring causal structure solely based on synthetic data generated by a 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., , 2011; 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-specific 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.

Background and Related Work

In this work, we follow Mooij et al., (2016) and define the causal structure $G$ of a set of $d$ variables $\mathbf{x}$ $=$ $(x_{1},\dots,x_{d})$ as the directed graph over $\mathbf{x}$ whose edges represent all direct causal effects among the variables. A variable $x_{i}$ has a direct causal effect on $x_{j}$ if intervening on $x_{i}$ affects the outcome of $x_{j}$ independent of the other variables $\smash{\mathbf{x}_{\backslash ij}:=\mathbf{x}\,\backslash\{x_{i},x_{j}\}}$, i.e., there exists $a\neq a^{\prime}$ such that

for some $\mathbf{c}$. An intervention $\smash{\text{do}(\cdot)}$ denotes any active manipulation of the generative process of $\mathbf{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 our 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 $\mathbf{x}$ contains all common causal parents of the variables $x_{i}$ (Peters et al., , 2017).

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-Deml et al., , 2018). Other methods combine constraint- and score-based ideas (Tsamardinos et al., , 2006) or use the noise properties of an SCM that is postulated to underlie the data-generating process (Shimizu et al., , 2006; Hoyer et 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 et 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.

AVICI: Amortized Variational Inference for Causal Discovery

To amortize causal structure learning, we define a data-generating distribution $p(D)$ that models the domain in which we infer causal structures. The observations $D$ $=$ $\{\mathbf{x}^{1},\dots,\mathbf{x}^{n}\}$ $\sim$ $p(D)$ are generated by sampling from a 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 it 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 $q(G;\theta)$. To amortize this inference task for the domain distribution $p(D)$, we optimize an inference model $f_{\phi}$ to predict the variational parameters $\theta$ by minimizing the expected forward KL divergence from the intractable posterior $p(G\,|\,D)$ to $q(G;\theta)$ for $D$ $\sim$ $p(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:

From an information-theoretic viewpoint, the objective (2) maximizes a variational lower bound on the mutual information $I[G;D]$ between the causal structure $G$ and the observations $D$ (Barber and Agakov, , 2004). Starting from the definition of mutual information, we obtain

2 Likelihood-Free Inference using the Forward KL

Inference Model

In the following section, we describe a choice for the variational distribution $q(G;\theta)$ and the inference model $f_{\phi}$ that predicts $\theta$ given $D$. After that, we detail our training procedure for optimizing the model parameters $\phi$ and for learning causal graphs with acyclicity constraints.

While any inference model that defines a density is feasible for maximizing the objective in (3), we opt to use a factorized variational family in this work.

The inference model $f_{\phi}$ maps a dataset $D$ corresponding to $n$ samples $\smash{\{\bm{o}^{1},\dots,\bm{o}^{n}\}}$ to a $d$-by-$d$ matrix $\theta$ parameterizing the variational approximation of the causal graph posterior. In addition to the joint observation $\mathbf{x}^{i}=(x^{i}_{1},\dots,x^{i}_{d})$, each sample $\smash{\bm{o}^{i}=(o^{i}_{1},\dots,o^{i}_{d})}$ may contain interventional information for each variable. When interventions or gene knockouts are performed, we set $\smash{o^{i}_{j}=(x^{i}_{j},u^{i}_{j})}$ and $\smash{u^{i}_{j}\in\{0,1\}}$ 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.

2 Model Architecture

To maximize statistical efficiency, $f_{\phi}$ should satisfy the symmetries inherent to the task of causal structure learning. Firstly, $f_{\phi}$ should be permutation invariant across the sample dimension (axis $n$). Shuffling the samples should not influence the prediction, i.e., for any permutation $\pi$, we have $f_{\phi}(\pi(\{\bm{o}\}))=f_{\phi}(\{\bm{o}\})$. Moreover, $f_{\phi}$ should be permutation equivariant across the variable dimension (axis $d$). Reordering the variables should permute the predicted causal edge probabilities, i.e., $f_{\phi}(\{\bm{o}_{\pi(1:d)}\})_{i,j}=f_{\phi}(\{\bm{o}_{1:d}\})_{\pi(i),\pi(j)}$. Lastly, $f_{\phi}$ should apply to any $d,n\geq 1$.

In the following, we show how to parameterize $f_{\phi}$ as a neural network that encodes these properties. After first mapping each $\smash{o^{i}_{j}}$ to a real-valued vector using a position-wise linear layer, $f_{\phi}$ operates over a continuous, three-dimensional tensor of $n$ rows for the observations, $d$ columns for the variables, and feature size $k$. Figure 1 illustrates the key components of the architecture.

Attending over axes $d$ and $n$ The core of $f_{\phi}$ is composed of $L=8$ identical layers. Each layer consists 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 to the Transformer encoder (Vaswani et al., , 2017). To enable information flow across all $n\times d$ 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 $d$ as a batch dimension. Since modules are shared across non-attended axes, the representation is permutation equivariant over axes $n$ and $d$ at all times (Lee et al., 2019b, ).

where $\sigma$ is the logistic function, $b$ a learned bias, and $\tau$ a positive scale that is learned in $\log$ space. Since max-pooling is invariant to permutations and since (6) permutes with respect to axis $d$, $f_{\phi}$ satisfies the required permutation invariance over axis $n$ and permutation equivariance over axis $d$.

3 Acyclicity

Cyclic causal effects often occur, e.g., when modeling stationary distributions of dynamical systems, and thus loops in a causal structure are possible. However, certain domains may be more accurately modeled by acyclic structures (Rubenstein et al., , 2017). While the variational family in (5) cannot enforce it, we can optimize for acyclicity through $\phi$. Whenever the acyclicity prior is justified, we amend the optimization problem in (2) with the constraint that $q$ only models acyclic graphs in expectation:

The function $h$ is zero if and only if the predicted edge probabilities induce an acyclic graph. We use the insight by Lee et al., 2019a , who show that acyclicity is equivalent to the spectral radius $\rho$, 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 $f_{\phi}(D)$ (Golub and Van der Vorst, , 2000; Lee et al., 2019a, ):

4 Optimization

Combining the objective in (3) with our inference model (5), we can directly use stochastic optimization to train the parameters $\phi$ of the inference model. The expectations over $p(G,D)$ inside $\mathcal{L}$ and $\mathcal{F}$ 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 $\phi$ of a neural network, so we rely on methods specifically tailored for deep learning and solve the constrained program $\max_{\phi}\mathcal{L}(\phi)$ s.t. $\mathcal{F}(\phi)=0$ through its dual formulation (Nandwani et al., , 2019):

Algorithm 1 summarizes the general optimization procedure for $q_{\phi}$, which converges to a local optimum under regularity conditions on the learning rates (Jin et al., , 2020). Without an acyclicity constraint, training reduces to the primal updates of $\phi$ with $\lambda=0$.

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 et 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, we additionally report results on a real-world proteomics dataset (Sachs et al., , 2005).

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-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.

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 well these uncertainty estimates are calibrated is quantified with the expected calibration error (ECE) (DeGroot and Fienberg, , 1983). More details on the metrics are given in Appendix B.

3 Inference Model Configuration

We train three inference models overall, one for each domain, and perform all experiments on 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 $n=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.

Experimental Results

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 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.1, we additionally report results for generalization from Linear to Rff 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 5 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\leq 50$).

2 Benchmarking

Next, we benchmark AVICI against existing algorithms. Using only observational data, we compare with the PC algorithm (Spirtes et 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.

In the Grn domain, where inductive biases are most difficult to specify, classical methods fail to 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 Dibaeinia and Sinha, 2020; Chen and Mar, 2018).

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 a real proteomics dataset (Sachs et al., , 2005).

Linear Rff Grn GES∗ 0.031 (0.00) 0.068 (0.02) 0.092 (0.01) LiNGAM∗ 0.066 (0.02) 0.054 (0.01) 0.053 (0.01) PC∗ 0.036 (0.00) 0.033 (0.01) 0.065 (0.01) DAG-GNN∗ 0.078 (0.01) 0.063 (0.01) 0.063 (0.01) GraN-DAG∗ 0.046 (0.01) 0.042 (0.01) 0.199 (0.05) AVICI (ours) 0.013 (0.00) 0.024 (0.01) 0.018 (0.00) GIES∗ 0.027 (0.00) 0.074 (0.02) 0.094 (0.01) IGSP∗ 0.042 (0.01) 0.083 (0.01) 0.077 (0.01) DCDI∗ 0.068 (0.01) 0.087 (0.02) 0.170 (0.03) DiBS 0.056 (0.02) 0.035 (0.01) 0.093 (0.01) AVICI (ours) 0.011 (0.00) 0.022 (0.01) 0.024 (0.01) $*$ Nonparametric DAG bootstrap (Friedman et al., , 1999) (b)

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 a predicted probability $\smash{p}$ is calibrated if we empirically observe an event in $\smash{(p\cdot 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 nonparametric DAG bootstrap to estimate edge probabilities (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 5 gives the calibration plots for AVICI and Table 5b 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 5a 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 5 for the baselines (Figure 8), which often show vastly overconfident predictions where the calibration line is far below the diagonal.

3 Ablations

Finally, we analyze the importance of key architecture components of the inference network $\smash{f_{\phi}}$. Focusing on the Rff domain, we train several additional models and ablate single architecture components. We vary the network depth $L$, the axes of attention, the representation of $\theta$, and the number of training steps for $\phi$. 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 $d$ and aggregating information over axis $n$ only once through pooling after the $2L$ self-attention layers. Attending only over axis $n$ is not sensible since variable interactions are not processed until the prediction of $\theta$, but we still include the results for completeness.

We also test an alternative variational parameter model given by $\smash{\theta_{i,j}=\phi^{\top}_{\theta}\tanh\left(\phi_{u}^{\top}\mathbf{u}^{i}+\phi_{v}^{\top}\mathbf{v}^{j}\right)}$ that uses an additional, learned vector $\smash{\phi_{\theta}}$ and matrices $\smash{\phi_{u},\phi_{v}}$. 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 $[\mathbf{u}^{i},\mathbf{v}^{j}]$. Ultimately, we find no statistically significant difference in performance to our simpler model in Eq. (6), hence we opt for less parameters and a lower memory requirement.

Lastly, Table 2 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$ or shorter training times achieve an accuracy that is on par with most baselines (cf. Table 1). Our main models ($\star$) have a moderate size of $4.2$ $\times$ $10^{6}$ parameters, which amounts to only $17.0$ MB at f$32$ precision. Performing causal discovery (computing a forward pass) given on a trained model takes only a few seconds on CPU.

Discussion

We proposed AVICI, a method for inferring causal structure by performing amortized variational inference over an arbitrary data-generating distribution. Our approach leverages the insight that inductive biases crucial for statistical efficiency in structure learning might be more easily encoded in a simulator 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-generating 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 a pretrained AVICI model on labeled real-world datasets is a promising avenue for future work. To facilitate this, our code and models are publicly available at: https://github.com/larslorch/avici.

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 51NF40 180545. Jonas Rothfuss was supported by an Apple Scholars in AI/ML fellowship.

References

Appendix A Domain Specification and Simulation

In Erdős-Rényi graphs, each edge is sampled independently with a fixed probability (Erdős and Rényi, , 1959). We scale this probability to obtain $O(d)$ edges in expectation. Scale-free graphs are generated by a sequential preferential attachment process, where in- or outgoing edges of node $i$ to the previous $i-1$ nodes are sampled with probability $\propto$ $\text{deg}(j)^{\alpha}$ (Barabási and Albert, , 1999). Watts-Strogatz graphs are $k$-dimensional lattices, whose edges get rewired globally to random nodes with a specified probability (Watts and Strogatz, , 1998). The stochastic block model generalizes Erdős-Rényi to capture community structure. Splitting the nodes into a random partition of so-called blocks, the inter-block edge probability is dampened by a multiplying factor compared to the intra-block probability, also tuned to result in $O(d)$ edges in expectation (Holland et al., , 1983). Lastly, geometric random graphs model connectivity based on two-dimensional Euclidian distance within some radius, where nodes are randomly placed inside the unit square (Gilbert, , 1961).

For undirected random graph models, we orient edges by selecting the upper-triangular half of the adjacency matrix. The classes of random graph models are sampled in equal proportion when generating a set of evaluation datasets (Tables 3 and 4).

A.1.2 Subgraph Extraction from Real-World Networks

For the evaluation in the Grn domain, we sample realistic causal graphs by extracting subgraphs from the known E. coli and S. cerevisiae regulatory networks. For this, we rely on the procedure by Marbach et al., (2009), which is also used by Schaffter et al., (2011) and Dibaeinia and Sinha, (2020). Their graph extraction method is carefully designed to capture the structural properties of biological networks by preserving the functional and structural properties of the source network.

The procedure extracts a random subgraph of the source network by selecting a subset of nodes $\mathcal{V}$, and then returning the graph containing all edges from the source network covered by $\mathcal{V}$. Starting from a random seed node, the algorithm proceeds by iteratively adding new nodes to $\mathcal{V}$. In each step, this new node is selected from the set of neighbors of the current set $\mathcal{V}$. The neighbor to be added is selected greedily such that the resulting subgraph has maximum modularity (Marbach et al., , 2009).

To introduce additional randomness, Marbach et al., (2009) propose to randomly draw the new node from the set of neighbors inducing the top-$p$ percent of the most modular graphs. In our experiments, we adopt the latter with $p=20$ percent, similar to Schaffter et al., (2011). The original method of Marbach et al., (2009) is intended for undirected graphs. Thus, we use the undirected skeleton of the source network for the required modularity and neighborhood computation.

Empirically, individual genes in E. coli tend to predominantly have either up- or down-regulating effects on their causal children. To capture this aspect in S. cerevisiae also, we fit the probability of an up-regulating effect caused by a given gene in E. coli to a Beta distribution. For each node $j$ in an extracted subgraph of S. cerevisiae, we draw a probability $p_{j}$ from this Beta distribution and then sample the effect signs for the outgoing edges of node $j$ using $p_{j}$. As a result, the genes in the subgraphs of S. cerevisiae individually also have mostly up- or down-regulating effects. Maximum likelihood estimation for this Beta distribution yielded $\alpha=0.2588$ and $\beta=0.2499$.

The E. coli and S. cerevisiae graphs and effect signs used in the experiments are taken from the GeneNetWeaver repository (Schaffter et al., , 2011) (MIT License).

A.2 Data-Generating Processes

In the Linear and Rff domains, the data-generating processes are modeled by structural causal models (SCMs). In this work, we consider SCMs with causal mechanisms that model each causal variable $x_{j}$ given its parents $\mathbf{x}_{\text{pa}(j)}$ as

where the noise $\epsilon_{j}$ is additive and may be heteroscedastic through an input-dependent noise scale $\smash{h_{j}(\mathbf{x}_{\text{pa}(j)})}$. Even in the homogeneous noise setting, the scale of each noise distribution $p(\epsilon_{j})$ is random and thus different for each variable $x_{j}$. We write $\mathbf{x}_{\text{pa}(j)}$ when indexing $\mathbf{x}$ at the parents of node $j$. In the heteroscedastic setting, we parameterize the noise scales as $h_{j}(\mathbf{x})=\log(1+\exp(g_{j}(\mathbf{x}))$ for a set of nonlinear functions $g_{j}$.

Prior to performing inference with AVICI or any baseline, each set of SCM observations $D$ is standardized variable-wise by subtracting its mean and dividing by its standard deviation, so that each $x_{j}$ has mean and variance $1$, avoiding potential varsortability bias (Reisach et al., , 2021).

In the Linear domain, the functions $f_{j}$ are given by affine transforms

whose weights $\mathbf{w}_{j}$ and bias $b_{j}$ are sampled independently for each $f_{j}$. In the Rff domain, the functions $f_{j}$ modeling each causal variable $x_{j}$ given its parents $\mathbf{x}_{\text{pa}(j)}$ are drawn from a Gaussian Process

with $\alpha^{(m)}\sim\mathcal{N}(0,1)$, $\bm{\omega}^{(m)}\sim\mathcal{N}(0,\mathbf{I})$, and $\delta^{(m)}\sim\text{Unif}(0,2\pi)$. Throughout this work, we use $M$ $=$ $100$. The function draws become faithful GP samples as $M\rightarrow\infty$ (Rahimi and Recht, , 2007). When $x_{j}$ is a root node and thus has no parents, $f_{j}$ is a constant.

A.2.2 Single-Cell Gene Expression Data

In the Grn domain, our goal is to evaluate causal discovery from realistic gene expression data. There exist several models to simulate the mechanisms, intervention types, and technical measurement noise underlying single-cell expression data of gene regulatory networks (Schaffter et al., , 2011; Huynh-Thu and Sanguinetti, , 2019; Dibaeinia and Sinha, , 2020). We use the simulator by Dibaeinia and Sinha, (2020) (SERGIO) because it resembles the data collected by modern high-throughput single-cell RNA sequencing (scRNA-seq) technologies. Related genomics simulators, for example, GeneNetWeaver (Schaffter et al., , 2011), were developed for the simulation of microarray gene expression platforms. In the following, we give an overview of how to simulate scRNA-seq data with SERGIO. Dibaeinia and Sinha, (2020) provide all the details and additional background from the related literature.

Given a causal graph over $d$ genes and a specification of the simulation parameters, SERGIO generates a synthetic scRNA-seq dataset $D$ in two stages. The $n$ observations in $D$ correspond to $n$ cell samples, that is, the expressions of the $d$ genes recorded in a single cell corresponds to one row in $D$.

In the first stage, SERGIO simulates clean gene expressions by sampling randomly-timed snapshots from the steady state of a dynamical system. In this regulatory process, the genes are expressed at rates influenced by other genes using the chemical Langevin equation, similar to Schaffter et al., (2011) and (Dibaeinia and Sinha, , 2020). The source nodes in the causal graph $G$ are denoted master regulators (MRs), whose expressions evolve at constant production and decay rates. The expressions of all downstream genes evolve nonlinearly under production rates caused by the expression of their causal parents in $G$. Cell types are defined by specifications of the MR production rates, which significantly influence the evolution of the system. Thus, the dataset contains variation due to biological system noise within collections of cells of the same type and due to different cell types. Ultimately, we generate single-cell samples collected from five to ten cell types (Dibaeinia and Sinha, , 2020).

In the second stage, the clean gene expressions sampled previously are corrupted with technical measurement error that resembles the noise phenomena found in real scRNA-seq data:

outlier genes: a small set of genes have unusually high expression across measurements

library size: different cells have different total UMI counts, following a log-normal distribution

dropouts: a high percentage of genes are recorded with zero expression in a given measurement

unique molecule identifier (UMI) counts: we observe Poisson-distributed count data rather than the clean expression values

To configure these noise modules, we use the parameters calibrated by Dibaeinia and Sinha, (2020) for datasets from different scRNA-seq technologies. We extend SERGIO to allow for the generation of knockout intervention experiments. For this, we force the production rate of knocked-out genes to zero during simulation. Our implementation uses the public source code by (Dibaeinia and Sinha, , 2020), which is available under a GNU General Public License v3.0.https://github.com/PayamDiba/SERGIO

Given a causal graph $G$, the parameters SERGIO requires to simulate $c$ cell types of $d$ genes are:

The technical noise components are configured by:

$p_{\text{outlier}}\in$: probability that a gene is an outlier gene

In our experiments, the simulator parameters are selected in the ranges suggested by Dibaeinia and Sinha, (2020).

There are several ways to preprocess and normalize single-cell transcriptomic data for downstream use (Robinson et al., , 2010). For simplicity, we employ $\log_{2}$ counts-per-million (CPM) normalization, which normalizes the total UMI counts per sample and then $\log_{2}$-transforms the relative count values. Specifically, the CPM value for gene $j$ in sample $i$ is defined as

For zero expressions $\smash{x_{j}^{i}}$, the $\log_{2}$-CPM values are imputed with zero. The remaining $\log_{2}$-CPM values range between $10$ and $19$, so we shift and scale the values before performing causal discovery. To replicate the sparsity pattern and the relative ordering of values within samples in the original dataset $D$, we standardize the nonzero $\log_{2}$-CPM values by subtracting the minimum (instead of the mean) and dividing by the overall standard deviation. All methods considered in Section 6, including AVICI, work with Grn data in this standardized $\log_{2}$-CPM format.

Appendix B Evaluation Metrics

We report several metrics to assess how well the predicted causal structures reflect the ground-truth graph. We measure the overall accuracy of the predictions and how well-calibrated the estimated uncertainties in the edge predictions are, since AVICI predicts marginal probabilities $q(g_{i,j};\theta_{i,j})$ for every edge. Unless evaluating these edge probabilities, we use a decision threshold of $0.5$ to convert the AVICI prediction to a single graph $G$.

The structural hamming distance (SHD) (Tsamardinos et al., , 2006) reflects the graph edit distance between two graphs, i.e., the edge changes required to transform $G$ into $G^{\prime}$. By contrast, the structural intervention distance (SID) (Peters and Bühlmann, , 2015) quantifies the closeness of two DAGs in terms of their valid adjustment sets, which more closely resembles our intentions of using the inferred graph for downstream causal inference tasks.

SHD and SID capture global and structural similarity to the ground truth, but notions like precision and recall at the edge level are not captured well. SID is zero if and only if the true DAG is a subgraph of the predicted graph, which can reward dense predictions (Prop. 8 by Peters and Bühlmann, (2015): SID$(G,G^{\prime})=0$ when $G$ is empty and $G^{\prime}$ is fully connected). Conversely, the trivial prediction of an empty graph achieves highly competitive SHD scores for sparse graphs.

For this reason, we report additional metrics that quantify both the trade-off between precision and recall of edges as well as the calibration of their uncertainty estimates. Specifically, given the binary predictions for all $d^{2}$ possible edges in the graph $G$, we compute the edge precision, edge recall, and their harmonic mean (F1-score) for each test case and estimate their means and standard errors across the test cases. Since the F1-score is high only when precision and recall are high, both empty and dense predictions are penalized and no trivial prediction scores well, making it a reliable metric for structure learning.

To evaluate the edge probabilities predicted by AVICI and the baselines, we compute the areas under the precision-recall curve (AUPRC) and receiver operating characteristic (AUROC) when converting the probabilities into binary predictions using varying decision thresholds (Friedman and Koller, , 2003). Both statistics capture different aspects of the confidence estimates. The AUROC is insensitive to changes in class imbalance (edge vs. no-edge) for a given $d$. However, when the number of variables $d$ in sparse graphs of $O(d)$ edges increases, AUROC increasingly discounts the accuracy on the shrinking proportion of edges present in the ground truth, which makes AUPRC more suitable for comparisons ranging over different $d$. The AUROC is equivalent to the probability that the method ranks a randomly chosen positive instance (i.e., an edge $i\rightarrow j$ present in the ground truth) higher than a randomly chosen negative instance (i.e., an edge $i\dots j$ absent in the ground truth) (Fawcett, , 2004).

To assess the true correctness likelihood implied by the predicted edge probabilities, we use the concept of calibration (DeGroot and Fienberg, , 1983; Guo et al., , 2017). A classifier is said to be calibrated if a predicted edge probability of $\hat{p}_{i,j}$ empirically results in the observation of an edge in $(\hat{p}_{i,j}\times 100)$% of the cases, i.e.,

Following Guo et al., (2017), we can estimate the degree to which this property is satisfied for the predicted probabilities by defining $M$ intervals $I_{m}=(\tfrac{m-1}{M},\tfrac{m}{M})$ and binning all instances $i,j$ where $\hat{p}_{i,j}\in I_{m}$ into a set $S_{m}$. The empirical confidence and accuracy per bin $S_{m}$ are then defined as

where a calibrated classifier has $\text{predicted}~{}\hat{p}(S_{m})=\text{empirical}~{}p(S_{m})$, analogous to (15). Thus, a calibrated edge classifier induces a diagonal line when plotting the empirical $p(S_{m})$ against the predicted $\hat{p}(S_{m})$. The expected calibration error (ECE) is a scalar summary of this calibration plot and amounts to the weighted average of the vertical deviation from the perfect calibration line, i.e.,

where $n$ is the total number of evaluated samples (i.e., edges). The ECE does not capture accuracy in the sense of being able to predict all classes with high certainty, for which the other metrics are more suitable, but rather whether predicted probabilities are reflective empirical likelihood (Guo et al., , 2017). In this work, we use $M=10$ bins to compute the calibration plot lines and the ECE. The plotted calibration lines compute the calibration statistics in aggregate over all test cases to reduce the variance of the empirical counts within the bins, thus not showing standard errors.

Appendix C Inference Model Details

Each AVICI model is trained as described in Algorithm 1. The objective $\mathcal{L}(\phi)$ relies on samples from the domain distribution $p(G,D)$ to perform Monte Carlo estimation of the expectations. During training, the number of variables $d$ in the simulated systems are chosen randomly from

The datasets $D$ in the training distributions always have $n=200$ samples, where with probability $0.5$, the observations in a given dataset contain $50$ interventional samples. The dimensionality of these training instances $G,D$ varies significantly with the number of variables $d$ and, therefore, so do the memory requirements of the forward passes of the inference model $f_{\phi}$.

Given these differences in problem size, we make efficient use of the GPU resources during training by performing individual primal updates in Algorithm 1 using only training instances $(G,D)$ with exactly $d$ variables, where $d$ is randomly sampled in each update step. Fixing the number of observations to $n=200$, this allows us to increase the batch size for each considered $d$ to the maximum possible given the available GPU memory (in our case ranging from batch sizes of $27$ for $d=2$ down to $6$ for $d=50$, per $24$ GiB GPU device).

During training, we tune the sampling probability of a given $d$ to ensure that $f_{\phi}$ sees roughly the same number of training data sets for each $d$, i.e., we oversample higher $d$, for which the effective batch size per update step is smaller. We also scale $\mathcal{L}(\phi)$ by dividing by $d^{2}$ to ensure an approximately equal loss and hence gradient scale across the different $d$ seen at training time.

The penalty $\mathcal{F}(\phi)$ for the acyclicity constraint is estimated using the same minibatch as for $\mathcal{L}(\phi)$.

Since we have access to the complete data-generating process rather than only a fixed dataset, we approximate $\mathcal{L}(\phi)$ with minibatches that are sampled uniformly randomly from a buffer, which is continually updated with fresh data from $p(G,D)$. Specifically, we initialize a first-in-first-out buffer that holds $200$ pairs $(G,D)$ for each unique number of variables $d$ considered during training. A pool of asynchronous single-CPU workers then constantly generates novel training data and replaces the oldest instances in the buffer using a producer-consumer workflow. We implement this buffer using an Apache PyArrow Plasma object store (Apache Licence 2.0). During training, we used $128$ CPU workers (Appendix E).

The workers balance the data generation for different buffers to ensure an equal sample-to-insert ratio across $d$, accounting for the oversampling of higher $d$ as well as the longer computation time needed for generating data $D$ of larger $d$, for instance, in the Grn domain. In addition, the dataset $D$ of each element $(G,D)$ in the buffer contains four times more observations than $n=200$ used during training. These observations are subsampled to obtain $n=200$ each time a given buffer element $(G,D)$ is drawn to introduce additional diversity in the training data in case buffer elements are sampled more than once.

When adding the acyclicity contraint in Linear and Rff, we use a dual learning rate of $\eta=10^{-4}$ and perform a dual update every $500$ primal steps. The dual learning rate $\eta$ is warmed up with a linear schedule from zero over the first $50000$ primal steps. To reduce the variance in the dual update, we use an exponential moving average of $\mathcal{F}(\phi)$ with step size $10^{-4}$ maintained during the updates of the primal objective. To approximate the spectral radius in Eq. (8), we perform $t=10$ power iterations initialized at $\mathbf{u},\mathbf{v}\sim\mathcal{N}(0,\mathbf{I}_{d})$.

C.2 Architecture

As described in Section 4.2, the core of our model consists of $L=8$ layers, each containing four residual sublayers. Different from the vanilla Transformer encoder, we employ layer normalization before each multi-head attention and feedforward module and after the last of the $L$ layers (Radford et al., , 2019). The multi-head attention modules have a model size of $128$, key size of $32$, and $8$ attention heads. The feedforward modules have a hidden size of $512$ and use ReLU activations. In held-out tasks of Rff and Grn, we found that dropout in the Transformer encoder does not hurt performance in-distribution, so we increased the dropout rates from $0.0$ to $0.1$ and $0.3$, respectively, to help generalization o.o.d. Dropout, when performed, is done before the residual layers are added, as in the vanilla Transformer (Vaswani et al., , 2017).

We implement AVICI with Haiku in JAX (Hennigan et al., , 2020; Bradbury et al., , 2018). We converged to the above optimization and architecture specifications through experimentation on held-out instances from the training distributions $p(D)$, i.e., in-distribution.

Appendix D Baselines

DCDI (Brouillard et al., , 2020): sparsity regularizer $\lambda\in\{10^{-2},10^{-1},1\}$, size of hidden layer in MLPs modeling the conditional distributions $\in\{8,32\}$

DAG-GNN (Yu et al., , 2019): graph thresholding parameter $\in\{0.1,0.2,0.3\}$, size of hidden layer in MLP encoder and decoder $\in\{16,64\}$

GraN-DAG (Lachapelle et al., , 2020): preliminary neighborhood selection threshold $\in\{0.5,2\}$, size of hidden layer $\in\{8,32\}$, pruning cutoff $\in\{10^{-3},10^{-5}\}$

IGSP (Wang et al., , 2017): significance $\alpha\in\{10^{-2},10^{-3},10^{-4}\}$, CI test \in\{\text{Gaussian, HSIC-\gamma}\}

PC (Spirtes et al., , 2000): significance $\alpha\in\{10^{-2},10^{-3},10^{-4}\}$, CI test \in\{\text{Gaussian, HSIC-\gamma}\}

DAG-GNN, DCDI, DiBS, and GraN-DAG use 80% of the available data to perform inference and compute held-out log likelihood or ELBO scores on the other 20% of the data. The best hyperparameters are then selected by averaging the metric over five held-out instances of $d=30$ variables. DiBS draws $10$ samples from $p(G\,|\,D)$ using the interventional BGe score for Linear and a nonlinear Gaussian interventional likelihood with MLP means for Rff and Grn. DiBS assumes an observation noise of $1$, uses a scale-free graph prior, and anneals acyclicty and relaxation parameters with rate $1$. All remaining parameters are kept at the settings suggested by the authors.

For the PC algorithm and IGSP, there is no held-out score, so we compute the SID and F1 scores using the ground-truth causal graphs to select their optimal parameters. This would not be possible in practice and thus favors these methods. The HSIC-$\gamma$ CI test did not scale to $d=100$ variables, so in these cases PC and IGSP always use the Gaussian CI test. For Grn $d=30$, IGSP also uses the Gaussian CI test because it OOMs at 100GB when using HSIC-$\gamma$. GES and GIES use the linear Gaussian BIC score function and thus do not require calibrating a sparsity parameter (Chickering, , 2003; Hauser and Bühlmann, , 2012). LiNGAM is based on independent component analysis and requires no regularization tuning either (Shimizu et al., , 2006).

To estimate edge probabilities for the non-Bayesian methods in Section 6.2, we use the nonparametric DAG bootstrap (Friedman et al., , 1999). We bootstrap ten datasets $D^{\prime}$ from $D$ by sampling with replacement and then run each baseline individually on each bootstrapped dataset $D^{\prime}$. The nonparametric probability estimate for an edge then amounts to the proportion of predicted graphs $G^{\prime}$ that contain the edge.

For GES, GIES, PC, and LiNGAM, we run the original R implementations of the authors using an extended version of the software by Kalainathan et al., (2020) (MIT Licence). For DCDI, DAG-GNN, GraN-DAG, and DiBS, we use the Python implementations provided by the authors (Brouillard et al., , 2020; Yu et al., , 2019; Lachapelle et al., , 2020; Lorch et al., , 2021) (MIT License, Apache License 2.0, MIT License, MIT Licence). For IGSP, we use the implementation provided as part of the CausalDAG package (Squires et al., , 2018) (3-Clause BSD license).

LiNGAM relies on the inversion of a covariance matrix, which frequently fails in the Grn domain due to the high sparsity in $D$. Thus, to benchmark LiNGAM in Grn, we add small Gaussian noise to the standardized count matrix $D$. For the IGSP and PC algorithms, the same numerical adjustment is needed to avoid crashes in the CI tests on Grn. Single IGSP runs that still failed for $d=100$ were ignored when computing the metrics. In the Grn results, we ignored a small number of single runs of DCDI for $d=100$ and PC for $d=30$ that failed to terminate after $24$ hours walltime (on a GPU machine for the former). Lastly, the CAM pruning post-processing procedure of the author’s implementation of GraN-DAG (Lachapelle et al., , 2020) crashes in a few instances. We skip the post-processing step in these cases.

Appendix E Extended Results

To carry out the experiments in this work, we trained three main AVICI models and several ablations. Each model was optimized for approximately four days using $8$ Quadro RTX 6000 or NVIDIA GeForce RTX 3090 GPUs ($24$ GiB memory each) and $128$ CPUs. To perform the benchmarking experiments, all baselines were run on four to eight CPUs each for up to $24$ hours, depending on the method. DCDI required one GPU to ensure a computation time of less than one day per task instance. In all experiments, test-time inference with AVICI is done on eight CPUs and no GPU.

E.1 AVICI generalization between Linear and Rff

In this section, we provide additional out-of-distribution generalization results for AVICI. Specifically, we test the AVICI model trained on the Linear domain on inference from Rff data, and vice versa. This means that the AVICI models not only operate under distributional shifts on the parameters of their respective data-generating processes, but also on the function classes of causal mechanisms themselves. The models infer causal structure from data generated from function classes never seen during training. As in all empirical analyses of Section 6, the graph and noise parameters are additionally o.o.d., that is, the Linear AVICI model is tested on the o.o.d. Rff data, and vice versa.

Table 5 summarizes the results. Even under this distributional shift, the performance of both AVICI models decreases reasonably and remains on par with most baselines (Table 1). On Linear data, the baselines achieve F1 scores of $0.15$ - $0.54$ with observational and $0.33$ - $0.74$ with interventional data, similar to the Rff AVICI model with $0.19$ and $0.45$, respectively. Conversely, on Rff data, the baselines achieve F1 scores of $0.22$ - $0.42$ with observational and $0.34$ - $0.41$ with interventional data, which is also matched by the Linear AVICI model here with $0.27$ and $0.42$, respectively. Overall, the Linear AVICI model generalizes marginally better to Rff data as vice versa. We do not report the SID here because the R code of Peters and Bühlmann, (2015) runs out of memory.

E.2 In-Distribution Benchmarking Results for d=30𝑑30d=30

Table 6 gives the benchmarking results for in-distribution data of $d=30$ variables given the otherwise unchanged setup of Section 6.2. Contrary to the o.o.d. setting, the data is generated under homogeneous, additive noise and the parameters of their generative processes are sampled from the training domains of AVICI (cf. Table 3). However, as throughout all experiments, the datasets and its data-generating parameters themselves are unique and have not been used by AVICI during training.

Compared to the o.o.d. setting, most baselines perform roughly the same. Since the data-generating processes are sampled from its training distribution, AVICI significantly improves by moving to the easier in-distribution setting, in particular in the SCM domains, which are less noisy. In the Grn domain, some baselines achieve slightly better F1 scores compared to the o.o.d. setting, which is most likely explained by a change in the graph rather than the simulator parameter distribution, since there is no reason to believe that different generative parameters are more challenging to the baselines.

E.3 Benchmarking Results for d=100𝑑100d=100

Table 7 shows the benchmarking results for $d=100$ variables given $n=1000$ observations and the experimental setup of Section 6.2. We highlight that in this evaluation regime, AVICI operates under distribution shift in terms of the causal structures, mechanisms or simulator parameters, and noise distributions, as well as the number of variables and the number of observations seen during training.

Overall, the qualitative ranking of the methods is very similar as for $d=30$. AVICI outperforms all baselines in the nonlinear Rff domain, with and without access to interventional data. Likewise, AVICI is the only method to achieve nontrivial edge accuracy in terms of F1 score on the challenging Grn domain. On the simpler Linear domain, there is no statistically significant difference between GES/GIES and AVICI, which perform overall most favorably.

E.4 Benchmarking Results on Real-World Proteomics Data

We additionaly evaluate all of the methods on the real-world dataset by Sachs et al., (2005), which contains continuous measurements of $d=11$ proteins involved in human immune system cells. Structure learning algorithms are commonly compared on this dataset, and for completeness, we report the performance of AVICI and the baselines here. However, the ground-truth network of $17$ edges put forward by Sachs et al., (2005) has been challenged by some experts (Mooij et al., , 2020) and the assumptions of causal sufficiency and acyclicity may not be justified even though assumed by most methods, which should be kept in mind when interpreting the results. A large part of the data are interventional, in which the measured proteins were activated or inhibited using specific reagents. Most interventions are likely not perfect and the intervention targets may not be completely accurate (Mooij et al., , 2020).

For this experiment, we follow Wang et al., (2017) and Brouillard et al., (2020) and discard data in which interventions were not targeted directly at one of $d=11$ measured proteins. Given this setup, we have $n$ $=$ $5846$ data points that contain $1755$ observational and $4091$ interventional measurements, which consist of five single-protein perturbations. In our results, methods that only use observational data take the concatenation of all of the data without the intervention target information as input. All baselines use the hyperparameters tuned for the nonlinear Rff domain. The data is standardized to have mean and variance $1$.

GES LiNGAM PC DAG-GNN GraN-DAG $\begin{subarray}{c}\text{\small{AVICI}}\\ \text{({Linear})}\end{subarray}$ $\begin{subarray}{c}\text{\small{AVICI}}\\ \text{({Rff})}\end{subarray}$

GIES IGSP DCDI $\begin{subarray}{c}\text{\small{AVICI}}\\ \text{({Linear})}\end{subarray}$ $\begin{subarray}{c}\text{\small{AVICI}}\\ \text{({Rff})}\end{subarray}$

Table 8 summarizes the results of all methods with respect to the reference causal graph. Figure 6 visualizes the prediction of each method. Overall, the results are not very conclusive. GES and GIES perform best in terms of SID, GraN-DAG is most favorable in terms of F1, and together with DCDI and AVICI also in terms of SHD. More generally, the number of predicted edges varies greatly across methods. Almost all F1 scores fall between $0.25$ and $0.30$.

E.5 Uncertainty quantification for d=30𝑑30d=30

Figure 8 gives the calibration plots for all methods considered in the uncertainty analysis of Section 6.2 of the main text. In the SCM domains, AVICI closely traces the diagonal calibration line, both when given access to observational and mixed data. Here, the nonparametric bootstraps of the PC, GIES, and IGSP algorithms as well as DiBS are similarly well-calibrated. These baselines achieve worse expected calibration error (ECE) than AVICI because a significantly larger total proportion of AVICI’s predictions are well-calibrated (cf. Equation 17). DCDI, LiNGAM, and DAG-GNN are highly overconfident, that is, they predict edges with high probability when empirically only few edges exist.

Table 9 summarizes the probabilistic AUROC and AUPRC metrics for all methods in the experiment of Figure 5. Explanations and interpretations for both metrics are given in Section B. The relative performance of the bootstrap baselines and AVICI is similar to the point estimate benchmark. Overall, AVICI performs favorably across the three domains, with GES and GIES on par in Linear. However, since AUROC and AUPRC metrics evaluate the full spectrum of decision thresholds, we additionally see that AVICI achieves nontrivial accuracy in Grn even without access to gene knockout data, indicating that AVICI may provide useful information even in settings where only passive observations are available. This aspect is not apparent when converting the posterior probability estimates of AVICI based on a single threshold and then comparing SID and F1 scores as in Table 1.