Original: Amortized Inference for Causal Structure Learning
Revised: Amortized Inference for Causal Structure Learning

Original: Anonymous Author(s)

Original: AffiliationAddress email

Original: Abstract
Revised: Abstract

Original: Learning causal structure poses a combinatorial search problem that typically involves evaluating structures with a score or independence test.
Revised: Inferring causal structure poses a combinatorial search problem that typically involves evaluating structures with a score or independence test.

Original: The resulting searchis costly, and designing suitable scores or tests that capture prior knowledge isdifficult.
Revised: The resulting search is costly, and designing suitable scores or tests that capture prior knowledge is difficult.

Original: In this work, we propose to amortize causal structure learning .
Revised: In this work, we propose to amortize causal structure learning .

Original: Ratherthan searching over structures directly, we train a variational inference model topredict the causal structure from observational or interventional data.
Revised: Rather than searching over structures, we train a variational inference model to predict the causal structure from observational or interventional data.

Original: This allowsus to bypass both the search over graphs and the hand-engineering of suitable scorefunctions.
Revised: This allows us to bypass both the search over graphs and the hand-engineering of suitable score functions.

Original: Instead, our inference model acquires domain-specific inductive biasesfor causal discovery solely from data generated by a simulator.
Revised: Instead, our inference model acquires domain-specific inductive biases for causal discovery solely from data generated by a simulator.

Original: The architecture ofour inference model emulates permutation invariances that are crucial for statisticalefficiency in structure learning, which facilitates generalization to significantlylarger problem instances than seen during training.
Revised: 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.

Original: On synthetic data and semisynthetic gene expression data, our models exhibit robust generalization capabilitieswhen subject to substantial distribution shifts and significantly outperform existingalgorithms, especially in the challenging genomics domain.
Revised: 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.

Revised: Our code and models are publicly available at: https://github.com/larslorch/avici .

Original:  Introduction
Revised:  Introduction

Original: Learning the causal structure among the variableswe observe is a fundamental task in variousscientific disciplines (Spirtes et al., 2000; Pearl, 2009).
Revised: Learning the causal structure among a set of observed variables is a fundamental task in various scientific disciplines (Spirtes et al., 2000; Pearl, 2009).

Original: However, inferring this causal structure fromobservations of the variables is a difficult inverse problem.
Revised: However, inferring this causal structure from observations of the variables is a difficult inverse problem.

Original: The solution space of potential causalstructures, usually modeled as directed graphs, grows superexponentially with the number of variables.
Revised: The solution space of potential causal structures, usually modeled as directed graphs, grows superexponentially with the number of variables.

Original: To infer a causal structure, standard methods have to search over potential graphs, usually maximizingeither a graph scoring function or testing for conditional independences (Heinze-Deml et al., 2018).
Revised: 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).

Original: A key limitation of search approaches to causal discovery is the difficulty of specifying realisticinductive biases.
Revised: Specifying realistic inductive biases is universally difficult for existing approaches to causal discovery.

Original: Score-based methods use strong assumptions about the data-generating process, suchas linearity (Shimizu et al., 2006), specific noise models (Hoyer et al., 2008; Peters and Bühlmann,2014), and the absence of measurement error (e.g., Scheines and Ramsey 2016; Zhang et al.
Revised: 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.

Original: 2017),which are difficult to verify (Dawid, 2010; Reisach et al., 2021).
Revised: Scheines and Ramsey 2016; Zhang et al. 2017), which are difficult to verify (Dawid, 2010; Reisach et al., 2021).

Original: Conversely, constraint-based methodsdo not have enough domain-specific inductive bias.
Revised: Conversely, constraint-based methods do not have enough domain-specific inductive bias.

Original: Even with oracle conditional independence tests,they are limited to identifying equivalence classes that may be exponentially large (He et al., 2015b).
Revised: Even with an arbitrarily large dataset, they are limited to identifying equivalence classes that may be exponentially large (He et al., 2015b).

Original: Moreover, the search over directed graphs itself introduces unwanted bias and artifacts (Colomboet al., 2014).
Revised: Moreover, the search over directed graphs itself introduces unwanted bias and artifacts (Colombo et al., 2014).

Original: The intractable search space ultimately imposes hard constraints on the causal structure,e.g., the node degree (Spirtes et al., 2000; Koivisto and Sood, 2004), which limits the suitability ofsearch in real-world domains.
Revised: 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.

Original: In the present work, we propose to amortize causal structure learning.
Revised: In the present work, we propose to amortize causal structure learning.

Original: In other words, our goal isto optimize an inference model to solve the inverse problem of predicting a causal structure from aprovided dataset.
Revised: In other words, our goal is to optimize an inference model to solve the inverse problem of predicting a causal structure from a provided dataset.

Original: This proposal is motivated by the dichotomy that, given a causal model, it is easy togenerate realistic data; the inverse direction, that is, inferring the causal structure from data, is hard.
Revised: This proposal is motivated by the dichotomy that, given a causal model, it is easy to generate realistic data; the inverse direction, that is, inferring the causal structure from data, is hard.

Original: Many high-impact domains of causal discovery can be realistically forward-simulated, for example,the chemical ODEs and technical noise of gene regulatory networks (Schaffter et al., 2011; Dibaeiniaand Sinha, 2020), fMRI brain responses (Buxton, 2009; Bassett and Sporns, 2017), and chemical kinetics (Anderson and Kurtz, 2011; Wilkinson, 2018).
Revised: Specifically, many high-impact domains of causal discovery can be realistically forwardsimulated, for example, the chemical ODEs and technical noise of 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).

Original: We can leverage such simulators to trainan inference model in causal discovery for specific data-generating processes.
Revised: We can leverage such simulators to train an inference model in causal discovery for specific data-generating processes.

Original: In doing so, amortizedinference allows us to specify domain-specific inductive biases that are not easily represented bya graph scoring function and bypass the problems of structure search.
Revised: In doing so, amortized inference allows us to specify domain-specific inductive biases that are not easily represented by graph scoring functions and bypass the problems of structure search.

Original: The central challenge of amortizing causal discovery consists in designing a learning rule thatgeneralizes out-of-distribution. This generalization is essential because any simulated data distributionis unlikely to cover all real-world distribution shifts.
Revised: The central challenge of amortizing causal discovery consists in generalizing out-of-distribution, since any simulated data distribution is unlikely to cover all real-world distribution shifts.

Original: To quantify the epistemic uncertainty inthis inference process, we propose to train an amortized variational inferencemodel for causaldiscovery (AVICI) by approximating the posterior over causal graphs given data from a specified datagenerating distribution, such as a simulator.
Revised: To quantify the epistemic uncertainty in this inference process, we propose to perform amortized variational inference for causal discovery (AVICI) and train an inference model to infer the posterior over causal graphs given data from a specified data-generating distribution.

Original: The model architectureof AVICI is permutation invariantand 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.
Revised: 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.

Original: On synthetic data and semisyntheticexpression data of gene regulatory networks, our approachsignificantly outperforms existing algorithms for causal discovery, often by a large margin.
Revised: On synthetic data and semisynthetic gene expression data, our approach significantly outperforms existing algorithms for causal discovery, often by a large margin.

Original: Moreover,we demonstrate that our inference models induce calibrated uncertainties and robust behavior whensubject to substantial distribution shifts of graphs, mechanisms, noise, and problem sizes, whichmakes our trained models  both reliable and versatile for future downstream use.
Revised: 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 trained models are both reliable and versatile for future downstream use.

Original: In particular, AVICIwas the only method to infer plausible causal structure from noisy gene expression data, advancingthe frontiers of structure discovery in fields such as molecular biology.
Revised: 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 molecular biology.

