Zero-Shot Learning of Causal Models

We thank the reviewer for their positive and insightful feedback! We appreciate your recognition of the soundness of our framework and the diversity of our experiments. We now address the concerns raised by the reviewer below.

Access to noise samples

Thank you for raising this point. We agree with the reviewer that most real-world problems do not provide such supervised signals. However, it is important to note that during inference, Cond-FiP does not require access to noise samples. Instead, it only needs the observations (and a predicted or true causal graph) to infer the functional mechanisms. This allows Cond-FiP to be applied to real-world datasets, as demonstrated in our experiment on the Sachs dataset (Appendix C).

An interesting extension of this work would be to explore a semi-supervised setting, where synthetic and real-world data are mixed during training. However, we believe this is outside the scope of the current paper.

Interventional results on Sachs

Regarding the lack of interventional generation results on the Sachs dataset, the main issue is that Cond-FiP (along with the other baselines considered in this work) only supports hard interventions, whereas the interventional data available for Sachs involves soft interventions (i.e., the exact interventional operations are unknown). As a result, we are unable to provide a comprehensive evaluation of Cond-FiP, or the other baselines, for interventional predictions on Sachs.

Additive Noise (ANM) Assumption

While the ANM assumption may be seen as a limitation, we would like to clarify that our method relies on the ANM assumption only for training the encoder. This is because we need the encoder to predict the noise from the data in order to obtain embeddings, which is simplified under the ANM assumption, as explained in Appendix A.2. However, it is important to emphasize that, while the ANM assumption is required for training the encoder, it is not necessary for training the decoder.

An interesting avenue for future work would be to explore a more general dataset encoding approach, potentially using self-supervised techniques. However, we believe this falls outside the scope of the current work.

Knowledge of causal graph

We agree with the reviewer that assuming knowledge of the true causal graphs is a strong assumption. However, as outlined in the manuscript (line 406), we can relax this assumption by inferring the causal graphs in a zero-shot manner via state-of-the-art prior works, such as AVICI. In Appendix D, we provide experimental results where we do not assume prior knowledge of the true graphs and infer them via AVICI. These results demonstrate that Cond-FiP can be extended to zero shot infer the full SCMs using only observational data.

Cond-FiP performance against baselines

First, it is important to note that all baselines use the true causal graphs for a fair comparison with Cond-FiP, and we employ their original implementations. Additionally, the baselines are trained on each test dataset, serving as the gold standard that our zero-shot approach aims to match. In contrast, although Cond-FiP is trained on specific scales, it generalizes well to both smaller and, more importantly, larger instance problems, while maintaining performance comparable to the other baselines. Furthermore, in scarce data regimes, Cond-FiP demonstrates superior generalization (Appendix E).

FiP implementation

We use the original code provided by the authors from their paper. In their implementation, the authors offer an adaptation of FiP when the causal graph is known. For a fair comparison, we evaluate Cond-FiP against this variant of FiP in our work.

Alternating application of transformer blocks E

The alternating block transformer is a feedforward neural network that takes as input a tensor of shape $(B,n,d,d_h)$ where $B$ is the batch size, $n$ is the number of samples, $d$ is the number of nodes and $d_h$ is the hidden dimension. In practice, we first permute the second and third dimensions before applying an attention mechanism to perform attention over the sample dimension. Afterward, we permute them back to apply an attention layer over the node dimension.

Counterfactual inference results

Thank you for raising this point. Cond-FiP is indeed capable of performing counterfactual generation, and we have conducted an experiment to evaluate this. The results can be found via the anonymous link. We observe that Cond-FiP performs slightly worse than the baselines in this task, and we believe that improving its performance in counterfactual generation will be a valuable direction for future work.

Thank you once again for your constructive comments! We are open to further discussion and would be happy to address any remaining concerns.

