[논문 리뷰] X-CAL: Explicit Calibration for Survival Analysis

Mark Goldstein, Xintian Han외 3명|PubMed|2021. 01. 13.

Machine Learning in Healthcare참고 문헌 17인용 수 13

한 줄 요약

이 논문은 학습 중에 직접 분포 캘리브레이션(d-calibration)을 최적화하는 미분 가능하고 명시적인 캘리브레이션 목표인 X-CAL을 소개한다. 최대우도추정과 결합함으로써, 시뮬레이션 데이터, MNIST 기반 생존 작업, MIMIC-III 및 TCGA를 포함한 다양한 데이터셋에서 예측 정확도나 순서 일致성(concordance)을 크게 떨어뜨리지 않으면서도 모델의 캘리브레이션을 향상시켜 의료 응용 분야에서의 위험 예측 신뢰도를 높인다.

ABSTRACT

Survival analysis models the distribution of time until an event of interest, such as discharge from the hospital or admission to the ICU. When a model's predicted number of events within any time interval is similar to the observed number, it is called <i>well-calibrated</i>. A survival model's calibration can be measured using, for instance, distributional calibration (D-CALIBRATION) [Haider et al., 2020] which computes the squared difference between the observed and predicted number of events within different time intervals. Classically, calibration is addressed in post-training analysis. We develop explicit calibration (X-CAL), which turns D-CALIBRATION into a differentiable objective that can be used in survival modeling alongside maximum likelihood estimation and other objectives. X-CAL allows practitioners to directly optimize calibration and strike a desired balance between predictive power and calibration. In our experiments, we fit a variety of shallow and deep models on simulated data, a survival dataset based on MNIST, on length-of-stay prediction using MIMIC-III data, and on brain cancer data from The Cancer Genome Atlas. We show that the models we study can be miscalibrated. We give experimental evidence on these datasets that X-CAL improves D-CALIBRATION without a large decrease in concordance or likelihood.

연구 동기 및 목표

• 강력한 예측 성능에도 불구하고 학습 중 캘리브레이션 최적화가 부족한 생존 모델의 문제를 해결하기 위해.
• 후행 캘리브레이션에 의존하는 대신 학습 중에 직접 분포 캘리브레이션(d-calibration)을 최적화하는 방법을 개발하기 위해.
• 예측 정확도와 캘리브레이션의 균형을 조정할 수 있는 조정 가능한 하이퍼파rameter를 통해 해석 가능성과 임상 신뢰도를 향상시키기 위해.
• 얕은 모델과 깊은 모델, 실제 데이터 및 시뮬레이션 데이터를 포함한 다양한 생존 모델링 시나리오에서 X-CAL의 효과성을 평가하기 위해.
• 학습 중 명시적 캘리브레이션을 통해 예측 정확도나 우도의 상당한 손실 없이 더 잘 캘리브레이션된 생존 예측을 달성할 수 있음을 입증하기 위해.

제안 방법

• X-CAL은 d-calibration을 미분 가능한 목적 함수로 공식화하여 표준 생존 모델링 프레임워크 내에서 엔드 투 엔드 최적화를 가능하게 한다.
• 메트릭의 제곱 차이를 시간 간격 내 예측된 사건 수와 관측된 사건 수의 차이로 근사화하며, 경험 분포의 미분 가능한 완화를 사용한다.
• 최대우도추정과 함께 정규화 요소로 통합되어 예측 우도와 캘리브레이션을 동시에 최적화할 수 있도록 한다.
• 계산의 스케일러빌리티 확보를 위해 데이터 서브샘플링을 활용한다.
• 파라미터 모델(예: Weibull, Log-Normal)과 비모수 모델(예: Cat, mtlr) 모두 지원하여 광범위한 적용 가능성을 확보한다.
• 하이퍼파rameter λ는 우도와 캘리브레이션 간의 트레이드오���을 제어하여 전문가가 원하는 캘리브레이션-예측 균형을 설정할 수 있도록 한다.

실험 결과

연구 질문

• RQ1 학습 중 명시적 캘리브레이션 최적화가 예측 성능을 떨어뜨리지 않고도 생존 모델의 d-calibration을 크게 향상시킬 수 있는가?
• RQ2 다양한 모델 아키텍처와 데이터셋에서 표준 최대우도 학습 대비 X-CAL의 캘리브레이션 및 순서 일치성(concordance) 성능은 어떻게 비교되는가?
• RQ3 MIMIC-III 및 The Cancer Genome Atlas와 같은 실제 생존 데이터셋에서 X-CAL은 얼마나 많은 오차 캘리브레이션을 줄일 수 있는가?
• RQ4 X-CAL 통합은 특히 체류 기간 예측 및 생존 예측 작업에서 임상 의사결정에 더 신뢰할 수 있는 위험 추정치를 제공하는가?
• RQ5 얕은 모델과 깊은 모델, 복잡한 신경망 아키텍처를 포함한 생존 모델 모두에 X-CAL을 효과적으로 적용할 수 있는가?

주요 결과

• X-CAL은 모든 데이터셋에서 d-calibration 오차를 크게 감소시켰다: MNIST 기반 생존 작업에서 d-calibration은 기준값 0.018에서 λ=1000일 때 0.002로 감소했으며, 순서 일치성은 0.80 이상 유지되었다.
• MIMIC-III 체류 기간 예측에서 X-CAL을 적용한 Cat-i 모델은 λ=1000일 때 d-calibration이 0.001로 나타났고, X-CAL 없이 학습한 경우 0.002였으며, 순서 일치성은 0.748(λ=1000)에서 0.758(λ=0)로 약간 감소했다.
• TCGA 간질종 데이터에서 X-CAL은 Weibull 모델의 d-calibration을 0.035에서 Cat-i 모델(λ=1000)로 0.003으로 감소시켰고, 대부분의 모델에서 순서 일치성이 0.80 이상 유지되었다.
• TCGA 데이터에서 X-CAL을 적용한 Log-Normal 모델은 d-calibration을 70% 감소시켜 0.059에서 0.004로 줄였고, 순서 일치성은 약 0.80 수준에서 안정을 유지했다.
• X-CAL은 파라미터 모델, 비모수 모델, 딥 모델 모두에서 캘리브레이션 성능을 향상시켜 강건성과 일반화 능력을 입증했다.
• 메서드를 통해 제어 가능한 트레이드오프가 가능하다: λ를 증가시킬수록 캘리브레이션 성능은 향상되지만 우도는 약간 떨어지며, 순서 일치성은 대부분 안정된 편이다.