Causal Foundation Models

Christopher Stith, Hossein Rahmani, Jesse C. Cresswell

Causal Foundation Models (CFMs) use pretrained transformers to estimate causal effects from observational data without weight updates.

Can we replace bespoke, per-problem causal inference pipelines with a single, pre-trained foundation model that performs causal estimation via in-context learning?

Traditional causal inference requires a bespoke, multi-step pipeline for every new dataset: proposing a mechanism, selecting an estimator, and tuning hyperparameters. This process is slow, requires significant expertise, and prevents the transfer of knowledge between tasks. Causal Foundation Models (CFMs) solve this by using prior-data fitted networks (PFNs) pretrained on a vast range of synthetic causal tasks. They treat causal estimation as an in-context learning problem, mapping observational datasets directly to causal quantities in a single forward pass. CFMs achieve performance on par with or better than task-specific models while drastically reducing inference time and the need for manual pipeline design.

Paper Primer

CFMs function as amortized Bayesian inference engines: they are trained on synthetic data generated from structural causal models (SCMs) to learn a mapping from observational samples to a posterior predictive distribution (PPD) of causal effects. By pretraining on diverse, synthetic causal priors, the model learns to handle the fundamental problem of causal inference—the inability to observe counterfactuals—without requiring further training on the target dataset.

CFMs match or exceed the performance of traditional, task-specific causal estimators.

Benchmarking against popular models like causal forests and double machine learning shows CFMs provide top-tier accuracy while eliminating the need for per-task model tuning. High performance across diverse causal settings (backdoor, frontdoor, and instrumental variables) with near-instant inference.

The core mechanism hinges on the "causal prior-data loss," which forces the model to predict ground-truth causal quantities (like the conditional expected potential outcome) from purely observational data. Unlike standard predictive transformers, CFMs are trained on a structurally harder task: they must infer interventional distributions from factual (x, t, y) triplets, effectively learning to perform causal adjustment internally.

Why are synthetic priors necessary for training these models?

Real-world observational data lacks the ground-truth interventional labels (e.g., counterfactual outcomes) required to compute the causal prior-data loss. Synthetic priors allow for the generation of diverse, high-quality tasks with known causal ground truths, which is essential for training the model to estimate causal effects.

How do CFMs handle the "fundamental problem of causal inference" at inference time?

CFMs leverage in-context learning to perform amortized Bayesian inference. By conditioning on the provided observational dataset, the model produces a posterior predictive distribution that quantifies uncertainty about the causal effect, accounting for both finite-sample noise and structural uncertainty about the underlying data-generating process.

CFMs shift causal inference from a manual, expert-driven pipeline to a scalable, plug-and-play model architecture. Practitioners can now perform rapid causal analysis on small datasets without the overhead of designing and tuning custom estimators.

Introduction

We expose the bottleneck of handcrafted causal pipelines and propose amortized inference via Causal Foundation Models.

Traditional causal inference demands a handcrafted pipeline for every new problem—designing a causal mechanism, picking a compatible estimator, tuning hyper‑parameters, and finally training the model. In contrast, the foundation‑model paradigm in machine learning replaces per‑task training with a single pretrained network that can be applied to new tasks without fine‑tuning. Causal Foundation Models (CFMs) bring this amortized approach to causal inference, enabling direct estimation of causal effects via in‑context learning.

A repeatable but manual sequence of steps that must be rebuilt for each causal analysis task.

The shift from manual, per‑problem pipelines to amortized foundation models promises faster, reusable causal inference.

Background and Motivation

Key concepts and formalisms that underlie causal inference are introduced.

Predictive models can achieve high accuracy, yet they cannot answer “what would happen if we intervene?” – a distinction first illustrated by the classic stork‑birth‑rate example.

The framework treats each possible treatment as a separate “what‑if” world, defining for every individual a random variable that records the outcome that would occur under that treatment.

Estimands are the quantities we ultimately want to compute from a causal model, such as the expected outcome under a specific treatment or the average effect of changing a treatment.

To reason about these estimands we must model how the data are generated, distinguish what is observable, and formalize interventions.

