[논문 리뷰] Selective prediction-set models with coverage guarantees

Jean Feng, Arjun Sondhi외 2명|arXiv (Cornell University)|2019. 06. 13.

Machine Learning in Healthcare참고 문헌 25인용 수 12

한 줄 요약

이 논문은 신뢰도를 높이기 위해 높은 불확실성 예측을 피하기 위해 예측 집합을 출력하거나 기피하는 선택적 예측집합(SPS) 모델을 훈련하기 위한 페널라이제이션 손실 최소화 프레임워크를 제안한다. K-폴드 교차검증을 사용해 모델을 앙상블함으로써 명목 수준에 가까운 커버리지 비율과 더 좁은 신뢰구간을 달성하며, MNIST 이미지 분류 및 ICU 환자 위험 예측 작업 모두에서 기존 방법을 능가한다.

ABSTRACT

Though black-box predictors are state-of-the-art for many complex tasks, they often fail to properly quantify predictive uncertainty and may provide inappropriate predictions for unfamiliar data. Instead, we can learn more reliable models by letting them either output a prediction set or abstain when the uncertainty is high. We propose training these selective prediction-set models using an uncertainty-aware loss minimization framework, which unifies ideas from decision theory and robust maximum likelihood. Moreover, since black-box methods are not guaranteed to output well-calibrated prediction sets, we show how to calculate point estimates and confidence intervals for the true coverage of any selective prediction-set model, as well as a uniform mixture of K set models obtained from K-fold sample-splitting. When applied to predicting in-hospital mortality and length-of-stay for ICU patients, our model outperforms existing approaches on both in-sample and out-of-sample age groups, and our recalibration method provides accurate inference for prediction set coverage.

연구 동기 및 목표

• 불확실성이 높을 경우 예측을 기피할 수 있는 일반적인 기계학습 모델 훈련 프레임워크를 개발함으로써 신뢰도를 높이고 임상의의 부담을 줄이는 것.
• 특히 고위험 의료 환경에서 기존 기계학습 기반 임상 예측 모델의 이론적 커버리지 보장을 부족하게 하는 문제를 해결하는 것.
• 점추정이 아닌 예측집합을 출력함으로써 임상 의사결정에 더 나은 불확실성 정량화를 제공하는 것.
• K-폴드 교차검증을 활용한 모델에 종속되지 않는 통계적 추론 절차를 개발하여 마진 커버리지 비율을 정밀하게 추정하는 것.
• 실세계 의료 데이터셋에서 방법을 경험적으로 검증하여 정확도 향상과 커버리지 안정성을 입증하는 것.

제안 방법

• 예측 정확도와 기피 행동을 동시에 최적화하는 페널라이제이션 경험 손실 함수를 제안하여, 분포 외부 또는 예측이 어려운 케이스에서 모델이 기피하도록 유도한다.
• 다양한 데이터 하위집단 간 일관된 조건부 커버리지 달성을 위해 균일 수용 페널티를 도입한다.
• K-폴드 교차검증을 사용해 다수의 SPS 모델을 훈련한 후 이를 앙상블하여 마진 커버리지 비율 추정을 향상시킨다.
• 유한 표본 커버리지 보장을 갖는 유효한 예측집합을 구성하기 위해 공명 추론 원리를 적용한다.
• K-폴드 모델의 집계된 예측을 기반으로 앙상블의 마진 커버리지 비율에 대한 통계적 추론을 수행한다.
• MNIST 및 MIMIC-III 데이터셋에서 pSPS 목적함수를 사용해 딥 네ural 네트워크를 훈련하여 성능과 커버리지 성능을 평가한다.

실험 결과

연구 질문

• RQ1 통합된 페널라이제이션 손실 프레임워크는 기계학습 모델이 선택적으로 예측집합을 출력하거나 기피할 수 있도록 훈련시켜 의료 응용 분야의 신뢰도를 향상시킬 수 있는가?
• RQ2 단일 분할 검증에 비해 K-폴드 교차검증 기반 SPS 모델 앙상블은 마진 커버리지 비율 추정의 정확도와 정밀도에 어떤 영향을 미치는가?
• RQ3 pSPS 모델이 분포 외부 또는 복잡한 케이스에서 기피함으로써 예측 정확도가 얼마나 향상되는가, 특히 임상 환경에서의 경우에 대해 어느 정도인가?
• RQ4 제안된 방법은 기존 접근 방식에 비해 명목 수준에 더 가까운 커버리지 비율을 달성하는가, 특히 실세계 ICU 환자 예측 작업에서의 경우에 대해 어떻게 되는가?
• RQ5 모델에 종속되지 않는 추론 절차는 블랙박스 방법을 통해 훈련된 SPS 모델의 마진 커버리지 비율을 신뢰성 있게 추정할 수 있는가?

주요 결과

• pSPS 모델은 입원 사망 예측 작업에서 테스트 손실 0.3768 (0.006)을 기록했으며, 이는 베이스라인 방법보다 우수했다.
• 입원 사망 예측 작업에서 pSPS 모델은 수용 비율 51.44% (2.44)와 마진 커버리지 비율 추정치 0.768 (0.047)를 기록했으며, 명목 수준 0.75에 가까웠다.
• 입원 기간 예측 작업에서 pSPS 모델은 마진 커버리지 비율 추정치 0.759 (0.007)를 기록했고, 95% 신뢰구간 (0.752, 0.766)을 갖는 것으로 나타났으며, 테스트 커버리지 0.756과 매우 유사했다.
• 앙상블 모델의 마진 커버리지에 대한 신뢰구간 너비는 개별 모델의 약 절반으로 줄었으며, 이는 더 정밀한 추정을 의미한다.
• pSPS 모델은 두 작업 모두에서 젊은 환자에게 더 자주 기피 행동을 보여, 기각된 케이스에서 더 낮은 테스트 손실과 높은 정확도를 달성했다.
• MIMIC-III 데이터셋에서 입원 기간 예측을 위한 최종 pSPS 모델은 젊은 환자에게 더 낮은 수용 확률을 할당했으며, 이는 변동성이 더 높아서 발생했고, 보다 향상된 校정을 반영했다.