PaperlayerView exact revision on arXiv

Source-linked AI summary

Do-PFN: In-Context Learning for Causal Effect Estimation

Jake Robertson, Arik Reuter, Siyuan Guo, Noah Hollmann, Frank Hutter, Bernhard Schölkopf

arXiv:2506.06039v3cs.LG

TL;DR

Causal effect estimation from observational data is limited by unavailable interventions, unknown causal graphs, and difficult assumptions such as unconfoundedness. Do-PFN pre-trains a PFN on synthetic SCMs and interventions to predict interventional outcomes from observational data, with evaluations indicating strong performance and robustness across causal settings. Its applicability remains dependent on whether the synthetic SCM prior captures real-world causal complexity.

  • Problem

    Existing causal-effect methods may require interventional data, a known causal graph, or assumptions such as unconfoundedness, limiting applicability in real-world settings.

  • Method

    Do-PFN is a pre-trained transformer that meta-learns to predict interventional outcomes from observational data using synthetic SCMs and simulated interventions.

  • Results

    Do-PFN performs better across different causal scenarios and is more robust to unconfoundedness violations than strong tabular and causal baselines.

  • Takeaways & Limitations

    Do-PFN estimates causal effects without knowledge of the underlying causal graph and captures uncertainty arising from causal-effect unidentifiability.

  • Takeaways & Limitations

    Generalization depends critically on the synthetic SCM prior capturing real-world causal complexity, while validation currently relies on synthetic data.

Abstract

from arXiv · show

Estimation of causal effects is critical to a range of scientific disciplines. Existing methods for this task either require interventional data, knowledge about the ground truth causal graph, or rely on assumptions such as unconfoundedness, restricting their applicability in real-world settings. In the domain of tabular machine learning, Prior-data fitted networks (PFNs) have achieved state-of-the-art predictive performance, having been pre-trained on synthetic data to solve tabular prediction problems via in-context learning. To assess whether this can be transferred to the harder problem of causal effect estimation, we pre-train PFNs on synthetic data drawn from a wide variety of causal structures, including interventions, to predict interventional outcomes given observational data. Through extensive experiments on synthetic case studies, we show that our approach allows for the accurate estimation of causal effects without knowledge of the underlying causal graph. We also perform ablation studies that elucidate Do-PFN's scalability and robustness across datasets with a variety of causal characteristics.

1 Introduction

Causal effect estimation matters across science but is difficult from observational data because experiments may be infeasible and common assumptions can be hard to justify. Do-PFN extends prior-data fitted networks by pre-training on synthetic causal data with interventions to predict interventional outcomes from observational data alone.

  • Causal effects are central to questions in medicine, economics, and the social sciences.
  • RCTs directly estimate causal effects but can be unethical, expensive, or impossible, whereas observational data alone may require strict assumptions.
  • Unconfoundedness enables identification from observational data but is difficult to verify because relevant confounders must be observed and properly accounted for.
  • Do-PFN extends PFNs to conditional interventional distributions by meta-learning causal inference from mixtures of synthetic observational and interventional data.
  • The paper evaluates Do-PFN across synthetic datasets and benchmarks, including ablations, out-of-distribution analysis, uncertainty calibration, and comparisons with causal and tabular baselines.

2 Background and related work

The background introduces SCMs, interventions, CIDs, and PFN-based amortized inference as the conceptual foundations for estimating conditional causal effects from observational data.

  • An SCM combines a DAG with structural equations that generate each variable from its causal parents and noise.
  • An intervention do(t) fixes treatment T and removes its incoming edges, with its causal effect represented by p(y|do(t), ψ).
  • A CID conditions an interventional outcome distribution on covariates and supports CATE estimation through τ(x) := E[y|do(1), x] − E[y|do(0), x].
  • Do-calculus requires a believed or inferred SCM to determine whether and how causal effects can be estimated, using adjustments such as back-door and front-door methods.
  • PFN amortized inference learns a dataset-to-posterior mapping by training on simulated datasets, query inputs, and target outcomes.

3 Methodology: causal inference with PFNs

