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 with a score or independence test.
Revised: Learning causal structure poses a combinatorial search problem that typically involves evaluating structures using 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 amortize  causal structure learning .
Revised: In this work, we propose to amortize the process of 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 causal structures directly, we train a variationalinference model to predict the causal structure from observational/interventionaldata.

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

Original: This allowsus to bypass both the search over graphs and the hand-engineering of suitable scorefunctions.
Revised: This allows us to bypass both thesearch 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: Moreover,the architecture of our inference model is permutation invariant w.r.t. the datapoints and permutation equivariant w.r.t.

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 variables, facilitating generalization tosignificantly 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 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.

Original:  Introduction
Revised:  Introduction

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

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

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

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: 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 suitabilityofsearch in real-world domains.
Revised: 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.

Original: In the present work, wepropose to amortize  causal structure learning.
Revised: In the present work, we amortize the process of 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 goalis to optimize an inference model to solve the inverse problem of inferring a causal structure froma provided data set.

Original: This proposal ismotivated 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: 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.

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

Original: We can leverage such simulatorsto 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.
Revised: 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.

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 data distributionis unlikely to cover all real-world distribution shifts.
Revised: This generalization is essential because any simulated datadistribution is unlikely to cover all real-world distribution shifts.

Original: To quantify the epistemic uncertainty inthis inference process, wepropose to train an amortized variational inference model for causaldiscovery (AVICI) by approximating the posterior over causal graphs givendata from a specified datagenerating distribution, such as a simulator.
Revised: 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.

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

Original: On synthetic data and semisynthetic expression data of gene regulatory networks, our approachsignificantly outperforms existing algorithms for causal discovery, often by a large margin.
Revised: 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.

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

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

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

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

Original: , x d ) andwhose edges represent all direct causal effects among the variables.
Revised: , x d ) and 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 : = x \{ x i , x j } , i.e., there exists a ̸ = a ′ such that
Revised: 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

Original: [Equation]
Revised: [Equation]

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

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

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

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: We emphasize that G is not necessarily acyclic, however we later show how to apply our approachwhen 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: In this work, we assume causal sufficiency, i.e., that there are nounobserved variables beyond x that are causal parents of more than one x i (Peters et al., 2017).

Original: 2.2 Related Work
Revised: 2.2 Related Work

Original: Classical methodsfor causal structure learning search over causal graphs and evaluate them usinga likelihood or conditional independence tests (Chickering, 2003; Kalisch and Bühlman, 2007; 
Revised: 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;

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

Original: Other methods combineconstraint- and score-based ideas (Tsamardinos et al., 2006) or usethe noise properties ofan SCMthat is postulated to underlie the data-generating process (Shimizu et al., 2006; Hoyer et  al., 2008).
Revised: Others combine constraint and score-based ideas(Tsamardinos et al., 2006) or use noise properties of the SCM (Shimizu et al., 2006; Hoyer et al.,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: 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 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 applied to causal inference, e.g., for the estimation of 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: In causal structure learning, neural networks have been usedprimarily to model nonlinear causal mechanisms among the variables (Goudet et al., 2018; Yu et al.
Revised: In causal structure learning, neural networks have mostlybeen used to model nonlinear causal mechanisms among the variables (Goudet et al., 2018; Yuet al., 2019;

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: Lachapelle et al., 2020; Brouillard et al., 2020; Ke et al., 2020; Lorch et al., 2021).

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 mechanisms for eachtask (Löwe et al., 2022) or only used observational correlation coefficients for prediction (Li et al.,2020).
Revised: 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).

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

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

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

Revised: 3.1 The 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 the process of causal structure learning, we define a data-generating distribution p ( D )that models the domain 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 over causal 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 overcausal structures p ( G ) and then obtaining realizations of a data-generating mechanism p ( D | G ) .

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

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

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 theexpected forward KL divergence from the intractableposterior p ( G | D )to q ( G ; θ ) 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 forward KL divergence from the intractable posteriorp ( G | D ) for D ∼ p

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

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

Original: E p (  G )[ p (D | G )] can be arbitrarily complex,  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: (cid:80) G p ( D | G ) p ( G ) can be arbitrarily complex, the training 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 ) .

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 byuse of a simulator for the domain of interest.

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
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 [Equation] 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 requiresoptimizing E p ( G,D )
Revised: The AVICI objective in (3) intentionally targets the forward KL D KL ( p ∥ q ) , which requires optimizing 124 E p ( D | G )

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 andhave access to samples from the true data-generating distribution.

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

Original: Furthermore, the forward KL typically yields more reliable uncertaintyestimates since it does not suffer from the variance underestimation problems common to the reverse
Revised: 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: Access to p ( G, D ) is natural when a forward-model of the

Original: KL (Bishop and Nasrabadi, 2006).
Revised: [Figure] x L [Figure] domain exists, e.g., when sampling D synthetically from a structural equation model or a specializeddomain simulator.

Revised: We expect the quality of simulators in many domains to improve substantiallyover time, reducing the sim-to-real bias of the inference model f ϕ .

Original: In contrast, variational inferenceusually 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 optimizes the reverse KL D KL ( q ∥ p ) with respect to θ , whichinvolves 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 aparametric, 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 jointly 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).

