Robust estimation of heterogeneous treatment effects in randomized trials leveraging external data

1. Problem Statement

Randomized Controlled Trials (RCTs) are the gold standard for estimating causal effects but often suffer from limited sample sizes. While they are well-powered to detect Average Treatment Effects (ATE), they frequently lack the statistical power to accurately estimate Conditional Average Treatment Effects (CATE), which are essential for personalized medicine and understanding treatment effect heterogeneity.

To address this, researchers often attempt to augment RCT data with external data (from other trials or observational studies). However, this introduces two major challenges:

1. Misalignment/Non-Transportability: The external population may differ from the trial population (covariate shift), meaning the CATE is not stable across populations.
2. Unmeasured Confounding: External data (especially observational) may contain unmeasured confounders, violating the assumptions required for causal inference.

Existing methods often assume transportability (that the CATE is identical across populations) or require strong assumptions about the bias structure. If these assumptions fail, integrating external data can actually degrade estimation performance compared to using trial data alone. The paper aims to develop a method that leverages external data to improve CATE estimation without risking harm if the external data is misaligned or confounded.

2. Methodology: The QR-Learner

The authors propose the QR-learner (Quasi-optimized Randomization-aware learner), a model-agnostic framework designed to robustly estimate CATEs in the trial population by leveraging external data.

Core Concepts

• Randomization-Aware Pseudo-Outcomes: The method utilizes a pseudo-outcome formulation $\psi(O; \eta)$ derived from the trial's known propensity score $e(X)$. This formulation ensures that, under the trial's randomization conditions, the expected value of the pseudo-outcome equals the true CATE, regardless of the nuisance models used.
  $$\psi(O_i; \eta) = \frac{A_i - e(X_i)}{e(X_i)(1 - e(X_i))}(Y_i - h_{A_i}(X_i)) + h_1(X_i) - h_0(X_i)$$
  where $\eta = \{h_1, h_0\}$ are nuisance models estimating the conditional mean outcomes.

• Two-Stage Procedure with Cross-Fitting:

  1. Stage 1 (Quasi-Optimization): The algorithm estimates the nuisance models $\hat{\eta} = \{\hat{h}_1, \hat{h}_0\}$ by minimizing a weighted mean squared error objective. Crucially, this optimization uses both trial and external data. The weights are derived to minimize an upper bound on the variance of the model selection uncertainty. This step assumes transportability and strong ignorability in the external data to utilize the larger sample size for better nuisance estimation.
  2. Stage 2 (CATE Estimation): Using the estimated nuisance models $\hat{\eta}$ from Stage 1, the algorithm computes pseudo-outcomes $\psi$ for the trial data only. It then regresses these pseudo-outcomes on covariates $X$ to estimate the CATE $\hat{\tau}(X)$.

• Robustness Mechanism:

  • Because Stage 2 relies only on the trial data and the known propensity score, the final estimator is unbiased for the trial CATE even if the external data is misaligned or confounded (provided the trial's randomization holds).
  • If the external data is aligned, the improved nuisance estimation in Stage 1 reduces the variance of the pseudo-outcomes, leading to a lower Mean Squared Error (MSE) than using trial data alone.

The Combined Learner

To provide an additional safeguard against finite-sample risks where the QR-learner might underperform due to misalignment, the authors propose a Combined Learner:
$$\hat{\tau}_{combined}(x; \lambda) = \lambda \cdot \hat{\tau}_{QR}(x) + (1 - \lambda) \cdot \hat{\tau}_{DR}(x)$$
where $\hat{\tau}_{DR}$ is a standard Doubly Robust (DR) learner trained only on trial data. The weight $\lambda$ is selected via cross-validation. Theoretical results show this combined estimator asymptotically achieves an MSE no worse than the best of its components.

3. Key Contributions

1. The QR-Learner Algorithm: A novel, model-agnostic learner that integrates external data to improve nuisance model estimation while maintaining robustness to population misalignment and unmeasured confounding in the external data.
2. Theoretical Guarantees:
   • Proves that the QR-learner recovers the true CATE in the trial population even when external data is non-transportable or confounded (due to the randomization-aware pseudo-outcome structure).
   • Proves that the combined learner asymptotically matches or exceeds the performance of its component learners (QR and DR) in terms of MSE.
3. Model Selection Strategy: Introduces a strategy to select nuisance models by minimizing an upper bound on the variance of the pseudo-risk, explicitly leveraging external data to reduce finite-sample uncertainty.
4. Empirical Validation: Demonstrates through simulations and a real-world case study (STAR dataset) that the method improves CATE estimation accuracy and statistical power for detecting heterogeneity, even when transportability assumptions are violated.

4. Results

• Simulation Studies:

  • Robustness: When external data was misaligned (violating transportability) or confounded, the QR-learner and Combined Learner maintained MSE comparable to the trial-only DR-learner, whereas other integrative methods (like CFACE, R-OSCAR, and Kallus et al.'s bias correction) suffered significant performance degradation.
  • Efficiency: When external data was aligned, the QR-learner significantly outperformed trial-only methods, with MSE decreasing as external sample size increased.
  • Statistical Power: The method improved the power to detect treatment effect heterogeneity (interaction effects) by approximately 10–20 percentage points compared to trial-only baselines, while maintaining nominal Type I error rates.

• Case Study (STAR Dataset):

  • Applied to the Tennessee Student/Teacher Achievement Ratio project, treating rural and urban schools as distinct populations with violated transportability (due to omitted location variables).
  • The Combined Learner consistently outperformed component learners and other baselines (CFACE, R-OSCAR) across varying trial sample sizes.
  • The QR-learner showed improvements over the DR-learner in specific configurations (e.g., when rural schools were the target), demonstrating the utility of external data even in misaligned settings.

5. Significance

This work addresses a critical bottleneck in precision medicine: the inability of small RCTs to reliably identify subgroups that benefit most from treatment.

• Safety First: Unlike previous "data borrowing" methods that risk introducing bias if assumptions fail, the QR-learner is designed to be robust by construction. It guarantees that using external data will not harm the estimation compared to using trial data alone.
• Practical Applicability: It provides a practical framework for regulatory bodies and researchers to safely incorporate historical controls or observational data into RCT analyses to enhance subgroup analysis without compromising causal validity.
• Methodological Advancement: It bridges the gap between transfer learning (which often assumes domain shift is correctable) and causal inference (which requires strict identification conditions), offering a solution that works even when the "transfer" assumption is false.

In summary, the paper presents a rigorous, robust, and empirically validated approach to enhancing the statistical power of randomized trials for personalized decision-making by safely leveraging external data sources.