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
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 ). 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=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 c and treatment probability p .
Regressor
Hyperparameter Grid
Lasso
Alpha: {0.001, 0.01, 0.1, 1, 10}
Random Forest
Number of trees: {50, 100} Maximum depth: {3, 5, None}
XGBoost
Number of trees: {50, 100} Learning rate: {0.01, 0.1} Maximum depth: {3, 5}
Appendix
Table 9: Set of hyperparameters for outcome imputation methods
Parameter
Default Value
Number of trees grown
2000
Fraction of data per tree
0.5
Variables tried per split
min(⌈p+20⌉,p)
Minimum samples in a leaf
5
Maximum imbalance of splits
0.05
Penalty for imbalance at split
0
Appendix
Table 10: Set of hyperparameters for Causal Survival Forests
Parameter
Value
Model width zdim,hdim1,hdim2
{8,16,32} (synthetic, ACTG) or {16,32,64} (MIMIC, Twin)
Number of shared layers
2
Number of head layers
2
Activation function
ReLU
Dropout rate
0.3
IPM type
Wasserstein
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 c and treatment probability p .
Variable
Proportion
Sex
Male
0.512
Female
0.488
Race
White
0.699
Black
0.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 ± 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 ± 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 ± 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=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.
Adaptive experimentation enables efficient estimation of causal effects, but existing methods are not designed for survival data with censoring, where event times are only partially observed (e.g., overall survival in cancer trials but with dropout). In this paper, we develop a novel framework for adaptive experimentation to estimate causal effects under right censoring. For this, we derive the semiparametric efficiency bound for the average survival effect curve as a function of the treatment allocation policy and thereby obtain a closed-form efficiency-optimal allocation policy. The policy generalizes classical Neyman allocation to survival settings by prioritizing patient strata where both event and censoring dynamics induce high uncertainty. Building on this, we propose the Adaptive Survival Estimator (ASE), an adaptive framework that learns the allocation policy and estimates the average survival effect curve sequentially. Our framework has three main benefits: (i) it accommodates arbitrary machine learning models for nuisance estimation; (ii) it is guided by a closed-form efficiency-optimal allocation policy; and (iii) it admits strong theoretical guarantees, including asymptotic normality via a martingale central limit theorem. We demonstrate our framework across various numerical experiments to show consistent efficiency gains over uniform randomization and censoring-agnostic baselines.
Yuxin Wang, Dennis Frauen, Jonas Schweisthal +3
LMU Munich Munich Center of Machine Learning (MCML)
Estimating heterogeneous survival treatment effects (HSTEs) from time-to-event data under time-varying treatments is important in fields such as medicine (e.g., to decide which sequence of treatments is better in terms of survival probability). However, this task is challenging because treatment decisions, patient characteristics, and censoring probability often evolve together, requiring sequential adjustment for both treatments and censoring. Existing survival learners mostly consider a single baseline treatment assignment, while longitudinal treatment-effect learners typically assume uncensored outcomes. In this paper, we introduce a novel doubly robust learner for heterogeneous survival treatment effects under time-varying treatments and history-dependent censoring. Our framework separates the intervention window from the survival horizon, with treatment prescribed for a finite period, after which care follows its natural course. We further establish the favorable theoretical properties of Neyman-orthogonality, rate double robustness and quasi-oracle efficiency. The framework accommodates flexible machine learning models for nuisance estimation and treatment-effect regression. We evaluate our learners through several numerical experiments, demonstrating their effectiveness.
Tomàs Garriga, Munib Mesinovic, Luke De Bretton-Gordon +3
Novartis · Barcelona Supercomputing Center · Universitat Politècnica de Catalunya +1
We study the problem of selecting the best heterogeneous treatment effect (HTE) estimator from a collection of candidates in settings where the treatment effect is fundamentally unobserved. We cast estimator selection as a multiple testing problem and introduce a ground-truth-free procedure based on a cross-fitted, exponentially weighted test statistic. A key component of our method is a two-way sample splitting scheme that decouples nuisance estimation from weight learning and ensures the stability required for valid inference. Leveraging a stability-based central limit theorem, we establish asymptotic familywise error rate control under mild regularity conditions. Empirically, our procedure provides reliable error control while substantially reducing false selections compared with commonly used methods across ACIC 2016, IHDP, and Twins benchmarks, demonstrating that our method is feasible and powerful even without ground-truth treatment effects.
Jiayi Guo, Zijun Gao
Peking University · Marshall School of Business, University of Southern California