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

Original: Anonymous Author(s)
Revised: Anonymous Author(s)

Original: AffiliationAddress email
Revised: AffiliationAddress email

Original: Abstract
Revised: Abstract

Original: Learning causal structure poses a combinatorial search problem that typically involves evaluating structures using a score or independence test.
Revised: Learning 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 searchis costly, and designing suitable scores or tests that capture prior knowledge isdifficult.

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

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

Original: Our inference model acquires domain-specific inductive bias for causal discovery solely from data generated by a simulator.

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

Original: Moreover,the architecture of our inference model is permutation invariant w.r.t. the datapoints and permutation equivariant w.r.t.
Revised: Instead, our inference model acquires domain-specific inductive biasesfor causal discovery solely from data generated by a simulator.

Original: the variables, facilitating generalization tosignificantly larger problem instances than seen during training.
Revised: 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.

Original: On synthetic dataand semi-synthetic gene expression data, our models exhibit robust generalizationcapabilities under substantial distribution shift and significantly outperform existingalgorithms, especially in the challenging genomics domain.
Revised: 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.

Original:  Introduction
Revised:  Introduction

Original: A fundamental task in various scientific disciplines is to learn the causal structure among thevariables we observe  (Spirtes et al., 2000; Pearl, 2009).
Revised: Learning the causal structure among the variables we observe is a fundamental task in variousscientific 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 fromobservations 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 causalstructures, 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 independencies (Heinze-Deml et al., 2018).
Revised: 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).

Original: A key limitation of search approaches to causal discovery is the difficulty of specifying realisticinductive biases.
Revised: A key limitation of search approaches to causal discovery is the difficulty of specifying realisticinductive biases.

Original: On the one hand, score-based methods use strong assumptions like linearity(Shimizu et al., 2006), specific noise models (Hoyer et al., 2008; Peters and Bühlmann, 2014) or theabsence of measurement error (cf.
Revised: 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.

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

Original: On the other hand, constraint-based methods do nothave enough domain-specific inductive bias, since even with oracle conditional independence tests,they are limited to identifying equivalence classes that may be exponentially large (He et al., 2015).
Revised: Conversely, constraint-based methodsdo not have enough domain-specific inductive bias. Even with oracle conditional independence tests,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 (Colomboet al., 2014).

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

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

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

Original: Our key insight is  the dichotomy that, given a causal model, it is easy to generaterealistic data—the inverse direction, i.e., inferring the causal structure from data, is hard.
Revised: 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.

Original: Manyhigh-impact domains of causal discovery can be realistically forward-simulated, for example, thechemical 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), or chemicalkinetics (Anderson and Kurtz, 2011; Wilkinson, 2018).
Revised: 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).

Original: Leveraging such simulators for generatingtraining sets, we can use amortized inference  to both acquire useful domain-specific inductive biasesnot easily represented by a graph scoring function and bypass the problems of structure search.
Revised: We can leverage such simulators to trainan inference model in causal discovery for specific data-generating processes. 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.

Original: The central challenge of amortizing causal discovery consists in designing a learning rule thatgeneralizes out-of-distribution.
Revised: The central challenge of amortizing causal discovery consists in designing a learning rule thatgeneralizes out-of-distribution.

Original: This generalization is essential because any simulated datadistribution is unlikely to cover all real-world distribution shifts.
Revised: This generalization is essential because any simulated data distributionis unlikely to cover all real-world distribution shifts.

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

Original: The model architecture of AVICI is permutation invariant and equivariant with respectto the observation and variable dimensions, respectively, generalizing to significantly larger probleminstances than seen during training.
Revised: The model architecture of 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.

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

Original: Moreover,we demonstrate that our inference model induces calibrated uncertainties and robust behavior undersubstantial distribution shifts of graphs, mechanisms, noise, and problem sizes, which makes ourtrained models both reliable and versatile for future downstream use.
Revised: 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.

Original: In particular, AVICI was theonly one to infer plausible causal structure from realistic gene expression data, pushing the boundaryof what is possible in structure discovery of fields like molecular biology.
Revised: 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.

Original:  Background and Related Work2.1 Causal Structure
Revised:  Background and Related Work2.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 define a causal structure in the same way as Mooij et al.

