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

Original: Anonymous Author(s)

Original: AffiliationAddress email

Original: Abstract
Revised: Abstract

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

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

Original: Ratherthan searching over structures, we train a variational inference model to predictthe 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 tobypass both the search over graphs and the hand-engineering of suitable scorefunctions.
Revised: This allows us to bypass both the search over graphs and the hand-engineering of suitable score functions.

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

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

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

Original:  Introduction
Revised:  Introduction

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

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

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

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

Original: A key limitation of search approaches to causal discovery is the difficulty of specifying realistic inductive biases.
Revised: A key limitation of search approaches to causal discovery is the difficulty of specifying realistic inductive 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 (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), which are difficult toverify (Dawid, 2010; Reisach et al., 2021).
Revised: Scheines and Ramsey 2016), which are difficult to verify (Dawid, 2010; Reisach et al., 2021).

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

Original: Even with oracle conditional independence tests, they are limitedto 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, thesearch 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.

Original: Theintractable search space ultimately imposes hard constraints on the causal structure, e.g., the nodedegree (Koivisto and Sood, 2004), 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 isto optimize an inference model to solve the inverse problem of predicting a causal structure from aprovided dataset.
Revised: In other words, our goal is to optimize an inference model to solve the inverse problem of predicting a causal structure from a provided dataset.

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

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

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

Original: In doing so, amortizedinference allows us to specify domain-specific inductive biases that are not easily represented bygraph 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 quantifythe epistemic uncertainty in this inference process, we propose to train an amortized variationalinference model for causal discovery (AVICI) that learns to approximate the posterior over causalgraphs given data from a specified data-generating distribution.
Revised: To quantify the epistemic uncertainty in this inference process, we propose to train an amortized variational inference model for causal discovery (AVICI) that learns to approximate the posterior over causal graphs given data from a specified data-generating distribution.

Original: Our model architecture is permutationin- 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 outperformsexisting 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 ourinference models induce robust behavior under substantial distribution shifts of graphs, mechanisms,noise, and problem sizes, which makes them both reliable and versatile for future downstreamuse.
Revised: Moreover, we demonstrate that our inference models induce robust behavior under substantial distribution shifts of graphs, mechanisms, noise, and problem sizes, which makes them both reliable and versatile for future downstream use.

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

Original:  Related Work
Revised:  Related Work

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

Original: Other methods combineconstraint- and score-based ideas (Tsamardinos et al., 2006) or use the noise properties of an SCMthat is postulated to underlie the data-generating process (Shimizu et al., 2006; Hoyer et al., 2008).
Revised: Other methods combine constraint- and score-based ideas (Tsamardinos et al., 2006) or use the noise properties of an SCM that is postulated to underlie the data-generating process (Shimizu et al., 2006; Hoyer et al., 2008).

Original: Deep learning has been applied to causal inference, e.g., for the estimation of treatment effects (Shalitet al., 2017; Yoon et al., 2018) and in instrumental variable analysis (Hartford et al., 2017; Bennettet al., 2019).
Revised: Deep learning has been applied to causal inference, e.g., for the estimation of treatment effects (Shalit 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 causalmechanisms (Goudet et al., 2018; Yu et al., 2019; Lachapelle et al., 2020; Brouillard et al., 2020; Lorchet al., 2021) or to infer the structure for a single dataset (Zhu et al., 2020).
Revised: In structure learning, neural networks have primarily been used to model nonlinear causal mechanisms (Goudet et al., 2018; Yu et al., 2019; Lachapelle et al., 2020; Brouillard et al., 2020; Lorch et al., 2021) or to infer the structure for a single dataset (Zhu et al., 2020).

Original: Prior work applying amortized inference to causal discovery only studied narrowly defined subproblems such as the bivariatecase (Lopez-Paz et al., 2015) and fixed causal mechanisms (Löwe et al., 2022) or only used observational 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 only used observational correlation coefficients for prediction (Li et al., 2020).

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

Original: (2022) also perform causaldiscovery as a supervised learning task.
Revised: (2022) also perform causal discovery as a supervised learning task.

Original: There are significant differences in their approach, and ourmodels generalize to significantly larger problem sizes, even when trained on realistic genomics data.
Revised: There are significant differences in their approach, and our models generalize 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 Discovery 3.1 The Variational Objective

Original: To amortize causal structure learning, we define a data-generating distribution p ( D ) that models thedomain 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 of adata-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 of a data-generating mechanism p ( D | G ) .

Original: p ( D | G ) characterizes a data-generating process whose causalgraph G formalizes all direct causal effects in the system, i.e., how intervening on certain variablesinfluences the realizations of others (Mooij et al., 2016; Peters et al., 2017; Rubenstein et al., 2017).
Revised: p ( D | G ) characterizes a data-generating process whose causal graph G formalizes all direct causal effects in the system, i.e., how intervening on certain variables influences the realizations of others (Mooij et al., 2016; Peters et al., 2017; Rubenstein et al., 2017).

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

Original: To amortize this inference task forthe data-generating distribution p ( D ) , we optimize an inference model f ϕ to predict the variationalparameters θ by minimizing the expected forward KL divergence from the intractable posteriorp ( G | D ) to q ( G ; θ ) for D ∼ p
Revised: To amortize this inference task for the data-generating 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]
Revised: [Equation]

Original: Since it is not tractable to compute the true posterior in (1), we make use of ideas by Barber and
Revised: Since it is not tractable to compute the true posterior in (1) [Equation] 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: p

