Primary Submission Category: Heterogeneous Treatment Effects
Statistical Learning for Heterogeneous Treatment Effects: Pretraining, Prognosis, and Prediction
Authors: Maximilian Schuessler, Erik Sverdrup, Robert Tibshirani,
Presenting Author: Maximilian Schuessler*
Robust estimation of heterogeneous treatment effects is a fundamental task for optimal decision-making in many applications from personalized medicine to educational policies. In recent years, predictive machine learning has emerged as a valuable toolbox for causal estimation, enabling more flexible and rigorous effect estimation. Despite these advances, robust conditional average treatment estimation (CATE) remains highly challenging, especially in settings with complex interactions, low signal or high dimensions. In this article, we propose a new pretraining strategy that leverages a phenomenon in real-world applications: factors that are prognostic of the outcome, are frequently also predictive of treatment effects. Drawing on the R-loss and the lasso, we develop a suite of refined model architectures of R-Learners that achieve lower error rates in settings with shared support between the mean outcome function and the treatment effect function. Intuitively, if factors associated with a good baseline prognosis are also predictive of high treatment effects, pretraining will improve the estimation of the CATE. This also offers a data-driven way to discover a potential overlap between prognostic and predictive factors. We extend this approach to nonlinear models, basis function expansion, and settings with right-censoring, which allows us to demonstrate the utility of this framework to a series of settings and medical applications.