Can We Trust the Uncertainty of Causal Foundation Models?
Date:
Location: relAI blog
Causal foundation models promise fast, flexible treatment-effect estimation, but their uncertainty can be misleading and even carry a confounding bias. This post looks at prior-data fitted networks (PFNs) — foundation models trained on synthetic datasets to perform amortized Bayesian inference — and at how the priors they were trained on leak into the posteriors they report. The result is a prior-induced confounding bias: implicit assumptions baked into the training distribution make the model understate the confounding actually present in real observational data, so its credible intervals are narrower and more confident than they should be. I then discuss a one-step posterior correction, a calibration step that makes PFN-based uncertainty behave much more like the uncertainty of classical frequentist causal estimators, restoring trustworthiness in high-stakes settings such as medicine.
Most of causal machine learning stops at the conditional average treatment effect (CATE), but an average hides the inherent randomness in how individuals respond to a treatment — and in medicine that randomness is exactly what decides whether a therapy is safe. This post, based on our NeurIPS 2024 paper, explains why the distribution of the treatment effect is hard to get at: it is a counterfactual quantity (we never observe both potential outcomes for the same patient), and observational data is confounded on top of that. We introduce the AU-learner, which combines Makarov bounds with normalizing flows to estimate the range of plausible treatment-effect distributions instead of pretending a point estimate is enough — making visible the substantial minority of patients who may be harmed by a treatment that looks beneficial on average.