[논문 리뷰] Hypernet-Ensemble Learning of Segmentation Probability for Medical Image Segmentation with Ambiguous Labels

Sung‐Min Hong, Anna K. Bonkhoff외 7명|arXiv (Cornell University)|2021. 12. 13.

Advanced Neural Network Applications인용 수 4

한 줄 요약

이 논문은 3D 의료 영상 분할에서 단일 애매한 레이블만을 사용하여 분할 확률 추정을 향상시키기 위해 하이퍼네트워크-앙상블 방법을 제안한다. 다양한 Tversky 손실을 활용하여 과소분할 및 과다분할을 유도하고, 네트워크 및 하이퍼네트워크 앙상블을 결합함으로써, 해부학적 애매함을 반영하면서도 부드럽고 직관적인 확률 맵을 생성한다. 이는 최신 기술 수준의 이진 분할 성능를 유지한다.

ABSTRACT

Despite the superior performance of Deep Learning (DL) on numerous segmentation tasks, the DL-based approaches are notoriously overconfident about their prediction with highly polarized label probability. This is often not desirable for many applications with the inherent label ambiguity even in human annotations. This challenge has been addressed by leveraging multiple annotations per image and the segmentation uncertainty. However, multiple per-image annotations are often not available in a real-world application and the uncertainty does not provide full control on segmentation results to users. In this paper, we propose novel methods to improve the segmentation probability estimation without sacrificing performance in a real-world scenario that we have only one ambiguous annotation per image. We marginalize the estimated segmentation probability maps of networks that are encouraged to under-/over-segment with the varying Tversky loss without penalizing balanced segmentation. Moreover, we propose a unified hypernetwork ensemble method to alleviate the computational burden of training multiple networks. Our approaches successfully estimated the segmentation probability maps that reflected the underlying structures and provided the intuitive control on segmentation for the challenging 3D medical image segmentation. Although the main focus of our proposed methods is not to improve the binary segmentation performance, our approaches marginally outperformed the state-of-the-arts. The codes are available at \url{https://github.com/sh4174/HypernetEnsemble}.

연구 동기 및 목표

• 의료 영상에서 레이블의 본질적 애매함에도 불구하고 딥러닝에서 과신경성과 극단화된 분할 확률 문제를 해결한다.
• 다중 인간 레이블에 의존하지 않고도 이미지당 하나의 애매한 레이블만을 사용하여 분할 확률 맵 추정을 향상시킨다.
• 조정 가능한 확률 임계값을 통해 사용자에게 분할 결과에 대한 직관적인 제어를 제공한다.
• 다양한 Tversky 손실 설정을 학습하는 통합 하이퍼네트워크 아키텍처를 통해 앙상블 학습의 계산 비용을 감소시킨다.
• 확률 맵 품질을 향상시키면서도 이진 분할 성능을 유지하거나 약간 향상시킨다.

제안 방법

• 다양한 하이퍼파라미터를 가진 Tversky 손실을 사용하여 과소분할, 과다분할 또는 균형 잡힌 분할을 유도하는 다수의 네트워크를 훈련시킨다.
• 예측된 분할 맵을 근거로 이러한 네트워크의 출력을 앙상블하여 더 대표적인 확률 맵을 추정한다.
• 다양한 Tversky 손실 설정에 대해 주 분할 네트워크의 가중치를 학습하는 하이퍼네트워크 아키텍처를 도입하여, 여러 개의 전체 네트워크를 훈련할 필요를 줄인다.
• 다양한 Tversky 손실 설정 간에 파라미터를 공유하는 통합 하이퍼네트워크 앙상블을 구현하여 훈련 비용을 크게 감소시킨다.
• 하나의 입력 조건(Tversky 손실 하이퍼파라미터)에 기반해 주 네트워크의 가중치를 생성하도록 하이퍼네트워크를 훈련시켜 다양한 출력을 효율적으로 추론할 수 있도록 한다.
• 다양한 병원에서의 임상 수준의 DWI 데이터를 사용하여 3D 급성 뇌졸중 병변 분할에 이 방법을 적용한다.

실험 결과

연구 질문

• RQ1 이중 레이블이 없이도 이미지당 하나의 애매한 레이블만을 사용하여 분할 확률 맵을 향상시킬 수 있는가, 이때 이진 분할 성능은 손상되지 않는가?
• RQ2 다양한 Tversky 손실을 가진 네트워크 앙상블은 레이블의 애매함이 존재하는 상황에서 분할 확률 맵 추정을 어떻게 향상시킬 수 있는가?
• RQ3 하이퍼네트워크 기반의 앙상블 접근법은 여러 네트워크를 훈련시키는 데 드는 계산 비용을 줄일 수 있는가, 동시에 확률 맵 추정 품질은 유지되는가?
• RQ4 제안된 방법은 어느 정도까지 확률 임계값 조절을 통해 분할 결과에 대한 사용자 직관적 제어를 가능하게 하는가?
• RQ5 제안된 방법은 분할 확률 맵의 부드러움과 불확실성 표현 측면에서 최신 기술 수준의 접근법과 비교하여 어떻게 성능을 내는가?

주요 결과

• 다양한 Tversky 손실을 적용한 제안된 ResUNet 앙상블은 최고의 집계된 Dice 점수(0.82)와 ROC-AUC(0.884)를 기록하여 최신 기술 수준의 방법을 초월했다.
• 하이퍼네트워크 앙상블는 최고의 균형 잡힌 정확도(0.891)와 재현율(0.783)을 기록했으며, ROC-AUC는 두 번째로 높은 수준(0.869)을 기록했다.
• 제안된 방법의 Dice 점수는 확률 임계값에 따라 부드럽게 변화하여 과소분할 및 과다분할에 대한 직관적인 제어를 가능하게 했다.
• 제안된 방법의 ROC 곡선은 확률 추정이 더 균형 잡혀 있었으며, FPR-TPR 점들이 더 조밀하게 군집되어 있어 보다 캘리브레이션된 확률 맵을 반영했다.
• 하이퍼네트워크 접근법은 훈련 시간을 13 GPU 시간으로 줄였으며, 이는 전체 앙상블 방법에 비해 뚜렷한 감소였다. 동시에 경쟁 가능한 성능를 유지했다.
• 추정된 확률 맵은 애매한 영역(예: 부분 볼륨 효과, 병변 확산)을 성공적으로 반영하여 불확실한 영역에 근접한 0.5 확률을 할당했으며, 시각적 비교에서 이를 확인할 수 있었다.