[논문 리뷰] A Mixed Focal Loss Function for Handling Class Imbalanced Medical Image Segmentation.

Michael Yeung, Evis Sala외 2명|arXiv (Cornell University)|2021. 02. 08.

Radiomics and Machine Learning in Medical Imaging인용 수 9

한 줄 요약

이 논문은 혼합 포칼 손실 함수를 제안하여, 수정된 포칼 손실과 포칼 딱시 손실을 조합함으로써 클래스 불균형 문제를 개선한다. 이는 KiTS19 및 BraTS20 데이터셋에서 이진 및 다중 클래스 분할 작업 모두에서 더 뛰어난 재현율-정밀도 균형과 강건성을 달성하며, 표준 손실 함수 및 구성 요소 손실 함수를 능가한다.

ABSTRACT

Automatic segmentation methods are an important advancement in medical imaging analysis. Machine learning techniques, and deep neural networks in particular, are the state-of-the-art for most automated medical image segmentation tasks, ranging from the subcellular to the level of organ systems. Issues with class imbalance pose a significant challenge irrespective of scale, with organs, and especially with tumours, often occupying a considerably smaller volume relative to the background. Loss functions used in the training of segmentation algorithms differ in their robustness to class imbalance, with cross entropy-based losses being more affected than Dice-based losses. In this work, we first experiment with seven different Dice-based and cross entropy-based loss functions on the publicly available Kidney Tumour Segmentation 2019 (KiTS19) Computed Tomography dataset, and then further evaluate the top three performing loss functions on the Brain Tumour Segmentation 2020 (BraTS20) Magnetic Resonance Imaging dataset. Motivated by the results of our study, we propose a Mixed Focal loss function, a new compound loss function derived from modified variants of the Focal loss and Focal Dice loss functions. We demonstrate that our proposed loss function is associated with a better recall-precision balance, significantly outperforming the other loss functions in both binary and multi-class image segmentation. Importantly, the proposed Mixed Focal loss function is robust to significant class imbalance. Furthermore, we showed the benefit of using compound losses over their component losses, and the improvement provided by the focal variants over other variants.

연구 동기 및 목표

• 작은 종양 대비 큰 배경 영역과 같은 의료 영상 분할에서의 클래스 불균형 문제를 해결하기 위해.
• 실제 의료 영상 데이터셋에서 다양한 딱시 기반 및 교차 엔트로피 기반 손실 함수의 강건성과 성능을 평가하기 위해.
• 심한 클래스 불균형 하에서 기존 손실 함수의 한계를 규명하고, 더 효과적인 복합 손실 함수 개발의 필요성을 제기하기 위해.
• 포칼 손실과 포칼 딱시 손실의 강점을 결합한 새로운 혼합 포칼 손실 함수를 제안하여 분할 성능을 향상시키기 위해.

제안 방법

• 저자들은 먼저 KiTS19 CT 데이터셋에서 교차 엔트로피 및 딱시 기반 변형을 포함한 7종의 기존 손실 함수를 평가하여 최고 성능을 보인 후보를 선별한다.
• 그 후, 상위 3개의 손실 함수를 선별하여 BraTS20 MR 영상 데이터셋에서 추가 평가함으로써 다양한 영상 모odalities 간의 일반화 능력을 검증한다.
• 제안된 혼합 포칼 손실은 수정된 포칼 손실과 수정된 포칼 딱시 손실을 조합하여 구성되며, 쉬운 음성 예측치를 가중치를 낮추고 어려운, 잘못 분류된 예측치에 집중하도록 설계된다.
• 복합 손실 함수는 모델 훈련 중에 엔드 투 엔드로 최적화되며, 이중 및 다중 클래스 분할 작업에서 재현율과 정밀도의 균형을 맞추기 위해 하이퍼파rameter를 조정한다.
• 이 방법은 포칼 모odulation을 활용하여 쉬운 배경 중심 샘플의 영향을 감소시켜, 종양과 같은 드문 중요한 구조물에 대한 학습을 향상시킨다.

실험 결과

연구 질문

• RQ1 의료 영상 분할에서 심한 클래스 불균형 상황에서 표준 교차 엔트로피 및 딱시 기반 손실 함수는 어떻게 성능을 발휘하는가?
• RQ2 포칼 손실과 포칼 딱시 손실을 조합한 복합 손실 함수는 개별 구성 요소보다 더 뛰어난 재현율-정밀도 균형을 달성할 수 있는가?
• RQ3 제안된 혼합 포칼 손실은 CT 및 MR와 같은 다양한 영상 모달리티 간에 일반화 가능한가?
• RQ4 표준 손실 함수에 비해 포칼 메커니즘이 종양과 같은 작은 드문 구조물에 대해 성능 향상에 얼마나 기여하는가?
• RQ5 불균형 분할 작업에서 복합 손실 함수는 구성 요소 손실 함수보다 더 강건하고 효과적인가?

주요 결과

• 혼합 포칼 손실 함수는 KiTS19 및 BraTS20 데이터셋에서 이진 및 다중 클래스 분할 작업 모두에서 모든 베이스라인 손실 함수를 뛰어넘는 성능을 보였다.
• 제안된 손실 함수는 특히 흔히 표준 손실 함수에 의해 간과되는 작은 드문 종양 영역의 검출을 향상시켜 재현율과 정밀도 간의 균형을 더 잘 달성하였다.
• 복합 손실 함수는 심한 클래스 불균형에 대해 강건성을 보였으며, 종양 영역이 총 부피의 1% 미만일 경우에도 높은 성능을 유지하였다.
• 손실 함수의 포칼 변형은 비포칼 대비 일관되게 뛰어난 성능을 보였으며, 어려운 잘못 분류된 샘플에 집중하는 것이 유의미한 이점을 제공한다는 점을 확인하였다.
• 연구는 복합 손실 함수를 사용할 경우 개별 구성 요소 손실 함수를 사용할 경우보다 측정 가능한 성능 향상이 이루어지며, 혼합 포칼 접근 방식의 설계가 타당하다는 점을 확인하였다.