Original: Other popularmodels such as causal Bayesian networks and structural causal models (Pearl, 2009; Peters et al. ,2017) are less well-suited to describing systems with feedback loops, which we consider practicallyrelevant.
Revised: 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: We define the causal structure G of a directed graph whose nodes correspond to d variablesx = ( x 1 , . . .
Revised: Specifically, we define the causal structure G as the directed graph whose nodes correspond to dvariables x = ( x 1 , . . .

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

Original: A variable x i hasa direct causal effect on x j if intervening on x i affects the outcome of x j independent of the othervariables x \ ij := x \{ x i , x j } , i.e., there exists a ̸ = a ′ such that
Revised: 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 : = x \{ x i , x j } , i.e., there exists a ̸ = a ′ such that

Original: [Equation]
Revised: [Equation]

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

Original: An intervention do ( · ) denotes any active manipulation of the generative process of theobservations x , like gene knockdowns in genomics, where the transcription rate of a gene is externallyreduced.
Revised: An intervention do ( · ) denotes any active manipulation of the generative process of theobservations x . An example of interventions are gene knockouts in genomics, in which transcriptionrates of genes are externally set to zero.

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

Original: We emphasize that G is not necessarily acyclic, however we later show how to apply our approachwhen G is constrained to be acyclic.
Revised: We emphasize that G is not necessarily acyclic; however, we latershow how to apply our approach when G is constrained to be acyclic.

Original: In this work, we assume causal sufficiency, i.e., thatthere are nounobserved variables beyond x that are causal parents of more than one x i (Peters et al., 2017).
Revised: In this work, we assume causalsufficiency, 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  search over causal graphs and evaluate them using a likelihood or conditionalindependence tests (Chickering, 2003; Kalisch and Bühlman, 2007;Hauser and Bühlmann, 2012;
Revised: 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;

Original:  Zheng et al., 2018; Heinze-Deml et al., 2018).
Revised: Hauser and Bühlmann, 2012; Zheng et al., 2018; Heinze-Deml et al., 2018).

Original: Others combine constraint and score-based ideas(Tsamardinos et al., 2006) or use  noise properties of  the SCM (Shimizu et al., 2006; Hoyer etal.,2008). In addition, there is work on leveraging multi-domain data (Tian and Pearl, 2001; Tillmanand Eberhardt, 2014; Ghassami et al., 2018; Zhang et al., 2017b; Huang et al., 2019, 2020).
Revised: 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).

Original: However,these multi-domain approaches assume a shared causal model across datasets and thus do not studythe problem of generalizing to unseen causal systems.

Original: Deep learning has been applied to causal inference, e.g., for the estimation of treatment
Revised: Deep learning has been applied to causal inference, e.g., for the estimation of treatment

Original: effects(Shalit et al., 2017;
Revised: effects(Shalit et al., 2017;

Original: Louizos et al., 2017; Yoon et al., 2018) and in instrumental variable analysis(Hartford et al., 2017; Bennett et al., 2019).
Revised: Louizos et al., 2017; Yoon et al., 2018) and in instrumental variable analysis(Hartford et al., 2017; Bennett et al., 2019).

Original: In causal structure learning, neural networks have mostlybeen used  to model nonlinear causal mechanisms among the variables (Goudet et al., 2018; Yuet al., 2019;
Revised: In causal structure learning, neural networks havebeen used primarily to model nonlinear causal mechanisms among the variables (Goudet et al. ,2018; Yu et al., 2019;

Original: Lachapelle et al., 2020; Brouillard et al., 2020;Ke et al., 2020; Lorch et al., 2021).
Revised: Lachapelle et al., 2020; Brouillard et al., 2020; Lorch et al., 2021).

Original: While amortized inference has gained increasing popularity in recent years (Amos, 2022), prior work applying it to causal discovery only studied narrowly defined subproblems like the bivariate caseor fixed causal mechanisms (Lopez-Paz et al., 2015; Löwe et al., 2022).
Revised: Whileamortized inference has gained increasing popularity in recent years (Amos, 2022), prior workapplying it to causal discovery only studied narrowly defined subproblems such as the bivariatecase (Lopez-Paz et al., 2015) and fixed causal mechanisms for each task (Löwe et al., 2022) or onlyused observational correlation coefficients for prediction (Li et al., 2020).

Original: In parallel work, Ke et al.(2022) independently perform causal discovery as a supervised learning task.
Revised: In parallel work, Ke et al.(2022) independently perform causal discovery as a supervised learning task.

Original: There are significantdifferences in their approach, in particular, our inference models generalizes to significantly largerproblem sizes, even when trained on realistic genomics data.
Revised: There are significant differences in their approach, in particular, our inference models generalize to significantly largerproblem sizes, even when trained on realistic genomics data.

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

Original: 3.1 The Variational Objective

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

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

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

Original: Itis important to highlight that p ( D | G ) does not need to represent an ancestral sampling procedureover a directed acyclic graph.
Revised: It is important to highlightthat p ( D | G ) does not need to represent an ancestral sampling procedure over a directed acyclicgraph.

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

