cs.LGMar 5, 2026

SurvHTE-Bench: A Benchmark for Heterogeneous Treatment Effect Estimation in Survival Analysis

Authors: Shahriar Noroozizadeh, Xiaobin Shen, Jeremy C. Weiss, George H. Chen

Organizations: Machine Learning Department & Heinz College, Carnegie Mellon University · Heinz College, Carnegie Mellon University · National Library of Medicine, National Institutes of Health

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Estimating heterogeneous treatment effects (HTEs) from right-censored survival data is critical in high-stakes applications such as precision medicine and individualized policy-making. Yet, the survival analysis setting poses unique challenges for HTE estimation due to censoring, unobserved counterfactuals, and complex identification assumptions. Despite recent advances, from Causal Survival Forests to survival meta-learners and outcome imputation approaches, evaluation practices remain fragmented and inconsistent. We introduce SurvHTE-Bench, the first comprehensive benchmark for HTE estimation with censored outcomes. The benchmark spans (i) a modular suite of synthetic datasets with known ground truth, systematically varying causal assumptions and survival dynamics, (ii) semi-synthetic datasets that pair real-world covariates with simulated treatments and outcomes, and (iii) real-world datasets from a twin study (with known ground truth) and from an HIV clinical trial. Across synthetic, semi-synthetic, and real-world settings, we provide the first rigorous comparison of survival HTE methods under diverse conditions and realistic assumption violations. SurvHTE-Bench establishes a foundation for fair, reproducible, and extensible evaluation of causal survival methods. The data and code of our benchmark are available at: https://github.com/Shahriarnz14/SurvHTE-Bench .

Figures & tables

