[Paper Review] BITES: Balanced Individual Treatment Effect for Survival data

Stefan Schrod, Andreas Schäfer+8 more|arXiv (Cornell University)|Jan 1, 2022

Statistical Methods in Clinical Trials4 citations

TL;DR

BITES is a novel deep learning framework that estimates individual treatment effects in survival data by combining a treatment-specific Cox proportional hazards model with integral probability metric (IPM)-based balancing of treated and control groups in a latent representation space. It outperforms state-of-the-art methods in simulation and real-world breast cancer data, improving treatment recommendation accuracy and survival outcomes while enabling interpretable predictions via SHAP values.

ABSTRACT

Estimating the effects of interventions on patient outcome is one of the key aspects of personalized medicine. Their inference is often challenged by the fact that the training data comprises only the outcome for the administered treatment, and not for alternative treatments (the so-called counterfactual outcomes). Several methods were suggested for this scenario based on observational data, i.e.~data where the intervention was not applied randomly, for both continuous and binary outcome variables. However, patient outcome is often recorded in terms of time-to-event data, comprising right-censored event times if an event does not occur within the observation period. Albeit their enormous importance, time-to-event data is rarely used for treatment optimization. We suggest an approach named BITES (Balanced Individual Treatment Effect for Survival data), which combines a treatment-specific semi-parametric Cox loss with a treatment-balanced deep neural network; i.e.~we regularize differences between treated and non-treated patients using Integral Probability Metrics (IPM). We show in simulation studies that this approach outperforms the state of the art. Further, we demonstrate in an application to a cohort of breast cancer patients that hormone treatment can be optimized based on six routine parameters. We successfully validated this finding in an independent cohort. BITES is provided as an easy-to-use python implementation.

Motivation & Objective

• To address the challenge of estimating individual treatment effects (ITE) in time-to-event data where counterfactual outcomes are unobserved.
• To overcome limitations of existing methods that do not adequately handle censored survival data or fail to balance treatment groups in high-dimensional settings.
• To develop a scalable, interpretable, and robust framework for personalized treatment optimization in observational survival studies.
• To validate the method on real-world breast cancer data and demonstrate improved survival outcomes through optimized hormone therapy recommendations.

Proposed method

• BITES integrates a treatment-specific semi-parametric Cox model with a deep neural network to estimate individualized treatment effects on survival outcomes.
• It applies integral probability metrics (IPM) to regularize the latent representations of treated and non-treated patients, promoting distributional balance.
• The model uses a dual-head architecture: one head predicts survival outcomes per treatment group, and the other learns balanced representations via adversarial training.
• It employs a hybrid loss combining Cox partial likelihood for survival modeling and IPM-based regularization to minimize distributional discrepancies between treatment groups.
• The framework supports end-to-end training with backpropagation, enabling joint optimization of predictive performance and treatment balance.
• Interpretability is enhanced using SHAP (SHapley Additive exPlanations) to assess feature contributions to individual risk predictions under different treatment options.

Experimental results

Research questions

• RQ1 Can a deep learning model effectively estimate individual treatment effects in time-to-event data while accounting for censored outcomes and treatment selection bias?
• RQ2 How does IPM-based balancing of treated and non-treated latent representations improve the accuracy and fairness of ITE estimation in survival analysis?
• RQ3 To what extent can BITES outperform existing state-of-the-art methods in terms of Harrell’s C-index, PEHE, and treatment recommendation correctness on both simulated and real-world survival data?
• RQ4 Does the model’s treatment recommendation strategy lead to improved survival outcomes in independent validation cohorts, particularly in breast cancer hormone therapy?
• RQ5 How transparent and interpretable are the model’s predictions, and can SHAP values effectively reveal the impact of clinical features on individual treatment decisions?

Key findings

• In simulation studies, BITES achieved superior performance in terms of Harrell’s C-index, treatment assignment accuracy, and PEHE compared to state-of-the-art methods.
• On the GBSG-2 breast cancer dataset, BITES recommended hormone therapy for 83.4% of patients, significantly outperforming Cox regression and DeepSurv, which recommended treatment for 100% of patients.
• The group following BITES recommendations showed better survival than the treated group in the test set, though the difference was not statistically significant (p = 0.000016 for log-rank test on controlled test set).
• BITES demonstrated strong interpretability: SHAP analysis revealed that hormone therapy mitigated the negative survival impact of post-menopausal status and amplified the risk associated with high tumor grade (grade 3).
• The model’s C-index of 0.666 on the independent GBSG-2 test set was among the highest, indicating strong predictive performance.
• Validation in an independent cohort confirmed that BITES-based treatment recommendations were associated with improved survival outcomes, suggesting clinical utility.