Original: Rather, p ( D | G ) characterizes a data-generating process whose causal graph G formalizesall direct causal effects (Eq. 1) of the system’s mechanisms—how intervening on certain variablesinfluences the realizations of others (Mooij et al., 2016; Peters et al., 2017; Rubenstein et al., 2017).
Revised: 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: 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 ; θ ) .

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  forward KL divergence from the intractable posteriorp ( G | D )  for D ∼ p
Revised: 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

Original: ( D ) (Ambrogioni et al., [Equation]
Revised: ( D ) :

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 arrive at an equivalent, tractable objective.
Revised: Agakov (2004) and rewrite the expected forward KL to obtain an equivalent, tractable objective:

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 ))]to amortize variational inference for causal discovery (AVICI).
Revised: [log q ( G ; f ϕ ( D ))] ,which allows us to amortize variational inference for causal discovery (AVICI).

Original: While the data-generating distributionp ( D ) =
Revised: While the datagenerating distribution p ( D ) =

Original: (cid:80) G p (D | G )  p (  G ) can be arbitrarily complex, thetraining objective L is tractable wheneverwe have access to the ground truth causal graph G underlying the generative process of D , i.e., tosamples from the joint distribution p ( G, D ) .
Revised: E p ( G ) [ p ( D | G )] can be arbitrarily complex, the objective L istractable whenever we have access to the ground-truth causal graph G underlying the generativeprocess 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 byuse of a simulator for the domain of interest.
Revised: In practice, p ( G ) and p ( D | G ) canthus be specified by a simulator for the domain of interest.

Revised: Since the quality of simulators in manydomains may improve substantially over time, the sim-to-real bias of the inference model f ϕ canreduce along the way.

Original: From an information-theoretic view, the objective (2) maximizes a variational lower bound on themutual information
Revised: From an information-theoretic view, the objective (2) maximizes a variational lower bound on themutual information

Original: I [ G ; D ] between the causal structure G and the observations D (Barber and
Revised: 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 p ( D ) D KL ( p ( G | D )
Revised: The bound is tight if E p ( D ) D KL ( p ( G | D )

Original: ∥ q ( G ; f ϕ ( D )))
Revised: ∥ q ( G ; f ϕ ( 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 requires optimizing 124 E p ( D | G )
Revised: The AVICI objective in (3) intentionally targets the forward KL D KL ( p ∥ q ( · ; θ )) , which requiresoptimizing 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 andhave access to samples from the true data-generating distribution.
Revised: This choice implies that we both model the density q ( G ; θ ) explicitly and have access to samples from the true data-generating distribution.

Original: As long as samples from p ( G, D )are available for estimating L ( ϕ ) , minimizing the forward KL in (2) enables us to infer causalstructures in arbitrarily complex domains—that is, even domains where it is difficult to specify anexplicit functional form for p ( D | G ) .
Revised: As long as samples fromp ( G, D ) are available to estimate L ( ϕ ) , minimizing the forward KL in (2) enables us to infer causalstructures in arbitrarily complex domains—that is, even domains where it is difficult to specify anexplicit likelihood p ( D | G ) .

Original: In addition, the forward KL typically yields more reliableuncertainty estimates since it does not suffer from the variance underestimation problems commonfor the backward KL (MacKay, 2003).
Revised: 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).

Original: Access to p ( G, D  ) is natural when a forward-model of the 
Revised: 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 )] (Blei et al., 2017). This objective

Original: [Figure]
Revised: [Figure]

Original: x L

Original: [Figure]
Revised: [Figure]

Original: domain exists, e.g., when sampling D synthetically from a structural equation model or a specializeddomain simulator. We expect the quality of simulators in many domains to improve substantiallyover time, reducing the sim-to-real bias of the inference model f ϕ . In contrast, variational inference optimizes the reverse KL D KL ( q ∥ p ) with respect to θ , whichinvolves the reconstruction term E q ( G ; θ ) [log p ( D | G )] (Blei et al., 2017).
Revised: [Figure]

Original: This objective requires aparametric, tractable marginal likelihood p ( D | G ) .
Revised: [Figure] [Figure] [Figure] requires a tractable marginal likelihood p ( D | G ) .

Original: Unlessjointly inferring the mechanism parameters(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 is not warranted in practice (Friston et al., 2000; Schaffteret al., 2011; Runge et al., 2019; Dibaeinia and Sinha, 2020).
Revised: 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; Schaffteret al., 2011; Runge et al., 2019; Dibaeinia and Sinha, 2020).

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

Original: 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 inferencemodel f ϕ that predicts θ given D .

Original: After that, we detail our training procedureof AVICI for optimizingthe model parameters ϕ as well as for learning the inductive bias of acyclicity.
Revised: After that, we detail our training procedure for optimizing themodel parameters ϕ as well as for learning the inductive bias of acyclicity.

Original:  Inference Model
Revised:  Inference Model

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  (3), we optto 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), weopt to use a factorized variational family in this work.

