Source-linked AI summary

Causal Foundation Models

Christopher Stith, Hossein Rahmani, Jesse C. Cresswell

arXiv:2609.03003v1cs.LGstat.ML

TL;DR

Traditional causal inference requires a new mechanism, estimator, and training pipeline for each problem, motivating pretrained causal foundation models. This work explains how CFMs use synthetic-prior pretraining and in-context learning, reviews their competitive performance and speed, and identifies current scope limitations.

  • Problem

    Traditional causal inference requires repeating data analysis, causal-mechanism selection, estimator tuning, and training for each new problem.

  • Method

    CFMs are pretrained on synthetic causal priors and use transformer-based in-context learning to estimate causal quantities on new observational datasets without model updates.

  • Results

    CFMs achieve competitive causal-benchmark performance, while all three tested models run 1–2 orders of magnitude faster than traditional training, tuning, and inference on CPU.

  • Takeaways & Limitations

    CFMs can reduce deployment time and provide comparable-quality predictions with amortized inference effort, including speedups of up to 100× versus a highly optimized T-Learner.

  • Takeaways & Limitations

    First-generation CFMs cannot condition on additional domain knowledge such as partial knowledge of the causal structure.

Abstract

from arXiv · show

Causal inference is the practice of estimating the effect of a treatment or intervention from data. It traditionally requires a bespoke pipeline for every new problem: first proposing a causal mechanism, selecting a compatible estimator, and finally training it. Meanwhile, across diverse settings and modalities, much of machine learning has shifted to the paradigm of foundation models: networks pretrained once at scale and applied to new tasks without fine-tuning. Causal foundation models (CFMs) bring this paradigm to causal inference. CFMs are pretrained neural networks that estimate causal quantities, such as the average treatment effect, on entirely new datasets using in-context learning without requiring model updates. This work provides a practical introduction to this emerging area. We summarize the necessary background in causal inference and machine learning before discussing CFMs. Throughout, we include example code and Jupyter notebooks.

1 Introduction

Causal inference estimates causal effects while avoiding confusion between association and causation, but traditional methods require a bespoke pipeline for each new problem. Causal foundation models instead use pretrained networks and in-context learning to estimate effects on unseen datasets, and this work introduces their methods and practical use.

  • Causal inference estimates causal effects between variables across domains including policy, marketing, and medicine.
  • Traditional causal workflows require studying the data, proposing a causal mechanism, selecting and tuning an estimator, and retraining for each new problem.
  • Causal foundation models are pretrained neural networks that estimate causal effects in unseen settings without further training or fine-tuning.
  • At inference time, CFMs use labeled examples as context to perform amortized Bayesian inference and predict causal effects on new query data.
  • This work provides a practical introduction to CFMs with example code, Jupyter notebooks, design comparisons, benchmarks, and discussion of developments and applications.

2 Background

Causal inference distinguishes predictive relationships from intervention effects using potential outcomes, causal estimands, and structural assumptions. Observational data reveal only factual outcomes, so identification depends on assumptions such as consistency, ignorability, positivity, and SUTVA.

  • Why causal inference?: Predictive models estimate outcomes as observed, whereas causal models estimate outcomes under intervention and therefore require reasoning about the data-generating mechanism.
  • Potential outcomes: The potential-outcomes framework defines Y(t) as the outcome that would occur under treatment t, but only one treatment outcome is observed for each individual.
  • Causal estimands: Causal estimands summarize potential outcomes through quantities including CEPO, ATE, CATE, and individual treatment-response curves.
  • Observed and counterfactual data: Observational data contain covariates, treatment, and factual outcome, while counterfactual potential outcomes remain unobserved.
  • Interventions: The do-operator represents setting treatment by intervention, whose distribution can differ from the observational distribution conditional on treatment.
  • Identifiability: In the backdoor setting, consistency, ignorability, and positivity support identification of causal quantities from observational data, but ignorability is untestable in practice.

3 Causal Foundation Models

Causal foundation models apply prior-data fitted networks to causal inference, using in-context learning to estimate causal quantities on new datasets without updating model weights. They compress much of the conventional causal-estimation workflow into pretrained inference, while their performance and validity depend on the causal priors, targets, and identifiability conditions used.

  • 3.1 What are CFMs?: CFMs are pretrained networks that estimate causal quantities on new datasets through in-context learning without updating model weights.Their working definition emphasizes pretraining across causal tasks and prediction on unseen tasks.
  • 3.1 What are CFMs?: CFMs learn mappings from observational datasets to posterior predictive distributions for causal quantities by training on tasks sampled from priors over data-generating processes.Pretraining pairs labeled observational data with causal queries computed from the sampled structural causal model.
  • 3.1 What are CFMs?: A CFM compresses dataset analysis, causal specification, estimator selection, nuisance-model fitting, overlap diagnosis, and estimation into a single pretrained forward pass.This workflow can reduce inference friction and support rapid analysis across many small datasets.
  • 3.1 What are CFMs?: Current CFMs have achieved performance above or on par with leading task-specific causal models and better than off-the-shelf tabular foundation models on causal problems.The comparison reflects benchmarking discussed for the first openly available CFMs and traditional causal models.
  • 3.1.1 The causal prior-data loss: Causal prior-data losses train CFMs to approximate causal posterior predictive distributions using simulated observational data and ground-truth causal quantities.The losses can be evaluated when interventional data are simulable and do not require closed-form posterior predictive distributions.
  • 3.1.1 The causal prior-data loss: Under suitable conditions, causal prior-data training converges to the true posterior predictive distribution, but accurate point estimates require causal identifiability under the prior.Non-identifiable priors leave posterior uncertainty because multiple data-generating processes can explain the same observational data.