Original: n [Figure] ) n … … (c)
Revised: In addition, 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 procedure  for optimizing themodel parameters ϕ as well as for learning the inductive bias of acyclicity.
Revised: After that, we detail our training procedure of AVICI for optimizingthe model 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 objectivein (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 (3), we optto 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 ϕ thus maps a data set 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- dmatrix θ parameterizing the variational approximation of the causal graph.

Original: In addition tothe joint observation x i =
Revised: In addition to the jointobservation 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 ) contains interventional informationfor each variable when available.

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

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 statistical efficiency, f ϕ should be constrained to satisfy the invariances inherent to thetask of learning a causal structure from data.
Revised: To maximize sample efficiency, f ϕ is constrained to satisfy the invariances inherent to the task oflearning a causal structure from data.

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

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

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

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

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

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 axes n and d The core of f ϕ is composed of L = 8 identical blocks.
Revised: Attending over axes n and d The core of our model is composed of L = 8 identical blocks.

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 the vanilla Transformer encoder (Vaswani et al.,2017).
Revised: 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).

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, themodel alternates in attending over the observation and the variable dimension.

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 message-passing schema keeps memory  in checkand results in a moderate model size because parameters are shared across the batched axes.
Revised: This design keeps memory requirements in check andresults 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 using the attentionblocks, we max-pool over the observation axisn to obtain a representation ( z 1 , . .
Revised: Causal graph model After building up a representation tensor from the input, we max-pool overthe observation axis 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 eachvariable.

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

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.

Original: Since max-pooling is invariant to permutations and since the graph model in (6) permutes withrespect to axis d , theinference modelf ϕ satisfies the required permutation invariance over axis n andpermutation equivariance over axis d .
Revised: 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.

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: However, certain domains may bemore accurately modeled by acyclic causal structures (Peters et al., 2017; Rubenstein et al., 2017).
Revised: Nonetheless, 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:

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 ) (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, I ) (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: (8) and a , b ∈ R d are initialized randomly.
Revised: and a , b ∈ R d are randomly initialized.

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:  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: Algorithm Parameters: ϕ variational, λ dual, η step size while not converged do for l steps do

Original: + η F ( ϕ )
Revised: [Equation]

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 stochasticoptimization to train the parameters ϕ of ourneural network 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 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 furtheraugmented using bootstrapping.
Revised: In practice,the available data can be further augmentedusing bootstrapping.

Original: When imposingan acyclicity prior on G , causal discovery algorithms often use the augmented Lagrangian method forconstrained optimization (Zheng et al., 2018;
Revised: In case we impose an acyclicity prior on G , we solve the constrained programmax [Equation]

Original: Yu et al., 2019; Brouillard et al., 2020). 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: ϕ [Equation]

Original: s.t. F ( ϕ ) = 

Revised: t. F [Equation] ϕ ) =

Original: through its dual formulation (Nandwani et al., 2019):min λ max ϕ
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 .Experimental Setup
Revised: Without an acyclicityconstraint, training reduces to the primal updates of ϕ with λ = 0 .

Original: Evaluating causal discovery algorithms is difficult since there are few interesting real-world datasetsthat 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).

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.

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

Original: Erdos-Renyi
Revised: 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]

Revised:  Experimental Setup

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

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

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

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

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

Original: ( D ) that incurs substantialshift from p ( D ) in terms of the underlying causal structures, mechanisms, and noise components.
Revised: ( D ) that incurs substantial shift from p ( D ) in terms of theunderlying 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 exogenous 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 endogenous noisedistributions p j ( ϵ j ) for each variable.

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

Original: 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.
Revised: For the out-of-distribution (o.o.d.)setting ˜ p ( D ) , we draw the linear function weights and the kernel scales from the tails of p ( D ) oreven unseen value ranges.

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

Original: ,where f j ( x Pa ( j ) , ϵ j ) =
Revised: ( D ) with 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: ϵ j for some random, input-dependent function h j .
Revised: ϵ jfor some random function h j .

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

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 by

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

Original: Contrary to SCMs, single-cellgene 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, single-cell samples of gene expressions correspondto draws from the steady-state of a stochastic dynamical system that varies across 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 modelsseveral technical noise effects, resulting inthe observation of highly zero-inflatedcount matrices D .
Revised: The simulator additionally modelstechnical noise effects, resulting in highly zero-inflated matrices D .

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

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

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

Original: Whenever a method inferscausal structure from interventional data in our experiments, half of thedataset consists of observational data and half of single-variable interventions.
Revised: 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.

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