Original:  Background and Related Work2.1 Causal Structure
Revised:  Background and Related Work 2.1 Causal Structure

Original: In this work, we define a causal structure in the same way as Mooij et al.
Revised: In this work, we follow Mooij et al.

Original: Models suchas causal Bayesian networks and structural causal models (Pearl, 2009; Peters et al., 2017) areless well-suited to describing systems with feedback loops, which we consider practically relevant.

Original: Specifically, we define the causal structure G as the directed graph whose nodes correspond to dvariables x = ( x 1 , . ..
Revised: and define the causal structure G of a set of d variables x = ( x 1 , . .

Original: , x d ) and whose edges represent all direct causal effects among the variables.
Revised: , x d ) as the directed graph over x whose edges represent all direct causal effects among the variables.

Original: Avariable x i has a direct causal effect on x j if intervening on x i affects the outcome of x j independentof the other variables x \ ij :
Revised: 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 x \ ij :

Original: = x \{ x i , x j } , i.e., there exists a ̸ = a ′ such that
Revised: = x \{ x i , x j } , i.e., there exists a ̸ = a ′ such that

Original: [Equation]
Revised: [Equation]

Original: for some c .
Revised: for some c .

Original: An intervention do ( · ) denotes any active manipulation of the generative process oftheobservations x . An example of interventions are gene knockouts in genomics, in which transcriptionrates of genes are externally set to zero.
Revised: 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.

Original: G is an abstraction of the physical processes generating thedata that describes the system’s behavior under arbitrary interventions (Peters et al., 2017; Rubensteinet al., 2017; Schölkopf, 2019).
Revised: 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.

Original: We emphasize that G is not necessarily acyclic; however, we latershow how to apply our approach when G is constrained to be acyclic.
Revised: 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.

Original: In this work, we assume causalsufficiency, i.e., that x contains all common causal parents of the variables x i (Peters et al., 2017).
Revised: We assume causal sufficiency, i.e., that x contains all common causal parents of the variables x i (Peters et al., 2017).

Original: 2.2 Related Work
Revised: 2.2 Related Work

Original: Classical methods for causal structure learning search over causal graphs and evaluate them usinga likelihood or conditional independence tests (Chickering, 2003; Kalisch and Bühlman, 2007; Hauser and Bühlmann, 2012; Zheng et al., 2018; Heinze-Deml et al., 2018).
Revised: 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).

Original: Other methods combineconstraint- and score-based ideas (Tsamardinos et al., 2006) or use the noise properties of an SCMthat is postulated to underlie the data-generating process (Shimizu et al., 2006; Hoyer et al., 2008).
Revised: 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).

Original: Deep learning has been applied to causal inference, e.g., for the estimation of treatment effects (Shalitet al., 2017; Louizos et al., 2017; Yoon et al., 2018) and in instrumental variable analysis (Hartfordet al., 2017; Bennett et al., 2019).
Revised: 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).

Original: Incausal structure learning, neural networks have  been usedprimarily to model nonlinear causal mechanismsamong the variables (Goudet et al., 2018;Yu et al.
Revised: In structure learning, neural networks have primarily been used to model nonlinear causal mechanisms (Goudet et al., 2018;

Original: ,2019; Lachapelle et al., 2020; Brouillard et al., 2020; Lorch et al., 2021) or to infer the structurefor a single dataset (Zhu et al., 2020).
Revised: 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).

Original: While amortized inference has gained increasing popularity inrecent years (Amos, 2022), prior work applying it to causal discovery only studied narrowly definedsubproblems such as the bivariate case (Lopez-Paz et al., 2015) and fixed causal mechanismsfor eachtask (Löwe et al., 2022) oronly usedobservational correlation coefficients for prediction (Li et al.,2020).
Revised: 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).

Original: In parallel work, Ke et al.
Revised: Ke et al.

Original: (2022) independently perform causal discovery as a supervisedlearning task.
Revised: (2022) also frame causal discovery as supervised learning, but with significant differences.

Original: There are significant differences in their approach, in particular, our inference modelsgeneralize to significantly larger problem sizes, evenwhen trained on realistic genomics data.
Revised: Most importantly, we optimize a variational objective under a model class that captures the symmetries of structure learning. Empirically, our models generalize to significantly larger problem sizes, even on realistic genomics data.

Original:  AVICI: Amortized Variational Inference for Causal Discovery3.1 The Variational Objective
Revised:  AVICI: Amortized Variational Inference for Causal Discovery 3.1 Variational Objective

Original: To amortize causal structure learning, we define a data-generating distribution p ( D ) that models thedomain over which structure learning is performed.
Revised: To amortize causal structure learning, we define a data-generating distribution p ( D ) that models the domain in which we infer causal structures.

Original: From a generative perspective, the observations D = { x 1 , . . .
Revised: The observations D = { x 1 , . . .

Original: , x n } ∼ p ( D ) are produced by sampling from a distribution over causal structures p ( G )and then obtaining realizations of a data-generating mechanism p ( D | G ) .
Revised: , x n } ∼ 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 ) .

Original: It is important to highlightthat p ( D | G ) does not need to represent an ancestral samplingprocedure over a directed acyclicgraph.
Revised: 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.

Original: In particular, real-world systems are often more naturally modeled at different granularitiesor as dynamical systems (Mooij et al., 2013; Hoel et al., 2013;  Schölkopf, 2019).
Revised: 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).

Original: Rather, p ( D | G )characterizes a data-generating process whose causal graph G formalizes all direct causal effects (Eq.1) of the system’s mechanisms, that is, how intervening on certain variables influences the realizationsof others (Mooij et al., 2016; Peters et al., 2017; Rubenstein et al., 2017).

Original: Having observed a set of observations D in our domain, our goal is to approximate the true posteriorover causal structures p ( G | D ) with a variational distribution q ( G ; θ ) .
Revised: 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 ; θ ) .

Original: To amortize this inferencetask for the data-generating distribution p ( D ) , we optimize an inference model f ϕ to predict thevariational parameters θ that minimize the expected forward KL divergence from the intractableposterior p ( G | D ) to q ( G ; θ ) for D ∼ p
Revised: To amortize this inference task for the domain distribution p ( D ) , we optimize an inference model f ϕ to predict the variational parameters θ by minimizing the expected forward KL divergence from the intractable posterior p ( G | D ) to q ( G ; θ ) for D ∼ p

Original: ( D ) :
Revised: ( D ) :

Revised: [Equation]

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Since it is not tractable to compute the true posterior in (2), we make use of ideas by Barber and
Revised: Since it is not tractable to compute the true posterior in (2), we make use of ideas by Barber and

Original: Agakov (2004) and rewrite the expected forward KL to obtain an equivalent, tractable objective:
Revised: Agakov (2004) and rewrite the expected forward KL to obtain an equivalent, tractable objective:

Revised: E p ( D ) D KL (cid:0) p ( G | D ) (cid:13)(cid:13) q

Revised: ( G ; f ϕ ( D ))

Revised: (cid:1) =

Revised: E p ( D ) E p ( G | D ) [log p ( G | D ) − log q ( G ; f ϕ ( D ))]

Revised: = − E p

Original: [Equation]
Revised: [Equation]

Original: The constant does not depend on ϕ , so we can maximize L ( ϕ ) :
Revised: The constant does not depend on ϕ , so we can maximize L ( ϕ ) :

Original: = E p ( G ) E p ( D | G )
Revised: = E p ( G ) E p ( D | G )

