PaperlayerView exact revision on arXiv

Source-linked AI summary

CausalPFN: Amortized Causal Effect Estimation via In-Context Learning

Vahid Balazadeh, Hamidreza Kamkari, Valentin Thomas, Benson Li, Junwei Ma, Jesse C. Cresswell, Rahul G. Krishnan

arXiv:2506.07918v2cs.LGstat.ML

TL;DR

Selecting among many specialized causal estimators for observational data requires substantial domain expertise. CausalPFN uses a transformer trained on simulated ignorability-satisfying DGPs to amortize causal effect estimation, matching or often outperforming specialized estimators across diverse domains. Its scope depends on strong ignorability, an untestable assumption in practice.

  • Problem

    Choosing and tuning among dozens of specialized causal estimators for observational data requires substantial domain expertise.

    Source: §1 · pp. 1–2

  • Method

    CausalPFN trains a transformer on a prior over simulated DGPs to infer causal effects from observational contexts without task-specific tuning.

    Source: §1 · p. 2

  • Results

    CausalPFN matches and often outperforms specialized causal estimators across diverse real-world domains, achieving the best average CATE rank across IHDP, ACIC, and Lalonde and competitive ATE performance.

    Source: §1 · p. 2, §7 · p. 10

  • Takeaways & Limitations

    Amortization reduces the burden of estimator selection at inference time, and the model is released with code and presets to foster adoption.

    Source: §7 · p. 10

    CausalPFN fundamentally assumes strong ignorability, which is untestable in practice, and has no validity guarantees when that condition fails.

    Source: §7 · p. 10

Abstract

from arXiv · show

