Unit-level aggregation breaks coverage in causal foundation models

Gabe Schulman; Lucius E. J. Bynum; Yoav Wald

stat.MLOct 16, 2026

Unit-level aggregation breaks coverage in causal foundation models

Authors: Gabe Schulman, Lucius E. J. Bynum, Yoav Wald

Organizations: New York University · Technion

Abstract

Causal foundation models (CFMs) show promising results for treatment-effect estimation, with unit-level estimates that are competitive on standard benchmarks. Many causal analyses, however, inform decisions through effects averaged over groups of units, such as the full sample, the treated units, or a subgroup, and for such decisions a CFM is useful only if its credible intervals are calibrated. We examine these intervals and find that their coverage degrades, both in simple synthetic settings and on standard benchmarks. On ACIC 2016, for example, temperature-calibrated 95% intervals from CausalPFN fall below nominal coverage beyond groups of 30 out of 4802 units, and intervals for the full-sample average contain the true value in only 20% of datasets. We trace this problem to how these models are trained and how their intervals are formed. They output a predictive distribution for one unit at a time and, at inference time, combine the unit-level variances as if the errors on different units were independent. We show formally that posterior uncertainty over the data-generating process makes these errors correlated. As a result, the combined variance is too small by a factor that grows with the size of the group. Empirically, this factor predicts where coverage fails and by how much. We then show that two neural process models, which make joint predictions for a group of units, account for this correlation and substantially improve coverage on synthetic tasks and, as heads on a frozen CausalPFN, on all six benchmarks.

Figures & tables

Figure 1: Coverage of nominal

95\%

subset-CATE intervals from temperature-calibrated CausalPFN on standard benchmarks, with unit-level aggregation. Subsets are drawn uniformly at random. The horizontal axis is the fraction of the sample aggregated, and its right edge is the full sample. Thin dotted lines are the prediction of Equation ( 1 ), with

m_{S}

and

\rho_{S}

measured once per benchmark and held fixed across

k

, and the horizontal dashed line marks nominal coverage.

dataset	covariates	source	$n$	$d$
Synthetic family	simulated $\mathcal{N}(0,I)$	Appendix A	500	3
Hill-A	simulated $\mathcal{N}(0,I)$	Hill (2011) , setting A	1000	25
Friedman	simulated $\mathrm{U}(0,1)$	Friedman (1991) , confounded	1000	10
IHDP	real, IHDP study	Hill (2011) ; splits of Shalit et al. (2017)	747	25
Hill-B	simulated $\mathcal{N}(0,I)$	Hill (2011) , setting B	1000	25
Kang–Schafer	simulated, nonlinear in 4 latents	Kang and Schafer (2007)	1000	4

Table 1: Datasets used in our experiments. Our models are trained and evaluated on the synthetic family, our heads on CausalPFN are evaluated on the six benchmarks, and CausalPFN is evaluated on all seven. Hill-A, Hill-B, Friedman, and Kang–Schafer are our reimplementations (Appendix A ).

d

is the number of covariates supplied to the estimator.

Figure 2: Unit-level and SCM-marginalized aggregation from one CNP and one latent NP, with uniformly drawn subsets at

s=1

, showing interval width and coverage against subset size.

Figure 3: Coverage of nominal

95\%

subset-CATE intervals on the six benchmarks, with subsets drawn uniformly at random. Left: CausalPFN with unit-level aggregation, as in Figure 1 . Middle and right: SCM-marginalized aggregation from the CNP and latent NP heads on the frozen CausalPFN, with lines showing the mean over three seeds and bands their range. The heads condition on the whole sample, while the CausalPFN curves come from its public interface, which on ACIC 2016 conditions groups of units on their nearest neighbours (Appendix F ). The figure compares aggregation rules on the same benchmarks, not the accuracy of the models.

Figure 4: Measured coverage against the variance ratio

R

of Proposition 1 , pooled across models, values of

s

, subset selection rules, and subset sizes. The curve is Equation ( 1 ), with

m_{S}

and

\rho_{S}

measured from unit-level errors. SCM-marginalized points are plotted at the

R

of unit-level aggregation.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

model	data	$m_{S}$	$k_{\max}$	interval	cov, smallest $k$	cov, $k_{\max}$
CausalPFN	$s=1$ , all units	21.90	500	reported	1.00	0.94
				rescaled	0.91	0.41
				$+$ design effect	0.91	0.93
CausalPFN	$s=4$ , bottom 10%	3.63	44	reported	0.99	0.79
				rescaled	0.96	0.58
				$+$ design effect	0.96	0.93

Appendix

Table 2: Coverage of the three intervals of Appendix C at the smallest and largest subset in each setting.

k_{\max}

is the largest evaluated subset size within the pool. The

