SurvivalPFN: Amortizing Survival Prediction via In-Context Bayesian Inference¶

会议: NeurIPS 2026
arXiv: 2605.15488
OpenReview: https://openreview.net/forum?id=bJ348fNOVg
代码: https://github.com/rgklab/SurvivalPFN
领域: 其他
关键词: 生存分析, 先验数据拟合网络, 上下文学习, 贝叶斯推断, 表格基础模型

一句话总结¶

针对生存分析中右删失数据建模难、高度依赖特定算法假设与繁琐超参数调优的痛点，本文提出首个基于上下文贝叶斯推断的生存分析基础模型 SurvivalPFN，通过在大规模可辨识合成生存数据生成过程上离线预训练，仅需单次前向传播即可自适应输出校准且精准的个体后验生存分布。

研究背景与动机¶

生存分析（Survival Analysis）是统计学与机器学习中建模时间-事件（Time-to-Event）结果的核心分支，在临床医学风险分层、重症监护、可靠性工程、用户流失预测以及金融风控等领域扮演着至关重要的角色。这类任务的关键特质在于普遍存在“右删失”（Right-Censoring）现象：对部分样本而言，在有限的随访或观测窗口截止前目标事件仍未发生，算法仅能获知真实事件发生时间严格大于记录的删失时间。为了处理这一不完整观测问题，传统统计学发展出了以 Cox 比例风险模型为代表的半参数方法，但其严格受限于恒定风险比和线性协变量效应的强假设；随后的机器学习与深度生存分析模型（如随机生存森林 RSF、DeepSurv、DeepHit、MTLR 以及连续时间神经模型 DSM、SuMoNet）大幅提升了拟合非线性的灵活性，却不可避免地引入了繁重的超参数调优负担，并在小样本场景下极易过拟合，同时依然受制于预设的时间离散网格或特定参数化分布族。

对于追求不确定性量化的从业者而言，贝叶斯生存分析能够通过对模型参数赋予先验并积分后验分布来刻画认知不确定性，然而在复杂非线性模型中，似然项往往需要数值积分求解，后验归一化常数与后验预测积分高度难解；依赖 MCMC 采样或变分推断（VI）的近似算法必须针对每一个新数据集从头重新迭代拟合数千轮，计算开销极其高昂且极难推广至现实高吞吐部署流程中。更为严峻的是，Tsiatis 非可辨识性定理指出，在缺乏先验假设时，仅从观测到的二元元组数据中根本无法唯一辨识潜在事件时间与删失时间的联合分布；若先验构造不当引入了不可辨识的相依删失机制，哪怕算力无限也无法收敛至真实风险分布。

面对灵活性与高效性不可兼得、以及每个下游数据集均需耗费大量领域经验进行调参适配的核心矛盾，本文的切入角度是将生存预测的贝叶斯后验推断过程彻底“分摊”（Amortize）到预训练阶段。受先验数据拟合网络（PFN）理念启发，既然 Transformer 具备在上下文中逼近贝叶斯后验预测分布的元学习潜能，便可以通过模拟器离线生成涵盖数百万个具有理论可辨识性保障的合成生存数据生成过程（DGP）。核心 idea：将右删失生存分析转化为上下文贝叶斯后验预测任务，在满足条件独立删失的可辨识先验空间上离线预训练 Transformer 基础模型 SurvivalPFN，在推理阶段将任意观测数据集与查询样本作为上下文输入，仅凭单次前向传播直接输出高精度、校准的完整生存分布。

方法详解¶

整体框架¶

SurvivalPFN 的核心思想是将原本针对单一数据集的耗时贝叶斯后验计算转移至离线合成先验预训练阶段。在推理部署时，模型接收包含 $N$ 个观测样本的右删失数据集 $\mathcal{D} = \{(x_i, t_i, \delta_i)\}_{i=1}^N$（其中 $x_i \in \mathbb{R}^d$ 为协变量，$t_i = \min(e_i, c_i)$ 为观测时间，$\delta_i = \mathbb{I}[e_i \le c_i]$ 为事件发生指示标）作为上下文，以及待预测个体的协变量查询 $x^*$。通过引入一个显式的查询目标指示标 $\widetilde{\delta}^* = 1$（指定输出事件时间分布），网络直接在连续时间上输出后验预测生存分布（PPSD）的概率质量向量，并通过单调逆映射得到精确的个体生存曲线 $\widehat{S}(t \mid x^*, \mathcal{D})$。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["右删失观测数据集上下文<br/>D = {(x_i, t_i, δ_i)} 及查询 x*"] --> B["自适应单调时间变换<br/>lognormal2normal 连续坐标映射"]
    B --> C["非对称注意力 Transformer 编码器<br/>上下文双向交互 + 查询单向交叉汇聚"]
    C --> D["目标可控预测头<br/>输入控制标 δ~*=1 请求事件分布"]
    D --> E["L=1024 离散分桶 Softmax<br/>直方图尾部概率累加求和"]
    E --> F["输出校准的个体生存曲线<br/>S_hat(t|x*, D) 单次前向秒级生成"]

