Amortized Active Causal Induction with Deep Reinforcement Learning
Authors: Yashas Annadani
Panagiotis Tigas
Stefan Bauer
Adam Foster
Venue: NeurIPS 2024 poster (2024)
Keywords: Active Causal Structure Learning
Adaptive Intervention Design
Reinforcement Learning
Primary area: active_learning
TL;DR: We propose an amortized and adaptive intervention strategy that results in a sample efficient estimate of the true causal graph on the distribution of training environment as well as on test-time environments with distribution shifts.

Abstract: We present Causal Amortized Active Structure Learning (CAASL), an active intervention design policy that can select interventions that are adaptive, real-time and that does not require access to the likelihood. This policy, an amortized network based on the transformer, is trained with reinforcement learning on a simulator of the design environment, and a reward function that measures how close the true causal graph is to a causal graph posterior inferred from the gathered data. On synthetic data and a single-cell gene expression simulator, we demonstrate empirically that the data acquired through our policy results in a better estimate of the underlying causal graph than alternative strategies.  Our design policy successfully achieves amortized intervention design on the distribution of the training environment while also generalizing well to distribution shifts in test-time design environments. Further, our policy also demonstrates excellent zero-shot generalization to design environments with dimensionality higher than that during training, and to intervention types that it has not been trained on.

