Demystifying amortized causal discovery with transformers

TL;DR

This paper analyzes how transformer-based amortized causal discovery methods implicitly rely on priors and identifiability, revealing their limitations in generalizing to unseen causal models and aligning with classical theory.

Contribution

It bridges the gap between amortized causal discovery and identifiability theory, and analyzes conditions under which training on multiple models improves generalization.

Findings

Training distribution defines a prior on causal models.

CSIvA cannot generalize to unseen causal classes without multiple models.

Amortized methods still adhere to identifiability constraints.

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., 2023) 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 theoretically and empirically analyze \textit{when} training CSIvA on datasets generated by multiple identifiable causal models with different structural assumptions improves its generalization at test time. Overall, we find that amortized causal discovery with transformers still adheres to identifiability theory, violating the previous hypothesis from Lopez-Paz et al. (2015) that supervised learning methods could overcome its restrictions.