Statistical Learning for Heterogeneous Treatment Effects: Pretraining, Prognosis, and Prediction

TL;DR

This paper introduces pretraining strategies that leverage the relationship between prognostic factors and treatment effect heterogeneity to improve the estimation of conditional average treatment effects using the R-learner framework.

Contribution

It proposes a novel pretraining approach that exploits the correlation between prognosis and treatment effects to enhance causal effect estimation in machine learning models.

Findings

Pretraining improves the accuracy of CATE estimates.

The approach reduces false discovery rates in heterogeneity detection.

Models show higher power in identifying treatment effect heterogeneity.

Abstract

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.