Original: [Equation]
Revised: [Equation]

Original: The inference model f ϕthus maps a data set 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.
Revised: , o n } to a d -by- dmatrix θ parameterizing the variational approximation of the causal graph posterior.

Original: In addition to the jointobservation x i =
Revised: In addition tothe 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 ) contains interventional informationfor each variable when available.
Revised: , o id ) may contain interventionalinformation for each variable.

Original: For example, in the case of gene knockout interventions, we haveo i j = ( x i j , u i j ) and u i j ∈ { 0 , 1 } indicating whether variable j was intervened upon in sample i .
Revised: 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 .

Original: Othersettings could be encoded analogously, e.g., when the intervention targets are unknown.
Revised: Othersettings could be encoded analogously, e.g., when the intervention targets are unknown.

Original: 4.2 Model Architecture
Revised: 4.2 Model Architecture

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

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

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

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

Original: In addition, f ϕ should be permutation equivariant acrossthe variable dimension (axis d ).
Revised: Moreover, f ϕ should be permutation equivariant acrossthe 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 edgeprobabilities, 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: We parameterize f ϕ as adeep neural networkand introduce a novel architecture that satisfies theabove invariants.
Revised: In the following, we show how to parameterize f ϕ as a neural network that encodes these properties.

Original: After  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 forthe 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 overa 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.
Revised: Figure 1 illustrates the key components of the architecture.

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

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

Original: To enable information flow across all n × d tokens of the representation, themodel alternates in attending over the observation and the variable dimension.
Revised: 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).

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

Original: This design keeps memoryrequirements in check andresults in a moderate model size because parameters are shared across the batched axes.
Revised: 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: 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).

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

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

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

Original: (2021), we use two linear layers to map each z i into embeddingsu i , v i ∈ R k and finally model the probability of each edge in the causal graph as:
Revised: (2021), we use two position-wiselinear layers to map each z i to two embeddings u i , v i ∈ R k . 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:

Original: [Equation]
Revised: [Equation]

Original:  where σ is the logistic function.
Revised: (6) 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 since the graphmodel in (6) permutes with respect to axis d , the  model  satisfies the invariants of Section 4.2.
Revised: Since max-pooling is invariant to permutations and since the 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 .

Original: 4.3 Acyclicity
Revised: 4.3 Acyclicity

Original: There is no reason to generally exclude feedback loops in G .
Revised: There is no reason to generally exclude feedback loops in G .

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

Original: Nonetheless, certain domains may bemore accurately modeled by acyclic causal structures (Peters et al., 2017; Rubenstein et al., 2017).
Revised: However, certain domains may bemore accurately modeled by acyclic causal structures (Peters et al., 2017; 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) does not allow enforcing it directly, we can optimize for acyclicitythrough ϕ .

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) withthe constraint that q only models acyclic graphs in expectation: [Equation]

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

Original: We use power iteration toapproximate and differentiate through the largest eigenvalue of f ϕ ( D, I ) (Golub and Van der Vorst,2000; Lee et al., 2019a):
Revised: We use power iteration toapproximate 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:

Original: [Equation]
Revised: [Equation]

Original:  and a , b ∈ R d are randomly initialized.
Revised: 8) 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 than alternative 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 alternative 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: 4.4 Optimization
Revised: Optimization

Original: Algorithm Parameters: ϕ variational, λ dual, η step size while not converged do for l steps do
Revised: while not converged do for l steps do

Original: [Equation]
Revised: [Equation]

Original: Combining the objective in (3) with our inference model (5), we can directly use stochasticoptimization to train the parameters ϕ of ourneural network model.
Revised: Combining the objective in (3) with our inference model (5), we can directly use stochasticoptimization to train the parameters ϕ of ourneural network model.

Original: The expectations overp ( G, D ) inside L and F are approximated using samples from the data-generating processof the domain, e.g, a simulator.
Revised: The expectations overp ( G, D ) inside L and F are approximated using samples from the data-generating processof the domain, e.g., a simulator.

Original: In practice,the available data can be further augmentedusing bootstrapping.
Revised: In practice,the available data can be further augmentedusing bootstrapping.

