临床预测模型里有一个问题经常被低估我们真正想回答的往往不是“这个治疗平均有效”而是“眼前这个病人到底该不该用这个治疗”。Surv-IPTB 这个方向把问题聚焦在生存数据下的个体治疗获益概率Individual Probability of Treatment Benefit上并用注意力机制把病人特征、治疗方案和风险结果联系起来。生存数据和普通分类数据最大的不同在于我们看到的不是“是否发生事件”而是“在多久之后发生事件”中间还夹着删失。这里不假设你已经拿到某篇论文的完整源码而是从建模思路出发用 PyTorch 实现一个最小可运行的 Surv-IPTB 示例覆盖数据构造、模型训练、获益评估和常见问题排查。1. 先理解生存数据里的“个体获益”到底在算什么1.1 生存数据与普通监督数据的差别普通监督学习的样本通常是(特征, 标签)标签是类别或者数值。生存数据里标签变成两个字段的组合观察时间time和事件指示event。例如在一个入组病人样本中time 18.5表示这个病人被跟踪了 18.5 个月event 1表示在 18.5 个月时观察到了终点事件event 0表示 18.5 个月时还没有发生事件病人仍然存活、失访或研究结束时仍未发生事件。event 0的样本就是删失样本。删失不等于没有事件只是我们在有限观察窗口内没有看到事件。如果把删失样本当作“未发生事件”的普通负样本会低估事件风险如果把删失样本直接丢弃又会损失大量信息。生存分析要处理的核心问题就是在存在删失的情况下估计事件发生时间或风险。在代码层面数据通常长这样字段含义示例x0 ... x7病人基线特征年龄、血压、化验指标treatment治疗方案0 表示对照1 表示治疗time观察时间12.5event终点事件1 表示发生0 表示删失1.2 平均治疗效果不能直接回答个体问题随机对照试验报告里经常出现“治疗组中位生存期比对照组延长 3 个月”或者“治疗组死亡风险降低 20%”。这些是平均治疗效果ATEAverage Treatment Effect是两组人群层面的比较。但临床决策发生在个体层面。有的病人可能从治疗中获益很大有的病人可能没有获益甚至可能承受毒副作用却没有获益。平均治疗效果可能来自少数极端获益者也可能被大量无获益者稀释。个体治疗获益就是要估计给定病人特征x治疗让这个病人的风险下降了多少获益概率有多大。这和 ITEIndividual Treatment Effect的思路一致只不过结果变量是生存时间和事件指示而不是单一标签。1.3 IPTB 的建模对象IPTB 在生存数据中的常见定义方式是预测病人在接受治疗 t 时的风险率 λ(t)预测同一个病人在接受对照或另一种治疗时的风险率 λ(c)计算风险差异或风险比用多次预测或不确定性估计把“差异为负”的概率作为获益概率。如果 λ(t) 小于 λ(c)说明治疗降低了该病人的风险可以认为治疗对这个病人有正向获益。这里有一个关键点真实世界里同一个病人只能接受一种治疗我们永远无法同时观测到接受治疗和接受对照的两个结局。这就是反事实缺失。建模时必须依赖可忽略性、正值性等因果假设否则模型拟合的只是相关关系而不是真正的治疗获益。这个在后面的最佳实践里会单独展开。2. 为什么这里要用注意力机制而不是只加一个交互项2.1 Cox 回归处理交互项的局限传统生存分析中最常用的是 Cox 比例风险模型h(t | x) h0(t) * exp(beta * x)要估计个体治疗获益一种常见办法是在模型里加入treatment与某些特征的交互项h(t | x) h0(t) * exp(beta1 * treatment beta2 * x beta3 * treatment * x)这么做有两个工程问题交互项需要人为指定。治疗效应可能同时受五六个特征影响而且影响方式不是简单线性乘积特征之间也存在相互作用。比如“年龄大 肾功能差 治疗强度高”才出现高风险这种高阶交互靠手工构造几乎不可能穷尽。2.2 注意力机制在表格数据里能做什么注意力机制最早来自序列模型核心操作是 query 去查询一组 key-value 中哪些信息更重要。在表格数据里我们可以把每个特征当成一个“token”让模型自己学会在预测风险时应该更关注哪些特征。这和传统全连接网络的区别是注意力权重是动态计算的。同一个特征在某个病人身上可能重要在另一个病人身上可能不重要在对照组可能重要在治疗组可能不重要。这正是个体化预测所需要的表达能力。注意力熵值用于衡量注意力分布退化程度。如果所有注意力权重都接近均匀分布熵值很高说明模型没有真正利用注意力机制退化成普通加权平均如果熵值过低说明注意力过度集中到某一个特征也可能是过拟合。后面评估部分会回来检查这个指标。2.3 Surv-IPTB 的交叉注意力设计Surv-IPTB 的直觉是治疗方案应该知道该“看”病人的哪些特征。具体实现采用交叉注意力query来自治疗方案向量key和value来自病人特征编码后的 token 序列注意力输出是一个加权后的病人表示同时保留治疗信息。这样做的好处是模型不再把治疗方案当作一个普通离散特征而是用治疗信息主动去提取病人特征中与获益相关的部分。传统模型里治疗只是一个0/1输入模型很难学习“治疗如何改变了对特征的解读”。3. 环境准备与模拟数据构造3.1 依赖环境这个示例只需要常见的 Python 科学计算和深度学习库。如果原始环境里没有装全先补齐依赖再继续。依赖用途安装方式示例Python 3.9运行环境建议使用虚拟环境numpy数据生成和数组运算pip install numpypandas数据处理pip install pandasscikit-learn训练验证集划分pip install scikit-learnPyTorch模型实现与训练按官网安装对应版本lifelines可选用于 C-index 评估pip install lifelines注意不同机器的 PyTorch 安装命令不同建议到官网确认与 CUDA 版本匹配的方案。CPU 环境也能运行这个最小示例只是训练会慢一些。3.2 用指数分布生成带真实治疗效应的模拟数据模拟数据的好处是“真实获益”是已知的可以验证模型是否学到了个体治疗获益结构。这里生成一个最简单的指数生存模型基线风险率由部分特征决定治疗效应由另外两个特征决定治疗组的风险率发生变化观察时间是事件时间和删失时间的最小值。import numpy as np import pandas as pd def simulate_survival_data(n_samples2000, n_features8, seed42): rng np.random.default_rng(seed) X rng.normal(size(n_samples, n_features)) treatment rng.integers(0, 2, sizen_samples) # 基线 log 风险率由 x0, x1 决定 log_lambda0 -1.5 0.3 * X[:, 0] - 0.2 * X[:, 1] # 个体获益评分由 x2, x3 决定 benefit_score 0.5 * X[:, 2] - 0.4 * X[:, 3] # 治疗组 log 风险率 基线风险率 获益评分 log_lambda log_lambda0 treatment * benefit_score lambda_ np.exp(log_lambda) # 指数分布生成事件时间 u rng.random(sizen_samples) event_time -np.log(1.0 - u) / lambda_ # 独立生成删失时间 censoring_time rng.exponential(scale5.0, sizen_samples) time np.minimum(event_time, censoring_time) event (event_time censoring_time).astype(np.float32) df pd.DataFrame(X, columns[fx{i} for i in range(n_features)]) df[treatment] treatment df[time] time df[event] event # 保留真实获益标记用于后续验证 df[true_benefit] (benefit_score 0).astype(float) return df if __name__ __main__: data simulate_survival_data() print(data.head()) print(data[event].mean(), data[treatment].mean())数据生成时有一个容易忽略的细节治疗效应特征必须进入模型输入否则模型不可能估计个体获益。很多真实项目在特征筛选时提前把“单变量分析不显著”的特征删掉结果把真正的异质性特征也删了这是常见错误。3.3 数据切分与张量化模型训练前需要把 DataFrame 转成 PyTorch 的 Tensor并划分训练集和验证集。import torch from torch.utils.data import DataLoader, TensorDataset from sklearn.model_selection import train_test_split def prepare_dataloaders(data, batch_size128, test_size0.2, seed7): feature_cols [c for c in data.columns if c.startswith(x)] X data[feature_cols].values.astype(np.float32) T data[treatment].values.astype(np.int64) time data[time].values.astype(np.float32) event data[event].values.astype(np.float32) X_tr, X_va, T_tr, T_va, time_tr, time_va, event_tr, event_va train_test_split( X, T, time, event, test_sizetest_size, random_stateseed ) train_ds TensorDataset( torch.from_numpy(X_tr), torch.from_numpy(T_tr), torch.from_numpy(time_tr), torch.from_numpy(event_tr), ) valid_ds TensorDataset( torch.from_numpy(X_va), torch.from_numpy(T_va), torch.from_numpy(time_va), torch.from_numpy(event_va), ) train_dl DataLoader(train_ds, batch_sizebatch_size, shuffleTrue) valid_dl DataLoader(valid_ds, batch_sizebatch_size, shuffleFalse) return train_dl, valid_dl实际项目中不要把train_test_split直接用于嵌套随机或纵向数据否则同一个病人的多条记录可能同时进入训练集和验证集造成信息泄漏。这里的数据是一人一行所以简单切分是可行的。4. 实现 Surv-IPTB 模型并完成训练4.1 模型总览模型分成三部分特征编码把每个原始特征映射成独立的 token让每个特征在后续注意力中都可以被单独加权交叉注意力用治疗方案向量作为 query去查询病人特征 token 中哪些特征与治疗风险最相关风险头把注意力输出和治疗向量拼接输出该样本的 log 风险率。import torch import torch.nn as nn class CrossAttention(nn.Module): def __init__(self, hidden): super().__init__() self.q_proj nn.Linear(hidden, hidden) self.k_proj nn.Linear(hidden, hidden) self.v_proj nn.Linear(hidden, hidden) self.scale hidden ** -0.5 def forward(self, tokens, t_emb): # tokens: (batch, n_features, hidden) # t_emb: (batch, hidden) q self.q_proj(t_emb).unsqueeze(1) # (batch, 1, hidden) k self.k_proj(tokens) # (batch, n_features, hidden) v self.v_proj(tokens) # (batch, n_features, hidden) attn torch.softmax( q k.transpose(-2, -1) * self.scale, dim-1, ) # (batch, 1, n_features) out attn v # (batch, 1, hidden) return out.squeeze(1), attn.squeeze(1) class SurvIPTB(nn.Module): def __init__(self, n_features, hidden64, dropout0.1): super().__init__() self.n_features n_features self.feature_proj nn.Linear(1, hidden) self.feature_pos nn.Parameter(torch.randn(n_features, hidden)) self.treatment_emb nn.Embedding(2, hidden) self.encoder nn.Sequential( nn.Linear(hidden, hidden), nn.ReLU(), nn.Dropout(dropout), nn.Linear(hidden, hidden), nn.ReLU(), ) self.cross_attn CrossAttention(hidden) self.risk_head nn.Sequential( nn.Linear(hidden * 2, hidden), nn.ReLU(), nn.Linear(hidden, 1), ) def forward(self, x, t): # x: (batch, n_features) tokens self.feature_proj(x.unsqueeze(-1)) # (batch, n_features, hidden) tokens tokens self.feature_pos.unsqueeze(0) # 特征位置信息 tokens self.encoder(tokens) t_emb self.treatment_emb(t) # (batch, hidden) context, attn self.cross_attn(tokens, t_emb) # (batch, hidden), (batch, n_features) fused torch.cat([context, t_emb], dim-1) log_lambda self.risk_head(fused).squeeze(-1) return log_lambda, attn这里有一点需要解释为什么把每个原始特征单独投影成 token因为只有每个特征独立成为一个位置注意力权重才能对应到具体特征名上。如果把整个特征向量压缩成一个向量再做注意力注意力就只能返回一个标量可解释性会明显下降。4.2 指数生存模型的损失函数示例使用连续时间指数生存模型假设风险率恒定。给定 log 风险率log_lambda风险率为lambda exp(log_lambda)事件时间服从指数分布。负对数似然为def exp_survival_loss(log_lambda, time, event): lambda_ torch.exp(log_lambda) loss event * log_lambda - lambda_ * time return -loss.mean()这个损失函数很简洁但是指数函数的数值稳定性需要注意。如果log_lambda过大torch.exp(log_lambda)可能溢出。训练时可以对log_lambda做截断或者在模型输出后限制范围log_lambda torch.clamp(log_lambda, min-10.0, max10.0)真实项目里指数风险率假设往往太强。更常见的做法是把时间分成多个桶对每个桶估计一个基线风险或者使用离散时间生存模型。这里的指数形式是为了让最小示例足够清晰生产环境不要直接照搬。4.3 训练循环与模型保存训练循环主要做四件事正向传播、计算损失、反向传播、在验证集上评估并保存最优模型。def evaluate(model, valid_dl): model.eval() total_loss 0.0 n 0 with torch.no_grad(): for xb, tb, ub, eb in valid_dl: log_lambda, _ model(xb, tb) log_lambda torch.clamp(log_lambda, min-10.0, max10.0) loss exp_survival_loss(log_lambda, ub, eb) total_loss loss.item() * xb.size(0) n xb.size(0) return total_loss / n def train_model(model, train_dl, valid_dl, epochs60, lr1e-3): optimizer torch.optim.Adam(model.parameters(), lrlr) best_loss float(inf) for epoch in range(epochs): model.train() total_loss 0.0 n 0 for xb, tb, ub, eb in train_dl: optimizer.zero_grad() log_lambda, _ model(xb, tb) log_lambda torch.clamp(log_lambda, min-10.0, max10.0) loss exp_survival_loss(log_lambda, ub, eb) loss.backward() optimizer.step() total_loss loss.item() * xb.size(0) n xb.size(0) train_loss total_loss / n valid_loss evaluate(model, valid_dl) if valid_loss best_loss: best_loss valid_loss torch.save(model.state_dict(), best_surviptb.pt) if (epoch 1) % 10 0: print(fepoch {epoch 1}: train_loss{train_loss:.4f}, valid_loss{valid_loss:.4f}) print(fbest valid loss: {best_loss:.4f}) model.load_state_dict(torch.load(best_surviptb.pt)) return model训练时要注意 batch size 和特征数量。模拟数据只有 8 个特征用一个较小的 MLP 就能拟合真实数据特征成百上千时特征 token 数量过多注意力矩阵占用的内存会显著增加。4.4 常用训练超参数超参数示例值影响hidden64隐藏层维度过小时注意力表达能力不足过大容易过拟合dropout0.1防止风险头过拟合越大模型越保守learning_rate0.001过大会导致 loss 抖动或 NaN过小收敛很慢batch_size128影响训练稳定性样本少时不要用太大的 batchepochs60用验证集 loss 判断是否早停不要只看训练集在真实项目中建议先固定一个小模型跑通完整流程再做超参数搜索。不要在第一次跑通前就上大模型。5. 如何从模型输出计算个体治疗获益概率5.1 治疗组与对照组的风险率预测模型在训练时只输入了病人实际接受的治疗方案所以一个样本只有一条 log 风险率输出。但在决策时我们需要回答同一个病人“如果接受治疗”和“如果不接受治疗”两种情况的差异。做法是对同一个特征向量分别输入treatment0和treatment1def predict_risk_difference(model, X, clip10.0): model.eval() X torch.as_tensor(X, dtypetorch.float32) n X.size(0) t0 torch.zeros(n, dtypetorch.long) t1 torch.ones(n, dtypetorch.long) with torch.no_grad(): log_lambda0, _ model(X, t0) log_lambda1, _ model(X, t1) log_lambda0 torch.clamp(log_lambda0, min-clip, maxclip) log_lambda1 torch.clamp(log_lambda1, min-clip, maxclip) lambda0 torch.exp(log_lambda0) lambda1 torch.exp(log_lambda1) risk_diff lambda0 - lambda1 return lambda0.numpy(), lambda1.numpy(), risk_diff.numpy()这里risk_diff 0表示治疗组风险率低于对照组也就是治疗降低了风险risk_diff 0表示治疗反而增加了风险。5.2 用 Monte Carlo Dropout 得到获益概率点估计告诉我们的只是“治疗和对照的风险率差多少”还不能直接回答“获益概率多大”。一个工程上比较容易实现的思路是保留模型中的 dropout预测时多次前向统计输出差异的概率分布。def predict_benefit_probability(model, X, n_mc50, clip10.0): model.train() X torch.as_tensor(X, dtypetorch.float32) n X.size(0) t0 torch.zeros(n, dtypetorch.long) t1 torch.ones(n, dtypetorch.long) benefit_flags [] with torch.no_grad(): for _ in range(n_mc): log_lambda0, _ model(X, t0) log_lambda1, _ model(X, t1) log_lambda0 torch.clamp(log_lambda0, min-clip, maxclip) log_lambda1 torch.clamp(log_lambda1, min-clip, maxclip) lambda0 torch.exp(log_lambda0) lambda1 torch.exp(log_lambda1) benefit_flags.append((lambda0 lambda1).float()) benefit_prob torch.stack(benefit_flags, dim0).mean(dim0) return benefit_prob.numpy()注意这里用model.train()模式让 dropout 生效但外层使用torch.no_grad()这样既保留随机性又不更新梯度。Monte Carlo Dropout 不是完整的贝叶斯推断只是一个工程近似适合用于可解释性报告和风险分层不能替代严格的贝叶斯深度生存模型。5.3 与真实获益对比验证模拟数据里保存了true_benefit字段可以直接评估模型是否学到个体化的治疗获益结构。data simulate_survival_data() X_cols [c for c in data.columns if c.startswith(x)] X data[X_cols].values model SurvIPTB(n_featureslen(X_cols)) train_dl, valid_dl prepare_dataloaders(data) model train_model(model, train_dl, valid_dl) benefit_prob predict_benefit_probability(model, X, n_mc50) pred_benefit (benefit_prob 0.5).astype(float) true_benefit data[true_benefit].values acc (pred_benefit true_benefit).mean() print(fbenefit prediction accuracy: {acc:.3f})如果数据生成时治疗效应与特征的关系较强这一准确率通常会明显高于随机水平。如果准确率接近 0.5优先检查治疗效应是否真的进入了输入特征其次检查模型容量和训练是否收敛。5.4 注意力权重与注意力熵值注意力权重可以直接观察到。每个样本的attnshape 是(n_features,)表示治疗方案对这个特征的关注程度。def inspect_attention(model, X, t): model.eval() X torch.as_tensor(X, dtypetorch.float32) t torch.as_tensor(t, dtypetorch.long) with torch.no_grad(): _, attn model(X, t) return attn.numpy() attn inspect_attention(model, X[:5], data[treatment].values[:5]) feature_names [fx{i} for i in range(len(X_cols))] print(pd.DataFrame(attn, columnsfeature_names).round(3))注意力熵值等于对每行权重计算信息熵再除以最大可能熵log(n_features)做归一化def attention_entropy(attn, eps1e-8): attn np.clip(attn, eps, 1.0) entropy -np.sum(attn * np.log(attn), axis-1) max_entropy np.log(attn.shape[-1]) return entropy / max_entropy如果大多数样本的归一化熵值都很接近 1说明注意力分布近似均匀模型没有从注意力机制中充分获益。这个现象通常提示特征投影、位置编码或隐层维度设置有问题而不是注意力机制本身无效。6. 常见问题与排查路径6.1 训练过程异常问题现象常见原因检查方式处理建议loss 变成 NaN学习率过大、时间值过大、指数溢出打印 log_lambda 的均值和最大值降低学习率对时间做归一化对 log_lambda 做 clamploss 下降缓慢特征未标准化检查特征均值方差对连续特征做 StandardScaler验证 loss 持续低于训练 loss可能是 dropout 影响评估方式对比 train/eval 模式确认评估时调用model.eval()训练 loss 很低IPTB 准确率仍然不高模型过拟合训练集但未学到异质性结构查看验证集 loss 和 C-index增加数据量简化模型加入早停6.2 预测结果不合理问题现象常见原因检查方式处理建议所有样本的获益概率都接近 0 或 1治疗效应本身太均匀打印风险差分布检查训练数据中治疗效应是否有异质性同一个样本在前两次预测中获益方向不稳定Monte Carlo 样本太少增加 n_mc 观察波动调大 n_mc 到 100 以上注意力权重几乎均匀模型没有学会特征交互计算注意力熵值增大 hidden检查特征 token 编码删失率过高导致模型效果差事件样本太少统计 event 均值考虑离散时间模型或修改研究终点6.3 排查优先级遇到问题时按下面顺序排查输入数据是否正确特征列、治疗列、时间和事件的 dtype 是否符合预期预处理是否一致训练集和验证集是否用了同一个标准化器模型输出是否合理log_lambda是否被 clampexp是否溢出损失函数是否正确删失样本是否也参与了计算训练是否收敛验证集 loss 是否持续下降评价指标是否匹配生存模型应该用 C-index 或时间相关 AUC而不是普通分类准确率。7. 最佳实践与扩展方向7.1 从最小示例到真实数据的差距模拟数据里没有缺失值、没有异常值、删失机制也完全独立于特征真实数据不会这么干净。落地时需要额外处理缺失特征不能直接填 0要结合业务含义选择中位数填补、模型填补或单独建模连续特征建议标准化治疗分配不平衡时不要直接用原始比例时间变量要统一单位避免不同中心的数据使用不同计量单位外部验证必须和训练数据使用完全一样的预处理流程最好把 StandardScaler 连同模型一起保存。import pickle with open(preprocess.pkl, wb) as f: pickle.dump({scaler: scaler, feature_cols: feature_cols}, f)7.2 因果推断假设用生存数据估计个体治疗获益不能只靠注意力模型解决所有问题。模型的预测能力再强如果数据不满足基本因果假设得到的“获益概率”也可能只是相关关系。需要满足的假设至少包括一致性观测到的结果等于在相同治疗下的潜在结果可忽略性给定特征治疗分配与潜在结果独立正值性每个特征组合下接受治疗的概率在 0 和 1 之间。随机对照试验数据最接近这些假设。观察性数据需要加入倾向得分加权、工具变量或更严格的反事实建模。生产环境里如果知道数据来自回顾性队列必须在结论里明确说明获益概率是相关意义上的估计不能直接声称因果效应。7.3 模型结构扩展最小示例的模型结构可以沿多个方向扩展离散时间生存模型把时间分成 K 个桶对每个桶输出一个风险结果可以去掉恒定风险率假设Transformer 化把治疗向量也作为 token 拼接到特征序列中直接用 Transformer Encoder 堆叠替代单层交叉注意力多治疗臂treatment_emb从Embedding(2, hidden)扩展为Embedding(n_arms, hidden)纵向数据把同一病人多次随访的记录按时间排列用 LSTM 或注意力序列模型处理贝叶斯深度生存模型在风险头上加入变分推断得到更可靠的不确定性估计。7.4 可复用的发布前检查清单训练完成并准备对外发布结果前建议按清单逐项确认是否检查过事件比例和删失比例是否了解删失机制特征标准化器是否只在训练集上拟合训练集、验证集、测试集是否完全独立模型是否保存了预处理参数和特征名是否报告了 C-index 和时间相关 AUC而不是只报告 loss获益概率是否带不确定性或多次预测结果是否明确说明了因果假设的成立条件是否存在数据泄漏比如把未来信息当作基线特征。这个最小示例的核心价值在于它把 Surv-IPTB 这个复杂的名字拆解成了可运行的工程流程。先构造带真实效应的生存数据再实现交叉注意力模型然后通过治疗组和对照组两组预测来计算个体获益概率最后用注意力权重和注意力熵值检查模型是否真的学到了特征交互。把这个流程跑通之后再往真实数据、离散时间模型和因果推断方向扩展会清晰很多。