[논문 리뷰] Improved Trainable Calibration Method for Neural Networks on Medical Imaging Classification

Gongbo Liang, Yu Zhang외 2명|arXiv (Cornell University)|2020. 09. 09.

COVID-19 diagnosis using AI참고 문헌 36인용 수 29

한 줄 요약

이 논문은 신경망의 캘리브레이션을 향상시키면서도 분류 정확도를 훼손하지 않는 새로운 학습 가능한 校정 방법인 Confidence와 Accuracy의 차이(Difference in Confidence and Accuracy, DCA)를 제안한다. DCA를 보조 손실로 추가함으로써 정확도가 정점에 도달했을 때 교차 엔트로피 손실에 대한 과적합을 방지하여, 네 개의 의료 영상 데이터셋과 네 가지 아키텍처에서 평균적으로 기대 캘리브레이션 오차(Expected Calibration Error, ECE)를 65.72% 감소시켰으며, 동시에 높은 정확도를 유지한다.

ABSTRACT

Recent works have shown that deep neural networks can achieve super-human performance in a wide range of image classification tasks in the medical imaging domain. However, these works have primarily focused on classification accuracy, ignoring the important role of uncertainty quantification. Empirically, neural networks are often miscalibrated and overconfident in their predictions. This miscalibration could be problematic in any automatic decision-making system, but we focus on the medical field in which neural network miscalibration has the potential to lead to significant treatment errors. We propose a novel calibration approach that maintains the overall classification accuracy while significantly improving model calibration. The proposed approach is based on expected calibration error, which is a common metric for quantifying miscalibration. Our approach can be easily integrated into any classification task as an auxiliary loss term, thus not requiring an explicit training round for calibration. We show that our approach reduces calibration error significantly across various architectures and datasets.

연구 동기 및 목표

• 의료 영상 분류에 사용되는 딥 네트워크에서 과도하게 자신감을 갖는 예측을 보이는 모델들이 자주 발생하는 심각한 캘리브레이션 문제를 해결하기 위해.
• 예측된 자신감과 실제 정확도 간의 격차를 줄임으로써 의료 AI 시스템의 불확실성 정량화를 향상시키기 위해.
• 모델 캘리브레이션을 크게 향상시키면서도 높은 분류 정확도를 유지하는 캘리브레이션 방법을 개발하기 위해.
• 추가 학습 라운드나 복잡한 후처리가 필요 없는 단순하고 학습 가능하며 통합 가능한 캘리브레이션 솔루션을 제공하기 위해.

제안 방법

• 교차 엔트로피 손실 감소 시 정확도가 정체되는 경우를 방지하기 위해, 보조 손실 항목인 Confidence와 Accuracy의 차이(Difference in Confidence and Accuracy, DCA)를 도입한다.
• DCA는 예측 확률의 각 밴드 내 평균 자신감과 정확도 간의 절대 차이로 정의되며, M개의 밴드를 사용한 ECE를 통해 근사된다.
• 표준 교차 엔트로피 손실과 함께 DCA 손실을 결합하여, 별도의 캘리브레이션 단계 없이 엔드 투 엔드 최적화를 가능하게 한다.
• DCA 손실의 스케일링을 위해 하이퍼파rameter β를 사용하며, 캘리브레이션 향상과 학습 안정성 간의 균형을 맞춘다.
• 이 방법은 아키텍처에 종속되지 않으며, 최소한의 수정으로도 어떤 신경망 분류기에도 적용 가능하다.

실험 결과

연구 질문

• RQ1 학습 가능한 엔드 투 엔드 캘리브레이션 방법이 의료 영상 분류에서 분류 정확도를 떨어뜨리지 않고 모델의 잘못된 캘리브레이션을 줄일 수 있는가?
• RQ2 DCA 손실이 다양한 의료 영상 데이터셋과 네트워크 아키텍처에서 캘리브레이션 향상에 얼마나 효과적인가?
• RQ3 기존의 캘리브레이션 기법들과 비교해 복구된 진정한 기저 확률 분포가 더 정확한가?
• RQ4 DCA 방법의 성능가 β 값의 선택에 얼마나 민감한가?

주요 결과

• 제안된 DCA 방법은 네 개의 의료 영상 데이터셋과 네 가지 CNN 아키텍처에서 평균적으로 기대 캘리브레이션 오차(Expected Calibration Error, ECE)를 0.1006에서 0.0345로 65.72% 감소시켰다.
• 모델는 높은 분류 정확도를 유지하였으며, 캘리브레이션되지 않은 모델의 경우 83.08%였고, DCA로 캘리브레이션된 모델의 경우 83.58%로 성능 저하 없이 유지되었다.
• t-SNE 시각화 결과, DCA로 학습된 특징는 특히 Kather 5000 데이터셋에서 온도 스케일링보다 더 잘 분류 가능하고 조밀하게 패킹되어 있음을 보여주었다.
• DCA 방법에 의해 복구된 확률 분포는 진정한 기저의 대각선을 잘 따르며, 캘리브레이션되지 않은 모델과 온도 스케일링보다 더 나은 캘리브레이션을 보였다.
• β ≥ 10일 경우 ECE는 β 값에 대해 상대적으로 민감도가 낮았으며, 대부분의 데이터셋에서 10–15 범위에서 최적의 성능을 보였다.