[논문 리뷰] Density-Aware Personalized Training for Risk Prediction in Imbalanced Medical Data

Zepeng Huo, Xiaoning Qian외 3명|arXiv (Cornell University)|2022. 07. 23.

Artificial Intelligence in Healthcare인용 수 4

한 줄 요약

이 논문은 표현 학습과 분류를 분리하고 밀도 인식 손실 및 학습 가능한 비용 행렬을 도입하여, 불균형한 의료 데이터에서 위험 예측을 위한 밀도 인식 개인화 학습 프레임워크를 제안한다. 이 방법은 TOPCAT 및 MIMIC-III 데이터셋에서 기준 모델 대비 우수한 校정성( calibration )을 달성하면서 AUC-ROC, AUC-PRC 및 Brier Skill Score를 향상시킨다.

ABSTRACT

Medical events of interest, such as mortality, often happen at a low rate in electronic medical records, as most admitted patients survive. Training models with this imbalance rate (class density discrepancy) may lead to suboptimal prediction. Traditionally this problem is addressed through ad-hoc methods such as resampling or reweighting but performance in many cases is still limited. We propose a framework for training models for this imbalance issue: 1) we first decouple the feature extraction and classification process, adjusting training batches separately for each component to mitigate bias caused by class density discrepancy; 2) we train the network with both a density-aware loss and a learnable cost matrix for misclassifications. We demonstrate our model's improved performance in real-world medical datasets (TOPCAT and MIMIC-III) to show improved AUC-ROC, AUC-PRC, Brier Skill Score compared with the baselines in the domain.

연구 동기 및 목표

• 희귀한 부정적 사건으로 인해 편향된 모델 성능을 초래하는 의료 위험 예측에서의 클래스 불균형 문제를 해결하기 위해.
• 데이터 밀도의 내재적 차이를 제거하는 대신 이를 활용하여 모델의 校정성( calibration )과 분류 능력을 향상시키기 위해.
• 고위험 환자 내의 클래스 내 이질성을 고려한 개인화된 학습 제도를 개발하기 위해.
• 수동 조정이 필요한 히وري스틱 리샘플링 또는 재가중 기법에 대한 의존도를 줄이기 위해.
• 실제 의료 데이터셋인 TOPCAT 및 MIMIC-III와 같은 환경에서 예측 성능과 확률적 校정성( calibration )을 동시에 향상시키기 위해.

제안 방법

• 클래스 밀도 차이로 인한 편향을 줄이기 위해 특징 추출과 분류를 분리한다.
• 저위험 환자(밀도 높음)와 고위험 환자(산산이 흩어져 있음)의 특징 공간 내 기하학적 분포를 모델링하는 밀도 인식 손실 함수를 도입한다.
• 클래스별 위험 프로파일에 기반해 오분류에 대한 개인화된 벌점(penalty)을 제공하기 위해 학습 가능한 비용 행렬을 사용한다.
• 교차 엔트로피 손실에 밀도 인식 손실 및 학습 가능한 비용 행렬을 결합하여 모델을 학습시킨다.
• 강력한 기준 모델과 함께 랜덤화된 임상 시험(TOPCAT)과 관찰 기반 전자 의무기록 데이터셋(MIMIC-III)에 이 프레임워크를 적용한다.
• 학습 가능한 비용 행렬의 파라미터 θ에 대한 추론 실험과 하이퍼파라미터 튜닝을 통해 성능을 최적화한다.

실험 결과

연구 질문

• RQ1 표현 학습을 분류와 분리하는 것이 불균형한 의료 데이터에서의 편향을 완화하는 데 효과적인가?
• RQ2 클래스 내 밀도 차이(밀도 높은 저위험 대비 산산이 흩어진 고위험)를 모델링하면 위험 예측에 어떤 영향을 미치는가?
• RQ3 학습 가능한 비용 행렬은 외부 校정성( calibration ) 데이터가 없이도 모델의 校정성( calibration )과 성능을 향상시킬 수 있는가?
• RQ4 제안된 프레임워크는 AUC-ROC와 Brier Skill Score 모두에서 표준 리샘플링 및 재가중 기법을 능가하는가?
• RQ5 이 프레임워크는 일반적으로 과신이 발생하는 고위험 확률 영역에서 얼마나 높은 수준의 校정성( calibration )을 향상시키는가?

주요 결과

• 전체 프레임워크는 TOPCAT 및 MIMIC-III 데이터셋에서 모든 기준 모델을 능가하는 최고의 AUC-ROC 및 AUC-PRC 성능을 달성했다.
• 분리 학습만으로도 AUC-PRC는 유의미하게 향상되었지만, 낮은 Brier Skill Score(BSS) 값으로 인해 校정성( calibration )은 열악한 편이었다.
• 밀도 인식 손실만으로도 校정성( calibration )이 향상되어 양의 BSS를 달성하여 더 나은 확률적 신뢰성( reliability )을 보였다.
• 분리 학습과 밀도 인식 손실, 학습 가능한 비용 행렬을 조합한 결과 최고의 BSS와 校정성( calibration )을 달성했으며, 校정성( calibration ) 플롯에서 예측 확률이 대각선에 더 가까워졌다.
• 하이퍼파라미터 튜닝 결과, θ=5가 AUC-ROC를 최대화했고, θ=10이 AUC-PRC에 최적임을 확인했으며, 이들 간 통계적 유의미한 차이는 없었다.
• 추가 校정성( calibration ) 데이터 없이도 뛰어난 校정성( calibration )을 달성하여 고위험 확률 영역에서의 강건성을 입증했다.