关键设计¶

1. 可辨识先验生成空间与混合合成机制：杜绝虚假关联并提供广义归纳偏置

根据 Tsiatis 非辨识性定理，若潜在事件时间 $E$ 与删失时间 $C$ 存在任意相依关系，无限样本的观测分布 $(X, T, \Delta)$ 仍对应无数个等价的潜在事件分布。为使模型学习到的后验分布具备渐近一致性，SurvivalPFN 严格在条件独立删失假设（$E \perp C \mid X$）与正性假设约束下构建仿真生成器 $\pi(\cdot)$。先验通过随机多层感知机（Random MLP + 逐层高斯噪声扰动）首先生成多维非线性协变量特征 $X$，随后分别构建独立的网络块生成潜在事件时间 $E$ 与删失时间 $C$，保证二者在以协变量为条件时完全独立。为覆盖极其异构的现实生存数据特征，模拟器融合了四类先验机制：朴素先验（Naive Prior，直接取随机网络输出）、生存分布先验（Survival-Distribution Prior，利用 Bernstein 单调多项式映射将均匀噪声变换为光滑正值时间）、混合先验（Mixture Prior，依赖协变量的 Weibull 和对数正态有限混合模型）以及大杂烩元先验（Kitchen-Sink Prior，动态加权混合上述先验）。同时，模拟器混合了均匀删失、随机删失、行政截断和条件独立删失四类机制。理论证明（Proposition 3.1 与 C.1）表明，在可辨识先验支撑集下，贝叶斯后验预测生存分布在样本量 $N \to \infty$ 时几乎必然收敛到真实的条件生存函数： $$S_{E \mid X, \mathcal{D}}(t \mid x^*, \mathcal{D}) = \int_{\Theta} S_{E \mid X, \Theta}(t \mid x^*, \theta) f_{\Theta \mid \mathcal{D}}(\theta \mid \mathcal{D}) d\theta \xrightarrow[N \to \infty]{\text{a.s.}} S_{E \mid X, \Theta}(t \mid x^*, \theta^*)$$ 若 SurvivalPFN 充分拟合该目标，其输出曲线将完全继承该一致性保证。

2. 目标可控查询指示器与非对称注意力编码：统一双向预测并保障置换不变性

为使单个模型既能理解上下文中的删失关系，又能灵活解耦预测事件与删失分布，SurvivalPFN 在查询端引入了二元目标控制指示标 $\widetilde{\delta}^* \in \{0, 1\}$。当 $\widetilde{\delta}^* = 1$ 时，模型明确被要求输出潜在事件时间 $E$ 的后验预测分布，其尾部累加即为生存曲线；当 $\widetilde{\delta}^* = 0$ 时，模型则输出潜在删失时间 $C$ 的分布。在预训练过程中，算法采用了三种调度方案（只监督事件、成对复制查询同时监督两者、或按经验删失率随机采样），迫使网络内部建立起事件风险与删失风险的完整统计表征。在网络架构上，模型采用无位置编码的 20 层 Transformer 编码器（隐藏维度 384，RMS QK 归一化，SwiGLU 前馈块，约 20M 参数）。通过非对称注意力掩码设计，上下文 token 之间互相全连接注意力交互，而查询 token 仅单向聚焦于上下文 token，彼此之间无任何信息交互。这种设计严格保证了输出对于上下文样本顺序的置换不变性（Permutation Invariance），并在单次前向传播中支持任意数量的测试样本并行批处理推断。

3. 上下文自适应单调时间变换与离散直方图参数化：克服尺度漂移并实现高分辨率建模