Original: Since G RN is more challenging, the interventional data here contains single-nodeinterventions on all variables.
Revised: Since G RNis more challenging, the interventional data here contains single-node interventions 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 modelsand known biologicalnetworks to generate ground-truth causal structures.
Revised: Following prior work, we use random graph models to generate groundtruth causal structures.

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

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

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, weuse subgraphs of the known S. cerevisiae and E. coli GRNs and the 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 thedistribution shift from p ( D ) to ˜ p
Revised: To illustrate the 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 noisedistribution 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: Details on the simulator by Dibaeinia and Sinha (2020) andsubgraph extraction are given in Appendix A. The specific parameter configurations and functionsdefining p ( D ) and all ˜ p ( D ) in the three domains are provided in Appendix B.

Original: 5.2 Evaluation Metrics
Revised: 5.2 Evaluation Metrics

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

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 quantifiesthe 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 . 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: In addition, we followexisting work and evaluate the predicted 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.

Original: More detailson the metrics are given in Appendix B.
Revised: More details 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 by training models for more narrowsettings, e.g., specializing  solely on inference from observational data.
Revised: Performance could be further improved, e.g., by training models for more narrow settings specializing them solely on inference from observational data.

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

Original: Allexperiments 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 modelsin the three domains share identical hyperparameters for the architecture andoptimization, exceptfor the dropout rate. 
Revised: 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: generalisation.

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

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

Original: 6.1 Out-Of-Distribution Generalization

Original: 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.
Revised: We study the generalization capabilities of our inference models across 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 test the model’s sensitivity to queries with o.o.d.
Revised: First, we consider the model’ssensitivity to queries with o.o.d.

Original: data.
Revised: data.

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

Original: Sensitivity to distribution shift We perform causal discovery from n = 1000 observations insystems of d = 30 variables.
Revised: Sensitivity to distribution shift We consider causal discovery from n = 1000 observations ofsystems 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 empiricalsensitivity analysis.
Revised: The top row of Figure 3 visualizes the results of an empirical o.o.d.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 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.
Revised: Given this, all metrics arein a similar value range and higher values indicate better performance.

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

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

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

Original: Prec.

Original: [Figure]

Original: SID

Original: [Figure]

Original: F
Revised: F

Original: [Figure]
Revised: [Figure]

Original: F
Revised: F

Original: [Figure]
Revised: [Figure]

Original: Error bars indicate the interquartile range.
Revised: Error bars indicate the interquartile range.

Revised: As L INEAR is the easiest domain, AVICI performs very well across all metrics and hardly suffersfrom performance degradation due to distribution shift.

Revised: In contrast, G RN is the most challenging, withperformance degrading more significantly far o.o.d..

Revised: We observe that AVICI sometimes performsbetter o.o.d.

Revised: than in-distribution, e.g., in G RN .

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

Original:  We observe that AUPRC and AUROC do not suffer as much from distributional shift as the metricsbased on thresholded point estimates.
Revised: The acyclicity constraint is perfectly satisfied for L INEAR and R FF . 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 allaspects of the data distribution  at the same time.
Revised: Overall, AVICI is surprisingly robust, even undersubstantial distribution shifts of all data distribution aspects at the same time.

Original: [-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 andthe number of data points n . 
Revised: [-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.

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 AVICI’s edge prediction when varying d andn .

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: Next, we compare AVICI with existing causal discovery algorithms that use purely observational dataor a combination of  observational and interventional data.
Revised: Finally, we compare AVICI with several causal discovery algorithms using either purely observationaldata or a mix of both observational and interventional data.

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

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

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

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

Original: In the observational setting, we compare AVICI with the PC algorithm (Spirtes et al., 2000), greedyequivalence search (GES) (Chickering, 2003), LiNGAM
Revised: In the observational setting, we compare AVICI with the PC algorithm (Spirtes et al. ,2000), greedy equivalence search (GES) (Chickering, 2003), LiNGAM

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

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

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

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, wefavor methods that only predict (interventional) Markov equivalence classes by orienting undirectededges 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] 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 for causal discovery from n = 1000 observationsin systems of d = 30 variables.

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

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

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

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

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 efficiencyin structure learning mightbe more easily  encoded ina domain simulator than in an inference technique.
Revised: Our approach leverages the insight thatinductive biases crucial for statistical efficiency might more easily be encoded in a domain simulatorthan 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 structurelearning problems on 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 settingswhere we cannot assume causal sufficiency (Bhattacharya et al., 2021).

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

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

Original: Moreover, similar to most amortizedinference techniques (Amos, 2022), AVICI gives no theoretical guarantees of performance. Someclassical methods cando so in the infinite sample limitgiven specific assumptions on the datagenerating process (Peters et al., 2017).
Revised: 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).

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: Further work might achieve similar guarantees for AVICI as previouslyobtained in the bivariate setting (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.
Revised: Finally, while our trained models are fast atinference time, they can be intensive to train.

Original: To mitigate this training cost and make our resultsaccessible to other researchers, we release thetrained models studied in this work.
Revised: To mitigate this training cost and make our results accessible to other researchers, we release the 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: 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.

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