Figure 1: (top) Borda count rankings of the top 10 estimator variants (out of 53 total), based on CATE RMSE across 40 datasets and averaged over 10 repeats (lower is better). (bottom) Family-level rankings, where for each dataset the best method variant within each method family is chosen using validation performance and then ranked on the held-out test set. Black bands connect methods without statistically significant differences (Wilcoxon signed-rank test, FDR-corrected at α=0.05\alpha=0.05 ). Shaded regions indicate the standard error of the rank across datasets.
Figure 2: CATE RMSE in Scenario C across 10 experimental repeats.
Figure 3: CATE RMSE for twin birth data with h=30h=30 days across 10 experimental runs.
Figure 4: CATE estimation comparison between baseline and high-censoring conditions under ZDV vs. ZDV+ddI treatments. Each point represents an individual patient, with the dashed diagonal line indicating perfect consistency between baseline CATE estimation and that with the additional censoring injected.
Appendix figures & tables27 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: (Synthetic datasets) Kaplan-Meier curves across causal configurations (rows) and survival scenarios (columns). Solid lines show event-time survival under control (blue) and treatment (orange); dotted lines show censoring-time survival for each arm. Each panel reports the empirical censoring rate cc and treatment probability pp .
RegressorHyperparameter Grid
LassoAlpha: {0.001, 0.01, 0.1, 1, 10}
Random ForestNumber of trees: {50, 100} Maximum depth: {3, 5, None}
XGBoostNumber of trees: {50, 100} Learning rate: {0.01, 0.1} Maximum depth: {3, 5}
Appendix
Table 9: Set of hyperparameters for outcome imputation methods
ParameterDefault Value
Number of trees grown2000
Fraction of data per tree0.5
Variables tried per splitmin⁡(⌈p+20⌉,p)\min(\lceil\sqrt{p}+20\rceil,p)
Minimum samples in a leaf5
Maximum imbalance of splits0.05
Penalty for imbalance at split0
Appendix
Table 10: Set of hyperparameters for Causal Survival Forests
ParameterValue
Model width zdim,hdim1,hdim2z_{\text{dim}},h_{\text{dim1}},h_{\text{dim2}}{8,16,32}\{8,16,32\} (synthetic, ACTG) or {16,32,64}\{16,32,64\} (MIMIC, Twin)
Number of shared layers2
Number of head layers2
Activation functionReLU
Dropout rate0.3
IPM typeWasserstein
Appendix
Table 11: Set of hyperparameters SurvITE.
Figure 6: Average ranking of each model for each Survival Scenario. Shaded regions indicate the standard error of the rank across datasets.
Figure 7: Average ranking of each model for each causal configuration. Shaded regions indicate the standard error of the rank across datasets.
Figure 8: CATE RMSE across different experiments in Scenario A .
Figure 9: CATE RMSE across different experiments in Scenario B .
Figure 10: CATE RMSE across different experiments in Scenario C .
Figure 11: CATE RMSE across different experiments in Scenario D .
Figure 12: CATE RMSE across different experiments in Scenario E .
Figure 13: ATE Bias across different experiments in Scenario A .
Figure 14: ATE Bias across different experiments in Scenario B .
Figure 15: ATE Bias across different experiments in Scenario C .
Figure 16: ATE Bias across different experiments in Scenario D .
Figure 17: ATE Bias across different experiments in Scenario E .
Figure 18: Convergence properties: CATE RMSE in Scenario C as number of training data increases.
Figure 19: (Semi-synthetic datasets) Kaplan-Meier curves. Solid lines show event-time survival under control (blue) and treatment (orange); dotted lines show censoring-time survival for each arm. Each panel reports the empirical censoring rate cc and treatment probability pp .
VariableProportion
Sex 
Male0.512
Female0.488
Race 
White0.699
Black0.141
Appendix
Table 30: Demographic and categorical distributions in MIMIC semi-synthetic datasets. Reported values are proportions.
Figure 20: Correlation heatmap of the 36 semi-synthetic MIMIC covariates. Variables include demographic features, admission descriptors, insurance and marital status indicators, and laboratory measurements. Most correlations are weak to moderate, with stronger dependencies visible among related laboratory values (e.g., hematocrit, hemoglobin, and red blood cell count).
Table 32: CATE RMSE (mean ±\pm std over 10 repeats) on MIMIC semi-synthetic datasets, using survival probability at 25th quantile event time of each dataset as the estimand. Best method per dataset is bolded .
Table 33: CATE RMSE (mean ±\pm std over 10 repeats) on MIMIC semi-synthetic datasets, using survival probability at 50th quantile event time of each dataset as the estimand. Best method per dataset is bolded .
Table 34: CATE RMSE (mean ±\pm std over 10 repeats) on MIMIC semi-synthetic datasets, using survival probability at 75th quantile event time of each dataset as the estimand. Best method per dataset is bolded .
Figure 21: CATE RMSE for twin birth data using different estimator families with h=180h=180 days across 10 experimental runs.
Figure 22: CATE Estimation comparison between baseline and high-censoring conditions under ZDV vs. ZDV+Zal treatments (HIV2). Each point represents an individual patient in test sets, with the dashed diagonal line indicating perfect consistency between baseline CATE estimation and that with the additional censoring injected.
Figure 23: CATE Estimation comparison between baseline and high-censoring conditions under ZDV vs. ddI treatments (HIV3). Each point represents an individual patient in test sets, with the dashed diagonal line indicating perfect consistency between baseline CATE estimation and that with the additional censoring injected.
Figure 24: CATE RMSE (left) and ATE bias (right) under informative censoring induced by unobserved confounding.
More in this topic

Individualized Treatment Effect Estimation

Individualized treatment effect estimation predicts how a treatment’s effects vary across people to support personalized decisions.

See all 44 papers

Explore similar work

CardsList

  1. Adaptive Experimentation for Censored Survival Outcomes

    May 18, 2026Yuxin Wang, Dennis Frauen, Jonas Schweisthal +3Causal Effect EstimationSurvival Analysis

    LMU Munich Munich Center of Machine Learning (MCML)

  2. Doubly Robust Learning of Heterogeneous Survival Treatment Effects under Time-Varying Treatments

    Oct 19, 2026Tomàs Garriga, Munib Mesinovic, Luke De Bretton-Gordon +3Causal Effect EstimationIndividualized Treatment Effect Estimation

    Novartis · Barcelona Supercomputing Center · Universitat Politècnica de Catalunya +1

  3. Reliable Selection of Heterogeneous Treatment Effect Estimators

    Nov 23, 2025Jiayi Guo, Zijun GaoCausal Effect EstimationConditional Average Treatment Effect Estimation

    Peking University · Marshall School of Business, University of Southern California