真实生存数据集的时间跨度差异巨大（从几天到数十年），且时间分布高度右偏重尾，若直接在原始时间域上设置固定网格分桶，不仅无法泛化到跨度相异的未知数据集，还会导致大量尾部稀疏区域浪费计算分辨率。SurvivalPFN 在将时间输入网络和进行目标离散化前，设计了基于上下文自适应拟合的单调递增时间变换 $g_{\mathcal{D}}(\cdot)$。模型重点采用了基于对数正态参数化映射的 lognormal2normal： $$g_{\mathrm{LN}}(t) = \frac{\log t - \mu}{\sigma}, \quad \sigma^2 = \log\left(1 + \frac{s^2}{m^2}\right), \quad \mu = \log(m) - \frac{1}{2}\sigma^2$$ 其中 $m$ 和 $s^2$ 为上下文观测时间的经验均值与方差。该变换在保持时间单调排序的同时平滑拉伸了右偏长尾，并自然赋予了模型外推到观测最大时间之外的能力。网络在变换后的标准化潜空间中划分 $L = 1024$ 个等宽分桶，通过 Softmax 输出概率质量向量 $q_\omega$。测试时的个体生存概率直接通过截断累加求和并经逆变换映射回真实时间轴： $$\widehat{S}_\omega(t \mid x^*, \mathcal{D}) = \sum_{\ell: g_{\mathcal{D}}^{-1}(\mathcal{I}_\ell) > t} q_{\omega, \ell}(x^*, \widetilde{\delta}^* = 1, \mathcal{D})$$

损失函数 / 训练策略¶

SurvivalPFN 在预训练时利用模拟器提供的潜变量真实事件时间 $e_\theta^*$ 与真实删失时间 $c_\theta^*$ 构建监督信号。给定查询指示标 $\widetilde{\delta}^*$，目标变量被定义为 $r_\theta^*(\widetilde{\delta}^*) = \widetilde{\delta}^* e_\theta^* + (1 - \widetilde{\delta}^*) c_\theta^*$。设 $\kappa_{\mathcal{D}^{tr}}(r) \in \{1, \ldots, L\}$ 为变换时间落入的目标分桶索引，模型通过最小化离散负对数似然（Negative Log-Likelihood, NLL）展开端到端训练： $$\mathcal{L}_{\mathrm{NLL}}(\omega) = \mathbb{E}_{\theta \sim \pi(\cdot)} \mathbb{E}_{\mathcal{D}^{tr}, x^*, e^*, c^*, \widetilde{\delta}^*} \left[ -\log q_{\omega, \kappa_{\mathcal{D}^{tr}}(r^*(\widetilde{\delta}^*))}(x^*, \widetilde{\delta}^*, \mathcal{D}^{tr}) \right]$$ 在群体期望最优解下，最小化该似然损失所得到的贝叶斯最优预测器，正是以观测上下文与查询特征为条件的真实后验预测分布。

训练策略分为两个阶段： 1. 预测预训练阶段（Predictive Phase）：基于通用表格预测先验初始化 Transformer 骨干，在 4 张 NVIDIA A100 GPU 上训练约 1 周，最大上下文长度设为 16,384，使用 AdamW 优化器配合学习率预热与余弦衰减调度； 2. 生存专业化阶段（Survival Phase）：切换到可辨识生存合成先验与离散生存似然损失，在 1 张 NVIDIA H100 GPU 上训练 2.5 天，最大上下文长度设为 2,048，并引入 Schedule-free 优化器以实现平稳快速收敛。

实验关键数据¶

主实验¶

论文构建了涵盖 81 个真实世界生存数据集的大规模评测基准，其中 20 个数据集仅用于验证模型检查点选择，其余 61 个独立数据集用于完全留出评估。评估对比了来自 5 大阵营的 21 种代表性模型，并在 5 大核心指标上通过 10 次重复划分（70% 训练 / 30% 测试）计算统计表现： - 综合 Brier 分数（IBS，概率分布准确性，越低越好） - 一致性指数（CI，排序辨别能力，越高越好） - D-calibration 卡方值（分布校准度，越低越好） - 伪观测平均绝对误差（MAE-PO，时间点估计精度，越低越好） - Log-Rank 统计量（与真实时序事件的符合度，越低越好）

下表汇总了跨全部 61 个数据集汇聚的中位数排名（Median Rank，中位数越小代表性能越好，带 95% Bootstrap 置信区间）及运行时间对比：