Because many DGPs can produce the same observational distribution, we must ask when a causal estimand is identifiable from $P_{\text{obs}}$ alone.

A causal quantity is identifiable when every DGP that looks the same observationally yields the same value for that quantity.

These three assumptions together constitute the classic “backdoor” setting that makes many causal estimands recoverable from observational data.

When the backdoor assumptions fail, alternative identification strategies such as instrumental variables or front‑door adjustment are sometimes available, though they are less commonly used.

A Bayesian network encodes conditional independencies via a DAG; a structural causal model augments the DAG with explicit functional equations that generate data.

PFNs are transformer models trained on synthetic causal priors that directly output posterior predictive distributions via in‑context learning, bypassing explicit posterior computation.

**Figure 2.** An example DAG arising from the SCM given in Example 2.1, with endogenous variables $\{X, T, Y\}$ and exogenous variables $\{U_1, U_2, U_3\}$. Dashed nodes denote exogenous variables.

**Figure 3.** Traditional vs. CFM workflows. (Top) For each new problem, the traditional pipeline involves analyzing the dataset, proposing a model, and training it before inference. (Bottom) CFMs are applied off-the-shelf to new problems, without training, fine-tuning, or hyperparameter optimization. The reader is invited to experiment with several implemented CFMs for themselves.

Defining Causal Foundation Models

Defines Causal Foundation Models and shows how they turn causal inference into a single forward pass.

Causal foundation models repurpose PFNs to perform causal inference without any weight updates. By treating causal inference as an in‑context learning problem, they replace the multi‑step traditional pipeline with a single forward pass.

A CFM is a transformer that has been pretrained on many synthetic causal tasks so that, given a new observational dataset, it can directly output a posterior distribution over the desired causal quantity.

How does a CFM differ from a standard PFN trained for predictive tasks?

A standard PFN learns to predict outcomes given inputs, while a CFM learns to map an observational dataset plus a causal query to a posterior over a causal effect. The training objective incorporates causal identification constraints rather than a pure predictive loss, and inference is performed via in‑context learning instead of direct prediction.

Observational data: $(X,T,Y) = \{(0,0,0.3), (1,1,2.7), (0,1,2.1), (1,0,0.5)\}$.

Query presented to the CFM: “What is the average treatment effect (ATE) given this dataset?”.

The CFM processes the dataset and query in‑context and outputs a posterior $P(\text{ATE}\mid\text{data}) = \mathcal{N}(1.9,\,0.2^2)$.

From the posterior the user reads the mean 1.9 as the estimated ATE and the variance 0.04 as uncertainty.

The model infers the causal effect directly from the observational data, bypassing separate nuisance‑model fitting and overlap diagnostics.

**Figure 1.** Causal foundation models predict on each new dataset using in-context learning without training or fine-tuning. Labeled context data (`X_ctx`, etc.) are provided as examples so that the model can predict causal effects on unlabeled query data (`X_qry`). A quickstart example is available in this notebook: Jupyter.

Adapting PFNs to causal inference requires three updates: the prior‑data loss must reflect causal objectives, synthetic priors must encode identifiable causal mechanisms, and the transformer architecture must support in‑context causal queries.

The Causal Prior-Data Loss

We define a causal prior‑data loss that lets CFMs learn causal predictive distributions from purely observational data.

Predictive PFNs train on supervised pairs $(x,y)$, but causal inference must predict a causal quantity from observational triples $(x,t,y)$. This mismatch makes the training task structurally harder, motivating a new loss.

Instead of matching a model to supervised labels, we ask it to match the model’s predicted causal distribution to the true causal quantity that would arise under an intervention, using only observational data.

Sample $ψ$ and compute the true CEPO for treatment $t=1$: $μ_{1}(0.5;P_{ψ}) = 2 \cdot 1 + 0.5 = 2.5$.

Form the loss term $-\,\log q_{θ}(2.5\mid x=0.5, t=1, D_{obs})$ and take its expectation (here a single sample).

