# BITES: Balanced Individual Treatment Effect for Survival data

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

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 19 of its 72 references that are in the map.
- First posted on arXiv 5 Jan 2022.

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

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [A Unified Approach to Interpreting Model Predictions](https://mlsanity.com/paper/1705.07874/) (2017)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Estimation and Inference of Heterogeneous Treatment Effects using Random Forests](https://mlsanity.com/paper/1510.04342/) (2015)
- [Tune: A Research Platform for Distributed Model Selection and Training](https://mlsanity.com/paper/1807.05118/) (2018)
- [DeepSurv: Personalized Treatment Recommender System Using A Cox Proportional Hazards Deep Neural Network](https://mlsanity.com/paper/1606.00931/) (2016)
- [Learning Representations for Counterfactual Inference](https://mlsanity.com/paper/1605.03661/) (2016)
- [Learning Generative Models with Sinkhorn Divergences](https://mlsanity.com/paper/1706.00292/) (2017)
- [Generalized Random Forests](https://mlsanity.com/paper/1610.01271/) (2016)
- [On Wasserstein Two Sample Testing and Related Families of Nonparametric Tests](https://mlsanity.com/paper/1509.02237/) (2015)
- [Bayesian Inference of Individualized Treatment Effects using Multi-task Gaussian Processes](https://mlsanity.com/paper/1704.02801/) (2017)
- [A Scalable Discrete-Time Survival Model for Neural Networks](https://mlsanity.com/paper/1805.00917/) (2018)
- [Generalization Bounds and Representation Learning for Estimation of Potential Outcomes and Causal Effects](https://mlsanity.com/paper/2001.07426/) (2020)
- [Continuous and Discrete-Time Survival Prediction with Neural Networks](https://mlsanity.com/paper/1910.06724/) (2019)
- [Estimating Treatment Effects with Causal Forests: An Application](https://mlsanity.com/paper/1902.07409/) (2019)
- [SurvITE: Learning Heterogeneous Treatment Effects from Time-to-Event Data](https://mlsanity.com/paper/2110.14001/) (2021)
- [Local Similarity-Aware Deep Feature Embedding](https://mlsanity.com/paper/1610.08904/) (2016)

## Built on by (most cited first, 7 of 7)

- [Orthogonal Survival Learners for Estimating Heterogeneous Treatment Effects from Time-to-Event Data](https://mlsanity.com/paper/2505.13072/) (2025)
- [BENK: The Beran Estimator with Neural Kernels for Estimating the Heterogeneous Treatment Effect](https://mlsanity.com/paper/2211.10793/) (2022)
- [Assessing the robustness of heterogeneous treatment effects in survival analysis under informative censoring](https://mlsanity.com/paper/2510.13397/) (2025)
- [Adversarial Distribution Balancing for Counterfactual Reasoning](https://mlsanity.com/paper/2311.16616/) (2023)
- [ConfoundingSHAP: Quantifying confounding strength in causal inference](https://mlsanity.com/paper/2605.10533/) (2026)
- [Adaptive Experimentation for Censored Survival Outcomes](https://mlsanity.com/paper/2605.18459/) (2026)
- [Surv-IPTB: An Attention-Based Model for Estimating Individual Probability of Treatment Benefit with Survival Data](https://mlsanity.com/paper/2608.06288/) (2026)

## Abstract (arXiv)

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.