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

arXiv: https://arxiv.org/abs/2505.00310 · HTML: https://mlsanity.com/paper/2505.00310/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 22 of its 71 references that are in the map.
- First posted on arXiv 1 May 2025.

## Builds on (most cited first, 22 of 22)

- [XGBoost: A Scalable Tree Boosting System](https://mlsanity.com/paper/1603.02754/) (2016)
- [Estimation and Inference of Heterogeneous Treatment Effects using Random Forests](https://mlsanity.com/paper/1510.04342/) (2015)
- [Meta-learners for Estimating Heterogeneous Treatment Effects using Machine Learning](https://mlsanity.com/paper/1706.03461/) (2017)
- [Quasi-Oracle Estimation of Heterogeneous Treatment Effects](https://mlsanity.com/paper/1712.04912/) (2017)
- [Generalized Random Forests](https://mlsanity.com/paper/1610.01271/) (2016)
- [Recursive Partitioning for Heterogeneous Causal Effects](https://mlsanity.com/paper/1504.01132/) (2015)
- [Orthogonal Statistical Learning](https://mlsanity.com/paper/1901.09036/) (2019)
- [Panning for Gold: Model-X Knockoffs for High-dimensional Controlled Variable Selection](https://mlsanity.com/paper/1610.02351/) (2016)
- [Transfer Learning for High-dimensional Linear Regression: Prediction, Estimation, and Minimax Optimality](https://mlsanity.com/paper/2006.10593/) (2020)
- [Some methods for heterogeneous treatment effect estimation in high-dimensions](https://mlsanity.com/paper/1707.00102/) (2017)
- [Iterative Random Forests to detect predictive and stable high-order interactions](https://mlsanity.com/paper/1706.08457/) (2017)
- [False Discoveries Occur Early on the Lasso Path](https://mlsanity.com/paper/1511.01957/) (2015)
- [Evaluating Treatment Prioritization Rules via Rank-Weighted Average Treatment Effects](https://mlsanity.com/paper/2111.07966/) (2021)
- [Combining T-learning and DR-learning: a framework for oracle-efficient estimation of causal contrasts](https://mlsanity.com/paper/2402.01972/) (2024)
- [Experimental Evaluation of Individualized Treatment Rules](https://mlsanity.com/paper/1905.05389/) (2019)
- [Robust angle-based transfer learning in high dimensions](https://mlsanity.com/paper/2210.12759/) (2022)
- [Machine Learning Who to Nudge: Causal vs Predictive Targeting in a Field Experiment on Student Financial Aid Renewal](https://mlsanity.com/paper/2310.08672/) (2023)
- [Pretraining and the Lasso](https://mlsanity.com/paper/2401.12911/) (2024)
- [Covariate powered cross-weighted multiple testing](https://mlsanity.com/paper/1701.05179/) (2017)
- [Qini Curves for Multi-Armed Treatment Rules](https://mlsanity.com/paper/2306.11979/) (2023)
- [Optimal Multiple Testing Under a Gaussian Prior on the Effect Sizes](https://mlsanity.com/paper/1504.02935/) (2015)
- [Univariate-Guided Sparse Regression](https://mlsanity.com/paper/2501.18360/) (2025)

## Abstract (arXiv)

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.