[论文解读] BigSurvSGD: Big Survival Data Analysis via Stochastic Gradient Descent

Aliasghar Tarkhan, Noah Simon|arXiv (Cornell University)|Feb 28, 2020

Stochastic Gradient Optimization Techniques参考文献 42被引用 5

一句话总结

该论文提出 BigSurvSGD，一种用于在大规模生存数据上拟合 Cox 比例风险模型的新型随机梯度下降框架。通过将部分似然函数重新表述为解耦目标，该方法实现了高效、稳定且可扩展的优化——尤其适用于大规模或流式数据集，同时保持与标准 Cox 回归的统计等价性，并支持复杂模型（如神经网络）。

ABSTRACT

In many biomedical applications, outcome is measured as a ``time-to-event'' (eg. disease progression or death). To assess the connection between features of a patient and this outcome, it is common to assume a proportional hazards model, and fit a proportional hazards regression (or Cox regression). To fit this model, a log-concave objective function known as the ``partial likelihood'' is maximized. For moderate-sized datasets, an efficient Newton-Raphson algorithm that leverages the structure of the objective can be employed. However, in large datasets this approach has two issues: 1) The computational tricks that leverage structure can also lead to computational instability; 2) The objective does not naturally decouple: Thus, if the dataset does not fit in memory, the model can be very computationally expensive to fit. This additionally means that the objective is not directly amenable to stochastic gradient-based optimization methods. To overcome these issues, we propose a simple, new framing of proportional hazards regression: This results in an objective function that is amenable to stochastic gradient descent. We show that this simple modification allows us to efficiently fit survival models with very large datasets. This also facilitates training complex, eg. neural-network-based, models with survival data.

研究动机与目标

• 解决传统 Cox 比例风险回归在大规模生物医学数据集上存在的计算不稳定性与内存限制问题。
• 克服标准部分似然函数因结构未解耦而与随机梯度优化不兼容的问题。
• 实现使用随机优化在生存时间数据上高效训练复杂模型（包括神经网络）。
• 提供一种可扩展至无法完全装入内存的数据集的框架，同时保持与标准 Cox 回归的统计等价性。
• 支持实时或大规模生存分析应用中的流式与非流式实现。

提出的方法

• 将 Cox 部分似然函数重新表述为一种在观测子集上解耦的目标函数，以支持随机梯度下降。
• 基于 U-统计量的公式推导出一个一致且可微分的目标函数，适用于小批量优化。
• 通过自适应随机梯度方法（如 AMSGrad）实现优化，并采用学习率调度 $ \gamma_m = C / \sqrt{m} $。
• 引入基于分层的采样方法，采样大小为 $ S $，以对个体进行分组，提升计算效率并降低方差。
• 提出两种渐近有效的置信区间估计方法：插值法与非参数自助法。
• 通过迭代值平均（AveAMSGrad）提升实际应用中的收敛性与稳定性。

实验结果

研究问题

• RQ1 能否设计一种改进的 Cox 比例风险回归目标函数，使其适用于随机梯度下降？
• RQ2 在模型假设下，所提方法是否保持与标准 Cox 回归的统计等价性？
• RQ3 该框架能否在全数据优化不可行的大数据与流式场景中实现高效扩展？
• RQ4 BigSurvSGD 在 concordance index 和系数估计方面与标准 coxph() 方法相比表现如何？
• RQ5 该方法能否通过随机优化支持神经网络等复杂模型在生存分析中的应用？

主要发现

• 在 FLCHAIN 数据集中，BigSurvSGD 在所有协变量上的风险比估计值与 coxph() 几乎完全一致，95% 置信区间的绝对宽度差异小于 0.01。
• 当分层大小 $ S=2 $ 时，BigSurvSGD 的 concordance index 达到 0.794，略优于 coxph()（0.792），并与最先进方法相当。
• 在 VETERAN 和 GBSG 数据集中，BigSurvSGD 的 concordance index 与 coxph() 相当，当使用更小的分层大小时，性能提升最高达 0.006。
• 插值法与自助法在标准误估计方面结果非常接近，所有协变量的 95% 置信区间差异均小于 0.01。
• 随着分层大小减小，该方法表现出良好的可扩展性，表明更小的 $ S $ 通过更准确捕捉生存时间数据中的成对 concordance 关系，提升了 concordance 指标。
• BigSurvSGD 通过使目标函数适用于随机优化，成功实现了基于神经网络的生存模型训练，而这是标准 Cox 回归所不具备的能力。