4 Benchmarking CFMs

The benchmark evaluates three frozen CFMs against tuned classical estimators on semi-synthetic RealCause-Lalonde data, measuring heterogeneous effects, population effects, runtime, and average rank. CFMs are competitive despite no task-specific training, with substantial runtime advantages and model-dependent ATE recovery.

  • Experimental setup: The study establishes a standardized comparison of three CFMs and classical estimators on semi-synthetic observational data with known ground truth under strong confounding.RealCause-Lalonde supports evaluation of both CATE and ATE while preserving empirical covariate and selection distributions.
  • Experimental setup: The benchmark compares tuned meta-learners, propensity and doubly robust methods, and Double Machine Learning against frozen, pretrained Do-PFN, CausalPFN, and CausalFM checkpoints.Classical baselines include S-, T-, and X-Learners, IPW, DR-Learner, and Double Machine Learning; CFMs receive no parameter or hyperparameter tuning.
  • Benchmark performance: CFMs are remarkably competitive without training on Lalonde distributions; CausalPFN has the lowest average rank among CFMs and outperforms the tuned T-Learner on Lalonde-CPS.CFMs producing CID-PPD or CDTE-PPD estimates remain competitive with tuned X- and S-Learners.
  • Runtime: CFMs are 1–2 orders of magnitude faster on CPU than traditional training, tuning, and inference, with amortization providing up to 100× speedup for comparable-quality predictions.CausalPFN is the fastest tested CFM on CPU, and GPU execution provides additional speedups.
  • ATE recovery: CausalPFN achieves ATE relative error 0.17 on Lalonde-CPS, whereas Do-PFN and CausalFM remain near relative error 1 on both cohorts.The latter models’ nearly identical errors and small standard errors suggest systematic shrinkage toward zero rather than unstable estimation.

5 Broader Directions and Applications

The field has expanded beyond first-generation binary-treatment, tabular effect estimators toward broader causal quantities, treatment and data regimes, causal discovery, and domain-specific interventions. Figure 7 highlights the accompanying trade-off between competitive CATE estimation and much faster inference.

  • Broader directions: First-generation CFMs established amortized Bayesian causal inference for binary-treatment, tabular causal effect estimation, while later work broadens the formulation along several axes.The broader directions include new identification settings, treatment and data regimes, causal discovery, and domain-specific interventional questions.
  • Applications and evaluation: Figure 7 compares CATE estimation ranking with wall-clock runtime on RealCause-Lalonde, showing competitive CFM estimation alongside orders-of-magnitude faster inference than tuned classical estimators.CFM runtimes are reported on matched CPU hardware and on an A100 GPU.
  • Scope: The paper defines CFMs as PFNs that estimate causal quantities such as the CEPO, CID, or CDTE, while related models may extend the core ideas without meeting that definition.The review therefore covers neighboring work in causal discovery and related scientific tasks separately from CFMs under the working definition.

5.1 Previous work

Before modern CFMs, learning-to-learn approaches applied meta-learning ideas to causal inference, but some precursor methods used comparatively narrow, low-dimensional synthetic structural settings. Other work established theoretical and zero-shot transformer-based connections between attention and causal inference.

  • Learning-to-learn precursors: CaML and BBCI introduced learning-to-learn approaches for causal inference before the first generation of CFMs.These were algorithms applicable in principle to trainable model families rather than specific foundation models.
  • Learning-to-learn precursors: BBCI trains models to predict causal estimands on synthetic structural causal models and deploys them on unseen data.Its training class is comparatively narrow and low-dimensional, with no indication in the passage that it scales to real-world big-data settings.
  • Attention-based precursors: Causal Inference with Attention linked attention theoretically to causal inference and proposed a zero-shot transformer-based model.The work is presented as an ideological precursor to modern CFMs.

5.2 Extending causal inference capabilities

