Can We Trust the Uncertainty of Causal Foundation Models?

Key takeaway: Causal foundation models promise fast, flexible treatment-effect estimation, but their uncertainty can be misleading and even contain a confounding bias. In our ICML 2026 paper, we show how one-step posterior correction can calibrate these models so that their uncertainty becomes reliable.

Imagine a public health team wants to know whether a strict lockdown reduced infection rates, or a clinician wants to know whether a new treatment improves patient outcomes. In both cases, the central question is causal: What would happen if we choose a different intervention?

Machine learning can help answer such questions from observational data, but only if we are careful. It is not enough to output a single number such as “the treatment reduces infections by 10%”. In high-stakes settings, we should only trust a model's output when we can be sufficiently certain, that it is correct. Therefore, we need to know how uncertain the model is. A method that makes wrong predictions but is very confident about them might be more dangerous than a method that is less confident but acknowledges what it does not know.

Given the importance of quantifying uncertainty, I ask a simple but important follow-up question:

Even if methods are able to report uncertainty, can we actually trust it?

Our recent ICML 2026 paper, “Frequentist Consistency of Prior-Data Fitted Networks for Causal Inference”¹, aims to answer this question wrt. the new class of foundation models for tabular and causal inference, called prior-data fitted networks (PFNs). The short answer is: Not always out of the box. But with a suitable calibration step, we can make PFN-based uncertainty behave much more like the uncertainty from classical causal frequentist estimators.

Foundation Models, but for Datasets

When people hear “foundation model”, they often think of large language models that take text as input and generate text as output. PFNs follow a related idea, but the “prompt” is not a sentence. Instead, the prompt is an entire dataset.

At a high level, a PFN is trained as follows:

1. One specifies a prior over possible data-generating processes.
2. Many synthetic datasets are sampled from this prior.
3. A neural network is trained, often a transformer, to make predictions from these datasets.
4. At test time, we feed a new dataset into the trained network and obtain predictions for a new datapoint in a single forward pass. Here, the predictions approximate posterior predicitive density (PPD) of the outcome -- a quantifier of the total uncertainty.

This is attractive for tabular data and causal inference because the model behaves like an amortized Bayesian inference engine: Rather than fitting a new model from scratch for every dataset, the PFN has already learned from many simulated datasets how inference should look.

For causal inference, recent models such as TabPFN²-style learners, CausalPFN³, and CausalFM⁴ suggest a compelling possibility: Feed in an observational dataset and get treatment-effect estimates, together with uncertainty, almost instantly (see a generic procedure in Fig. 1).

But there is a catch. PFNs are trained on synthetic data sampled from a prior. This prior is helpful, because it gives the model inductive bias. Yet it can also be too opinionated.

Fig. 1. Naїve application of prior-data fitted networks (PFNs) for average treatment effect estimation.

The Causal Quantity: Average Treatment Effect

Here, we focus on the average treatment effect (ATE). Suppose each unit of measurement (e.g., a patient) has covariates (X), receives a binary treatment (A), and has an observed outcome (Y). In potential-outcomes notation, (Y[1]) is the outcome that would be observed under treatment and (Y[0]) the outcome that would be observed without treatment. The ATE is

$$
E[Y[1] - Y[0]].
$$

It answers a population-level question: What is the average effect of switching everyone from control to treatment?

The challenge is that, for each person, we observe only one of the two potential outcomes. If a patient received the treatment, we observe what happened under treatment, but not what would have happened without it. This missing counterfactual is the reason causal inference requires assumptions.

Under standard assumptions such as consistency, overlap, and no unmeasured confounding, the ATE can be identified from observational data [1, 4] using two important nuisance functions:

• the outcome model, the expected outcome given covariates and treatment;
• the propensity score, the probability of receiving treatment given covariates.

They are called nuisance functions: On the one hand, they are not the main target of inference, and, on the other, good ATE estimation heavily depends on learning these nuisance functions well.

Why Uncertainty Is Harder Than Prediction

A PFN can produce predictive uncertainty, for example a posterior predictive distribution for an outcome. This is already useful. However, causal uncertainty is more subtle. Specifically, we want to know the uncertainty about a causal functional, the ATE. This uncertainty depends on how errors in the outcome model and the propensity model propagate into the final treatment-effect estimate.

Classical causal inference has long studied how to quantify uncertainty around treatment-effect estimates. A useful reference point is the family of efficient frequentist estimators, such as augmented inverse probability weighting (A-IPTW). At a high level, these methods do two things: They first estimate the relevant aspects of the data, and then apply a correction that makes the final treatment-effect estimate statistically well-behaved as the dataset grows.

Bayesian methods approach uncertainty differently. Instead of producing a single estimate plus a confidence interval, they produce a posterior distribution: A range of plausible treatment effects given the observed data and the model’s prior assumptions. Ideally, when we have enough data, this Bayesian uncertainty should tell a similar story as the uncertainty from classical frequentist methods. This agreement is formalized by the Bernstein–von Mises theorem.

