[論文レビュー] 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

ひとこと要約

本稿では、右側打ち切りおよび区間打ち切りデータを伴う生存予測に一般化された連続ランク確率スコア（CRPS）の新しいスコアリングルールであるSurvival-CRPSを提案する。大規模なEHRデータセット（STARRおよびMIMIC-III）上で深層ニューラルネットワークを用いて訓練したモデルは、特に高打切り率の状況下でも、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.

研究の動機と目的

• MLEで訓練された生存モデルにおける高い分散と過信の問題を、特に高打切り率の現実世界のEHRデータにおいて解決すること。
• 鋭さとキャリブレーションのバランスを取る適切なスコアリングルールにMLEを置き換えることで、確率的生存予測の質を向上させること。
• ポイント推定やランク付けのみを出力するのではなく、患者ごとの完全な生存曲線を出力するスケーラブルな深層学習フレームワークを構築すること。
• 数百万件の患者を含む現実世界のEHRデータセットにおいて、Survival-CRPSの性能を評価し、予測の鋭さとキャリブレーションの両面でMLEを上回ることを示すこと。
• 真のイベント発生時刻の周囲における予測確率の集中度を捉える新しい評価指標「Survival-AUPRC」を導入すること。

提案手法

• 右側打ち切りおよび区間打ち切り生存データに対するCRPSスコアリングルールの一般化であるSurvival-CRPSを提案し、確率的生存モデルの適切な訓練を可能にする。
• 時間経過に伴う生存分布をモデル化するために、深層ニューラルネットワーク（STARRではRNN、MIMIC-IIIではFCN）で対数正規分布をパラメータ化する。
• 不適切なキャリブレーションや過度な分散をペナルティとして課す適切なスコアリングルールの目的関数を採用し、集中的かつキャリブレートされた予測分布を好む。
• 標準的なMLEベースの尤度最大化に代わり、訓練中にSurvival-CRPS目的関数を適用してモデルパラメータを最適化する。
• 真のイベント発生時刻の周囲における予測生存曲線の集中度を定量化する新しい評価指標「Survival-AUPRC」を導入する。
• キャリブレーションプロットとAUPRCを用いてモデルの性能を評価し、分位数ベースの評価において打ち切り観測を特別に処理する。

実験結果

リサーチクエスチョン

• RQ1 CRPSのような適切なスコアリングルールは、高打切り率のEHRデータにおいて、キャリブレーションを損なうことなく生存予測の鋭さを向上させることができるか？
• RQ2 現実世界の生存予測タスクにおいて、Survival-CRPSはMLEに比べて予測分布の質がどのように向上するか？
• RQ3 Survival-CRPSの使用は、臨床意思決定の文脈において、より正確で実用的な患者ごとの生存曲線をもたらすか？
• RQ4 提案手法は、打ち切りデータにおけるMLEで訓練されたモデルで見られる過信と高い分散をどの程度軽減するか？
• RQ5 Survival-AUPRC指標は、真のイベント発生時刻の周囲における予測生存確率の集中度を効果的に捉えられるか？

主な発見

• Survival-CRPSで訓練されたモデルは、MLEで訓練されたモデルに比べて顕著に鋭い予測分布を達成し、STARRデータセットでは不適切な領域（例：120歳以上）に割り当てられる確率質量が75％削減された。
• STARRデータセットでは、区間打ち切りを伴うSurvival-CRPSのSurvival-AUPRCは0.348 ± 4e-3を達成し、MLE（0.329 ± 2e-3）および他のベースラインを上回った。
• MIMIC-IIIデータセットでは、6か月、1年、5年予測期間においてもSurvival-CRPSは高いキャリブレーションを維持し、生存患者のAUPRC値は1.0に近い水準を示した。
• 高打切り率の設定下で過信を軽減し、MLEモデルが打ち切り例の優位性により、不実在の長期生存確率に75％以上の予測質量を割り当てていたのを是正した。
• Survival-CRPSで訓練されたモデルは、すべての分位数において良好にキャリブレーションされており、予測確率と観測されたイベント頻度の一致度が顕著に高いことをキャリブレーションプロットで確認した。
• 提案されたSurvival-AUPRC指標は、分布の鋭さを効果的に捉えており、高い値は真のイベント発生時刻の周囲における予測確率の集中度が高いかを示している。