论文提出了 Surv-IPTB,用注意力机制做个体治疗获益估计,比传统 T-learner、S-learner 在复杂非线性数据上更稳,代码已开源。
Surv-IPTB 是一种用于生存分析的新框架,直接量化患者在接受治疗相较于对照时生存期延长的概率。该方法将 IPTB 估计重构为二分类问题,利用治疗组与对照组之间的成对患者比较,并通过不精确概率表示处理右删失观测。模型采用带可学习查询-键变换的注意力机制,灵活聚合成对比较,同时为删失病例学习软类别概率。在螺旋、钟形和圆形等非线性合成数据集上的实验中,Surv-IPTB 在多种删失率和治疗效果强度下表现稳健,优于配备随机生存森林、Cox 比例风险和 Beran 估计器的元学习基线。
Surv-IPTB: An Attention-Based Model for Estimating Individual Probability of Treatment Benefit with Survival Data
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.