A Causal Foundation Model for Structure and Outcome Prediction

Abstract: We introduce TabPFN-CFM, a causal foundation model that can handle multiple causal problems. TabPFN-CFM predicts both causal structure and outcomes from observational data, supports queries on all three levels of Pearl's Causal Hierarchy and uses known graph structure when available to improve predictions. TabPFN-CFM is trained on synthetic datasets, and generalises to real datasets, demonstrating improved performance over both structural and outcome prediction baselines.

• The paper introduces TabPFN-CFM which jointly estimates causal structure and predicts outcomes using a transformer-based architecture.
• It integrates ADMG representations and cross-attention mechanisms to efficiently handle observational, interventional, and counterfactual queries.
• Empirical benchmarks show superior performance in structure estimation and outcome prediction over traditional causal inference methods.

TabPFN-CFM: Joint Causal Structure and Outcome Prediction

Motivation and Problem Formulation

TabPFN-CFM addresses a central challenge in causal inference: reconstructing underlying causal structure and accurately predicting outcomes for multiple types of causal queries from observational data. The model operates in the context of structural causal models (SCMs), with both observed and unobserved variables. Observed variables are decomposed into covariates, binary treatment, and outcome; prior knowledge of the causal graph is optionally incorporated. The architecture is designed to answer observational, interventional, and counterfactual queries within Pearl's causal hierarchy, enabling practical prediction across a broad range of interventions and what-if scenarios.

Model Architecture

TabPFN-CFM leverages a transformer-based architecture extending Do-PFN and TabPFNv2, integrating explicit encoding and cross-attention over acyclic directed mixed graphs (ADMGs) to model confounders (bidirected edges). Separate embedding streams are constructed for fit data, prediction targets, and graph structure. Covariates and outcomes are positionally embedded, randomized per pass, and concatenated with their cumulative distribution functions (CDFs), aiding robustness to outliers and preserving information about variable distributions. The graph encoding encompasses adjacency, ancestral, and confounding matrices, exposing parent, child, ancestor, and descendant relationships directly to the model. Cross-attention and row/column-wise self-attention allow prediction queries to access fit-set and graph-relational information efficiently. Notably, the predicted graph structure is made independent of the prediction query, ensuring modularity and correctness in structure estimation.

Additional architecture optimizations—such as $\mathrm{relu}^2$ activations, layer normalization, Muon optimizer, increased parameter width, and QK normalization—have been adopted for improved convergence and training efficiency.

Figure 1: Ablation results showing sequential application of architectural changes and their impact on prediction and adjacency matrix loss.

Figure 2: Loss curves indicating that TabPFN-CFM achieves significantly lower prediction and adjacency matrix loss than baseline architectures during extended training.

Training Procedure

Synthetic SCMs are sampled from a diverse prior over random DAG/ADMG structures, nonlinearities, noise distributions, and node types. Observational datasets are generated for fit and prediction purposes, including explicit counterfactual samples by intervening on treatment values and re-simulating descendant nodes while maintaining fixed exogenous noise inputs. Training is supervised via cross-entropy losses for both outcome prediction and elementwise binary losses for structural matrices (adjacency, ancestral, confounding), with graph priors randomized for robustness.

Causal Query Handling

TabPFN-CFM is directly trained to estimate $$P(y^*|\mathbf{x}^*, \mathrm{do}(T=1-t^*), \mathcal{D}^{\mathrm{obs}}, \mathcal{G})$$ and $P(\mathcal{G}|\mathcal{D}^{\mathrm{obs}})$, thereby bypassing explicit Bayesian posterior inference over SCMs and facilitating amortized causal query answering. Joint training over all three causal query types induces broad causal reasoning ability and increased sample efficiency. Theoretical analysis in the appendix proves that including the true graph structure as input cannot worsen estimation and may yield strictly improved posterior accuracy—unless the observational data uniquely identifies the graph.

Evaluation: Synthetic and Real Data

Synthetic Toy Examples

Using linear and nonlinear SCMs, including instrument variable problems, TabPFN-CFM demonstrates high fidelity in both structure and outcome prediction. Model predictions closely track exact distributions for observational, interventional, and counterfactual queries—both with and without graph priors. The model robustly identifies causal edges and confounders.

Figure 3: IV example distributions predicted by TabPFN-CFM, showing strong alignment with exact solutions across observational, interventional, and counterfactual settings, without graph prior.

Figure 4: Distributions from the IV SEM with model access to the graph prior, evidencing improved accuracy.

Structural and Outcome Prediction Benchmarks

TabPFN-CFM is benchmarked against meta-learners (S-/T-/X-learner, DR-learner), Do-PFN, and structure learning baselines (AVICI, FCI, GES, LiNGAM, PC). On both in-distribution and out-of-distribution synthetic SCMs, TabPFN-CFM achieves:

• Lower mean squared error (MSE) for interventional and counterfactual outcome queries compared to baselines (see Table synth_ood_acc_short).
• Superior AUROC and accuracy in adjacency, ancestral, and confounding matrix estimation (see Table synth_ood_graph_short).
• Direct prediction of ancestral matrices outperforms estimation by Monte Carlo generation from adjacency, indicating architectural advantages in relation modeling.

Figure 5: Observational, interventional, and counterfactual distributions in nonlinear SEMs, demonstrating improved predictive accuracy with graph prior incorporation.

Figure 6: Shifts in predicted counterfactual distributions as the number of fit samples increases in nonlinear SEMs.

Real-World Dataset Evaluation

On real datasets (Amazon Sales; Law School Admissions), TabPFN-CFM matches or outperforms baselines in observational settings and delivers substantial gains for interventional and counterfactual predictions. In Law School Admissions, counterfactual reasoning substantially enhances model accuracy—a capability absent in all baseline learners. Structural prediction metrics (AUROC, accuracy) on real-world graphs also favor TabPFN-CFM over established structure learning algorithms.

Theoretical and Practical Implications

TabPFN-CFM exemplifies a unified, amortized approach to causal inference, reconstructing graph structure and answering arbitrary causal queries without explicit posterior sampling or separate meta-learner architectures. The adoption of ADMGs for confounder representation advances the capacity to model latent variable effects, crucial for real-world causal reasoning. The transformer-based architecture enables rapid, scalable inference once trained, and theoretical guarantees ensure optimal leveraging of available graph information.

Practically, TabPFN-CFM is positioned as a foundation model for tabular and causal tasks: structural learning, interventional analysis, and counterfactual evaluation. It mitigates the need for separate models per task and learns causal generalizations that transfer robustly to OOD settings.

Future Directions

Potential extensions include time series causal modeling, cyclic causal structures, integration with richer real-world datasets, and further architectural optimizations for scalability. The modular separation of structure and outcome prediction components suggests avenues for hierarchical causal modeling, multi-modal input extension, and application to broader scientific inference tasks.

Conclusion

TabPFN-CFM provides a general-purpose, efficient, and theoretically principled model for joint causal structure and outcome prediction across Pearl's causal hierarchy. Its empirical performance, architectural refinements, and theoretical guarantees mark a substantial step forward in amortized causal inference, offering scalable practical utility and strong foundation for further AI causal reasoning developments.