+

design effect interval uses

\rho_{S}

and

m_{S}

measured against the true effects and is not available in practice.

subset	$\rho_{S}$	$m_{S}$	$k_{\max}$	$1+(k_{\max}\!-\!1)\rho_{S}$	$\hat{k}^{*}$
$s=1$ , propensity stratum
all units (baseline)	0.034	21.90	500	17.9	398
bottom 50%	0.069	23.23	250	18.3	216
bottom 25%	0.120	14.10	125	15.9	83
bottom 10%	0.224	8.07	44	10.6	34
$s=1$ , nested covariate cell

Appendix

Table 3: Subset selection rules on CausalPFN at

n=500

, with intervals formed by unit-level aggregation, over

64

tasks with

48

subsets per task.

k_{\max}

is the largest evaluated subset size within the pool and

1+(k_{\max}\!-\!1)\rho_{S}

the design effect at that size; “—” means coverage never falls below nominal within the pool.

\hat{k}^{*}

is defined in Section 5.1 .

benchmark	$n$	$\rho_{S}$	$m_{S}$	nRMSE	cov at $k$ = $n$	se	$\hat{k}^{}$ ( $\hat{k}^{}/n$ )
Hill-A	1000	0.0771	35.13	0.07	0.82	0.06	502 (50%)
Friedman	1000	0.0334	11.75	0.10	0.70	0.07	383 (38%)
IHDP	747	0.0420	5.43	0.11	0.60	0.08	85 (11%)
Hill-B heavy-tailed	1000	0.0077	2.03	0.30	0.53	0.08	62 (6.2%)
Hill-B	1000	0.0038	1.42	0.28	0.53	0.08	58 (5.8%)
Kang–Schafer	1000	0.0185	1.43	0.09	0.30	0.07	21 (2.1%)

Appendix

Table 4: The benchmarks of Figure 1 and the heavy-tailed variant of Hill-B, which the figure omits: CausalPFN with its temperature calibration, unit-level aggregation, subsets drawn uniformly at random, 40 datasets per benchmark.

\hat{k}^{*}

is the subset size at which observed coverage first falls below nominal, interpolated in

\log k

, with

\hat{k}^{*}/n

as a percentage of the sample. nRMSE is the root-mean-square error of the per-unit effect estimates divided by the standard deviation of the observed outcome, within each dataset and averaged over datasets. Rows are ordered by

\hat{k}^{*}/n

model	aggregation	$k$ =5	$k$ =50	$k$ =500	$k$ =2000
CNP	unit-level	0.93	0.61	0.29	0.17
latent NP	SCM-marginalized	1.00	0.98	0.97	0.97
CNP	SCM-marginalized	0.95	0.98	0.98	0.98

Appendix

Table 5: Coverage of nominal

95\%

subset-CATE intervals at

n=2000

over 96 tasks with 24 subsets each. The two CNP rows are separately trained models. Each model was trained once, with optimization settings that differ from Appendix B and from each other, so the comparison is not controlled for optimization.

$n$	$\rho_{S}$	$m_{S}$	$\hat{k}^{*}$	$1+(n-1)\rho_{S}$	cov at $k$ = $n$
125	0.0385	7.03	—	6	0.95 $\pm$ 0.02
250	0.0422	14.25	—	12	0.96 $\pm$ 0.02
500	0.0339	21.90	398	18	0.94 $\pm$ 0.03
1000	0.0230	32.67	980	24	0.95 $\pm$ 0.02
2000	0.0230	40.08	1376	47	0.93 $\pm$ 0.03

Appendix

Table 6: CausalPFN on the synthetic family of Appendix A at five sample sizes, with subsets drawn uniformly at random, 64 tasks and 48 subsets per task; coverage at

k=n

is given

\pm

one standard error clustered by task. The same CausalPFN model is used throughout, so no training choice varies down the table.

1+(n-1)\rho_{S}

is the design effect at the full sample; “—” means coverage never falls below nominal.

\hat{k}^{*}

is defined in Section 5.1 .

	CausalPFN	CNP head		latent NP head
benchmark	unit-level	unit-level	SCM-marg.	unit-level	SCM-marg.
Hill-A	0.82	0.18 [.12, .20]	0.94 [.88, .97]	0.62 [.57, .70]	0.92 [.85, .97]
Hill-B	0.53	0.30 [.23, .35]	0.83 [.78, .95]	0.42 [.25, .57]	0.85 [.68, .97]
Friedman	0.70	0.23 [.20, .25]	0.97 [.90, 1.00]	0.68 [.68, .70]	0.90 [.82, .97]
Kang–Schafer	0.30	0.38 [.35, .42]	1.00 [1.00, 1.00]	0.72 [.68, .78]	0.98 [.97, 1.00]
IHDP	0.60	0.30 [.20, .42]	0.94 [.90, 1.00]	0.48 [.47, .50]	0.89 [.80, .95]

