Demystifying amortized causal discovery with transformers

Authors: Francesco Montagna, Max Cairney-Leeming, Dhanya Sridhar, Francesco Locatello
Primary area: causal reasoning
Keywords: causal discovery, amortized inference, transformers, identifiability

Abstract:
Supervised learning for causal discovery from observational data often achieves competitive performance despite seemingly avoiding the explicit assumptions that traditional methods require for identifiability. In this work, we analyze CSIvA (Ke et al., 2023b) on bivariate causal models, a transformer architecture for amortized inference promising to train on synthetic data and transfer to real ones. First, we bridge the gap with identifiability theory, showing that the training distribution implicitly defines a prior on the causal model of the test observations: consistent with classical approaches, good performance is achieved when we have a good prior on the test data, and the underlying model is identifiable. Second, we find that CSIvA can not generalize to classes of causal models unseen during training: to overcome this limitation, we show that learning on datasets generated from different types of causal models, unambiguously identifiable in isolation, improves the test generalization. We analyze this empirical evidence with theory, illustrating that the ambiguous cases resulting from the mixture of identifiable causal models are unlikely to occur. Overall, we find that amortized causal discovery still adheres to identifiability theory, violating the previous hypothesis from Lopez-Paz et al. (2015) that supervised learning methods could overcome its restrictions.

Decision: reject
Reviewer scores: 6, 3, 5, 6

Reviews and discussion:
<thread>
creation date:
1729260582329

thread_type:
official_review

<thread_object>
creator: reviewer_vygz

<content>
summary:
The paper presents a comprehensive analysis of CSIvA in the context of bivariate causal models. The authors first demonstrate that while CSIvA generalizes well in distribution, it fails to generalize out-of-distribution, regardless of changes in mechanisms or noise distributions. By hypothesizing that CSIvA can identify the SCM underlying the data generation process, the authors challenge the previous claim that “supervised learning methods can address complex data generation and reduce the need for explicit assumptions.” Additionally, the paper takes an empirical step forward by showing that CSIvA has the potential to capture multiple identifiable causal structures within a single network, supported by theoretical claims that post-ANM models are generally identifiable.

strengths:
The paper is well-organized and clearly presented. The authors articulate their points effectively, with implications interwoven throughout the sections and the corresponding background provided in the appendix. The experiments on synthetic datasets are comprehensive, demonstrating that the performance of CSIvA is non-trivial. The empirical results support the proposed hypothesis. The experiment demonstrates that CSIvA’s generalization ability can be significantly enhanced by incorporating multiple distributions into the training dataset. 

This paper potentially addresses the limitations of supervised learning in causal discovery, particularly the challenge of applying pre-trained models to datasets that may be out-of-distribution. Specifically,  the highlight of this paper is that it offers a possible way to make supervised learning-based causal discovery practical: if the training dataset is sufficiently comprehensive, the model’s predictions are somewhat reliable.

weaknesses:
In my view, the largest weakness is the lack of evaluation on real-world datasets. Since the experiments already demonstrate the potential of achieving a relatively robust predictor by incorporating multiple distributions in the training dataset, an investigation and comparison with traditional methods on any real-world datasets would be valuable and appreciated. For example, https://www.cmu.edu/dietrich/causality/projects/causal_learn_benchmarks/. This is just an example, the author does not need to be constrained by this and can use any dataset that is valid for causal discovery. If there are any reasons why experiments on real-world datasets may not be feasible or practical, I would be very interested in hearing them and would welcome the explanation with an open mind.



Another limitation, which is also discussed by the authors, is the scope of analysis is rather limited in a bivariate model without considering intervention.

questions:
I would appreciate clarification on how the method handles independent variable pairs. Does it require a training dataset of independent variables, or is it primarily designed to identify a DAG from the Markov equivalent class, meaning independent cases don't need to be considered? Addressing this point explicitly could be helpful for readers with varying levels of expertise in the field.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730720120940

thread_type:
official_review

<thread_object>
creator: reviewer_qpwy

<content>
summary:
This work explores supervised learning for causal discovery, a promising direction to potentially replace traditional score-based or constraint-based methods. However, many gaps remain unclear when considering the supervised learning paradigm for this task. To clarify these gaps, we conduct extensive experiments using a specialized method, CSIvA, on a variety of synthetic datasets. The results of these experiments are presented, along with a detailed discussion of their implications.

strengths:
Some insights from the experiments are provided.


1) For in-distribution generalization, CSIvA effectively generalizes to unseen samples from the training distribution.

2) For out-of-distribution generalization, it is unsurprising that CSIvA performs poorly.

3) CSIvA's performance depends on whether the data is sampled from an identifiable SCM.

4) Training CSIvA on mixtures of SCMs leads to improved generalization.

