SurvivalPFN：通过上下文贝叶斯推断进行生存预测的摊销

# SurvivalPFN: 通过上下文贝叶斯推理进行生存预测的摊销  
来源：https://arxiv.org/html/2605.15488  

Shi\-ang Qi¹，Vahid Balazadeh¹,²，Michael Cooper¹,²，Russell Greiner³,⁴，Rahul G\. Krishnan¹,²  
¹Vector Institute  
²多伦多大学  
³阿尔伯塔大学  
⁴阿尔伯塔机器智能研究所  

###### 摘要  

生存分析为在存在删失的情况下建模时间至事件结果提供了强大的统计框架。然而，从众多专门的生存方法中选择合适的估计量通常需要大量的方法和领域专业知识。我们引入了 SurvivalPFN，一种先验数据拟合网络，通过上下文学习摊销删失观测的贝叶斯推理。SurvivalPFN 在多样化的合成、可识别且右删失的数据生成过程族上进行预训练，从而能够在推理时通过单次前向传播摊销生存分析。因此，该模型能够适应每个数据集的有效复杂性，无需特定任务的训练或超参数调整，避免了限制性参数假设，并生成校准的生存分布。在一个涵盖 61 个数据集、21 种方法和 5 个评估指标的大规模基准测试中，SurvivalPFN 表现出强大的预测性能，并且常常优于已有的生存模型。这些结果表明 SurvivalPFN 为生存分析提供了一个有原则且实用的基础模型，在医疗保健、金融和工程等高级别影响领域具有潜在应用 (https://github.com/rgklab/SurvivalPFN)。  

## 1 引言  

参考图例  
图 1：跨 61 个数据集和 5 个指标的计算效率与性能。SurvivalPFN 在匹配经典模型速度的同时取得了最佳中位数排名。  

生存分析建模事件发生时间分布，应用涵盖医学[56 (https://arxiv.org/html/2605.15488#bib.bib2)、92 (https://arxiv.org/html/2605.15488#bib.bib3)、81 (https://arxiv.org/html/2605.15488#bib.bib153)、10 (https://arxiv.org/html/2605.15488#bib.bib6)、9 (https://arxiv.org/html/2605.15488#bib.bib95)、20 (https://arxiv.org/html/2605.15488#bib.bib94)]、电子商务[59 (https://arxiv.org/html/2605.15488#bib.bib14)、74 (https://arxiv.org/html/2605.15488#bib.bib15)、16 (https://arxiv.org/html/2605.15488#bib.bib22)]、工程[72 (https://arxiv.org/html/2605.15488#bib.bib16)、6 (https://arxiv.org/html/2605.15488#bib.bib20)、51 (https://arxiv.org/html/2605.15488#bib.bib23)]和金融[60 (https://arxiv.org/html/2605.15488#bib.bib25)、23 (https://arxiv.org/html/2605.15488#bib.bib24)、25 (https://arxiv.org/html/2605.15488#bib.bib26)]。这些模型通常在表现出右删失的数据上进行学习和评估：对于某些实例，事件在随访期间未被观测到，因此我们只知道事件时间超过删失时间。  

已经提出了多种生存分析方法来处理右删失数据，但每种方法都施加了不同的归纳偏差。经典模型如 Cox 比例风险 (CoxPH) [11 (https://arxiv.org/html/2605.15488#bib.bib38)] 通常依赖于恒定风险比和线性协变量效应。集成方法和深度生存模型提高了灵活性，但通常需要仔细调整，并且常常通过参数形式[85 (https://arxiv.org/html/2605.15488#bib.bib65)、71 (https://arxiv.org/html/2605.15488#bib.bib129)]、比例风险[45 (https://arxiv.org/html/2605.15488#bib.bib29)]、固定时间/分位数离散化[103 (https://arxiv.org/html/2605.15488#bib.bib82)、55 (https://arxiv.org/html/2605.15488#bib.bib30)、73 (https://arxiv.org/html/2605.15488#bib.bib127)]或基于混合的连续时间分布[68 (https://arxiv.org/html/2605.15488#bib.bib32)、32 (https://arxiv.org/html/2605.15488#bib.bib125)]来保留结构性假设。因此，从业者必须驾驭大量具有不同假设和局限性的估计量；模型选择、训练和验证需要大量的领域和方法专业知识。  

本研究旨在设计一种生存估计量，它能够：(i) 避免僵化的简化假设；(ii) 适应观测数据的有效复杂性；(iii) 实现高效推理，无需大量训练或超参数调整。  

参考图例  
图 2：传统生存分析 vs. SurvivalPFN。(左)：传统生存分析需要分析人员为观测数据选择并拟合合适的估计量。(右)：SurvivalPFN 在多样化合成、可识别的 DGP 上进行预训练。推理时，观测数据集作为上下文提供，通过单次前向传播即可获得查询实例的生存分布。  

为此，我们基于先验数据拟合网络 (PFNs) [66 (https://arxiv.org/html/2605.15488#bib.bib87)] 进行构建：这些是基于 Transformer 的模型[100 (https://arxiv.org/html/2605.15488#bib.bib106)]，通过使用合成任务学习后验预测分布的上下文近似。SurvivalPFN 不是为每个数据集拟合一个新的生存模型，而是将计算转移到离线的先验数据预训练阶段。推理时，将观测到的右删失数据集作为上下文，单次前向传播即可返回新个体的后验生存分布。这种方法提供了一条实用的贝叶斯生存预测路径，避免了数据集特定的优化和广泛的超参数调整。  

我们提出 SurvivalPFN，一种通过上下文学习进行生存预测的 Transformer 模型。我们的框架使用条件独立删失下的通用先验来生成数百万个模拟数据生成过程 (DGP)。通过在这些多样化的 DGP 上进行训练，SurvivalPFN 学会直接从观测的右删失数据中推断条件生存分布，从而产生一个易于使用且高效的估计量，具有强大的经验性能；参见图 1 (https://arxiv.org/html/2605.15488#S1.F1)。图 2 (https://arxiv.org/html/2605.15488#S1.F2) 对比了 SurvivalPFN 工作流与传统生存建模。  

我们的主要贡献：  

1. 我们引入了一个通过大规模预训练进行摊销贝叶斯生存预测的框架。SurvivalPFN 使用单次前向传播适应数据复杂性，无需特定任务训练或超参数调整，同时避免了限制性参数假设。  
2. 我们为 SurvivalPFN 提供了理论依据，证明其在可识别右删失数据生成过程下是一个渐近一致的估计量。  
3. 我们进行了大规模基准测试，比较了 21 个模型、61 个数据集和 5 个评估指标，据我们所知，这是迄今为止最大的生存模型基准研究之一。SurvivalPFN 取得了最佳中位数排名。  
4. 我们发布了 SurvivalPFN 的训练代码，以及一个 scikit-learn 风格的 API（参见补充材料）。  

## 2 背景  

生存分析与预测。让 \(X \in \mathbb{R}^d\) 表示协变量向量，\(E, C \in \mathbb{R}_+\) 表示事件时间和删失时间。我们假设在元组 \((X, E, C, T, \Delta)\) 上存在一个真实的联合分布 \(P\)，其中 \(T = \min(E, C)\) 且 \(\Delta = \mathbb{1}[E \leq C]\)。具体来说，我们从观测变量 \(P_{\mathrm{obs}}(X, T, \Delta)\) 的分布中观测到 \(N\) 个样本 \(\mathcal{D} = \{(x_i, t_i, \delta_i)\}_{i=1}^N\)；\(E\) 和 \(C\) 是潜在变量；仅观测到 \(T\) 和 \(\Delta\)。生存预测器旨在学习条件密度或生存函数：
\[
f_{E \mid X}(t \mid x) = \Pr(E = t \mid X = x), \quad \text{或（等价地）} \quad S_{E \mid X}(t \mid x) = \Pr(E > t \mid X = x).
\]

可识别生存分析。如果任何两个候选数据生成过程在 \((X, T, \Delta)\) 上诱导相同的观测定律，则它们必然也会诱导相同的事件时间生存函数，那么我们称条件生存函数 \(S_{E \mid X}\) 是（非参数）*可识别的* [99 (https://arxiv.org/html/2605.15488#bib.bib89)、98 (https://arxiv.org/html/2605.15488#bib.bib45)]：
\[
P_{\mathrm{obs}}^{(1)}(X, T, \Delta) = P_{\mathrm{obs}}^{(2)}(X, T, \Delta) \quad \Longrightarrow \quad S_{E \mid X}^{(1)}(t \mid x) = S_{E \mid X}^{(2)}(t \mid x),
\]
对于几乎所有 \(x\) 和所有在可识别支撑集内的 \(t\)。非参数可识别性的一个充分条件由以下标准假设给出。  

###### 假设 2.1 (条件独立删失)。\(E \perp C \mid X\)。  

###### 假设 2.2 (正性)。对于感兴趣的时间区域 \(\mathcal{T}\)，\(\Pr(C \geq t \mid X = x) > 0\)，\(\forall t \in \mathcal{T}\)。  

当 \(E \not\perp C \mid X\) 时，事件分布通常不能仅从 \((X, T, \Delta)\) 非参数地识别；此时需要额外的假设，例如指定 \(E\) 和 \(C\) 之间依赖性的特定 Copula 族 [24 (https://arxiv.org/html/2605.15488#bib.bib1)、105 (https://arxiv.org/html/2605.15488#bib.bib96)]。附录 B (https://arxiv.org/html/2605.15488#A2) 包含更多关于可识别性的理论。  

贝叶斯生存预测。考虑一个由 \(\theta \in \Theta\) 索引的可识别生存数据生成过程族，在 \(\Theta\) 上具有先验 \(\pi(\cdot)\)。每个 \(\theta\) 为事件时间和删失时间都诱导条件密度和生存函数：\(f_{E \mid X, \Theta}(e \mid x, \theta)\) 和 \(S_{E \mid X, \Theta}(e \mid x, \theta)\) 用于事件；\(f_{C \mid X, \Theta}(c \mid x, \theta)\) 和 \(S_{C \mid X, \Theta}(c \mid x, \theta)\) 用于删失。在贝叶斯生存建模中，我们放置先验密度 \(f_{\Theta}(\theta)\)，并通过贝叶斯规则推断后验密度：
\[
f_{\Theta \mid \mathscr{D}}(\theta \mid \mathcal{D}) \propto f_{\mathscr{D} \mid \Theta}(\mathcal{D} \mid \theta) f_{\Theta}(\theta). \tag{2.1}
\]
在条件独立删失下，似然可以分解为：
\[
f_{\mathscr{D} \mid \Theta}(\mathcal{D} \mid \theta) = \prod_{i=1}^N \left[ f_{E \mid X, \Theta}(t_i \mid x_i, \theta) \, S_{C \mid X, \Theta}(t_i \mid x_i, \theta) \right]^{\delta_i} \left[ f_{C \mid X, \Theta}(t_i \mid x_i, \theta) \, S_{E \mid X, \Theta}(t_i \mid x_i, \theta) \right]^{1-\delta_i}
\]
其中，对于 \(A \in \{E, C\}\)，\(S_{A \mid X, \Theta}(t_i \mid x_i, \theta) = \int_{t_i}^{\infty} f_{A \mid X, \Theta}(\tau \mid x_i, \theta) \, d\tau\)。  
给定一个新的协变量向量 \(x^*\)，事件时间的贝叶斯后验预测分布 (PPD) 为：
\[
f_{E \mid X, \mathscr{D}}(t \mid x^*, \mathcal{D}) = \int_{\Theta} f_{E \mid X, \Theta}(t \mid x^*, \vartheta) f_{\Theta \mid \mathscr{D}}(\vartheta \mid \mathcal{D}) \, d\vartheta. \tag{2.2}
\]
类似地，后验预测生存分布 (PPSD) 为：
\[
S_{E \mid X, \mathscr{D}}(t \mid x^*, \mathcal{D}) = \int_{\Theta} S_{E \mid X, \Theta}(t \mid x^*, \vartheta) f_{\Theta \mid \mathscr{D}}(\vartheta \mid \mathcal{D}) \, d\vartheta. \tag{2.3}
\]
这个框架很有吸引力，因为它整合了多种可能的生存机制，而不是承诺一个单一的拟合模型。然而，将其直接用于灵活的生存模型是困难的：评估似然可能需要数值积分来获得 \(S_{E \mid X, \Theta}\)；方程 2.1 (https://arxiv.org/html/2605.15488#S2.E1) 中的归一化常数以及方程 2.2 (https://arxiv.org/html/2605.15488#S2.E2) 和 2.3 (https://arxiv.org/html/2605.15488#S2.E3) 中的后验预测积分通常是难处理的；并且诸如马尔可夫链蒙特卡洛 (MCMC) [70 (https://arxiv.org/html/2605.15488#bib.bib100)、1 (https://arxiv.org/html/2605.15488#bib.bib101)、102 (https://arxiv.org/html/2605.15488#bib.bib102)] 或变分推理 (VI) [43 (https://arxiv.org/html/2605.15488#bib.bib103)、101 (https://arxiv.org/html/2605.15488#bib.bib104)、35 (https://arxiv.org/html/2605.15488#bib.bib105)、79 (https://arxiv.org/html/2605.15488#bib.bib133)] 等近似推理方法必须为每个新数据集重新运行。因此，我们寻求一种摊销过程，该过程保留贝叶斯生存预测的后验预测解释，同时避免数据集特定的后验计算。  

先验数据拟合网络与摊销贝叶斯推理。先验数据拟合网络 (PFNs) 通过从先验数据生成过程中采样的合成任务上训练 Transformer 来摊销贝叶斯后验预测 [66 (https://arxiv.org/html/2605.15488#bib.bib87)、67 (https://arxiv.org/html/2605.15488#bib.bib86)]。每个任务由一个上下文集和查询输入组成，PFN 被训练根据先验诱导的后验预测分布来预测查询目标。预训练后，后验推理不再是显式的：Transformer 的上下文计算将新数据集和查询点直接映射到单次前向传播中的预测分布，取代了数据集特定的 MCMC 或 VI。这将 PFNs 与元学习 [17 (https://arxiv.org/html/2605.15488#bib.bib88)] 联系起来，但用上下文推理取代了任务特定的适应。PFNs 在表格预测 [36 (https://arxiv.org/html/2605.15488#bib.bib91)、37 (https://arxiv.org/html/2605.15488#bib.bib92)、84 (https://arxiv.org/html/2605.15488#bib.bib154)]、因果效应估计 [2 (https://arxiv.org/html/2605.15488#bib.bib93)、88 (https://arxiv.org/html/2605.15488#bib.bib164)] 和时间序列预测 [38 (https://arxiv.org/html/2605.15488#bib.bib163)、3 (https://arxiv.org/html/2605.15488#bib.bib165)] 中取得了强大的迁移性能，这激励我们使用这种范式进行摊销贝叶斯生存预测。  

## 3 方法  

### 3.1 SurvivalPFN：摊销后验预测推理  

概述。SurvivalPFN 学习一个关于右删失生存数据的贝叶斯后验预测推理的上下文近似。我们没有为每个数据集指定一个可处理的似然并单独进行后验推理，而是通过一个关于可识别右删失 DGP 的模拟器来指定先验。一次抽样 \(\theta \sim \pi(\cdot)\) 决定了在 \((X, E, C, T, \Delta)\) 上的联合分布 \(P^\theta\)。对于每个合成任务，模拟器产生一个观测到的右删失上下文数据集 \(\mathcal{D}^{tr}_\theta = \{(x_i, t_i, \delta_i)_\theta\}_{i=1}^N\)，以及留出的查询协变量 \(x^*_\theta\) 及其潜在事件和删失时间 \((e^*_\theta, c^*_\theta)\)。潜在时间仅用于先验数据训练；推理时，SurvivalPFN 接收与普通生存预测中相同的信息：一个观测数据集 \(\mathcal{D}\) 和查询协变量 \(x^*\)。  

架构。让 \(q_\omega\) 表示一个具有参数 \(\omega\) 的 Transformer。给定一个上下文数据集和一个查询协变量，SurvivalPFN 输出一个关于时间的预测分布。我们还提