We thank the reviewer for their insightful feedback! We appreciate that they found our claim regarding Cond-FiP’s performance relative to state-of-the-art methods well justified. Additionally, we highlight our experiments in scarce data regimes (Appendix E), where Cond-FiP demonstrates superior generalization compared to baselines.

We now address the reviewer’s concerns below.

Diverse noise distributions

Thank you for raising this point! Following the reviewer’s recommendation, we experiment with a mixture of gaussian noise for the Large Backdoor and Weak Arrow datasets from CSuite. Specifically, the noise is sampled with equal probability from either $N(0,1)$ or $N(0,2)$. Results accessible via the anonymous link show that Cond-FiP is competitive with baselines for sample generation, and slightly worse on other tasks (still competitive with DECI). We emphasize that the baselines are trained from scratch specifically for the mixture of gaussian noise, while Cond-FiP has been pretrained only on gaussian noise.

We also want to clarify the noise distribution choices for the main experiments follow the prior works (Lorch et al. 2022, Scetbon et al. 2024). Finally, our ablation studies in Appendix F.3 also evaluate Cond-FiP’s performance under varying noise distribution complexity. We assess sensitivity to distribution shifts by adjusting noise parameters, controlling the shift magnitude. Results in Tables 25–27 show Cond-FiP’s OOD generalization drops (as expected) as the severity of the shift increases.

Assumption of known causal graphs

We agree that assuming knowledge of the true causal graphs is a strong assumption. However, standard causal discovery methods are not necessary for inferring causal graphs at inference time. As outlined in the manuscript (line 406), we can relax this assumption by inferring the causal graphs via state-of-the-art amortized causal discovery techniques, such as AVICI. In Appendix D, we provide experimental results where we do not assume prior knowledge of the true graphs and infer them via AVICI. These results demonstrate that Cond-FiP can be extended to infer full SCMs (without updating parameters) using only observational data.

Additive Noise Model (ANM) assumption

While the ANM assumption may be seen as a limitation, we would like to clarify that our method relies on the ANM assumption only for training the encoder. This is because we need the encoder to predict the noise from the data in order to obtain embeddings, which is simplified under the ANM assumption, as explained in Appendix A.2. However, it is important to emphasize that, while the ANM assumption is required for training the encoder, it is not necessary for training the decoder.

An interesting avenue for future work would be to explore a more general dataset encoding approach, potentially using self-supervised techniques. However, we believe this falls outside the scope of the current work.

Justification behind predicting noise variables

We agree with the reviewer that the map between the noise variables and causal variables must be invertible. In our setting, we adopt the ANM assumption, which ensures invertibility since the jacobian w.r.t noise is a triangular matrix with nonzero diagonal. Please check Appendix A.1 for more details. Notably, this does not require assuming a specific noise distribution.

Assumption of noise variables during training

Thank you for raising this point. We agree with the reviewer that most real-world problems do not provide such supervised signals. However, it is important to note that during inference, Cond-FiP does not require access to noise samples. Instead, it only needs the observations (and a predicted or true causal graph) to infer the functional mechanisms. This allows Cond-FiP to be applied to real-world datasets, as demonstrated in our experiment on the Sachs dataset (Appendix C).

An interesting extension of this work would be to explore a semi-supervised setting, where synthetic and real-world data are mixed during training. However, we believe this is outside the scope of the current paper.

Zero-shot terminology

We agree with the reviewer that the term "zero-shot" in the context of in-context learning and LLM literature typically refers to the number of examples a model sees. However, in our work, we adopt this terminology following the literature on amortized causal learning (Zhang et al. 2023, Nilforoshan et al. 2023, Gupta et al. 2023), where "zero-shot" refers to making predictions without updating the model parameters. We are open to adjusting our notation and adopting the terminology "amortized causal learning" if the reviewer prefers this.

Thank you once again for your constructive comments! We are open to further discussion and would be happy to address any remaining concerns. If you believe your concerns have been addressed, kindly increase your score.