weaknesses:
I am generally open to, encouraging, and expecting new paradigms that go beyond traditional score-based or constraint-based methods for causal discovery. However, this work raises several key concerns, prompting me to review this draft with caution.

1) As an experimental paper, it is expected to provide a wide range of existing methods applied to various datasets, not just simulated data. However, the experimental settings in this work are relatively simple and do not appear to be robust enough for a comprehensive experimental study.

2) The insights from the experiments do not yield particularly exciting results and seem somewhat predictable. I would expect a deeper discussion of these experiments or even some theoretical progress. For example, it would be helpful to explore why training data sampled from identifiable SCMs performs well or how the learned model (e.g., a transformer) leverages this information. Additionally, while the learned model may not explicitly satisfy identifiable conditions (such as those in linear non-Gaussian additive noise models), it would be interesting to investigate whether the model implicitly learns such information.

3) For data sampled from a non-identifiable SCM, there may be multiple graphs within a single Markov equivalence class. In this case, do supervised learning paradigms truly offer no advantage? For instance, could we apply a Bayesian approach to the model to yield the most probable result? Additionally, given the existence of multiple possible graphs, could we frame this as a one-to-many problem, perhaps using techniques like Mixture Density Networks?

Overall, although the insights from this work are not surprising, it is still interesting and valuable, and I believe it contributes to the community. However, the current version is not yet ready for publication. I encourage the authors to incorporate more robust experimental results, provide deeper insights, and offer rigorous justifications.

questions:
See above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730722375597

thread_type:
official_review

<thread_object>
creator: reviewer_anxc

<content>
summary:
The paper tries to analysis the CSIvA model on bivariate causal models. The authors proposed an hypothesis to attempt to bridge the gap between model free (i.e. end-to-end learning without given knowledge of causal generating process) causal discovery algorithms and the identifiability theory, but only use some examples, without a proper proof.

strengths:
The problem that the authors are trying to attack is actually important. Identifiability is central to causal discovery because it determines whether the true causal structure can be uniquely recovered from the data. Without identifiability, any conclusions about causality are inherently uncertain, making the validation and application of causal models problematic. Meanwhile, models such as CSIvA does not provide any identifiability results, but achieves good performance in practice. Thus to ensure the reliability of such models, identifiability results would be required.

weaknesses:
1. The hypothesis in the paper is not proved.
2. Actually, for a causal discovery algorithm, if it is differentiable at each steps, then it would be trivial to use a RNN to mimic  the action of the algorithm exactly. Then a large enough transformer can be used to fit the RNN easily. Thus it would be trivial that a transformer would perform good in identifiable causal discovery cases if such algorithm exists, For example, for linear gaussian with equal variance, each steps of NOTEARS is differentiable if the objective is differentiable. Thus we can use a differentiable neural network to mimic NOTEARS. In other cases, such explanation may also holds.

3. The experimental on 2-variables dataset are too simple.

questions:
See above weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730722982242

thread_type:
official_review

<thread_object>
creator: reviewer_ed52

<content>
summary:
This paper analyzes amortized causal discovery, where synthetic data allows supervised learning of causal discovery. The investigation focuses on whether learning-based methods can overcome traditional identifiability constraints in causal discovery. Extensive out-of-distribution experiments in a bivariate setting with a CSIvA transformer demonstrate that the method cannot generalize well to unseen mechanisms or noise families. This finding establishes an important link to traditional causal discovery literature, revealing how identifiability assumptions are transferred from algorithm design to training data selection. The work further demonstrates empirically that model generalization can be significantly improved by training on diverse families of mechanisms and noise distributions. A theoretical justification for this approach is provided by proving that the set of non-identifiable structures arising from such unions is negligible, particularly extending known results to post-additive noise models.

strengths:
- Establishes a clear connection between how assumptions in classical causal discovery are translated to training data assumptions in amortized approaches.
- Provides strong empirical validation and theoretical justification of generalization in amortized learning improves with mixed training, something that previously has not been studied to this extent, providing an important result and discussion point that will benefit the causal discovery community. 
- Establishes a testing methodology for amortized causal discovery that if adopted stands to benefit the field in improving the methods.

weaknesses:
- Limited Scope:
  - Analysis restricted to bivariate cases
  - Not clear how well the results extend to more complex graphs, nor is there a clear path for extending results to larger graphs
  - Doesn't address practical challenges like confounding or selection bias
- Experimental Limitations:
  - Only evaluated on purely synthetic data
  - Limited exploration of how results scale with dataset size

questions:
- Have you explored how your results might extend to larger graphs? What are the main challenges you foresee?
- Can you provide more concrete guidance on choosing optimal ratios when mixing different types of causal models during training?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732266685980

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Common response