Original: In case we impose an acyclicity prior on G , we solve the constrained programmax[Equation] ϕ [Equation]
Revised: When imposing an acyclicity prior on G , we solve the constrained programmax ϕ L ( ϕ ) s.t. F ( ϕ )

Original: t. F

Original: [Equation]

Original: ϕ )

Original: = 

Revised: =

Original: through its dual formulation (Nandwani et al., 2019):
Revised: through its dual formulation (Nandwani et al., 2019):

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 localoptimum 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 acyclicityconstraint, training reduces to the primal updates of ϕ with λ = 0 .

Original: Erdos-Renyi
Revised: Evaluating causal discovery algorithms is difficult since there are few interesting real-world datasetsthat come with ground-truth causal structure. Often, the believed ground truths may be incompleteor change as expert knowledge improves (Schaffter et al., 2011; Mooij et al., 2020). 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: In Appendix E, we additionally reportresults on a real-world proteomics dataset (Sachs et al., 2005).

Original: Scale-free
Revised: 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: Experimental Setup

Original: Evaluating causal discovery algorithms is difficult since there are there few interesting real-worlddata sets that come with ground-truth causal structure.

Original: Often the believed ground truths may beincomplete or change as expert knowledge improves (Schaffter et al., 2011).

Original: Following prior work,we deal with this issue by using simulated data with known causal structure.

Original: 5.1 Domains and Simulated Components
Revised: 5.1 Domains and Simulated Components

Original: Our goal is to evaluate our inference model’s ability to learn inductive biases in previously intractabledomains and to generalize robustly beyond the data distribution p ( D ) seen during training.
Revised: 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.

Original: Weconsider three domains: two classes of structural causal models (SCMs) as well as semi-syntheticsingle-cell expression data of gene regulatory networks (GRNs).
Revised: We studythree domains: two classes of structural causal models (SCMs) as well as semisynthetic single-cellexpression data of gene regulatory networks (GRNs).

Original: To study generalization, wecarefully construct a test distribution ˜ p
Revised: To study the generalization of AVICI beyondthe training distribution p ( D ) , we carefully construct a test distribution ˜ p

Original: ( D ) that incurs substantial shift from p ( D ) in terms of theunderlying causal structures, mechanisms, and noise components.
Revised: ( D ) that incurs substantialshift from p ( D ) in terms of the underlying causal structures, mechanisms, and noise components.

Original: Data-generating processes p ( D | G )
Revised: Data-generating processes p ( D | G )

Original: 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 endogenous noisedistributions p j ( ϵ j ) for each variable.
Revised: 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.

Original: We consider the two domains of linear functions (L INEAR ) andnonlinear functions of random Fourier features (R FF ), which correspond to samples from a Gaussianprocess with squared exponential kernel (Rahimi and Recht, 2007).
Revised: 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 Gaussian process with squared exponential kernel (Rahimi and Recht, 2007).

Original: For the out-of-distribution (o.o.d.)setting ˜ p ( D ) , we draw the linear functionweights andthe kernel scales from the tails of p ( D ) oreven unseen value ranges.
Revised: In the out-of-distribution(o.o.d.) setting ˜ p ( D ) , we sample the linear function and kernel parameters from the tails of p ( D )and unseen value ranges.

Original: While using additive homoscedastic Gaussian noise for p ( D ) , we employheteroscedastic Cauchy or Laplacian noise in ˜ p
Revised: Moreover, while simulating additive, homoscedastic Gaussian noise ϵ jin the training distribution p ( D ) , we employ heteroscedastic Cauchy and Laplacian noise o.o.d.

Original: ( D ) with f j ( x Pa ( j ) , ϵ j ) =
Revised: ,where f j ( x Pa ( j ) , ϵ j ) =

Original: f j ( x pa ( j ) )
Revised: f j ( x pa ( j ) )

Original: + h j ( x pa ( j ) )
Revised: + h j ( x pa ( j ) )

Original: ϵ jfor some random function h j .
Revised: ϵ j for some random, input-dependent function h j .

Original: Interventions are performed by fixing selected variables x j to arandom value during sampling.
Revised: Interventions are performed by clipping a variable x j to a random value during data generation.

Original: In addition to SCMs, we consider the challenging domain of GRNs (G RN ) using the simulator by
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-cellsamples of gene expressions correspondto draws from the steady-state of a stochastic dynamical system that varies across cell types(Huynh-Thu and Sanguinetti, 2019).
Revised: 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).

