PaperlayerView exact revision on arXiv

Source-linked AI summary

Unit-level aggregation breaks coverage in causal foundation models

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

arXiv:2610.22770v1stat.MLcs.LG

TL;DR

Causal foundation models estimate unit-level effects, but many decisions depend on group averages whose credible intervals need calibrated coverage. This paper traces falling coverage to correlated errors from shared SCM uncertainty and evaluates neural process models that predict group effects jointly. Those models substantially improve full-sample coverage across six benchmarks, though they do not always attain nominal coverage.

  • Problem

    Many causal decisions concern averages over groups, but current CFM intervals for these effects lack a coverage guarantee and fall below nominal on tested benchmarks.

    Source: §1 · pp. 1–2

  • Method

    The paper analyzes how shared SCM uncertainty correlates unit-effect errors and tests latent and direct-decode neural process models that jointly account for dependence across subset units.

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

  • Results

    Across six benchmarks, CNP and latent-NP heads on frozen CausalPFN covered the full-sample average in 90% and 89% of 240 datasets, versus 25% and 52% for unit-level aggregation and 53% for CausalPFN.

    Source: §1 · p. 3

  • Takeaways & Limitations

    Joint SCM-marginalized predictions substantially improve aggregate-effect coverage on synthetic tasks and six benchmarks, but do not always reach nominal coverage.

    Source: §7 · p. 10

    As heads on frozen CausalPFN, the models recover most but not all lost coverage, and their unit-level estimates are less accurate than CausalPFN’s on three of six benchmarks.

    Source: §6 · p. 10

Abstract

from arXiv · show

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.

1 INTRODUCTION

CFMs estimate unit-level effects efficiently, but many decisions require calibrated uncertainty for averages over groups. The paper links degrading group coverage to correlated errors from shared SCM uncertainty and proposes joint-prediction models.

1 INTRODUCTION

  • Group averages matter for decisions, yet credible intervals from current CFMs fall below nominal coverage on every tested benchmark, including for the full sample.Averages over the full sample, treated units, or subgroups can support applied decisions such as subgroup analyses.Source: §1 · pp. 1–2
  • Shared posterior uncertainty over the SCM makes unit-effect errors correlated, so averaging shrinks only their unshared component and unit-level aggregation understates subset-CATE variance by the design effect.For fixed reported variance and Gaussian errors, coverage tends to zero as group size grows when ρS > 0.Source: §1 · p. 2
  • On ACIC 2016, with mS = 1.63 and ρS = 0.02, predicted full-sample coverage was 20%, matching the observed 20%; coverage fell below nominal at 32 predicted versus 30 observed units.Temperature-calibrated CausalPFN intervals for the full-sample average contained the truth in only 8 of 40 datasets.Source: §1 · p. 2
  • The proposed latent neural process and direct-decode conditional neural process integrate over the SCM jointly across subset units, accounting for their dependence.They are trained on a synthetic family and as heads on a frozen CausalPFN.Source: §1 · p. 3
  • Across six benchmarks, CNP and latent-NP heads on frozen CausalPFN covered the full-sample average in 90% and 89% of 240 datasets, versus 25% and 52% for unit-level aggregation and 53% for CausalPFN.These models substantially improved coverage but did not always reach nominal coverage.Source: §1 · p. 3
  • On the synthetic family, Proposition 1 predicted coverage with a median absolute error of 0.05.The prediction uses the design-effect analysis of unit-level aggregation.Source: §1 · p. 3

2 RELATED WORK

The related work spans group-effect inference, design effects, amortized uncertainty quantification, and causal foundation models. This paper distinguishes its focus on correlated estimation errors in subset-CATE intervals from prior approaches to aggregate inference.

