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 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 amortize causal structure learning .
Revised: In this work, we propose to amortize causal structure learning .

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

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

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

Original: The architecture ofour inference model emulates permutation invariances that are crucial for statisticalefficiency in structure learning, which facilitates generalization to significantlylarger problem instances than seen during training.
Revised: The architecture 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 data and semisynthetic gene expression data, our models exhibit robust generalization capabilitieswhen subject to substantial distribution shifts and significantly outperform existingalgorithms, especially in the challenging genomics domain.
Revised: On synthetic data and semisynthetic gene expression data, our models exhibit robust generalization capabilitieswhen subject to substantial distribution shifts and significantly outperform existingalgorithms, especially in the challenging genomics domain.

Original:  Introduction
Revised:  Introduction

Original: Learning the causal structure among the variables we observe is a fundamental task in variousscientific disciplines (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 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 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: Score-based methods use strong assumptions about the data-generating process, suchas linearity (Shimizu et al., 2006), specific noise models (Hoyer et al., 2008; Peters and Bühlmann,2014) and the absence of measurement error (e.g., Scheines and Ramsey 2016; Zhang et al.
Revised: Score-based methods use strong assumptions about the data-generating process, 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: 2017),which are difficult to verify (Dawid, 2010; Reisach et al., 2021).
Revised: 2017),which are difficult to verify (Dawid, 2010; Reisach et al., 2021).

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

Original: Even with oracle conditional independence tests,they are limited to identifying equivalence classes that may be exponentially large (He et al., 2015b).
Revised: Even with 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: The intractable search space ultimately imposes hard constraints on the causal structure,e.g., the node degree (Spirtes et al., 2000; Koivisto and Sood, 2004), which limits the suitability ofsearch in real-world domains.
Revised: The intractable search space ultimately imposes hard constraints on the causal structure,e.g., the node degree (Spirtes et al., 2000; Koivisto and Sood, 2004), which limits the suitability ofsearch in real-world domains.

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

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

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

Original: Many high-impact domains of causal discovery can be realistically forward-simulated, for example,the chemical ODEs and technical noise of gene regulatory networks (Schaffter et al., 2011; Dibaeiniaand Sinha, 2020), fMRI brain responses (Buxton, 2009; Bassett and Sporns, 2017), and chemical kinetics (Anderson and Kurtz, 2011; Wilkinson, 2018).
Revised: 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: We can leverage such simulators to trainan inference model in causal discovery for specific data-generating processes.
Revised: We can leverage such simulators to trainan inference model in causal discovery for specific data-generating processes.

Original: In doing so, amortizedinference allows us to specify domain-specific inductive biases that are not easily represented bya graph scoring function and bypass the problems of structure search.
Revised: In doing so, 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 data distributionis 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 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.
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 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 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 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 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 models induce calibrated uncertainties and robust behavior whensubject to substantial distribution shifts of graphs, mechanisms, noise, and problem sizes, whichmakes our trained models both reliable and versatile for future downstream use.
Revised: Moreover,we demonstrate that our inference models induce calibrated uncertainties and robust behavior 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, 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, 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: 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: Models suchas causal Bayesian networks and structural causal models (Pearl, 2009; Peters et al., 2017) areless well-suited to describing systems with feedback loops, which we consider practically relevant.

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

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

Original: Avariable x
Revised: Avariable x

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

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

Original: [Equation]
Revised: [Equation]

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

Original: An intervention do ( · ) denotes any active manipulation of the generative process of theobservations x .
Revised: An intervention do ( · ) denotes any active manipulation of the generative process of theobservations x .

Original: An example of interventions are gene knockouts in genomics, in which transcriptionrates of genes are externally set to zero.
Revised: 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 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 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 latershow how to apply our approach when 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 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 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 for causal structure learning search over causal graphs and evaluate them usinga likelihood or conditional independence tests (Chickering, 2003;
Revised: Classical methods for causal structure learning search over causal graphs and evaluate them usinga likelihood or conditional independence tests (Chickering, 2003;

Original: Kalisch and Bühlman, 2007;
Revised: Kalisch and Bühlman, 2007;

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

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

Original: ,2018; Yu et al., 2019; Lachapelle et al., 2020; Brouillard et al., 2020; Lorch et al., 2021).
Revised: ,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).

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

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

Original: There are significant differences in their approach, in particular, our inference models generalize to significantly largerproblem sizes, even when trained on realistic genomics data.
Revised: There are significant differences in their approach, in particular, our inference modelsgeneralize to significantly larger problem 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 Discovery3.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 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
Revised: From a generative perspective, the observations