Original: Like L INEAR and R FF , the simulator parameters consideredo.o.d. are drawn from significantly wider ranges than for p ( D ) .
Revised: 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 modelstechnical noise effects, resulting in  highly zero-inflated  matrices D .
Revised: 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: When moving o.o.d.

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

Original: In G RN , interventions correspond to gene knockouts and are performed by forcingthe expression rate of a variable x j to zero.
Revised: In G RN , interventions correspond togene knockouts and are performed by forcing the transcription rate of a variable x j to zero.

Original: Whenever a method infers  structure from interventional data in our experiments, half of the dataset consists of observational data and half of single-variable interventions.
Revised: 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 singlevariable interventions are performed on arandom subsetcontaining 50% of  the nodes.
Revised: 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 RNis more challenging, the interventional data here contains single-node interventions on all variables.
Revised: Since G RN is more challenging, the interventional data here contains single-nodeinterventions on all variables.

Revised: 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  to generate groundtruth causal structures.
Revised: Following prior work, we use random graph models and known biologicalnetworks to generate ground-truth causal structures.

Original: In all three domains, the training data distribution p ( D ) is induced by simple Erd˝os-Rényi and scale-free causal structures of various degrees (Erd˝os and Rényi, 1959; Barabási and Albert, 1999).
Revised: In all three domains, the training data distributionp ( D ) is induced by simple Erd˝os-Rényi and scale-free graphs of various degrees (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, the data distributions ˜ p

Original: ( D ) of the L INEAR and
Revised: ( D ) of the L INEAR and

Original: R FF domains are simulated using causal structures from different 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: R FF domains are simulated using causal structures from unseen 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).

Original: In the G RN domain, weuse subgraphs of the known S. cerevisiae and E. coli GRNs and the effect signs whenever known.
Revised: In the G RN domain, weuse 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 structuralpatterns like motifs and modularity (Ravasz et al., 2002; Shen-Orr et al., 2002).

Original: To illustrate the  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 noisedistribution 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: Details on the simulator by Dibaeinia and Sinha (2020) andsubgraph extraction are given in Appendix A. Thespecific parameter configurations and functionsdefining p ( D ) andall ˜ p ( D ) in the three domains are provided in Appendix B.
Revised: 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).

Original: 5.2 Evaluation Metrics
Revised: 5.2 Evaluation Metrics

Original: To asses how well the the predicted causal structure reflects the ground truth graph, we report thestructural hamming distance (SHD) as well as the structural intervention distance (SID) (Petersand Bühlmann, 2015).
Revised: To asses how well a predicted causal structure reflects the ground-truth graph, 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 quantifiesthe closeness of two graphs in terms of their interventional adjustment sets.
Revised: While the SHD simply reflects the graph edit distance, the SID quantifies the closeness of two graphsin terms of their interventional adjustment sets.

Original:  In addition, we followexisting work and evaluate thepredicted uncertainty estimates q ( g i,j ; θ i,j ) over the set of possibleedges by computing the areas under the precision-recall curve (AUPRC) and the receiver operatingcharacteristic (AUROC) (Friedman and Koller, 2003; Lorch et al., 2021). For comparison withclassical methods, we also collapse these uncertainties to hard predictions using a decision thresholdof 0.5 and compute precision, recall, and the F1 score.
Revised: 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 . 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).

Original: More details  are given in Appendix B.
Revised: 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.
Revised: We train three inference models overall, one for each domain, and perform all experiments on thesethree trained models.

Original: Performance could be further improved, e.g., by training models for more narrow settings specializingthem solely on inference from observational data.
Revised: Performance could be further improved by training models for more narrowsettings, e.g., specializing solely on inference from observational data.

Original: However, to demonstratethe versatility of AVICI, we use the same modelsboth for  evaluation with only observational dataas well as combinations of observations and interventional data.
Revised: 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.

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

Original: All  inference models  share identicalhyperparameters for the architecture and optimization except  the dropout rate.In held-out tasks of R FF and G RN , we found that dropout does not hurt performance in-distribution, so we increased thedropout rates from 0.0 to 0.1 and 0.3, respectively, to help o.o.d.
Revised: 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. The inference models in the three domains share identical hyperparameters for the architecture andoptimization, except for the dropout rate.

Original: generalisation.

Original: We add the acyclicityconstraint 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: The data sets from p ( D ) during training haved = 5 to 50 variables and n = 200 samples, where with probability 0.5, the observations contain50 interventional samples.

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 Results6.1 Out-Of-Distribution Generalization
Revised:  Experimental Results

Revised: 6.1 Out-Of-Distribution Generalization

Original: We study the generalization capabilities of our inference models across the spectrum of test distributions described in Section 5.1.
Revised: In the first part of our experiments, we study the generalization capabilities of our inference modelsacross the spectrum of test distributions described in Section 5.1.

