[论文解读] Countdown Regression: Sharp and Calibrated Survival Predictions

Anand Avati, Tony Duan+4 more|arXiv (Cornell University)|Jun 21, 2018

Machine Learning in Healthcare参考文献 33被引用 18

一句话总结

本文提出了一种名为 Survival-CRPS 的新型评分规则，该规则将连续概率评分（Continuous Ranked Probability Score, CRPS）推广至右删失和区间删失的生存预测任务。在大规模电子健康记录（EHR）数据集（STARR 和 MIMIC-III）上使用深度神经网络进行训练，采用 Survival-CRPS 的模型相较于最大似然估计（MLE）训练的模型，不仅显著提升了预测分布的锐度，同时保持了良好的校准性，尤其在高删失率场景下表现更优。

ABSTRACT

Probabilistic survival predictions from models trained with Maximum Likelihood Estimation (MLE) can have high, and sometimes unacceptably high variance. The field of meteorology, where the paradigm of maximizing sharpness subject to calibration is popular, has addressed this problem by using scoring rules beyond MLE, such as the Continuous Ranked Probability Score (CRPS). In this paper we present the \emph{Survival-CRPS}, a generalization of the CRPS to the survival prediction setting, with right-censored and interval-censored variants. We evaluate our ideas on the mortality prediction task using two different Electronic Health Record (EHR) data sets (STARR and MIMIC-III) covering millions of patients, with suitable deep neural network architectures: a Recurrent Neural Network (RNN) for STARR and a Fully Connected Network (FCN) for MIMIC-III. We compare results between the two scoring rules while keeping the network architecture and data fixed, and show that models trained with Survival-CRPS result in sharper predictive distributions compared to those trained by MLE, while still maintaining calibration.

研究动机与目标

• 为解决在高删失率的真实世界 EHR 数据中，MLE 训练的生存模型存在的高方差和过度自信问题。
• 通过用一种能平衡锐度与校准性的合理评分规则替代 MLE，提升概率生存预测的质量。
• 开发一种可扩展的深度学习框架用于生存分析，输出完整的患者特异性生存曲线，而非仅点估计或排序结果。
• 在包含数百万名患者的现实世界 EHR 数据集中评估 Survival-CRPS 的性能，证明其在锐度与校准性方面均优于 MLE。
• 提出 Survival-AUPRC 作为新的评估工具，用于捕捉预测概率在真实事件时间附近的集中程度。

提出的方法

• 提出 Survival-CRPS，将 CRPS 评分规则推广至右删失和区间删失的生存数据，使概率生存模型能够进行合理训练。
• 在深度神经网络中采用对数正态参数化（STARR 使用 RNN，MIMIC-III 使用全连接网络），以随时间建模生存分布。
• 采用合理的评分规则目标，同时惩罚校准偏差与过度方差，偏好集中但校准良好的预测分布。
• 在训练过程中应用 Survival-CRPS 目标以优化模型参数，替代传统的基于 MLE 的似然最大化方法。
• 提出一种新颖的评估指标 Survival-AUPRC，用于量化预测生存曲线在真实事件时间附近的集中程度。
• 使用校准图与 AUPRC 评估模型性能，对分位数评估中的删失观测值进行特殊处理。

实验结果

研究问题

• RQ1 在高删失率的 EHR 数据中，CRPS 等合理评分规则是否能在不牺牲校准性的情况下提升生存预测的锐度？
• RQ2 在真实世界的生存预测任务中，Survival-CRPS 与 MLE 相比，在预测分布质量方面表现如何？
• RQ3 使用 Survival-CRPS 是否能生成更准确且更具临床决策实用性的患者特异性生存曲线？
• RQ4 所提出的方法在多大程度上减少了 MLE 训练模型在删失数据中常见的过度自信与高方差问题？
• RQ5 Survival-AUPRC 指标是否能有效捕捉预测生存概率在真实事件时间附近的集中程度？

主要发现

• 使用 Survival-CRPS 训练的模型相较于 MLE 训练的模型，显著提升了预测分布的锐度，在 STARR 数据集中，分配给不合理区域（如 >120 岁）的概率质量减少了 75%。
• 在 STARR 数据集中，采用区间删失的 Survival-CRPS 获得 0.348 ± 4e-3 的 Survival-AUPRC，优于 MLE（0.329 ± 2e-3）及其他基线方法。
• 在 MIMIC-III 数据集中，Survival-CRPS 在 6 个月、1 年和 5 年的预测时间点均保持了高水平的校准性，对存活患者的 AUPRC 值接近 1.0。
• 该方法显著降低了在高删失率设置下的过度自信问题，MLE 模型因删失样本占主导，将超过 75% 的预测质量分配给不切实际的长期生存。
• Survival-CRPS 训练的模型在所有分位数上均保持良好校准性，校准图显示预测事件频率与实际事件频率高度一致。
• 所提出的 Survival-AUPRC 指标有效捕捉了分布的锐度，数值越高表示预测概率在真实事件时间附近的集中程度越好。