Demystifying amortized causal discovery with transformers

2024-05-27

arXiv.org

AI Summary

Transformer-based supervised causal discovery methods (e.g., CSIvA) trained on synthetic data exhibit implicit reliance on distributional priors, undermining their transferability to real-world data—despite empirical success, they remain subject to theoretical identifiability constraints. Method: We propose a unified framework integrating amortized causal discovery with classical identifiability theory, introducing a hybrid identifiable causal generative model that relaxes stringent assumptions on structural mechanisms while improving robustness to diverse noise types. Contribution/Results: We prove that, under the hybrid model, the probability of causal ambiguity asymptotically vanishes. Empirically, training across heterogeneous synthetic model classes significantly boosts causal graph prediction accuracy on real-world data, with markedly improved generalization stability compared to baseline approaches.

Technology Category

• Machine Learning: Causal Learning
• Reasoning under Uncertainty: Causality
• Constraint Satisfaction and Optimization: Satisfiability Modulo Theories

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Supervised learning approaches for causal discovery from observational data often achieve competitive performance despite seemingly avoiding explicit assumptions that traditional methods make for identifiability. In this work, we investigate CSIvA (Ke et al., 2023), a transformer-based model promising to train on synthetic data and transfer to real data. First, we bridge the gap with existing identifiability theory and show that constraints on the training data distribution implicitly define a prior on 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. At the same time, we find new trade-offs. Training on datasets generated from different classes of causal models, unambiguously identifiable in isolation, improves the test generalization. Performance is still guaranteed, as the ambiguous cases resulting from the mixture of identifiable causal models are unlikely to occur (which we formally prove). Overall, our study finds that amortized causal discovery still needs to obey identifiability theory, but it also differs from classical methods in how the assumptions are formulated, trading more reliance on assumptions on the noise type for fewer hypotheses on the mechanisms.

Problem

• Investigates transformer-based causal discovery model CSIvA
• Bridges gap between supervised learning and identifiability theory
• Explores trade-offs in training with mixed identifiable causal models

Innovation

• Transformer-based model for causal discovery
• Training on synthetic data transfers to real data
• Trade noise assumptions for mechanism hypotheses