Amortized Inference of Causal Models via Conditional Fixed-Point Iterations

TMLR Paper6059 Authors

Abstract: Structural Causal Models (SCMs) offer a principled framework to reason about interventions and support out-of-distribution generalization, which are key goals in scientific discovery. However, the task of learning SCMs from observed data poses formidable challenges, and often requires training a separate model for each dataset. In this work, we propose an amortized inference framework that trains a single model to predict the causal mechanisms of SCMs conditioned on their observational data and causal graph. We first use a transformer-based architecture for amortized learning of dataset embeddings, and then extend the Fixed-Point Approach (FiP) to infer the causal mechanisms conditionally on their dataset embeddings. As a byproduct, our method can generate observational and interventional data from novel SCMs at inference time, without updating parameters. Empirical results show that our amortized procedure performs on par with baselines trained specifically for each dataset on both in and out-of-distribution problems, and also outperforms them in scare data regimes.

Submission Length: Regular submission (no more than 12 pages of main content)

Changes Since Last Submission:

(Rebuttal Revision 2)
- Added experiments for analyzing the effect of pretraining scale on Cond-FiP's performance in Appendix H

(Rebuttal Revision 1)
- Added discussion on ANM assumption and potential extension at end of Section 3.2
- Expanded real world experiments using the ecoli dataset from the bnlearn repository (Table 2)
- Moved the discussion on computational complexity advantages from appendix to end of Section 4
- Added pseduo code for the dataset encoder and cond-fip decoder in Appendix B.5
- Added additional experiments in Appendix G to analyze the robustness of cond-fip to errors in input causal graph

Assigned Action Editor: ~Arash_Mehrjou1

Submission Number: 6059