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

Original: Abstract
Revised: Abstract

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

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

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

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

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

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

Original: The architecture of our inference model emulates permutation invariances that are crucial for statistical efficiency in structure learning, which facilitates generalization to significantly larger problem instances than seen during training.
Revised: The architecture of our inference model emulates permutation invariances that are crucial for statistical efficiency in structure learning, which facilitates generalization to significantly larger problem instances than seen during training.

Original: On synthetic data and semisynthetic gene expression data, our models exhibit robust generalization capabilities when subject to substantial distribution shifts and significantly outperform existing algorithms, especially in the challenging genomics domain.
Revised: On synthetic data and semisynthetic gene expression data, our models exhibit robust generalization capabilities when subject to substantial distribution shifts and significantly outperform existing algorithms, especially in the challenging genomics domain.

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

Original:  Introduction
Revised:  Introduction

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

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

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

Original: To infer a causal structure standard methods have to search over potential graphs, usually maximizing either a graph scoring function or testing for conditional independences (Heinze-Deml et al., 2018).
Revised: To infer a causal structure standard methods have to search over potential graphs, usually maximizing either a graph scoring function or testing for conditional independences (Heinze-Deml et al., 2018).

Original: Specifying realistic inductive biases is universally difficult for existing approaches to causal discovery.
Revised: Specifying realistic inductive biases is universally difficult for existing approaches to causal discovery.