模型分类	模型名称	综合排名 (Overall)	IBS 排名	CI 排名	D-calib 排名	MAE 排名	运行耗时特性
表格基础模型 (本文)	SurvivalPFN (Ours)	3.8 [3.2, 4.4]	4.0 [3.0, 5.0]	6.0 [5.0, 7.5]	6.0 [4.5, 8.0]	4.0 [3.0, 5.5]	秒级推断 (仅次于CoxNet)
表格基础模型包装	StaticSurvivalTFM (TabDPT)	12.0 [11.0, 13.5]	13.0 [11.5, 14.5]	13.0 [11.0, 14.5]	15.0 [13.0, 16.5]	10.0 [8.5, 12.0]	受限于展开样本上下文膨胀
经典统计模型	CoxPH	14.0 [12.5, 15.0]	14.0 [12.0, 16.0]	12.0 [10.0, 14.0]	14.0 [12.0, 16.0]	13.0 [11.0, 15.0]	极快但因奇异Hessian失败13次
经典统计模型	CoxNet (Elastic-net)	8.5 [7.5, 10.0]	8.0 [6.5, 10.0]	9.0 [7.5, 11.0]	11.0 [8.5, 13.0]	8.0 [6.5, 10.0]	最快基线但表达能力受限
树模型集成	Random Survival Forest (RSF)	5.5 [4.5, 7.0]	6.0 [4.5, 7.5]	3.5 [2.5, 5.0]	8.0 [6.0, 10.0]	7.0 [5.5, 8.5]	耗时中等，CI 表现强势
树模型集成	CWGB (Boosting)	6.5 [5.0, 8.0]	7.0 [5.5, 8.5]	4.0 [3.0, 6.0]	9.0 [7.0, 11.5]	7.5 [6.0, 9.5]	耗时较高，依赖充分样本
离散神经生存模型	DeepSurv	7.0 [5.5, 8.5]	8.0 [6.0, 9.5]	5.0 [3.5, 7.0]	10.0 [7.5, 12.0]	8.0 [6.0, 10.0]	需调参，在大表上可能退化
离散神经生存模型	DeepHit	9.0 [7.5, 10.5]	9.5 [8.0, 11.5]	5.5 [4.0, 7.5]	12.0 [9.5, 14.0]	10.0 [8.0, 12.0]	排序项强但校准度较弱
离散神经生存模型	MTLR	7.5 [6.0, 9.0]	7.0 [5.0, 8.5]	8.0 [6.5, 10.0]	7.5 [5.5, 10.0]	8.0 [6.5, 10.0]	性能平衡但需精细网格设置
离散神经生存模型	IWSG	13.0 [11.0, 14.5]	12.0 [10.0, 14.0]	11.0 [9.0, 13.5]	13.0 [10.5, 15.0]	12.0 [10.0, 14.0]	训练极其缓慢，5个数据集超时
连续神经生存模型	DSM (Weibull Mixture)	10.0 [8.0, 11.5]	10.0 [8.0, 12.0]	10.0 [8.0, 12.0]	9.0 [7.0, 11.5]	11.0 [9.0, 13.0]	混合参数难优化，多次NaN
连续神经生存模型	SuMoNet (Monotonic NN)	11.0 [9.5, 12.5]	11.0 [9.0, 13.0]	12.0 [10.0, 14.0]	8.5 [6.5, 11.0]	11.0 [9.5, 13.0]	梯度反向传播计算代价大
连续神经生存模型	SurvivalMDN	9.5 [8.0, 11.0]	9.0 [7.0, 11.0]	9.5 [7.5, 12.0]	8.0 [6.0, 10.5]	10.0 [8.0, 12.0]	高斯混合耗显存，超显存崩溃

消融实验¶

论文在验证集上针对预训练初始化、先验族结构、时间变换映射、训练损失函数、查询指示标调度以及训练样本比例等核心设计开展了系统性消融评测。下表提取了最具代表性的消融配置对比表现：

变体版本	预测预训练	先验生成族	时间变换方式	优化损失	查询指示调度	综合中位数排名	核心结论与机制归因
v01 (Default)	✓	Surv.-dist.	lognormal2normal	NLL	Both	3.2	全场最优，IBS 与校准度最佳，具备强外推能力
v02	✓	Surv.-dist.	time2quantile	NLL	Both	5.8	换用经验分位数无法外推时间尾部，导致各指标全线下滑
v04	✓	Surv.-dist.	time2quantile	CE (平滑)	Random	6.4	平滑交叉熵损失在分位数空间下加剧了尾部分布失真
v06	✓	Surv.-dist.	lognormal2normal	NLL	Random	4.1	随机调度相比同时监督两者略微损失多任务泛化性
v08	✓	Kitchen-sink	lognormal2normal	NLL	Event-only	4.3	大杂烩元先验增加多样性但略微稀释了平滑先验的聚焦性
v10	✓	Surv.-dist.	lognormal2normal	NLL	Event-only	3.9	仅监督事件时间表现稳定，但略逊于同时预测事件与删失
v12	✓	Naive (MLP)	lognormal2normal	NLL	Event-only	5.2	朴素先验缺乏真实生存分布形态，下游概率校准度明显变差
v13	✗	Naive (MLP)	lognormal2normal	CE (平滑)	Random	7.9	缺失通用预训练初始化导致 Transformer 难以自适应多维表格