Causal effect estimation from observational data is fundamental across various applications. However, selecting an appropriate estimator from dozens of specialized methods demands substantial manual effort and domain expertise. We present CausalPFN, a single transformer that amortizes this workflow: trained once on a large library of simulated data-generating processes that satisfy ignorability, it infers causal effects for new observational datasets out of the box. CausalPFN combines ideas from Bayesian causal inference with the large-scale training protocol of prior-fitted networks (PFNs), learning to map raw observations directly to causal effects without any task-specific adjustment. Our approach achieves superior average performance on heterogeneous and average treatment effect estimation benchmarks (IHDP, Lalonde, ACIC). Moreover, it shows competitive performance for real-world policy making on uplift modeling tasks. CausalPFN provides calibrated uncertainty estimates to support reliable decision-making based on Bayesian principles. This ready-to-use model requires no further training or tuning and takes a step toward automated causal inference (https://github.com/vdblm/CausalPFN/).

1 Introduction

Estimating causal effects from observational data is difficult because confounding obscures effects and choosing among many specialized estimators requires domain expertise. CausalPFN amortizes this workflow with a transformer trained on simulated causal processes, achieving strong benchmark performance without task-specific tuning.

1 Introduction

  • Dozens of specialized causal estimators make selecting, tuning, or designing an appropriate method for each application burdensome.Examples include Meta-Learners, doubly robust methods, DML, and neural-network approaches.Source: §1 · pp. 1–2
  • CausalPFN bridges Bayesian causal inference and PFN-style in-context learning, using a prior over ignorability-satisfying DGPs to infer effects from observational data.PFNs amortize posterior predictive inference through transformer training, but prior PFNs were designed for regression and classification rather than causal inference.Source: §1 · p. 2
  • CausalPFN achieves the best average CATE rank across IHDP, ACIC, and Lalonde, with competitive ATE performance and no task-specific tuning.The comparison covers 310 tasks; Figure 1 also reports faster wall-clock time from data to estimates.Source: Figure 1 · p. 1, §1 · p. 2
  • A single transformer trained on diverse simulated DGPs can match or surpass specialized estimators across datasets without task-specific tuning.The method is ready for inference on new datasets without further training, fine-tuning, or hyperparameter optimization.Source: §1 · p. 2

2 Background

The background defines causal effects using potential outcomes and explains why observational data require assumptions for identification. It then situates Bayesian inference and prior-fitted networks as approaches to predictive inference and uncertainty quantification.

2 Background

  • Causal effects are defined through conditional expected potential outcomes, which yield ATE and CATE for binary treatments.ATE averages the difference between treatment-specific conditional means; CATE conditions that difference on covariates.Source: §2 · p. 3
  • Strong ignorability makes conditional expected potential outcomes identifiable by requiring unconfoundedness given covariates and positive treatment probability.Without further assumptions, distinct DGPs can produce the same observational distribution but different causal effects.Source: §2 · p. 3
  • Bayesian causal inference places a prior over DGPs, updates it with observed data, and derives a posterior predictive distribution for causal estimands.Point estimates and credible intervals follow from the induced posterior, while approximate inference methods such as MCMC or variational inference are commonly used.Source: §2 · p. 3
  • Prior-fitted networks use a transformer trained on simulated DGPs to map an observational context directly to a posterior predictive distribution.This shifts posterior inference costs from each new dataset to pretraining, but the PFNs described here target regression and classification rather than causal inference.Source: §2 · pp. 3–4

3 The Mathematical Framework of CausalPFN

CausalPFN targets the posterior predictive distribution of conditional expected potential outcomes, from which point estimates and uncertainty can be obtained. Its consistency result links recovery of true effects to identifiability under the prior, while training amortizes posterior predictive inference in a transformer.

3 The Mathematical Framework of CausalPFN

  • CausalPFN targets the CEPO posterior predictive distribution, whose mean estimates conditional expected potential outcomes and whose samples support uncertainty intervals for CEPOs, CATEs, and ATEs.CATEs follow from treatment-specific CEPO differences, and ATEs from empirical averaging across units.Source: §3 · pp. 4–5
  • Under mild regularity assumptions, CEPO-PPDs consistently recover true CEPOs if and only if the prior is CEPO-identifiable.CEPO-identifiability means that the CEPOs depend only on the observational distribution; the result applies to almost all prior DGPs and covariates.Source: §3 · p. 4
  • The consistency theorem itself applies beyond strong ignorability when CEPOs are identifiable, though the practical setting uses a prior enforcing strong ignorability.The prior must rule out non-identifiable cases; broadening it can increase the chance that a true DGP lies in its support.Source: §3 · p. 4
  • A single transformer is trained with the causal data-prior loss to approximate the full CEPO-PPD without explicitly computing the posterior at test time.Training samples DGPs from the prior, simulates observational contexts and query truths, and updates the model to assign probability to the true CEPO.Source: §3 · pp. 4–5

4 Implementing CausalPFN

The implementation focuses on binary treatments under strong ignorability and constructs a diverse prior by converting base tables into causal datasets. A PFN-style transformer predicts quantized CEPO distributions, trained with a histogram-loss approximation to the causal data-prior loss.

4 Implementing CausalPFN

  • The implementation uses binary treatments under strong ignorability, leaving more general treatment settings for future work.The mathematical framework is stated for arbitrary finite treatment sets and identifiability.Source: §4 · p. 5
  • The prior is designed to cover diverse scenarios while ensuring every DGP satisfies strong ignorability and exposes ground-truth CEPOs for training.Covariates are drawn from a mix of real and synthetic tables, without distributional assumptions on covariates and potential outcomes.Source: §4 · pp. 5–6
  • The prior converts real or synthetic base tables into causal data by selecting covariates and potential-outcome columns, then generating treatment and observed outcomes.A random treatment function depends on covariates; this construction ensures unconfoundedness and positivity by design and can extend to finite discrete treatments.Source: §4 · pp. 5–6, Figure 4 · p. 5
  • A PFN-style transformer encodes observational rows and query inputs, then outputs a quantized CEPO distribution over 1024 outcome bins.Query tokens cannot attend to one another and attend only to context tokens, making predictions conditionally independent given the context.Source: §4 · p. 6
  • The histogram loss approximates the causal data-prior loss and coincides with it in the limits σ → 0 and L → ∞.This histogram formulation provides a tractable proxy for the continuous CEPO-PPD.Source: §4 · p. 6

5 Experiments

CausalPFN achieves the best average CATE rank across standard benchmarks, performs competitively on marketing uplift tasks, and shows OOD overconfidence that regression-based temperature scaling addresses.

5 Experiments

  • The evaluation compares broad causal and neural-estimator baselines, tuning most with cross-validation; CausalPFN trains on context rows and predicts CEPO distributions for treatment–covariate queries.Its transformer uses asymmetric attention, and the CEPO-PPD mean supplies the inference-time point estimate.Source: §5 · p. 6, Figure 5 · p. 7

Benchmarks with Ground-Truth Effects

  • CausalPFN outperforms all baselines in average CATE rank across 310 realizations and is competitive in average ATE rank, despite training only on simulated data.The benchmarks cover IHDP, ACIC 2016, and Lalonde CPS/PSID; Table 1 reports mean PEHE and average rank for CATE, and mean ATE relative error and rank.Source: Benchmarks with Ground-Truth Effects · p. 7, Table 1 · p. 7

Policy Evaluation on Marketing Randomized Trials

  • On 50k stratified subsamples, CausalPFN achieves the best mean normalized Qini performance; on Hillstrom it closely matches X Learner, while full-table runs lose performance.The four larger campaigns use 50k subsamples; Hill(1) and Hill(2) use all 64,000 rows.Source: Policy Evaluation on Marketing… · p. 8, Table 2 · p. 8

Uncertainty & Calibration

  • In-distribution CATE intervals are reliably calibrated, but OOD intervals are severely overconfident; regression-ICE temperature scaling removes observed OOD overconfidence and yields calibrated or conservative intervals.Figure 7 compares coverage with nominal credibility and shows the calibration patterns before and after scaling.Source: Figure 7 · p. 9, Uncertainty & Calibration · p. 9
  • Because test-time CATEs are unavailable, the method fits softmax temperature by five-fold regression calibration on observational data; regression ICE is biased by outcome noise but empirically no greater than CATE ICE.The regression intervals combine epistemic uncertainty with irreducible outcome noise.Source: Uncertainty & Calibration · p. 9

6 Related Work

Causal inference methods include per-dataset estimators and amortized approaches, but the latter have not consistently surpassed specialized estimators on standard benchmarks. CausalPFN is presented as a ready-to-use amortized estimator that infers effects without test-time adaptation.

6 Related Work

  • Existing single-dataset estimators, including learner, doubly robust, and neural approaches, require per-dataset training and do not amortize across datasets.Examples include X-, S-, DR-, and RA-Learners, IPW, DML, TARNet, DragonNet, CEVAE, and NCMs.Source: §6 · pp. 9–10
  • Table 3 compares PEHE and ATE relative error for TabPFN⋆, the latest TabPFN model tuned with the authors’ prior, and other methods.The caption states that best numbers are highlighted but does not provide their values.Source: Table 3 · p. 10
  • Prior amortized methods recover causal graphs before computing interventions or infer effects end-to-end, while CausalPFN is trained once without adapting to test-time DGPs.The paper characterizes existing amortized methods as not providing a ready-to-use estimator that consistently surpasses specialized single-dataset estimators on standard benchmarks.Source: §6 · p. 10
  • TabPFN, TabDPT, and TabICL apply broad-prior in-context learning to tabular prediction; CausalPFN extends this approach to causal inference.The related-work discussion says increasing model size and training data have consistently led to stronger performance in transformer in-context learning.Source: §6 · p. 10

7 Conclusions, Limitations, and Future Work

CausalPFN combines Bayesian causal inference with large-scale tabular training to estimate effects across diverse domains, reducing estimator-selection burden at inference. Its scope is constrained by ignorability, theoretical assumptions, scalability, and the current binary-treatment implementation.

7 Conclusions, Limitations, and Future Work

  • CausalPFN matches and often outperforms specialized causal estimators across diverse real-world domains despite learning solely from simulated data.The approach combines Bayesian causal inference with large-scale tabular training.Source: §7 · p. 10
  • Amortization reduces the burden of selecting an estimator at inference time, and the authors open-source code and presets to foster adoption.Source: §7 · p. 10
  • Strong ignorability is untestable in practice, and CausalPFN has no validity guarantees when that assumption does not hold.The method also relies on idealized well-specified-prior and asymptotically large-dataset assumptions, with no finite-sample theory characterizing practical behavior.Source: §7 · p. 10
  • CausalPFN’s performance degrades on the largest marketing tables, reflecting the size-scalability trade-off inherent to PFN-style models.Source: §7 · p. 10
  • Although the framework supports finite multi-arm discrete treatments, implementation is binary-only; continuous treatments remain unexplored.The authors also identify scalable priors for richer settings such as instrumental variables as future work.Source: §7 · p. 10

NeurIPS Paper Checklist

This section documents the paper’s claims, limitations, theory, and experiments for verification, with pointers to supplementary material. It also formalizes CEPOs and CEPO-identifiability and provides the theorem result and consistency limit used in its proof.

NeurIPS Paper Checklist

  • The authors say the paper’s stated contributions match its scope, discuss limitations and future directions, and provide complete assumptions and proof for Theorem 1.The main paper includes a proof sketch; the complete proof is in Section B.Source: NeurIPS Paper Checklist · pp. 19–20
  • The authors report releasing code, model checkpoints, notebooks, a PyPI package, and public training data with instructions to reproduce the experiments.They also say the paper and Section D specify important experimental details and inference hyperparameters.Source: NeurIPS Paper Checklist · pp. 20–21
  • The checklist reports 1-sigma standard errors across benchmark realizations and calibration samples, and identifies societal benefits as potential rather than directly achieved.It also states that training the base model used one A100 GPU for seven days.Source: NeurIPS Paper Checklist · pp. 21–23

Appendix Contents

  • The appendix contents list sections on dataset-size sensitivity, baseline hyperparameters without tuning, and marketing experiments.Source: Appendix Contents · p. 26

A Notation, Definitions, and Assumptions

  • The formal notation defines observed and counterfactual variables and imposes a measurability condition on the mapping from data-generating parameters to laws.Source: A Notation, Definitions, and Assumptions · p. 27
  • CEPO-identifiability requires each treatment’s conditional expected potential outcome to be determined by covariates and the observational-equivalence class of the data-generating process.Formally, the CEPO value must equal a map of x and [ψ] for prior-almost-all parameters and almost-all covariate values.Source: A Notation, Definitions, and Assumptions · p. 28

B.1 Re-Statement of Theorem 1

  • Under Assumptions 2 and 3, Proposition 1 states that its result holds if and only if the prior is CEPO-identifiable.Source: §B.1 · p. 28

B.2 Preliminaries for the Proof of Theorem 1

  • For almost every ϕ₀ under ν, the posterior expectation of g(ϕ) given n observations converges to g(ϕ₀), P^ϕ₀-almost surely.The stated corollary of Doob’s consistency theorem is introduced as a tool for proving Theorem 1.Source: §B.2 · p. 28

B.3 Proof of Theorem 1

The proof connects consistency of the CEPO posterior predictive distributions to CEPO-identifiability. It establishes the two implications and uses full-measure sets to complete the argument.

B.3 Proof of Theorem 1

  • Under Assumptions 2 and 3, full-measure sets X0 and Ψ0 exist; for every t∈T and x∈X0, Eπ[|µt(x;ψ)|] < ∞.Source: §B.3 · p. 29
  • Theorem 3 establishes consistency in the observational quotient space, while identifiability connects that result to the original CEPOs.The proof applies the quotient-space result to all treatments and covariates in the stated full-measure set.Source: §B.3 · pp. 29–30
  • CEPO-identifiability implies consistency, and consistency in turn implies CEPO-identifiability.The proof establishes each direction on sets of prior probability one.Source: §B.3 · p. 30

C Validity of the Causal Data-Prior Loss

The section justifies the causal data-prior loss by relating its optima to an expected forward-KL objective, then describes how CausalPFN’s prior and inference workflow are constructed. The model uses diverse simulated causal data and retrieval-based context selection for large tables.

C Validity of the Causal Data-Prior Loss

D.1 Prior Generation & Simulating DGPs

  • Input-dependent noise preserves conditional outcome means while allowing heteroscedastic outcome distributions.The noise scale is chosen to yield training data with a more informative signal-to-noise ratio.Source: §D.1 · p. 32
  • Prior generation combines real or synthetic base tables with sampled covariates, CEPOs, treatments, and outcomes to create DGPs satisfying strong ignorability and positivity.The design also varies treatment-effect heterogeneity and positivity across generated DGPs.Source: §D.1 · p. 32

D.2 Model Details

  • CausalPFN’s query outputs form a discrete CEPO posterior, from which inference returns a posterior-mean estimate and sampled credible intervals.The model samples 10,000 times to estimate intervals at the desired significance level.Source: §D.2 · p. 33
  • For tables exceeding the context window, weak CATE estimates stratify context rows and queries so each query batch uses a nearby contiguous context window.The method uses bisection search to select the window for each batch.Source: §D.2 · pp. 33–34

D.3 Sensitivity to Dataset Size

On synthetic datasets, CausalPFN’s PEHE falls faster as sample size increases, while its advantage narrows as covariate count grows. For large tables, its inference workflow uses fixed forward passes and avoids deployment-time refitting or tuning.

D.3 Sensitivity to Dataset Size

  • CausalPFN shows faster PEHE decay as sample size increases, with a slight plateau at very large contexts.The comparison fixes covariates at 10 and reports mean ± SE across 50 datasets.Source: §D.3 · p. 34, Table 4 · p. 34
  • CausalPFN outperforms the baselines across tested covariate sizes, but the gap narrows as dimensions grow beyond its training exposure of up to 100 dimensions.The experiments fix sample size at 1,000 and report PEHE across 50 datasets.Source: §D.3 · p. 34, Table 5 · p. 34

D.4 Discussion on Inference Speed

  • For large-table deployment, CausalPFN uses a small fixed number of forward passes and parallelizes inference, while baselines may incur refitting or tuning costs.The comparison concerns practitioner-facing time from data arrival to returned effects, not intrinsic speed against every baseline.Source: §D.4 · p. 34

D.5 Baseline Hyperparameters and Results without Hyperparameter Tuning

Without hyperparameter tuning, CausalPFN attains the best average rank on CATE and the second-best on ATE. The comparison uses recommended or default baseline settings and reports PEHE alongside ATE relative error.

D.5 Baseline Hyperparameters and Results without Hyperparameter Tuning

  • Without hyperparameter tuning, CausalPFN attains the best average rank on CATE and second-best on ATE.Table 6 reports PEHE for CATE and ATE relative error, including its overall average.Source: §D.5 · p. 35, Table 6 · p. 35
  • Untuned baselines use recommended EconML settings, default CATE Nets settings, and GRF with internal tuning disabled.BART’s scores remain unchanged because it offers no comparable hyperparameter-tuning option.Source: §D.5 · p. 35
  • All inference, including baseline inference, was performed on an 80 GB H100 GPU, 32 CPUs, and 256 GB RAM.This supplies the stated hardware context for the comparison.Source: §D.5 · p. 35

D.6 Marketing Experiments

The marketing experiments evaluate CausalPFN’s CATE rankings for uplift policy assessment across several promotion datasets, while full-table runs require retrieval because of context-length limits.

D.6 Marketing Experiments

  • The uplift datasets cover Criteo ad exposures, X5 personalized coupons, Lenta SMS promotions, and MegaFon telecom offers, with conversion, purchase, visit, and acceptance outcomes.Criteo uses a 2.5M stratified subset; the other datasets are described as randomized offer experiments or domain-faithful synthetic data.Source: §D.6 · pp. 35–36
  • Qini scores evaluate CATE-based rankings using five-fold outcome-and-treatment-stratified splits, with 80% of rows as context or baseline training data and 20% held out.Rows are ranked by estimated CATE to compute Qini curves and scores.Source: §D.6 · p. 36
  • The 50,000-token context limit is insufficient for full-table runs, so those evaluations use retrieval, which the authors conjecture may explain weaker performance than subsampled runs.The model’s default context length was increased from 4,096 to 50,000 tokens; retrieval is used when the full table exceeds that capacity.Source: §D.6 · p. 36
  • Full-table Qini runs cover Hill(1) and Hill(2) (64K rows), Criteo (2.5M), X5 (200K), Lenta (687K), and Mega (600K).Table 7 reports normalized Qini scores, with each dataset’s top-performing model normalized to 1.0.Source: Table 7 · p. 36

D.7 Calibration, Coverage, and Credible Intervals

Post-hoc regression calibration brings CausalPFN’s CATE calibration closer to ideal across synthetic and standard benchmarks, including out-of-distribution settings. The concurrent-work discussion contrasts its estimand and identifiability design with Do-PFN.

D.7 Calibration, Coverage, and Credible Intervals

  • The synthetic calibration tests vary polynomial degree and sinusoidal frequency, including tests on higher degrees or frequencies not seen during training.The polynomial setup trains on one degree family and tests across others; the sinusoidal setup includes L3 as an unseen-frequency test family.Source: §D.7 · p. 36
  • On sinusoidal DGPs, uncalibrated CATE estimates can be overconfident out of distribution, while regression-curve calibration produces well-calibrated or conservative results even there.The synthetic test families include L3 frequencies unseen during training.Source: §D.7 · pp. 36–37, Figure 8 · p. 37, Figure 9 · p. 38
  • On polynomial DGPs, regression calibration yields near-perfect CATE calibration after uncalibrated models show overconfidence when tested on out-of-distribution degrees such as Cubic after Quadratic training.Degrees 1–4 define Linear, Quadratic, Cubic, and Quartic families; each higher degree includes the lower-degree terms.Source: §D.7 · pp. 36–37, Figure 10 · p. 38, Figure 11 · p. 39
  • Post-hoc regression calibration improves reliability across synthetic and standard benchmarks, nearly perfecting calibration on synthetic data and correcting conservatism on IHDP and ACIC 2016 while nearing ideal alignment on Lalonde.The large-scale model generally appears conservative, which the authors suggest may be related to Gaussian smoothing in the histogram loss, a choice they say is needed for training stability.Source: §D.7 · p. 37, Figure 12 · p. 39, Figure 13 · p. 39

E Concurrent Work on PFNs for Causal Inference

  • Do-PFN targets the interventional posterior predictive distribution over outcomes, whereas CausalPFN targets the expected interventional outcome to isolate epistemic uncertainty about causal effects.The paper notes that this target removes outcome, or aleatoric, noise from the prediction target.Source: E Concurrent Work on PFNs for Causal… · p. 40
  • CausalPFN’s ignorability-based prior avoids observationally equivalent DGPs with different effects; the paper reports better PEHE and ATE relative error than Do-PFN on standard benchmarks.The comparison is reported in Table 8.Source: E Concurrent Work on PFNs for Causal… · p. 40, Table 8 · p. 40

Related papers

Loading 2506.07918v2…