Amortized Active Causal Induction with Deep Reinforcement Learning

Thanks a lot for the detailed feedback. We answer your questions and concerns below:

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

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

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

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

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

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

We hope that addresses your concerns and questions. We would be glad to answer any further questions you may have.

Thanks a lot for the detailed feedback. We address your questions and concerns below:

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

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

• “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].

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

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

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

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

We hope that addresses your concerns and questions. We would be glad to answer any further questions you may have.

Dear reviewer,

Thanks a lot for the comments and the feedback.

I don't think this really addresses the point since the random strategy can also handle all these novelties. The performance compared to random on SERGIO seems like a very significant limitation that is glossed over in both the paper and the general Openreview discussion.

We would like to highlight that we perform only 10 experiments with 50 initial observational samples in this setting, which might also be less samples for the AVICI model (which is used for evaluation and calculating the returns) to show significant improvements. An alternative is to train a policy which is either a batch of designs, or as mentioned in the rebuttal, by increasing the number of experiments. Of course, it is practically harder to train a policy which has these properties with RL, which can be counted as a limitation. Another point to highlight is that the existing methods for (Bayesian) active causal structure learning evaluate exclusively on synthetic data, where we perform significantly better. Besides, all the results are provided with 95% confidence intervals over 100 random environments, including for the SERGIO environment. The improvements we see in SERGIO, though not as substantial as that for the synthetic environment, is still noteworthy given that no other active causal structure learning algorithm is applicable in such a setting. Hence, our contribution should be viewed as a step forward in the ongoing discussion about advancing active causal structure learning.

I think showing that an adaptive strategy is worth pursuing can only be done by showing that it is actually performing better than the ablated non-adaptive strategy (eg a bandit/setting time horizon to one). It doesn't really matter that the strategy proposed indeed ends up being adaptive, if it doesn't turn out that it is helping. This seems like a very fundamental experiment being missed when proposing an adaptive method.

Theoretically, as shown in Greenwald et al. 2019 [23] and Choo and Shiragur 2023 [13], performing interventions that are adaptive can be more sample efficient. Empirically, given the limited space, we focus our experiments on studying the amortization and generalization properties of the transformer based policy, which forms the core of our contribution. However, we are prepared to add the Bandit setting/ setting the time horizon to one for the revision. We have also provided the code, which is compatible to train the policy for time horizon one, and we will make it available on acceptance.

Thanks a lot for the very positive feedback of our work. We provide detailed response to your questions below:

• “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.
• “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.
• 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).
• 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.
• “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.
• We will add the limitations about OOD generalization guarantees for amortized algorithms.

We hope that addresses your concerns and questions. We would be glad to answer any further questions you may have.

Thanks a lot for your detailed and very positive review of our work. We provide answers and clarifications to your questions below:

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

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

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

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

• “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].

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

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

• "Why CAASL and not CASL": We would like to emphasize both the active learning part and the amortization part, hence the acronym CAASL.

We hope that addresses your concerns and questions. We would be glad to answer any further questions you may have.