Original: [Equation] [Equation] [Equation] 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 domaindistribution 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 haveaccess to the causal graph G underlying the generative process of D , i.e., to samples from the jointp ( 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 p ( G, D ) .

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

Original:  Inference Model 
Revised:  Inference Model 4.1 Variational Family

Original: [Figure] 4.1 Variational Family

Original: [Figure]

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

Original: ∼ Bern ( θ

Original: [Equation]
Revised: [Equation]

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

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

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

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

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

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

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

Original: When interventions or gene knockouts are performed, we seto i j
Revised: When interventions or gene knockouts are performed, we set o i j

Original: =
Revised: =

Original: ( x i j , u i j ) and u i j ∈ { 0 , 1 } indicating whether variable j was intervened upon in sample i .[Figure]
Revised: ( x i j , u i j ) and u i j ∈ { 0 , 1 } indicating whether variable j was intervened upon in sample i .

Original: 4.2 Model Architecture
Revised: 4.2 Model Architecture

Original: [Figure]
Revised: n d k

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

Original: Specifically, f ϕ should be permutation invariant across thesample dimension (axis n ) as shuffling the samples shouldnot influence the prediction.
Revised: Specifically, f ϕ should be permutation invariant across the sample dimension (axis n ) as shuffling the samples should not influence the prediction.

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

Original: Reorderingthe variables should permute the causal edge probabilities.
Revised: Reordering the variables should permute the causal edge probabilities.

Original: In the following, we show how to parameterize f ϕ as a neuralnetwork 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 eacho ij to a real-valued vector using a position-wise linear layer,f ϕ operates over a three-dimensional tensor of n rows for theobservations, 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 three-dimensional tensor of n rows for the observations, d columns for the variables, and feature size k .

Original: Attending over axes d and n The core of f ϕ is composed of
Revised: Attending over axes d and n The core of f ϕ is composed of

Original: L = 8 blocks.
Revised: L = 8 blocks.

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

Original: To enable information flow acrossall n × d tokens of the representation, the model alternates in attending over the observation and thevariable 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: The first self-attention block attends over axis d , treating axisn as a batch dimension; the second attends over axis n , treating axis d as a batch dimension (Figure 1).
Revised: The first self-attention block attends over axis d , treating axis n as a batch dimension; the second attends over axis n , treating axis d as a batch dimension (Figure 1).

Original: Variational parameters After building up a representation tensor from the input using the attentionblocks, we max-pool over the observation axis n to obtain a representation ( z 1 , . .
Revised: Variational parameters After building up a representation tensor from the input using the attention blocks, 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 , 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 model theprobability of each edge in the causal graph with an inner product:
Revised: We model the probability of each edge in the causal graph with an inner product:

Revised: [Equation]

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 (4) 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 (4) 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: While feedback loops in G may generally exist, certain domains may be more accurately modeledby acyclic causal structures.
Revised: While feedback loops in G may generally exist, certain domains may be more accurately modeled by acyclic causal structures.

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

Original: Whenever the acyclicity prior is justified, we amend theoptimization problem in (1) with the constraint that q only models acyclic graphs in expectation:
Revised: Whenever the acyclicity prior is justified, we amend the optimization problem in (1) 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 edge probabilities induce an acyclic graph.
Revised: The function h is zero if and only if the edge probabilities induce an acyclic graph.

Original: We build on Leeet al.
Revised: We build on Lee et al.

Original: (2019), who show that acyclicity is equivalent to the spectral radius ρ , i.e., the largest absoluteeigenvalue, of the adjacency matrix being zero.
Revised: (2019), who show that acyclicity is equivalent to the spectral radius ρ , i.e., the largest absolute eigenvalue, of the adjacency matrix being zero.

Original: Following Lee et al.
Revised: Following Lee et al.

Original: (2019), we use h ( W ) :
Revised: (2019), we use h ( W ) :

Original: = ρ ( W )and use power iteration to approximate and differentiate through the largest eigenvalue of f ϕ ( D ) ,which is more efficient than constraints based on matrix powers (Zheng et al., 2018).
Revised: = ρ ( W ) and use power iteration to approximate and differentiate through the largest eigenvalue of f ϕ ( D ) , which is more efficient than constraints based on matrix powers (Zheng et al., 2018).

Original: 4.4 Optimization
Revised: 4.4 Optimization

Original: Parameters: ϕ variational, λ dual, η step size while not converged do for l steps do ∆ [Equation] (cid:0) 
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 (2) with our inferencemodel (3), we can directly use stochastic optimization to train the parameters ϕ .
Revised: Combining the objective in (2) with our inference model (3), we can directly use stochastic optimization to train the parameters ϕ .

Original: The expectationsover p ( G, D ) inside L and F are approximatedusing samples from the data-generating processof the domain.
Revised: The expectations over p ( G, D ) inside L and F are approximated using samples from the data-generating process of the domain.

Original: When enforcing acyclicity, we donot use the augmented Lagrangian to solve theconstrained program max ϕ L ( ϕ )
Revised: When enforcing acyclicity, we do not use the augmented Lagrangian to solve the constrained program max ϕ L ( ϕ )

Original: s.t. F ( ϕ ) = 0 ,which is commonly used in causal discovery (Zheng et al., 2018).
Revised: s.t. F ( ϕ ) = 0 , which is commonly used in causal discovery (Zheng et al., 2018).

Original: Instead, we rely on methodsspecifically tailored for deep learning and solve its dual formulation (Nandwani et al., 2019): [Equation]
Revised: Instead, we rely on methods specifically tailored for deep learning and solve its dual formulation (Nandwani et al., 2019): min λ max ϕ

Original: max
Revised: L ( ϕ ) −

Original: [Equation]
Revised: [Equation]

Revised: F ( ϕ ) (6)

Original: Algorithm 1 summarizes the training procedure, which converges to a local optimum under regularityconditions (Jin et al., 2020).
Revised: Algorithm 1 summarizes the training procedure, which converges to a local optimum under regularity conditions (Jin et al., 2020).

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

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

Original: 5.1 Domains and Simulated Components

Original: We study three domains: two classes of structural causal models (SCMs) as well as semisyntheticsingle-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
Revised: To study the generalization of

Original: AVICI beyond the training distribution p ( D ) , we carefully construct a test distribution ˜ p
Revised: AVICI beyond the training distribution p ( D ) , we carefully construct a test distribution ˜ p

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

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

Original: We consider SCMs with linear functions (L INEAR ) andwith nonlinear functions of random Fourier features (R FF ) that correspond to functions drawn froma Gaussian process (Rahimi and Recht, 2007).
Revised: We consider SCMs with linear functions (L INEAR ) and with nonlinear functions of random Fourier features (R FF ) that correspond to functions drawn from a Gaussian process (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 ) , wesample 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, we simulate homoscedastic Gaussian noise in the training distribution p ( D ) but employheteroscedastic Cauchy and Laplacian noise o.o.d.
Revised: Moreover, we simulate homoscedastic Gaussian noise in the training distribution p ( D ) but employ heteroscedastic Cauchy and Laplacian noise o.o.d.

Original: Interventions clip variables to random values.
Revised: Interventions clip variables to random values.

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

Original: Like L INEAR and R FF , the simulator parameters considered o.o.d.
Revised: Like L INEAR and R FF , the simulator parameters considered o.o.d.

Original: are drawnfrom significantly wider ranges.
Revised: are drawn from significantly wider ranges.

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

Original: , we use the noise levels of different single-cell
Revised: , we use the noise levels of different single-cell

Original: RNA sequencing technologies as calibrated on real datasets.
Revised: RNA sequencing technologies as calibrated on real datasets.

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

Original: Whenever a method uses interventional data in our experiments, half of the dataset consists ofobservational data and half of single-variable interventions.
Revised: Whenever a method uses interventional data in our experiments, half of the dataset consists of observational data and half of single-variable interventions.

Original: In L INEAR and R FF , these interventionsare performed on a subset of target variables containing 50 % of the nodes and in G RN on all nodes.
Revised: In L INEAR and R FF , these interventions are performed on a subset of target variables containing 50 % of the nodes and in G RN on all nodes.

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

Original: Following prior work, we use random graph models and known biologicalnetworks to generate ground-truth causal structures.
Revised: Following prior work, we use random graph models and known biological networks to 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 (Erd˝os and Rényi, 1959; Barabásiand Albert, 1999).
Revised: In all three domains, the training data distribution p ( D ) is induced by simple Erd˝os-Rényi and scale-free graphs (Erd˝os and Rényi, 1959; Barabási and Albert, 1999).

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

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

Original: ( D ) of the L INEAR and R FF domains are simulated usingcausal structures from the Watts-Strogatz model (Watts and Strogatz, 1998), the stochastic blockmodel (Holland et al., 1983), and geometric random graphs (Gilbert, 1961).
Revised: ( D ) of the L INEAR and R FF domains are simulated using causal structures from the Watts-Strogatz model (Watts and Strogatz, 1998), the stochastic block model (Holland et al., 1983), and geometric random graphs (Gilbert, 1961).

Original: In G RN , we extractsubgraphs of the known S. cerevisiae and E. coli GRNs and their effect signs (Marbach et al., 2009).
Revised: In G RN , we extract subgraphs of the known S. cerevisiae and E. coli GRNs and their effect signs (Marbach et al., 2009).

Original: The parameter configurations defining p ( D ) and ˜ p
Revised: The parameter configurations defining p ( D ) and ˜ p

Original: ( D ) are given in Appendix A. We also provide 
Revised: ( D ) are given in Appendix A. We also provide details

Original: Prec
Revised: Prec

Original: [Figure]
Revised: [Figure]

Original: F
Revised: F

Original: [Figure]
Revised: [Figure]

Original: F
Revised: F

Original: [Figure]
Revised: [Figure]

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

Original: 5.2 Metrics and Model Configuration
Revised: 5.2 Metrics and Model Configuration

Original: To asses how well a predicted structure reflects the ground truth, we report the structural Hammingdistance (SHD) and the structural intervention distance (SID) (Peters and Bühlmann, 2015).
Revised: To asses 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: For thesemetrics and for single-edge precision, recall, and F1 score, we collapse the posterior probabilitiespredicted by AVICI to hard predictions using a threshold of 0 . 5 .
Revised: For these metrics and for single-edge precision, recall, and F1 score, we collapse the posterior probabilities predicted by AVICI to hard predictions using a threshold of 0 . 5 .

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

Original: More details on the metrics are given in Appendix B.
Revised: More details on the metrics are given in Appendix B.

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

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

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

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

Original: All experiments throughout this paper are conducted on datasets that AVICI hasnever seen during training.
Revised: All experiments throughout this paper are conducted on datasets that AVICI has never seen during training.

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

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

Original: 6.1 Out-Of-Distribution Generalization

Original: 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 in systems 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 introduce the 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, where fully o.o.d.

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

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

Original: The top row of radar plots in Figure 2 visualizes the results of anempirical sensitivity analysis, disentangling the effects of the three o.o.d.
Revised: The top row of radar plots in Figure 2 visualizes the results of an empirical sensitivity analysis, disentangling the effects of the three o.o.d.

Original: aspects.
Revised: aspects.

Original: In addition tothe 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
Revised: In the

Original: L INEAR domain, AVICI performs very well in all metrics and hardly suffers under distribution shift.
Revised: 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 and the performance degrades more significantly far o.o.d.
Revised: In contrast, G RN is the most challenging and the performance degrades more significantly far o.o.d.

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

Original: G RN .
Revised: G RN .

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

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: [Table]
Revised: [Table]

Original: Generalization to unseen problem sizes We additionally study the ability to generalize to unseenproblem sizes.
Revised: Generalization to unseen problem sizes We additionally study the ability to generalize to unseen problem sizes.

Original: The bottom row of Figure 2 illustrates the performance of AVICI when varying d andn on unseen in-distribution data.
Revised: The bottom row of Figure 2 illustrates the performance of AVICI when varying d and n on unseen in-distribution data.

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

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

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

Original: 6.2 Benchmarking
Revised: 6.2 Benchmarking

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

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

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

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

Original: We favor methods that predict (interventional) Markov equivalence classesby orienting undirected edges correctly when present in the ground truth.
Revised: We favor methods that predict (interventional) Markov equivalence classes by orienting undirected edges correctly when present in the ground truth.

Original: See Appendix D for details.
Revised: See Appendix D for details.

Original: The benchmark is performed fully o.o.d.
Revised: The benchmark is performed fully o.o.d.

Original: , i.e., under distribution shifts on causal graphs, mechanisms,and noise distributions w.r.t.
Revised: , 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 scoresof all methods given n = 1000 observations for d = 30 variables.
Revised: Table 1 shows the SID and F1 scores of all methods given n = 1000 observations for d = 30 variables.

Original: The AVICI model trained on
Revised: The AVICI model trained on

Original: L INEAR outperforms all baselines, both given observational or interventional data, despite operatingunder significant distribution shift.
Revised: 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
Revised: The same holds for

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

Original: In G RN , the inductive bias encoded by classical methods reaches its limits.
Revised: In G RN , the inductive bias encoded by classical methods reaches its limits.

Original: However, providedinterventional data, AVICI can use its learned inductive bias to infer plausible structures from the noisygene expressions, even under distribution shift, which is a promising step towards reliable structurediscovery in fields like molecular biology.
Revised: However, provided interventional data, AVICI can use its learned inductive bias to infer plausible 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: Even without knockout data, AVICI achieves nontrival
Revised: Even without knockout data, AVICI achieves nontrival

Original: AUROC and AUPRC while classical methods predict close to randomly (Table 7 in Appendix E).
Revised: AUROC and AUPRC while classical methods predict close to randomly (Table 7 in Appendix E).

Original: Results for in-distribution data and for larger graphs of d = 100 variables are given in Appendices
Revised: Results for in-distribution data and for larger graphs of d = 100 variables are given in Appendices

Original: E.1 and E.2.
Revised: E.1 and E.2.

Original: In Appendix E.3, we also report results for a real proteomics dataset (Sachs et al., 2005).
Revised: In Appendix E.3, we also report results for a real proteomics dataset (Sachs et al., 2005).

Original:  Discussion
Revised:  Discussion

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

Original: Our approach leverages the insight that inductive biasescrucial for statistical efficiency in structure learning might be more easily encoded in a simulator thanin 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 learningproblems 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: 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 empirically induces acyclicity, this constraint is not satisfied with certainty.

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

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

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

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

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

Original: A Domain Specification and Simulation
Revised: A Domain Specification and Simulation

Original: In this section, we define p ( D ) and ˜ p
Revised: In this section, we define p ( D ) and ˜ p

Original: ( D ) concretely in terms of the parameters and notation introducedin the main text.
Revised: ( D ) concretely in terms of the parameters and notation introduced in the main text.

Original: Based on these definitions, Table 2 summarizes all parameters of the data-generatingprocesses for L INEAR and R FF and specifies how they are sampled for a random task instance.
Revised: Based on these definitions, Table 2 summarizes all parameters of the data-generating processes for L INEAR and R FF and specifies how they are sampled for a random task instance.

Original: Table3 lists the same specifications for the G RN domain.
Revised: Table 3 lists the same specifications for the G RN domain.

Original: A.1 Causal StructuresA.1.1 Random graph models
Revised: A.1 Causal Structures A.1.1 Random graph models

Original: In Erd˝os-Rényi graphs, each edge is sampled independently with a fixed probability (Erd˝os and
Revised: In Erd˝os-Rényi graphs, each edge is sampled independently with a fixed probability (Erd˝os and

Original: Rényi, 1959).
Revised: Rényi, 1959).

Original: We scale this probability to obtain O ( d ) edges in expectation.
Revised: We scale this probability to obtain O ( d ) edges in expectation.

Original: Scale-free graphs aregenerated by a sequential preferential attachment process, where in- or outgoing edges of node ito the previous i − 1 nodes are sampled with probability ∝ deg ( j ) α (Barabási and Albert, 1999).
Revised: Scale-free graphs are generated by a sequential preferential attachment process, where in- or outgoing edges of node i to the previous i − 1 nodes are sampled with probability ∝ deg ( j ) α (Barabási and Albert, 1999).

Original: Watts-Strogatz graphs are k -dimensional lattices, whose edges get rewired globally to random nodeswith a specified probability (Watts and Strogatz, 1998).
Revised: Watts-Strogatz graphs are k -dimensional lattices, whose edges get rewired globally to random nodes with a specified probability (Watts and Strogatz, 1998).

Original: The stochastic block model generalizes
Revised: The stochastic block model generalizes

Original: Erd˝os-Rényi to capture community structure.
Revised: Erd˝os-Rényi to capture community structure.

Original: Splitting the nodes into a random partition of so-calledblocks, the inter-block edge probability is dampened by a multiplying factor compared to the intrablock probability, also tuned to result in O ( d ) edges in expectation (Holland et al., 1983).
Revised: Splitting the nodes into a random partition of so-called blocks, the inter-block edge probability is dampened by a multiplying factor compared to the intrablock probability, also tuned to result in O ( d ) edges in expectation (Holland et al., 1983).

Original: Lastly,geometric random graphs model connectivity based on two-dimensional Euclidian distance withinsome radius, where nodes are randomly placed inside the unit square (Gilbert, 1961).
Revised: Lastly, geometric random graphs model connectivity based on two-dimensional Euclidian distance within some radius, where nodes are randomly placed inside the unit square (Gilbert, 1961).

Original: For undirected random graph models, we orient edges by selecting the upper-triangular half ofthe adjacency matrix.
Revised: For undirected random graph models, we orient edges by selecting the upper-triangular half of the adjacency matrix.

Original: The classes of random graph models are sampled in equal proportion whengenerating a set of evaluation datasets (Tables 2 and 3).
Revised: The classes of random graph models are sampled in equal proportion when generating a set of evaluation datasets (Tables 2 and 3).

Original: A.1.2 Subgraph Extraction from Real-World Networks
Revised: A.1.2 Subgraph Extraction from Real-World Networks

Original: For the evaluation in the G RN domain, we sample realistic causal graphs by extracting subgraphsfrom the known E. coli and S. cerevisiae regulatory networks.
Revised: For the evaluation in the G RN domain, we sample realistic causal graphs by extracting subgraphs from the known E. coli and S. cerevisiae regulatory networks.

Original: For this, we rely on the procedure by
Revised: For this, we rely on the procedure by

Original: Marbach et al.
Revised: Marbach et al.

Original: (2009), which is also used by Schaffter et al.
Revised: (2009), which is also used by Schaffter et al.

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

Original: Their graph extraction method is carefully designed to capture the structural properties of biologicalnetworks by preserving the functional and structural properties of the source network.
Revised: Their graph extraction method is carefully designed to capture the structural properties of biological networks by preserving the functional and structural properties of the source network.

Original: Algorithm The procedure extracts a random subgraph of the source network by selecting a subsetof nodes V , and then returning the graph containing all edges from the source network covered by
Revised: Algorithm The procedure extracts a random subgraph of the source network by selecting a subset of nodes V , and then returning the graph containing all edges from the source network covered by

Original: V .
Revised: V .

Original: Starting from a random seed node, the algorithm proceeds by iteratively adding new nodes to
Revised: Starting from a random seed node, the algorithm proceeds by iteratively adding new nodes to

Original: V .
Revised: V .

Original: In each step, this new node is selected from the set of neighbors of the current set V .
Revised: In each step, this new node is selected from the set of neighbors of the current set V .

Original: The neighborto be added is selected greedily such that the resulting subgraph has maximum modularity (Marbachet al., 2009).
Revised: The neighbor to be added is selected greedily such that the resulting subgraph has maximum modularity (Marbach et al., 2009).

Original: To introduce additional randomness, Marbach et al.
Revised: To introduce additional randomness, Marbach et al.

Original: (2009) propose to randomly draw the new nodefrom the set of neighbors inducing the top- p percent of the most modular graphs.
Revised: (2009) propose to randomly draw the new node from the set of neighbors inducing the top- p percent of the most modular graphs.

Original: In our experiments,we adopt the latter with p = 20 percent, similar to Schaffter et al.
Revised: In our experiments, we adopt the latter with p = 20 percent, similar to Schaffter et al.

Original: The original method of
Revised: The original method of

Original: Marbach et al.
Revised: Marbach et al.

Original: (2009) is intended for undirected graphs.
Revised: (2009) is intended for undirected graphs.

Original: Thus, we use the undirected skeleton of thesource network for the required modularity and neighborhood computation.
Revised: Thus, we use the undirected skeleton of the source network for the required modularity and neighborhood computation.

Original: Real GRNs We take the E. coli and S. cerevisiae regulatory networks as provided by the
Revised: Real GRNs We take the E. coli and S. cerevisiae regulatory networks as provided by the

Original: GeneNetWeaver repository, 1 which have 1565 and 4441 nodes (genes), respectively.
Revised: GeneNetWeaver repository, 1 which have 1565 and 4441 nodes (genes), respectively.

Original: For E. coli , wealso know the true signs of a large proportion of causal effects.
Revised: For E. coli , we also know the true signs of a large proportion of causal effects.

Original: When extracting a random subgraphfrom E. coli , we take the true signs of the effects and map them onto the randomly sampled interactionterms k ∈ R d × d used by SERGIO; cf.
Revised: When extracting a random subgraph from E. coli , we take the true signs of the effects and map them onto the randomly sampled interaction terms k ∈ R d × d used by SERGIO; cf.

Original: Section A.2.2.
Revised: Section A.2.2.

Original: When the interaction signs are unknown oruncertain in E. coli , we impute a random sign in the interaction terms k of SERGIO based on thefrequency of known positive and negative signs in the E. coli graph.
Revised: When the interaction signs are unknown or uncertain in E. coli , we impute a random sign in the interaction terms k of SERGIO based on the frequency of known positive and negative signs in the E. coli graph.

Original: Empirically, individual genes in E. coli tend to predominantly have either up- or down-regulatingeffects on their causal children.
Revised: Empirically, individual genes in E. coli tend to predominantly have either up- or down-regulating effects on their causal children.

Original: To capture this aspect in S. cerevisiae also, we fit the probabilityof an up-regulating effect caused by a given gene in E. coli to a Beta distribution.
Revised: To capture this aspect in S. cerevisiae also, we fit the probability of an up-regulating effect caused by a given gene in E. coli to a Beta distribution.

Original: For each node j in an extracted subgraph of S. cerevisiae , we draw a probability p j from this Beta distribution and
Revised: For each node j in an extracted subgraph of S. cerevisiae , we draw a probability p j from this Beta distribution and

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

Original: then sample the effect signs for the outgoing edges of node j using p j .
Revised: then sample the effect signs for the outgoing edges of node j using p j .

Original: As a result, the genes in thesubgraphs of S. cerevisiae individually also have mostly up- or down-regulating effects.
Revised: As a result, the genes in the subgraphs of S. cerevisiae individually also have mostly up- or down-regulating effects.

Original: Maximumlikelihood estimation for this Beta distribution yielded α = 0 .
Revised: Maximum likelihood estimation for this Beta distribution yielded α = 0 .

Original: 2588 and β = 0 .
Revised: 2588 and β = 0 .

Original: The E. coli and S. cerevisiae graphs and effect signs used in the experiments are taken from the
Revised: The E. coli and S. cerevisiae graphs and effect signs used in the experiments are taken from the

Original: GeneNetWeaver repository (Schaffter et al., 2011) (MIT License).
Revised: GeneNetWeaver repository (Schaffter et al., 2011) (MIT License).

Original: A.2 Data-Generating ProcessesA.2.1 Structural Causal Models
Revised: A.2 Data-Generating Processes A.2.1 Structural Causal Models

Original: In the L INEAR and R FF domains, the data-generating processes are modeled by structural causalmodels (SCMs).
Revised: In the L INEAR and R FF domains, the data-generating processes are modeled by structural causal models (SCMs).

Original: In this work, we consider SCMs with causal mechanisms that model each causalvariable x j given its parents x pa ( j ) as
Revised: In this work, we consider SCMs with causal mechanisms that model each causal variable x j given its parents x pa ( j ) as

Original: [Equation]
Revised: [Equation]

Original: (7) where the noise ϵ j is additive and may be heteroscedastic through an input-dependent noise scaleh j ( x pa ( j ) ) .
Revised: where the noise ϵ j is additive and may be heteroscedastic through an input-dependent noise scale h j ( x pa ( j ) ) .

Original: We write pa ( j ) when indexing x
Revised: We write pa ( j ) when indexing x

Original: at the parents of node j .
Revised: at the parents of node j .

Original: In the heteroscedastic setting,we parameterize the noise scales as h j ( x ) =
Revised: In the heteroscedastic setting, we parameterize the noise scales as h j ( x ) =

Original: log(

Revised: log(

Original: + exp( g j ( x )) for a set of nonlinear functions
Revised: + exp( g j ( x )) for a set of nonlinear functions

Original: g j .
Revised: g j .

Original: In this work, the scale of each noise distribution p ( ϵ j ) is random and thus different for each variablex j .
Revised: In this work, the scale of each noise distribution p ( ϵ j ) is random and thus different for each variable x j .

Original: Prior to performing inference with AVICI or any baseline, each set of SCM observations D isstandardized variable-wise by subtracting its mean and dividing by its standard deviation, so thateach x j has mean 0 and variance 1 (Reisach et al., 2021).
Revised: Prior to performing inference with AVICI or any baseline, each set of SCM observations D is standardized variable-wise by subtracting its mean and dividing by its standard deviation, so that each x j has mean 0 and variance 1 (Reisach et al., 2021).

Original: In the L INEAR domain, the functions f
Revised: In the L INEAR domain, the functions f

Original: j are given by affine transforms
Revised: j are given by affine transforms

Original: [Equation]
Revised: [Equation]

Original: (8) whose weights w j and bias b j are sampled independently for each f j .
Revised: 8) whose weights w j and bias b j are sampled independently for each f j .

Original: In the R FF domain, thefunctions f j modeling each causal variable x j given its parents x pa ( j ) are drawn from a Gaussian
Revised: In the R FF domain, the functions f j modeling each causal variable x j given its parents x pa ( j ) are drawn from a Gaussian

Original: Process
Revised: Process

Original: [Equation]
Revised: [Equation]

Original: with bias b j and squared exponential (SE) kernel k j ( x , x ′ )
Revised: with bias b j and squared exponential (SE) kernel k j ( x , x ′ )

Original: =
Revised: =

Original: c 2 j exp (cid:0) −∥ x − x ′ ∥
Revised: c 2 j exp (cid:0) −∥ x − x ′ ∥

Original: 22 / 2 ℓ 2 j (cid:1) with lengthscale ℓ j and output scale c j .
Revised: 22 / 2 ℓ 2 j (cid:1) with length scale ℓ j and output scale c j .

Original: The parameters b j , c j , and ℓ j are sampled independently for eachvariable j .
Revised: The parameters b j , c j , and ℓ j are sampled independently for each variable j .

Original: To obtain explicit function draws f j from the GP, we approximate f j with random Fourierfeatures (Rahimi and Recht, 2007).
Revised: To obtain explicit function draws f j from the GP, we approximate f j with random Fourier features (Rahimi and Recht, 2007).

Original: Specifically, we can obtain f j ∼ GP ( b j , k
Revised: Specifically, we can obtain f j ∼ GP ( b j , k

Original: ( x , x ′ )) for a SE kernelk with length scale ℓ j and output scale c j by sampling
Revised: ( x , x ′ )) for a SE kernel k with length scale ℓ j and output scale c j by sampling

Original: [Equation]
Revised: [Equation]

Original: with α ( m ) ∼ N (0 , 1) , ω ( m ) ∼ N (0 , I ) , and δ ( m ) ∼ Unif (0 , 2 π ) .
Revised: with α ( m ) ∼ N (0 , 1) , ω ( m ) ∼ N (0 , I ) , and δ ( m ) ∼ Unif (0 , 2 π ) .

Original: Throughout this work, we use
Revised: Throughout this work, we use

Original: M = 100 .
Revised: M = 100 .

Original: The function draws become faithful GP samples as M → ∞ (Rahimi and Recht, 2007).
Revised: The function draws become faithful GP samples as M → ∞ (Rahimi and Recht, 2007).

Original: A.2.2 Single-Cell Gene Expression Data
Revised: A.2.2 Single-Cell Gene Expression Data

Original: In the G RN domain, our goal is to evaluate causal discovery from realistic gene expression data.
Revised: In the G RN domain, our goal is to evaluate causal discovery from realistic gene expression data.

Original: Thereexist several models to simulate the mechanisms, intervention types, and technical measurement noiseunderlying single-cell expression data of gene regulatory networks (Schaffter et al., 2011; Huynh-Thuand Sanguinetti, 2019; Dibaeinia and Sinha, 2020).
Revised: There exist several models to simulate the mechanisms, intervention types, and technical measurement noise underlying single-cell expression data of gene regulatory networks (Schaffter et al., 2011; Huynh-Thu and Sanguinetti, 2019; Dibaeinia and Sinha, 2020).

Original: We use the simulator by Dibaeinia and Sinha(2020) (SERGIO) because it resembles the data collected by modern high-throughput single-cell RNAsequencing (scRNA-seq) technologies.
Revised: We use the simulator by Dibaeinia and Sinha (2020) (SERGIO) because it resembles the data collected by modern high-throughput single-cell RNA sequencing (scRNA-seq) technologies.

Original: Related genomics simulators, for example, GeneNetWeaver(Schaffter et al., 2011), were developed for the simulation of microarray gene expression platforms.
Revised: Related genomics simulators, for example, GeneNetWeaver (Schaffter et al., 2011), were developed for the simulation of microarray gene expression platforms.

Original: In the following, we give an overview of how to simulate scRNA-seq data with SERGIO.
Revised: In the following, we give an overview of how to simulate scRNA-seq data with SERGIO.

Original: Dibaeiniaand Sinha (2020) provide all the details and additional background from the related literature.
Revised: Dibaeinia and Sinha (2020) provide all the details and additional background from the related literature.

Original: Simulation Given a causal graph over d genes and a specification of the simulation parameters,
Revised: Simulation Given a causal graph over d genes and a specification of the simulation parameters,

Original: SERGIO generates a synthetic scRNA-seq dataset D in two stages.
Revised: SERGIO generates a synthetic scRNA-seq dataset D in two stages.

Original: The n observations in Dcorrespond to n cell samples, that is, the expressions of the d genes recorded in a single cellcorresponds to one row in D .
Revised: The n observations in D correspond to n cell samples, that is, the expressions of the d genes recorded in a single cell corresponds to one row in D .

Original: In the first stage, SERGIO simulates clean gene expressions by sampling randomly-timed snapshotsfrom the steady state of a dynamical system.
Revised: In the first stage, SERGIO simulates clean gene expressions by sampling randomly-timed snapshots from the steady state of a dynamical system.

Original: In this regulatory process, the genes are expressed at ratesinfluenced by other genes using the chemical Langevin equation, similar to Schaffter et al.
Revised: In this regulatory process, the genes are expressed at rates influenced by other genes using the chemical Langevin equation, similar to Schaffter et al.

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

Original: The source nodes in the causal graph G are denoted master regulators(MRs), whose expressions evolve at constant production and decay rates.
Revised: The source nodes in the causal graph G are denoted master regulators (MRs), whose expressions evolve at constant production and decay rates.

Original: The expressions of alldownstream genes evolve nonlinearly under production rates caused by the expression of their causalparents in G .
Revised: The expressions of all downstream genes evolve nonlinearly under production rates caused by the expression of their causal parents in G .

Original: Cell types are defined by specifications of the MR production rates, which significantlyinfluence the evolution of the system.
Revised: Cell types are defined by specifications of the MR production rates, which significantly influence the evolution of the system.

Original: Thus, the dataset contains variation due to biological systemnoise within collections of cells of the same type and due to different cell types.
Revised: Thus, the dataset contains variation due to biological system noise within collections of cells of the same type and due to different cell types.

Original: Ultimately, wegenerate single-cell samples collected from five to ten cell types (Dibaeinia and Sinha, 2020).
Revised: Ultimately, we generate single-cell samples collected from five to ten cell types (Dibaeinia and Sinha, 2020).

Original: In the second stage, the clean gene expressions sampled previously are corrupted with technicalmeasurement error that resembles the noise phenomena found in real scRNA-seq data:
Revised: In the second stage, the clean gene expressions sampled previously are corrupted with technical measurement error that resembles the noise phenomena found in real scRNA-seq data:

Original: • outlier genes : a small set of genes have unusually high expression across measurements • library size : different cells have different total UMI counts, following a log-normal distribution • dropouts : a high percentage of genes are recorded with zero expression in a given measurement• unique molecule identifier (UMI) counts : we observe Poisson-distributed count data rather than the clean expression values
Revised: • outlier genes : a small set of genes have unusually high expression across measurements • library size : different cells have different total UMI counts, following a log-normal distribution • dropouts : a high percentage of genes are recorded with zero expression in a given measurement • unique molecule identifier (UMI) counts : we observe Poisson-distributed count data rather than the clean expression values

Original: To configure these noise modules, we use the parameters calibrated by Dibaeinia and Sinha (2020)for datasets from different scRNA-seq technologies.
Revised: To configure these noise modules, we use the parameters calibrated by Dibaeinia and Sinha (2020) for datasets from different scRNA-seq technologies.

Original: We extend SERGIO to allow for the generationof knockout intervention experiments.
Revised: We extend SERGIO to allow for the generation of knockout intervention experiments.

Original: For this, we force the production rate of knocked-out genes tozero during simulation.
Revised: For this, we force the production rate of knocked-out genes to zero during simulation.

Original: Our implementation uses the public source code by (Dibaeinia and Sinha,2020), which is available under a GNU General Public License v3.0.
Revised: Our implementation uses the public source code by (Dibaeinia and Sinha, 2020), which is available under a GNU General Public License v3.0.

Original: Parameters Given a causal graph G , the parameters SERGIO requires to simulate c cell types of dgenes are:
Revised: Parameters Given a causal graph G , the parameters SERGIO requires to simulate c cell types of d genes are:

Original: • k ∈ R d × d : interaction strengths (only used if edge i → j exists in G ) • b ∈ R d × c + : MR production rates (only used if gene j is a source node in G ) • γ ∈ R d × d
Revised: • k ∈ R d × d : interaction strengths (only used if edge i → j exists in G ) • b ∈ R d × c + : MR production rates (only used if gene j is a source node in G ) • γ ∈ R d × d

Original: + : Hill function coefficients controlling nonlinearity of interactions• λ ∈ R d : decay rates per gene• ζ ∈ R d + : scales of stochastic process noise per gene for chemical Langevin equations
Revised: + : Hill function coefficients controlling nonlinearity of interactions • λ ∈ R d : decay rates per gene • ζ ∈ R d + : scales of stochastic process noise per gene for chemical Langevin equations

Original: The technical noise components are configured by:
Revised: The technical noise components are configured by:

Original: • p outlier ∈
Revised: • p outlier ∈

Original: [0 , 1] : probability that a gene is an outlier gene • µ outlier ∈ R + , σ outlier ∈ R + : parameters of the log-normal distribution for the outlier multipliers• µ lib ∈ R + , σ lib ∈ R + : parameters of the log-normal distribution for the library size multipliers • δ ∈ [0 , 100] , ξ ∈ R + : dropout percentile and temperature of the logistic function parameterizing the dropout probability of a recorded expression
Revised: [0 , 1] : probability that a gene is an outlier gene • µ outlier ∈ R + , σ outlier ∈ R + : parameters of the log-normal distribution for the outlier multipliers • µ lib ∈ R + , σ lib ∈ R + : parameters of the log-normal distribution for the library size multipliers • δ ∈ [0 , 100] , ξ ∈ R + : dropout percentile and temperature of the logistic function parameterizing the dropout probability of a recorded expression

Original: In our experiments, the simulator parameters are selected in the ranges suggested by Dibaeinia and
Revised: In our experiments, the simulator parameters are selected in the ranges suggested by Dibaeinia and

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

Original: Standardization There are several ways to preprocess and normalize single-cell transcriptomicdata for downstream use (Robinson et al., 2010).
Revised: Standardization There are several ways to preprocess and normalize single-cell transcriptomic data for downstream use (Robinson et al., 2010).

Original: For simplicity, we employ log 2 counts-per-million(CPM) normalization, which normalizes the total UMI counts per sample and then log 2 -transformsthe relative count values.
Revised: For simplicity, we employ log 2 counts-per-million (CPM) normalization, which normalizes the total UMI counts per sample and then log 2 -transforms the relative count values.

Original: Specifically, the CPM value for gene j in sample i is defined asx cpm ,i j := x i j · 10l i with library sizel i =
Revised: Specifically, the CPM value for gene j in sample i is defined as [Equation] library size

Original: d (cid:88)
Revised: [Equation]

Original: j = x ij . For zero expressions x ij , the log 2 -CPM values are imputed with zero.
Revised: For zero expressions x ij , the log 2 -CPM values are imputed with zero.

Original: The remaining log 2 -CPMvalues range between 10 and 19 , so we shift and scale the values before performing causal discovery.
Revised: The remaining log 2 -CPM values range between 10 and 19 , so we shift and scale the values before performing causal discovery.

Original: To replicate the sparsity pattern and the relative ordering of values within samples in the originaldataset D , we standardize the nonzero log 2 -CPM values by subtracting the minimum (instead of themean) and dividing by the overall standard deviation.
Revised: To replicate the sparsity pattern and the relative ordering of values within samples in the original dataset D , we standardize the nonzero log 2 -CPM values by subtracting the minimum (instead of the mean) and dividing by the overall standard deviation.

Original: All methods considered in Section 6, including
Revised: All methods considered in Section 6, including

Original: AVICI, work with G RN data in this standardized log 2 -CPM format.
Revised: AVICI, work with G RN data in this standardized log 2 -CPM format.

Original: B Evaluation Metrics
Revised: B Evaluation Metrics

Original: We report several metrics to assess how well the predicted causal structures reflect the ground-truthgraph.
Revised: We report several metrics to assess how well the predicted causal structures reflect the ground-truth graph.

Original: We measure the overall accuracy of the predictions and how well-calibrated the estimateduncertainties in the edge predictions are, since AVICI predicts marginal probabilities q ( g i,j ; θ i,j ) forevery edge.
Revised: We measure the overall accuracy of the predictions and how well-calibrated the estimated uncertainties in the edge predictions are, since AVICI predicts marginal probabilities q ( g i,j ; θ i,j ) for every edge.

Original: Unless evaluating these edge probabilities, we use a decision threshold of 0 .
Revised: Unless evaluating these edge probabilities, we use a decision threshold of 0 .

Original: 5 to convertthe AVICI prediction to a single graph G .
Revised: 5 to convert the AVICI prediction to a single graph G .

Original:  The structural hamming distance (SHD) (Tsamardinos et al., 2006)reflects the graph edit distance between two graphs, i.e., the edge changes required to transform
Revised: Structural and edge accuracy The structural hamming distance (SHD) (Tsamardinos et al., 2006) reflects the graph edit distance between two graphs, i.e., the edge changes required to transform

Original: G into G ′ .
Revised: G into G ′ .

Original: By contrast, the structural intervention distance (SID) (Peters and Bühlmann, 2015)quantifies the closeness of two DAGs in terms of their valid adjustment sets, which more closelyresembles our intentions of using the inferred graph for downstream causal inference tasks.
Revised: By contrast, the structural intervention distance (SID) (Peters and Bühlmann, 2015) quantifies the closeness of two DAGs in terms of their valid adjustment sets, which more closely resembles our intentions of using the inferred graph for downstream causal inference tasks.

Original: SHD and SID capture global and structural similarity to the ground truth, but notions like precisionand recall at the edge level are not captured well.
Revised: SHD and SID capture global and structural similarity to the ground truth, but notions like precision and recall at the edge level are not captured well.

Original: SID is zero if and only if the true DAG is a subgraphof the predicted graph, which can reward dense predictions (Prop. 8 by Peters and Bühlmann (2015):
Revised: SID is zero if and only if the true DAG is a subgraph of the predicted graph, which can reward dense predictions (Prop. 8 by Peters and Bühlmann (2015):

Original: SID ( G, G ′ )
Revised: SID ( G, G ′ )

Original: = 0 when G is empty and G ′ is fully connected).
Revised: = 0 when G is empty and G ′ is fully connected).

Original: Conversely, the trivial prediction of anempty graph achieves highly competitive SHD scores for sparse graphs.
Revised: Conversely, the trivial prediction of an empty graph achieves highly competitive SHD scores for sparse graphs.

Original: For this reason, we report additional metrics that quantify both the trade-off between precision andrecall of edges as well as the calibration of their uncertainty estimates.
Revised: For this reason, we report additional metrics that quantify both the trade-off between precision and recall of edges as well as the calibration of their uncertainty estimates.

Original: Specifically, given the binarypredictions for all d 2 possible edges in the graph G , we compute the edge precision, edge recall, andtheir harmonic mean (F1-score) for each test case and estimate their means and standard errors acrossthe test cases.
Revised: Specifically, given the binary predictions for all d 2 possible edges in the graph G , we compute the edge precision, edge recall, and their harmonic mean (F1-score) for each test case and estimate their means and standard errors across the test cases.

Original: Since the F1-score is high only when precision and recall are high, both empty anddense predictions are penalized and no trivial prediction scores well.
Revised: Since the F1-score is high only when precision and recall are high, both empty and dense predictions are penalized and no trivial prediction scores well.

Original: Edge confidence To evaluate the edge probabilities predicted by AVICI, we compute the areasunder the precision-recall curve (AUPRC) and receiver operating characteristic (AUROC) whenconverting the probabilities into binary predictions using varying decision thresholds (Friedman and
Revised: Edge confidence To evaluate the edge probabilities predicted by AVICI, we compute the areas under the precision-recall curve (AUPRC) and receiver operating characteristic (AUROC) when converting the probabilities into binary predictions using varying decision thresholds (Friedman and

Original: Koller, 2003).
Revised: Koller, 2003).

Original: Both statistics capture different aspects of the confidence estimates.
Revised: Both statistics capture different aspects of the confidence estimates.

Original: The AUROCis insensitive to changes in class imbalance (edge vs. no-edge) for a given d .
Revised: The AUROC is insensitive to changes in class imbalance (edge vs. no-edge) for a given d .

Original: However, when thenumber of variables d in sparse graphs of O ( d ) edges increases, AUROC increasingly discounts theaccuracy on the shrinking proportion of edges in the ground truth, making AUPRC more suitable forcomparisons ranging over different d .
Revised: However, when the number of variables d in sparse graphs of O ( d ) edges increases, AUROC increasingly discounts the accuracy on the shrinking proportion of edges in the ground truth, making AUPRC more suitable for comparisons ranging over different d .

Original: The AUROC is equivalent to the probability that the methodranks a randomly chosen positive instance (i.e., an edge i → j present in the ground truth) higher thana randomly chosen negative instance (i.e., an edge i . . .
Revised: The AUROC is equivalent to the probability that the method ranks a randomly chosen positive instance (i.e., an edge i → j present in the ground truth) higher than a randomly chosen negative instance (i.e., an edge i . . .

Original: j absent in the ground truth)
Revised: j absent in the ground truth)

Original: (Fawcett, 2004).
Revised: (Fawcett, 2004).

Original: C Inference Model Details 
Revised: C Inference Model Details C.1 Optimization

Original: C.1 Optimization

Original: Batch sizes Each AVICI model is trained as described in Algorithm 1.
Revised: Batch sizes Each AVICI model is trained as described in Algorithm 1.

Original: The objective L ( ϕ ) relies onsamples from the domain distribution p ( G, D ) to perform Monte Carlo estimation of the expectations.
Revised: The objective L ( ϕ ) relies on samples from the domain distribution p ( G, D ) to perform Monte Carlo estimation of the expectations.

Original: During training, the number of variables d in the simulated systems are chosen randomly fromd ∈ { 2 , 5 , 10 , 20 , 30 , 40 , 50 } (12)
Revised: During training, the number of variables d in the simulated systems are chosen randomly from d ∈ { 2 , 5 , 10 , 20 , 30 , 40 , 50 } (12)

Original: The datasets D in the training distributions have n = 200 samples, where with probability 0 .
Revised: The datasets D in the training distributions have n = 200 samples, where with probability 0 .

Original: 5 ,the observations in a given dataset contain 50 interventional samples.
Revised: 5 , the observations in a given dataset contain 50 interventional samples.

Original: The dimensionality of thesetraining instances G, D varies significantly with the number of variables d and, therefore, so do thememory requirements of the forward passes of the inference model f ϕ .
Revised: The dimensionality of these training instances G, D varies significantly with the number of variables d and, therefore, so do the memory requirements of the forward passes of the inference model f ϕ .

Original: Given these differences in problem size, we make efficient use of the GPU resources during training byperforming individual primal updates in Algorithm 1 using only training instances ( G, D ) with exactly d variables, where d is randomly sampled in each update step.
Revised: Given these differences in problem size, we make efficient use of the GPU resources during training by performing individual primal updates in Algorithm 1 using only training instances ( G, D ) with exactly d variables, where d is randomly sampled in each update step.

Original: Fixing the number of observations ton = 200 , this allows us to increase the batch size for each considered d to the maximum possiblegiven the available GPU memory (in our case ranging from batch sizes of 27 for d = 2 down to 6 ford = 50 , per 24 GiB GPU device).
Revised: Fixing the number of observations to n = 200 , this allows us to increase the batch size for each considered d to the maximum possible given the available GPU memory (in our case ranging from batch sizes of 27 for d = 2 down to 6 for d = 50 , per 24 GiB GPU device).

Original: During training, the sampling probability of a given d is tuned toensure that f ϕ sees the same number of examples for each d , i.e., we oversample higher d , for whichthe effective batch size per update step is smaller.
Revised: During training, the sampling probability of a given d is tuned to ensure that f ϕ sees the same number of examples for each d , i.e., we oversample higher d , for which the effective batch size per update step is smaller.

Original: When considering the acyclicity constraint, thepenalty F ( ϕ ) is estimated using the same minibatch as for L ( ϕ ) .
Revised: When considering the acyclicity constraint, the penalty F ( ϕ ) is estimated using the same minibatch as for L ( ϕ ) .

Original: Buffer Since we have access to the complete data-generating process rather than only a fixed dataset,we approximate L ( ϕ ) with minibatches that are sampled uniformly randomly from a buffer, which iscontinually updated with fresh data from p ( G, D ) .
Revised: Buffer Since we have access to the complete data-generating process rather than only a fixed dataset, we approximate L ( ϕ ) with minibatches that are sampled uniformly randomly from a buffer, which is continually updated with fresh data from p ( G, D ) .

Original: Specifically, we initialize a first-in-first-out bufferthat holds 200 pairs ( G, D ) for each unique number of variables d considered during training.
Revised: Specifically, we initialize a first-in-first-out buffer that holds 200 pairs ( G, D ) for each unique number of variables d considered during training.

Original: A poolof asynchronous single-CPU workers then constantly generates novel training data and replaces theoldest instances in the buffer using a producer-consumer workflow.
Revised: A pool of asynchronous single-CPU workers then constantly generates novel training data and replaces the oldest instances in the buffer using a producer-consumer workflow.

Original: We implement this buffer usingan Apache PyArrow Plasma object store (Apache Licence 2.0).
Revised: We implement this buffer using an Apache PyArrow Plasma object store (Apache Licence 2.0).

Original: During training, we used 128 CPUworkers (Appendix E).
Revised: During training, we used 128 CPU workers (Appendix E).

Original: The workers balance the data generation for different buffers to ensure an equal sample-to-insertratio across d , accounting for the oversampling of higher d as well as the longer computation timeneeded for generating data D of larger d , for instance, in the G RN domain.
Revised: The workers balance the data generation for different buffers to ensure an equal sample-to-insert ratio across d , accounting for the oversampling of higher d as well as the longer computation time needed for generating data D of larger d , for instance, in the G RN domain.

Original: In addition, the dataset
Revised: In addition, the dataset

Original: D of each element ( G, D ) in the buffer contains four times more observations than n = 200 usedduring training.
Revised: D of each element ( G, D ) in the buffer contains four times more observations than n = 200 used during training.

Original: These observations are subsampled to obtain n = 200 each time a given bufferelement ( G, D ) is drawn to introduce additional diversity in the training data in case buffer elementsare sampled more than once.
Revised: These observations are subsampled to obtain n = 200 each time a given buffer element ( G, D ) is drawn to introduce additional diversity in the training data in case buffer elements are sampled more than once.

Original: Parameter updates The primal updates of the inference model parameters ϕ are performed usingthe LAMB optimizer with a constant base learning rate 3 · 10 − 5 and adaptive square-root scaling bythe maximum effective batch size 3 (You et al., 2019).
Revised: The primal updates of the inference model parameters ϕ are performed using the LAMB optimizer with a constant base learning rate 3 · 10 − 5 and adaptive square-root scaling by the maximum effective batch size 3 (You et al., 2019).

Original: Gradients with respect to ϕ are clipped at aglobal ℓ 2 norm of one (Pascanu et al., 2013).
Revised: Gradients with respect to ϕ are clipped at a global ℓ 2 norm of one (Pascanu et al., 2013).

Original: In all three domains, we optimize ϕ for a total numberof 300 , 000 primal steps, reducing the learning rate by a factor of ten after 200 , 000 steps.
Revised: In all three domains, we optimize ϕ for a total number of 300 , 000 primal steps, reducing the learning rate by a factor of ten after 200 , 000 steps.

Original: When adding the acyclicity contraint in L INEAR and R FF , we use a dual learning rate of η = 10 − 4and perform a dual update every 500 primal steps.
Revised: When adding the acyclicity contraint in L INEAR and R FF , we use a dual learning rate of η = 10 −and perform a dual update every 500 primal steps.

Original: The dual learning rate η is warmed up with a linearschedule from zero over the first 50 , 000 primal steps.
Revised: The dual learning rate η is warmed up with a linear schedule from zero over the first 50 , 000 primal steps.

Original: To reduce the variance in the dual update, weuse an exponential moving average of F ( ϕ ) with step size 10 − 4 maintained during the updates of theprimal objective.
Revised: To reduce the variance in the dual update, we use an exponential moving average of F ( ϕ ) with step size 10 − 4 maintained during the updates of the primal objective.

Original: To approximate the spectral radius in Eq.
Revised: To approximate the spectral radius in Eq.

Original: (5), we perform t = 10 power iterationsinitialized at u , v ∼ N (0 , I d ) .
Revised: (5), we perform t = 10 power iterations initialized at u , v ∼ N (0 , I d ) .

Original: C.2 Architecture
Revised: C.2 Architecture

Original: As described in Section 4.2, the core of our model consists of L = 8 blocks of four residual layers.
Revised: As described in Section 4.2, the core of our model consists of L = 8 blocks of four residual layers.

Original: Different from the vanilla Transformer encoder, we employ layer normalization before each multihead attention and feedforward module and after the last of the L layers (Radford et al., 2019).
Revised: Different from the vanilla Transformer encoder, we employ layer normalization before each multihead attention and feedforward module and after the last of the L layers (Radford et al., 2019).

Original: Themulti-head attention modules have a model size of 128 , key size of 32 , and 8 attention heads.
Revised: The multi-head attention modules have a model size of 128 , key size of 32 , and 8 attention heads.

Original: Thefeedforward modules have a hidden size of 512 and use ReLU activations.
Revised: The feedforward modules have a hidden size of 512 and use ReLU activations.

Original: In held-out tasks of R FFand G RN , we found that dropout in the Transformer encoder does not hurt performance in-distribution,so we increased the dropout rates from 0 . 0 to 0 .
Revised: In held-out tasks of R FF and G RN , we found that dropout in the Transformer encoder does not hurt performance in-distribution, so we increased the dropout rates from 0 . 0 to 0 .

Original: 1 and 0 . 3 , respectively, to help generalization o.o.d.
Revised: 1 and 0 . 3 , respectively, to help generalization o.o.d.

Original: Dropout, when performed, is done before the residual layers are added, as in the vanilla Transformer(Vaswani et al., 2017).
Revised: Dropout, when performed, is done before the residual layers are added, as in the vanilla Transformer (Vaswani et al., 2017).

Original: The position-wise linear layers that map the two-dimensional representation ( z 1 , . .
Revised: The position-wise linear layers that map the two-dimensional representation ( z 1 , . .

Original: . , z d ) ∈ R d × k tou i and v i , respectively, apply layer normalization prior to their transformations.
Revised: . , z d ) ∈ R d × k to u i and v i , respectively, apply layer normalization prior to their transformations.

Original: We use Kaiminguniform initialization for the weights (He et al., 2015a).
Revised: We use Kaiming uniform initialization for the weights (He et al., 2015a).

Original: The bias term inside the logistic function of
Revised: The bias term inside the logistic function of

Original: Eq.
Revised: Eq.

Original: (4) is initialized at − 3 and learned alongside all other parameters ϕ .
Revised: (4) is initialized at − 3 and learned alongside all other parameters ϕ .

Original: Likewise, the scale parameterτ is learned but optimized in log space to ensure positivity, i.e., τ = exp( τ log ) where τ log is updatedas part of ϕ and initialized at 2 .
Revised: Likewise, the scale parameter τ is learned but optimized in log space to ensure positivity, i.e., τ = exp( τ log ) where τ log is updated as part of ϕ and initialized at 2 .

Original: When optimizing models under the acyclicity constraint, we ignorethe diagonal predictions θ ii and mask the corresponding loss terms.
Revised: When optimizing models under the acyclicity constraint, we ignore the diagonal predictions θ ii and mask the corresponding loss terms.

Original: We implement AVICI with Haiku in JAX (Hennigan et al., 2020; Bradbury et al., 2018).
Revised: We implement AVICI with Haiku in JAX (Hennigan et al., 2020; Bradbury et al., 2018).

Original: Weconverged to the above optimization and architecture specifications through experimentation onheld-out instances from the training distributions p ( D ) , i.e., in-distribution.
Revised: We converged to the above optimization and architecture specifications through experimentation on held-out instances from the training distributions p ( D ) , i.e., in-distribution.

Original: D Baselines
Revised: D Baselines

Original: Algorithms and hyperparameter tuning We calibrate important hyperparameters for all methodson held-out problem instances from the test data distributions ˜ p ( D ) of L INEAR , R FF and G RN ,individually in each domain.
Revised: Algorithms and hyperparameter tuning We calibrate important hyperparameters for all methods on held-out problem instances from the test data distributions ˜ p ( D ) of L INEAR , R FF and G RN , individually in each domain.

Original: For the following algorithms, we search over the parameters relevant forcontrolling sparsity and the complexity of variable interactions:
Revised: For the following algorithms, we search over the parameters relevant for controlling sparsity and the complexity of variable interactions:

Original: • DCDI
Revised: • DCDI

Original: (Brouillard et al., 2020): sparsity regularizer λ ∈ { 10 − 2 , 10 − 1 , 1 } , size of hidden layerin MLPs modeling the conditional distributions ∈ { 8 , 32 }• DAG-GNN (Yu et al., 2019): graph thresholding parameter ∈ { 0 . 1 , 0 . 2 , 0 .
Revised: (Brouillard et al., 2020): sparsity regularizer λ ∈ { 10 − 2 , 10 − 1 , 1 } , size of hidden layer in MLPs modeling the conditional distributions ∈ { 8 , 32 } • DAG-GNN (Yu et al., 2019): graph thresholding parameter ∈ { 0 . 1 , 0 . 2 , 0 .

Original: 3 } , size of hiddenlayer in MLP encoder and decoder ∈ { 16 , [-
 }•-] DiBS (Lorch et al., 2021): latent kernel length scale ℓ z ∈ { 3 , 10 , 30 } ,.. .with BGe marginal likelihood (L INEAR ): effective sample size (sparsity) α BGe µ ∈ { 0 . 1 , 1 . 0 }...with nonlinear Gaussian likelihood (R FF , G RN ):
Revised: 3 } , size of hidden layer in MLP encoder and decoder ∈ { 16 , [-
 }•-]64 } • DiBS (Lorch et al., 2021): latent kernel length scale ℓ z ∈ { 3 , 10 , 30 } , .. .with BGe marginal likelihood (L INEAR ): effective sample size (sparsity) α BGe µ ∈ { 0 . 1 , 1 . 0 } ...with nonlinear Gaussian likelihood (R FF , G RN ):

Original: parameter length scale ℓ θ ∈ { 30 , 300 , 3000 }• GraN-DAG (Lachapelle et al., 2020): preliminary neighborhood selection threshold ∈ { 0 . 5 , 2 } ,size of hidden layer ∈ { 8 , 32 } , pruning cutoff ∈ { 10 − 3 , 10 − 5 }• IGSP (Wang et al., 2017): significance α ∈ { 10 − 2 , 10 − 3 , 10 − 4 } , CI test ∈ { Gaussian, HSIC- γ }• PC (Spirtes et al., 2000): significance α ∈ { 10 −  , 10 − 3 , 10 − 4 } , CI test ∈ { Gaussian, HSIC- γ } 
Revised: parameter length scale ℓ θ ∈ { 30 , 300 , 3000 } • GraN-DAG (Lachapelle et al., 2020): preliminary neighborhood selection threshold ∈ { 0 . 5 , 2 } , size of hidden layer ∈ { 8 , 32 } , pruning cutoff ∈ { 10 − 3 , 10 − 5 } • IGSP (Wang et al., 2017): significance α ∈ { 10 − 2 , 10 − 3 , 10 − 4 } , CI test ∈ { Gaussian, HSIC- γ } • PC (Spirtes et al., 2000): significance α ∈ { 10 − 2 , 10 − 3 , 10 − 4 } , CI test ∈ { Gaussian, HSIC- γ } DAG-GNN, DCDI, DiBS

Original: DAG-GNN, DCDI, DiBS, and GraN-DAG use 80% of the available data to perform inferenceand compute held-out log likelihood or ELBO scores on the other 20% of the data.
Revised: , and GraN-DAG use 80% of the available data to perform inference and compute held-out log likelihood or ELBO scores on the other 20% of the data.

Original: The besthyperparameters are then selected by averaging the metric over five held-out instances of d = 30variables.
Revised: The best hyperparameters are then selected by averaging the metric over five held-out instances of d =variables.

Original: DiBS draws 10 samples from p ( G | D ) using the interventional BGe score for L INEAR anda nonlinear Gaussian interventional likelihood with MLP means for R FF and G RN .
Revised: DiBS draws 10 samples from p ( G | D ) using the interventional BGe score for L INEAR and a nonlinear Gaussian interventional likelihood with MLP means for R FF and G RN .

Original: DiBS assumes anobservation noise of 1 , uses a scale-free graph prior, and anneals acyclicty and relaxation parameterswith rate 1 .
Revised: DiBS assumes an observation noise of 1 , uses a scale-free graph prior, and anneals acyclicty and relaxation parameters with rate 1 .

Original: All remaining parameters are kept at the settings suggested by the authors.
Revised: All remaining parameters are kept at the settings suggested by the authors.

Original: For the PC algorithm and IGSP, there is no held-out score, so we compute the SID and F1 scoresusing the ground-truth causal graphs to select their optimal parameters.
Revised: For the PC algorithm and IGSP, there is no held-out score, so we compute the SID and F1 scores using the ground-truth causal graphs to select their optimal parameters.

Original: This would not be possiblein practice and thus favors these methods.
Revised: This would not be possible in practice and thus favors these methods.

Original: The HSIC- γ CI test did not scale to d = 100 variables,so in these cases PC and IGSP always use the Gaussian CI test.
Revised: The HSIC- γ CI test did not scale to d = 100 variables, so in these cases PC and IGSP always use the Gaussian CI test.

Original: For G RN d = 30 , IGSP also usesthe Gaussian CI test because it OOMs at 100GB when using HSIC- γ .
Revised: For G RN d = 30 , IGSP also uses the Gaussian CI test because it OOMs at 100GB when using HSIC- γ .

Original: GES and GIES use the linear
Revised: GES and GIES use the linear

Original: Gaussian BIC score function and thus do not require calibrating a sparsity parameter (Chickering,2003; Hauser and Bühlmann, 2012).
Revised: Gaussian BIC score function and thus do not require calibrating a sparsity parameter (Chickering, 2003; Hauser and Bühlmann, 2012).

Original: LiNGAM is based on independent component analysis andrequires no regularization tuning either (Shimizu et al., 2006).
Revised: LiNGAM is based on independent component analysis and requires no regularization tuning either (Shimizu et al., 2006).

Original: DAG bootstrap To estimate edge probabilities for the non-Bayesian methods in Section 6.2, weuse the nonparametric DAG bootstrap (Friedman et al., 1999).
Revised: DAG bootstrap To estimate edge probabilities for the non-Bayesian methods in Section 6.2, we use the nonparametric DAG bootstrap (Friedman et al., 1999).

Original: We bootstrap ten datasets D ′ from Dby sampling with replacement and then run each baseline individually on each bootstrapped dataset
Revised: We bootstrap ten datasets D ′ from D by sampling with replacement and then run each baseline individually on each bootstrapped dataset

Original: D ′ .
Revised: D ′ .

Original: The nonparametric probability estimate for an edge then amounts to the proportion of predictedgraphs G ′ that contain the edge.
Revised: The nonparametric probability estimate for an edge then amounts to the proportion of predicted graphs G ′ that contain the edge.

Original: Implementation For GES, GIES, PC, and LiNGAM, we run the original R implementations ofthe authors using an extended version of the software by Kalainathan et al.
Revised: Implementation For GES, GIES, PC, and LiNGAM, we run the original R implementations of the authors using an extended version of the software by Kalainathan et al.

Original: (2020) (MIT Licence).
Revised: (2020) (MIT Licence).

Original: For DCDI, DAG-GNN, GraN-DAG, and DiBS, we use the Python implementations provided by theauthors (Brouillard et al., 2020; Yu et al., 2019; Lachapelle et al., 2020; Lorch et al., 2021)
Revised: For DCDI, DAG-GNN, GraN-DAG, and DiBS, we use the Python implementations provided by the authors (Brouillard et al., 2020; Yu et al., 2019; Lachapelle et al., 2020; Lorch et al., 2021)

Original: (MIT
Revised: (MIT

Original: License, Apache License 2.0, MIT License, MIT Licence).
Revised: License, Apache License 2.0, MIT License, MIT Licence).

Original: For IGSP, we use the implementationprovided as part of the CausalDAG package (Squires et al., 2018)
Revised: For IGSP, we use the implementation provided as part of the CausalDAG package (Squires et al., 2018)

Original: (3-Clause BSD license).
Revised: (3-Clause BSD license).

Original: LiNGAM relies on the inversion of a covariance matrix, which frequently fails in the G RN domaindue to the high sparsity in D .
Revised: LiNGAM relies on the inversion of a covariance matrix, which frequently fails in the G RN domain due to the high sparsity in D .

Original: Thus, to benchmark LiNGAM in G RN , we add small Gaussian noise tothe standardized count matrix D .
Revised: Thus, to benchmark LiNGAM in G RN , we add small Gaussian noise to the standardized count matrix D .

Original: For the IGSP and PC algorithms, the same numerical adjustmentis needed to avoid crashes in the CI tests on G RN .
Revised: For the IGSP and PC algorithms, the same numerical adjustment is needed to avoid crashes in the CI tests on G RN .

Original: Single IGSP runs that still failed for d = 100were ignored when computing the metrics.
Revised: Single IGSP runs that still failed for d =were ignored when computing the metrics.

Original: In the G RN results, we ignored a small number of single
Revised: In the G RN results, we ignored a small number of single

Original: [Table]
Revised: [Table]

Original: runs of DCDI for d = 100 and PC for d = 30 that failed to terminate after 24 hours walltime (on a
Revised: runs of DCDI for d = 100 and PC for d = 30 that failed to terminate after 24 hours walltime (on a

Original: GPU machine for the former).
Revised: GPU machine for the former).

Original: Lastly, the CAM pruning post-processing procedure of the author’simplementation of GraN-DAG (Lachapelle et al., 2020) crashes in a few instances.
Revised: Lastly, the CAM pruning post-processing procedure of the author’s implementation of GraN-DAG (Lachapelle et al., 2020) crashes in a few instances.

Original: We skip thepost-processing step in these cases.
Revised: We skip the post-processing step in these cases.

Original: E Extended Results
Revised: E Extended Results

Original: Compute Resources To carry out the experiments in this work, we trained three main AVICImodels and several ablations.
Revised: Compute Resources To carry out the experiments in this work, we trained three main AVICI models and several ablations.

Original: Each model was optimized for approximately four days using 8 Quadro
Revised: Each model was optimized for approximately four days using 8 Quadro

Original: RTX 6000 GPUs ( 24 GiB memory each) and 128 CPUs.
Revised: RTX 6000 GPUs ( 24 GiB memory each) and 128 CPUs.

Original: To perform the benchmarking experiments,all baselines were run on four to eight CPUs each for up to 24 hours, depending on the method.
Revised: To perform the benchmarking experiments, all baselines were run on four to eight CPUs each for up to 24 hours, depending on the method.

Original: DCDI required one GPU to ensure a computation time of less than one day per task instance.
Revised: DCDI required one GPU to ensure a computation time of less than one day per task instance.

Original: In allexperiments, test-time inference with AVICI is done on eight CPUs and no GPU.
Revised: In all experiments, test-time inference with AVICI is done on eight CPUs and no GPU.

Original: E.1 In-Distribution Benchmarking Results for d =
Revised: E.1 In-Distribution Benchmarking Results for d =

Original: Table 4 gives the benchmarking results for in-distribution data of d = 30 variables given the otherwiseunchanged setup of Section 5.
Revised: Table 4 gives the benchmarking results for in-distribution data of d = 30 variables given the otherwise unchanged setup of Section 5.

Original: Contrary to the o.o.d.
Revised: Contrary to the o.o.d.

Original: setting, the data is generated under homogeneous,additive noise and the parameters of their generative processes are sampled from the training domainsof AVICI (cf.
Revised: setting, the data is generated under homogeneous, additive noise and the parameters of their generative processes are sampled from the training domains of AVICI (cf.

Original: Table 2).
Revised: Table 2).

Original: However, as throughout all experiments, the datasets and its data-generatingparameters themselves are unique and have not been used by AVICI during training.
Revised: However, as throughout all experiments, the datasets and its data-generating parameters themselves are unique and have not been used by AVICI during training.

Original: Compared to the o.o.d.
Revised: Compared to the o.o.d.

Original: setting, most baselines perform roughly the same.
Revised: setting, most baselines perform roughly the same.

Original: Since the data-generatingprocesses are sampled from its training distribution, AVICI significantly improves by moving to theeasier in-distribution setting, in particular in the SCM domains, which are less noisy.
Revised: Since the data-generating processes are sampled from its training distribution, AVICI significantly improves by moving to the easier in-distribution setting, in particular in the SCM domains, which are less noisy.

Original: In the G RNdomain, some baselines achieve slightly better F1 scores compared to the o.o.d.
Revised: In the G RN domain, some baselines achieve slightly better F1 scores compared to the o.o.d.

Original: setting, which is mostlikely explained by a change in the graph rather than the simulator parameter distribution, since thereis no reason to believe that different generative parameters are more challenging to the baselines.
Revised: setting, which is most likely explained by a change in the graph rather than the simulator parameter distribution, since there is no reason to believe that different generative parameters are more challenging to the baselines.

Original: E.2 Benchmarking Results for d =100
Revised: E.2 Benchmarking Results for d =

Original: Table 5 shows the benchmarking results for d = 100 variables given n = 1000 observations and theexperimental setup of Section 6.2.
Revised: Table 5 shows the benchmarking results for d = 100 variables given n = 1000 observations and the experimental setup of Section 6.2.

Original: We highlight that in this evaluation regime, AVICI operates underdistribution shift in terms of the causal structures, mechanisms or simulator parameters, and noisedistributions, as well as the number of variables and the number of observations seen during training.
Revised: We highlight that in this evaluation regime, AVICI operates under distribution shift in terms of the causal structures, mechanisms or simulator parameters, and noise distributions, as well as the number of variables and the number of observations seen during training.

Original: Overall, the qualitative ranking of the methods is very similar as for d = 30 .
Revised: Overall, the qualitative ranking of the methods is very similar as for d = 30 .

Original: AVICI outperforms allbaselines in the nonlinear R FF domain, with and without access to interventional data.
Revised: AVICI outperforms all baselines in the nonlinear R FF domain, with and without access to interventional data.

Original: Likewise,
Revised: Likewise,

Original: AVICI is the only method to achieve nontrivial edge accuracy in terms of F1 score on the challenging
Revised: AVICI is the only method to achieve nontrivial edge accuracy in terms of F1 score on the challenging

Original: [Table]
Revised: [Table]

Original: G RN domain.
Revised: G RN domain.

Original: On the simpler L INEAR domain, there is no statistically significant difference between
Revised: On the simpler L INEAR domain, there is no statistically significant difference between

Original: GES/GIES and AVICI, which perform overall most favorably.
Revised: GES/GIES and AVICI, which perform overall most favorably.

Original: E.3 Benchmarking Results on Real-World Proteomics Data
Revised: E.3 Benchmarking Results on Real-World Proteomics Data

Original: We additionaly evaluate all of the methods on the real-world dataset by Sachs et al.
Revised: We additionaly evaluate all of the methods on the real-world dataset by Sachs et al.

Original: (2005), whichcontains continuous measurements of d = 11 proteins involved in human immune system cells.
Revised: (2005), which contains continuous measurements of d = 11 proteins involved in human immune system cells.

Original: Structure learning algorithms are commonly compared on this dataset, and for completeness, wereport the performance of AVICI and the baselines here.
Revised: Structure learning algorithms are commonly compared on this dataset, and for completeness, we report the performance of AVICI and the baselines here.

Original: However, the ground-truth network of 17edges put forward by Sachs et al.
Revised: However, the ground-truth network ofedges put forward by Sachs et al.

Original: (2005) has been challenged by some experts (Mooij et al., 2020)and the assumptions of causal sufficiency and acyclicity may not be justified even though assumed bymost methods, which should be kept in mind when interpreting the results.
Revised: (2005) has been challenged by some experts (Mooij et al., 2020) and the assumptions of causal sufficiency and acyclicity may not be justified even though assumed by most methods, which should be kept in mind when interpreting the results.

Original: A large part of the dataare interventional, in which the measured proteins were activated or inhibited using specific reagents.
Revised: A large part of the data are interventional, in which the measured proteins were activated or inhibited using specific reagents.

Original: Most interventions are likely not perfect and the intervention targets may not be completely accurate(Mooij et al., 2020).
Revised: Most interventions are likely not perfect and the intervention targets may not be completely accurate (Mooij et al., 2020).

Original: For this experiment, we follow Wang et al.
Revised: For this experiment, we follow Wang et al.

Original: (2017) and Brouillard et al.
Revised: (2017) and Brouillard et al.

Original: and discard data inwhich interventions were not targeted directly at one of d = 11 measured proteins.
Revised: and discard data in which interventions were not targeted directly at one of d = 11 measured proteins.

Original: Given this setup,we have n = 5846 data points that contain 1755 observational and 4091 interventional measurements,which consist of five single-protein perturbations.
Revised: Given this setup, we have n = 5846 data points that contain 1755 observational and 4091 interventional measurements, which consist of five single-protein perturbations.

Original: In our results, methods that only use observationaldata take the concatenation of all of the data without the intervention target information as input.
Revised: In our results, methods that only use observational data take the concatenation of all of the data without the intervention target information as input.

Original: Allbaselines use the hyperparameters tuned for the nonlinear R FF domain.
Revised: All baselines use the hyperparameters tuned for the nonlinear R FF domain.

Original: The data is standardized tohave mean 0 and variance 1 .
Revised: The data is standardized to have mean 0 and variance 1 .

Original: Table 6 summarizes the results of all methods with respect to the reference causal graph.
Revised: Table 6 summarizes the results of all methods with respect to the reference causal graph.

Original: Figure 3visualizes the prediction of each method.
Revised: Figurevisualizes the prediction of each method.

Original: Overall, the results are not very conclusive.
Revised: Overall, the results are not very conclusive.

Original: GES and GIESperform best in terms of SID, GraN-DAG is most favorable in terms of F1, and together with DCDIand AVICI also in terms of SHD.
Revised: GES and GIES perform best in terms of SID, GraN-DAG is most favorable in terms of F1, and together with DCDI and AVICI also in terms of SHD.

Original: More generally, the number of predicted edges varies greatly acrossmethods.
Revised: More generally, the number of predicted edges varies greatly across methods.

Original: Almost all F1 scores fall between 0 .
Revised: Almost all F1 scores fall between 0 .

Original: 25 and 0 . 30 .
Revised: 25 and 0 . 30 .

Original: E.4 Uncertainty quantification for d =30
Revised: E.4 Uncertainty quantification for d =

Original: Table 7 summarizes the probabilistic AUROC and AUPRC metrics for all methods.
Revised: Table 7 summarizes the probabilistic AUROC and AUPRC metrics for all methods.

Original: Explanationsand interpretations for both metrics are given in Section
Revised: Explanations and interpretations for both metrics are given in Section

Original: B. The relative performance of the bootstrapbaselines and AVICI is similar to the point estimate benchmark.
Revised: B. The relative performance of the bootstrap baselines and AVICI is similar to the point estimate benchmark.

Original: Overall, AVICI performs favorablyacross the three domains, with GES and GIES on par in L INEAR .
Revised: Overall, AVICI performs favorably across the three domains, with GES and GIES on par in L INEAR .

Original: However, since AUROC and
Revised: However, since AUROC and

Original: AUPRC metrics evaluate the full spectrum of decision thresholds, we additionally see that AVICIachieves nontrivial accuracy in G RN even without access to gene knockout data, indicating that
Revised: AUPRC metrics evaluate the full spectrum of decision thresholds, we additionally see that AVICI achieves nontrivial accuracy in G RN even without access to gene knockout data, indicating that

Original: AVICI may provide useful information even in settings with passive observations.
Revised: AVICI may provide useful information even in settings with passive observations.

Original: This aspect is notapparent when converting the posterior probability estimates of AVICI based on a single thresholdand then comparing SID and F1 scores as in Table 1.
Revised: This aspect is not apparent when converting the posterior probability estimates of AVICI based on a single threshold and then comparing SID and F1 scores as in Table 1.

Original: [Table]
Revised: [Table]

Original: E.5 AVICI generalization between L INEAR and RFF
Revised: E.5 AVICI generalization between L INEAR and R

Revised: In

Original: In this section, we provide additional out-of-distribution generalization results for AVICI.
Revised: this section, we provide additional out-of-distribution generalization results for AVICI.

Original: Specifically,we test the AVICI model trained on the L INEAR domain on inference from R FF data, and vice versa.
Revised: Specifically, we test the AVICI model trained on the L INEAR domain on inference from R FF data, and vice versa.

Original: This means that the AVICI models not only operate under distributional shifts on the parameters oftheir respective data-generating processes, but also on the function classes of causal mechanismsthemselves.
Revised: This means that the AVICI models not only operate under distributional shifts on the parameters of their respective data-generating processes, but also on the function classes of causal mechanisms themselves.

Original: The models infer causal structure from data generated from function classes neverseen during training.
Revised: The models infer causal structure from data generated from function classes never seen during training.

Original: As in all empirical analyses of Section 6, the graph and noise parameters areadditionally o.o.d., that is, the L INEAR AVICI model is tested on the o.o.d.
Revised: As in all empirical analyses of Section 6, the graph and noise parameters are additionally o.o.d. , that is, the L INEAR AVICI model is tested on the o.o.d.

Original: R FF data, and vice versa.
Revised: R FF data, and vice versa.

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

Original: Even under this distributional shift, the performance of both AVICImodels decreases reasonably and remains on par with most baselines (Table 1).
Revised: Even under this distributional shift, the performance of both AVICI models decreases reasonably and remains on par with most baselines (Table 1).

Original: On L INEAR data,the baselines achieve F1 scores of 0 .
Revised: On L INEAR data, the baselines achieve F1 scores of 0 .

Original: 54 with observational and 0 .
Revised: 54 with observational and 0 .

Original: 33 - 0 . 74 with interventionaldata, similar to the R FF AVICI model with 0 .
Revised: 74 with interventional data, similar to the R FF AVICI model with 0 .

Original: 19 and 0 . 45 , respectively.
Revised: 19 and 0 . 45 , respectively.

Original: Conversely, on R FF data, thebaselines achieve F1 scores of 0 .
Revised: Conversely, on R FF data, the baselines achieve F1 scores of 0 .

Original: 42 with observational and 0 .
Revised: 42 with observational and 0 .

Original: 41 with interventional data,which is also matched by the L INEAR AVICI model here with 0 .
Revised: 41 with interventional data, which is also matched by the L INEAR AVICI model here with 0 .

Original: 27 and 0 .
Revised: 27 and 0 .

Original: 42 , respectively.
Revised: 42 , respectively.

Original: Overall,the L INEAR AVICI model generalizes marginally better to R FF data as vice versa.
Revised: Overall, the L INEAR AVICI model generalizes marginally better to R FF data as vice versa.

Original: We do not reportthe SID here because the R code of Peters and Bühlmann (2015) runs out of memory.
Revised: We do not report the SID here because the R code of Peters and Bühlmann (2015) runs out of memory.

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

Original: F
Revised: F

Original: [Table]
Revised: [Table]