Original: Score-based methods use strong assumptions about the data-generating process, such as linearity (Shimizu et al., 2006), specific noise models (Hoyer et al., 2008; Peters and Bühlmann, 2014), and the absence of measurement error (cf.
Revised: Score-based methods use strong assumptions about the data-generating process, such as linearity (Shimizu et al., 2006), specific noise models (Hoyer et al., 2008; Peters and Bühlmann, 2014), and the absence of measurement error (cf.

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

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

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

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

Original: The intractable search space ultimately imposes hard constraints on the causal structure, e.g., the node degree (Spirtes et al., 2000), which limits the suitability of search in real-world domains.
Revised: The intractable search space ultimately imposes hard constraints on the causal structure, e.g., the node degree (Spirtes et al., 2000), which limits the suitability of search in real-world domains.

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

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

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

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

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

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

Original: The central challenge of amortizing causal discovery consists in generalizing out-of-distribution, since any simulated data distribution is unlikely to cover all real-world distribution shifts.
Revised: The central challenge of amortizing causal discovery consists in generalizing out-of-distribution, since any simulated data distribution is unlikely to cover all real-world distribution shifts.

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

Original: Our model architecture is permutation in- and equivariant with respect to the observation and variable dimensions of the provided dataset, respectively, and generalizes to significantly larger problem instances than seen during training.
Revised: Our model architecture is permutation in- and equivariant with respect to the observation and variable dimensions of the provided dataset, respectively, and generalizes to significantly larger problem instances than seen during training.

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

Original: Moreover, we demonstrate that our inference models induce calibrated uncertainties and robust behavior when subject to substantial distribution shifts of graphs, mechanisms, noise, and problem sizes.
Revised: Moreover, we demonstrate that our inference models induce calibrated uncertainties and robust behavior when subject to substantial distribution shifts of graphs, mechanisms, noise, and problem sizes.

Original: This suggests that our trained models are both reliable and versatile for future downstream use.
Revised: This suggests that our trained models are both reliable and versatile for future downstream use.

Original: In particular, AVICI was the only method to infer plausible causal structures from noisy gene expression data, advancing the frontiers of structure discovery in fields such as molecular biology.
Revised: In particular, AVICI was the only method to infer plausible causal structures from noisy gene expression data, advancing the frontiers of structure discovery in fields such as molecular biology.

Original:  Background and Related Work 2.1 Causal Structure
Revised:  Background and Related Work 2.1 Causal Structure

Original: In this work, we follow Mooij et al.
Revised: In this work, we follow Mooij et al.

Original: and define the causal structure G of a set of d variables x = ( x 1 , . .
Revised: and define the causal structure G of a set of d variables x = ( x 1 , . .

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

Original: A variable x i has a direct causal effect on x j if intervening on x i affects the outcome of x j independent of the other variables x \ ij :
Revised: A variable x i has a direct causal effect on x j if intervening on x i affects the outcome of x j independent of the other variables x \ ij :

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

Original: [Equation]
Revised: [Equation]

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

Original: An intervention do ( · ) denotes any active manipulation of the generative process of x , like gene knockouts, in which the transcription rates of genes are externally set to zero.
Revised: An intervention do ( · ) denotes any active manipulation of the generative process of x , like gene knockouts, in which the transcription rates of genes are externally set to zero.

Original: Other models such as causal Bayesian networks and structural causal models (Peters et al., 2017) are less well-suited for describing systems with feedback loops, which we consider practically relevant.
Revised: Other models such as causal Bayesian networks and structural causal models (Peters et al., 2017) are less well-suited for describing systems with feedback loops, which we consider practically relevant.

Original: However, we note that our approach does not require any particular formalization of causal structure.
Revised: However, we note that our approach does not require any particular formalization of causal structure.

Original: In particular, we later show how to apply our approach when G is constrained to be acyclic.
Revised: In particular, we later show how to apply our approach when G is constrained to be acyclic.

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

Original: 2.2 Related Work
Revised: 2.2 Related Work

Original: Classical methods for causal structure learning search over causal graphs and evaluate them using a likelihood or conditional independence test (Chickering, 2003; Kalisch and Bühlman, 2007; Hauser and Bühlmann, 2012; Zheng et al., 2018; Heinze-Deml et al., 2018).
Revised: Classical methods for causal structure learning search over causal graphs and evaluate them using a likelihood or conditional independence test (Chickering, 2003; Kalisch and Bühlman, 2007; Hauser and Bühlmann, 2012; Zheng et al., 2018; Heinze-Deml et al., 2018).

Original: Other methods combine constraint- and score-based ideas
Revised: Other methods combine constraint- and score-based ideas

Original: (Tsamardinos et al., 2006) or use the noise properties of an SCM that is postulated to underlie the data-generating process (Shimizu et al., 2006; Hoyer et al., 2008).
Revised: (Tsamardinos et al., 2006) or use the noise properties of an SCM that is postulated to underlie the data-generating process (Shimizu et al., 2006; Hoyer et al., 2008).

Original: Deep learning has been used for causal inference, e.g., for estimating treatment effects (Shalit et al., 2017; Louizos et al., 2017; Yoon et al., 2018) and in instrumental variable analysis (Hartford et al., 2017; Bennett et al., 2019).
Revised: Deep learning has been used for causal inference, e.g., for estimating treatment effects (Shalit et al., 2017; Louizos et al., 2017; Yoon et al., 2018) and in instrumental variable analysis (Hartford et al., 2017; Bennett et al., 2019).

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

Original: Yu et al., 2019; Lachapelle et al., 2020; Brouillard et al., 2020; Lorch et al., 2021) or to infer the structure of a single dataset (Zhu et al., 2020).
Revised: Yu et al., 2019; Lachapelle et al., 2020; Brouillard et al., 2020; Lorch et al., 2021) or to infer the structure of a single dataset (Zhu et al., 2020).

Original: Prior work applying amortized inference to causal discovery only studied narrowly defined subproblems such as the bivariate case (Lopez-Paz et al., 2015) and fixed causal mechanisms (Löwe et al., 2022) or used correlation coefficients for prediction (Li et al., 2020).
Revised: Prior work applying amortized inference to causal discovery only studied narrowly defined subproblems such as the bivariate case (Lopez-Paz et al., 2015) and fixed causal mechanisms (Löwe et al., 2022) or used correlation coefficients for prediction (Li et al., 2020).

Original: Ke et al.
Revised: Ke et al.

Original: (2022) also frame causal discovery as supervised learning, but with significant differences.
Revised: (2022) also frame causal discovery as supervised learning, but with significant differences.

Original: Most importantly, we optimize a variational objective under a model class that captures the symmetries of structure learning.
Revised: Most importantly, we optimize a variational objective under a model class that captures the symmetries of structure learning.

Original: Empirically, our models generalize to significantly larger problem sizes, even on realistic genomics data.
Revised: Empirically, our models generalize to significantly larger problem sizes, even on realistic genomics data.

Original:  AVICI: Amortized Variational Inference for Causal Discovery 3.1 Variational Objective
Revised:  AVICI: Amortized Variational Inference for Causal Discovery 3.1 Variational Objective

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

Original: The observations D = { x 1 , . . .
Revised: The observations D = { x 1 , . . .

Original: , x n } ∼ p ( D ) are generated by sampling from a distribution over causal structures p ( G ) and then obtaining realizations from a data-generating mechanism p ( D | G ) .
Revised: , x n } ∼ p ( D ) are generated by sampling from a distribution over causal structures p ( G ) and then obtaining realizations from a data-generating mechanism p ( D | G ) .

Original: The data-generating process p ( D | G ) characterizes all direct causal effects (1) in the system, but it is not necessarily induced by ancestral sampling over a directed acyclic graph.
Revised: The data-generating process p ( D | G ) characterizes all direct causal effects (1) in the system, but it is not necessarily induced by ancestral sampling over a directed acyclic graph.

Original: Real-world systems are often more naturally modeled at different granularities or as dynamical systems (Mooij et al., 2013;
Revised: Real-world systems are often more naturally modeled at different granularities or as dynamical systems (Mooij et al., 2013;

Original: Hoel et al., 2013; Rubenstein et al., 2017; Schölkopf, 2019).
Revised: Hoel et al., 2013; Rubenstein et al., 2017; Schölkopf, 2019).

Original: Given a set of observations D , our goal is to approximate the posterior over causal structures p ( G | D ) with a variational distribution q ( G ; θ ) .
Revised: Given a set of observations D , our goal is to approximate the posterior over causal structures p ( G | D ) with a variational distribution q ( G ; θ ) .

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

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

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: [Equation]
Revised: [Equation]

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

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

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

Original: ( G ; f ϕ ( D ))
Revised: ( G ; f ϕ ( D ))

Original: (cid:1) =
Revised: (cid:1) =

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

Original: = − E p
Revised: = − E p

Original: [Equation]
Revised: [Equation]

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

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

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

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

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

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

Original: In practice, p ( G ) and p ( D | G ) can thus be specified by a simulator.
Revised: In practice, p ( G ) and p ( D | G ) can thus be specified by a simulator.

Original: From an information-theoretic viewpoint, the objective (2) maximizes a variational lower bound on the mutual information I [ G ; D ] between the causal structure G and the observations D (Barber and
Revised: From an information-theoretic viewpoint, the objective (2) maximizes a variational lower bound on the mutual information I [ G ; D ] between the causal structure G and the observations D (Barber and

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

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

Original: [Equation]
Revised: [Equation]

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

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

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

Original: [Equation]
Revised: [Equation]

Original: G | D )
Revised: G | D )

Original: ∥ q ( G
Revised: ∥ q ( G

Original: [Equation]
Revised: [Equation]

Original: f ϕ
Revised: f ϕ

Original: [Equation]
Revised: [Equation]

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

Original: = 0 .
Revised: = 0 .

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

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

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

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

Original: Minimizing the forward KL enables us to infer causal structures in arbitrarily complex domains—that is, even domains where it is difficult to specify an explicit likelihood p
Revised: Minimizing the forward KL enables us to infer causal structures in arbitrarily complex domains—that is, even domains where it is difficult to specify an explicit likelihood p

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

Original: Moreover, the forward
Revised: Moreover, the forward

Original: KL typically yields more reliable uncertainty estimates since it does not suffer from the variance underestimation problems common to the reverse KL (Bishop and Nasrabadi, 2006).
Revised: KL typically yields more reliable uncertainty estimates since it does not suffer from the variance underestimation problems common to the reverse KL (Bishop and Nasrabadi, 2006).

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

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

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

Original: Unless inferring the mechanism parameters jointly (e.g. Brouillard et al. 2020; Lorch et al. 2021), this requirement limits inference to conjugate models with linear
Revised: Unless inferring the mechanism parameters jointly (e.g. Brouillard et al. 2020; Lorch et al. 2021), this requirement limits inference to conjugate models with linear

Original: Gaussian or categorical mechanisms that assume zero measurement error (Geiger and Heckerman, 1994; Heckerman et al., 1995), which are not justified in practice (Friston et al., 2000; Schaffter et al., 2011; Runge et al., 2019; Dibaeinia and Sinha, 2020).
Revised: Gaussian or categorical mechanisms that assume zero measurement error (Geiger and Heckerman, 1994; Heckerman et al., 1995), which are not justified in practice (Friston et al., 2000; Schaffter et al., 2011; Runge et al., 2019; Dibaeinia and Sinha, 2020).

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

Original:  Inference Model
Revised:  Inference Model

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

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

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

Original: 4.1 Variational Family
Revised: 4.1 Variational Family

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

Original: [Equation]
Revised: [Equation]

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

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

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

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

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

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

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

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

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

Original: 4.2 Model Architecture
Revised: 4.2 Model Architecture

Original: To maximize statistical efficiency, f ϕ should satisfy the symmetries inherent to the task of causal structure learning.
Revised: To maximize statistical efficiency, f ϕ should satisfy the symmetries inherent to the task of causal structure learning.

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

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

Original: =
Revised: =

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

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

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

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

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

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

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

Original: Figure 1 illustrates the key components of the architecture.
Revised: Figure 1 illustrates the key components of the architecture.

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

Original: Each layer consists of four residual sublayers, where the first and third apply multi-head self-attention and the second and fourth position-wise feed-forward networks, similar to the Transformer encoder (Vaswani et al., 2017).
Revised: Each layer consists of four residual sublayers, where the first and third apply multi-head self-attention and the second and fourth position-wise feed-forward networks, similar to the Transformer encoder (Vaswani et al., 2017).

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

Original: Specifically, the first self-attention sublayer attends over axis d , treating axis n as a batch dimension; the second attends over axis n , treating axis d as a batch dimension.
Revised: Specifically, the first self-attention sublayer attends over axis d , treating axis n as a batch dimension; the second attends over axis n , treating axis d as a batch dimension.

Original: Since modules are shared across non-attended axes, the representation is permutation equivariant over axes n and d at all times (Lee et al., 2019b).
Revised: Since modules are shared across non-attended axes, the representation is permutation equivariant over axes n and d at all times (Lee et al., 2019b).

Original: Variational parameters After building up a representation tensor from the input using the attention layers, we max-pool over the observation axis n to obtain a representation ( z 1 , . .
Revised: Variational parameters After building up a representation tensor from the input using the attention layers, we max-pool over the observation axis n to obtain a representation ( z 1 , . .

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

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

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

Original: We then model the probability of each edge in the causal graph with an inner product:
Revised: We then model the probability of each edge in the causal graph with an inner product:

Original: [Equation]
Revised: [Equation]

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

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

Original: 4.3 Acyclicity
Revised: 4.3 Acyclicity

Original: Cyclic causal effects often occur, e.g., when modeling stationary distributions of dynamical systems, and thus loops in a causal structure are possible.
Revised: Cyclic causal effects often occur, e.g., when modeling stationary distributions of dynamical systems, and thus loops in a causal structure are possible.

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

Original: While the variational family in (5) cannot enforce it, we can optimize for acyclicity through ϕ .
Revised: While the variational family in (5) cannot enforce it, we can optimize for acyclicity through ϕ .

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

Original: [Equation]
Revised: [Equation]

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

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

Original: We use power iteration to approximate and differentiate through the largest eigenvalue of f ϕ ( D )
Revised: We use power iteration to approximate and differentiate through the largest eigenvalue of f ϕ ( D )

Original: (Golub and Van der Vorst, 2000; Lee et al., 2019a):
Revised: (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] [Equation]
Revised: [Equation] [Equation]

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

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

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

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

Original: 4.4 Optimization
Revised: 4.4 Optimization

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

Original: [Equation]
Revised: [Equation]

Original: ϕ ) −
Revised: ϕ ) −

Original: [Equation]
Revised: [Equation]

Original: F ( ϕ
Revised: F ( ϕ

Original: [Equation]
Revised: [Equation]

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

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

Original: The expectations over p ( G, D ) inside L and
Revised: The expectations over p ( G, D ) inside L and

Original: F are approximated using samples from the datagenerating process of the domain.
Revised: F are approximated using samples from the datagenerating process of the domain.

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

Original: Brouillard et al. 2020).
Revised: Brouillard et al. 2020).

Original: In this work, we optimize the parameters ϕ of a neural network, so we rely on methods specifically tailored for deep learning and solve the constrained program max ϕ L ( ϕ )
Revised: In this work, we optimize the parameters ϕ of a neural network, so we rely on methods specifically tailored for deep learning and solve the constrained program max ϕ L ( ϕ )

Original: s.t. F ( ϕ ) = 

Revised: s.t. F ( ϕ ) =

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

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

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

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

Original:  Experimental Setup
Revised:  Experimental Setup

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

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

Original: Following prior work, we deal with this difficulty by evaluating our approach using simulated data with known causal structure and by controlling for various aspects of the task.
Revised: Following prior work, we deal with this difficulty by evaluating our approach using simulated data with known causal structure and by controlling for various aspects of the task.

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

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

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

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

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

Original: Whenever we consider interventional data in our experiments, half of the dataset consists of observational data and half of single-variable interventions.
Revised: Whenever we consider interventional data in our experiments, half of the dataset consists of observational data and half of single-variable interventions.

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

Original: We consider SCMs with linear functions (L INEAR ) and with nonlinear functions of random Fourier features (R FF ) that correspond to functions drawn
Revised: We consider SCMs with linear functions (L INEAR ) and with nonlinear functions of random Fourier features (R FF ) that correspond to functions drawn

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]

Original: from a Gaussian process with squared exponential and Recht,
Revised: from a Gaussian process with squared exponential and Recht,

Original: In the out
Revised: In the out

Original: ˜ we the and unseen
Revised: ˜ we the and unseen

Original: Moreover, we simulate homoscedastic Gaussian noise in the training but on heteroscedastic and that is induced by
Revised: Moreover, we simulate homoscedastic Gaussian noise in the training but on heteroscedastic and that is induced by

Original: In L
Revised: In L

Original: , interventions to random values and are performed on a subset of target variables containing half of the nodes.
Revised: , interventions to random values and are performed on a subset of target variables containing half of the nodes.

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

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

Original: Contrary to SCMs, gene expression samples correspond to draws from the steady state of a stochastic dynamical system that varies between cell types (Huynh-Thu and
Revised: Contrary to SCMs, gene expression samples correspond to draws from the steady state of a stochastic dynamical system that varies between cell types (Huynh-Thu and

Original: Sanguinetti, 2019).
Revised: Sanguinetti, 2019).

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

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

Original: In addition, we use the noise levels of different single-cell RNA sequencing technologies, which were calibrated on real datasets.
Revised: In addition, we use the noise levels of different single-cell RNA sequencing technologies, which were calibrated on real datasets.

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

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

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

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

Original: Albert, 1999).
Revised: Albert, 1999).

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

Original: setting, ˜ p
Revised: setting, ˜ p

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

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

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

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

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

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

Original: In Appendix A, we give the detailed parameter configurations and functions defining p ( D ) and ˜ p
Revised: In Appendix A, we give the detailed parameter configurations and functions defining p ( D ) and ˜ p

Original: ( D ) in the three domains.
Revised: ( D ) in the three domains.

Original: We also provide details on the simulator by Dibaeinia and Sinha (2020) and subgraph extraction (Marbach et al., 2009).
Revised: We also provide details on the simulator by Dibaeinia and Sinha (2020) and subgraph extraction (Marbach et al., 2009).

Original: 5.2 Evaluation Metrics
Revised: 5.2 Evaluation Metrics

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

Original: To assess how well a predicted structure reflects the ground truth, we report the structural Hamming distance (SHD) and the structural intervention distance (SID) (Peters and Bühlmann, 2015).
Revised: To assess how well a predicted structure reflects the ground truth, we report the structural Hamming distance (SHD) and the structural intervention distance (SID) (Peters and Bühlmann, 2015).

Original: While the
Revised: While the

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

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

Original: We evaluate the uncertainty estimates by computing the areas under the precision-recall curve (AUPRC) and the receiver operating characteristic (AUROC) (Friedman and Koller, 2003).
Revised: We evaluate the uncertainty estimates by computing the areas under the precision-recall curve (AUPRC) and the receiver operating characteristic (AUROC) (Friedman and Koller, 2003).

Original: How
Revised: How

Original: F
Revised: F

Original: [Figure]
Revised: [Figure]

Original: task instances.
Revised: task instances.

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

Original: well these uncertainty estimates are calibrated is quantified with the expected calibration error (ECE) (DeGroot and Fienberg, 1983).
Revised: well these uncertainty estimates are calibrated is quantified with the expected calibration error (ECE) (DeGroot and Fienberg, 1983).

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 these three trained models, both when predicting from only observational and from interventional data.
Revised: We train three inference models overall, one for each domain, and perform all experiments on these three trained models, both when predicting from only observational and from interventional data.

Original: During training, the datasets sampled from p ( D ) have d = 2 to 50 variables and n = 200 samples.
Revised: During training, the datasets sampled from p ( D ) have d = 2 to 50 variables and n = 200 samples.

Original: With probability 0.5, these training datasets contain 50 interventional samples.
Revised: With probability 0.5, these training datasets contain 50 interventional samples.

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

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

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

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

Original: Sensitivity to distribution shift In our first set of experiments, we study the generalization capabilities of our inference models across the spectrum of test distributions described in Section 5.1.
Revised: Sensitivity to distribution shift In our first set of experiments, we study the generalization capabilities of our inference models across the spectrum of test distributions described in Section 5.1.

Original: We perform causal discovery from n = 1000 observations in systems of d = 30 variables.
Revised: We perform causal discovery from n = 1000 observations in systems of d = 30 variables.

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

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

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

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

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

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

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

Original: In the L INEAR domain, AVICI performs very well in all metrics and hardly suffers under distribution shift.
Revised: In the L INEAR domain, AVICI performs very well in all metrics and hardly suffers under distribution shift.

Original: In contrast, G RN is the most challenging problem domain and the performance degrades
Revised: In contrast, G RN is the most challenging problem domain and the performance degrades

Original: [Table] [Table]
Revised: [Table] [Table]

Original: more significantly for the o.o.d.
Revised: more significantly for the o.o.d.

Original: scenarios.
Revised: scenarios.

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

Original: This is because AVICI empirically performs better at predicting edges adjacent to large-degree nodes, a common feature of the E. coli and S. cerevisiae graphs not present in the Erd ˝ os-Rényi training structures.
Revised: This is because AVICI empirically performs better at predicting edges adjacent to large-degree nodes, a common feature of the E. coli and S. cerevisiae graphs not present in the Erd ˝ os-Rényi training structures.

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

Original: In Appendix E.1, we additionally report results for generalization from L INEAR to R FF and vice versa, i.e., to entirely unseen function classes of causal mechanisms in addition to the previous o.o.d.
Revised: In Appendix E.1, we additionally report results for generalization from L INEAR to R FF and vice versa, i.e., to entirely unseen function classes of causal mechanisms in addition to the previous o.o.d.

Original: shifts.
Revised: shifts.

Original: Generalization to unseen problem sizes In addition to the sensitivity to distribution shift, we study the ability to generalize to unseen problem sizes.
Revised: Generalization to unseen problem sizes In addition to the sensitivity to distribution shift, we study the ability to generalize to unseen problem sizes.

Original: The bottom row of Figure 3 illustrates the
Revised: The bottom row of Figure 3 illustrates the

Original: AUPRC for the edge predictions of AVICI when varying d and n on unseen in-distribution data.
Revised: AUPRC for the edge predictions of AVICI when varying d and n on unseen in-distribution data.

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

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

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

Original: 6.2 Benchmarking
Revised: 6.2 Benchmarking

Original: Next, we benchmark AVICI against existing algorithms.
Revised: Next, we benchmark AVICI against existing algorithms.

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

Original: LiNGAM (Shimizu et al., 2006),
Revised: LiNGAM (Shimizu et al., 2006),

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

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

Original: We tune the important hyperparameters of each baseline on held-out task instances of each domain.
Revised: We tune the important hyperparameters of each baseline on held-out task instances of each domain.

Original: When computing the evaluation metrics, we favor methods that only predict (interventional) Markov equivalence classes by orienting undirected edges correctly when present in the ground truth.
Revised: When computing the evaluation metrics, we favor methods that only predict (interventional) Markov equivalence classes by orienting undirected edges correctly when present in the ground truth.

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

Original: The benchmarking is performed on the fully o.o.d.
Revised: The benchmarking is performed on the fully o.o.d.

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

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

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

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

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

Original: The same holds for R FF , where GraN-DAG and DCDI perform well but ultimately do not reach the accuracy of AVICI.
Revised: The same holds for R FF , where GraN-DAG and DCDI perform well but ultimately do not reach the accuracy of AVICI.

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

Original: However, provided interventional data, AVICI can use its learned inductive bias to infer plausible causal structures from the noisy gene expressions, even under distribution shift.
Revised: However, provided interventional data, AVICI can use its learned inductive bias to infer plausible causal structures from the noisy gene expressions, even under distribution shift.

Original: This is a promising step towards reliable structure discovery in fields like molecular biology.
Revised: This is a promising step towards reliable structure discovery in fields like molecular biology.

Original: Even without gene knockout data, AVICI achieves nontrival AUROC and AUPRC while classical methods predict close to randomly (Table 9 in Appendix E; see also
Revised: Even without gene knockout data, AVICI achieves nontrival AUROC and AUPRC while classical methods predict close to randomly (Table 9 in Appendix E; see also

Original: [Table]
Revised: [Table]

Original: m
Revised: m

Original: [Table]
Revised: [Table]

Original: i r i c a l p
Revised: i r i c a l p

Original: [Table]
Revised: [Table]

Original: i
Revised: i

Original: [Table]
Revised: [Table]

Original: j =
Revised: j =

Original: [Table]
Revised: [Table]

Original: D
Revised: D

Original: [Table]
Revised: [Table]

Original: L INEAR
Revised: L INEAR

Original: R FF
Revised: R FF

Original: G RN
Revised: G RN

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: Uncertainty quantification Using metrics of calibration, we can evaluate the degree to which predicted edge probabilities are consistent with empirical edge frequencies (DeGroot and Fienberg, 1983; Guo et al., 2017).
Revised: Uncertainty quantification Using metrics of calibration, we can evaluate the degree to which predicted edge probabilities are consistent with empirical edge frequencies (DeGroot and Fienberg, 1983; Guo et al., 2017).

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

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

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

Original: For further explanations, see Appendix B.
Revised: For further explanations, see Appendix B.

Original: Since the baseline algorithms only infer point estimates of the causal structure, we use the nonparametric DAG bootstrap to estimate edge probabilities (Friedman et al. 1999, Appendix D).
Revised: Since the baseline algorithms only infer point estimates of the causal structure, we use the nonparametric DAG bootstrap to estimate edge probabilities (Friedman et al. 1999, Appendix D).

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

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

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

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

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

Original: In Appendix E.5, we additionally report AUROC and AUPRC metrics for all methods.
Revised: In Appendix E.5, we additionally report AUROC and AUPRC metrics for all methods.

Original: We also provide calibration plots analogous to Figure 4 for the baselines (Figure 6), which often show vastly overconfident predictions where the calibration line is far below the diagonal.
Revised: We also provide calibration plots analogous to Figure 4 for the baselines (Figure 6), which often show vastly overconfident predictions where the calibration line is far below the diagonal.

Original: 6.3 Ablations
Revised: 6.3 Ablations

Original: Finally, we analyze the importance of key architecture components of the inference network f ϕ .
Revised: Finally, we analyze the importance of key architecture components of the inference network f ϕ .

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

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

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

Original: Table 2 summarizes the results.
Revised: Table 2 summarizes the results.

Original: Most noticeably, we find that the performance drops significantly when attending only over axis d and aggregating information over axis n only once through pooling after the 2 L self-attention layers.
Revised: Most noticeably, we find that the performance drops significantly when attending only over axis d and aggregating information over axis n only once through pooling after the 2 L self-attention layers.

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

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

Original: ⊤ θ tanh (cid:0) ϕ
Revised: ⊤ θ tanh (cid:0) ϕ

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

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

Original: This variant also satisfies permutation equivariance (cf.
Revised: This variant also satisfies permutation equivariance (cf.

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

Original: Ultimately, we
Revised: Ultimately, we

Original: [Table]
Revised: [Table]

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

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

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

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

Original: Table 1).
Revised: Table 1).

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

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

Original: MB at f 32 precision.
Revised: MB at f 32 precision.

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

Original:  Discussion
Revised:  Discussion

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

Original: Our approach leverages the insight that inductive biases crucial for statistical efficiency in structure learning might be more easily encoded in a simulator than in an inference technique.
Revised: Our approach leverages the insight that inductive biases crucial for statistical efficiency in structure learning might be more easily encoded in a simulator than in an inference technique.

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

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

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

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

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

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

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

Original: Future work might obtain guarantees for AVICI that are similar to learning theory results for the bivariate causal discovery case (Lopez-Paz et al., 2015).
Revised: Future work might obtain guarantees for AVICI that are similar to learning theory results for the bivariate causal discovery case (Lopez-Paz et al., 2015).

Original: Nevertheless, our experiments demonstrate that our inference models are highly robust to distributional shift, suggesting that the trained models could be useful out-of-the-box in causal structure learning tasks outside the domains studied in this paper.
Revised: Nevertheless, our experiments demonstrate that our inference models are highly robust to distributional shift, suggesting that the trained models could be useful out-of-the-box in causal structure learning tasks outside the domains studied in this paper.

Original: In this context, fine-tuning a pretrained
Revised: In this context, fine-tuning a pretrained

Original: AVICI model on labeled real-world datasets is a promising avenue for future work.
Revised: AVICI model on labeled real-world datasets is a promising avenue for future work.

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

Original: Acknowledgments and Disclosure of Funding
Revised: Acknowledgments and Disclosure of Funding

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

Original: This research was supported by the European Research Council (ERC) under the European Union’s Horizon 2020 research and innovation program grant agreement no. 815943 and the Swiss National Science Foundation under NCCR Automation, grant agreement 51NF40 180545.
Revised: This research was supported by the European Research Council (ERC) under the European Union’s Horizon 2020 research and innovation program grant agreement no. 815943 and the Swiss National Science Foundation under NCCR Automation, grant agreement 51NF40 180545.

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