Amortized Active Causal Induction with Deep Reinforcement Learning
Yashas Annadani
Panagiotis Tigas Stefan Bauer Adam Foster
Helmmoltz AI, Munich
Technical University of Munich
3 OATML, University of Oxford
Abstract
We present Causal Amortized Active Structure Learning (CAASL), an active intervention design policy that can select interventions that are adaptive, real-time and that does not require access to the likelihood. This policy, an amortized network based the transformer, is trained with reinforcement learning on a simulator of the design environment, and reward function that measures how close the true causal graph is to causal graph posterior inferred from the gathered data. On synthetic data and a single-cell gene expression simulator, we demonstrate empirically that the data acquired through our policy results in better estimate of the underlying causal graph than alternative strategies. Our design policy successfully achieves amortized intervention design the distribution the training environment while also generalizing well to distribution shifts test-time design environments. Further, our policy also demonstrates excellent zero-shot generalization to design environments with dimensionality higher than that during training, and intervention types that has not been trained on.
Introduction
world
Infer, design and experiment is three step loop in the empirical scientific discovery paradigm. Causal induction (a.k.a. causal structure learning), the problem of finding causal relationships present in data, also falls under this paradigm when experiments in the form of interventions are permissi- ble [52 28). Causal structure learning has gained increasing importance in empirical sciences, for example in single-cell biology, where perturbation experiments like gene knockouts can be carried out with high-precision (54). Such interventions are not only more informative to infer the underlying causal graph than just observational data, but in certain cases essential to beyond the Markov equivalence class 43), making the problem of design of interventions both relevant and important. For the problem of structure learning with interventions, how- ever, inference and design both involve significant challenges. For instance, inference of the causal graph from data usually involves search over the space of graphs with likelihood (usu- ally weighted by prior) score function [4 9 (2), which is slow and not robust to violations of data generation assump- the
experiments
history inseventonn deservaional dana
poloy
transformet hilooyy anoodder
Figure Causal Amortized Struc- ture Learning (CAASL) is an active intervention design method that di- rectly proposes the next intervention to perform by just forward-pan pass of the transformer based policy.
tions [41) The design of informative interventions, other hand, utilizes the inferred causal graph from existing data to select promising designs and rank them according scoring criterion. This scoring criterion is usually based an approximation of mutual information between the unknown causal graph and the interventional data [56] [53], which
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

also involves the (interventional) data likelihood. In problems related to empirical sciences where causal structure learning is essential, like inferring gene regulatory network with gene knockouts or knockdowns, the likelihood of the data is typically intractable. While progress has been made in terms of fikelihood-rree inference of causal graphs [38 32], existing intervention design algorithms have been largely restricted to likelihood- based strategies.
With a focus on addressing practical intervention design challenges that arise in empirical sciences like inferring the gene regulatory network, in this work, we propose an intervention design method called CAASL that significantly differs from existing approaches. Instead of following the infer, design and experiment loop, we amortize the intervention design procedure by training single design network policy, based on the transformer [57]. which encodes key design space symmetries. During test-time, our trained policy directly predicts the next intervention to perform by just a forward -pass of the data collected so far, without the need to undergo slow and expensive inference of the causal graph corresponding to that data. We train the transformer policy with Soft Actor-Critic (SAC) (26 to maximize cumulative rewards over a fixed number of design iterations (budget), thereby making the policy adaptive. The choice of a good reward function is essential for informative designs. We discuss various reward function choices, primarily based on an estimate of the true causal graph obtained from a likelihood-frre amortized causal structure learning approach. Both our policy and the reward function only require access to a simulator of the design environment. Further, we present connections of our approach to amortized sequential Bayesian experimental design [22]. We demonstrate that the reward function is related to an approximation of expected information gain based on the amortized posterior distribution over causal graphs. As such, CAASL is an intervention design method for performing sample efficient causal structure learning, but is not a new causal structure learning method in itself.
On synthetic data and the single-cell gene expression simulator SERGIO [7], we empirically study various aspects of our trained policy-the amortization performance on training distribution of the design environment as well as on design environments with distribution shifts from the training environment. We find that our policy obtains better causal structure learning performance for given budget than alternate intervention strategies. Overall, we observe excellent generalization capability of the transformer for intervention design, similar to what has been demonstrated in other domains [TO, BT [88. The robustness of the amortized policy opens up the possibility for lab-in- the-loop intervention design for single-cell data, wherein single network can propose informative interventions across different cell lines and experimental conditions.
Our contributions are:
We propose an amortized active intervention design method for causal structure learning based on transformer parameterized policy that encodes key design space symmetries.
Based the AVICI [3X amortized causal structure inference model, we propose a reward
function for training the policy with reinforcement learning that does not require access to the likelihood.
We demonstrate the superiority of CAASL by performing extensive evaluation various in-distribution and out-of-distributio settings in the synthetic and SERGIO [17 gene regulatory network simulator environments.
2 Background and Related Work
Structural Causal Models. Let y {. : be the random variables of interest associated with the vertices of a graph G. Let A E 1)dxd be the adjacency matrix corresponding to G. A Structural Causal Model (SCM) [3 is framework for causality which consists of a set of equations in which each variable y, is deterministic function of its direct causes Mppg(() as well as an exogenous noise variable e; with a distribution Pc,
y: f (3mm
(I)
The functions fi, with parameters 0,, are mechanisms that relate how the direct causes affect the variable y;- The structural assignments are typically assumed to be acyelic, with G being directed acyelic graph whose edges indicate direct causes. In addition, an SCM defines the likelihood of any data sample y under this model, denoted as p(y {A, 0)). Further, we assume that the SCM is

---

causally sufficient, i.c. all the variables are measurable (but can be missing at random), and the noise variables are mutually independent.
Interventions. The SCM framework admits reasoning about effects of interventions on any variable in y. Most notable types of intervention include perfect (do) intervention, and shift interven- any variable y; corresponds to changing the structural equation of that variable to the desired value, y/i == 0t. is denoted by the do-operator [42] as do = v.). In a shift intervention, the conditional mean of the interventional variable Blys (ra is shifted by us
tion [48]. A perfect intervention
The likelihood of any data under an intervention 7 is denoted as p(y {A, A, 7). For perfect and
shift interventions, I can be parameterized as a d x 2 dimensional matrix, where the first column corresponds to one-hot encoding of whether a particular variable is intervened or not, and the second column corresponds to the value (or the shift) of the intervention corresponding to each potential intervention target.
Causal Structure Learning. The problem causal structure learning corresponds to estimating A (and other parameters of the SCM 0) given samples from Pdata [28). In general, there could be multiple models (and hence graphs) that can be consistent with a given joint distribution over y. which necessitates causal structure learning with interventional data [43] There are various approaches, either based on independence tests [1]] [52,, or graph search by maximizing a score function (likelihood of the data with certain assumptions on the SCM) 9 (7) Reinforcement learning has also been used for search over graphs with a score function [60], however it differs entirely from our approach wherein tractable likelihood, there are also causal structure learning methods that estimate the posterior distribution 4(A D) of graphs [4 [6 for a dataset D that sampled from Pdata.
RL for intervention design. Alternately, based on the
Likelihood-Free Amortized Causal Structure Learning (AVICI) (38]. More recently, instead of inferring causal graphs over specific datasets, amortized posterior inference of causal graphs has
also been studied [38 32). In particular, the amortized posterior from Lorch et al. (38]. called AVICI, makes use of a transformer to directly predict the posterior (A D) by just forward- -asss of any dataset. The amortized posterior, parameterized as a product of independent Bernoulli random variables over the presence of edges in the causal graph, is trained from simulator without having access to the likelihood of the data. Since the simulator provides the ground truth value of A, the amortized posterior can be trained by maximum likelihood of graph edges with a combination of observational and interventional data. Empirically it has been shown that the AVICI model can amortize over datasets with different dimensionalities d, while also generalizing to new datasets that have not been seen during training. Since it is computationally cheap obtain the posterior distribution with AVICI, we use it for computing the reward for our intervention design policy.
Active Intervention Design. Active intervention design is the problem of designing interventional experiments to obtain data that enables causal structure learning in a sample efficient manner (under a fixed budget). While adaptive strategies have been explored [4 25). these approaches still require intermediate inference of the SCM, and are also not amortized. Intervention design based on Bayesian optimal experimental design B 12 has also been considered, although only with additive noise models, which enable likelihood evaluation [55] 56, I 59 (33]. Reinforcement learning has also been used in intervention design 49 35) however, they have been limited to non-amortized small scale settings. (50) highlight the usefulness of intervention design in an amortized framework for causal discovery. In contrast earlier work, demonstrate the applicability of our method to single-cell simulated gene expression data, wherein the mechanisms are defined by differential equations and also include technical noise (section 5.2).
3 Amortized Intervention Design
We first present our active intervention design strategy with reinforcement learning, the corresponding amortized network and its training. In Section 4.1 we then present connections of our reward to sequential Bayesian experimental design.
Setting. Given a budget T, intervention design is the problem of finding a sequence of informative interventions with policy 41.... Ir ~ that results in an estimate of the causal graph that is
3

---

close to A. For any intervention I, a causal model defines a generative model of the data with likelihood p [A.0).1 1) and prior p(A, 9). We indicate initial (observational) data, available, as yo {vo y and the corresponding interventions with Io. where Iu {0)-0 if the initial data is fully observational. Let h4 E [R(no+t)xxx2 denote the interventional history (yo, (y,, Tt), obtained by concatenation of y and first column of that correspond to interventional targets. We do not explicitly encode intervention values in history, since for do intervention, the intervention values are already present in 3 Existing intervention design strategies like [56] approximate a posterior on {A, 0} at each step f, approximate expected information gain (EIG) [37 45 and greedily maximize it to compute I41. Details of this greedy approach are given in Appendix A.1
3.1 Intervention Design with Reinforcement Learning
In this work, we instead treat intervention design as a Reinforcement Learning (RL) problem and train a single policy network F with parameters @ to obtain a sequence of adaptive interventions I,,..., Ir for any underlying causal graph with adjacency matrix A. In order to do so, first describe the RL environment under which the interventions are performed.
Intervention Design Environment. Similar to Blau et al. 77, we define interventional de- sign environment as a Hidden-Parameter Markov Decision Process (HiP-MDPP (18]. The HiP- MDP use, ((A,00) ()), has hidden parameters (A,0) and can be fully described by the tu- ple (S.A, p, 3, T, R1 7, Pa). The state-space S consists of the histories p ho = yoo Ia) corresponds to initial data, the action-space A corresponds interventions a₂ 4e and / describes the space of all causal models (graphs and parameters) with prior Pa = p(A, 0). The hidden parameters are sampled for each episode at the beginning from the prior. y is the discount factor. In a HiP-MDP, the transition function T and reward R depend on the hidden parameters. The transition function T(h, ht-1, Ie, {A, 0)) is Markovian, and involves two operations: (1) sampling interventional data y/: p(y [A, 0), Tt), and (2) updating the history state h = Concat [ 1, (ye, T.)]. For a reward function R(he, Je. 1 (,, that we define below, intervention design corresponds to finding the parameters @ of the amortized policy that maximizes the expected cumulative reward of all interventions:
he, the initial state
max
E 0. T.P(A.)
Σ 1-2 Rhe, Te, h1-1, (A, 0})
with Ie ~ 1 (h+-)-
Reward Function. For the purpose of amortized intervention design, a good reward function should be cheap to evaluate while leading to informative interventions. In this work, we propose to utilize the estimate of the causal graph from an amortized causal graph posterior {(A hx). In particular, we use the pretrained AVICI model 38 AVICI a transformer based neural network trained with (interventional) data from simulator to directly predict the probability of presence or absence of any edge in the causal graph by just a forward pass of the data, without requiring access the likelihood. For any history 11,1, we define the reward for performing intervention L and reaching state he as the improvement in the number of correct entries in the predicted adjacency matrix of the AVICI model:
E1(ANs AAno]
R(he, h, (A,0))
E o(A(b))
R(he-1146-1, he-2 (A,
(3)
where I[:] is the indicator function and R(ho, Lo. {A, 0)) = E By(A)l)) EEn [Ars - We note that our choice of reward function revolves around obtaining good estimate for the causal graph, A; we do not (directly) reward learning about 0.
The above RL problem for intervention design is intuitive: reward the intervention in proportion to the improvement it brings in terms of number of correct entries of the adjacency matrix from the amortized posterior. Also, for any t, the cumulative reward, eq. (2), for y = 1 of all interventions including To. which includes an additional term R(ho, Io, {A, B)) ), is simply the number of correct
We train our policy only on do interventions.

---

Variables
CAASI Aa
1
Pu 1) Samples
Attn samples
varibbles
Causal Intervention Simulator
Y
d
Pretrained IVICI Model
<(1 Aa)
Reward
ng +
History he
Figure 2: Schematic diagram illustrating the proposed CAASL policy along with the AVICI for computing the reward for interventions designed.
model 38
entries of the adjacency matrix predicted by the amortized posterior model for h4- This reward telescoping was inspired by Blau et al. 77]. We also show Section[4.T]tha this reward function also related approximation of multi-step EIG, the quantity of interest in sequential Bayesian experimental design (2
3.2 Policy
Architecture. In order for the policy to achieve amortization and generalize to new environments not during training, it should encode key design space symmetries. In partieular, for the problem intervention design, the interventions should be permutation equivariant to ordering of the variables and permutation invariant to the ordering of the history. This can be ensured by transformer architecture [57] wherein self-attention is applied alternately-one
over the variable axis and next over the samples axis [34]. More precisely, we input h E (t)++t)xdx2 and apply self-attentior(] over first the no + t axis and next over the d axis. This ensures that the history representation is permutation equivariant over both the axes [36. After multiple layers of alternating self-attention, we apply
pooling over the samples (dim. mi axis, which gives encoding of size 1 of the history B, E dR that respects the desired symmetries. The same symmetries apply for amortized causal structure learning, hence the reward model AVICI also leverages the alternate attention architecture. The history embedding B, is then passed through a multi-layer perceptron, whose outputs parameterize the logits of the Caussian-Tanil distribution (26), from which the interventions are sampled. In our setting, we model both the intervention targets and intervention values, hence I is 2 dimensional. Gaussian-Tanh samples range from -1 We use [: 0] to encode the interventional targets by discretizing the values to a binary mask (0 and 1) by thresholding where indicates intervention on the variable ys. If an intervention on yi active, the value to intervene with is given by TV 1].
Training. Training the policy involves addressing two main challenges: computing the reward in eq. 3 since {A, 0} would be unknown for real datasets, and optimizing this reward, which discrete. In order to address the first challenge, we simulate interventional data p(y {A', 0*), I1) for sample {A', 6'} p(A, 9) from the prior using a simulator. Such simulators exist for single cell gene regulatory networks (e.g.
51 and are becoming increasingly widespread in other
Every self-attention is multi-headed followed by position-wiz wise feedforward layer and layer normalization. in standard transformer block.
5

---

properties that domains [24] 2 The reward model AVICI is pretrained on datasets from the prior P({A,0) using the same simulator. During training of the policy, we only use the pretrained reward model for inference and do not update its parameters. To address the second challenge, we train our policy using Soft-Actor Critic (SAC) [26], an off-policy reinforcement learning algorithm that does not require rewards to be differentiable. We use the REDQ version of SAC to improve sample efficiency [13] REDQ trains multiple Q-function networks to optimize the reward. For each Q-function network, we use transformer based history state encoder with architecture similar to that in the policy, but the weights are not shared. This is beneficial because the same equivariance-invariee hold for the policy should also hold for the Q-function.
Inference. Deploying the policy in a real (i.c. not simulated) environment amounts to a rollout of the policy through interaction with the real environment. This requires just forward pass of policy network for each time step t. Note that we do not need intermediate Bayesian inference or other estimation of the causal graph the collected data.
4 Choice of Reward Function
4.1 Connection to Sequential Bayesian Experimental Design
As discussed in Appendix A the problem tackle has connection to likelihood-free sequential Bayesian experimental design [2 30. With the aim of gathering data to learn about the causal graph
A, the multi-step expected information gain (EIG) can be written
EIG(A;
E T. T(((110))
[lo [+]] + const.
(4)
Since the posterior P(A he) is intractable, could replace it by approximate posterior ((A h2). This gives rise to the Aarber-Agaoov (BA) bound 66 21.. which was recently explored in a sequential context by Blau et al. 8. This tells that we have an EIG lower bound by using q in place of pt:
EIG(A:T<) 2
E
[log ((A [h)] + const.
To...T.P(AA.
We can interpret eq. 5 in simple ferms-takeen log q(A fu)) as reward function is equivalent to optimizing a lower bound on the EIG. Although eq. (5) implies that we only receive a reward on the final state he, it is possible to rewrite this using telescoping rewards (7] exactly
do in eq. (3). The BA bound therefore represents the closest point of comparison between the method we outline in Section 3 and sequential Bayesian experimental design. As with the BA bound, we make use of an amortized approximate posterior distribution g ( he) that works backwards from data hy to predict the graph that might have generated it. Unlike the BA bound, however, we use the adjacency matrix accuracy to compare the true A samples Å from the amortized posterior, rather than computing the log-likelihood of the true graph under that amortized posterior, log ((A At). We found that this worked better in practice. Nevertheless, we see close relationship between the approach we take and the methods of sequential Bayesian experimental design.
4.2 Other Possible Reward Functions
Any target metric for causal structure learning like structural hamming distance computed on the amortized posterior could be used as a reward function. Depending on the application, domain specific causal graph objectives could also be considered. While large number of possibilities exist, we use expected number of correct entries of adjacency matrix, eq. (3), as the reward for training CAASL. As opposed to Structural Hamming Distance (SHD) and Area Under Precision Recall Curve (AUPRC), eq. (3) is straightforward to compute and parallelize.
5 Experiments
We train CAASL policy on two challenging environment domains: 1. Synthetic design environment with a causal model defined by linear mechanisms and additive noise, and 2. SERGIO [17] single-cell simulator corresponding to any gene regulatory network. For each domain, we train single CAASL policy on distribution of design environments with d 10. A distribution of

---

CAASL Policy Rallott
Ry MI
1=
Ground Truth
6
Vf Insevsstion
-03.5-501
Figure 3: Visualization of the rollout of the trained CAASL policy a randomly sampled en- vironment with no = 50 initial observational samples. Colored circles indicate nodes with a do intervention. The policy selects interventions that mostly correspond to the variables with a child in the ground truth graph. At t = 2, the policy selects the only child yı. which breaks all direct causal effects. This gives lesser information about the overall causal model. After this, y11 is never chosen. Initially, the policy is exploratory wrt targets and exploitative wrt values. This trend is reversed as the episode progresses. The policy is trained environments with d 2, therefore it has not seen any graphs with d = 3 before.
intervention design environments is defined by the choice of prior over causal models p(A, 6), which includes priors over graphs A (e.g. Erdös-Rényi (20)), mechanism parameters 0 and noise. We define Out-of-Distribution (OOD) environment as any environment with the choice of prior (either over graphs, mechanisms parameters or noise) that different from training. addition to these distribution shifts, we also consider OOD environments wherein the priors remain the same, but either the dimensionality the data increases (i.c. 10), or the performed intervention type changes. given in Appendix [ All evaluation
Precise choice of training and OOD testing distributions experiments are conducted on environments with causal model parameters that CAASL has never seen during training, regardless of whether the environment is in-distribution OOD. In addition, all evaluation done by just forward passing the history through the policy.
Baselines. We compare our approach with amortized strategies: Random and Observational. Random corresponds to obtaining data from random interventions, while Observational corresponds collecting more observational data. The random baseline is shown be very competitive for the problem of active causal structure learning, especially when the experimental budget or the density of graphs is sufficiently high [19 55 (6). For the synthetic design environment domain, we also compare with DifCCBED (56) and SS Finite (53). These intervention strategies use likelihood of the data to perform designs. So in certain OOD synthetic design environments and the single-cell simulator SERGIO where the likelihood is not available,
omit these baselines. DiffCBED and SS Finite rely on an approximate causal graph posterior to design interventions. As suggested (56). we use bootstrapped GIES 27 23 the approximate posterior distribution for these baselines. For an evaluation task on 100 random design environments with a budget of 10, DiffCBED and SS Finite methods require approximate posterior inference of the causal graph 1000 times. The performance of DiffCBED and SS Finite are limited by the performance the underlying approximate posterior inference method they rely on. In general, approximate posterior inference for causal structure is computationally expensive problem with significant limitations [40
All evaluation done on 100 random test enviromments. As CAASL is an intervention design method, we measure the cumulative rewards with = 1 (returss) obtained from the graph predicted by the amortized posterior. However, structure learning related metrics like the Structural Hamming Distance (SHD), the Area under Precision Recall Curve (AUPRC) and Fl score (Edge FI) between the graph predicted by the amortized posterior and the true graph 38, 4. Precise definition of these metrics is provided in
Metrics.
the sake of completeness, we also measure
7

---

132
AUPRC (1)
Edge
Banden
28
26
0.55
0.50
0.35
0.30
22
10
6
10
10
Intervention Herratio
Intervantion Beration
Durration
Beration
Figure 4: Amortization results of various intervention strategies on 100 random test environments. CAASL significantly outperforms other intervention strategies. Shaded area represents 95% CI.
Zeco Ont Goneralltaiio Goayle Marta Nokie
Jes Graghh
21 Cragph Merhaanions
1d Ieterobedaate Naker
Ert betecvation Tage
70
00
66
Figure 5: Zero-shot OOD returns of CAASL on 100 random environments with distribution shift coming from (a) graphs, (b) graphs and mechanisms, (c) graphs, mechanisms and noise, (d) noise changes from homoskedastic to heteroskedastic, and finally (e) intervention changes from do to shift intervention. CAASL outperforms other intervention strategies. Shaded area represents 95% CL.
Appendix G.2
We find that in most cases all the metrics are correlated. Therefore, unless otherwise mentioned, we only report the returns and relegate the other metrics to Appendix G.2
5.1 Synthetic Design Environment
Training Distribution of the Design Environment. We train CAASL on synthetically generated data, wherein p(A, 8) consists of linear SCMs with additive homoskedastic Gaussian noise. The dimensionality during training is d 10. The prior over causal graphs is Erdōs-Rényi (20], with 3 edges per node expectation. The prior over linear coefficients is chosen such that the marginal variance of each variable is close to 1. This is done to ensure that structure learning algorithms are not sensitive to the scale of the data [46). During training, an intervention exclusively corresponds to do intervention. Further, we set no 50 and the budget T' is fixed to 10.
Training Details. We train CAASL with 4 layers of alternating attention for the transformer, followed by max pooling operation over the history, to give embedding with size = 32. SAC related hyperparameters are tuned based on performance on held-out design environments. Details of the architecture, hyperparameter tuning and optimizer is given in Appendix D For the reward model, we use AVICI that is pretrained on random linear additive noise datasets.
Amortization Performance. We test on novel environments with hidden parameters sampled from the training prior p(A, 0). Results are provided in fig. [4] We find that our policy significantly outperforms the random baseline in terms of returns as well more common structure learning metrics like the SHD, AUPRC and Edge F1. For instance, our method achieves returns close to 76 with just 10 interventional samples. while the random baseline achieves close 72. Other intervention strategies like DifICBED (56) and SS Finite [533 perform worse, while still making use of the likelihood and performing intermediate inference of causal structure.
Zero-Shot OOD Generalization. We also test the trained CAASL policy on environments when the prior changes. All results correspond to zero-shot performance, obtained by just a forward forward pass of the trained policy. fig. presennts the returns of CAASL alongside other applicable baselines. We consider shifts which become increasingly different from training: (1) the graph prior changes from Erdōs-Rény 20 to Scale-Free [5] (fig. (a)), (2) apart from the graph, prior over
8

---

mechanisms also change (fig. (b)), (3) apart from graph and mechanisms, the noise distribution changes from Gaussian to Gumbel (fig. (c)). We find that our policy achieves better performance than random strategy by significant margin. Further, our method also outperforms DifCCBED and SS Finite which explicitly optimize for designs corresponding to these environments. In ad- dition to these OOD settings, we also consider OOD environments in which the prior remains the same, but the noise is heteroskedastic instead of homoskedastic (fig. (d)) Although the random strategy is very competitive, CAASL performs better. Finally, we consider OOD environments wherein the intervention design suggested by the policy during testing is used for performing a shift intervention instead of do (fig. (d)). CAASL performs better than baselines even in this setting.
Zero-Shot OOD: Dimeusionality
Slightly different to the above OOD environments, we also consider OOD environments in which the dimension- ality of the data changes during test- ing, but the prior remains the same. fig. presents the results, with further details in fig. g CAASL obtains bet- ter returns on average than random at all points of acquisition. The relative performance of CAASL decreases as d increases (up to d = 30) from train- ing, although it still performs better than random.
Relative Perf, net Ramdom
315.0
0.020
312.5
Rel. Reture
0.015
310.0
0.010
307.5
0.005
0.000
302.5
0
5
9 10
2 3
8
10
Intervention Heration
Interveation Iteration
Figure 6: Zero-Shot OOD generalization results when dimen- sionality d changes for synthetic environment. For training, d = 10. Left: Zero-Shot test returns with d = 20. Right: Relative mean zero-shot returns of CAASL wrt random for different d. Results on 100 random environments. Shaded area represents 95% CI.
5.2 Single-Cell Gene Regulatory Network Environment
In this setting, we train CAASL policy based on the single-cell gene expression simulator SER- GIO [I7]. Given causal graph that corresponds to interaction between different genes terms of their transcription regulation, SERGIO simulates expressions of genes that correspond to steady state of differential equations that govern the interaction between the genes. Each variable entry indicates the count of mRNA that is produced corresponding to that gene, similar to the output of modern single-cell RNA sequencing (scRNA-seq) technological platforms [39]. In addition, SERGIO can be extended to support interventions. Interventions in this setting correspond to either gene knockouts, wherein the transcription rate of the intervened gene is actively set to 0, or gene knockdown, wherein the intervened gene's transcription rate is actively halved. Since there is value selection in this setting, the dimensionality of the policy is d instead of d x 2. SERGIO also simulates technical noise such that the statisties of the data match that obtained from real scRNA-seq platforms. Some of the technical noise includes dropouts (missingness of the data), library size effects and random outlier effects. Most notably, atleast 70% of the data is missing in most single cell platforms. Therefore, in this domain, not only is the likelihood intractable, but also there is high amount of missing data. We do not impute the missing data, but just encode it with O.
Training Distribution of the Design Environment. For training, we set 10, with nig = 50 observational (wild-type) data with budget T 10. The statistics of the data corresponds to 10X Chromium platform [17] wherein around 74% of the data is dropped out. The prior over causal graphs is set to Erdōs-Rény. (20] with 3 edges per node on average. An intervention exclusively corresponds to gene knockout, We provide details of the simulator in appendix B.2 and the training prior parameters in appendix C.2
Training Details. We train CAASL with layers of alternate attention, followed by a max pooling operation, giving an embedding of size I = 32. Just like in the synthetic linear domain, SAC related hyperparameters are tuned based on performance on held-out design environments. Details are given in appendix D Once trained, we perform forward pass of the history through the policy to obtain intervention designs for all test environments. For the reward model, we use AVICI that is pretrained on this simulator with post-noise data statistics matching that of 10X chromium platform.

---

0s4 Pr
let In Dia..
Ih 000 Gragls
to 0000 Nebe
141 OOE: Jetervcttee bye
80 000 Naw Intercation
18.5
15.0
14.3
13.0
2-5
17.5
7400
142
17.
71.5
340
Figure 7: Results SERGIO environment with 100 random environments. (a) corresponds to in-distribution performance, (b)-(e) correspond to zero-shot OOD performance with distribution shift coming from either (b) graphs, (c) technical noise, (d) intervention changing to gene-knockdowg (c) Noisy interventions, which include off-target effects. Shaded area represents 95% CL.
Amortization Performance. The in-distribution amortization performance is presented in fig. 7a). After 5 acquisitions, CAASL obtains better returns than random.
Zero-Shot OOD Generalization. We test the CAASL policy when the environment is subject to various test-time distribution shifts. Robustness to distribution shifts important in real world- settings, where experimental conditions change. We consider 4 different OOD environments: (1) the prior over graphs changes from Frdōs-Rényí to Scale-free (fig. (b)), (2) The perturbation platform changes to Drop-Seq [39], wherein among other noise parameters, the amount of missing data increases from 74 to 85% (fig. (c), (3) The intervention type changes from knockout to knockdown (fig. d)) and, (4) Noisy knockout interventions, where there a 10% chance that either the intended gene does not get knocked out, or an off-target gene is knocked out (fig. (e)). We find that CAASL shows excellent robustness to these distribution shifts, and obtains better returns than baselines. When the intervention type changes, the random baseline is still competitive, An interesting observation is that the for the OOD graph and the OOD noise setting, the model shows exploratory behavior in the beginning where the returns decrease, but later becomes better than random. Robustness to various distribution-shi shifts demonstrates the generality of the policy.
Limitations. In the SERGIO gene regulatory network environment, we found that the random strategy is very competitive in general. In particular, for the setting of zero-shot OOD generalization with increasing data dimensionality, the performance (fig. 9p. We hypothesize that since almost 74% of data is missing, the incorporated design space symmetries might not be as relevant, which might limit the extent of zero-shot generalization. In addition, while the empirical results in the shot OOD settings extremely encouraging, there are no theoretical guarantees on the performance of CAASL in this setting.
random strategy is on par with CAASL
Conclusion
We have presented amortized and adaptive intervention design strategy CAASL, that does not require intermediate inference of the causal graph. CAASL is based on a policy parameterized by transformer which is permutation equivariant to ordering of the variables and permutation invariant to ordering of the collected data. Through various experiments, including on simulator which respects the data statistics of real gene-expression readouts, find that method shows excellent amortized intervention design and zero-shot generalization to significant distribution shifts. The achieved performance motivates intervention design in complex settings high-throughput eeprrugnneeepermmennt with large batch sizes and utilization existing real offline data for designing interventions.
Acknowledgements
The authors gratefully acknowledge the Gauss Centre for Supercomputing e. V.V. for funding this project by providing computing time on the GCS Supercomputer JUWELS at Jülich Supercomputing Centre (JSC) [3].
www.aringgeeww.wwwugsssseoomm (www.gauss-centre.og
10

---

References
[1] Raj Agrawal, Chandler Squires, Karren Yang, Karthikeyan Shanmugam, and Caroline Uhler. Abed-strategy: Budgeted experimental design for targeted causal structure discovery. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 3400-3409. PMLR, 2019.
[2] Ossama Ahmed, Frederik Träuble, Anirudh Goyal, Alexander Neitz, Yoshua Bengio, Bern- hard Schölkopf, Manuel Wüthrich, and Stefan Bauer. Causalworld: A robotic manipulation benchmark for causal structure and transfer learning. arXiv preprint 0rXiv:2010.042290 2020.
[3] Damian Alvarez. Juwels cluster and booster: Exascale pathfinder with modular supercomputing architecture at juelich supercomputing centre. Journal of large-scale research facilities JLSRF, 7:A183-A188 183, 2021.
[4] Yashas Annadani, Nick Pawlowski, Joel Jennings, Stefan Bauer, Cheng Zhang, and Wenbo Gong. Bayesdag: Gradient-ban -ased posterior sampling for causal discovery. arXiv preprint arXiv:2307.1311, 2023.
[5] Albert-LászIó Barabási and Réka Albert. Emergence of scaling in random networks. science, 286(5439):509-555 1999.
[6] David Barber and Felix Agakov. The im algorithm: variational approach information maximization. Advances in neural information processing systems, 16(320):201, 2004.
[7] Tom Blau, Edwin V Bonilla, Iadine Chades, and Amir Dezfouli. Optimizing sequential experimental design with deep reinforcement learning. In International conference on machine learning, pages 2107-2128. PMLR, 2022.
[8] Tom Blau, Edwin Bonilla, ladine Chades, and Amir Dezfouli. Cross-entropy estimators for sequential experiment design with reinforcement learning, arXiv preprint arXiv:2305.1843: 2023.
[9] Philippe Brouillard, Sébastien Lachapelle, Alexandre Lacoste, Simon Lacoste-Julien, and Alexandre Drouin. Differentiable causal discovery from interventional data. Advances in Neural Information Processing Systems, 33:21 835-18655,,, 2020.
[10] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. Advances in neural information processing systems, 33:18 33:1877-1901, 2020.
[11] Robrecht Cannoodt, Wouter Saelens, Louise Deconinck, and Yvan Saeys. Spearheading future omics analyses using dyngen, multi-modal simulator of single cells. Nature Communications, 12(1):3942, 2021.
[12] Kathryn Chaloner and Isabella Verdinelli. Bayesian experimental design: A review. Statistical science, pages 273-304, 1995.
[13] Xinyue Chen, Che Wang, Zijian Zhou, and Keith Ross. Randomized ensembled double q- learning: Learning fast without a model. arXiv preprint arXiv:2101.0590 2021.
[14] Davin Choo and Kirankumar Shiragur. Adaptivity complexity for causal graph discovery. In Uncertainty in Artificial Intelligence, pages 391-402. PMLR, 2023.
[15] Haoyue Dai, Ignavier Ng. Gongxu Luo, Peter Spirtes, Petar Stojanov, and Kun Zhang. Gene regulatory network inference in the presence of dropouts: a causal view. arXiv preprint arXiv:2403. 15500, 2024.
[16) Tristan Deleu, Mizu Nishikawa-Toomee Toomey, Jithendaraa Subramanian, Nikolay Malkin, Laurent Charlin, and Yoshua Bengio. Joint bayesian inference of graphical structure and parameters with single generative flow network. Advances in Neural Information Processing Systems, 36. 2024.
[17] Payam Dibaeinia and Saurabh Sinha. Sergio: single-cell expression simulator guided by gene regulatory networks. Cell systems, 11(3):252-271,
2020. 2020.

---

[18] Finale Doshi- Velez and George Konidaris. Hidden parameter markov decision processes; A semiparametric regression approach for discovering latent task parametrizations. IJCAI: proceedings of the conference, volume 2016, page 1432. NIH Public Access, 2016.
[19] Frederick Eberhardt, Clark Glymour, and Richard Scheines. On the number of experiments sufficient and the worst case necessary identify all causal relations among variables. arXiv preprint arXiv:
2012.
[20] Paul Erdös, Alfréd Rényi, et al. On the evolution of random graphs. Publ. math. inst. hung. acad. sci, 511)117-60, 1960.
[21] Adam Foster, Martin Jankowiak, Matthew O'Meara, Yee Whye Teh, and Tom Rainforth. A unified stochastic gradient approach to designing bayesian-optimal experiments. In International Conference on Artificial Intelligence and Statistics, pages 2959-2969. PMLR, 2020.
[22] Adam Foster, Desi R Ivanova, Ilyas Malik, and Tom Rainforth. Deep adaptive design: Amortiz- ing sequential bayesian experimental design. In International conference on machine learning, pages 3384-3395. PMLR, 2021.
[23] Nir Friedman, Moises Goldszmidt, and Abraham Wyner. Data analysis with bayesian networks: A bootstrap approach. arXiv preprint 1301. 2013.
[24] Juan L Gamella, Jonas Peters, and Peter Bühlmann. The causal chambers: Real physical systems testbed for ai methodology. arXiv preprint arXiv: 11341, 2024.
[25] Kristjan Greenewald, Dmitriy Katz, Karthikeyan Shanmugam, Sara Magliacane, Murat Ko- caoglu, Enric Boix Adsera, and Guy Bresler. Sample efficient active learning of causal trees. Advances in Neural Information Processing Systems, 32, 2019.
[26] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference machine learning, pages 1861-1870. PMLR, 2018
[27] Alain Hauser and Peter Bühlmann. Characterization and greedy learning of interventional markov equivalence classes of directed acyelic graphs. The Journal of Machine Learning Research, (3(1):2409-2466) 2012.
[28] Christina Heinze- -emll Marloes H Maathuis, and Nicolai Meinshausen. Causal structure learning. Annual Review of Statistics and Its Application. 5:371-391, 2018.
[29] Xun Huan and Youssef M Marzouk. Sequential bayesian optimal experimental design via approximate dynamic programming. arXiv preprint arXiv:1604.0088 160.008300 2016.
[30] Desi R Ivanova, Adam Foster, Steven Kleinegesse, Michael U Gutmann, and Thomas Rainforth. Implicit deep adaptive design: Policy-based experimental design without likelihoods. Advances in Neural Information Processing Systems, 34:25785-25798, 2021.
[31] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint ar 2tt::000..883.1 2020.
[32] Nan Rosemary Ke, Silvia Chiappa. Jane Wang, Anirudh Goyal, Jorg Bornschein, Melanie Rey, Theophane Weber, Matthew Botvinic, Michael Mozer, and Danilo Jimenez Rezende. Learning to induce causal structure. arXiv preprint arXiv:2204.04875, 2022.
[33] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 11269914 2014.
[34] Jannik Kossen, Neil Band, Clare Lyle, Aidan N Gomez, Thomas Rainforth, and Yarin Gal. Self-attention between datapoints: Going beyond individual input-output pairs in deep learning. Advances in Neural Information Processing Systems, 34:28742-28755 2021.
[35] Andrew Lampinen, Stephanie Chan, Ishita Dasgupta, Andrew Nam, and Jane Wang. Pas- sive learning of active causal strategies in agents and language models. Advances in Neural Information Processing Systems, 36, 2024.
12

---

[36] Juho Lee, Yoonho Lee, Jungtaek Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set transformer: A framework for uttention-based permutation-invant invvariant neural networks. In International conference machine learning, pages 3744-3753. PMLR, 2019.
[37] Dennis V Lindley. On measure of the information provided by an experiment. The Annals of Mathematical Statistics, 27(4):986-1005, 1956.
[38] Lars Lorch, Scott Sussex, Jonas Rothfuss, Andreas Krause, and Bernhard Schölkopf. Amortized inference for causal structure learning. Advances in Neural Information Processing Systems, 35: 13104-13118, 2022.
[39] Evan Macosko, Anindita Basu, Rahul Satija, James Nemesh, Karthik Shekhar, Melissa Goldman, Itay Tirosh, Allison R Bialas, Nolan Kamitaki, Emily M Martersteck, et al. Highly parallel genome-wide wide expression profiling of individual cells using nanoliter droplets. Cell, 161 (5):1222-121,, 2015.
[40] Amir Mohammad Karimi Mamaghan, Panagiotis Tigas, Karl Henrik Johansson, Yarin Gal, Yashas Annadani, and Stefan Bauer. Challenges and considerations in the evaluation of bayesian causal discovery. arXiv preprint arXiv:2406. 03209, 2024.
[41] Francesco Montagna, Atalanti Mastakouri, Elias Eulig, Nicoletta Noceti, Lorenzo Rosasco, Dominik Janzing, Bryon Aragam, and Francesco Locatello. Assumption violations in causal discovery and the robustness of score matching. Advances in Neural Information Processing Systems, 36, 2024.
[42] Judea Pearl. Causality. Cambridge university press, 2009.
[43] Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Elements of causal inference: founda- tions and learning algorithms. The MIT Press, 2017.
[44] Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. Advances in neural information processing systems, 20, 2007.
[45] Tom Rainforth, Adam Foster, Desi R Ivanova, and Freddie Bickford Smith. Modern bayesian experimental design. Statistical Science, 39(1):100-1144 2024.
[46] Alexander Reisach, Christof Seiler, and Sebastian Weichwald. Beware of the simulated dag! causal discovery benchmarks may be easy to game. Advances in Neural Information Processing Systems, 34:27772-27784 2021.
[47] Paul Rolland, Volkan Cevher, Matthäus Kleindessner, Chris Russell, Dominik Janzing, Bernhard Schölkopf, and Francesco Locatello. Score matching enables causal discovery of nonlinear additive noise models. International Conference on Machine Learning, pages 18741-18753. PMLR, 2022.
[48] Dominik Rothenhäusler, Christina Heinze, Jonas Peters, and Nicolai Meinshausen. Backshift: Learning causal cyclic graphs from unknown shift interventions. advances in neural information processing systems, 28, 2015.
[49] Andreas WM Sauter, Erman Acar, and Vincent François-Lavet, A meia-reiiforremment learning algorithm for causal discovery. In Conference Causal Learning and Reasoning, pages 602-619, PMLR, 2023.
[50] Andreas WM Sauter, Nicolò Botteghi, Erman Acar, and Aske Plaat. Core: Towards scalable and efficient causal discovery with reinforcement learning. arXiv preprint 4rX1v:2401.16974, 2024.
[51] Thomas Schaffter, Daniel Marbach, and Dario Floreano. Genenetweaver: in silico benchmark generation and performance profiling of network inference methods. Bioinformatics, 27(16): 2263 2270, 2011.
[52] Peter Spirtes, Clark Glymour, and Richard Scheines. Causation, prediction, and search. MIT press, 2001.
13

---

[53] Scott Sussex, Caroline Uhler, and Andreas Krause. Near-optimal multi-perturbation experimen- tal design for causal structure learning. Advances in Neural Information Processing Systems, 34:777-788, 2021.
[54] Alejandro
Paul Bertin, Stefan Bauer, Hananeh Alie,, Yoshua Bengio, and Fabian Theis. Causal machine learning for single-cell genomics. arXiv preprint arXiv:2310. /4935, 2023.
[55] Panagiotis Tigas, Yashas Annadani, Andrew Jesson, Bernhard Schölkopf, Yarin Gal, and Stefan Bauer. Interventions, where and how? experimental design for causal models at scale, Advances in neural information processing systems, 15:24430-221144 2022.
(56) Panagiotis Tigas, Yashas Annadani, Desi R Ivanova, Andrew Jesson, Yarin Gal, Adam Fos- ter, and Stefan Bauer. Differentiable multi-target causal bayesian experimental design. In International Conference Machine Learning, pages 34263-34279. PMLR, 2023.
[57] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
[58] Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transform- ers. Proceedings of the IEEE/CVE conference on computer vision and pattern recognition, pages 12104-12113, 2022.
[59] Zeyu Zhang, Chaozhuo Li, Xu Chen, and Xing Xie. Bayesian active causal discovery with multi-fidelity experiments. Advances in Neural Information Processing Systems, 36, 2024.
[60] Shengyu Zhu, Ignavier Ng. and Zhitang Chen. Causal discovery with reinforcement learning. arXiv preprint 1rXiv 2019.
A
Connections to Bayesian Experimental Design using Expected Information Gain
A.1 Greedy Approaches
We consider the model with unknown parameters {A, 0}, prior pl (A,0}) and likelihood of the data
p(y {A, 8). I) under an intervention I. The Expected Information Gain (EIG) is given by:
EIG(I) (((AL1))p((KAA)))1
E
[logp(y) (A,0),1) logp(y []]
(6)
In the standard greedy approach to Bayesian experimental design [45]. given a history h4-1, we replace the prior p( ((A,, 0}) with the posterior conditional existing data p( P([A,)) 0} hh-1) and then estimate the one-step EIG
EIG(T)
E Ml(A,0) (..0yiiNn-1)0(l(A.oo (P(U)(A.))..)
log (A,0), -)og 14-1, T)]
where p(y I) SA p( [A, he-1)p(y (A,0), ; 7). The EIG is estimated for each candidate design 1, and the one with the largest EIG is selected. This gives rise the policy Tgreedy- which was applied to causal structure learning by e.g. Tigas et al. [56].
A.2 Non-greedy Approaches
Non-greedy approaches to experimental design using EIG were also explored 29 22 Using the parameters {A, 0} of our model and the notation of Foster et al. [22, the EIG of a sequence of t experiments generated using policy as about (A,0) 0} is given by
E EIG((A,0):=a) (((0))))
[log p(he [A.0). TO) logp(he TT)]
(8)
where p(h, (A.0), = II p(I- Th (4,0). I.)
(9)
14

---

and p(h, TT) is the marginal of this quantity over (A, Equation 8 cannot be computed exactly, so likelihood-baste based [22] and likelihood- freee [3) approximations have both been explored. The likelihood-basee sPCE lower bound on EIG 22 was also used as reward function to train an
RI policy 7
The problem we consider in this paper is decidedly likelihood-free for two reasons: (1) some simulators do not have explicit likelihoods, (2) even where likelihood available, is generally conditional on both A and 0. We made the choice to focus on experimental design to learn A (ignoring information gain about 8). this case, the relevant EIG is
FIG A; kT)
E pl ((4.0))((1((000))...)
[log A,)
log p(g [T)]
(10)
We would have to perform costly marginalization over 0 to obtain the relevant likelihood, p(h,, A, na).
Equation
10 can be rearranged using Bayes Theorem to read
EIG(A: ((4
E
[log((A ht) log p(4)
(11)
E p((A,))) (b.4)))))))
[lggp(A hc) + const,
(12)
where we make the observation that E|- log p(4) is constant with respect to the design policy TT. The form of EIG eq. (12) is the jumping off point for the BA bound, in which we replace p(A h₂) with an approximate posterior.
B Details of Design Environments
B.1 Synthetic Design Environment
We consider linear additive noise models. For homoskedastic noise, they can written as
y/ 0.¹ (f) + c4
(13)
where 0, ~ p(0) and e; ~ Pnoise which can be either Gaussian or gumbel distribution.
For heteroskedastic noise, the above equation can be written as:
= 0.¹¹ yf + oz (1)) -E;
(14)
where a;(-) is scaling factor that is obtained by squash operation (x) log(1 [1 + exp(g:
on any nonlinear function J-- Similar [8]. we implement di with 100 random Fourier feature features 44 Random Fourier feature functions require a kernel, for which we use Squared-
Exponential Kernel with length Is 10 and output scale os 2.
B.2 Single-Cell Gene Regulatory Network Environment
SERGIO 17 is single-cell simulator of gene expression for any user provided gene-regulatory
network that resembles the data obtained with modern single-cell RNA sequencing (seRNA-seq)
technologies technologies like Drop-Seq 39 and 10X Chromium [17].
We provide brief overview of simulation procedure of SERGIO. Our simulation is based on the original simulator provided by Dibaeinia and Sinha 17 and hence further details can be found the paper. This simulator further vectorize the simulator to produce datasets for multiple regulatory networks parallely. The simulator of Dibaeinia and Sinha [17] is publicly available under GPL-3.0 license.
extended by Lorch et al. 38 to support knockouts and knockdowns. We
B.2.1 Simulation of Gene Expressions
We first describe the data that is generated without interventions, also called as wild-type measure- ments. Later, we then describe interventional simulations.
For an observational dataset of size n x d, the data produced from the simulator corresponds to the count of mRNA corresponding to each gene of different single cells. In particular, given
15

---

10X Chromium
Drop-Seq
Poutlier
0.01
0.01
Moutlier outller
3.0, 1.0
3.0, 1.0
Hibb Tib
6.0, 0.3
4.4, 0.8
6, 17
74,
85,
Table I: Technical noise parameters for 10X Chromium and Drop-Seq Single-Cell RNA sequencing platforms that is used for experiments in this work.
regulatory network A, steady-state of the regulatory differential equation simulated for n single cells that is regulated according to A. SERGIO allows for biological variations within this pool of single cells, such as varying basal rates of master regulator genes. master regulator gene a gene with upstream genes in A. The transcription rate of master regulator gene is usually constant, called the basal rate. Usually, cell of the same type have the same basal rates. For simulation, we consider cell types and ne single cells of each type such that c- nle In this work, we fix 5. is less than 5, we sample n single-cells at random after simulating single-cells corresponding to each cell type. The expression of all the downstream genes is effected nonlinearly by the mRNA production and decay of their respective regulatory genes. This expression simulated according to Langevin equation. Finally, the clean data is the continuous valued mRNA concentration that is measured at random- time points of the steady-state Langevin simulation.
The clean data is then subject to technical measurement noise. The series of simulated noise is as follows:
1. With probability Poutlier [0, 1]. gene is converted outlier gene that has unusually high expression across different cells. This is done by multiplying the current expression with values from a log-normal with mean Poutlier and scale Joutlier-
2. Based on the single-cell pool considered, different cells have different count distribution data. This is called as library-size effect, which is modeled log-normal distribution with mean jiih and scale dlib
3. The dropouts are simulated with parameters dropout percentile [0, 100] and the temper- ature of the logistic function rj E R,+
The actual values of the noise parameters differs across different scRNA-seq technologies. The final scRNA-seq resembling simulated data is obtained by sampling from a poisson distribution that is parameterized by the post-noise mRNA concentration levels.
During a gene knockout, the upstream genes do effect the knocked out gene. The activity of the gene is set to 0 and is propagated downstream as before. In gene knockdowns, the upstream genes still work the same way as before, however, the expression of the knocked down gene is multiplied by 0.5 for every time step of the steady state simulation. The reduced gene expression of the knocked down gene is propagated downstream genes before.
B.2.2 Simulation Parameters
For generating the clean data, we use the following parameters across all settings, which similar to what is used for training the reward model 38
Number of cell types 5.
Basal rates b ~ Uniform(1, (1,3).
Rate of decay of each gene À = 0.8.
Langevin equation related parameters: Hill function coefficient y system noise scale e, = 1.0, interaction strength ~ Uniform(1,5 (1,5) and the sign of the interaction which indicates a promotive repressive regulation sgn(k) Bernoulli(ps) with pk Beta(0.5, 0.5).
For technical noise, we consider two different platforms: 10X Chromium and Drop-Seq. The noise parameters used are suggested in [17]. These parameters are presented in table [
16

---

C Training and OOD Distributions of Design Environments
C.1 Synthetic Design Environment
For training distribution, we make the following choices:
Prior over graphs p(4) = (kin 3) is an Erdós Erdös-Rényí 20 with 3 edges per node in expectation.
Prior over parameters p(0) = (() where oî is chosen such that marginal variance of each variable is [46].
Noise c (O0 of)) where o, InvGamma(10,, 1).
For an OOD distribution, all the priors except for the parameter that undergoes distribution shift remain the same as during training. We define the following OOD environments and their corresponding distribution shifts:
a Graphs: p(A) = PSF = 3) is Scale-Free S with 3 edges per node in expectation.
Graphs+Mechanisms:: Prior over graphs is p(A) psr (kin 3) and prior over parame- ters p(0) (0,1 oğ) with o2 chosen as during training.
Graphs-Mech.+NNoist Prior over graphs is p(A) = PSF (kin = 3) and prior over parame- ters p() = (0 oğ) with o² chosen as during training. Noise Gumbel(O, o.) where 0, ~ InvGamma(10,, 1).
Heteroskedastic Noise: The causal model changes from equation 13
equation
14
Intervention Type: The performed intervention the environment changes from a do to shift intervention.
f Dimensionality d: The dimensionality of the environment increases from training distribu- tion (d 10).
C.2 Single-Cell Gene Regulatory Network Environment
For training distribution, we make the following choices:
Prior over graphs p(4) = (kin 3) is an Erdōs 20 with 3 edges per node in expectation.
Prior over mechanisms given in appendix B.2.2
For technical noise, we consider the 10X Chromium platform whose parameters are given in [able
For an OOD distribution, all the priors except for the parameter that undergoes distribution shift remain the same as during training. We define the following OOD environments and their corresponding distribution shifts:
Graphs: P(A) (kr = 3) is Scale-Free with edges per node in expectation.
(Technical) Noise: The single-cell RNA sequencing platform changes from 10X Chromium Drop-Seq. thereby changing the noise levels table
The performed intervention in the environment changes from gene knockout to gene knockdown.
Intervention Type:
With a 10% probability, the gene suggested by the policy for knockout is either does not happen, or there is an off-target that is knocked-out. We achieve this by flipping the one hot encoding of the intervention target labels with 10% probability. But the history only appended with the intervention sampled from the policy. Therefore, the policy has knowledge of the noisy intervention.
Noisy Interventions:
17

---

Table 2: Hyperparameters used for training in CAASL.
Synthetic Environment
SERGIO Environment
Hyperameter Search
No. uttention layers (for policy, Q-Function) No. atlention heads (for policy, Q-Function)
3
8
8
Transformer parameters (History state encoder)
32
32
Dropout (Policy)
0.1
0.1
Pooling (Policy)
Max pool over samples
Max pool over samples
Max pool over samples, sum pool over variables
Max pool over samples, sum pool over variables
Pooling (Q Function)
Decoder parameters
Hidden sizes (for policy and Q)
(128, 128)
(128, 128)
Non-linearity
ReLU
ReLU
M
(2,3,5)
2
(1,3,5)
1
REDQ/ SAC
Y
0.95)
0.95
training
Buffer Size
(10e6, 10e73
10c7
10e6
parameters
Policy LR
(0.01, 0.001)
0.001
0.01
Q-Function LR
{3e 5, 3e
6)
3e 5
3e
0.01
0.01
D Training Details
D.1 Architecture Details
We use the alternating attention based transformer for both the policy and the Q-function approxima- tion. We maintain the same architecture for the transformer for both the policy and the Q-function, which we describe below.
For the transformer, we use a standard transformer block 57 with 8 heads of self-attention. As our transformer has alternating attention, each layer has two such self self-attention operations. Each self-attention is followed by feedforward layer, whose dimension set to 4 * | where the size of the state representation. We choose = 32. After L layers of alternate attention, we perform max pooling over ordering of the data to obtain the state representation. The state representation is passed through a two hidden layer MIP with 128 hidden dimensions each and ReLU nonlinearity.
D.2 Hyperparameter Tuning
algorithm based on SAC [26] trains M different Q-function networks and updates the gradients of each of them G times before updating the policy. We treat both these quantities as hyperaprameters, All the parameters are updated with the Adam optimizer [33] and the learning rate is tuned. We list all the hyperparameters and the corresponding grid search in table
REDQ KK
E Computational Resources
train all models on 40GB NVIDIA A100 GPU accelerators. We provide wall time of days, which results in total computational budget of 216 GPU hours for each model. We also tune hyperparameters outlined in table 2 for both environments, resulting in a total usage of 70,000 hours. For testing, just rollout the policy CPU, which can be completed in seconds.
18

---

Generaltration (Synthetic Emstrnment)
182.
605
7300
150.0
600
177.5
515
175.0
720.0
172.5
300
717.5
.90
3
6
Intervention Eierations
Intervantion
Figure Results of zero-shot OOD generalization when dimensionality of the data increases in the synthetic environment. Results are performed on 100 random test environments. Shaded area represents 95% CI.
SERGIO Sinulator 0OD Rosalli
Dyp
108.2 25
108.00
Metumnc
90.6
107.75
90.4
90.2
Intervention
Figure 9: Results of zero-shot OOD generalization when dimensionality of the data increases in the SERGIO environment. We notice that the random baseline is very competitive. We hypothesize that the symmetries encoded in the policy, which are crucial for generalization, might not be so relevant this setting due to high amount of missing data. Results are performed on 100 random test environments. Shaded area represents 95% CI.
Licenses
For the single-cell gene simulator, we make use of the publicly available repository which is released under GPL-3.0 License For the reward model AVICI, we make use of the publicly released code and trained models. These are released under MIT License For baselines, we use the DiffCBED open source repository, which is released under MIT License
G
Full Results
G.1 Results on Zero-Shot OOD Generalization to Higher Dimensions
The results for Zero-shot OOD generalization to problems of higher dimensions is available in figs. and 9
G.2 Results on all Metrics
Herein include all the results that correspond other metrics omitted in the main text. In particular, apart from returns, we measure Struetural Hamming Distance (SHD), Area under Precision Recall Curve (AUPRC), and Edge F1 (Edge F1) score. These additional metrics are defined follows:
SHD: Structural Hamming Distance (SHD) measures the hamming distance between graphs. In particular, is a measure of number of edges that are be added, removed or reversed to get the ground truth from the estimated graph. Since we have posterior distribution q(A he) over graphs, measure the expected SHD:
100
SHD+
[SHD(A, A R
Σ (A") ACT with A g(A) he)
100
19

---

Synthetic Gragh
032
AUPRC (12
Edge
Banden (iammtional
0.60
0.35
24
0.50
22
4
10
6
10
5
10
Intervemtio Beration
Intervantion Beration
Durrttio
Intervantion Buration
Figure 10: Results of zero-shot OOD graph setting with various intervention strategies on 100 random test environments. Shaded area represents 95% CI.
Synthetic OUD: Graph Mechanlsms
AUPRIC
Edge
Realen Oburrtiral
0.60
26
DUTETRED
0.55
0.40
0.50
0.35
10
0
9
10
5
9
10
4
Figure 11: Results of zero-shot OOD graph and mechanisms setting with various intervention strategies on 100 random synthetic test environments, Shaded area represents 95% CI.
where ACT is the ground-truth causal graph.
Edge F1: is F1 score of each edge being present or absent in comparison to the true edge
set, averaged over all edges.
AUPRC: It is the area under the precision recall curve obtained by thresholding the edge probabilities of the amortized graph posterior qAA he).
H Broader Impact Statement
This paper is concerned with applying deep reinforcement learning for causal experimental design. All the experiments are conducted either on simulator or synthetic data. Advances in the problem being addressed are impactful for drug discovery and better understanding of mechanisms that lead to diseases. The authors do not forsee any negative societal impacts of this work beyond what might be enabled due to general advancements machine learning.
Sy Synthetic OOD: Graph Mech. Noise
74
Reiurs (7)
SHD
:1:
AUPRCC (1)
Edge P1 (T)
CAASL
0.55
Radian
Oburrttral
0.54
0.36
71
25
0.52
0.36
0.50
0.48
0.32
Intervention Heration
Intervotion Ttenatin
Intervention Berratinn
Intervention Inerstion
Figure 12: Results of zero-shot OOD graph, mechanisms and noise setting with various intervention strategies on 100 random synthetic test environments. Shaded area represents 95% CI.
20

---

Syothetic OOD: Intervention Type Shift Interventions
Returns (1)
SIID 14:
0.60
AUPRC (1)
Edge (1)
CANSL
26
0.45
Redies
0.55
26
0.50
0.35
24
0.45
Intervention Ieeraton
Intervcation Beratioo
Inteveation Iteration
Interveation Berution
Figure 13: Results of zero-shot OOD intervention type setting with various intervention strategies on 100 random synthetic test environments. Shaded area represents 95% CI.
Synthetir 0OD:
Heterndndaatic Noine
Returas (7)
S1ID
AUPRCI (1)
Edge (1)
CAARI
0.525
0.375
0.500
0.350
26
0.475
25
0.450
0.300
Intervention Itrration
Interveatixn Dieration
Intervvaatio Dicration
Intervaation Iterstion
Figure 14: Results of zero-shot OOD heteroskedastic noise setting with various intervention strategies on 100 random synthetic test environments. Shaded area represents 95% CI.
SERGIO Amortization Performance
Returas (?)
AUPRC (13
Edge (1)
Sanden
0.63
0.54
0.52
74.5
74.0
23.5
23.0
73.5
Intervention Hteration
Interveution Iterodion
Intervention Iteration
Infervention Iteration
Figure 15: Results of amortization with various intervention strategies on 100 random SERGIO test environments. Shaded represents 95% CI.
SERGIO Graph
Returns
SHDD :4:
AUPROC (1)
Edge F1 (?)
Raalen
22.0
0.68
73.0
21.5
0.67
0.58
77.5
21.0
0.55
77.0
0.54
0.64
76.5
20.0
Intersention Herraionn
Intervcution Ineruiio
Intervationn Iteration
Interveation Eteration
Figure 16: Results of zero-shot OOD graph setting with various intervention strategics on 100 random SERGIO test environments. Shaded area represents 95% CI.
21

---

Drop-Sey
maa-Tw
Tanden
0.52
74.5
24.0
0.60
23.5
0.45
73.5
0.58
0.45
10
3
10
Intervantion
Intervention
Figure 17: Results of zero-shot OOD scRNA-seq platform and their noise setting with various intervention strategies on 100 random SERGIO test environments. Shaded area represents 95% CI.
SERGIO OOD: Intervention
Gene Knockdown
CAASL
AUPRC
Randoa
34.5
0.83
75.0
0.02
24.0
0.61
74.55
0.50
23.5
74.0
0.59
0.45
23.0
0.58
0.46
Figure 18: Results of zero-shot OOD intervention type changing to gene knockdown with various intervention strategies on 100 random SERGIO test environments. Shaded area represents 95% CI.
SERGIO OOD: Noicy Interventions
CANAA
Returas (1)
SHD
AUPRC :13
Edge F1
Tital Syue
24.5
0.54
24.0
0.61
0.52
74.5
29.5
0.60
0.44
74.0
23.0
0.58
Intervention
Intervention
Intervattinn
Infervention Heration
Figure 19: Results of zero-shot OOD noisy gene knockouts with various intervention strategies 100 random SERGIO test environments. Shaded area represents 95% CI.
NeurIPS 2024 Author Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
All the claims made are justified in the method and experiments.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: The limitations are presented at the end of experiments section.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [NA]
Justification: [NA]
22

---

4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: All details provided in appendix
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
The code is provided along with supplementary material.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
See appendix D
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: All experiments are performed on 100 random cases and 95% Confidence intervals are provided.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: See appendix
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification: The research conducted in this paper conforms to the policy set out in NeurIPS Code of Ethics.
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes].
Justification:
See appendix
H
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: [NA]
12. Licenses for existing assets
23

---

Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes)
Justification: See appendix
13. New Assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: [NA]
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: [NA]
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: [NA]
24


Reviews:
[{"review_id": "M1wo6BQlR2", "reviewer": "Reviewer_rMit", "summary": "The authors propose an intervention design method, Causal Amortized Structure Learning, **CAASL**, which works by using a transformer to directly predict the next intervention to perform given a simulator of the environment (in which the intervention is to be performed) using the Soft Actor Critic (SAC) policy. The rewards are obtained by using the amortized posterior of the adjacency matrix of the causal graph given interventional data using AVICI, a recently proposed likelihood-free approach for obtaining an estimate of the causal graph. The authors demonstrate the amortization performance of CAASL on synthetic data and SERGIO (a single cell gene expression simulator) and its generalization to distribution shifts. Importantly, CAASL doesn't require inference of the causal graph.", "questions": "- just curious, why is it CAASL and not CASL?", "limitations": "The authors discuss some of the limitations of their work, but I think the limitations with respect to experimental results on changing intervention type, and types of simulation environments where this method is applicable can be clarified.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The paper tackles the important problem of designing an intervention policy in environments where the likelihood-based inference of a causal graph is intractable and it might be difficult to utilize the data likelihood to select interventions.\n- The use of the pretrained AVICI model as a reward function can potentially be more generally applicable--\n   - it formalises the reward of performing an intervention in the environment as the increase in the number of correct entries in the adjacency matrix $A$ of the causal graph resulting from the intervention-- and this is computable since we have access to the simulator, and hence $A$.\n   - this reward function is also shown to be kind of optimizing the lower bound on the multi-step expected information gain (EIG) by learning the adjacency matrix $A$ for the policy\n- The notation and writing of the paper is very clear, making it approachable.", "weaknesses": "Experiments:\n- How dense/sparse can $A$ be?\n- What are some other intervention design approaches and can they be used as baselines? \n- Can you provide some insights on why the random baselines is competitive when the intervention type changes on SERGIO?\n\nWriting:\n- on what kinds of datasets is CAASL applicable-- what are some other datasets where we can access $A$?\n- can you clearly establish the link between sequential Bayesian experimental design and the multi-step EIG-- where does the related BA bound appear in Bayesian experimental design (such that it doesn't involve a reading of ref.17)?\n- If the BA bound uses the log-likelihood of the true causal graph under the amortized posterior, how does comparing $A$ and $\\hat{A}$ make it still connected to the BA bound?\n- (minor): can you clarify what is meant by line 59 \"As such, CAASL is an intervention design method for performing sample efficient causal structure learning, but is not a new causal structure learning method in itself.\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks a lot for your detailed and very positive review of our work. We provide answers and clarifications to your questions below:\n\n- \"How dense/sparse can $A$ be?\" For training the policy, we set the density of graphs to 3 edges on average per node. Amortization can also be achieved across different density of graphs by simply sampling from a prior over graphs which is more denser/ sparser during training. \n\n- \"What are some other intervention design approaches and can they be used as baselines?\": For the synthetic design environment setting, any of the approaches listed in L119 can be potential baselines. However, we compare with diffCBED as it is the most recent and is shown to perform better than other possible baselines. For GRN environment, we are not aware of any baselines apart from random as it requires access to the likelihood.\n\n- \"Why is random strategy competitive when intervention type changes in sergio\"?: The input to the policy is the experimental history, which only contains the interventional sample and the corresponding {0,1}-vector interventional target. When the intervention type changes, the semantics of the data generating process also changes significantly, which the policy is unaware of since the input to the policy is still the same. Likewise, there are no specific inductive biases that are encoded in either the training procedure or the policy architecture that might result in a good generalization. We will clarify this limitation in the revision. In order to get better results in a different intervention type, we can amortize over different intervention types by training the policy on data generated by different intervention types. However, in this work, we deliberately train on a narrow domain to demonstrate the OOD generalization properties.\n\n- \"on what kinds of datasets is CAASL applicable-- what are some other datasets where we can access ?\" CAASL is applicable to any datasets/ settings that has a corresponding simulator. Some of the examples of causal structure learning simulators and applications are given in line 199. \n\n- “can you clearly establish the link between sequential Bayesian experimental design and the multi-step EIG-- where does the related BA bound appear in Bayesian experimental design (such that it doesn't involve a reading of ref.17)?”: The link between sequential Bayesian experimental design and related BA bound is discussed in Blau et al. Ref. [7].\n\n- \"can you clarify what is meant by line 59 \"As such, CAASL is an intervention design method for performing sample efficient causal structure learning, but is not a new causal structure learning method in itself.\" In this sentence, we mean to emphasize that this work focuses on the active learning/ intervention design part rather than focusing on how to do just causal structure learning given data.\n\n- \"If the BA bound uses the log-likelihood of the true causal graph under the amortized posterior, how does comparing  and make it still connected to the BA bound?\" As mentioned in line 230, although the proposed reward function is not exactly connected to BA bound, comparing the normalized accuracy, i.e. log-probs of the true graph would result in a BA bound.  \n\n- \"Why CAASL and not CASL\": We would like to emphasize both the active learning part and the amortization part, hence the acronym CAASL. \n\nWe hope that addresses your concerns and questions. We would be glad to answer any further questions you may have.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0YqZfBmWO8", "reviewer": "Reviewer_ydzf", "summary": "**Update after rebuttal:**\nI personally think that the authors have sufficiently addressed all criticism. Adding an extended discussion regarding the strong performance of a random baseline in many domains is a good idea and would situate the work even better. I remain in favor of accepting the paper, and have raised my confidence after taking into account all information.\n\nThe paper aims at the problem of designing efficient intervention strategies for causal induction. In causal induction, the goal is to infer an underlying structural causal model from observational and interventional data - and while this inference problem in itself is hard and a very active area of research, the problem of how to optimally choose interventions given a limited intervention-budget is of great importance in practice. Proposed solutions range from evolutionary and uncertainty-based active learning approaches to reinforcement learning, which is the method this paper pursues. The paper implements the intervention strategy via a (history-based) transformer that is trained via off-policy RL (via soft actor-critic) and is capable of incorporating some reasonable inductive biases via its architecture (permutation equivariance via attention). The trained policy is referred to as an amortized solution, since after training, a single forward pass determines the next intervention (without any parameter updates or RL). The paper also designs a novel reward function that is based on AVICI, a previously published transformer-based solution to obtain an amortized predictor (posterior) over the structural causal model, given a history of interventions and observations. Putting all of this together, yields a policy for interventional strategies that well aligns with practices in modern machine learning, and uses transformers’ great empirical generalization qualities to do the heavy lifting. Accordingly, out-of-distribution performance and generalization are a concern (since theoretical guarantees are scarce or vacuous) - but as the paper shows, the method generalizes remarkably well to non-trivial distributional shifts on synthetic data (10 variable linear SCMs with additive homoskedastic Gaussian noise for training) and data from a single-cell gene regulatory network simulator (which simulates realistic measurements of gene expressions including noise and significant amounts of missing values). On-distribution and mainly out-of-distribution experiments in both domains show promising performance and good generalization of the method.", "questions": "**Minor comments:**\n\n1. L108-109 (nit): L108 states that the posterior is trained “without having access to the likelihood of the data’, but L109 states that “the posterior can be trained by maximum likelihood with…”. Maybe rewrite to avoid possible confusion.\n2. L111: Maybe add that “Empirically it has been shown that the AVICI model can generalize to…”.\n3. L120: typo “they”.\n4. L123: What is “technical noise”? (maybe add a reference to 5.2 where this is explained)\n5.Zero-shot OOD (Fig. 4 and 5) - just to clarify: performance of the policy in this regime is only indirectly/implicitly reliant on good generalization of AVICI (since there is no RL at this stage; the reward model got “baked” into the amortized policy)? To improve OOD performance via additional learning it is plausible that fine-tuning the policy with fixed AVICI would be sufficient?", "limitations": "The paper clearly states limitations throughout and in a short limitations section. One general limitation (that is obvious to readers familiar with amortized algorithms) is that amortization typically only comes with in-distribution guarantees, and no ood-guarantees. Perhaps it is worth adding this to the limitations section (and clearly, the paper puts a lot of emphasis on empirically investigating this generalization).", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "* Very well written paper, with good introduction for non-domain experts, making the paper very accessible to a wide NeurIPS audience.\n* Careful combination of previously published approaches with a novel intervention policy parametrization (based on a transformer) and a novel formulation of a reward function (based on a previously proposed method for amortized causal inference). The focus is on forward-pass simplicity and efficiency, and using transformer-based architectures throughout to achieve good generalization of amortized solutions.\n* Experiments show good performance, and most of the attention is devoted to evaluating out-of-distribution performance under various OOD conditions, which greatly supports the generality and reliability of the results.", "weaknesses": "* Experimental ablations investigate the data-generating distribution / ground-truth, but not the network architecture and various design choices (such as, e.g., the network size, the use of tanh vs. sigmoid vs. ReLU, the sensitivity to the underlying AVICI model and its training regime, …). It would be nice to see some ablations to empirically justify the design choices, but I do consider the OOD experiments more important for publication and am pleased to see a fairly thorough investigation there.\n* The discussion around amortization and how it relies on generalization to “just work” could be expanded a bit to make the paper even more broadly accessible.\n* Though the paper is very clear about which contributions are novel compared to what is taken off the shelf, it may help to have all the novel contributions in one place (e.g., end of the intro).\n\n**Verdict:**\nThough I am not a domain expert on causal inference for gene regulation, I greatly enjoyed reading the paper - it is very well written. The great introduction and background section makes the paper easily accessible to a wide ML audience. The method itself is somewhat intricate and composed of a number of parts that each come with several design choices - the paper strikes a good balance between laying out these options and explaining the ones taken in the paper, and then separating the novel parts from previously published parts. In the end this section can get quite dense, and it might benefit from some additional summary or overview diagram at the end of Section 3 (though I really enjoyed Sec. 4.1, the paper would also work if it were pushed to the appendix if more space is needed). The experiments are sensible and of small to medium-sized scale, and I appreciate the use of a realism-focused simulator for gene-regulation networks instead of a larger-scale synthetic experiment. The focus on out-of-distribution performance by investigating several OOD settings is particularly important for amortized methods, and I think it is well executed in the paper and results show the good performance and generalizability of the method. As a non-domain expert I cannot comment on the strength of the baselines that the method is being compared against (and have lowered my confidence accordingly) - though they seem to be non-trivial and sensible baselines. Overall I think the paper is very well written, makes a novel and original contribution (that includes putting many existing techniques together), the empirical investigation is thorough and convincing, and the work is of interest to a wide audience. I currently recommend to accept the paper, and have some (mostly optional) suggestions for improvement.\n\n**Improvements:**\n1. Ablations for design choices: to support the design choices and get a sense of the sensitivity of the results some additional ablations would be interesting (I do not consider these to be critical for publication, but they would further strengthen the paper; the rebuttal period may be too short to conduct these):\n    1. How sensitive is the performance of the method to the underlying AVICI and the regime it has been trained on, ie., does AVICI first need to be carefully tuned for the application domain or can a relatively general AVICI model be used to easily obtain good results on many domains?\n    2. How sensitive are results in terms of the network architecture of the policy (network size, other activation functions)?\n2. Another interesting experiment could be to test how well the amortized policy generalizes when increasing the intervention budget beyond the maximal training range (does the policy continue to perform well / degrade gracefully, or does it eventually break down catastrophically?). Like above, I consider this to be optional, but interesting.\n3. To make the paper even more widely accessible consider having a brief paragraph in the background section on amortization. What is it? What are the advantages? What are the potential drawbacks (in-distribution vs. ood generalization “guarantees”).\n4. Maybe have a brief summary or overview figure, etc. at the end of Sec. 3 that summarizes the choices made in this paper (as opposed to the potential alternatives for the various parts which are pointed out nicely already) and identifies the novel parts and/or consider having a summary of contributions (e.g., at the end of the intro).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks a lot for the very positive feedback of our work. We provide detailed response to your questions below:\n\n- “Sensitivity to AVICI”: That’s a good point. In our work, we use a pre trained AVICI which has been trained on datasets drawn from wide range of SCM and noise distributions. So a relatively general AVICI model sufficed. But if amortization performance is more important than the OOD, having an AVICI model which is trained on a more narrow domain should further improve the results of the policy due to the more informative reward function.\n- “Increasing the intervention budget”: Thanks a lot for the great suggestion. The performance of the policy is based on the AVICI model. So as we give AVICI more interventional data, we expect performance to improve, but the returns become more incremental. But it is nevertheless interesting to think about increasing the intervention budget during training, especially to investigate until when the policy remains competitive as compared to random. We will include this experiment in the camera-ready.\n- Given the extra page for camera-ready, we will expand a bit on amortization, summary of contributions and also a schematic figure of the training setup (which was requested by other reviewer as well). \n- Typos and minor comments: Thanks a lot for the correction. We will correct these typos and sentences. We agree that it will help make the paper more readable.\n- “performance of the policy in this regime is only indirectly/implicitly reliant on good generalization of AVICI (since there is no RL at this stage; the reward model got “baked” into the amortized policy)?”Indeed, the evaluation still relies on a causal structure learning model, in which case we still use AVICI. As a result, the performance of the policy in OOD regime still implicitly relies on good generalization of AVICI. One could also consider the possibility of updating the AVICI model when updating the policy, but we only study the simpler setting in this work. \n- We will add the limitations about OOD generalization guarantees for amortized algorithms. \n\nWe hope that addresses your concerns and questions. We would be glad to answer any further questions you may have.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "M7yS7t7NFL", "reviewer": "Reviewer_o26x", "summary": "The paper proposes to use RL to train policies that can adaptively select interventions/experiments for obtaining data to do active causal structure learning. The approach amortizes policy training by training the policy on synthetically generated data, an approach that has shown some promise in the amortized training of causal structure learning algorithms. These pretrained causal structure learning inference models are used as the reward function in the training of the policy used for experiment design. The policy takes in a dataset and outputs an experiment (intervention on the casual model), and uses a network architecture that bakes in the relevant invariances/equivariances. Evaluation is done on synthetic environments including a gene regulatory network environment.", "questions": "- Section 5.1. What is the intervention? I got from the appendix that it is a do() intervention but for each method how do you select what value to set the intervention targets to in the case of Random and SS Finite? \n- Line 88: \"where the first column corresponds to one-hot encoding of whether a particular variable is intervened or not\". This made me believe you only deal with singleton intervention targets, but reading further it seems like the approach also handles and makes use of multi-target interventions. I think it is not 'one-hot' in your case, but confirmation from the authors would be appreciated.", "limitations": "The method requires inputting an entire dataset into the policy. It would be helpful to have experiments and discussion that reflect the ability for the approach to scale not just over graph size (included already) but over dataset size/number of rounds of experiment. Besides this I have no complaints about the transparency of limitations.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The approach seems reasonable and I think members of the NeurIPS community will find the experiments evaluating this approach interesting \n- The experiments in section 5.1, particularly the results on OOD performance are encouraging\n- The methodology is generally clearly presented\n- It is a novel and interesting combination of the ideas in [6], [36]", "weaknesses": "- Section 5.2. I think more information needs to be given to interpret the results, perhaps with additional plots showing a measure more easily interpretable than 'Returns'. My interpretation from the plots currently presented in the main paper in section 5.2 is that the proposed method technically improves over a random baseline, but by such a small amount that the improvement is practically insignificant\n- Section 5. It would be helpful to show experiments demonstrating that the approach actually converges to learning the true graph if enough experiments are done, since even simple approaches like random have this property. \n- I think more background needs to be given on what the baselines in the experiments are and how they are run. Both perform much worse than a random strategy, which is surprising. It also seems that both baselines are batch approaches (at each iteration select multiple distinct interventions), while the approach here is always using batch size 1. How are / do the baselines need to be modified for the batch-size 1 setting and why do they perform so poorly, even in-distribution in Figure 3? \n- Since greedy approaches are so common in this field, I think a natural ablation would have been to train the RL policy with an effective horizon of 1 (effectively a greedy/bandit policy trained under the same methodology presented). This can be done by rolling out the policy for the full horizon $H$, but then updating the policy by considering each timestep as it's own 'episode'. This would answer the very important question of whether this extra RL/adaptive machinery is really necessary compared to just looking one experiment ahead. I couldn't find what the horizon length used during training of the policy was. Apologies if I just simply missed this.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks a lot for the detailed feedback. We address your questions and concerns below:\n\n- “Section 5.2… More information needs to be given to interpret the results”: In Appendix, Fig 14-18 presents results with regards to metrics other than returns. We argue that returns is an interpretable metric, as it is an indicator of the test-time rewards of the policy for unseen environments. Regarding the practical significance of the performance of the approach wrt random, we would like to emphasize that our approach can handle unknown likelihoods and high amounts of missing data, which has not been studied in active causal structure learning before. The results that we have obtained can be further improved by also updating the reward model AVICI as more interventional data is acquired during training. We hope that through this contribution, further active structure learning methods aimed at practical scenarios can be developed. \n\n- “Section 5. It would be helpful to show experiments demonstrating that the approach actually converges to learning the true graph.” Thanks a lot for the suggestion. While this is a desirable property of any method to converge to the true graph when acquiring more interventional data, we would like to highlight that the convergence property is also influenced by the structure learning algorithm that is used to evaluate the policy. In our case, we use a pre trained AVICI model, which does not have necessarily the property that it converges to the true graph when given more data. Hence, even a random policy, when evaluated with AVICI also need not converge to the true graph. \n\n- “More background regarding baselines”: The chosen baselines are the state-of-the-art in Bayesian Optimal Experimental Design for causal discovery which can select multi-target interventions. For the intervention value, it is selected at random between the range [-1,1] for random strategy.  For SS-Finite, we set the intervention value to 5, as suggested in previous approaches like [52]. \n\n- \"Baselines are batch strategies and perform poorly\": Indeed, both of the baselines are batch strategies. We set the batch size to 1 in order to make a fair comparison. Most baselines work well when the underlying graphs are very sparse (for instance when the average edges per node is less than 1). However, in the settings we consider, the sparsity of graphs is set that it is more realistic in terms of real-world applications, wherein a random policy is still competitive. In addition, baselines still use an approximate posterior, thereby making the induced policy at times worse than random. We will make these details clear in the revision. \n\n- “Possibility of training a greedy policy”: Thanks for bringing this great point. In figure 2, one can see the adaptiveness of the policy, esp. wrt optimal intervention values, at play. We agree that if adaptiveness is not necessarily a requirement, one can just train the policy with a time horizon of 1. In such a setting, we expect that our approach still offers advantages in terms of amortization through design space symmetries. For the experiments, the horizon is always set to 10, both during training and evaluation. \n\n- “where the first column corresponds to one-hot encoding of whether a particular variable is intervened or not”: We apologize for the poor choice of wording. By one-hot encoding, we mean to say the corresponding vector is a {0,1} vector. We willl make this correction of terminology to make it clear that we select multi-target interventions. \n\n- “Scalability wrt dataset size and number of experiments”: With regards to dataset size, one can scale it depending on the compute availability, by performing data parallel training approach. With regards to scaling to number of experiments, our method might show limitations as any baseline RL algorithm does not perform well over long horizons. However, as noted to the other reviewer (ydzf), it is indeed interesting to look at OOD performance when the number of experiments during testing changes. We will include this experiment for the camera-ready in appendix.\n\nWe hope that addresses your concerns and questions. We would be glad to answer any further questions you may have.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3tF5DCgGAp", "reviewer": "Reviewer_EZYL", "summary": "The authors propose CAASL, which amortizes an intervention policy in the setting of causal structure learning. They apply this algorithm in the setting of a synthetic environment and a gene expression simulator. They use AVICI as a reward model to estimate an approximate posterior over the graph’s adjacency matrix. The RL algorithm is SAC, and the policy and value networks are transformers. On their two environments, CAASL outperforms several baselines (in terms of uncovering the true adjacency matrix) and performs better in out-of-distribution settings.", "questions": "Why does the random intervention policy outperform the other baselines, and why does CAASL revert to this performance as the dimensionality increases?\n\nHave previous works explored the environments used here? I.e., are these standard benchmarks?", "limitations": "Yes.", "rating": 5, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The authors combine several lines of work to tackle the setting of causal structure learning: soft actor-critic, AVICI, transformer policies. From what I can tell, this is a novel combination of these methods. These results seem reasonably significant, enabling fewer interventions to obtain comparable adjacency matrix estimation accuracy.\n\nThe paper itself is fairly clear in terms of writing and notation. The empirical analyses are clear and helpful, exploring in-distribution vs. out-of-distribution generalization and ablations. All plots are clearly labeled.", "weaknesses": "The authors put considerable effort into explaining the setup of causal structure learning, etc. However, perhaps due to my relative unfamiliarity with this area, the particular details of the setup were still unclear to me. Figure 1 attempts to present an overview of where their method fits into a causal structure learning pipeline, however, I would have instead (or in addition) preferred to see how CAASL links up with AVICI, the simulator, etc. From Figure 1, we only see the amortized intervention policy without any of the other machinery.\n\nAgain, this may be due to my unfamiliarity with this area, but the novelty of CAASL was unclear to me. For instance, how does AVICI estimate the intervention policy? Is this not an applicable baseline? Is the main benefit of CAASL to improve the efficiency of policy estimation over AVICI? What are the natural baselines to consider in comparison with CAASL? In the experiments, the authors present two other intervention baselines, but these do not perform well. Is the benefit, then, some type of accuracy/performance?\n\nThe experiments seemed slightly toyish, and the baseline comparisons did not ease my concerns. The most performant baseline is a random intervention policy. Are we to believe that the other baselines are actually *worse* than random? And are these problems simple enough to just use random search? Further, as the data dimensionality increases, CAASL is not able to outperform the random baseline. I would have, naively, expected the opposite, in which amortizing the intervention policy should yield an even larger improvement as the problem becomes harder.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks a lot for the detailed feedback. We answer your questions and concerns below:\n\n- \"preferred to see how CAASL links up with AVICI, the simulator, etc\": Thanks a lot for the great suggestion. For the camera-ready version, which provides an additional content page, we will include another figure which will highlight the individual components of the CAASL training procedure. We would like to clarify that AVICI is an amortized structure learning framework and does not itself do intervention design. We use AVICI to design rewards for our intervention design policy, which can be trained with a simulator. \n\n- \"Difference between AVICI and CAASL\": We would like to clarify that AVICI is an amortized causal structure learning algorithm and CAASL is an intervention design policy, and both are addressing different problems. In particular, AVICI predicts (distribution of) causal graphs given any dataset, and CAASL predicts the most informative intervention to perform given experimental history. We use the prediction of AVICI to design rewards for training the intervention policy. The benefit of our method over all other applicable baselines is that we can not only achieve amortized intervention design without parameter updates, we also do not require any intermediate posterior estimation. \n\n- \"The experiments seemed slightly toyish\": Compared to prior work, the experimental settings we consider are some of the more challenging ones when it comes to the problem of active causal structure learning. If there is any specific concrete experimental setting that you believe could improve our work, we are happy to run those experiments and include them.\n\n- \"The most performant baseline is a random strategy\": In active causal structure learning, random strategy is a very well regarded baseline due to its strong empirical performance. All the related works listed in the paper do compare with random, and its competitive performance can also be seen in prior works. We would like to clarify that the reason random baseline is competitive is not because the problems are simple, rather the opposite, that they are in fact hard. Other baselines like DiffCBED use an explicit posterior model to predict interventions. Getting an explicit causal graph posterior is a hard problem in itself, even for very simple settings, which in turn affects the performance of other baselines as compared to random. \n\n- \"as the data dimensionality increases, CAASL is not able to outperform the random baseline\": Note that the dimensionality increase experiments are OOD experiments, i.e. CAASL has never seen any data with that dimensionality before during training. Hence, more the dimensionality difference as compared to the dimensionality during training, the more is the distribution shift. With more distribution shift, CAASL is expected to slowly degrade in performance, as is expected of other ML/ RL methods.\n\n- “Have previous works explored the environments used here? I.e., are these standard benchmarks?“: Synthetic environments are very common across both causal structure learning and active causal structure learning. GRN environment SERGIO has been used in AVICI for structure learning before, but not in active causal structure learning. This is due to the fact that SERGIO is a very challenging environment for the existing baselines, and for most baselines it is not applicable due to no likelihood.\n\nWe hope that addresses your concerns and questions. We would be glad to answer any further questions you may have.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all the reviewers for the detailed feedback regarding our work. We are delighted to find that in general the reviewers find our work tackles an important problem (rMit), well-written (ydzf, rMit, EZYL), novel and interesting combination of existing ideas (o26x, EZYL) and provides clear empirical analysis (ydzf, EZYL).  

In addition, for the camera-ready,  we will incorporate the following changes as suggested by the reviewers:
1. Summarize the contributions at the end of the introduction. 
2. Include a schematic figure of the training of CAASL, which shows the network architecture, simulator and the AVICI model used for reward function.
3. Experiments with respect to generalization when the experiment budget is more than that during training. 
4. We will expand the limitations section to include the lack of theoretical guarantees wrt OOD generalization.
5. Expanded explanation of the baselines and their limitatiions. 

We have answered to each reviewer's questions and concerns that answers the above points in more detail. We would be happy to answer further questions if there are any.

Meta-review: The reviewers are all positive on the paper, with an average score of 6. They appreciate the generality of the approach and feel that it is well-motivated and a novel combination of ideas. They also applaud the inclusion of out-of-distribution evaluations, which shows the generalizability of the approach. However, there are several (not major) weaknesses that they agreed needed to be improved upon before publication; for example, they feel that further discussion on the good performance of the random baseline should be included. Since the authors promise to address these weaknesses in the camera-ready, I feel that this paper can be accepted as a poster.
Decision: Accept (poster)