Calibrated Inference for the Conditional Average Treatment Effect in the Few-Placebo Regime via Gaussian Processes

Technical Summary: Calibrated Inference for the CATE in the Few-Placebo Regime via Gaussian Processes

Problem Statement
The paper addresses the challenge of estimating the Conditional Average Treatment Effect (CATE), $\tau(x) = E[Y(1) - Y(0) | X=x]$, with calibrated uncertainty intervals in the "few-placebo regime." This regime is defined by a severe imbalance in sample sizes between treatment arms ($N_0 \ll N_1$), a scenario common in unequal-allocation clinical trials, budget-constrained A/B tests, and rare-regime studies.

While the X-Learner is the standard meta-learner for imbalanced arms, it provides only a point estimate. The authors investigate the natural extension of making the X-Learner's second stage Bayesian to generate credible intervals. They demonstrate that this approach fails to provide calibrated coverage (i.e., the intervals contain the true effect less often than their nominal level) specifically in the few-placebo regime.

Diagnosis of Failure
The paper identifies a structural cause for the under-coverage of Bayesian X-Learner intervals:

1. First-Order Nuisance Bias: The X-Learner constructs pseudo-outcomes (e.g., $D_1 = Y - \hat{\mu}_0(X)$) to estimate the treatment effect. Because machine learning nuisance models (like regularized trees) trade bias for variance, the estimate $\hat{\mu}_0$ is not conditionally unbiased. Consequently, the pseudo-outcome's conditional mean is $\tau(x) - b_0(x)$, where $b_0(x)$ is the regularization bias of the nuisance model fitted on the small control arm.
2. Posterior Displacement: A Bayesian regression on these biased pseudo-outcomes produces a posterior centered on the biased target ($\tau - b_0$) rather than the true effect. Since the bias is a location shift, increasing the posterior variance (widening the interval) does not correct the coverage; the interval remains displaced from the truth.
3. Failure of Orthogonal Scores: The standard remedy for nuisance bias is regressing a Neyman-orthogonal, doubly-robust (DR) score. The authors show this also fails in the few-placebo regime. The regime inherently induces poor overlap (small propensity for the control arm, $1-\pi \approx 0$). The DR score's "second-order" remainder is amplified by a factor of $1/(1-\pi)^2$, rendering the remainder non-negligible. Furthermore, applying inverse-variance weighting to stabilize the high variance of the control arm disrupts the cross-arm cancellation mechanism that ensures unbiasedness, effectively reverting the estimator to the biased X-Learner form.

Methodology: GP-CATE
To resolve these issues, the authors propose GP-CATE, which abandons the point-estimate-plus-post-hoc-variance paradigm in favor of a fully Bayesian approach for the nuisance functions themselves.

• Modeling: Instead of fitting point estimates $\hat{\mu}_0$ and $\hat{\mu}_1$, GP-CATE models the outcome surface for each arm ($\mu_0$ and $\mu_1$) using independent Gaussian Processes (GPs).
• Posterior Construction: The GP returns a full posterior distribution over the function, consisting of a mean $m_w(x)$ and a pointwise standard deviation $s_w(x)$. The uncertainty $s_w(x)$ is naturally large where data is sparse (the control arm) and small where data is dense.
• CATE Inference: The CATE posterior is derived as the difference of the two arm posteriors:
  $$\tau(x) | D \sim \mathcal{N}(m_1(x) - m_0(x), s_1(x)^2 + s_0(x)^2)$$
  This construction ensures that the uncertainty of the small control arm enters the CATE posterior directly as variance, rather than as an unmodelled bias.
• Computational Strategy: While exact GPs are $O(N^3)$, the authors note that in the few-placebo regime, the control arm ($N_0$) is small, making its exact fit computationally negligible. For the large treatment arm ($N_1$), a scalable surrogate (e.g., sparse/inducing-point GP) is sufficient because the control arm's uncertainty dominates the total CATE variance.

Key Contributions

1. Diagnosis: The paper establishes that under-coverage in Bayesian X-Learners for imbalanced data is caused by first-order nuisance bias from non-orthogonal pseudo-outcomes, which displaces the posterior mean.
2. Negative Result: It demonstrates that the doubly-robust/orthogonal-score remedy fails in this regime due to overlap amplification and a fundamental conflict between efficiency weighting and unbiasedness, even when the propensity score is known.
3. Root-Cause Analysis: The authors argue that attaching a variance to a point estimate of a "hard-to-learn" nuisance (one fitted on sparse data) is structurally flawed because regularization bias is invisible to post-hoc variance estimates.
4. Method and Empirical Validation: GP-CATE is proposed as the solution, replacing point estimates with Bayesian posteriors. Empirical results show it is the only method tested that achieves calibrated coverage across linear, non-linear, and semi-synthetic (IHDP) benchmarks in the few-placebo regime.

Experimental Results
The authors evaluated GP-CATE against Causal Forest, BART, and standard meta-learners (X-Learner, DR-Learner) on three designs:

• Coverage: In the few-placebo regime ($N_0=30$), standard interval estimators severely under-cover. Causal Forest coverage dropped to as low as 0.34 on the IHDP benchmark, and BART dropped to 0.59. In contrast, GP-CATE maintained coverage between 0.94 and 0.99 across all designs and sample sizes.
• Interval Width: The calibrated intervals produced by GP-CATE are significantly wider than those of the alternatives in high-dimensional settings with few controls. The authors argue this is a feature, not a bug: the wide intervals correctly reflect the high uncertainty inherent in estimating a surface from sparse data, whereas narrower intervals from other methods are miscalibrated.
• Point Accuracy: GP-CATE remains competitive in point estimation accuracy (measured by $\sqrt{\epsilon PEHE}$), often outperforming or matching Causal Forest and BART, particularly as $N_0$ increases.

Significance and Claims
The paper claims that the few-placebo regime represents a structural boundary where standard asymptotic oracle properties (which assume fixed overlap) break down. The significance of the work lies in identifying that calibrated inference in this regime requires a fully Bayesian nuisance model rather than a point estimate with a corrected variance.

The authors modestly frame their contribution not as a novel estimator architecture (noting that a Bayesian T-learner with GPs is similar to prior work by Alaa and van der Schaar, 2017), but as a diagnostic insight: the failure of meta-learners in this regime is due to the separation of bias and variance estimation. GP-CATE is presented as the estimator that the root-cause analysis necessitates. The paper concludes that practitioners should prefer GP-CATE in imbalanced settings, accepting wider intervals as the honest reflection of data scarcity, rather than relying on narrower, miscalibrated intervals from standard methods.