A Causal Foundation Model for Structure and Outcome Prediction

Determining causal relationships and predicting the outcomes of interventions from observational data remains a significant challenge in machine learning. Existing methods often focus on either structural discovery or outcome prediction separately, and many struggle to account for unobserved confounding variables. This paper introduces TabPFN-CFM, a causal foundation model designed to handle both tasks simultaneously while supporting observational, interventional, and counterfactual queries.

TabPFN-CFM extends the Prior-Data Fitted Network (PFN) framework to causal inference. The model is trained on a vast array of synthetic structural causal models (SCMs) with varying graph structures, noise distributions, and nonlinearities. By training on these synthetic datasets, the model learns to perform amortized Bayesian inference. The architecture uses a transformer-based encoder that processes observational data, optional prior graph knowledge, and query types to predict both the causal graph (represented as an Acyclic Directed Mixed Graph to account for unobserved confounding) and the outcome distribution. The training procedure is optimized for efficiency, resulting in a 3-4x speedup compared to previous iterations.

TabPFN-CFM demonstrates strong performance across both synthetic and real-world datasets. In synthetic experiments, it outperforms established baselines in interventional and counterfactual prediction tasks. Notably, the model effectively leverages known graph structure when available to improve prediction accuracy. On real-world benchmarks, such as the Law School Admissions dataset, the model shows a significant advantage in counterfactual reasoning. The authors also show that directly predicting the ancestral matrix of the causal graph yields higher accuracy than deriving it from the adjacency matrix, highlighting the benefits of the model's specialized structural decoder.