[论文解读] Deep Cox Mixtures for Survival Regression

Chirag Nagpal, Steve Yadlowsky+2 more|arXiv (Cornell University)|Jan 16, 2021

Statistical Methods and Inference参考文献 9被引用 5

一句话总结

本文提出 Deep Cox Mixtures，一种灵活的生存回归模型，通过将患者建模为具有各自比例风险结构的潜在子群组，利用深度神经网络和非参数基线风险函数进行估计，从而提升校准性和判别性能。该方法在少数族裔群体中显著改善了校准性能，优于经典方法和深度学习基线模型。

ABSTRACT

Survival analysis is a challenging variation of regression modeling because of the presence of censoring, where the outcome measurement is only partially known, due to, for example, loss to follow up. Such problems come up frequently in medical applications, making survival analysis a key endeavor in biostatistics and machine learning for healthcare, with Cox regression models being amongst the most commonly employed models. We describe a new approach for survival analysis regression models, based on learning mixtures of Cox regressions to model individual survival distributions. We propose an approximation to the Expectation Maximization algorithm for this model that does hard assignments to mixture groups to make optimization efficient. In each group assignment, we fit the hazard ratios within each group using deep neural networks, and the baseline hazard for each mixture component non-parametrically. We perform experiments on multiple real world datasets, and look at the mortality rates of patients across ethnicity and gender. We emphasize the importance of calibration in healthcare settings and demonstrate that our approach outperforms classical and modern survival analysis baselines, both in terms of discriminative performance and calibration, with large gains in performance on the minority demographics.

研究动机与目标

• 解决医疗环境中生存模型校准性差的问题，特别是在比例风险假设不成立时。
• 提升代表性不足的种族和性别少数群体等人群的预测性能，特别是校准性。
• 开发一种灵活的生存模型，通过学习潜在患者子群组，超越比例风险假设的限制。
• 证明提升校准性不会以牺牲判别性能为代价，判别性能通过 concordance index（C-index）衡量。
• 提供一个实用的、开源的深度学习框架，适用于真实世界临床数据集的生存分析。

提出的方法

• 使用多个 Cox 回归组件的混合模型来建模生存结果，其中每个组件代表一个具有独立风险函数的潜在患者子群组。
• 利用深度神经网络基于患者特征估计各组特定的风险比，以实现灵活的非线性关系。
• 在每组内使用基于样条的平滑方法非参数化估计基线生存函数，避免参数分布假设。
• 采用近似蒙特卡洛期望最大化（EM）算法，并对混合组件实施硬分配，以提高优化效率。
• 通过 Kaplan-Meier 估计近似 E-step，通过随机梯度下降优化各组特定参数实现 M-step。
• 在模型评估中整合不确定性与校准度量，尤其关注按种族和性别划分的子群组性能。

实验结果

研究问题

• RQ1 与标准 Cox 模型和深度学习基线相比，基于深度神经网络风险比的 Cox 模型混合是否能改善生存预测的校准性？
• RQ2 所提出的模型是否能减少代表性不足的少数族裔群体（如非裔美国人和女性患者）在生存预测中的性能差异？
• RQ3 通过潜在子群组放松比例风险假设，在多大程度上能同时提升判别性能和校准性？
• RQ4 在真实世界数据集（如 SUPPORT）中，模型估计的子群组特定生存曲线与经验 Kaplan-Meier 估计相比如何？
• RQ5 模型可解释性和子群组发现对可操作的临床决策支持有何影响？

主要发现

• Deep Cox Mixtures 在所有人群中的校准性均显著提升，其中在少数族裔群体（如非裔美国人和女性患者）中性能提升最为显著。
• 该模型保持或超过最先进基线模型的 concordance index（C-index），表明在判别性能与校准性之间不存在权衡。
• 模型估计的子群组特定生存曲线存在交叉，表明存在非比例风险，而该模型通过潜在聚类成功捕捉了这一现象。
• 子群组的 Kaplan-Meier 估计与模型预测的生存曲线高度一致，验证了模型恢复有意义患者子群组的能力。
• 该模型减少了对非裔美国人群体心血管风险的过度估计，与现有工具（如 2013 ACC/AHA Pooled Cohort Equations）中已知的偏差一致。
• GitHub 上的开源实现支持可复现性，并为未来医疗和研究应用中的临床生存分析提供了扩展基础。