CFMs are being extended beyond basic tabular, binary-treatment settings to incorporate structural knowledge, richer treatment and temporal regimes, partial identifiability, and reliability safeguards.

  • Structural knowledge: First-generation CFMs cannot condition on additional causal-structure knowledge, but partial ancestral information with learnable attention biases yields sizeable performance gains.Partial ancestral information represents potentially incomplete knowledge of each variable’s ancestors.
  • Expanded causal capabilities: TabPFN-CFM predicts causal structures and effects across observational, interventional, and counterfactual levels of Pearl’s causal hierarchy.Known causal graphs can be supplied as inputs to improve causal-effect estimates.
  • Expanded causal capabilities: CCPFN extends CFMs to continuous treatments by reconstructing an individual treatment-response curve over the treatment range.This requires a higher-dimensional prior over treatment and outcome mechanisms than binary or multi-arm settings.
  • Temporal settings: CausalLongPFN estimates future counterfactual outcomes from historical trajectories and proposed treatment sequences under longitudinal dynamics and time-dependent confounding.Related temporal-prior work samples time-dependent structural causal models coupled with observational and interventional time series.
  • Partial identification: IV-ICL and a causal-sensitivity PFN extend CFMs to partially identifiable settings by producing effect bounds or sensitivity bounds rather than precisely identified effects.IV-ICL estimates the posterior predictive distribution of causal effects and derives bounds from its quantiles.
  • Reliability and calibration: Reliability studies report prior-induced bias, possible low credible-interval coverage, and errors from post-treatment covariates, motivating diverse priors and correction methods.Proposed responses include posterior correction, task-specific fine-tuning, and filtering inappropriate inputs.

5.3 Foundation models for causal discovery

Foundation-model methods are being applied to causal discovery by mapping observational or interventional data toward causal structures or their posterior distributions, with synthetic-prior diversity shaping generalization.

  • Foundations: Causal discovery seeks causal relationships from observational data and commonly faces non-identifiability, with data often determining only a Markov equivalence class rather than a unique graph.It is formulated as recovering the underlying causal graph over observed variables.
  • Foundations: A prior over structural causal models induces graph labels for synthetic datasets, enabling prior-fitted networks to amortize causal discovery into posterior prediction over causal structures.Causal-effect CFMs are argued to perform this discovery implicitly.
  • Early approaches: Early models such as CSIvA and AVICI directly map data to causal structures, while AVICI additionally learns a posterior over causal graphs.CSIvA generalizes poorly to unseen causal structures, whereas broader structural diversity improves generalization.
  • Early approaches: The Bayesian Causal Neural Process amortizes Bayesian causal discovery but trains across multiple smaller priors rather than one highly diverse prior.Its experiments examine combining those separate priors.
  • Early approaches: SEA aggregates weak predictors and summary statistics into a graph, whereas ADAG directly maps data to graphs but generalizes well only when unseen tasks share structure or ordering with training data.These designs differ from posterior-predicting causal discovery foundation models.
  • Recent foundation models: Later models include zero-shot Arrow, graph-oriented TabCausal, local-to-global DCD-PFN, dual-submodule DAG-FM, and temporal models predicting time-dependent adjacency matrices.Several are pretrained on large, diverse synthetic priors, while Arrow outputs a single DAG rather than a posterior.

5.4 Domain-specific and adjacent applications

The prior-fitting methodology also supports specialized scientific models, including biological perturbation prediction and applications to survival analysis with censored or truncated data.

  • Domain-specific models: Domain-specific causal models can use specially constructed synthetic priors when the target scientific mechanisms are highly complex but sufficiently structured for scalable pretraining.The methodology is presented as applicable beyond general causal-effect estimation.
  • Biological applications: MapPFN predicts the post-perturbation distribution of single-cell systems from observational data and interventional experiments.It is trained on a synthetic biological prior and estimates the response to a new intervention.
  • Adjacent applications: Survival analysis targets time-to-event prediction, but its data are often censored or truncated, so observed information may be partial or excluded.Examples include patient death, user churn, and device failure.

5.5 Future directions

Future CFM development emphasizes broader support for treatment regimes, identification assumptions, estimands, information sources, and prior coverage while preserving transparent uncertainty and assumption handling.

  • General-purpose models: A future general-purpose CFM could support multiple treatment regimes, identification assumptions, and causal estimands within one model.The proposed scope includes binary, multi-arm, and continuous treatments; backdoor and IV settings; and estimands such as CEPO, CID, and CDTE.
  • General-purpose models: Greater generality should retain an interface for specifying estimands and assumptions, checking identification, and propagating causal-structure uncertainty into downstream estimates.Common internal representations could support several downstream tasks without separate pretrained models.
  • Flexible inputs: Future CFMs could incorporate experimental data, partial causal knowledge, and domain constraints, then revise estimates when new experiments or structural information become available.Recent models using partial ancestral knowledge or graphical structure provide examples of this direction.
  • Prior coverage: CFM performance and generalization depend directly on synthetic-prior quality, motivating quantitative prior-coverage measures and automatic mismatch detection.The causal setting makes mismatch difficult because observational data do not fully characterize the distributions of interest.

6 Conclusion

Causal foundation models are introduced as pretrained, transformer-based methods that estimate causal effects on unseen settings through in-context learning. This work provides a hands-on entry point with notebooks, empirical results, and discussion of their competitive performance and broader scope.

  • CFMs are pretrained on a synthetic causal prior and use in-context learning to estimate causal effects in unseen settings.
  • The CFM paradigm reduces deployment time compared with traditional estimators trained and tuned for each new problem.
  • The paper provides Jupyter notebooks and empirical results on the effectiveness and speed of current CFMs.
  • The work is an introduction to emerging causal foundation modeling techniques intended to demystify them and broaden their audience.
Loading 2609.03003v1…