Foundation Models for Causal Inference via Prior-Data Fitted Networks

2025-06-12

AI Summary

This work addresses the challenge of robust Bayesian estimation of causal effects—such as conditional average treatment effects (CATE)—under diverse identification strategies (e.g., backdoor, frontdoor, and instrumental variable settings). We propose CausalFM, the first foundational model framework for Bayesian causal inference. It grounds validity priors in structural causal models (SCMs), enabling formal verification, and systematically introduces prior-data fitting networks (PFNs) to causal inference—a novel application. We further design a family of causally informed Bayesian neural network priors, supported by SCM-guided prior construction and theoretical validation. Built upon a Transformer architecture, CausalFM leverages in-context learning and synthetic-data pretraining. Evaluated on multiple synthetic and semi-synthetic benchmarks, it achieves state-of-the-art CATE estimation accuracy while establishing a reusable foundational model training paradigm for causal inference.

Technology Category

• Machine Learning: Causal Learning
• Reasoning under Uncertainty: Causality
• Cognitive Modeling & Cognitive Systems: Conceptual Inference and Reasoning

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Prior-data fitted networks (PFNs) have recently been proposed as a promising way to train tabular foundation models. PFNs are transformers that are pre-trained on synthetic data generated from a prespecified prior distribution and that enable Bayesian inference through in-context learning. In this paper, we introduce CausalFM, a comprehensive framework for training PFN-based foundation models in various causal inference settings. First, we formalize the construction of Bayesian priors for causal inference based on structural causal models (SCMs) in a principled way and derive necessary criteria for the validity of such priors. Building on this, we propose a novel family of prior distributions using causality-inspired Bayesian neural networks that enable CausalFM to perform Bayesian causal inference in various settings, including back-door, front-door, and instrumental variable adjustment. Finally, we instantiate CausalFM and explicitly train a foundation model for estimating conditional average treatment effects (CATEs) using back-door adjustment. We show that CausalFM performs competitively for CATE estimation using various synthetic and semi-synthetic benchmarks. In sum, our framework can be used as a general recipe to train foundation models for various causal inference settings. In contrast to the current state-of-the-art in causal inference, CausalFM offers a novel paradigm with the potential to fundamentally change how practitioners perform causal inference in medicine, economics, and other disciplines.

Problem

• Develops PFN-based foundation models for causal inference
• Formalizes Bayesian priors using structural causal models
• Enables Bayesian causal inference in diverse adjustment settings

Innovation

• PFNs enable Bayesian inference via transformers
• CausalFM uses causality-inspired Bayesian neural networks
• Trains foundation models for diverse causal settings