[논문 리뷰] BEDS-Bench: Behavior of EHR-models under Distributional Shift--A Benchmark

Anand Avati, Martin Seneviratne외 4명|arXiv (Cornell University)|2021. 07. 17.

Machine Learning in Healthcare인용 수 5

한 줄 요약

BEDS-Bench는 분포 이탈 상황에서 EHR 기반 기계학습 모델의 행동을 평가하기 위한 벤치마크로, 두 개의 공개 EHR 데이터셋을 사용하여 다양한 외부 분포(OutOfDistribution, OOD) 설정을 구축한다. 연구 결과, 평가된 모든 모델이 OOD 조건 하에서 일반화 및 校정 능력이 열악함을 드러내며, 실제 임상 환경에 배포될 때 더 높은 견고성 확보가 절실한 필요성을 시사한다.

ABSTRACT

Machine learning has recently demonstrated impressive progress in predictive accuracy across a wide array of tasks. Most ML approaches focus on generalization performance on unseen data that are similar to the training data (In-Distribution, or IND). However, real world applications and deployments of ML rarely enjoy the comfort of encountering examples that are always IND. In such situations, most ML models commonly display erratic behavior on Out-of-Distribution (OOD) examples, such as assigning high confidence to wrong predictions, or vice-versa. Implications of such unusual model behavior are further exacerbated in the healthcare setting, where patient health can potentially be put at risk. It is crucial to study the behavior and robustness properties of models under distributional shift, understand common failure modes, and take mitigation steps before the model is deployed. Having a benchmark that shines light upon these aspects of a model is a first and necessary step in addressing the issue. Recent work and interest in increasing model robustness in OOD settings have focused more on image modality, while the Electronic Health Record (EHR) modality is still largely under-explored. We aim to bridge this gap by releasing BEDS-Bench, a benchmark for quantifying the behavior of ML models over EHR data under OOD settings. We use two open access, de-identified EHR datasets to construct several OOD data settings to run tests on, and measure relevant metrics that characterize crucial aspects of a model's OOD behavior. We evaluate several learning algorithms under BEDS-Bench and find that all of them show poor generalization performance under distributional shift in general. Our results highlight the need and the potential to improve robustness of EHR models under distributional shift, and BEDS-Bench provides one way to measure progress towards that goal.

연구 동기 및 목표

• 분포 이탈 상황에서 EHR 모델 성능을 평가하기 위한 벤치마크 부족 문제를 해결하기 위해, 특히 실제 임상 환경에서의 배포 시나리오를 고려한다.
• 학습 데이터 분포와 크게 다를 경우 EHR 모델이 어떻게 행동하는지 수량화한다.
• OOD 조건 하에서 모델의 校정 및 신뢰도 추정에서 나타나는 일반적인 실패 유형을 규명한다.
• 다양한 OOD 설정에서 EHR 모델의 견고성 측정 및 비교를 위한 표준화된 평가 프레임워크를 제공한다.
• 임상 AI 응용 분야에서 모델의 신뢰성과 안정성 향상을 위한 연구를 촉진한다.

제안 방법

• 벤치마크는 환자 인구통계학적 특성, 임상 상태 및 시간 창을 조작하여 두 개의 익명화된 공개 EHR 데이터셋을 사용해 다수의 OOD 데이터 설정을 구성한다.
• 시간-사건 정의를 일관되게 적용하여 질병 발병 예측 및 사망 위험 추정과 같은 다수의 임상 예측 과제에서 모델을 평가한다.
• OOD 탐지 평가는 OOD-AUC, ECE(기대 校정 오차), 그리고 신뢰도 캘리브레이션을 포함한 지표를 사용하며, 단일 예측 모델의 경우 예측 엔트로피를 통해 신뢰도를 측정한다.
• 벤치마크는 내부 분포(IND) 및 OOD 데이터 분할을 포함하며, OOD 데이터는 다른 환자 집단, 시간대, 또는 의료 시스템 기관에서 유래한다.
• 모델은 OOD 예시에 대해 낮은 신뢰도(높은 불확실성)를 부여할 수 있는 능력, 특히 캘리브레이션 및 분류 성능에 중점을 두고 평가된다.
• 재현 가능한 평가를 지원하며, 코드는 https://github.com/Google-Health/records-research/tree/master/beds-bench 에 공개되어 있다.

실험 결과

연구 질문

• RQ1 학습 데이터 분포와 다소 다를 경우 EHR 기반 기계학습 모델은 어떻게 성능을 내는가?
• RQ2 실제 EHR 환경의 OOD 조건에서 모델은 OOD 예시에 대해 얼마나 잘 캘리브레이션되며, 신뢰도 추정을 유지하는가?
• RQ3 어떤 모델 아키텍처나 학습 전략이 EHR 데이터에서 분포 이탈에 대해 더 뛰어난 견고성을 보이는가?
• RQ4 EHR 환경에서 클래스 불균형 및 OOD 탐지 시나리오에서 일반적인 신뢰도 추정 방법(예: 엔트로피)은 어떻게 작동하는가?
• RQ5 기존의 OOD 탐지 지표인 OOD-AUC 및 ECE는 임상 테이블 데이터에서 분포 이탈 상황에서 모델 행동을 신뢰성 있게 반영할 수 있는가?

주요 결과

• 모든 평가된 EHR 모델, OOD에 대한 견고성을 고려해 설계된 모델까지도 여러 OOD 설정에서 분포 이탈 상황에서 일반화 성능이 열악함을 보였다.
• 특정 모델 아키텍처가 모든 OOD 시나리오에서 뚜렷이 우월한 성능을 보이지 않아, EHR 모델링에서의 견고성에 대한 명백한 해결책이 없음을 시사한다.
• 모델들은 OOD 예측에 대해 높은 신뢰도를 부여하는 경향이 있으며, 이는 잘못된 예측임에도 불구하고 신뢰도 캘리브레이션이 열악함을 시사한다.
• OOD-AUC 지표만으로는 충분하지 않으며, 클래스 불균형 상황에서는 오해의 소지가 있다. 만약 OOD 예측이 잘못 분류되어 양성으로 간주될 경우, 높은 신뢰도를 부여하는 모델을 보상하는 결과를 낳을 수 있다.
• 기대 校정 오차(ECE) 및 신뢰도 캘리브레이션은 여전히 모델 신뢰성의 강력한 지표로 남아 있으나, 테스트된 모든 모델이 OOD 조건 하에서 심각한 캘리브레이션 오류를 보였다.
• 본 연구는 현재의 EHR 모델이 OOD로 식별되기 어렵다는 점을 확인하였으며, 분포 이탈이 발생하는 실제 임상 환경에 배포될 경우 임상적 위험을 초래할 수 있음을 시사한다.