Original: This analysis consists of two parts.
Revised: This analysis consists of two parts.

Original: First, we consider the model’ssensitivity to queries with o.o.d.
Revised: First, we test the model’s sensitivity to queries with o.o.d.

Original: data.
Revised: data.

Original: Second, we investigate how well it generalizes to differentproblem sizes.
Revised: Second, we investigate how well itgeneralizes to different problem sizes.

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

Original: Starting from the training distribution p ( D ) , we incrementally introducethe described o.o.d.
Revised: Starting from the training distribution p ( D ) , we incrementally introducethe described o.o.d.

Original: shifts in the causal structures, causal mechanisms, and finally noise, wherefully o.o.d.
Revised: shifts in the causal structures, causal mechanisms, and finally noise, wherefully 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 empirical o.o.d.sensitivity analysis.
Revised: The top row of Figure 3 visualizes the results of an empiricalsensitivity 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 the AVICI models differently onthe three domains L INEAR , R FF , and G RN .
Revised: aspects, i.e.,graphs, mechanisms, and noise, affect the empirical performance of the AVICI models differently onthe 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 reportthe percentage of predicted graphs that are acyclic.

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

Original: SID in-dist.
Revised: SID in-dist.

Original: / SID and SHD ratio , respectively.
Revised: / SID and SHD ratio , respectively.

Original: Given this, all metrics arein a similar value range and higher values indicate better performance.F [Figure]
Revised: Given this, all metrics arein a similar value range and higher values indicate better performance.

Original: F [Figure] Error bars indicate the interquartile range. As L INEAR is the easiest domain, AVICI performs very well across all metrics and hardly suffersfrom performance degradation due to distribution shift.
Revised: As L INEAR is the easiest domain, AVICI performs very well in all metrics and hardly suffers fromperformance degradation due to 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, withperformance degrading more significantly far o.o.d.

Original: We observe that AVICI sometimes performsbetter o.o.d. than in-distribution, e.g., in G RN .
Revised: We observe that AVICI can perform better undercertain distribution shifts than in-distribution, e.g., in G RN .

Original: This is because AVICI is empirically better atpredicting edges adjacent to large-degree nodes, a predominant feature of the GRN graph distributionextracted from E. coli and S. cerevisiae that is not present in Erd˝os-Rényi graphs of the trainingdistribution.
Revised: 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.

Original: The acyclicity constraint is perfectly satisfied for L INEAR and R FF . 
Revised: The acyclicity constraint is perfectly satisfied for L INEAR and R FF . Prec [Figure] [Figure]

Revised: [Figure]

Revised: Error bars indicate the interquartile range.

Original: We observe that AUPRC and AUROC do not suffer as much from distributional shift as the metricsbased on thresholded point estimates.
Revised: We observe that AUPRC and AUROC do not suffer as much from distributional shift as the metricsbased 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: 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.

Original: Overall, AVICI is surprisingly robust, even undersubstantial distribution shifts of all  data distributionaspects at the same time.
Revised: Overall, AVICI is surprisingly robust, even undersubstantial distribution shifts of all aspects of the data distribution at the same time.

Original: {+In Appendix E.
+} Generalization to different problem sizes Next, we investigate the ability of our inference modelsto generalize to unseen problem sizes in terms of number of variables d and  number of data pointsn .To this end, we test the models on increasingly larger sets of variables and various numbers ofobservations per task instance, here, with in-distribution data.
Revised: {+In Appendix E.
+} Generalization to different problem sizes In addition to the sensitivity to distribution shift, weare interested in the ability of our inference models to generalize to unseen problem sizes in termsof the number of variables d and the number of data points n .

Original: The bottom row of Figure 3 illustrates the AUPRC for AVICI’s edge prediction when varying d andn .
Revised: The bottom row of Figure 3 illustratesthe AUPRC for the edge predictions of AVICI when varying d and n on unseen in-distributiondata.

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 while exhibiting diminishing marginalimprovement 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 thenumber of variables 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 behaviorcan be observed well beyond the settings ( n = 200 and d ≤ 50 ) used during training.

Original: 6.2 Benchmarking
Revised: 6.2 Benchmarking

Original: Finally, we compare AVICI with several causal discovery algorithms using either purely observationaldata or a mix ofboth observational and interventional data.
Revised: Next, we compare AVICI with existing causal discovery algorithms that use purely observational dataor a combination of observational and interventional data.

Original: The benchmark is performed on thefully o.o.d.
Revised: The benchmark is performed on the fullyo.o.d.

