Source-linked AI summary
Some methods for heterogeneous treatment effect estimation in high-dimensions
Scott Powers, Junyang Qian, Kenneth Jung, Alejandro Schuler, Nigam H. Shah, Trevor Hastie, Robert Tibshirani
TL;DR
The paper addresses personalized treatment-effect estimation from high-dimensional observational data, where doctors often lack quantitative evidence and patient similarity is difficult to define. It proposes three methods based on conditional mean regression and shared bases, and reports promising simulation performance and a kidney-health relationship in SPRINT data. Confidence-interval construction remains future work.
Problem
High-dimensional observational medical data make it difficult to identify similar patients and infer heterogeneous treatment effects for treatment decisions.
Method
The paper proposes three methods that prefer shared-basis conditional mean regression and adapt the basis to treatment-effect heterogeneity.
Results
Causal boosting and bagged causal MARS perform promisingly in simulations, and both identify a kidney-health relationship with treatment effect in SPRINT data.
Takeaways & Limitations
The proposed methods provide a basis for estimating heterogeneous treatment effects from high-dimensional data, including observational settings.
Takeaways & Limitations
The paper does not develop confidence intervals, leaving their construction for future work.
Abstract
from arXiv · showhide
When devising a course of treatment for a patient, doctors often have little quantitative evidence on which to base their decisions, beyond their medical education and published clinical trials. Stanford Health Care alone has millions of electronic medical records (EMRs) that are only just recently being leveraged to inform better treatment recommendations. These data present a unique challenge because they are high-dimensional and observational. Our goal is to make personalized treatment recommendations based on the outcomes for past patients similar to a new patient. We propose and analyze three methods for estimating heterogeneous treatment effects using observational data. Our methods perform well in simulations using a wide variety of treatment effect functions, and we present results of applying the two most promising methods to data from The SPRINT Data Analysis Challenge, from a large randomized trial of a treatment for high blood pressure.
1 Introduction
The paper targets personalized treatment recommendations from high-dimensional observational medical data, where treatment effects may differ across patient subgroups. It formulates this problem using covariates, treatment assignments, outcomes, and potential outcomes.
- Motivation: Doctors often lack quantitative evidence for choosing among treatment options when no randomized controlled trial compares them.The paper motivates using electronic medical records to support these decisions.
- Problem setting: High-dimensional patient data make it difficult to define which past patients are sufficiently similar to a new patient.Thousands of features may describe each patient and contribute to similarity.
- Problem setting: The goal is to estimate heterogeneous treatment effects and identify predictor-space regions where treatment effects are relatively large or small.Such regions support personalized medicine when average effects obscure subgroup-specific benefits.
- Problem setting: Each observation contains covariates X, a binary treatment indicator T, and a quantitative response Y under an i.i.d. sampling model.The framework uses the Neyman–Rubin potential outcomes model, in which only one potential outcome is observed for each patient.
- Scope: The paper studies both randomized trials and observational studies, where treatment assignment may depend on covariates.It develops methods for estimating heterogeneous effects in these settings.
2 Related work
Prior work includes subgroup comparisons, adaptive trees and forests, transformed-outcome methods, and propensity-score approaches. The paper addresses a gap because many existing methods assume randomized treatment or do not target covariate-dependent treatment-effect heterogeneity.
- Earlier subgroup methods: Earlier methods compared predefined patient subgroups or modeled treatment–covariate interactions using STEPP and MFPI.These approaches were developed for randomized experiments and flexible subgroup modeling.
- Tree-based methods: Causal trees adaptively split patients to estimate treatment effects, using within-leaf treatment-minus-control means rather than ordinary outcome means.Splits are selected to maximize heterogeneity in the treatment-effect function, with the paper using a T-statistic criterion.
- Machine-learning methods: Causal forests extend causal trees with random forests, while transformed-outcome methods apply supervised learning to outcomes transformed for treatment-effect estimation.The related work includes causal forests, support-vector methods, linear transformed-covariate models, and outcome-weighted learning.
- Observational-data gap: Most prior methods assume randomized treatment assignment, whereas propensity forests address nonrandomized treatment without targeting treatment-effect heterogeneity.Other approaches model effects through propensity scores or test constant effects, which can miss covariate-dependent heterogeneity or high-dimensional structure.
- Propensity-score methods: Propensity-score matching, stratification, and inverse-probability weighting reduce bias in observational estimates of population-average treatment effects.These methods rely on unconfoundedness, requiring treatment and potential outcomes to be conditionally independent given observed covariates.
3 Transformed outcome regression and conditional mean regression
The paper contrasts transformed outcome regression with conditional mean regression for heterogeneous treatment effects. It prefers shared-basis conditional mean methods because transformed outcomes can have high variance and separate basis selection can create spurious differences.
- Two regression approaches: The paper organizes heterogeneous treatment-effect estimators into transformed outcome regression and conditional mean regression.Causal forests, causal boosting, and causal MARS use conditional mean regression, while propensity transformed outcome combines both approaches.
- Transformed outcome regression: Transformed outcome regression creates Z with E[Z|X = x] = τ(x), then applies supervised learning to estimate the conditional mean of Z.Its variance can be large when propensity scores approach zero or one.
- Conditional mean regression: Conditional mean regression estimates µ1(x) and µ0(x) separately and obtains τ(x) from their difference.Ignoring propensity scores can introduce bias when treatment and control covariate distributions differ and the conditional means are imperfectly estimated.
- Variance comparison: Both ˆτCM and ˆτTO are unbiased for τ ≡µ1−µ0, but they have different variances.The comparison is developed first for estimating an average treatment effect in randomized data.
- Variance comparison: For small n, ˆτTO can have slightly smaller variance near zero main effect, but this advantage disappears as n increases and reverses when the main effect is large.The paper therefore prefers conditional mean regression because transformed outcome estimation can have much greater variance.
- Shared-basis conditional mean regression: Separate basis selection for treatment and control can produce conditional-mean differences caused by selection randomness rather than heterogeneous treatment effects.The proposed methods use a shared basis and adapt it to treatment-effect heterogeneity rather than response heterogeneity.
4 Pollinated transformed outcome (PTO) forests
The PTO forest estimates heterogeneous treatment effects by learning structure from transformed outcomes, then replacing noisy leaf estimates with treatment-specific means. An optional final forest adds regularization and interpretability.
- Pollination lowers treatment-effect estimate variance by replacing a transformed-outcome estimator with a conditional-mean estimator.
- The PTO forest builds a depth-controlled random forest on X to predict the transformed treatment-effect outcome.
- Pollination replaces each tree leaf’s transformed-outcome estimate with the difference between treated and control mean responses.Each observation is passed through each tree to obtain treatment-specific leaf means.
- An optional additional forest predicts the pollinated treatment effects, adding regularization and variable-importance interpretability.
- Figure 2 compares raw and pollinated forests in randomized and non-randomized simulations using treatment-effect mean square error and standard errors.The passage reports that pollination improves estimates in each case.
5 Causal boosting
Causal boosting adapts boosting to heterogeneous treatment-effect estimation by iteratively fitting causal trees separately across treatment arms. For observational data, propensity-score stratification adjusts within-leaf treatment-effect estimates, while pollinated validation supports tree-number selection.
- Causal boosting adapts least-squares boosting by fitting causal trees to residuals and updating a sequence of treatment-specific functions.The algorithm initializes residuals with observed responses, fits causal trees iteratively, and updates the model with a learning rate.
- For an observation x, the estimated treatment effect is the difference between the boosted functions evaluated at treatment and control.
- Cross-validation: Gradient boosting can over-fit as trees increase, making the number of trees an important tuning parameter.The usual cross-validation framework is complicated because individual treatment effects are unobserved.
- Cross-validation: Validation pollinates the fitted tree sequence with held-out treatment and outcome data to evaluate treatment-effect estimates across tree counts.This supplies a validation procedure for choosing among models with different numbers of boosting trees.
- Within-leaf propensity adjustment: Propensity-score methods do not immediately extend to individualized treatment effects under stratification, because predictions across strata would extrapolate beyond their training sets.Inverse probability weighting remains viable, but its volatility is exacerbated for varying treatment effects.
- Within-leaf propensity adjustment: The method handles non-randomized treatment by estimating propensity scores and grouping patients into strata of similar scores.Within causal-tree leaves, propensity-score stratification reduces bias in the estimated within-leaf average treatment effect.
6 Causal MARS
Causal MARS adapts multivariate adaptive regression splines to treatment-effect estimation by fitting parallel treatment and control models with shared basis functions. Bagging is used to reduce its variance, while the paper notes unresolved bias and confidence-interval challenges for tree-based comparisons.
- Tree-based treatment-effect estimates may have high bias because each leaf predicts using its within-leaf average treatment effect.The manuscript does not develop confidence intervals and identifies bias as a challenge for future work.
- Causal MARS fits parallel MARS models in treatment and control arms while adding the same basis functions to both.
- Its basis-selection criterion favors functions that explain treatment effects through different coefficients across the two arms.The criterion compares error reduction from arm-specific coefficients with error reduction from a shared coefficient.
- The parameter D controls the maximum dimension of the regression basis, and the examples use D = 11.
- Bagging averages treatment-effect estimates from B bootstrap-fitted causal MARS models to reduce variance.When bagging, backward deletion and model selection can be skipped to save computation time.
- For observational data, propensity strata use common basis functions but separate regression coefficients, combined through a propensity-adjusted criterion.
7 Simulation study
The simulation study evaluates heterogeneous treatment-effect estimators across randomized and observational settings spanning varied feature dimensions, functional forms, noise levels, and treatment-assignment mechanisms. Shared-basis methods perform strongly overall, with causal boosting and causal MARS standing out across the full set of scenarios.
- Simulation design: The simulations cover univariate and multivariate, additive and interactive, and linear and piecewise-constant functions.The generating functions are centered and scaled to have means near zero and roughly comparable variances.
- Randomized experiments: Most estimators outperform the null estimator except when treatment effects are constant or causal forest performs worst.The transformed-outcome and different-basis estimators are not competitive with the proposed methods, supporting shared-basis conditional mean regression.
- Simulation design: 16 scenarios vary sample size, dimensionality, mean-effect and treatment-effect functions, noise level, and treatment-assignment bias.Scenarios 1–8 use randomized assignment with π(x) = 1/2; scenarios 9–16 use biased assignment.
- Observational studies: Propensity-adjusted causal boosting performs strongly under biased treatment assignment, ranking among the two lowest-error methods in six of eight simulations.Across all 16 scenarios, causal boosting and causal MARS stand out overall; propensity adjustment can nevertheless hurt causal MARS in scenario 13.
- Bias and variance: In scenario 8, causal MARS has much lower bias than causal forest for a complex treatment-effect function with n = 1000 and p = 100.Its greater flexibility increases variance, but the bias reduction motivates causal MARS for future confidence-interval construction.
8 Application
Using causal boosting and bagged causal MARS on SPRINT data, the paper estimated personalized effects of intensive blood-pressure control and examined whether kidney health marked treatment-effect heterogeneity. Training results suggested less benefit for patients with CKD, but causal boosting did not reproduce that relationship in validation.
- Data and methods: 7344 uncensored patients were split into equally sized training and validation sets for personalized treatment-effect estimation.The training and validation sets each contained 3672 patients.
- Training results: 0.4% of the training sample had bagged causal MARS estimates outside the plotted region.These estimates contributed to the method’s heavier-tailed distribution.
- Training results: Causal boosting inferred that older patients with high triglycerides could gain more from intensive treatment than younger patients with high triglycerides.Among patients with low triglycerides and high glucose, lower creatinine was associated with greater estimated benefit.
- Kidney-health finding: Patients with CKD, defined here by eGFR below 60, were estimated to gain less from intensive blood-pressure control.The analysis emphasized eGFR and creatinine as kidney-health variables related to personalized treatment effects.
- Validation: In validation, bagged causal MARS again showed a negative relationship between eGFR and treatment effect, whereas causal boosting showed no clear eGFR relationship.The authors note that insufficient power could explain the lack of agreement by causal boosting.
9 Discussion
The paper concludes that causal boosting and bagged causal MARS are promising in simulations and recover a kidney-health relationship in SPRINT data also identified by other researchers. Confidence-interval construction remains future work.
- Causal boosting and bagged causal MARS appear particularly promising in simulations.
- Both methods found a relationship between kidney health and treatment effect in SPRINT data that other researchers also identified.
- Confidence-interval construction is an important next step left for future work.The authors developed causal MARS to facilitate confidence-interval construction but did not complete it here.