[논문 리뷰] Countdown Regression: Sharp and Calibrated Survival Predictions

Anand Avati, Tony Duan외 4명|arXiv (Cornell University)|2018. 06. 21.

Machine Learning in Healthcare참고 문헌 33인용 수 18

한 줄 요약

이 논문은 우연히 생존 예측을 위한 새로운 점수 규칙인 Survival-CRPS를 소개한다. 이는 우연히 케이싱된 데이터와 간격 케이싱된 데이터를 포함한 생존 예측에 연속적 순위 확률 점수(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.

연구 동기 및 목표

• 실제 EHR 데이터에서 고비율 케이싱 상황에서 MLE로 훈련된 생존 모델의 높은 분산과 과신 문제를 해결하기 위해.
• 예측 분포의 품질을 향상시키기 위해 MLE를 대체하는 적절한 점수 규칙을 도입하여 날카움과 校정성의 균형을 이루기 위해.
• 포인트 추정치나 순위만 출력하는 것이 아니라 환자별로 완전한 생존 곡선을 출력하는 확장 가능한 딥 러닝 프레임워크를 개발하기 위해.
• 수천만 명의 환자를 포함한 실제 EHR 데이터셋에서 Survival-CRPS의 성능을 평가하여, MLE 대비 날카움과 校정성 측면에서 뛰어난 성능을 입증하기 위해.
• 진짜 사건 발생 시간 주변의 예측 확률 집중 정도를 측정할 수 있는 새로운 평가 도구인 Survival-AUPRC 지표를 도입하기 위해.

제안 방법

• 우연히 케이싱된 데이터와 간격 케이싱된 생존 데이터에 대해 CRPS 점수 규칙을 일반화한 Survival-CRPS를 제안하여, 확률적 생존 모델의 적절한 훈련을 가능하게 한다.
• 시간에 따른 생존 분포를 모델링하기 위해 딥 뉴럴 네트워크(RNN은 STARR, FCN은 MIMIC-III에 사용)에서 로그정규 분포 파arameterization을 사용한다.
• 교정성과 과도한 분산을 모두 방지하는 적절한 점수 규칙 목적함수를 사용하여, 집중적이고 잘 校정된 예측 분포를 선호한다.
• 표준 MLE 기반의 우도 최대화를 대체하여, 훈련 중에 Survival-CRPS 목적함수를 적용하여 모델 파라미터를 최적화한다.
• 예측 생존 곡선이 진짜 사건 발생 시간 주변에 얼마나 집중되어 있는지 측정할 수 있는 새로운 평가 지표인 Survival-AUPRC를 도입한다.
• 교정성 플롯과 AUPRC를 사용하여 모델 성능을 평가하며, 분위수 기반 평가에서 케이싱 관측치에 특별한 처리를 적용한다.

실험 결과

연구 질문

• RQ1 CRPS와 같은 적절한 점수 규칙이 고비율 케이싱 EHR 데이터에서 교정성을 유지하면서도 생존 예측의 날카움을 향상시킬 수 있는가?
• RQ2 실제 생존 예측 작업에서 Survival-CRPS는 MLE 대비 예측 분포 품질 측면에서 어떻게 비교되는가?
• RQ3 Survival-CRPS의 사용이 임상 의사결정 상황에서 더 정확하고 실용적인 환자별 생존 곡선을 도출하는 데 기여하는가?
• RQ4 제안된 방법은 케이싱된 데이터에서 MLE로 훈련된 모델이 보이는 과신과 높은 분산 문제를 어느 정도 줄이는가?
• RQ5 Survival-AUPRC 지표는 진짜 사건 발생 시간 주변의 예측 생존 확률 집중 정도를 효과적으로 측정할 수 있는가?

주요 결과

• Surv命-CRPS로 훈련된 모델은 MLE로 훈련된 모델보다 예측 분포가 훨씬 날카롭게 나타났으며, STARR 데이터셋에서 비합리적인 영역(예: 120세 이상)에 할당된 질량이 75% 감소했다.
• STARR 데이터셋에서 간격 케이싱을 고려한 Survival-CRPS는 Survival-AUPRC 0.348 ± 4e-3을 기록하여 MLE(0.329 ± 2e-3) 및 기타 베이스라인을 초월했다.
• MIMIC-III 데이터셋에서 Survival-CRPS는 6개월, 1년, 5년 예측 수준 전반에서 높은 교정성을 유지했으며, 생존한 환자에 대해 AUPRC 값이 1.0에 가까웠다.
• 고비율 케이싱 환경에서 과신을 감소시켰으며, MLE 모델은 케이싱된 예제가 지배적으로 영향을 미쳐 120세 이상의 비현실적인 장기 생존 확률에 75% 이상의 예측 질량을 할당했다.
• Survival-CRPS로 훈련된 모델은 교정성 플롯을 통해 예측된 사건 빈도와 관측된 빈도가 밀접하게 일치함을 확인하여 모든 분위수에서 잘 교정된 상태를 유지했다.
• 제안된 Survival-AUPRC 지표는 분포의 날카움을 효과적으로 측정했으며, 높은 값일수록 진짜 사건 발생 시간 주변의 예측 확률 집중 정도가 높음을 나타냈다.