We thank the reviewer for their positive and insightful feedback! Thank you for acknowledging the soundness of our framework and diverse experiments. We now address the concerns raised by the reviewer below.

The paper claims that the proposed method performs well in low-data scenarios. Table 11 (Appendix E) suggests that DoWhy often matches or outperforms the proposed approach. Could you provide some insights into why this happens?

We agree with the reviewer's observation that Cond-FiP performs comparably to DoWhy in Table 11 for the LIN IN and LIN OUT cases. However, we emphasize that Cond-FiP significantly outperforms DoWhy in the RFF IN and RFF OUT cases, particularly when the total number of nodes is 50 or 100. This strong performance in non-linear settings reinforces our claim that Cond-FiP exhibits superior generalization in scarce data regimes. Note that learning effective solutions with limited data is more challenging for the non-linear functional relationships (RFF IN/OUT) scenario. As a result, DoWhy struggles in these scenarios, while Cond-FiP demonstrates a clear advantage.

A key advantage of the proposed method is that it learns a single model instead of multiple models. How does this compare in terms of memory requirements, inference time, and computational efficiency?

Memory Requirements. We trained Cond-FiP on a single L40 GPU with 48GB of memory (see line 323 in the paper), using an effective batch size of $8$ with gradient accumulation. Below, we outline the detailed memory computation:

• Each batch consists of $n=400$ samples with dimension $d=20$, requiring less than $1$ MiB of data in FP32 precision.
• Storing the model on the GPU requires under 100 MiB.
• Our transformer architecture has 4 attention layers, a 256-dimensional embedding space, and a 512-dimensional feedforward network. Using a standard (non-flash) attention implementation, a forward pass consumes approximately 30 GiB of GPU memory.

Compared to the baselines, Cond-FiP has similar memory requirements to DECI and FiP, as all three train neural networks of comparable size. The main exception is DoWhy, which fits simpler models for each node, but this approach does not scale well as the graph size increases.

Computational Cost. Like other amortized approaches, Cond-FiP has a higher training cost than the baselines, as it is trained across multiple datasets. While the cost of each forward-pass is comparable to FiP, we trained Cond-FiP over approximately 4M datasets in an amortized manner. However, Cond-FiP offers a significant advantage at inference time since it requires only a single forward pass to generate predictions, whereas the baselines must be retrained from scratch for each new dataset. Thus, while Cond-FiP incurs a higher one-time training cost, its substantially faster at inference.

Given that the primary evidence supporting the claims of the paper comes from experiments, are there any plans to open-source the code to facilitate replicability and transparency?

Thanks for this point! We plan to open-source the code along with comprehensive documentation to facilitate reproducibility of our experiments. For the rebuttal phase, we have prepared an anonymized version of the codebase, which can be accessed via this link.

Please note that while the codebase is not directly executable, it provides full access to the implementation of all components of our framework:

• cond_fip/models contains the implementation of the transformer-based encoder and the Cond-FIP architecture.
• cond_fip/tasks includes the training and inference methods associated with our framework.

The paper demonstrates an average level of originality as it builds on pre-existing ideas like the transformer architecture for SCM inference and FiP.

We agree with the reviewer that our framework builds upon prior works, specifically AVICI (Lorch et al. 2022) and FiP (Scetbon et al. 2024). However, we want to clarify that our main contribution relies on integrating these two frameworks to enable zero-shot inference of generative SCMs, a problem not previously addressed. To achieve this, we made substantial modifications to the FiP architecture, as it needed to be conditioned on datasets. Additionally, to facilitate this dataset conditioning we propose to learn dataset embeddings via the noise prediction task.

While most existing studies on amortized causal learning focus on treatment effect estimation or causal discovery (as discussed in Section 2), our work tackles the novel task of amortized learning to infer the causal mechanisms of SCMs.

Thanks again for your constructive comments! We are open to further discussion and would be happy to address any remaining concerns.