Original: [log q ( G ; f ϕ ( D ))] ,which allows us to amortize variational inference for causal discovery (AVICI).
Revised: [log q ( G ; f ϕ ( D ))] , which allows us to perform amortized variational inference for causal discovery (AVICI).

Original: While the datagenerating distribution p ( D ) =
Revised: While the domain distribution p ( D ) =

Original: E p ( G )
Revised: E p ( G )

Original: [ p ( D | G )] can be arbitrarily complex, L is tractable wheneverwe have access to theground-truth causal graph G underlying the generative process of D , i.e., tosamples from the joint distribution p ( G, D ) .
Revised: [ p ( D | G )] can be arbitrarily complex, L is tractable whenever we have access to the causal graph G underlying the generative process of D , i.e., to samples from the joint distribution p ( G, D ) .

Original: In practice, p ( G ) and p ( D | G ) can thus be specified bya simulator for the domain of interest.
Revised: In practice, p ( G ) and p ( D | G ) can thus be specified by a simulator.

Original: Since the quality of simulators in many domains may improvesubstantially over time, the sim-to-real bias of the inference model f ϕ can reduce along the way.

Original: From an information-theoretic view, the objective (2) maximizes a variational lower bound on themutual information I [ G ; D ] between the causal structure G and the observations D (Barber and
Revised: 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

Original: Agakov, 2004).
Revised: Agakov, 2004).

Original: Starting from the definition of mutual information, we obtain
Revised: Starting from the definition of mutual information, we obtain

Original: [Equation]
Revised: [Equation]

Original: where the entropy H
Revised: where the entropy H

Original: [ G ] is constant.
Revised: [ G ] is constant.

Original: The bound is tight if E 
Revised: The bound is tight if E p ( D ) D KL ( p

Original: [Equation]
Revised: [Equation]

Original:  D )
Revised: G | D )

Original: ∥ q ( G ; f ϕ( D )))
Revised: ∥ q ( G [Equation] f ϕ

Revised: [Equation]

Revised: D )))

Original: = 0 .
Revised: = 0 .

Original: 3.2 Likelihood-Free Inference using the Forward KL
Revised: 3.2 Likelihood-Free Inference using the Forward KL

Original: The AVICI objective in (3) intentionally targets the forward KL D KL ( p ∥ q ( · ; θ )) , which requiresoptimizing E p ( G,D )
Revised: The AVICI objective in (3) intentionally targets the forward KL D KL ( p ∥ q ( · ; θ )) , which requires optimizing E p ( G,D )

Original: [log q ( G ; θ )] .
Revised: [log q ( G ; θ )] .

Original: This choice implies that we both model the density q ( G ; θ ) explicitly and have access to samples from the true data-generating distribution.
Revised: This choice implies that we both model the density q ( G ; θ ) explicitly and assume access to samples from the true data-generating distribution p ( G, D ) .

Original: As long as samples fromp ( G, D ) are available to estimate L ( ϕ ) , minimizing the forward KLin (2) enables us to infer causalstructures in arbitrarily complex domains—that is, even domains where it is difficult to specify anexplicit likelihood p
Revised: Minimizing the forward KL enables us to infer causal structures in arbitrarily complex domains—that is, even domains where it is difficult to specify an explicit likelihood p

Original: ( D | G ) .
Revised: ( D | G ) .

Original: Furthermore, the forward KL typically yields more reliable uncertaintyestimates since it does not suffer from the variance underestimation problems common to the reverse KL (Bishop and Nasrabadi, 2006).
Revised: 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).

Original: In contrast, variational inference usually optimizes the reverse KL D KL ( q (· ; θ ) ∥ p )with respectto θ , which involves the reconstruction term E q ( G ; θ ) [log p ( D | G )]
Revised: In contrast, variational inference usually optimizes the reverse KL D KL ( q ( ∥ p ) , which involves the reconstruction term E q ( G ; θ ) [log p ( D | G )]

Original: (Blei et al., 2017).
Revised: (Blei et al., 2017).

Original: This objectiverequires a tractable marginal likelihood p ( D | G ) .
Revised: This objective requires a tractable marginal likelihood p ( D | G ) .

Original: Unless inferring the mechanism parameters jointly(Brouillard et al., 2020; Lorch et al., 2021), this requirement limits inference to conjugate models withlinear 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 n [Figure]
Revised: 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; Runge et al., 2019; Dibaeinia and Sinha, 2020).

Original: n
Revised: Furthermore, unless the noise scale is learned jointly, likelihoods can be sensitive to the measurement scale of x (Reisach et al., 2021).

Original:  [Figure] ) n
Revised:  Inference Model

Original: … … (c) In the following section, we describe a choice for the variational distribution q ( G ; θ ) and the inferencemodel f ϕ that predicts θ given D .
Revised: In the following section, we describe a choice for the variational distribution q ( G ; θ ) and the inference model f ϕ that predicts θ given D .

Original: After that, we detail our training procedure for optimizing themodel parameters ϕ as well as for learning the inductive bias of acyclicity.
Revised: After that, we detail our training procedure for optimizing the model parameters ϕ and for learning causal graphs with acyclicity constraints.

Original: Inference Model
Revised: n d k x L … … (a) n d k x L … … (b) n d k x L … … (c)

Original: 4.1 Variational Family
Revised: 4.1 Variational Family

Original: While any inference model that defines a density is feasible for maximizing the objective in (3), weopt to use a factorized variational family in this work.
Revised: 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.

Original: [Equation]
Revised: [Equation]

Original: The inference model f ϕ maps a dataset D corresponding to n samples { o 1 , . . .
Revised: The inference model f ϕ maps a dataset D corresponding to n samples { o 1 , . . .

Original: , o n } to a d -by- dmatrix θ parameterizing the variational approximation of the causal graph posterior.
Revised: , o n } to a d -by- d matrix θ parameterizing the variational approximation of the causal graph posterior.

Original: In addition tothe joint observation x i =
Revised: In addition to the joint observation x i =

Original: ( x i 1 , . . .
Revised: ( x i 1 , . . .

Original: , x id ) , each sample o i =
Revised: , x id ) , each sample o i =

Original: ( o i 1 , . . .
Revised: ( o i 1 , . . .

Original: , o id ) may contain interventionalinformation for each variable.
Revised: , o id ) may contain interventional information for each variable.

Original: When interventions or gene knockouts are performed, we seto ij = ( x ij , u ij ) and u ij ∈ { 0 , 1 } indicating whether variable j was intervened upon in sample i .
Revised: When interventions or gene knockouts are performed, we set o ij = ( x ij , u ij ) and u ij ∈ { 0 , 1 } indicating whether variable j was intervened upon in sample i .

Original: Othersettings could be encoded analogously, e.g., when the intervention targets are unknown.
Revised: Other settings could be encoded analogously, e.g., when the intervention targets are unknown or measurements incomplete.

Original: 4.2 Model Architecture
Revised: 4.2 Model Architecture

Original: To maximize statistical efficiency, f ϕ shouldbe constrained to satisfy the invariances inherent to thetask oflearning a causal structure from data.
Revised: To maximize statistical efficiency, f ϕ should satisfy the symmetries inherent to the task of causal structure learning.

Original: Firstly, f ϕ should be permutation invariant across thesample dimension (axis n ).
Revised: Firstly, f ϕ should be permutation invariant across the sample dimension (axis n ).

Original: Shuffling the samples should have no influenceon the prediction, i.e., forany permutation π ,  f ϕ ( π ( { o } )) =
Revised: Shuffling the samples should not influence the prediction, i.e., for any permutation π , we have f ϕ ( π ( { o } )) =

Original: f ϕ ( { o } ) .
Revised: f ϕ ( { o } ) .