Original: D = { x 1 , . . .
Revised: 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 over causal structures p ( G )and then obtaining realizations of a data-generating mechanism p ( D | G ) .

Original: It is important to highlightthat p
Revised: It is important to highlightthat p

Original: ( D | G ) does not need to represent an ancestral sampling procedure over a directed acyclicgraph.
Revised: ( 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 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 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 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 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 expected 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 expected forward KL divergence from the intractableposterior p ( G | D ) to q ( G ; θ ) for D ∼ p

Original: ( D ) :
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 obtain 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 ))] ,which allows us 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 datagenerating distribution p ( D ) =
Revised: While the datagenerating distribution p ( D ) =

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

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

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

Original: 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.
Revised: 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 p ( D ) D KL ( p ( G | D )
Revised: The bound is tight if E [Equation] 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 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 and have access to samples from the true data-generating distribution.
Revised: This choice implies that we both model the density q ( G ; θ ) explicitly and have 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
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

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

Original: Furthermore, the forward KL typically yields more reliable uncertaintyestimates since it does not suffer from the variance underestimation problems common to the reverse
Revised: Furthermore, the forward KL typically yields more reliable uncertaintyestimates since it does not suffer from the variance underestimation problems common to the reverse

Original: KL (Bishop and Nasrabadi, 2006).
Revised: KL (Bishop and Nasrabadi, 2006).

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

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

Original: This objective
Revised: This objectiverequires a tractable marginal likelihood p

Original: [Figure]
Revised: ( D | G ) .

Original: [Figure]
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; Schaffter n

Original: [Figure]
Revised: [Figure]

Original: [Figure] [Figure]
Revised: n

Original: [Figure]
Revised: [Figure]

Original: requires a tractable marginal likelihood p ( D | G ) .
Revised: ) n

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; Schaffteret al., 2011; Runge et al., 2019; Dibaeinia and Sinha, 2020). Furthermore, unless the noise scale islearned jointly, likelihoods can be sensitive to the measurement scale of x (Reisach et al., 2021).
Revised: … … (c)

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

Original: [Equation]
Revised: [Equation]

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

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

Original: In addition tothe joint observation 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 ) may contain interventionalinformation for each variable.
Revised: , o id ) may contain interventionalinformation for each variable.

Original: When interventions or gene knockouts are performed, we seto ij = ( x ij , u ij ) and u ij ∈ { 0 , 1 } indicating whether variable j was intervened upon in sample i .
Revised: When interventions or gene knockouts are performed, we 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 statistical efficiency, f ϕ should be constrained to satisfy the invariances inherent to thetask of learning 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 thesample dimension (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., forany permutation π , 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: Moreover, 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: In the following, we show how to parameterize f ϕ as a neural network that encodes these properties.
Revised: In the following, we show how to parameterize f ϕ as a neural network that encodes these properties.

Original: After first mapping each o ij to a real-valued vector using a position-wise linear layer, f ϕ operates overa continuous, three-dimensional tensor of n rows for the observations, d columns for the variables,and feature size k .
Revised: After first mapping each o ij to a real-valued vector using a position-wise linear layer, f ϕ operates 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, which consists ofan observation encoder followed by a generative graph model, similar to Zhu et al.

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 f ϕ 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: 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, the model alternates inattending over the observation and the variable dimension (Kossen et al., 2021).
Revised: To enable information flow across all n × d tokens of the representation, the model alternates 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 message-passing schema keeps memory in checkand results 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 using the attentionblocks, we max-pool over the observation axis n 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 ofone vector z i ∈ R k for each causal variable.
Revised: z d ) consisting of one vector z i ∈ R k for each causal variable.

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

Original: (2021), we use two position-wiselinear layers to map each z i to two embeddings u i , v i ∈ R k .
Revised: (2021), we use two position-wiselinear layers to map each z i to two embeddings u i , v i ∈ R k .

Original: After normalizing the embeddingsu i , v i to an ℓ 2 -norm of one (Radford et al., 2021), we model the probability of each edge in thecausal graph with an inner product:
Revised: 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: (6) where σ is the logistic function, b a learned bias, and τ a positive scale that is learned in log space.
Revised: ) where σ is the logistic function, b a learned bias, and τ a positive scale that is learned in log space.