Appendix

Table 7: Coverage of nominal

95\%

intervals for the full-sample average (

k=n

) on the six benchmarks, 40 datasets each. CausalPFN’s column is the full-sample column of Table 4 . The heads’ columns give the mean over three seeds, with the range of the per-seed values in brackets. Standard errors clustered by dataset are at most

0.06

for the heads. The last row counts datasets whose full-sample interval contains the true value.

benchmark	model	$m_{S}$	$\rho_{S}$	nRMSE	width, own	width, unit-level
Hill-A	CausalPFN	35.13	0.0771	0.07	—	0.055
	CNP head	2.83	0.2367	0.06	0.092	0.012
	latent NP head	24.19	0.1613	0.06	0.103	0.038
Hill-B	CausalPFN	1.42	0.0038	0.28	—	0.047
	CNP head	0.84	0.0183	0.52	0.184	0.050
	latent NP head	0.97	0.0122	0.52	0.198	0.068

Appendix

Table 8: Unit-level calibration, error correlation, accuracy, and interval width on the six benchmarks.

m_{S}

and

\rho_{S}

use the estimator of Appendix B , from each model’s unit-level means and standard deviations (for the CNP head, the head queried with each unit alone), and nRMSE is defined as in Table 4 . CausalPFN’s

m_{S}

\rho_{S}

, and nRMSE are those of Table 4 . Widths are mean widths of the full-sample interval in units of

\mathrm{sd}(y)

: own is the head’s SCM-marginalized interval, and unit-level is unit-level aggregation from the same model. CausalPFN’s width is computed with the whole sample as context, at the temperatures fitted for Table 4 . Heads: mean over three seeds.

Figure 5: Fixed-weight comparison for the CNP head on the six benchmarks, with subsets drawn uniformly at random: SCM-marginalized aggregation and unit-level aggregation from the same weights. Lines show the mean over three seeds and bands their range. Dotted lines average the three seeds’ predictions from Equation ( 1 ) for unit-level aggregation, each with

m_{S}

and

\rho_{S}

measured once per benchmark.

Figure 6: As Figure 5 , for the latent NP head.

Causal Foundation Models

Causal foundation models are pretrained models that discover causal graphs or estimate effects on new datasets using data and causal assumptions.

See all 22 papers

Causal Foundation Models
Sep 2, 2026Christopher Stith, Hossein Rahmani, Jesse C. CresswellCausal Foundation Models Causal Effect Estimation
Layer 6 AI, Toronto, Canada · TD Bank Group, Toronto, Canada
What You Observe Determines How You Identify Causal Effects: Evaluating Causal Models across Observational Views
Sep 29, 2026Heejin Jung, Gyeongdeok Seo, Hoyoon Byun +2Causal Effect Estimation Causal Foundation Models
Yonsei University · University of Illinois Urbana–Champaign
Use What You Know: Causal Foundation Models with Partial Graphs
Feb 16, 2026Arik Reuter, Anish Dhir, Cristiana Diaconu +6Causal Foundation Models Causal Discovery
University of Cambridge, Cambridge, United Kingdom · Max Planck Institute for Intelligent Systems, Tübingen, Germany · Gatsby Computational Neuroscience Unit, University College London +5

Sep 2, 2026J/K move · Enter open · S save

Causal Foundation Models

Christopher Stith, Hossein Rahmani, Jesse C. Cresswell

Layer 6 AI, Toronto, Canada · TD Bank Group, Toronto, Canada

Causal Foundation Models Causal Effect Estimation Causal Inference

Open paper AskPDF arXiv

Causal inference is the practice of estimating the effect of a treatment or intervention from data. It traditionally requires a bespoke pipeline for every new problem: first proposing a causal mechanism, selecting a compatible estimator, and finally training it. Meanwhile, across diverse settings and modalities, much of machine learning has shifted to the paradigm of foundation models: networks pretrained once at scale and applied to new tasks without fine-tuning. Causal foundation models (CFMs) bring this paradigm to causal inference. CFMs are pretrained neural networks that estimate causal quantities, such as the average treatment effect, on entirely new datasets using in-context learning without requiring model updates. This work provides a practical introduction to this emerging area. We summarize the necessary background in causal inference and machine learning before discussing CFMs. Throughout, we include example code and Jupyter notebooks.

Abstract

Causal Foundation Models

Explore similar work

Causal Foundation Models

What You Observe Determines How You Identify Causal Effects: Evaluating Causal Models across Observational Views

Use What You Know: Causal Foundation Models with Partial Graphs