Original: Moreover, f ϕ should be permutation equivariant acrossthe variable dimension (axis d ).
Revised: Moreover, f ϕ should be permutation equivariant across the variable dimension (axis d ).

Original: Reordering the variables should permute the predicted causal edgeprobabilities, i.e., f ϕ ( { o π (1: d ) } ) i,j = f ϕ ( { o 1: d } )
Revised: Reordering the variables should permute the predicted causal edge probabilities, i.e., f ϕ ( { o π (1: d ) } ) i,j = f ϕ ( { o 1: d } )

Original: π ( i ) ,π ( j ) .
Revised: π ( i ) ,π ( j ) .

Original: Lastly, f ϕ should apply to any d, n ≥ 1 .
Revised: Lastly, f ϕ should apply to any d, n ≥ 1 .

Original: In the following, we show how to parameterize f ϕ as a neural network that encodes these properties.
Revised: In the following, we show how to parameterize f ϕ as a neural network that encodes these properties.

Original: After first mapping each o ij to a real-valued vector using a position-wise linear layer, f ϕ operates overa continuous, three-dimensional tensor of n rows for the observations, d columns for the variables,and feature size k .
Revised: After first mapping each o ij to a real-valued vector using a position-wise linear layer, f ϕ operates over a continuous, three-dimensional tensor of n rows for the observations, d columns for the variables, and feature size k .

Original: Figure 1 illustrates the key components of the architecture, which consists ofan observation encoder followed by a generative graph model, similar to Zhu et al.
Revised: Figure 1 illustrates the key components of the architecture.

Original: Attending over axesn and d  The core of f ϕ is composed of L = 8 identical blocks.
Revised: Attending over axes d and n The core of f ϕ is composed of L = 8 identical layers.

Original: Each block consists of four residual layers, where the first and third apply multi-head self-attention and the second andfourth position-wise feed-forward networks, similar to thevanilla Transformer encoder (Vaswani et al.,2017).
Revised: 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).

Original: To enable information flow across all n × d tokens of the representation, the model alternates inattending over the observation and the variable dimension (Kossen et al., 2021).
Revised: To enable information flow across all n × d tokens of the representation, the model alternates in attending over the observation and the variable dimension (Kossen et al., 2021).

Original: Specifically, the firstself-attention block attends over axis n , treating axis d as a batch dimension; the second attends overaxis d , treating axis n as a batch dimension.
Revised: 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.

Original: This message-passing schema keeps memory in checkand results in a moderate model size because parameters are shared across the batched axes.

Original: Moreover,since multi-head attention is permutation equivariant and modules are shared across non-attendedaxes, the representation is permutation equivariant over axes n and d at all times (Lee et al., 2019b).
Revised: 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).

Original: Causal graph model After building up a representation tensor from the input using the attentionblocks, we max-pool over the observation axis n to obtain a representation ( z 1 , . .
Revised: 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 1 , . .

Original: z d ) consisting of one vector z i ∈ R k for each causal variable.
Revised: z d ) consisting of one vector z i ∈ R k for each causal variable.

Original: Following Lorch et al.
Revised: Following Lorch et al.

Original: (2021), we use two position-wiselinear layers to map each z i to two embeddings u i , v i ∈ R k .
Revised: (2021), we use two position-wise linear layers to map each z i to two embeddings u i , v i ∈ R k , which are ℓ 2 normalized.

Original: After normalizing the embeddingsu i , v i to an ℓ 2 -norm of one (Radford et al., 2021), we model the probability of each edge in thecausal graph with an inner product:
Revised: We then model the probability of each edge in the causal graph with an inner product:

Original: [Equation]
Revised: [Equation]

Original: ) where σ is the logistic function, b a learned bias, and τ a positive scale that is learned in log space.
Revised: where σ is the logistic function, b a learned bias, and τ a positive scale that is learned in log space.

Original: Since max-pooling is invariant to permutations and sincethe graph model in (6) permutes withrespect to axis d ,the inference model f ϕ satisfies the required permutation invariance over axis n andpermutation equivariance over axis d .
Revised: Since max-pooling is invariant to permutations and since (6) permutes with respect to axis d , f ϕ satisfies the required permutation invariance over axis n and permutation equivariance over axis d .

Original: 4.3 Acyclicity
Revised: 4.3 Acyclicity

Original: There is no reason to generally exclude feedback loops in G . Cyclic causal effects often occur, e.g.,when modeling stationary distributions of dynamical systems.
Revised: Cyclic causal effects often occur, e.g., when modeling stationary distributions of dynamical systems, and thus loops in a causal structure are possible.

Original: However, certain domains may bemore accurately modeled by acycliccausal structures (Peters et al., 2017; Rubenstein et al., 2017).
Revised: However, certain domains may be more accurately modeled by acyclic structures (Rubenstein et al., 2017).

Original: While the variational family in (5) does not allow enforcing it directly, we can optimize for acyclicitythrough ϕ .
Revised: While the variational family in (5) cannot enforce it, we can optimize for acyclicity through ϕ .

Original: Whenever the acyclicity prior is justified, we amend the optimization problem in (2) withthe constraint that q only models acyclic graphs in expectation:
Revised: Whenever the acyclicity prior is justified, we amend the optimization problem in (2) with the constraint that q only models acyclic graphs in expectation:

Original: [Equation]
Revised: [Equation]

Original: The function h is zero if and only if the predicted edge probabilities induce an acyclic graph.
Revised: The function h is zero if and only if the predicted edge probabilities induce an acyclic graph.

Original: Weuse the insight by Lee et al. (2019a), who show that acyclicity is equivalent to the spectral radius ρ ,i.e., the largest absolute eigenvalue, of the predicted matrix being zero.
Revised: We use the insight by Lee et al. (2019a), who show that acyclicity is equivalent to the spectral radius ρ , i.e., the largest absolute eigenvalue, of the predicted matrix being zero.

Original: We use power iteration toapproximate and differentiate through the largest eigenvalue of f ϕ ( D ) (Golub and Van der Vorst,2000; Lee et al., 2019a):
Revised: We use power iteration to approximate and differentiate through the largest eigenvalue of f ϕ ( D ) (Golub and Van der Vorst, 2000; Lee et al., 2019a):

Original: [Equation]
Revised: [Equation]

Original: where for t steps:
Revised: where for t steps:

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: (8) and a , b ∈ R d are initialized randomly.
Revised: and a , b ∈ R d are initialized randomly.

Original: Since a few steps t are sufficient in practice, (8) scales with
Revised: Since a few steps t are sufficient in practice, (8) scales with

Original: O ( d 2 ) and is significantly more efficient thanalternative O ( d 3 ) constraints based on matrix powers(Zheng et al., 2018; Yu et al., 2019).
Revised: O ( d 2 ) and is significantly more efficient than O ( d 3 ) constraints based on matrix powers (Zheng et al., 2018; Yu et al., 2019).

Original: We do not backpropagate gradients with respect to ϕ through a , b .
Revised: We do not backpropagate gradients with respect to ϕ through a , b .

Original:  Optimization
Revised: 4.4 Optimization

Original: Parameters: ϕ variational, λ dual, η step size while not converged do for l steps do ∆ ϕ ∝ ∇ ϕ (cid:0) L ( ϕ ) − λ F ( ϕ ) (cid:1)
Revised: Parameters: ϕ variational, λ dual, η step size while not converged do for l steps do ∆ ϕ ∝ ∇ ϕ (cid:0) L [Equation] ϕ ) − [Equation]