Repeat for $t=0$: $μ_{0}(0.5;P_{ψ}) = 2 \cdot 0 + 0.5 = 0.5$, yielding $-\,\log q_{θ}(0.5\mid x=0.5, t=0, D_{obs})$.

Sum the two terms to obtain $L_{t}(θ)$ for this query.

The example shows how the loss forces the model to predict the correct counterfactual mean for each treatment, even though the observational data only contains the outcome for $t=0$.

How does the causal prior‑data loss differ from the standard prior‑data loss used for predictive PFNs?

The standard loss matches $q_{θ}(y\mid x)$ to observed labels $y$ from supervised pairs $(x,y)$. The causal loss instead matches $q_{θ}$ to a true interventional quantity (either $μ_{t}(x)$ or an outcome $y$ under $do(T=t)$) while conditioning on an observational dataset $D_{obs}$ that never contains the counterfactual outcome. Thus the causal loss operates on a harder prediction task and uses simulated ground‑truth from a sampled DGP.

Designing Synthetic Causal Priors

Synthetic priors enable training CFMs when causal ground truth is unavailable.

Real‑world observational datasets never contain the interventional outcomes required by the prior‑data loss, making it impossible to evaluate a model on true causal quantities. Synthetic priors solve this by providing a controllable source of both observational and interventional data, guaranteeing the labels needed for training.

A synthetic causal prior is a probability distribution over structural causal models (SCMs) that can be sampled to produce complete observational and interventional datasets with known ground‑truth effects.

How does an identifiable prior differ from a non‑identifiable one in practice?

An identifiable prior restricts the sampled SCMs to satisfy conditions (e.g., backdoor) that make the causal effect uniquely recoverable from observational data; a non‑identifiable prior relaxes those constraints, so different causal graphs can produce the same observational statistics, and the model must learn to represent that ambiguity.

Sample 1: draw treatment probability $p=0.5$ and noise $\epsilon$ for each of 4 units, yielding observational pairs $(X,T,Y) =\{(0,0,0),(1,1,1),\dots\}$.

Intervention do$(T=1)$: replace the treatment equation with $T=1$ for a new batch of 2 units, keep the same noise distribution, and compute $Y = 1 + \epsilon$.

Sample 2: change the treatment probability to $p=0.8$; repeat observational and interventional generation, producing a different causal effect magnitude.

Because the full SCM is known, we have exact counterfactual outcomes for every unit in both samples.

This toy illustrates how a synthetic prior yields paired observational and interventional data with exact ground truth, something unattainable from real‑world data.

Sample an SCM from the synthetic prior $\pi$.

Generate an observational dataset $\mathcal{D}_{\text{obs}}$ by feeding exogenous noise through the SCM.

Simulate interventional targets (e.g., CEPO, CATE) by applying do‑operators to the treatment node and re‑sampling the remaining variables.

Provide $\mathcal{D}_{\text{obs}}$ as context to the CFM and train it to predict the simulated interventional targets.

**Figure 5.** Sampling data from a synthetic causal prior. *Step 1.* Sample an SCM from $\pi$. *Step 2.* Generate an observational dataset $\mathcal{D}_{\text{obs}}$ from the SCM. *Step 3.* Simulate interventional supervised learning targets, such as the CEPO or CATE. *Step 4.* Provide the observational data to the CFM as context and train it to estimate the target.

**Figure 6.** Observational and interventional dataset generation from an SCM $S^\psi$. (a) The full set of original structural equations map a batch of exogenous noise vectors $\{(u_{n,1}^{\text{obs}}, u_{n,2}^{\text{obs}}, u_{n,3}^{\text{obs}})\}_{n=1}^N$ to observational data $\mathcal{D}_{\text{obs}}$. (b) The intervention $\text{do}(T = t^*)$ sets the value of treatment for the $m$th query individual to be $t^*$, replacing the structural equation $f_T$ of $T$ with $T = t^*$. Using the noise samples $\{(u_{m,1}^{\text{int}}, u_{m,2}^{\text{int}}, u_{m,3}^{\text{int}})\}_{m=1}^M$, the remaining structural equations then produce an interventional dataset $\mathcal{D}_{\text{int}}$ with ground-truth potential outcomes $y_m(t^*)$.