2 RELATED WORK

  • Non-amortized methods retain unit dependence by averaging within posterior draws or targeting group effects with orthogonal scores (Semenova & Chernozhukov, 2021).Conformal methods target individual treatment effects rather than group averages (Lei & Cand`es, 2021; Alaa et al., 2023; Memmesheimer et al., 2025).Source: §2 · p. 3
  • The design effect, introduced for complex surveys, also describes variance inflation from correlated units in cluster sampling and cluster-randomized trials.Here, correlation arises from posterior uncertainty over the SCM rather than the sampling design.Source: §2 · p. 3
  • PFNs approximate posterior inference from synthetic tasks (M¨uller et al., 2022), while prior work notes pointwise procedures do not provide joint inference across query points (Nagler & R¨ugamer, 2025).Other work obtains joint draws with martingale posteriors and studies prior-induced bias in aggregate-effect posteriors (Melnychuk et al., 2026).Source: §2 · pp. 3–4
  • Causal foundation models perform competitively on standard benchmarks, but those reporting uncertainty typically provide unit-level intervals or aggregate independent unit-level draws.This paper quantifies coverage loss for arbitrary subsets and models dependence with a latent neural process or a conditional neural process that directly predicts subset averages.Source: §2 · p. 4

3 SETUP

The setup defines a subset-CATE as an average of unit-level CATEs over queried units and frames CFMs as amortized estimators trained on synthetic causal tasks. It evaluates uncertainty through interval coverage under stated identification assumptions.

3 SETUP

  • A subset-CATE averages unit-level CATEs at the realized covariates of a queried set, making it a finite-sample estimand distinct from the sample average treatment effect and the PATE.For the full sample it is a sample analogue of the PATE; for treated units it is analogous to the ATT.Source: §3 · p. 4
  • CFMs are trained across tasks containing an SCM, a dataset, and a queried subset; current models typically query one unit and return a predictive distribution for its treatment effect.Estimating a larger subset-CATE with an interval requires combining unit-level means and variances.Source: §3 · p. 4
  • For symmetric credible intervals, the reported subset-CATE estimate is expanded by a standard-normal quantile times the model-reported standard deviation.At α = 0.05, the stated quantile is approximately 1.96.Source: §3 · p. 4
  • Coverage is the fraction of queries whose interval contains the true subset-CATE, with evaluation probability taken over SCMs, datasets, and subsets.The subset may be sampled uniformly or through a selection rule.Source: §3 · p. 5
  • Identification assumes observed confounders and positivity; although benchmarks have no unobserved confounding, IHDP and some ACIC 2016 settings have limited overlap and possible extrapolation error.The synthetic family satisfies the identification assumptions by construction.Source: §3 · p. 5

4 WHY UNIT-LEVEL AGGREGATION BREAKS COVERAGE

Unit-level aggregation can understate subset-CATE uncertainty because posterior uncertainty over a shared SCM induces cross-unit error covariance. The section derives coverage consequences and describes SCM-marginalized joint prediction as an alternative.

4.1 THE DESIGN EFFECT

  • The analysis uses raw error second moments, including bias, and defines them over the queried subset because selected subgroups may have more correlated errors.For uniformly drawn subsets, the defined moments do not depend on subset size.Source: §4.1 · p. 5
  • Under Gaussian errors, coverage is nominal at R = 1 and below nominal for R < 1; calibrated unit-level intervals retain nominal coverage only for single units.Wider unit-level intervals can retain nominal coverage up to k*, with k*−1 inversely proportional to ρS.Source: §4.1 · p. 6
  • When reported variance varies across queries, Equation (1) is an approximation because R is a ratio of expectations.Equations (8) and (9) hold for any error distribution.Source: §4.1 · p. 6

4.2 THE SOURCE OF THE ERROR CORRELATION

  • Posterior uncertainty over a shared SCM induces dependence among unit effects; for the posterior mean, error correlation comes from covariance across posterior SCM draws.The units are conditionally i.i.d. given the SCM, but integrating over a random SCM yields a joint distribution that does not factorize across units.Source: §4.2 · p. 6
  • Unit-level marginals do not determine cross-unit covariance, so recovering that dependence requires a joint model of unit effects or repeated forward passes.Source: §4.2 · p. 6

4.3 SCM-MARGINALIZED AGGREGATION

  • For the exact posterior, SCM-marginalized aggregation includes cross terms and gives R = 1, yielding exact coverage at every subset size k.CausalPFN’s aggregate estimator omits those cross terms by sampling unit effects independently.Source: §4.3 · p. 7
  • A latent neural process uses a learned latent variable to capture dependence, trained on unit-level and subset-level likelihood terms and sampled jointly at inference.The subset-level term directly trains dependence through variation in within-draw averages.Source: §4.3 · p. 7
  • A direct-decode CNP predicts a Gaussian distribution for the selected subset effect and can handle new subsets without retraining.Its training objective yields the posterior mean and variance under its training prior.Source: §4.3 · p. 7

5 EXPERIMENTS

The experiments compare CausalPFN’s unit-level aggregation with SCM-marginalized aggregation on synthetic tasks and six benchmarks. Joint aggregation substantially improves group-interval coverage, while the variance-ratio analysis tests how well the predicted coverage tracks observed failures.

5.1 SETUP

  • CausalPFN is used without retraining, with built-in per-dataset temperature calibration on benchmarks; the CNP and latent NP heads are trained on a broader prior without benchmark training or checkpoint selection.The authors’ own models are also tested on a synthetic family with n = 500, where subset-CATE targets have a closed form.Source: §5.1 · pp. 7–8
  • The setup evaluates 40 datasets per benchmark and compares aggregation rules on shared subsets, weights, context, and targets in the fixed-weight experiment.The study reports nominal 95% interval coverage and the subset size at which observed coverage first falls below nominal.Source: §5.1 · p. 8

5.2 RESULTS

  • Coverage begins at or above nominal on every benchmark, then falls below nominal as subsets grow, fastest on ACIC 2016 and slowest on Hill-A.Hill-A falls below nominal only beyond half the sample, while Hill-B loses coverage at 5.8% of its sample despite a correlation twenty times smaller than Hill-A’s.Source: §5.2 · p. 8
  • At the largest synthetic subsets, unit-level intervals are 7.7 times narrower for the CNP and 5.5 times narrower for the latent NP than SCM-marginalized intervals, while unit-level coverage falls with subset size.SCM-marginalized coverage stays between 0.90 and 0.98; across 60 cells its largest-subset median is 0.89, versus 0.47 as the highest unit-level value.Source: §5.2 · p. 8, Figure 2 · p. 9
  • 90% and 89% of full-sample intervals cover the true value across 240 benchmark datasets for the CNP and latent NP heads, versus 25%, 52%, and 53% for unit-level aggregation and CausalPFN.The comparison uses heads on a frozen CausalPFN; the heads do not improve unit-level accuracy and have wider full-sample intervals than CausalPFN on every benchmark.Source: §5.2 · pp. 8–9, Figure 3 · p. 9
  • Unit-level coverage closely follows Equation (1): median absolute error is 0.043 for the authors’ models and 0.021 for CausalPFN on the synthetic family.On benchmarks, the unit-level prediction’s median absolute error is 0.029 for the CNP head and 0.042 for the latent NP head over points with R < 1.Source: §5.2 · pp. 9–10, Figure 4 · p. 10

6 DISCUSSION AND LIMITATIONS

The proposed models recover much of the benchmark coverage lost by unit-level aggregation, but the authors do not claim accuracy or generality comparable to CausalPFN. A design-effect correction also requires parameters that vary substantially by benchmark, model, and subgroup.

6 DISCUSSION AND LIMITATIONS

  • The small models are trained on a synthetic family, and the authors do not claim they match CausalPFN’s accuracy or generality.As heads on frozen CausalPFN, they recover most, but not all, lost benchmark coverage; unit-level estimates are less accurate on three of six benchmarks.Source: §6 · p. 10
  • Design-effect correction requires mS and ρS, which vary by more than an order of magnitude across benchmarks and depend on the model and subgroup.Rescaling existing models’ unit-level intervals did not help in the reported experiments and worsened coverage when those intervals were too wide.Source: §6 · p. 10

7 CONCLUSION

The paper concludes that unit-level aggregation misses posterior dependence across units, causing group-interval coverage to fall as groups grow. Jointly integrating over the SCM improves coverage on synthetic tasks and benchmarks, while the appendix describes the synthetic data-generating family used for evaluation.

7 CONCLUSION

  • Across six benchmarks, SCM-marginalized aggregation recovers much of the coverage lost by unit-level aggregation, though not always to nominal levels.The same conclusion holds on the synthetic family; the authors trace the variance shortfall to posterior dependence among unit effects under a shared SCM.Source: §7 · p. 10
  • The authors report using language models for writing assistance and related-work discovery, with authors reviewing drafted text and verifying cited work.They state that theoretical claims and experimental results are their own.Source: USE OF AI TOOLS · p. 11
  • The reproducibility materials specify the synthetic family, benchmark designs, model training and checkpoint selection, and CausalPFN code paths at a pinned commit.The paper states that code to reproduce every figure and table is available in its linked repository.Source: REPRODUCIBILITY STATEMENT · p. 11

A DATA-GENERATING PROCESSES

  • The synthetic family varies conditions under which coverage fails and permits exact computation of each subset-CATE target.Each task draws a fresh SCM and then n units independently conditional on that SCM, making units exchangeable but not independent marginally.Source: A DATA-GENERATING PROCESSES · p. 14, §A.1 · p. 14

A.1 SYNTHETIC FAMILY

  • The synthetic outcome combines a covariate response, treatment effect, and Gaussian noise; the propensity sharpness parameter s remains within positivity bounds.The models are trained only at s = 1, while the CNP is also evaluated at s = 4, outside its training distribution.Source: §A.1 · p. 14

A.2 BENCHMARKS

The appendix describes six benchmarks with simulated response surfaces, allowing interval scores to be compared with known unit-level conditional effects. The benchmark suite combines real-study covariates in IHDP and ACIC 2016 with simulated covariates in the other designs.

A.2 BENCHMARKS

  • Across the six benchmarks, response surfaces are simulated and intervals are scored against known conditional effects, with no unobserved confounding by construction.IHDP uses splits 1–40 of 100 and ACIC 2016 uses settings 1–40 of 77, one replication each; simulated designs use seeds 1–40.Source: §A.2 · pp. 14–15
  • Hill-A has 25 Gaussian covariates and a constant treatment effect of 4.Its outcome surfaces are specified as μ0 = Xβ and μ1 = Xβ + 4.Source: §A.2 · p. 15
  • Friedman uses ten uniform covariates and a nonlinear outcome surface, while Kang–Schafer uses transformed observed covariates and has zero treatment effect for every unit.The covariate designs are simulated for these benchmarks; the appendix also describes a heavy-tailed Hill-B variant using t3 noise.Source: §A.2 · p. 15

B TRAINING AND EVALUATION PROTOCOL

The CNP and latent NP both summarize sample information to predict subset effects, but differ in how they represent joint uncertainty. Their evaluations combine synthetic-task training, benchmark heads on frozen CausalPFN, and subgroup-selection tests.

Architectures

  • The CNP decodes a Gaussian for the subset effect from pooled sample and queried-subset summaries, while the latent NP samples a 16-dimensional latent and forms subset effects within each draw.Both use two-layer, width-128 ReLU MLPs without attention; the latent NP decodes unit effects as 64-bin histograms.Source: B TRAINING AND EVALUATION PROTOCOL · p. 15

Training

  • Both models train on synthetic tasks with n = 500; the CNP uses random subsets and Gaussian negative log-likelihood, while the latent NP combines unit-effect and subset-effect losses.The latent NP uses 16 latent draws, and its subset-effect loss has weight 2; training uses 30,000 steps and three seeds per model.Source: B TRAINING AND EVALUATION PROTOCOL · p. 15

Subset selection rules

  • The fixed-weight experiment tests six subset-selection rules for the CNP and three for the latent NP, with three seeds yielding 45 and 15 cells, respectively.The latent NP’s two intervals share unit-level marginals and differ only in aggregation; the CNP intervals also differ in point estimate.Source: B TRAINING AND EVALUATION PROTOCOL · p. 15

Heads on CausalPFN

  • On benchmarks, both models are heads on frozen CausalPFN and receive each unit’s hidden states under both treatments plus its unit-level effect mean and standard deviation.The CNP head has 461k parameters; the latent NP head has 621k, and each uses a 16-dimensional latent variable and 64-bin unit-effect histograms.Source: B TRAINING AND EVALUATION PROTOCOL · p. 16

Estimating coverage

  • Coverage uncertainty is clustered by task because subsets from one task share an error component; benchmark estimates use 40 datasets with 48 subsets each.The protocol also estimates ρS and mS from true effects, so those quantities are diagnostics rather than practice-ready measures.Source: B TRAINING AND EVALUATION PROTOCOL · p. 16

C CORRECTING FOR THE DESIGN EFFECT

The section compares reported intervals with unit-level rescaling and a design-effect correction. Rescaling can worsen coverage, while the design-effect correction depends on true effects and is unavailable in practice.

Correction comparison

  • For CausalPFN, rescaling its over-wide unit-level intervals lowered coverage at both the smallest subset and kmax, while the CNP’s near-calibrated unit intervals changed little and kmax coverage remained near 0.24.The CNP’s mS was 0.93.Source: C CORRECTING FOR THE DESIGN EFFECT · p. 16
  • CausalPFN temperature calibration reduced full-sample ACIC 2016 coverage from 18 of 40 datasets to 8 of 40.The calibration acts like rescaling unit-level variances.Source: C CORRECTING FOR THE DESIGN EFFECT · p. 17
  • The design-effect interval removes dependence on k but requires ρS and mS measured against true effects, making it unavailable in practice.Table 2 compares coverage of three intervals at the smallest and largest subset sizes in each setting.Source: C CORRECTING FOR THE DESIGN EFFECT · p. 17, Table 2 · p. 17

D SUBGROUP SELECTION

Subgroup-selection rules change the relationship between error correlation, pool size, and unit-level calibration. Existing CausalPFN subset intervals instead aggregate independently sampled unit predictions, and other cited models lack an aggregate estimator.

D SUBGROUP SELECTION

  • Narrower pools raise ρS but also shrink kmax, so the largest reachable design effect changes less than either quantity and falls for the narrowest pools.Narrow propensity strata lower mS, nested covariate cells leave it roughly unchanged, and covariate neighborhoods raise it.Source: D SUBGROUP SELECTION · p. 17
  • At s = 1, nested covariate cells reach ρS comparable to the bottom 25% propensity stratum yet never fall below nominal because mS stays high; at s = 4, the first two covariate cells do.The tested rules include propensity- and covariate-based subgroup selection, as well as other analyst-like selection rules.Source: D SUBGROUP SELECTION · p. 17

E HOW EXISTING MODELS CONSTRUCT SUBSET-CATE INTERVALS

  • CausalPFN constructs subset intervals by independently sampling each unit’s predictive distribution, averaging draws over the subset, and taking percentiles.The sampled aggregate distribution uses independent unit draws, so its variance omits cross-unit covariance terms.Source: E HOW EXISTING MODELS CONSTRUCT… · p. 17
  • Do-PFN and CausalFM provide no aggregate estimator; the documented uncertainty procedures also treat units independently or provide only unit-level intervals.CausalFM’s documented aggregate quantity is an average of unit-level point estimates without a credible interval.Source: E HOW EXISTING MODELS CONSTRUCT… · p. 18

F ADDITIONAL RESULTS

Additional results examine how error correlation and unit-level calibration shape coverage as sample size and benchmark conditions vary. The section also reports prediction accuracy, heavy-tailed-noise results, and benchmark comparisons for the neural-process heads.

Sample size

  • The full-sample design effect grew from 6 at n = 125 to 47 at n = 2000, as ρS fell by only 1.7-fold over the sixteenfold sample-size range.A power law fitted to five sample sizes had exponent −0.24, much slower decay than n^-1.Source: F ADDITIONAL RESULTS · p. 18, Table 6 · p. 19
  • Despite that growth, full-sample coverage stayed close to nominal across the synthetic-family sample sizes because mS rose from 7 to 40.CausalPFN’s per-unit errors shrank faster than its reported intervals; on the benchmarks, mS ranged from 1.42 to 35.1 and coverage fell below nominal on every benchmark.Source: F ADDITIONAL RESULTS · pp. 18–19, Table 6 · p. 19

Accuracy of the predictions

  • Equation (1) approximated the measured threshold ˆk* within 15% on Hill-A, Kang–Schafer, and ACIC 2016, but was about twice the measured value on Hill-B and its heavy-tailed variant.It also overestimated full-sample coverage on Hill-B, Hill-B heavy-tailed, and Kang–Schafer.Source: F ADDITIONAL RESULTS · p. 19, Table 4 · p. 19

The heavy-tailed variant

  • Replacing Hill-B’s Gaussian noise with t3 noise changed little: coverage fell below nominal at 6.2% of the sample versus 5.8% for Hill-B, and both reached 0.53 at full sample.The heavy-tailed variant is omitted from Figure 1 but included in Table 4.Source: F ADDITIONAL RESULTS · p. 19, Table 4 · p. 19

Heads on the benchmarks

  • At k = n, SCM-marginalized coverage ranged from 0.68 to 0.97 for the latent NP head on Hill-B and from 0.70 to 0.88 on ACIC 2016 across seeds.The fixed-weight comparison presents SCM-marginalized and unit-level aggregation from the same weights on six benchmarks.Source: F ADDITIONAL RESULTS · p. 20, Figure 5 · p. 21

Context on ACIC 2016

  • On ACIC 2016, CausalPFN’s full-sample coverage was 0.25 when rescored with the whole sample as context, compared with 0.20 under its public interface.On the other benchmarks, the two context settings agreed within 0.03.Source: F ADDITIONAL RESULTS · p. 20

G PROOFS

The proofs characterize when aggregate uncertainty is preserved or understated, and establish exact-coverage results under posterior and subset-selection conditions.

G PROOFS

  • The results assume data-selected subsets; subsets chosen by the true effect are not selected on the data.Source: G PROOFS · p. 21
  • With variable reported variance, coverage is E[f(R_D,S)]; f(R) upper-bounds it only when R_D,S and v_D,S are independent.Source: G PROOFS · p. 22
  • Posterior covariance between unit errors equals posterior covariance between unit effects, and unit-level marginals alone cannot determine aggregate posterior variance.Source: G PROOFS · pp. 22–23
  • SCM-marginalized aggregation uses the joint predictive law over selected units, whereas models returning only marginals support unit-level aggregation.Source: G PROOFS · p. 22
  • With data-selected S, SCM-marginalized aggregation has R = 1, and exact posterior equal-tailed intervals achieve C(k) = 1 − α for every k.Source: G PROOFS · p. 23
  • The exact-posterior coverage corollary averages over π and does not imply coverage at a fixed SCM.Source: G PROOFS · p. 23
  • Among measurable Gaussian-output maps, the population loss minimizer reports the posterior mean and variance of θS.Source: G PROOFS · p. 23

Related papers

Loading 2610.22770v1…