Source-linked AI summary
Prediction-Powered Data Fusion for Treatment Effect Estimation
Yonghan Jung, Shu Yang
TL;DR
RCTs identify treatment effects without confounding but are often small, while larger OBS data may be confounded; existing fusion approaches for ATE and CATE have limitations in assumptions or borrowing. The paper develops fusion estimators that retain the trial-based target while borrowing from OBS, and reports improved performance over trial-only estimation, PPI, and naive pooling.
Problem
Existing ATE fusion methods may assume properties of the OBS or borrow too little from them, while CATE fusion methods may require unconfounded OBS, model confounding, or trade bias for variance.
Source: §1 · p. 2Method
AIPW-Fusion adds OBS outcome regressions and a PPI correction to the trial’s AIPW score, while DR-Fusion and R-Fusion apply the borrowing principle to CATE learning.
Source: §1 · p. 2, §3.1 · p. 6Results
The estimators outperform trial-only estimation, PPI, and naive pooling in experiments on synthetic data, IHDP, and ACIC 2016.
Source: §1 · p. 2Takeaways & Limitations
The framework provides ATE and CATE estimators that preserve the trial target while using OBS data to improve precision.
Source: §6 · p. 26Selecting borrowing weights and evaluating the estimate on the same samples can bias the estimate, so the algorithm uses separate tuning and evaluation samples.
Source: §3.3 · p. 9
Abstract
from arXiv · showhide
Randomized controlled trials (RCTs) identify treatment effects without confounding but are often small, whereas observational studies (OBS) are large but may be confounded. Many estimators combining a small RCT with a large OBS have been developed for the average treatment effect (ATE) and the conditional ATE (CATE). However, existing ATE estimators either make assumptions on the OBS or do not borrow enough power from them. The CATE has been studied less than the ATE. Existing CATE methods either assume the OBS are unconfounded, rely on a model of the confounding function, or accept bias in exchange for lower variance. We therefore propose a framework that, without special assumptions on the OBS, fuses the OBS and the RCT by preserving the unbiasedness of RCT-based estimation while borrowing power from the large OBS to boost precision. Applying this principle, we build an ATE estimator, AIPW-Fusion, with closed-form weights and confidence intervals, and two CATE learners, DR-Fusion and R-Fusion. Experiments corroborate our findings.
1. Introduction
The paper seeks to improve precision for trial-population treatment effects by borrowing from potentially confounded observational data without sacrificing the trial-based target. It proposes an ATE estimator and two CATE learners, and reports that they outperform trial-only estimation, PPI, and naive pooling in its experiments.
1. Introduction
- AIPW-Fusion combines OBS outcome regressions in the trial’s AIPW score with a PPI correction, remaining unbiased under any OBS confounding and providing a valid interval under an estimated ratio.The two coefficients are chosen jointly in closed form.Source: §1 · p. 2
- DR-Fusion and R-Fusion preserve the trial’s CATE target under any OBS confounding, borrow without modeling the confounding function, and provably do almost no worse than the trial alone.Source: §1 · p. 2
- In experiments on synthetic data, IHDP, and ACIC 2016, the estimators outperform trial-only estimation, PPI, and naive pooling.Source: §1 · p. 2
1.1. Related Work
- For ATE fusion, existing approaches variously borrow OBS outcomes, use only OBS covariates, or preserve unbiasedness through predictions or PPI; AIPW-Fusion combines two previously separate borrowing channels.Source: §1.1 · p. 3
- Prior CATE fusion methods either require unconfounded OBS, model the OBS bias, or trade bias for variance; DR-Fusion and R-Fusion need no model of the trial–OBS CATE discrepancy.Source: §1.1 · pp. 3–4
2. Background and Setting
The paper defines its targets in the randomized-trial population, where treatment assignment is known, and contrasts that setting with observational data confounded by an unmeasured variable. Its estimation setup separates nuisance fitting, tuning, and evaluation samples.
2. Background and Setting
- The targets are the RCT-population conditional effect τ0(x) and average effect θ0; treatment depends on X through a known propensity in the RCT, while unmeasured U affects treatment and outcome in the OBS.Source: §2 · p. 5
- Trial identification assumes consistency, conditional randomization given X, and a known propensity bounded away from zero and one.Source: §2 · p. 5
- The sources are split into independent nuisance, tuning, and evaluation samples; the reported exact results use a single split, though roles can be rotated and estimates averaged.Source: §2 · p. 5
- The AIPW score is conditionally unbiased for τ0(X) under the trial assumptions, supporting the RCT-only estimator of θ0 and its use with alternative regressors.Source: §2 · p. 6
3. Fusion for the ATE
The ATE fusion method combines RCT-based estimation with observational-study information while retaining unbiasedness under a correct density ratio. It calibrates an estimated ratio to control drift and accounts for calibration noise in confidence intervals.
3.1. Estimator, Variance, and Oracles
- AIPWF uses a blended regressor and a transported control variate, with λ and ω controlling the two borrowing channels.The RCT-only AIPW estimator is recovered at (λ, ω) = (0, 0).Source: §3.1 · p. 6
- Oracle AIPWF variance subtracts C²/A and D²/B from RCT-only variance, making it no greater than PPI++ or RCT-only variance.Source: §3.1 · p. 7
- The ω channel can help when effect heterogeneity contributes substantially to score variance, while λ helps when the observational regressor differs from the trial regression.Source: §3.1 · p. 8
3.2. Estimating the Density Ratio
- Calibrating the density ratio to balance the observational mean of the control variate with its trial mean sets the moment driving AIPWF bias to zero.For AIPWF, balancing the control variate itself suffices.Source: §3.2 · pp. 8–9
3.3. Algorithm and Confidence Intervals
- The estimator’s conditional mean shifts by ω times the ratio-induced moment drift, while calibration noise is included in the confidence-interval variance.With calibrated ratios, the conditional bias is zero to first order after averaging over nuisance samples.Source: §3.3 · p. 10
4. Fusion for the CATE
DRF and RF extend the two OBS-borrowing channels of AIPW-Fusion to CATE learning. With the true density ratio, borrowing changes variance without shifting the target; with an estimated ratio, calibration and validation bound the resulting risks.
4. Fusion for the CATE
- DRF uses the DR-learner’s AIPW pseudo-outcome, while RF uses the R-learner loss; both borrow from OBS through the two AIPW-Fusion channels.The blended regressor is a convex combination of RCT and OBS regressions under the stated assumptions.Source: §4 · p. 11
4.1. CATE Risks
- With the true density ratio, each learner’s expected empirical risk equals CATE risk up to a candidate-independent constant, so both have the same minimizer.The true ratio also makes OBS borrowing change only the fitted learner’s variance, not its target.Source: §4.1 · p. 12
- When the density ratio is misspecified, the proxy-risk term can shift the minimizer for ω > 0; for linear sieves, calibration reduces that shift to nuisance-sampling error.The learning equation is orthogonal to the RCT regression for any ratio, but to the OBS regression only at the true ratio; ratio orthogonality requires bg = t.Source: §4.1 · pp. 12–13
4.2. Fitting and Selecting Coefficients
- Coefficient selection fits candidates over a finite grid and chooses the smallest held-out proxy-risk score, which ranks CATE risk when the ratio is correct.The grid includes (0, 0), the RCT-only learner; the selected learner’s risk is at most 2εeval(α) above the best grid candidate plus sampling and ratio-error terms.Source: §4.2 · pp. 13–14
4.3. Linear Sieves: Risk and Calibration
- For linear sieves, the ω-dependent bias term is the product of ratio error and the gap between OBS prediction and its best sieve fit, and vanishes if either factor is zero.At the true ratio, ratio-imbalance bias vanishes, leaving ridge bias; the risk theorem also accounts for approximation and estimation error.Source: §4.3 · pp. 15–16
- Calibration requires one feature for ATE but 1 + p + p(p + 1)/2 features for sieve CATE; balance gaps need only be within their standard errors.With calibrated ratios, the ratio-bias component is order n^-1_nuis and transport error is order n^-1/2_nuis.Source: §4.3 · p. 16
5. Experiments
Experiments compare AIPW-Fusion, DR-Fusion, and R-Fusion with trial-only estimators and alternative fusion methods across synthetic and semi-synthetic designs. Fusion has the lowest ATE MSE and CATE risk across the reported 45 cells, with additional tests examining bias, observational sample size, and estimator components.
5.1. Designs
- The synthetic design and IHDP and ACIC 2016 benchmarks include OBS confounding, with common or shifted covariate laws; benchmark samples reuse rows from finite covariate tables.The synthetic design sets confounding bias to 1.5 sdR{τ0(X)} for its main setting; the bias sweep varies its magnitude.Source: §5.1 · pp. 16–18
5.2. Estimators and Metrics
- Each trial is divided into nuisance, tuning, and evaluation samples; comparisons include trial-only estimators, PPI++, naive pooling, shrinkage, pretesting, and alternative CATE learners.Errors are measured relative to trial-only AIPW or the trial-only DR-learner, using 1,000 replications per cell.Source: §5.2 · pp. 18–19
5.3. Main Results
- 0.40 to 1.00 of trial-only AIPW MSE for AIPW-Fusion and 0.53 to 0.67 of trial-only DR-learner risk for DRF/RF across 45 cells; each achieved the lowest error or risk.The ATE MSE ratios are from Table 5 and the CATE risk ratios from Table 6.Source: Table 5 · p. 19, Table 6 · p. 20, §5.3 · p. 21
5.4. Bias and Observational Sample Size
- Across bias sweeps, ATE fusion ratios remain flat as confounding grows; at n = 3,000, the CATE IR ratio stays between 1.08 and 1.10, while CATE fusion risk falls as OBS N grows.At n = 3,000, the naive pool falls behind fusion at bias between 0 and 0.25 standard deviations for ATE, and between 0.25 and 0.5 for CATE.Source: Figure 5 · p. 21, Figure 6 · p. 21, §5.4 · p. 22
5.5. Ablations and Further Designs
- Under covariate shift, Eq. (11) lowers fusion MSE by 1% to 4% versus the rule that ignores the calibration term in all 30 cells.Without shift, fusion is at or below λ-only in 14 of 15 cells; under shift, λ-only is slightly better in 28 of 30 cells.Source: §5.5 · p. 22, Table 7 · p. 22
- Against practitioner trial-only learners using two thirds of the trial for nuisance fitting, DRF/RF risk is 0.54 to 0.71 of theirs in all 45 cells.Fusion remains below practitioner trial-only AIPW in 42 of 45 cells.Source: §5.5 · pp. 22, 24, Table 8 · p. 23
- With the trial sieve augmented by bg, most DRF/RF gains over trial-only learners come from fusion loss at small n, but from bg as a basis function at large n.Against trial-only DR risk, the sieve with bg has 0.75 to 0.93 risk at n = 300 and 0.55 to 0.67 at n = 3,000.Source: Table 9 · p. 23, §5.5 · p. 24
- With a misspecified ratio, calibrated fusion retains 0.48 to 0.99 of trial-only MSE and 0.94 to 0.96 coverage; the uncalibrated ratio is biased and exceeds trial-only MSE at n ≥ 900.In the large-heterogeneity design, fusion improves on λ-only without shift, but is worse than λ-only under strong shift.Source: §5.5 · pp. 24–25, Table 11 · p. 25
6. Conclusion
The paper’s fusion estimators preserve the RCT-based target while using observational data to improve precision. The conclusion reports lower error across every design and trial size.
6. Conclusion
- Fusion has the lowest error in every design and trial size.The conclusion also reports closed-form optimal weights and valid Wald intervals for AIPW-Fusion.Source: §6 · p. 26
- AIPW-Fusion has closed-form optimal weights and valid Wald intervals with a calibrated ratio.These properties are supported by Theorems 1–2 and Corollary 1.1.Source: §6 · p. 26
- DR-Fusion and R-Fusion target the trial’s CATE risk, with validation selecting weights relative to the RCT-only learner.For linear sieves, calibration on finitely many moments controls the ratio error in risk.Source: §6 · p. 26
A. Proofs
The proofs establish unbiasedness and inference for ATE fusion, then derive risk and selection guarantees for CATE fusion. They also address calibration and linear-sieve risk control.
A. Proofs
- RCT validity and exact transport establish unbiasedness for the ATE estimator using the true density ratio, with its variance derived from independent evaluation blocks.The variance expression substitutes the score decomposition and uses zero covariance between the RCT score difference and the balancing function.Source: A. Proofs · p. 29
- Conditioned on fitted objects, unbiased sample variances estimate the estimator variance, while a two-sample Lyapunov CLT supports Wald interval coverage under Assumption 4.Variance unbiasedness alone does not establish a central limit theorem.Source: A. Proofs · p. 31
- For DRF, the cross term vanishes and RDRF = QDRF(λ) + R0(t); for the R loss, randomization gives ER[U(A−e) | X] = 0.Source: A. Proofs · p. 31
- The validation-selected CATE pair is within a bound of the best grid risk, and that benchmark is no worse than the trial-only pair.The bound includes uniform evaluation deviation and transport error.Source: A. Proofs · p. 32
- For linear sieves, calibrated feature imbalances control the nuisance-sample contribution to ratio error under the theorem’s conditions.The fixed-p expansion conditions are assumed rather than verified in the paper.Source: A. Proofs · pp. 32–33
B. Additional Simulation Results
The appendix reports additional ATE and absolute-error simulation results, including metrics and interval summaries for the listed designs.
B. Additional Simulation Results
- The appendix also provides the ATE bias, standard deviation, and Wald intervals, plus absolute errors behind Figure 4.The absolute-error results are presented in Figure 7.Source: B. Additional Simulation Results · p. 33
- Table 13 reports ATE bias, standard deviation, RMSE, mean estimated SE, coverage, and mean 95% Wald interval width over 1,000 synthetic-design replications.Shrinkage, Pretest, and the naive pool have no interval.Source: Table 13 · p. 34
- Table 14 gives the same ATE summaries for IHDP and ACIC 2016.Its scope is specified as matching Table 13.Source: Table 13 · p. 34, Table 14 · p. 35