关键发现¶

小样本场景具备压倒性优势：在样本量 $N < 500$ 的 24 个小型数据集上，SurvivalPFN 相比所有经典方法、树模型与深度学习模型形成了断层式领先优势。传统深度学习模型因参数量过大在小样本下严重过拟合，而 SurvivalPFN 得益于预训练期间摄入的海量先验函数归纳偏置，单次前向即可自适应匹配数据有效复杂度。
删失率鲁棒性极强：在低删失（$<33\%$）、中等删失（$33\%-67\%$）与极高删失（$\ge 67\%$）各个子集切片下，SurvivalPFN 始终稳居第一梯队，表现出极其稳健的跨删失域泛化性。
计算效率与运行可靠性兼备：与耗时巨大、频繁遭遇梯度爆炸、矩阵奇异或超时崩溃的传统基线相比（如 CoxPH 奇异失败 13 次，IWSG 超时失败 5 次），SurvivalPFN 在全部 61 个数据集上 100% 成功运行，总推断时间比调参的深度网络快 2-3 个数量级，逼近最轻量的线性模型。
优于通用表格模型直接适配：对比将未删失数据直接送入通用表格模型（TabPFN、TabICL、MITRA）的方案，SurvivalPFN 在 IBS、D-calibration 与 CI 上全面胜出，证明为生存分析定制“可辨识删失先验”是构建可靠生存基础模型的决定性前提。

亮点与洞察¶

将高难贝叶斯推断成功分摊至前向传播：首次在生存分析领域成功落地 PFN 范式，通过“合成数据元学习”绕开了令人生畏的 MCMC 采样与变分后验近似，使任意真实数据集均可在毫秒级获得贝叶斯后验预测分布，极大拓宽了生存分析基础模型的实用化边界。
理论可辨识性与数据模拟器的完美闭环：不仅停留在经验拟合，更以 Tsiatis 定理为出发点，深入剖析了右删失可辨识性条件，并在仿真器中确立了 $E \perp C \mid X$ 的充分非参数识别约束，给出了严格的后验预测渐近一致性证明。
目标指示标 $\widetilde{\delta}^*$ 的优雅解耦设计：通过在查询 token 中附加轻量级二元控制位，以最小的结构改动实现了网络对潜在事件过程与潜在删失过程的联合感知，这一设计可天然扩展至竞争风险（Competing Risks）、区间删失与因果反事实推断。

局限与展望¶

受限于条件独立删失的核心假设：若现实数据中存在强烈的信息性相依删失（Dependent Censoring，例如患者因病情极度恶化而提前退出研究导致 $E$ 与 $C$ 在协变量外仍高度相依），非参数可辨识性不复存在，SurvivalPFN 将无法从单纯观测数据中还原真实边际生存概率。未来需探索结合 Copula 结构或工具变量等可辨识相依删失先验。
大样本长序列下的性能边际递减：当数据集规模极大（$N \ge 5000$）时，由于 Transformer 上下文长度（2048）的硬性限制以及长序列下信息聚合的压缩效应，SurvivalPFN 相比能够从海量数据中充分端到端微调的深度网络，相对优势有所收窄。未来需结合线形注意力、状态空间模型（SSM）或分层检索增强机制以突破长上下文瓶颈。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次将先验拟合网络引入生存分析领域，形式化了可辨识先验空间与单次前向贝叶斯分摊推断。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 81 个多领域数据集、21 种基线模型、5 大评测指标及全面的样本量/删失率切片分析。
写作质量: ⭐⭐⭐⭐⭐ 理论证明完整严密，方法叙述层次清晰，实验分析深刻自洽。
价值: ⭐⭐⭐⭐⭐ 为高风险医疗、金融及工业领域提供了一个开箱即用、校准精准的高可靠生存分析基础模型底座。