Statistical Learning for Heterogeneous Treatment Effects: Pretraining, Prognosis, and Prediction
This work addresses a fundamental problem in causal inference for domains like personalized medicine and educational policy, offering incremental improvements through cross-task learning within an existing framework.
The paper tackles the challenge of accurately estimating heterogeneous treatment effects, particularly with many covariates, by proposing pretraining strategies that leverage synergies between risk prediction and causal effect estimation, resulting in lower estimation error, reduced false discovery rates, and higher power for detecting heterogeneity.
Robust estimation of heterogeneous treatment effects is a fundamental challenge for optimal decision-making in domains ranging from personalized medicine to educational policy. In recent years, predictive machine learning has emerged as a valuable toolbox for causal estimation, enabling more flexible effect estimation. However, accurately estimating conditional average treatment effects (CATE) remains a major challenge, particularly in the presence of many covariates. In this article, we propose pretraining strategies that leverage a phenomenon in real-world applications: factors that are prognostic of the outcome are frequently also predictive of treatment effect heterogeneity. In medicine, for example, components of the same biological signaling pathways frequently influence both baseline risk and treatment response. Specifically, we demonstrate our approach within the R-learner framework, which estimates the CATE by solving individual prediction problems based on a residualized loss. We use this structure to incorporate side information and develop models that can exploit synergies between risk prediction and causal effect estimation. In settings where these synergies are present, this cross-task learning enables more accurate signal detection, yields lower estimation error, reduced false discovery rates, and higher power for detecting heterogeneity.