Original:  λ ← λ + η F ( ϕ )
Revised: F ( ϕ [Equation] (cid:1) λ ← λ + η F ( ϕ )

Original: Combining the objective in (3) with our inferencemodel (5), we can directly use stochastic optimization to train the parameters ϕ of our neural networkmodel.
Revised: Combining the objective in (3) with our inference model (5), we can directly use stochastic optimization to train the parameters ϕ of the inference model.

Original: The expectations over p ( G, D ) inside Land F are approximated using samples from thedata-generating process of the domain, e.g., a simulator.
Revised: The expectations over p ( G, D ) inside L and F are approximated using samples from the datagenerating process of the domain.

Original: In practice, the available data can be furtheraugmented using bootstrapping.

Original: When imposingan acyclicity prior on G , causal discovery algorithms often use the augmented Lagrangian method forconstrained optimization (Zheng et al., 2018;
Revised: When enforcing acyclicity, causal discovery algorithms often use the augmented Lagrangian method for constrained optimization (e.g., Zheng et al. 2018;

Original: Yu et al., 2019; Brouillard et al., 2020).
Revised: Brouillard et al. 2020).

Original: In this work, weoptimize the parameters ϕ of a neural network, which is a different constrained optimization problemin nature. For this reason, we rely on methods specifically tailored for deep learning and solve theconstrained program max ϕ L ( ϕ )
Revised: In this work, we optimize the parameters ϕ of a neural network, so we rely on methods specifically tailored for deep learning and solve the constrained program max ϕ L ( ϕ )

Original: s.t. F ( ϕ ) = 

Revised: s.t. F ( ϕ ) =

Original: through its dual formulation (Nandwani et al., 2019):min λ max ϕ
Revised: through its dual formulation (Nandwani et al., 2019): [Equation]

Original: [Equation]
Revised: [Equation]

Original: Algorithm 1 summarizes the general optimization procedure for q ϕ , which converges to a localoptimum under regularity conditions on the learning rates (Jin et al., 2020).
Revised: Algorithm 1 summarizes the general optimization procedure for q ϕ , which converges to a local optimum under regularity conditions on the learning rates (Jin et al., 2020).

Original: Without an acyclicityconstraint, training reduces to the primal updates of ϕ with λ = 0 .
Revised: Without an acyclicity constraint, training reduces to the primal updates of ϕ with λ = 0 .

Original:  Experimental Setup
Revised:  Experimental Setup

Original: Evaluating causal discovery algorithms is difficult since there are few interesting real-world datasetsthat come with ground-truth causal structure.
Revised: Evaluating causal discovery algorithms is difficult since there are few interesting real-world datasets that come with ground-truth causal structure.

Original: Often, the believed ground truths may be incompleteor change as expert knowledge improves (Schaffter et al., 2011; Mooij et al., 2020).
Revised: Often, the believed ground truths may be incomplete or change as expert knowledge improves (Schaffter et al., 2011; Mooij et al., 2020).

Original: Following priorwork, we deal with this difficulty by evaluating our approach using simulated data with known causalstructure and by controlling for various aspects of the task.
Revised: 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.

Original: In Appendix E, we additionally reportresults on a real-world proteomics dataset (Sachs et al., 2005).
Revised: In Appendix E, we additionally report results on a real-world proteomics dataset (Sachs et al., 2005).

Revised: 5.1 Domains and Simulated Components

Revised: We study three domains: two classes of structural causal models (SCMs) as well as semisynthetic single-cell expression data of gene regulatory networks (GRNs).

Revised: To study the generalization of AVICI beyond the training distribution p ( D ) , we carefully construct a spectrum of test distributions ˜ p

Revised: ( D ) that incur substantial shift from p ( D ) in terms of the causal structures, mechanisms, and noise, which we study in various combinations.

Original:  Erdos-Renyi
Revised: Whenever we consider interventional data in our experiments, half of the dataset consists of observational data and half of single-variable interventions. Data-generating processes p ( D | G ) We consider SCMs with linear functions (L INEAR ) and with nonlinear functions of random Fourier features (R FF ) that correspond to functions drawn Erdos-Renyi

Original: Scale-free
Revised: Scale-free

Original: Watts-Strogatz
Revised: Watts-Strogatz

Original: Stochastic Block Model
Revised: Stochastic Block Model

Original: Geometric Random x i 5 0 5 x j 50fc =15, =
Revised: Geometric Random x i 5 0 5 x j 50fc =15, =

Original: x i 5 0 5 x j 50fc =20, =
Revised: x i 5 0 5 x j 50fc =20, =

Original: x i 5 0 5 x j 50fc =20, =   5 x j   j
Revised: x i 5 0 5 x j 50fc =20, =   5 x j   j

Original: N(0,1)
Revised: N(0,1)

Original:  5 x j 5.0 2.
Revised:  5 x j 5.0 2.

Original: j
Revised: j

Original: Laplace(0, 2 ), = h j ( x pa ( j ) )
Revised: Laplace(0, 2 ), = h j ( x pa ( j ) )

Original: [Figure]
Revised: [Figure]

Original: 5.1 Domains and Simulated Components

Original: [Figure] Our goal is to evaluate our inference models’ ability to learn inductive biases in previously intractabledomains and to generalize robustly beyond the data distribution p ( D ) seen during training. We studythree domains: two classes of structural causal models (SCMs) as well as semisynthetic single-cellexpression data of gene regulatory networks (GRNs). To study the generalization of AVICI beyondthe training distribution p ( D ) , we carefully construct a test distribution ˜ p
Revised: from a Gaussian process with squared exponential and Recht,

Original: ( D ) that incurs substantialshift from p ( D ) in terms of the underlying causal structures, mechanisms, and noise components. Data-generating processes p ( D | G ) The data-generating process of SCMs is defined byrecursively sampling x j ← f j ( x pa ( j ) , ϵ j ) based on a set of mechanisms f j and exogenous noisedistributions p j ( ϵ j ) for each variable. We consider the two domains of linear functions (L INEAR ) andof nonlinear functions of random Fourier features (R FF ) that correspond to functions drawn from a
Revised: In the out

Original: Gaussian process with squared exponential kernel (Rahimi and Recht, 2007). In the out-of-distribution(o.o.d.) setting ˜p ( D ) , wesample thelinear function andkernel parameters from the tails of p ( D )and unseenvalue ranges.
Revised: ˜ we the and unseen

Original: Moreover, while simulating additive, homoscedastic Gaussian noise ϵ jin the training distribution p ( D ) , we employ heteroscedasticCauchy and Laplacian noise o.o.d.
Revised: Moreover, we simulate homoscedastic Gaussian noise in the training but on heteroscedastic and that is induced by

Original: ,where f j ( x Pa ( j ) , ϵ j ) = f j ( x pa ( j ) ) + h j ( x pa ( j ) )
Revised: In L

Original: ϵ j for some random, input-dependent function h j . Interventions are performed by clipping a variable x j to a random value during data generation.
Revised: , interventions to random values and are performed on a subset of target variables containing half of the nodes.

Original: In addition to SCMs, we consider the challenging domain of GRNs (G RN ) using the simulator of
Revised: In addition to SCMs, we consider the challenging domain of GRNs (G RN ) using the simulator of

Original: Dibaeinia and Sinha (2020).
Revised: Dibaeinia and Sinha (2020).

Original: Contrary to SCMs,single-cell gene expression samples correspondto draws from the steady state of a stochastic dynamical system that varies between cell types(Huynh-Thu and Sanguinetti, 2019).
Revised: 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).

Original: Like L INEAR and R FF , the simulator parameters consideredo.o.d. are drawn from significantly wider ranges than for p ( D ) .