So, for our purposes, the key question is simple: As the dataset becomes larger, does the uncertainty reported by a Bayesian PFN-based causal estimator behave like the uncertainty of a trusted classical estimator? We refer to this property as a frequentist consistency. If it fails, the model may still produce uncertainty intervals that look sophisticated, but those intervals may not be reliable for causal decision-making.

The Problem: Prior-Induced Confounding Bias

Fig. 2. Prior-induced confounding bias of different exisitng PFNs. Source: Melnychuk et al. 2026¹.

Our first finding is that existing PFNs, when used naively as Bayesian ATE estimators, can suffer from prior-induced confounding bias (see Fig. 2).

To understand the issue, recall that PFNs learn from synthetic datasets. The distribution of these synthetic datasets acts as an implicit prior. If this prior rarely contains strong confounding, then the PFN learns a world in which treatment assignment is usually not too strongly tied to the covariates that also affect the outcome.

Now suppose the real observational dataset is strongly confounded. For example, sicker patients may be more likely to receive a treatment, and they may also have worse outcomes regardless of treatment. In such a case, the raw difference in outcomes between treated and untreated groups can be very misleading.

A well-calibrated causal estimator should recognize this confounding and reflect the resulting uncertainty. A PFN with an overly strong implicit prior may instead shrink the degree of confounding toward zero. In other words, the model behaves as if the data were closer to an unconfounded randomized study than they really are.

This can bias the ATE posterior, and crucially, the bias may not disappear in the way we would like as the sample size grows. The synthetic pretraining prior can continue to influence the uncertainty of the causal estimate.

A useful analogy is a doctor who has only ever seen randomized clinical trials. If shown messy observational hospital data, the doctor might underestimate how different the treated and untreated groups really are. More patient records help, but only if the doctor’s mental model allows for such confounding in the first place.

Our Fix: One-Step Posterior Correction

To address this problem, we use a one-step posterior correction⁵ (OSPC). The idea is closely related to the bias correction used by frequentist A-IPTW estimators.

At a high level, OSPC does the following:

1. Start with a Bayesian posterior over the nuisance functions.
2. Push this posterior through the usual plug-in formula for the ATE.
3. Add a correction term based on the efficient influence function of the ATE.

The theoretical result in the paper shows that, under suitable assumptions, this correction restores frequentist consistency: The corrected ATE posterior asymptotically matches the normal distribution of the efficient A-IPTW estimator. This is the semi-parametric Bernstein-von Mises result for the calibrated PFN-based estimator.

A key benefit is that the correction can be applied on top of any pretrained PFN. We do not need to retrain the whole foundation model from scratch.

Why This Matters for Reliable AI

The broader lesson is that foundation models for causal inference should not only be judged by point-estimation accuracy. In many applications, the uncertainty is part of the decision.

A policymaker may ask: Is the estimated effect large enough to justify an intervention?
A clinician may ask: How confident are we that this treatment is beneficial?
A scientist may ask: Would another dataset lead to a meaningfully different conclusion?

For these questions, a narrow but miscalibrated uncertainty interval can create false confidence. A wider but honest interval can be more useful.

PFNs are a promising route toward fast, flexible causal inference. But our results show that their Bayesian uncertainty is not automatically reliable for causal estimands. The uncertainty must be calibrated to the target causal question. Our calibration framework is one way to do this: It combines the flexibility and finite-sample inductive bias of PFNs with the asymptotic reliability of classical semi-parametric estimators.

In short: causal foundation models should not only predict effects; they should know how uncertain those effects are.

References & Further Reading

• ¹ V. Melnychuk, V. Balazadeh, S. Feuerriegel, and R. G. Krishnan. “Frequentist Consistency of Prior-Data Fitted Networks for Causal Inference.” ICML 2026. https://arxiv.org/abs/2603.12037
• ² N. Hollmann et al. “TabPFN: A Transformer That Solves Small Tabular Classification Problems in a Second.” ICLR 2023.
• ³ V. Balazadeh et al. “CausalPFN: Amortized Causal Effect Estimation via In-Context Learning.” NeurIPS 2025.
• ⁴ Y. Ma, D. Frauen, E. Javurek, and S. Feuerriegel. “Foundation Models for Causal Inference via Prior-Data Fitted Networks.” ICLR 2026.
• ⁵ Yiu, A., Fong, E., Holmes, C., & Rousseau, J. (2025). Semiparametric posterior corrections. Journal of the Royal Statistical Society Series B: Statistical Methodology, 87(4), 1025-1054.

About the Author

Valentyn Melnychuk, MSc is PhD candidate specializing in causal machine learning and data-driven decision-making at Prof. Stefan Feuerriegel's group at the Institute of AI in Management (LMU Munich). As a member of relAI, his work focuses on the intersection of uncertainty quantification, robust causal inference, and reliable AI systems.