Do-PFN places a prior over SCMs and trains a transformer on synthetic observational–interventional pairs to approximate CIDs from observational data. Its predictive uncertainty reflects outcome noise, causal unidentifiability, and finite-data uncertainty, while infinite observations remove only the last component.

  • Do-PFN assumes a prior over SCMs capable of generating observational data and uses this prior to extend causal inference beyond fixed graphs or unconfoundedness.
  • The model is pre-trained on synthetic observational and interventional datasets to predict interventional outcomes from covariates, treatment values, and observational data.
  • Prior fitting minimizes negative log-likelihood so qθ approximates the conditional interventional distribution using observational data alone.
  • Do-PFN’s posterior predictive captures aleatoric outcome noise, uncertainty from observationally equivalent SCMs, and epistemic uncertainty about the equivalence class from finite data.
  • With infinitely many observational datapoints, uncertainty about the Markov-equivalence class vanishes, while other uncertainty sources remain.
  • The transformer has 7.3 million parameters and adds an indicator identifying treatment and covariate columns.

4 Experiments

Do-PFN is evaluated across synthetic causal case studies, out-of-distribution settings, and hybrid synthetic-real-world benchmarks. It generally estimates causal effects competitively or better than baselines, while showing robustness to dataset characteristics and causal assumptions.

  • Predicting conditional interventional distributions: Do-PFN significantly outperforms Random Forest, TabPFN (v2), and Dont-PFN on synthetic causal prediction tasks.The contrast with Dont-PFN indicates that intervention-aware pre-training contributes beyond observational posterior prediction.
  • Predicting conditional interventional distributions: Do-PFN performs competitively with DoWhy while automatically applying appropriate front-door and back-door adjustment without explicit causal-graph knowledge.Its shortest pre-trained variant remains competitive with a graph-informed version, supporting graph-free adjustment.
  • Estimating conditional average treatment effects: Do-PFN significantly outperforms meta-learners, causal-forest DML, DragonNet, and TARNet for CATE estimation, with a similar trend for ATE estimation.These baselines’ performance depends on unconfoundedness, whereas Do-PFN remains more consistent when that assumption is not met.
  • Ablation studies: Do-PFN remains competitive with DoWhy as dataset size grows, becomes more consistent with more samples, and performs competitively across graph complexities.It also uses additional data points to alleviate increasing noise, with slightly larger improvements on more complex graphs.
  • Ablation studies: Do-PFN is robust to many noise distributions, nonlinearities, base-rate treatment effects, and unidentifiable settings, but larger Gaussian noise and out-of-prior graph sizes reduce performance.Training Do-PFN-v1.1 on graphs up to 60 nodes significantly improves performance on graphs with 21–50 nodes.
  • Hybrid synthetic-real-world data: On RealCause, Do-PFN-v1.1 performs strongly among competitive baselines on three of four datasets despite the benchmark satisfying unconfoundedness.On Law School Admissions it achieves the strongest performance, while on Amazon Sales it is not significantly worse than the best method despite complex mediation.

5 Discussion

Do-PFN estimates interventional outcomes from observational data and performs well across several causal scenarios, including settings that violate unconfoundedness. Its scope remains bounded by dependence on the synthetic SCM prior, developing theory for amortized inference, and unaddressed causal configurations.

  • Do-PFN performs better across different causal scenarios and is more robust to unconfoundedness violations than strong tabular and causal baselines.The authors report this outcome from controlled synthetic experiments.
  • Do-PFN is robust to many unseen forms of noise and non-linearities, and performs competitively on two semi-synthetic benchmarks.Performance is especially strong in mediated scenarios where unconfoundedness is not satisfied.
  • Do-PFN’s generalization depends critically on whether its synthetic SCM prior captures real-world causal complexity.The paper states that robustness to prior–reality mismatch and performance on diverse real-world datasets require further systematic exploration.
  • Do-PFN offers efficient amortized inference, but its theoretical guarantees remain less developed than those of traditional estimators under known causal structures.The paper identifies the statistical theory for such amortized models as still developing.
  • The initial work does not cover broader intervention types, counterfactuals, non-i.i.d. inputs, or several observational data characteristics and modalities.The authors suggest incorporating these cases into the data-generating prior as future work.
  • The paper presents Do-PFN as a novel methodology intended to broaden access to causal effect estimation.This conclusion is stated as an outlook rather than an established real-world deployment result.

A Proof of Proposition 1