Original: The simulator additionally modelsseveral technical noise effects, resulting in the observation of highly zero-inflated count matrices D .

Original: When moving o.o.d.
Revised: In the o.o.d.

Revised: setting, the parameters sampled for the G RN simulator are drawn from significantly wider ranges.

Original: , we use the noise levels of different single-cell RNA sequencing technologies,as calibratedby Dibaeinia and Sinha (2020) on real datasets.
Revised: In addition, we use the noise levels of different single-cell RNA sequencing technologies, which were calibrated on real datasets.

Original: In G RN , interventionscorrespond togene knockouts and are performed by forcing the transcription rate of a variablex j to zero.
Revised: In G RN , interventions are performed on all nodes and correspond to gene knockouts, forcing the transcription rate of a variable to zero.

Original: Whenever a method infers causal structure from interventional data in our experiments, half of thedataset consists of observational data and half of single-variable interventions.

Original: In L INEAR and R FF ,these single-variable interventions are performed on a subset of target variables containing 50 %of the nodes.

Original: Since G RN is more challenging, the interventional data here contains single-nodeinterventions on all variables.

Original: All datasets are standardized (Reisach et al., 2021).

Original: Causal structures p ( G )
Revised: Causal structures p ( G )

Original: Following prior work, we use random graph models and known biologicalnetworks to generate ground-truth causal structures.
Revised: Following prior work, we use random graph models and known biological networks to sample ground-truth causal structures.

Original: In all three domains, the training data distributionp ( D ) is induced by simple Erd˝os-Rényi and scale-free graphsof various degrees (Erd˝os and Rényi,1959; Barabási and Albert, 1999).
Revised: In all three domains, the training data distribution p ( D ) is induced by simple Erd˝os-Rényi and scale-free graphs (Erd˝os and Rényi, 1959; Barabási and Albert, 1999).

Original: In the o.o.d.
Revised: In the o.o.d.

Original: setting,the data distributions ˜ p
Revised: setting, ˜ p

Original: ( D ) of the L INEAR and R FF domains are simulated using causal structures fromunseen graph models: the Watts-Strogatzmodel, capturing small-world phenomena (Watts and Strogatz, 1998); the stochastic block model,generalizing Erd˝os-Rényi to community structures (Holland et al., 1983); and geometric randomgraphs, modeling connectivity based on spatial distance (Gilbert, 1961).
Revised: ( D ) of the L INEAR and R FF domains are simulated using causal structures from the Watts-Strogatz model, capturing small-world phenomena (Watts and Strogatz, 1998); the stochastic block model, generalizing Erd ˝ os-Rényi to community structures (Holland et al., 1983); and geometric random graphs, modeling connectivity based on spatial distance (Gilbert, 1961).

Original: In the G RN domain, weuse subgraphs of the known S. cerevisiae and E. coli GRNs and their effect signs whenever known.
Revised: In the G RN domain, we use subgraphs of the known S. cerevisiae and E. coli GRNs and their effect signs whenever known.

Original: To extract these subgraphs, we use the procedure by Marbach et al.
Revised: To extract these subgraphs, we use the procedure by Marbach et al.

Original: (2009) to maintain structuralpatterns like motifs and modularity (Ravasz et al., 2002; Shen-Orr et al., 2002).
Revised: (2009) to maintain structural patterns like motifs and modularity (Ravasz et al., 2002; Shen-Orr et al., 2002).

Original: To illustrate the distribution shift from p ( D ) to ˜ p
Revised: To illustrate the distribution shift from p ( D ) to ˜ p

Original: ( D ) , Figure 2 shows a set of graph, mechanism,and noise distribution samples in the R FF domain.
Revised: ( D ) , Figure 2 shows a set of graph, mechanism, and noise distribution samples in the R FF domain.

Original: The parameter configurations and functionsdefining p ( D ) and ˜ p ( D ) in the three domains are given in Appendix A. We also provide details onthe simulator by Dibaeinia and Sinha (2020) and subgraph extraction (Marbach et al., 2009).
Revised: In 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 Dibaeinia and Sinha (2020) and subgraph extraction (Marbach et al., 2009).

Original: 5.2 Evaluation Metrics
Revised: 5.2 Evaluation Metrics

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

Original: To asses how well a predictedcausal structure reflects the ground-truth graph, we report the structural Hamming distance (SHD) and the structural intervention distance (SID) (Peters and Bühlmann, 2015).
Revised: 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).

Original: While the SHD simply reflects the graph edit distance, the SID quantifies the closeness of two graphsin terms of their interventional adjustment sets.
Revised: While the SHD simply reflects the graph edit distance, the SID quantifies the closeness of two graphs in terms of their interventional adjustment sets.

Original: For these metrics, as well as for edge precision, recall,and F1 score, we collapse the posterior probabilities predicted by AVICI to hard predictions using adecision threshold of 0 . 5 .
Revised: For these metrics and for single-edge precision, recall, and F score, we convert the posterior probabilities predicted by AVICI to hard predictions using a threshold of 0 . 5 .

Original: In addition, we follow existing work and evaluate the uncertainty estimatesby computing the areas under the precision-recall curve (AUPRC) and the receiver operating characteristic (AUROC) (Friedman and Koller, 2003). 
Revised: 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 F [Figure]

Original:  More details on the metrics are given in Appendix B.
Revised: task instances. Error bars indicate the interquartile range. 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.

Original: 5.3 Inference Model Configuration
Revised: 5.3 Inference Model Configuration

Original: We train three inference models overall, one for each domain, and perform all experiments on thesethree trained models. Performance could be further improved by training models for more narrowsettings, e.g., specializing solely on inference from  observational  data.
Revised: 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.

Original: However, to demonstratethe versatility of AVICI, we use the same models for the evaluation with only observational dataas well as with a combination of observational and interventional data.

Original: The datasets sampled fromp ( D )during training have d = 2 to 50 variables and n = 200 samples, where with probability 0.5,these observations contain 50 interventional samples.
Revised: 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.

Original: All experiments throughout this paper areconducted on datasets that AVICI has never seen during training, regardless of whether we evaluatethe predictive performance in-distribution or o.o.d.

Original: The inference models in the three domains share identical hyperparameters for the architecture andoptimization, except for the dropout rate.
Revised: The inference models in the three domains share identical hyperparameters for the architecture and optimization, except for the dropout rate.

Original: We add the acyclicity constraint for the SCM domains L INEAR and R FF .
Revised: We add the acyclicity constraint for the SCM domains L INEAR and R FF .

Original: Details on the optimization and architecture are given in Appendix C.
Revised: Details on the optimization and architecture are given in Appendix C.

Original:  Experimental Results 
Revised:  Experimental Results 6.1 Out-Of-Distribution Generalization

Original: 6.1 Out-Of-Distribution Generalization

Original:  In the first part ofour experiments, we study the generalization capabilities of our inference modelsacross the spectrum of test distributions described in Section 5.1.This analysis consists of two parts. First, we test the model’s sensitivity to queries with o.o.d. data.
Revised: 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.1.

Original: Second, we investigate how well itgeneralizes to different problem sizes.

Original: Sensitivity to distribution shift We perform causal discovery from n = 1000 observations insystems of d = 30 variables.
Revised: We perform causal discovery from n = 1000 observations in systems of d = 30 variables.

Original: Starting from the training distribution p ( D ) , we incrementally introducethe described o.o.d. shifts in the causal structures, causal mechanisms, and finally noise, wherefully o.o.d.
Revised: Starting from the training distribution p ( D ) , we incrementally introduce the described distribution shifts in the causal structures, causal mechanisms, and finally noise, where fully o.o.d.