Original: domain distributions ˜ p ( D ) , i e., under distribution shifts on causal graphs, mechanisms, [Table]
Revised: domain distributions ˜ p ( D ) , i.e., under distribution shifts on causal graphs, mechanisms, andnoise distributions w.r.t.

Original: and noise distributions.
Revised: the training distribution of AVICI.

Original: For each class of methods, we consider at least one constraint-based and onescore-based approach as well as one recent algorithm using neural networks for modeling nonlinearmechanisms.
Revised: 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 compare AVICI with the PC algorithm (Spirtes et al. ,2000), greedy equivalence search (GES) (Chickering, 2003), LiNGAM
Revised: In the observational setting, we compare AVICI with the PC algorithm (Spirtes et al., 2000), greedyequivalence search (GES) (Chickering, 2003), LiNGAM

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

Original: When using interventional data, we compare with greedy interventionalequivalence search (GIES) (Hauser and Bühlmann, 2012),IGSP (Wang et al., 2017), and DCDI(Brouillard et al., 2020).
Revised: 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).

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

Original: When computing the evaluation metrics, wefavor methods that only predict (interventional) Markov equivalence classes by orienting undirectededges correctly when present in the ground truth.
Revised: When computing the evaluation metrics, we favor [Table] methods that only predict (interventional) Markov equivalence classes by orienting undirected edgescorrectly 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 1 shows the SID and F1 scores of all methods for causal discovery from n = 1000 observationsin systems of d = 30 variables.
Revised: Table 1 shows the SID and F1 scores of all methods given n = 1000 observations for d = 30 variables.

Original: The AVICI model trained on L INEAR outperforms all baselines, bothwhen given only observational data and with interventional data, despite operating with significantdistribution shift.
Revised: The AVICI model trained on L INEAR outperforms all baselines, both given only observational andin combination with 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 , whereconstraint-based and nonlinear approaches outperform GES and GIES as expected, but ultimately donot reach the accuracy set by AVICI.
Revised: 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.

Original: Finally, for G RN , we find that the inductive bias encoded byclassical methods reaches its limits.
Revised: In the G RN domain, the inductive bias encoded by classical methodsreaches its limits.

Original: However, provided interventional data, AVICI is able to use itslearned inductive bias to infer plausible causal structure from the noisy gene expression data, whichis a promising step towards reliable structure discovery in fields like molecular biology.
Revised: 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.

Original: Extendedresults for d = 100 variables are given in Appendix E. Overall, the methods’ relative performanceremains unchanged except for L INEAR , where GIES slightly outperforms AVICI.
Revised: 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 Dibaeinia and Sinha 2020; Chen and Mar 2018). Results for d = 100 variables are given in Appendix E.2. 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 of AVICI and the baselines.

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 variationalinference over an arbitrary data-generating distribution.

Original: Our approach leverages the insight thatinductive biases crucial for statistical efficiency  might  more easilybe encoded in a domain simulatorthan in an inference technique.
Revised: 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.

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

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

Original: Using AVICI comes with several trade-offs that a practitioner must consider.
Revised: Using AVICI comes with several trade-offs.

Original: First, whileour experiments showed that optimizing the dual program helps induce acyclicity, this constraint cannot be satisfied with certainty.
Revised: First, while optimizing the dual program empiricallyinduces acyclicity, this constraint is not satisfied with certainty.

Original: Moreover, similar to most amortized inference techniques (Amos, 2022), AVICIprovides no theoretical guarantees of performance, whereas classical methods can  in the infinite sample limit  (Peters et al., 2017).
Revised: Moreover, similar to most amortizedinference techniques (Amos, 2022), AVICI gives no theoretical guarantees of performance. Someclassical methods can do so in the infinite sample limit given specific assumptions on the datagenerating process (Peters et al., 2017).

Original: Further work might achievesimilar guarantees for AVICI as previouslyobtained in the bivariate setting (Lopez-Paz et al., 2015).
Revised: 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).

Original: Finally, while our trained modelsare fast atinference time, they can be intensive to train.
Revised: 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 results accessible to other researchers, we release the  models studied in this work.
Revised: To mitigate this training cost and make our resultsaccessible to other researchers, we release the trained models studied in this work.

Original: Our experiments demonstratethat the trained models are highly robust to simulator misspecification, suggesting that these modelscould be useful even in causal structure learning tasks outside the domains studied in this paper.
Revised: Our experimentsdemonstrate that our inference models are highly robust to simulator misspecification, suggesting thatthe trained models could be useful even in causal structure learning tasks outside the domains studiedin this paper.

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