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

2026-08-06

AI Summary

This study addresses the challenge of estimating individual probability of treatment benefit (IPTB) in survival analysis, particularly under complex scenarios involving heavy right-censoring and nonlinear treatment effects. The authors propose a novel approach that reformulates IPTB estimation as a pairwise binary classification problem between treated and control patients. By introducing a learnable query-key attention mechanism, the method flexibly aggregates cross-group comparison information and, for the first time, explicitly models censoring-induced uncertainty through interval-based probability representation, enabling end-to-end soft probabilistic learning. Extensive experiments on diverse nonlinear synthetic datasets demonstrate that the proposed method significantly outperforms established baselines—including T- and S-learners combined with random survival forests, Cox proportional hazards models, and Beran estimators—exhibiting remarkable robustness especially under high censoring rates and weak treatment effects.

Abstract

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.

Problem

• Individual Probability of Treatment Benefit
• Survival Analysis
• Right-censoring
• Personalized Treatment Effect
• Treatment Benefit Estimation

Innovation

• attention mechanism
• individual probability of treatment benefit
• survival analysis
• right-censoring
• imprecise probability