Original: Since max-pooling is invariant to permutations and 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 .
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: However, 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:[Equation]
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 ) (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: (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: Optimization
Revised: Optimization

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

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

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 inferencemodel (5), we can directly use stochastic optimization to train the parameters ϕ of our neural networkmodel.

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 over p ( G, D ) inside Land F are approximated using samples from thedata-generating process of 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 furtheraugmented using bootstrapping.

Original: When imposing an acyclicity prior on G , we solve the constrained programmax ϕ L ( ϕ ) s.t. F ( ϕ )
Revised: When imposingan acyclicity prior on G , causal discovery algorithms often use the augmented Lagrangian method forconstrained optimization (Zheng et al., 2018;

Original:  = 

Revised: 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 ( ϕ ) s.t. F ( ϕ ) =

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

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

Original: Evaluating causal discovery algorithms is difficult since there are few interesting real-world datasetsthat come with ground-truth causal structure.
Revised: Evaluating causal discovery algorithms is difficult since there are few interesting real-world 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).
Revised: 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.
Revised: 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). 
Revised: In Appendix E, we additionally reportresults on a real-world proteomics dataset (Sachs et al., 2005). Erdos-Renyi

Original: free
Revised: Scale-free

Original: Watts-Strogatz
Revised: Watts-Strogatz

Original: Stochastic Block Model
Revised: Stochastic Block Model

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

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

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

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

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

Original: j
Revised: j

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

Original: [Figure]
Revised: [Figure]

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

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

Original: We studythree domains: two classes of structural causal models (SCMs) as well as semisynthetic single-cellexpression 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 the generalization of AVICI beyondthe training distribution p ( D ) , we carefully 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 substantialshift from p ( D ) in terms of the underlying 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 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 exogenous 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
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

Original: Gaussian process with squared exponential kernel (Rahimi and Recht, 2007).
Revised: Gaussian process with squared exponential kernel (Rahimi and Recht, 2007).

Original: In the out-of-distribution(o.o.d.) setting ˜
Revised: In the out-of-distribution(o.o.d.) setting ˜

Original: p
Revised: p

Original: ( D ) , we sample the linear function and kernel parameters from the tails of p ( D )and unseen value ranges.
Revised: ( D ) , we sample the linear function and kernel parameters from the tails of p ( D )and 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: Moreover, while simulating additive, homoscedastic Gaussian noise ϵ jin the training distribution p ( D ) , we employ heteroscedastic Cauchy and Laplacian noise o.o.d.

Original: ,where
Revised: ,where

Original: f j ( x Pa ( j ) , ϵ j ) =
Revised: 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: ϵ j for some random, input-dependent 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 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 of
Revised: In addition to SCMs, we consider the challenging domain of GRNs (G RN ) using the simulator of

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

Original: Contrary to SCMs, single-cell gene expression samples correspondto draws from the steady state of a stochastic dynamical system that varies between cell types(Huynh-Thu and Sanguinetti, 2019).
Revised: Contrary to SCMs, 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 modelsseveral technical noise effects, resulting in the observation of highly zero-inflated count 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 use the noise levels of different single-cell RNA sequencing technologies,as calibrated by Dibaeinia and Sinha (2020) on real datasets.
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 togene knockouts and are performed by forcing the transcription 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 causal 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 causal structure from interventional data in our experiments, half of thedataset consists of observational data and half of single-variable interventions.

Original: In L INEAR and R FF ,these single-variable interventions are performed on a subset of target variables containing 50 %of the nodes.
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 RN is more challenging, the interventional data here contains single-nodeinterventions on all variables.
Revised: Since G RN is more challenging, the interventional data here contains single-nodeinterventions on all variables.

Original: All datasets are standardized (Reisach et al., 2021).
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 and known biologicalnetworks to generate ground-truth 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 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 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 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 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 their 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 distribution shift from p ( D ) to ˜ p
Revised: To illustrate the distribution shift from p ( D ) to ˜ p

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

Original: The parameter configurations and functionsdefining p ( D ) and ˜ p ( D ) in the three domains are given in Appendix A. We also provide details onthe simulator by Dibaeinia and Sinha (2020) and subgraph extraction (Marbach et al., 2009).
Revised: 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 a predicted causal structure reflects the ground-truth graph, we report the structural
Revised: To asses how well a predicted causal structure reflects the ground-truth graph, we report the structural

Original: Hamming distance (SHD) and the structural intervention distance (SID) (Peters and Bühlmann, 2015).
Revised: Hamming distance (SHD) and the structural intervention distance (SID) (Peters and Bühlmann, 2015).

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

Original: For these metrics, as well as for edge precision, recall,and F1 score, we collapse the posterior probabilities predicted by AVICI to hard predictions using adecision threshold of 0 . 5 .
Revised: For these metrics, 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 .

Original: In addition, we follow existing work and evaluate the uncertainty estimatesby computing the areas under the precision-recall curve (AUPRC) and the receiver operating characteristic (AUROC) (Friedman and Koller, 2003).
Revised: 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 on the metrics 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 by training models for more narrowsettings, e.g., specializing 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 models for the 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 for the evaluation with only observational dataas well as with a combination of observational and interventional data.

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

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

Original: We add the acyclicity constraint for the SCM domains
Revised: We add the acyclicity constraint for the SCM domains

Original: L INEAR and R FF .
Revised: L INEAR and R FF .

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

Original:  Experimental Results
Revised:  Experimental Results

Original: 6.1 Out-Of-Distribution Generalization
Revised: 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: 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 test the model’s sensitivity 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 itgeneralizes to different problem sizes.
Revised: Second, we investigate how well itgeneralizes to different problem sizes.

Original: Sensitivity to distribution shift We perform causal discovery from n = 1000 observations insystems of d = 30 variables.
Revised: 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 empiricalsensitivity 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.
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.
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 can perform better undercertain distribution shifts 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 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.
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 .

Original: Prec
Revised: Prec.

Original: [Figure]
Revised: [Figure]

Revised: SID

Original: [Figure]
Revised: [Figure]

Revised: F

Revised: [Figure]

Revised: F

Original: [Figure]
Revised: [Figure]

Original: Error bars indicate the interquartile range.
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 aspects of the data distribution 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.

Revised: In Appendix E.

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

Original: and the number of data points n .
Revised: and the number of data points n .

Original: The bottom row of Figure 3 illustratesthe AUPRC for the edge predictions of AVICI when varying d and n on unseen in-distributiondata.
Revised: The bottom row of Figure 3 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: Next, we compare AVICI with existing causal discovery algorithms that use purely observational dataor a combination of 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 the fullyo.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, andnoise distributions w.r.t.
Revised: domain distributions ˜ p ( D ) , i.e., under distribution shifts on causal graphs, mechanisms, andnoise distributions w.r.t.

Original: the training distribution of AVICI.
Revised: the training distribution of AVICI.

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 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), greedyequivalence search (GES) (Chickering, 2003),
Revised: In the observational setting, we compare AVICI with the PC algorithm (Spirtes et al., 2000), greedyequivalence search (GES) (Chickering, 2003),

