[논문 리뷰] Maximum Entropy on Erroneous Predictions (MEEP): Improving model calibration for medical image segmentation

Agostina J. Larrazabal, César E. Martínez외 2명|arXiv (Cornell University)|2021. 12. 22.

Explainable Artificial Intelligence (XAI)인용 수 5

한 줄 요약

이 논문은 의료 영상 분할에서 모델 캘리브레이션을 향상시키기 위해 잘못 분류된 픽셀에서 낮은 엔트로피(과신뢰도)를 방지하는 전략인 오류 예측에 대한 최대 엔트로피(MEEP)를 제안한다. 오류 예측에 대해 엔트로피 또는 KL 정규화를 통해 엔트로피를 최대화함으로써 MEEP는 모델 복잡도를 증가시키지 않고도 불확실성 인식 능력을 향상시켜 다양한 데이터셋과 아키텍처에서 유의미하게 향상된 캘리브레이션과 분할 성능을 달성한다.

ABSTRACT

Modern deep neural networks achieved remarkable progress in medical image segmentation tasks. However, it has recently been observed that they tend to produce overconfident estimates, even in situations of high uncertainty, leading to poorly calibrated and unreliable models. In this work we introduce Maximum Entropy on Erroneous Predictions (MEEP), a training strategy for segmentation networks which selectively penalizes overconfident predictions, focusing only on misclassified pixels. Our method is agnostic to the neural architecture, does not increase model complexity and can be coupled with multiple segmentation loss functions. We benchmark the proposed strategy in two challenging segmentation tasks: white matter hyperintensity lesions in magnetic resonance images (MRI) of the brain, and atrial segmentation in cardiac MRI. The experimental results demonstrate that coupling MEEP with standard segmentation losses leads to improvements not only in terms of model calibration, but also in segmentation quality.

연구 동기 및 목표

• 딥 러닝 기반 의료 영상 분할에서 과신뢰도가 높고 캘리브레이션 수준이 낮은 예측이라는 중요한 문제를 해결하기 위해.
• 특히 높은 불확실성과 잘못 분류된 영역에서의 불확실성 추정을 향상시켜 모델 신뢰도를 높이기 위해.
• 모델 아키텍처에 종속되지 않고 모델 복잡도를 증가시키지 않으며, 기존의 표준 분할 손실과도 조합 가능한 훈련 전략을 개발하기 위해.
• 임상 AI 모델 평가를 위해 분할 지표 외에도 캘리브레이션 지표가 필수적임을 입증하기 위해.
• 전역적으로 정규화를 적용하는 것보다 오류 예측에만 정규화를 집중하는 것이 더 나은 캘리브레이션을 이끌어낼 수 있음을 보여주기 위해.

제안 방법

• 잘못 분류된 픽셀에서 직접 엔트로피를 최대화하는 정규화 항과, 잘못된 예측에서 높은 불확실성을 유도하기 위한 프록시로 KL 발산 항을 제안한다.
• 훈련 중에 잘못 분류된 픽셀에만 정규화를 적용하여 어려운, 불확실한 케이스에 집중한다.
• 네트워크 아키텍처를 수정하거나 추론 비용을 증가시키지 않고 주 손실 함수에 정규화를 통합한다.
• 엔트로피 기반 정규화를 통해 잘못 분류된 픽셀의 예측 확률 분포를 균일성(높은 엔트로피)으로 이동시켜 과신뢰도를 감소시킨다.
• 높은 불확실성과 잘못 분류된 픽셀에 집중하기 위해 가중치가 적용된 예측 마스크($\hat{Y}_w$)를 사용한다.
• 기존의 표준 분할 손실(예: Dice, 교차 엔트로피, 포칼 손실)과도 호환되며, 기존 훈련 파이프라인에 쉽게 통합할 수 있다.

실험 결과

연구 질문

• RQ1 잘못 분류된 픽셀에서 낮은 엔트로피를 방지하는 것이 의료 영상 분할에서 모델 캘리브레이션을 향상시킬 수 있는가?
• RQ2 오직 오류 예측에만 불확실성 정규화를 적용하는 것이 전역 정규화보다 더 나은 캘리브레이션을 이끌어내는가?
• RQ3 MEEP는 모델 복잡도를 증가시키지 않고도 다양한 아키텍처와 분할 손실 함수에 적용 가능한가?
• RQ4 MEEP를 통한 캘리브레이션 향상은 동시에 분할 성능 향상으로 이어지는가?
• RQ5 기존의 캘리브레이션 기법과 비교했을 때 MEEP는 신뢰성과 불확실성 추정 측면에서 더 낫게 작용하는가?

주요 결과

• MEEP는 Dice나 교차 엔트로피 손실과 조합했을 때 WMH 데이터셋에서 브리어 스코어를 최대 65% 감소시키고 ECE를 최대 60% 감소시켜 모델 캘리브레이션을 유의미하게 향상시켰다.
• 이 방법은 상태기반 캘리브레이션 성능를 달성하여, 교차 엔트로피 손실과 함께 KL 정규화를 사용했을 때 WMH에서 브리어 스코어(0.434 × 10⁻³)와 ECE(0.174) 모두 최고 성능 기록을 달성했다.
• 교차 엔트로피 손실과 함께 MEEP를 사용했을 때 WMH에서 분할 성능(Dice 계수)이 최대 0.016(1.6%) 향상되었고, LA 데이터셋에서는 0.013(1.4%) 향상되었다.
• 제안된 방법은 백본 아키텍처에 종속되지 않으며, UNet과 ResUNet 아키텍처 모두에서 유사한 향상이 관찰되어 일반화 능력을 확인했다.
• 정성적 결과에서는 MEEP로 훈련된 모델이 더 균형 잡힌 확률 지도를 생성하며, 어려운 분할 영역에서 값이 약 0.5에 수렴함으로써 높은 불확실성 수준을 반영하고 있음을 보여주었다.
• 이 연구는 단지 Dice와 같은 분할 지표에 의존하는 것이 부족하며, 브리어 스코어와 ECE와 같은 캘리브레이션 지표가 임상 신뢰성 평가에 필수적이라는 점을 입증했다.