[논문 리뷰] Estimating Uncertainty in Neural Networks for Cardiac MRI Segmentation: A Benchmark Study

Matthew Ng, Fumin Guo외 5명|arXiv (Cornell University)|2020. 12. 31.

Explainable Artificial Intelligence (XAI)인용 수 7

한 줄 요약

이 벤치마크 연구는 심장 MRI 분할을 위한 U-Net 모델에서 불확실성 추정 방법—베이지안 신경망(BBB), 몬테카를로 드롭아웃(MCD), 딥 앙상블(Ensemble), 스토하스틱 스펙트럴 정규화(SSN)—을 평가한다. 결과적으로 BBB와 MCD는 노이즈 및 왜곡 조건에서도 높은 분할 정확도를 유지하면서 신뢰할 수 있는 불확실성 정량화를 제공하는 것으로 나타났으며, 딥 앙상블은 특히 분포 외(scenario) 상황에서 가장 뛰어난 성능을 보였다.

ABSTRACT

Objective: Convolutional neural networks (CNNs) have demonstrated promise in automated cardiac magnetic resonance image segmentation. However, when using CNNs in a large real-world dataset, it is important to quantify segmentation uncertainty and identify segmentations which could be problematic. In this work, we performed a systematic study of Bayesian and non-Bayesian methods for estimating uncertainty in segmentation neural networks. Methods: We evaluated Bayes by Backprop, Monte Carlo Dropout, Deep Ensembles, and Stochastic Segmentation Networks in terms of segmentation accuracy, probability calibration, uncertainty on out-of-distribution images, and segmentation quality control. Results: We observed that Deep Ensembles outperformed the other methods except for images with heavy noise and blurring distortions. We showed that Bayes by Backprop is more robust to noise distortions while Stochastic Segmentation Networks are more resistant to blurring distortions. For segmentation quality control, we showed that segmentation uncertainty is correlated with segmentation accuracy for all the methods. With the incorporation of uncertainty estimates, we were able to reduce the percentage of poor segmentation to 5% by flagging 31--48% of the most uncertain segmentations for manual review, substantially lower than random review without using neural network uncertainty (reviewing 75--78% of all images). Conclusion: This work provides a comprehensive evaluation of uncertainty estimation methods and showed that Deep Ensembles outperformed other methods in most cases. Significance: Neural network uncertainty measures can help identify potentially inaccurate segmentations and alert users for manual review.

연구 동기 및 목표

• 심장 MRI 분할을 위한 딥 러닝 모델에서 다양한 불확실성 추정 기법의 성능을 평가하고 비교하는 것.
• 노이즈, 흐림, 기하학적 변형과 같은 다양한 영상 왜곡 조건에서 불확실성 방법이 분할 정확도를 얼마나 잘 유지하는지 평가하는 것.
• 예측 불확실성 측정치(예: 엔트로피, 상호정보량)와 실제 분할 오차 간의 상관관계를 조사하는 것.
• 임상 영상 환경에서 가장 신뢰할 수 있고 校정된 불확실성 추정을 제공하는 불확실성 추정 방법을 규명하는 것.
• 실제 심장 MRI 분할 작업에서 각 불확실성 추정 방법의 계산 비용 및 추론 효율성(추론 시간 및 학습 기간)을 비교 분석하는 것.

제안 방법

• 베이지안 신경망(BBB), 몬테카를로 드롭아웃(MCD, 드롭아웃 비율 0.1 및 0.5), 딥 앙상블(10개 모델), 스토하스틱 스펙트럴 정규화(SSN)를 포함한 다섯 가지 불확실성 추정 기법을 적용한 U-Net 아키텍처를 학습하였다.
• 강화 학습을 위해 랜덤 회전, 이동, 스케일링 등의 데이터 증강 기법을 적용하여 강건성을 향상시켰다.
• Adam 옵timizer를 사용하였으며, 초기 학습률(lr) = 1e-4로 설정하고, 30 에포크 후에 1e-5로 감소시키는 코즈인 감소 스케줄을 적용하였다.
• 이미지 전처리로 전체 데이터셋의 평균 및 분산을 사용하여 중심 자르기(160×160) 및 정규화를 수행하였다.
• 통제된 왜곡 조건이 가미된 테스트 이미지에서 예측 엔트로피, 상호정보량, 불확실성 인식 메트릭(Dice_WS, ASSD_WS)을 측정하여 불확실성 수준을 평가하였다.
• UK Biobank 및 ACDC 데이터셋에서 성능을 평가하였으며, UKBB에서 ACDC로의 제로샷 전이(zero-shot transfer)도 수행하였다.

실험 결과

연구 질문

• RQ1 노이즈가 증가하는 조건에서 다양한 불확실성 추정 기법(BBB, MCD, Ensemble, SSN) 간의 분할 정확도는 어떻게 비교되는가?
• RQ2 예측 불확실성(예: 엔트로피, 상호정보량)은 왜곡된 심장 MRI 스캔에서 실제 분할 오차와 얼마나 상관이 깊은가?
• RQ3 UK Biobank에서 ACDC 데이터셋으로 일반화할 경우, 어떤 불확실성 추정 방법이 가장 신뢰할 수 있고 잘 校정된 불확실성 추정을 제공하는가?
• RQ4 특히 추론 시간 및 학습 기간 측면에서 불확실성 추정 방법 간의 계산 비용은 어떻게 다름이 있는가?
• RQ5 불확실성 추정을 적용할 경우, 흐림 및 기하학적 변형과 같은 분포 외 왜곡에 대한 모델의 강건성이 향상되는가?

주요 결과

• 딥 앙상블는 청소년 UK Biobank 이미지에서 좌심실(LV) 분할의 최고 중앙값 Dice 스코어(0.965, 93.6–98.3)와 가장 낮은 ASSD(0.74, 0.43–1.12)를 기록하여 일반 U-Net 및 기타 방법들을 압도적으로 능가하였다.
• BBB와 MCD-0.1은 노이즈 및 흐림 조건에서 뛰어난 강건성을 보였으며, 일반 U-Net보다 높은 픽셀 단위 정확도(예: BBB: 노이즈 수준 1에서 0.9929)와 Dice 스코어(예: BBB: 노이즈 수준 1에서 0.947)를 유지하였다.
• 왜곡 수준이 증가함에 따라 상호정보량(MI)이 유의미하게 증가—예를 들어 SSN의 경우 UKBB에서 0.32×10⁻³에서 UKBB→ACDC로 이격된 경우 6.89×10⁻³로 증가—불확실성과 영상 열화 간 강한 상관관계를 시사하였다.
• ACDC 데이터셋에서 BBB와 MCD-0.1은 상호정보량과 ASSD 간에 높은 스피어만 상관계수(r > 0.7)를 보였으며, 이는 신뢰할 수 있는 불확실성 校정 성능을 입증하였다.
• SSN는 높은 불확실성 값을 기록했으나(예: UKBB→ACDC에서 6.89×10⁻³ MI), 분할 정확도는 가장 열악했으며(LV Dice: 0.961), 불확실성 추정에 대해 과신하는 경향을 보였다.
• 통계적 검정 결과, BBB vs. 일반 U-Net(p < 0.0001), MCD-0.1 vs. MCD-0.5(p < 0.0001), 앙상블 vs. SSN(p < 0.0001) 간 유의미한 성능 차이가 확인되어 각 방법의 신뢰성에 차이가 있음을 입증하였다.