Original: corresponds to ˜ p
Revised: corresponds to ˜ p

Original: ( D ) .
Revised: ( D ) .

Original: The top row of Figure 3 visualizes the results of an empiricalsensitivity analysis.
Revised: The top row of Figure 3 visualizes the results of an empirical sensitivity analysis.

Original: The radar plots disentangle how combinations of the three o.o.d.
Revised: The radar plots disentangle how combinations of the three o.o.d.

Original: aspects, i.e.,graphs, mechanisms, and noise, affect the empirical performance of theAVICI models differently onthe three domains L INEAR , R FF , and G RN .
Revised: aspects, i.e., graphs, mechanisms, and noise, affect the empirical performance in the three domains L INEAR , R FF , and G RN .

Original: In addition to the metrics in Section 5.2, we also reportthe percentage of predicted graphs that are acyclic.
Revised: In addition to the metrics in Section 5.2, we also report the percentage of predicted graphs that are acyclic.

Original: To make the structural distances comparable,we use their ratios SID ratio :=

Original: SID in-dist. / SID and SHD ratio , respectively. Given this, all metrics arein a similar value range and higher values indicate better performance. As L INEARis the easiest domain, AVICI performs very well in all metrics and hardly suffers fromperformance degradation due to distribution shift.
Revised: In the L INEAR domain, AVICI performs very well in all metrics and hardly suffers under distribution shift.

Original: In contrast, G RN is the most challenging, withperformance degrading more significantly far o.o.d.
Revised: In contrast, G RN is the most challenging problem domain and the performance degrades

Original: We observe that AVICI can perform better undercertain distribution shifts than in-distribution, e.g., in G RN . This is because AVICI is empiricallybetter at predicting edges adjacent to large-degree nodes, a predominant feature of the GRN graphdistribution extracted from E. coli and S. cerevisiae that is not present in Erd˝os-Rényi graphs of thetraining distribution. The acyclicity constraint is perfectly satisfied for L INEAR and R FF . Prec.
Revised: [Table]

Original: [Figure]
Revised: [Table]

Original: SID [Figure] F
Revised: more significantly for the o.o.d.

Original: [Figure]
Revised: scenarios.

Original: F
Revised: We observe that AVICI can perform better under certain distribution shifts than in-distribution, e.g., in G RN .

Original: [Figure]
Revised: 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 ˝ os-Rényi training structures.

Original: Error bars indicate the interquartile range. We observe that AUPRC and AUROC do not suffer as much from distributional shift as the metricsbased on thresholded point estimates.
Revised: We also find that acyclicity is perfectly satisfied for L INEAR and R FF and that AUPRC and AUROC do not suffer as much from distributional shift as the metrics based on thresholded point estimates.

Original: This indicates that the marginal posterior edge probabilitiesconvey useful information about how certain AVICI is about individual edge predictions and howlikely it is that the thresholded prediction is wrong.
Revised: In Appendix E.1, we additionally report results for generalization from L INEAR to R FF and vice versa, i.e., to entirely unseen function classes of causal mechanisms in addition to the previous o.o.d.

Original: [-Overall, AVICI is surprisingly robust, even undersubstantial distribution shifts of all aspects of the data distribution at the same time. In Appendix E.
-]
Revised: [-Overall, AVICI is surprisingly robust, even undersubstantial distribution shifts of all aspects of the data distribution at the same time. In Appendix E.
-]shifts.

Original: Generalization to different problem sizes In addition to the sensitivity to distribution shift, weare interested in the abilityof our inference models to generalize to unseen problem sizes in termsof the number of variables d
Revised: Generalization to unseen problem sizes In addition to the sensitivity to distribution shift, we study the ability to generalize to unseen problem sizes.

Original: and the number of data points n .

Original: The bottom row of Figure 3 illustratesthe AUPRC for the edge predictions of AVICI when varying d and n on unseen in-distributiondata.
Revised: 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.

Original: The predictions improve with the number of data points  while exhibiting diminishing marginalimprovement when seeing additional data.
Revised: The predictions improve with the number of data points n while exhibiting diminishing marginal improvement when seeing additional data.

Original: Moreover, the performance decreases smoothly as thenumber of variables  increases and the task becomes harder.
Revised: Moreover, the performance decreases smoothly as the number of variables d increases and the task becomes harder.

Original: Most importantly, this robust behaviorcan be observed well beyond the settings  ( n = 200 and d ≤ 50 ) used during training.
Revised: Most importantly, this robust behavior can be observed well beyond the settings used during training ( n = 200 and d ≤ 50 ).

Original: 6.2 Benchmarking
Revised: 6.2 Benchmarking

Original: Next, we compare AVICI with existing causal discovery algorithms that use purely observational dataor a combination of observational and interventional data. The benchmark is performed on the fullyo.o.d. domain distributions ˜ p ( D ) , i.e., under distribution shifts on causal graphs, mechanisms, andnoise distributions w.r.t. the training distribution of AVICI.
Revised: Next, we benchmark AVICI against existing algorithms.

Original: For each class of methods, we considerat least one constraint-based and one score-based approach as well as one recent algorithm usingneural networks to model nonlinear mechanisms.

Original: In the observational setting, we compareAVICI with the PC algorithm (Spirtes et al., 2000), greedyequivalence search (GES) (Chickering, 2003),
Revised: Using only observational data, we compare with the PC algorithm (Spirtes et al., 2000), GES (Chickering, 2003),

Original: LiNGAM (Shimizu et al., 2006),and DAG-GNN (Yuet al., 2019).
Revised: LiNGAM (Shimizu et al., 2006), DAG-GNN (Yu et al., 2019), and GraN-DAG (Lachapelle et al., 2020).

Original: When using interventional data, we compare with  IGSP (Wang et al., 2017),greedyinterventional equivalence search (GIES) (Hauser and Bühlmann, 2012), and DCDI (Brouillard et al.,2020).
Revised: Mixed with interventional data, we compare with GIES (Hauser and Bühlmann, 2012), IGSP (Wang et al., 2017), and DCDI (Brouillard et al., 2020).

Original: For each domain, we independently calibrate important hyperparameters of each baselinemethod in a grid search on held-out task instances.
Revised: We tune the important hyperparameters of each baseline on held-out task instances of each domain.

Original: When computing the evaluation metrics, we favormethods that only predict (interventional) Markov equivalence classes by orienting undirected edgescorrectly when present in the ground truth.
Revised: 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.

Original: Details on the baselines are given in Appendix D.
Revised: Details on the baselines are given in Appendix D.

Original: [Table]
Revised: The benchmarking is performed on the fully o.o.d. domain distributions ˜ p ( D ) , i.e., under distribution shifts on causal graphs, mechanisms, and noise distributions w.r.t. the training distribution of AVICI.

Original: Table 1 shows the SID and F1 scores of all methods given n = 1000 observations for d = 30 variables.
Revised: Table 1 shows the SID and F1 scores of all methods given n = 1000 observations for d =variables.

Original: The AVICI model trained on L INEAR outperforms all baselines, both givenonly observational andin combination with interventional data, despite operating under significant distribution shift.
Revised: The AVICI model trained on L INEAR outperforms all baselines, both given observational or interventional data, despite operating under significant distribution shift.

Original: Only GIES achieves comparable accuracy.
Revised: Only GIES achieves comparable accuracy.

Original: The same holds for R FF , where constraint-based and nonlinearapproaches outperform GES and GIES, whose scores assume linearity, but ultimately do not reachthe accuracy set by AVICI.
Revised: The same holds for R FF , where GraN-DAG and DCDI perform well but ultimately do not reach the accuracy of AVICI.