The proof shows that Do-PFN’s training objective is a forward-KL-optimal approximation of the interventional outcome distribution under the synthetic data-generating process, with consistency established through posterior convergence.

  • Minimizing the pre-training risk yields a forward Kullback–Leibler optimal approximation of the true conditional interventional distribution by Do-PFN’s predictive distribution.The expectation is taken over data simulated from the SCM prior and the resulting observational datasets.
  • Do-PFN’s posterior predictive interventional distribution is consistent up to observationally equivalent structural causal models as the observational sample size tends to infinity.The proof recovers the posterior over SCMs that induce the same observational distribution.
  • The proof uses conditional-independence identities, conditional expectations, and Doob’s martingale convergence theorem to establish the limiting posterior result.The argument also represents observationally equivalent SCM classes using measurable representatives.
  • The consistency proof establishes almost-sure convergence on every measurable outcome set, which implies weak convergence of the predicted interventional distribution.This conclusion follows from the stated convergence of the conditional probabilities.

B.2 Comments on the consistency theorem 3

The consistency result is interpreted relative to observational equivalence, while the prior-fitting procedure generates diverse SCMs and paired observational–interventional datasets for training.

  • B.2 Comments on the consistency theorem 3: For infinite observational data, Do-PFN recovers the interventional distribution only up to the set of SCMs that are observationally equivalent to the data-generating SCM.A unique observationally compatible SCM is a special case in which the target distribution is recovered directly.
  • B.2 Comments on the consistency theorem 3: The synthetic prior samples DAGs, nonlinear mechanisms, and additive-noise SCM parameters to expose Do-PFN to varied causal structures.Node nonlinearities are sampled from quadratic, tanh, and ReLU functions.
  • B.2 Comments on the consistency theorem 3: Observational datasets contain 10 to 2,200 samples, while interventional datasets use the remaining budget to provide intervention-conditioned outcomes for training.The model is trained on paired observational and interventional datasets generated from the same sampled SCM.
  • B.2 Comments on the consistency theorem 3: Holding the same noise realization across pre-intervention covariates and interventional outcomes also permits the generated pairs to represent counterfactuals or single potential outcomes.
  • B.2 Comments on the consistency theorem 3: Each training iteration computes negative log-likelihood for interventional outcomes conditioned on the intervention, covariates, and observational dataset, followed by stochastic gradient descent.The implementation uses mini-batch stochastic gradient descent with Adam.

D Experimental Details

The experimental prior varies SCM structure, nonlinearities, and noise while constructing synthetic causal case studies from explicit structural equations.

  • Synthetic SCM mechanisms use sampled nonlinearities, including quadratic, tanh, and ReLU functions, combined with additive noise.
  • The simulation holds the same exogenous noise fixed when generating pre-intervention covariates and interventional outcomes, enabling counterfactual-style paired outcomes.
  • Table 1 lists the structural equations defining all causal case studies used in the experiments.

D.2 Evaluation metric

Evaluation compares normalized MSE-based causal predictions against causal and tabular baselines, including Do-PFN variants and standard treatment-effect learners.

  • D.2 Evaluation metric: Results are evaluated using normalized mean squared error so performance can be compared across datasets.
  • D.2 Evaluation metric: The comparison includes Do-PFN variants trained with different graph sizes, durations, graph priors, and noise distributions.
  • D.2 Evaluation metric: Baselines include observational-prediction PFNs, Random Forest, DoWhy-based interventional and counterfactual models, and standard CATE learners.
  • D.2 Evaluation metric: Additional baselines use TabPFN within X-Learner and double-machine-learning frameworks, alongside tuned T-Learner and X-Learner implementations.
  • D.2 Evaluation metric: The CATE comparison covers TARNet, DragonNet, causal forests, DoWhy-CATE, TabPFN-based meta-learners, and hyperparameter-optimized S-Learner models.

D.4 Hybrid synthetic-real-world data

The paper evaluates Do-PFN on hybrid synthetic-real-world datasets built from agreed-upon causal graphs, using those graphs to simulate gold-standard interventional outcomes.

  • The agreed-upon graphs enable DoWhy to estimate causal mechanisms and simulate interventional data for evaluating Do-PFN and baselines.
  • Amazon Sales: The Amazon Sales dataset studies how special shopping events affect smartphone-sales profit using advertising, pricing, sales, demand, revenue, and cost variables.
  • The Amazon Sales and Law School Admissions datasets provide real-world observational data paired with widely agreed-upon causal graphs.
  • Law School Admissions: The Law School Admissions dataset comes from the 1998 LSAC study and includes Race as a protected attribute in prior counterfactual-fairness work.