Transformer Architectures for CFMs

How CFMs turn causal inference into an in‑context prediction problem using transformers.

Traditional causal estimators require a separate training phase and a hand‑crafted pipeline, which makes scaling to many tasks cumbersome. CFMs avoid this by treating causal queries as in‑context examples, letting a single transformer handle any downstream causal task.

CFMs feed a tabular observational dataset as a “context” to a transformer, then ask the model to answer causal queries (e.g., CATE or CEPO) directly, without any extra training.

How does this differ from a standard transformer language model that also uses masked attention?

In language modeling the mask blocks future tokens so the model cannot see words it is about to predict. In CFMs the mask blocks *query‑to‑query* attention, allowing the model to see the entire context but never to let one query influence another. This yields context‑only conditioning rather than autoregressive generation.

Tokenize each row: “treatment 0 | covariate 0.1 | outcome 1”, “treatment 1 | covariate 0.4 | outcome 0”.

Embed the two context rows and the query row (“treatment ? | covariate 0.3”) into vectors $c_1,c_2,q$.

Apply the masked attention matrix so $q$ attends to $c_1$ and $c_2$, while $c_1$ and $c_2$ attend only to each other (no self‑attention to $q$).

Aggregate the attended representations to produce a prediction vector $p$ for the query.

Take the mean of $p$ across the output distribution to obtain the point estimate $\hat{\tau}(x_{\text{q}})$.

The mask ensures that the query’s prediction cannot be altered by the presence or order of other queries, guaranteeing that each causal estimate is a pure function of the observed context.

**Table 2.** Model size and transformer depth for different CFMs and TFMs.

Tokenize and embed the observational context $D_{\text{obs}}$ and the query (CATE or CEPO) separately.

Feed both sets of embeddings into the transformer with the context‑query mask applied.

Read out the transformer’s output tokens corresponding to the query.

For CATE, average the output distribution to obtain $\hat{\tau}(x_m)=\mathbb{E}[q_\theta(\text{CATE}(x_m)\mid x_m,D_{\text{obs}})]$.

For CEPO, average similarly to obtain $\hat{\mu}_{t_m}(x_m)=\mathbb{E}[q_\theta(\mu_{t_m}(x_m)\mid x_m,t_m,D_{\text{obs}})]$.

Training Procedures

Training CFMs repeatedly samples causal models, creates observational and interventional data, and learns via a prior‑data loss.

The core difficulty is that a CFM must learn to infer causal quantities from purely observational context, which demands a training distribution that spans many possible causal mechanisms. A diverse synthetic prior over Structural Causal Models provides exactly that coverage.

Think of the model as a student who practices on countless toy worlds, each with its own hidden causal graph, so that when a real‑world problem appears the student can answer by recalling the right pattern.

How does this training differ from a standard predictive PFN that only sees observational data?

Standard PFNs generate both context and target from the same observational distribution, so the model learns to copy patterns. In CFM training the third step explicitly replaces the treatment equation with a fixed query, forcing the model to learn the mapping from observational context to interventional outcomes—a causal reasoning task rather than pure prediction.

Sample an SCM $S_ψ\simπ$ (choose graph edges, assign MLP parameters, and define noise distribution $PU$).

Generate $D_{obs}$ of size $N$ by drawing $U_k\sim PU$ and evaluating the structural equations for $(X,T,Y)$.

For each query treatment $t^*$, replace the $T$ equation with $T=t^*$, then generate $D_{int}$ of size $M$ using the same noise draws.

Compute the causal prior‑data loss on the pair $(D_{obs},D_{int})$ and back‑propagate to update the transformer parameters.

Repeat until convergence.

Sample the DAG and assign simple linear equations: $T = X_1 + U_T$, $Y = T + X_2 + U_Y$ with $U_T,U_Y\sim\mathcal{N}(0,0.1)$.