Original: LiNGAM
Revised: LiNGAM

Original: (Shimizu et al., 2006), and DAG-GNN (Yuet al., 2019).
Revised: (Shimizu et al., 2006), and DAG-GNN (Yuet 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 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 important hyperparameters of each baselinemethod 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, we favor [Table] methods that only predict (interventional) Markov equivalence classes by orienting undirected edgescorrectly when present in the ground truth.
Revised: When computing the evaluation metrics, we favormethods 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 given n = 1000 observations for d = 30 variables.
Revised: [Table] 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, 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, both given only observational andin combination with interventional data, despite operating under significant distribution shift.

Original: Only
Revised: Only

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

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

Original: In the G RN domain, the inductive bias encoded by classical methodsreaches its limits.
Revised: In the G RN domain, the inductive bias encoded by classical methodsreaches 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 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: Evenwithout gene knockout data, the edge probabilities predicted by AVICI achieve nontrival AUROCand AUPRC while classical methods predict close to randomly (Table 7 in Appendix E; see also
Revised: 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

Original: Dibaeinia and Sinha 2020; Chen and Mar 2018).
Revised: Dibaeinia and Sinha 2020; Chen and Mar 2018).

Original: Results for d = 100 variables are given in Appendix E.2.
Revised: Results for d = 100 variables are given in Appendix E.2.

Original: In Appendix E.3 and E.4, we also reportbenchmarking results on a real-world proteomics dataset (Sachs et al., 2005) and give additionalmetrics that evaluate the uncertainty quantification of AVICI and the baselines.
Revised: 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 in structure learning might be more easily encoded ina domain simulator than 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 AVICIsolves structure learning problems in complex domains intractable for existing approaches (Dibaeiniaand 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, 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, includingstructure learning settings where 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.

Original: First, while optimizing the dual program empiricallyinduces acyclicity, this constraint is not satisfied with certainty.
Revised: First, while optimizing the dual program empiricallyinduces acyclicity, this constraint is not satisfied with certainty.

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

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

Original: However, future work might achieve guarantees for AVICI thatare similar to learning theory results for the bivariate causal discovery case (Lopez-Paz et al., 2015).
Revised: 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 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 causal discovery with AVICI amounts to a single forward pass of f ϕ at inference time,the inference models can be intensive to train.

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

Original: Our experimentsdemonstrate that our inference models are highly robust to simulator misspecification, suggesting thatthe trained models could be useful even in causal structure learning tasks outside the domains studiedin this paper.
Revised: 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.

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