E Supplementary results

Do-PFN combines strong prediction performance with efficient inference, while the paper notes a terminology limitation in its treatment of Race in the law-school dataset.

  • E.1 Inference speed: Do-PFN achieves Pareto-optimal prediction speed and performance across all four RealCause dataset groups.The benchmark contains up to 100 realizations of each dataset, with 10–50 features and up to 10,000 samples.
  • E.1 Inference speed: Synthetic pre-training shifts computational cost upfront, allowing in-context examples to be processed in a single forward pass at inference.
  • The paper rejects treating Race in the law-school dataset as binary and states that ethnicity better describes the underlying social construct.

E.2 In-distribution analysis

In-distribution analyses show that Do-PFN is robust across treatment-effect rates and graph complexity, improves with larger datasets, and retains the strongest performance when unconfoundedness is violated.

  • Base rate average treatment effect: Do-PFN remains robust across varied base rates of the ATE, whereas TabPFN’s performance diverges as the base ATE grows.
  • Dataset size: Do-PFN’s CATE MSE decreases and its variance converges as dataset size increases from 5 to 2,000 samples.
  • Graph size and complexity: Do-PFN performs competitively with graph-aware DoWhy across graphs containing 4–10 nodes and 2–43 directed edges.
  • Standard deviation of additive noise: Performance decreases with higher irreducible additive noise but improves with dataset size at comparable noise levels, more strongly than DoWhy (Cntf.).
  • Causal criteria: Do-PFN maintains the strongest CATE performance both when unconfoundedness holds and when it is violated, although all methods degrade under violation.
  • Gold standard baselines: Do-PFN significantly outperforms Do-PFN-Graph and DoWhy (Int.) in conditional interventional distribution estimation.
  • Do-PFN variants: Do-PFN outperforms Do-PFN-Graph and Do-PFN-Short, while Do-PFN-Mixed performs similarly with half the pre-training time.

E.4 Common effect case study

The common-effect case study represents an RCT-like setting without intervention-induced distribution shift, where traditional tabular predictors match Do-PFN’s performance.

  • Common-Effect case study: The Common-Effect case perfectly represents a randomized controlled trial, making it the only evaluated setting where traditional tabular predictors achieve comparable CID performance.
  • Common-Effect case study: In the Common-Effect case, regression baselines perform similarly to Do-PFN because the intervention does not create a distribution shift between Dob and Din.
  • Uncertainty quantification: Do-PFN is highly uncertain for the unobserved-confounder case because the causal effect is unidentifiable.
  • Uncertainty calibration: Do-PFN is slightly under-confident for identifiable case studies, while its uncertainty calibration is evaluated using prediction-interval coverage.
  • Critical-difference comparisons: Across six causal case studies, Do-PFN significantly improves CID and ATE estimation over regression, meta-learner, double-machine-learning, and deep-learning baselines.
  • CATE comparisons: Do-PFN-CATE significantly outperforms the causal baselines across six synthetic case studies and remains competitive across all four RealCause dataset groups.

NeurIPS Paper Checklist

The checklist reports that the paper supports its claims, provides reproducibility materials and experimental details, and discusses limitations. It also states that the paper has no theoretical results requiring theorem proofs.

  • The paper states that its abstract and introduction claims are fully supported by its empirical and theoretical results.
  • Societal impact: The checklist indicates that the paper addresses potential societal impacts and expects discussion of harms, privacy, fairness, security, or unintended uses where relevant.
  • Limitations: The paper explicitly discusses four major limitations in Section 5.
  • Theory assumptions and proofs: The checklist indicates that the paper has no theoretical results requiring complete theorem assumptions and proofs.
  • Open access to data and code: The authors report releasing the prior, pre-trained model, inference code, and raw results, with reproducibility details in the appendices.
  • Experimental setting/details: Experimental details, compute requirements, and uncertainty reporting are described through training specifications, GPU and timing information, boxplots, quantiles, and a critical-difference plot.
Loading 2506.06039v3…