Draw two noise vectors for $D_{obs}$: $(U_T,U_Y)=(0.05, -0.02)$ and $( -0.03, 0.04)$. Compute $(X,T,Y)$ pairs, yielding $D_{obs}=\{(0,0,0.03),(1,1,1.01)\}$.

Fix the query treatment to $t^*=1$. Replace the $T$ equation with $T=1$ and reuse the same noise draws to produce $D_{int}=\{(0,1,1.03),(1,1,2.01)\}$.

Estimate $μ_{t=1}(x)$ by averaging the $Y$ values for each $x$: $μ_{t=1}(0)=1.03$, $μ_{t=1}(1)=2.01$.

This tiny example shows how a single forward pass after fixing $T$ yields both the interventional outcomes and the CEPO estimates, avoiding a separate Monte Carlo loop for each $(x,t)$ pair.

By repeatedly exposing the model to such synthetic causal worlds, the CFM learns a generic mapping from observational context to interventional predictions, enabling zero‑shot causal inference at test time.

Benchmarking and Evaluation

CFMs achieve near‑state‑of‑the‑art accuracy while being up to 100× faster at inference.

CausalPFN matches the best‑tuned T‑Learner in PEHE rank while delivering up to 100× faster inference than traditional estimators.

Table 3 shows CausalPFN has the lowest average rank among CFMs and is close to T‑Learner; Figure 7 demonstrates CFMs are orders of magnitude faster.

Do‑PFN learns a distribution over potential outcomes conditioned on observed data, enabling direct CATE estimation without any retraining.

CausalPFN approximates the CEPO‑PPD, producing calibrated conditional effect estimates directly from the pretrained network.

**Table 1.** Performance comparison of Causal Foundation Models (CFMs) and traditional estimators on Lalonde datasets.

**Figure 7.** CATE Estimation Ranking vs. Wall-Clock Runtime on RealCause-Lalonde. Causal foundation models deliver competitive estimation results while achieving orders-of-magnitude faster inference compared to tuned classical estimators because they do not need to be trained on RealCause. For CFMs, runtimes are shown both on CPU (matching the hardware used for all other models) and an A100 GPU.

Broader Directions and Applications

We survey how Causal Foundation Models are expanding beyond basic causal effect estimation.

The first generation of CFMs proved that Bayesian causal inference can be amortized, establishing a new class of foundation models that excel on standard benchmarks.

Early learning‑to‑learn approaches such as CaML, BBCI, and CInA introduced the idea of training models on synthetic causal priors, but they were limited by narrow structural assumptions and scalability concerns.

Subsequent work has broadened CFMs along several axes: incorporating partial structural knowledge, handling continuous and longitudinal treatments, addressing partial identification, and improving reliability through calibration techniques.

Methods such as Reuter et al.’s attention bias and TabPFN‑CFM’s graph conditioning demonstrate that supplying domain‑specific structural cues yields sizable performance gains.

New treatment regimes are tackled by CCFPN, which learns full continuous response curves, and by CausalLongPFN, which predicts counterfactual outcomes in longitudinal settings with time‑dependent confounding.

Partial‑identification research, exemplified by IV‑ICL and sensitivity‑analysis PFNs, provides interval bounds or sensitivity limits when exact effects cannot be recovered.

Reliability studies reveal prior‑induced biases and coverage issues; corrective strategies include one‑step posterior correction, task‑specific fine‑tuning, and input filtering.

Causal discovery foundation models such as Arrow, TabCausal, FoundCause, DCD‑PFN, and DAG‑FM extend the amortization principle to graph inference, often outputting probabilistic adjacency matrices rather than single deterministic graphs.

Domain‑specific applications illustrate the versatility of the prior‑fitting approach: MapPFN and PerturbPFN model single‑cell perturbations, while SIC, SurvivalPFN, and SurvPFN bring PFNs to survival analysis with censored data.

Future work aims to unify treatment regimes, identification assumptions, and estimands in a single CFM, to enable flexible incorporation of additional data sources, and to develop quantitative metrics for synthetic prior coverage and mismatch detection.

Read the original paper

Open the simplified reader on Paperglide

Browse all simplified papers