comment:
We thank the reviewers for their valuable feedback and the attention dedicated to our paper. Our work is unanimously recognized as *addressing* *an important problem* (AnXC), and praised for its *clear presentation* (vYGz, Ed52, qPWy).  Almost all reviewers acknowledge the quality of our study: our findings are supported by *theoretical justification and strong empirical validation* (Ed52, vYGz), and the experiments are *extensive with detailed discussion of their implications* (qPWy). Our work is a *valuable* *contribution to the community* (qPWy, Ed52) and offers *a way to make amortized causal discovery practical* (vYGz).

In the common response, we address concerns shared by multiple reviewers:

1. The limit of experiments to bivariate graphs (Ed52, vYGz, AnXC). 
2. The absence of real-world experiments (Ed52, vYGz, qPWy).

**The new experiments we discuss are presented in the updated manuscript pdf (Appendices C and D), where novel parts are colored in red**. The manuscript bibliography lists the papers cited in this response.

---

**Limit to bivariate experiments**. We agree with the reviewers on this point, also discussed in the paper (L207). We outline theoretical and empirical reasons for focusing on two variables.

1. **Theoretical perspective**: 
    - The **bivariate setting is the privileged ground for analyzing identifiability**: identifiability of additive noise and post-nonlinear models is established on bivariate graphs in the seminal works of Hoyer et al., 2008 and  Zhang & Hyvärinen, 2009, and extended to multi-variate graphs with a combinatorial argument, see next.
    - **Multivariate guarantees build on these bivariate results**: identifiability for multivariate models is ensured by iteratively verifying the identifiability of all bivariate subgraphs (Peters et al., 2014, Tu et al., 2022). Thus, generalizing our findings to multivariate models is about empirically verifying that transformers align with theoretical identifiability for each pair of nodes connected in the graph. For this, it suffices to know that transformers align with identifiability theory in the bivariate case, a fact that we show. This theoretical argument suggests that the empirical **identifiability guarantees we find similarly extend to multivariate graphs**. This discussion is now part of the updated pdf, Appendix C, and addresses **Ed52's query** on how our results would extend to more dimensions.
2. **Empirical perspective**:
    - Unfortunately, we find that **CSIvA performs poorly on multivariate graphs in purely observational settings.** This is also found in the original paper by Ke et al., 2023, Appendix A.7.4. Investigating OOD generalization or identifiability via experiments requires convergence during the training and reliable in-distribution performance, so that test errors can be imputed to distribution shifts or non-identifiability. In the multivariate setting, CSIvA lacks these properties.
    - **Training on multivariate graphs is computationally intensive:** training on 5-node graphs takes ~2 days to converge.  Without access to Ke et al.'s code (not released) we can’t compare our implementations. Yet, our algorithm converges in approximately the same number of iterations as theirs, taking more time per iteration: differences may stem from hardware.

**Summary remark.** CSIvA has limitations in scaling to high dimensions, but scalability is not our focus. Amortized causal discovery has **two components**: one is that of scalability to large graphs, and the other is relative to the **theoretical viability of the idea**. Despite the current limitations that we observe with CSIvA, transformers clearly could scale structure learning to large graphs (a point shared e.g. by AnXC, point two of the weaknesses). On the other hand, the causal discovery community did not have evidence that this idea is theoretically sound. We addressed this aspect conclusively, additionally finding surprising trade-offs (e.g., mixed training, assumptions on noise vs mechanisms), that differ from classical algorithms and are worthy of further research.

---

**Absence of real-world experiments and other benchmarks** (Appendix D.1 of the updated pdf)**.** We agree that real-world experiments could strengthen our conclusions. Given the challenges of scaling CSIvA, we present results on real-world pairwise data from the CMU benchmark.  We find that our **mixed training strategy on CSIvA improves SHD on real data**, consistent with our prior synthetic results. As suggested by vYGz and qPWy we **compare CSIvA** **with** **well-established causal discovery methods**: DirectLiNGAM, CAM (Buhlmann et al., 2014), NoGAM (Montagna et al., 2023, reaching SOTA accuracy on several synthetic benchmarks), and GraNDAG (Lachapelle et al., 2019). Interestingly, CSIvA with mixed training consistently matches with or outperforms these benchmarks.

We thank vYGz for pointing us to the CMU benchmark for the real-data analysis.
</content>
</thread_object>

</thread>


Metareview:
metareview:
While the manuscript targets an important question that why amortized causal discovery with transformers works, major concerns raised by the reviewers—particularly regarding the limited scope of experiments (mostly bivariate) and the lack of theoretical claims—still remain. 

Two reviewers were marginally positive but did not fully endorse the paper, while two others maintained their skepticism. Given the discussion, the additional experiments, and the persistent reservations, the balance of arguments weighs slightly against acceptance at this time. The paper provides a valuable perspective and represents a step forward in understanding amortized causal discovery, but it may require further theoretical development or more robust multivariate evidence to be fully convincing. Therefore, the final recommendation is to reject the paper in its current form.