# Surv-IPTB: An Attention-Based Model for Estimating Individual Probability of Treatment Benefit with Survival Data

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 24 of its 51 references that are in the map.
- First posted on arXiv 6 Aug 2026.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Effective Approaches to Attention-based Neural Machine Translation](https://mlsanity.com/paper/1508.04025/) (2015)
- [A Survey on Causal Inference](https://mlsanity.com/paper/2002.02770/) (2020)
- [Nonparametric Estimation of Heterogeneous Treatment Effects: From Theory to Learning Algorithms](https://mlsanity.com/paper/2101.10943/) (2021)
- [Causal Transformer for Estimating Counterfactual Outcomes](https://mlsanity.com/paper/2204.07258/) (2022)
- [Estimating the Effects of Continuous-valued Interventions using Generative Adversarial Networks](https://mlsanity.com/paper/2002.12326/) (2020)
- [VCNet and Functional Targeted Regularization For Learning Causal Effects of Continuous Treatments](https://mlsanity.com/paper/2103.07861/) (2021)
- [Residual Weighted Learning for Estimating Individualized Treatment Rules](https://mlsanity.com/paper/1508.03179/) (2015)
- [A unified survey of treatment effect heterogeneity modeling and uplift modeling](https://mlsanity.com/paper/2007.12769/) (2020)
- [Estimating heterogeneous treatment effects with right-censored data via causal survival forests](https://mlsanity.com/paper/2001.09887/) (2020)
- [Exploring Transformer Backbones for Heterogeneous Treatment Effect Estimation](https://mlsanity.com/paper/2202.01336/) (2022)
- [Estimating Individual Treatment Effects using Non-Parametric Regression Models: a Review](https://mlsanity.com/paper/2009.06472/) (2020)
- [Individualized Treatment Effects with Censored Data via Fully Nonparametric Bayesian Accelerated Failure Time Models](https://mlsanity.com/paper/1706.06611/) (2017)
- [Quantifying Aleatoric Uncertainty of the Treatment Effect: A Novel Orthogonal Learner](https://mlsanity.com/paper/2411.03387/) (2024)
- [Orthogonal Survival Learners for Estimating Heterogeneous Treatment Effects from Time-to-Event Data](https://mlsanity.com/paper/2505.13072/) (2025)
- [Targeted Estimation of Heterogeneous Treatment Effect in Observational Survival Analysis](https://mlsanity.com/paper/1910.08877/) (2019)
- [CETransformer: Casual Effect Estimation via Transformer Based Representation Learning](https://mlsanity.com/paper/2107.08714/) (2021)
- [BITES: Balanced Individual Treatment Effect for Survival data](https://mlsanity.com/paper/2201.03448/) (2022)
- [On the estimation of average treatment effects with right-censored time to event outcome and competing risks](https://mlsanity.com/paper/1907.12912/) (2019)
- [Enabling Counterfactual Survival Analysis with Balanced Representations](https://mlsanity.com/paper/2006.07756/) (2020)
- [Estimating Conditional Average Treatment Effects via Sufficient Representation Learning](https://mlsanity.com/paper/2408.17053/) (2024)
- [A new method for clustered survival data]{A new method for clustered survival data: Estimation of treatment effect heterogeneity and variable selection](https://mlsanity.com/paper/2206.08271/) (2022)
- [Survival Analysis as Imprecise Classification with Trainable Kernels](https://mlsanity.com/paper/2506.10140/) (2025)
- [BENK: The Beran Estimator with Neural Kernels for Estimating the Heterogeneous Treatment Effect](https://mlsanity.com/paper/2211.10793/) (2022)

## Abstract (arXiv)

This work presents a novel attention-based framework for estimating the Individual Probability of Treatment Benefit (IPTB) in survival analysis contexts. The proposed model, called Surv-IPTB, directly quantifies the probability that a specific patient will experience extended survival time under treatment versus control. We reformulate IPTB estimation as a binary classification problem, leveraging pairwise patient comparisons across treatment and control cohorts. The framework incorporates a principled handling of right-censored observations through imprecise probability representations, where uncertain treatment effects are characterized by interval-valued probabilities. An attention mechanism with learnable query-key transformations enables flexible, data-driven aggregation of pairwise comparisons, while simultaneously learning soft class probabilities for censored cases. Through extensive experiments on synthetic datasets with complex nonlinear structures, including spiral, bell-shaped, and circular feature spaces, we demonstrate that our approach maintains robust performance across varying censoring rates and treatment effect strengths. The model consistently outperforms meta-learner baselines (T-learner and S-learner) equipped with random survival forests, Cox proportional hazards, and Beran estimators, particularly in challenging nonlinear scenarios where conventional methods exhibit significant degradation. The results establish the proposed attention-based framework as a scalable and statistically principled solution for personalized treatment benefit assessment in survival settings. The code implementing the model is publicly available.