Original: In the G RN domain, the inductive bias encoded by classical methodsreaches its limits.
Revised: In the G RN domain, where inductive biases are most difficult to specify, classical methods fail to infer plausible graphs.

Original: However, provided interventional data, AVICI can use its learned inductive biasto infer plausible causal structures from the noisy gene expressions, even under distribution shift,which is a promising step towards reliable structure discovery in fields like molecular biology.
Revised: 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.

Original: Evenwithout gene knockout data,the edge probabilities predicted by AVICI achieve nontrival AUROCand AUPRC while classical methods predict close to randomly (Table 7 in Appendix E; see also 
Revised: 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 [Table] m [Table]

Original: Dibaeinia and Sinha 2020; Chen and Mar 2018).
Revised: i r i c a l p [Table] i [Table]

Original: Results for d = 100 variables are given in Appendix E.2.
Revised: j = [Table] D [Table]

Revised: L INEAR

Revised: R FF

Revised: G RN

Revised: [Table]

Revised: [Figure]

Revised: 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).

Revised: We say that a predicted probability p is calibrated if we empirically observe an event in ( p · 100)% of the cases.

Revised: When plotting the observed edge frequencies against their predicted probabilities, a perfectly calibrated algorithm induces a diagonal line.

Revised: The expected calibration error (ECE) represents the weighted average deviation from this diagonal.

Revised: For further explanations, see Appendix B.

Revised: 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).

Revised: We additionally compare AVICI with DiBS, which infers Bayesian posterior edge probabilities like

Revised: AVICI (Lorch et al., 2021).

Revised: Figure 4 gives the calibration plots for AVICI and Table 4b the ECE for all methods.

Revised: In each domain, the marginal edge probabilities predicted by AVICI are the most calibrated in terms of ECE.

Revised: Moreover, Figure 4a shows that AVICI closely traces the perfect calibration line, which highlights its accurate uncertainty calibration across the probability spectrum.

Original: In Appendix E.3 and E.4, we also reportbenchmarking results on a real-world proteomics dataset (Sachs et al., 2005) and give additionalmetrics that evaluate the uncertainty quantification ofAVICI and the baselines.
Revised: In Appendix E.5, we additionally report AUROC and AUPRC metrics for all methods. We also provide calibration plots analogous to Figure 4 for the baselines (Figure 6), which often show vastly overconfident predictions where the calibration line is far below the diagonal. 6.3 Ablations Finally, we analyze the importance of key architecture components of the inference network f ϕ .

Revised: Focusing on the R FF domain, we train several additional models and ablate single architecture components.

Revised: We vary the network depth L , the axes of attention, the representation of θ , and the number of training steps for ϕ .

Revised: All other aspects of the model, training and data simulation remain unchanged.

Revised: Table 2 summarizes the results.

Revised: 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 2 L self-attention layers.

Revised: Attending only over axis n is not sensible since variable interactions are not processed until the prediction of θ , but we still include the results for completeness.

Revised: We also test an alternative variational parameter model given by θ i,j = ϕ

Revised: ⊤ θ tanh (cid:0) ϕ

Revised: ⊤ u u i + ϕ ⊤ v v j (cid:1) that uses an additional, learned vector ϕ θ and matrices ϕ u , ϕ v .

Revised: This model has been used in related causal discovery work for searching over high-scoring causal DAGs (Zhu et al., 2020) and is an instance of a relational network (Santoro et al., 2017).

Revised: This variant also satisfies permutation equivariance (cf.

Revised: Section 4.2) since it applies the same MLP elementwise to each edge pair [ u i , v j ] .

Revised: Ultimately, we

Revised: [Table]

Revised: find no statistically significant difference in performance to our simpler model in Eq.

Revised: (6), hence we opt for less parameters and a lower memory requirement.

Revised: Lastly, Table 2 shows that the causal discovery performance of AVICI scales up monotonically with respect to network depth and training time.

Revised: Even substantially smaller models of L = 4 or shorter training times achieve an accuracy that is on par with most baselines (cf.

Revised: Table 1).

Original: {+Our main models ( ⋆ ) have a moderate size of 4 . 2 × 10 6 parameters, which amounts to only 17 . 
+}
Revised: {+Our main models ( ⋆ ) have a moderate size of 4 . 2 × 10 6 parameters, which amounts to only 17 . 
+}

Revised: MB at f 32 precision.

Revised: Computing test-time predictions via a forward pass takes only a few CPU seconds.

Original:  Discussion
Revised:  Discussion

Original: We proposed AVICI, a method for inferring causal structure by performing amortized variationalinference over an arbitrary data-generating distribution.
Revised: We proposed AVICI, a method for inferring causal structure by performing amortized variational inference over an arbitrary data-generating distribution.

Original: Our approach leverages the insight thatinductive biases crucial for statistical efficiency in structure learning might be more easily encoded ina domain simulator than in an inference technique.
Revised: 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.

Original: This is reflected in our experiments, where AVICIsolves structure learning problems in complex domains intractable for existing approaches (Dibaeiniaand Sinha, 2020).
Revised: This is reflected in our experiments, where AVICI solves structure learning problems in complex domains intractable for existing approaches (Dibaeinia and Sinha, 2020).

Original: Our method can likely be extended to other typically difficult domains, includingstructure learning settings where we cannot assume causal sufficiency (Bhattacharya et al., 2021).
Revised: Our method can likely be extended to other typically difficult domains, including settings where we cannot assume causal sufficiency (Bhattacharya et al., 2021).

Revised: Our approach will continually benefit from ongoing efforts in developing (conditional) generative models and domain simulators.

Original: Using AVICI  comes with several trade-offs.
Revised: Using AVICI still comes with several trade-offs.

Original: First, while optimizing the dual program empiricallyinduces acyclicity, this constraint is not satisfied with certainty.
Revised: First, while optimizing the dual program empirically induces acyclicity, this constraint is not satisfied with certainty using the variational family considered here.

Original: Moreover, similar to most amortizedinference techniques (Amos, 2022), AVICI gives no theoretical guarantees of performance.
Revised: Moreover, similar to most amortization techniques (Amos, 2022), AVICI gives no theoretical guarantees of performance.

Original: Someclassical methods can do so in the infinite sample limit given specific assumptions on the datagenerating process (Peters et al., 2017).
Revised: Some classical methods can do so in the infinite sample limit given specific assumptions on the data-generating process (Peters et al., 2017).

Original: However, future work might achieve guarantees for AVICI thatare similar to learning theory results for the bivariate causal discovery case (Lopez-Paz et al., 2015).
Revised: 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).

Original: Finally, while causal discovery with AVICI amounts to a single forward pass of f ϕ at inference time,the inference models can be intensive to train.

Original: To mitigate this training cost and make our resultsaccessible to other researchers, we release the trained models studied in this work.

Original: Our experimentsdemonstrate that our inference models are highly robust to distributional shift, suggesting that thetrained models could be useful out-of-the-box in causal structure learning tasks outside the domainsstudied in this paper.
Revised: Nevertheless, 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.

Original: In this context, fine-tuning a pretrained AVICI model on labeled real-worlddatasets is a promising avenue of future work.
Revised: In this context, fine-tuning a pretrained AVICI model on labeled real-world datasets is a promising avenue for future work.

Revised: Our code and models are publicly available at: https://github.com/larslorch/avici .

Revised: Acknowledgments and Disclosure of Funding

Revised: We thank Alexander Neitz, Giambattista Parascandolo, and Frederik Träuble for their feedback and the reviewers for their helpful comments.

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

Revised: Jonas Rothfuss was supported by an Apple Scholars in AI/ML fellowship.