[논문 리뷰] Improving Calibration and Out-of-Distribution Detection in Medical Image Segmentation with Convolutional Neural Networks

Davood Karimi, Ali Gholipour|arXiv (Cornell University)|2020. 04. 12.

Radiomics and Machine Learning in Medical Imaging인용 수 5

한 줄 요약

이 논문은 다양한 종류의 의료 영상 데이터를 활용한 다중 작업 학습을 제안하여 CNN 기반 분할 모델의 校정성과 분포 외(OOD) 탐지 성능을 향상시킨다. 서로 다른 기관과 영상 모odalities에서 학습한 단일 모델을 통해 보다 보정된 예측 성능을 달성하고, CNN 특징 맵에 대한 스펙트럼 분석 기법을 도입하여 기존의 불확실성 기반 OOD 탐지 방법을 능가하는 성능을 보이며, 핵심 실험에서 100% 정확도를 달성한다.

ABSTRACT

Convolutional Neural Networks (CNNs) have shown to be powerful medical image segmentation models. In this study, we address some of the main unresolved issues regarding these models. Specifically, training of these models on small medical image datasets is still challenging, with many studies promoting techniques such as transfer learning. Moreover, these models are infamous for producing over-confident predictions and for failing silently when presented with out-of-distribution (OOD) data at test time. In this paper, we advocate for multi-task learning, i.e., training a single model on several different datasets, spanning several different organs of interest and different imaging modalities. We show that not only a single CNN learns to automatically recognize the context and accurately segment the organ of interest in each context, but also that such a joint model often has more accurate and better-calibrated predictions than dedicated models trained separately on each dataset. Our experiments show that multi-task learning can outperform transfer learning in medical image segmentation tasks. For detecting OOD data, we propose a method based on spectral analysis of CNN feature maps. We show that different datasets, representing different imaging modalities and/or different organs of interest, have distinct spectral signatures, which can be used to identify whether or not a test image is similar to the images used to train a model. We show that this approach is far more accurate than OOD detection based on prediction uncertainty. The methods proposed in this paper contribute significantly to improving the accuracy and reliability of CNN-based medical image segmentation models.

연구 동기 및 목표

• 작은 데이터셋으로 학습된 CNN 기반 의료 영상 분할 모델에서 낮은 보정성과 신뢰할 수 없는 예측 성능을 해결한다.
• 테스트 시에 모델이 침묵적으로 실패하는 경우가 많기 때문에, 분포 외 데이터를 탐지함으로써 모델의 신뢰성을 향상시킨다.
• 다양한 기관과 영상 모달리티를 포함한 이질적 데이터셋을 함께 학습함으로써 전반적인 일반화 및 보정성 향상 여부를 탐색한다.
• 기존의 불확실성 기반 접근 방식보다 더 정확한 새로운 OOD 탐지 방법을 의료 영상 분할 분야에 개발한다.

제안 방법

• 다양한 기관과 영상 모달리티를 포함하는 이질적 의료 영상 데이터셋의 혼합물에 대해 단일 CNN 모델을 학습하며, 모달리티나 기관에 대한 추가 입력 없이 수행한다.
• 최종 합성곱층의 특징 맵에 대한 스펙트럼 분석을 통해 OOD 스코어(OODM)를 계산한다. 이는 활성화의 파wer 스펙트럼에 기반한다.
• 보정을 위해 분포 내 데이터를 사용하여 OODM에 임계값 τ를 정의하고, 테스트 이미지를 분포 내 또는 분포 외로 분류한다.
• AUC, 정확도, 민감도, 특이도 등의 지표를 사용하여 제안된 OODM 방법과 예측 불확실성 기반 탐지 방법을 비교한다.
• 예를 들어 간의 MRI, 태반 MRI에서의 코르티컬 플레이트 등 여러 데이터셋을 함께 학습하여 일반화 및 보정 성능를 평가한다.
• 분할 정확도와 신뢰도 보정성을 평가하기 위해 DSC(Dice Similarity Coefficient)와 ECE 등 보정 지표를 사용한다.

실험 결과

연구 질문

• RQ1 다양한 기관과 모달리티를 포함한 이질적 의료 영상 데이터로 학습한 단일 CNN 모델이 개별 데이터셋으로 학습된 전용 모델과 비교해 유사하거나 더 높은 분할 성능를 달성할 수 있는가?
• RQ2 이질적 데이터에서의 다중 작업 학습이 개별 데이터셋으로 학습한 경우보다 더 보정된 예측 성능를 제공하는가?
• RQ3 CNN 특징 맵의 스펙트럼 분석이 의료 영상 분할에서 분포 외 이미지를 효과적으로 탐지할 수 있는가?
• RQ4 스펙트럼 분석 기반으로 제안된 OOD 탐지 방법은 기존의 불확실성 기반 방법에 비해 정확도에서 얼마나 뛰어나게 성능을 내는가?
• RQ5 모델이 모달리티나 기관 정보를 명시적으로 제공받지 않더라도 다양한 데이터 분포(예: 다양한 연령대, 영상 프rotocol)에 일반화할 수 있는가?

주요 결과

• 다양한 기관과 모달리티를 포함한 이질적 데이터셋(예: MRI에서의 간, 태아 MRI에서의 코르티컬 플레이트)으로 학습한 단일 CNN 모델이 개별 데이터셋으로 학습된 전용 모델과 유사하거나 더 높은 분할 성능를 달성했다.
• 이질적 데이터로 학습한 모델는 더 낮은 기대 보정 오차(Expected Calibration Error, ECE)로 측정되어 보다 보정된 예측 성능를 보였다.
• CNN 특징 맵의 스펙트럼 분석 기반으로 제안된 OOD 탐지 방법은 간 분할 작업에서 Liver-MRI-SPIR 데이터로부터 OOD 데이터를 탐지할 때 100% 정확도, 100% 민감도, 100% 특이도, AUC 1.00의 성능을 달성했다.
• 반면, 불확실성 기반 OOD 탐지 방법은 동일한 실험에서 정확도 64%, 민감도 61%, 특이도 60%, AUC 0.65에 머물렀다.
• 분포 내 데이터(Liver-MRI-DUAL-In 및 Liver-MRI-DUAL-Out)의 OODM 값은 항상 임계값 τ=0.0057 이하로 일관되게 유지되었고, 분포 외 데이터(Liver-MRI-SPIR)는 유의미하게 높은 값을 보여 완벽한 분리가 가능했다.
• 테스트 시점에서 모델은 Liver-MRI-DUAL-Out 데이터셋을 성공적으로 분할했으며(DSC=0.886), OODM 값은 학습 시 분포 내 데이터와 유사하게 분포되어 있어 분포 이동에 대한 강건성을 확인했다.