[논문 리뷰] Unified Focal loss: Generalising Dice and cross entropy-based losses to handle class imbalanced medical image segmentation

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

Retinal Imaging and Analysis인용 수 9

한 줄 요약

이 논문은 의료 영상 분할에서 클래스 불균형 문제를 더 잘 다루기 위해 딱딱한 손실과 교차 엔트로피 손실을 일반화하는 계층적 프레임워크인 Unified Focal Loss를 제안한다. 다양한 가중치를 적용하는 탄성 있는 방법과 포커스 손실 원리를 결합함으로써, 다섯 개인 불균형 데이터셋에서 일관된 성능 향상을 달성하며, 이중 및 다중 클래스 작업 모두에서 표준 손실 함수와 복합 손실 함수를 능가한다.

ABSTRACT

Automatic segmentation methods are an important advancement in medical image analysis. Machine learning techniques, and deep neural networks in particular, are the state-of-the-art for most medical image segmentation tasks. Issues with class imbalance pose a significant challenge in medical datasets, with lesions often occupying a considerably smaller volume relative to the background. Loss functions used in the training of deep learning algorithms differ in their robustness to class imbalance, with direct consequences for model convergence. The most commonly used loss functions for segmentation are based on either the cross entropy loss, Dice loss or a combination of the two. We propose the Unified Focal loss, a new hierarchical framework that generalises Dice and cross entropy-based losses for handling class imbalance. We evaluate our proposed loss function on five publicly available, class imbalanced medical imaging datasets: CVC-ClinicDB, Digital Retinal Images for Vessel Extraction (DRIVE), Breast Ultrasound 2017 (BUS2017), Brain Tumour Segmentation 2020 (BraTS20) and Kidney Tumour Segmentation 2019 (KiTS19). We compare our loss function performance against six Dice or cross entropy-based loss functions, across 2D binary, 3D binary and 3D multiclass segmentation tasks, demonstrating that our proposed loss function is robust to class imbalance and consistently outperforms the other loss functions. Source code is available at: https://github.com/mlyg/unified-focal-loss

연구 동기 및 목표

• 손상 또는 기관이 영상 부피의 소수의 비율을 차지하는 의료 영상 분할에서 지속적인 클래스 불균형 문제를 해결한다.
• 표준 교차 엔트로피 손실과 딱딱한 손실의 한계를 극복하며, 이는 클래스 불균형에 민감하고 수렴이 불량하거나 편향된 예측을 초래할 수 있다.
• 딱딱한 손실과 교차 엔트로피 기반 손실을 일반화하는 통합된, 미분 가능한 프레임워크를 개발하여 불균형 조건에서의 강건성과 성능을 향상시킨다.
• 2D 이중 분할, 3D 이중 분할, 3D 다중 클래스 분할을 포함한 다양한 의료 영상 작업에서 제안된 손실 함수의 일관된 우수성을 입증한다.
• 기존 U-Net 및 유사 아키텍처에 쉽게 통합할 수 있는 유연하고 즉시 사용 가능한 손실 함수를 제공한다.

제안 방법

• 다양한 가중치를 적용하는 탄성 있는 메커니즘을 통해 딱딱한 손실과 교차 엔트로피 손실을 통합하는 계층적 손실 프레임워크를 제안한다.
• 클래스별 동적 가중치를 적용하여 쉬운 음성 예측을 낮추고 어려운 예측에 집중하는 일반화된 포커스 손실 구성 요소를 도입한다.
• 수정된 딱딱한 손실과 포커스 교차 엔트로피 손실의 가중치 합으로 Unified Focal Loss를 수식화하며, 학습 가능한 또는 고정된 초모수를 사용한다.
• 포커스 효과를 제어하기 위해 감마(γ) 매개변수를 사용하며, 다중 클래스 작업에서는 단순성과 함께 모든 클래스에 동일한 γ=0.5를 사용한다.
• U-Net 및 유사한 인코더-디코더 아키텍처에서 표준 역전파와 호환되도록 손실 함수를 미분 가능하게 보장한다.
• 표준 U-Net 및 3D U-Net 아키텍처를 사용하여 다섯 개의 공개 데이터셋에서 2D 및 3D 분할 네트워크 학습 시 손실 함수를 적용한다.

실험 결과

연구 질문

• RQ1 의료 영상 분할에서 클래스 불균형 문제에 강건한 동시에 딱딱한 손실과 교차 엔트로피 손실을 일반화하는 통합된 손실 함수를 설계할 수 있는가?
• RQ2 다양하고 불균형한 의료 영상 데이터셋에서 Unified Focal Loss는 표준 딱딱한 손실, 교차 엔트로피 손실, 복합 손실 함수(예: Combo loss)와 비교해 성능 면에서 어떻게 다른가?
• RQ3 작은 종양이 큰 기관에서 분할되는 고도로 불균형된 상황과 같이, 제안된 손실 함수는 안정적인 기울기를 유지하고 수렴을 향상시키는가?
• RQ4 다중 희귀 클래스가 포함된 다중 클래스 분할 작업, 예를 들어 신장 종양 분할과 같이 Unified Focal Loss는 일반화 가능한가?
• RQ5 다양한 영상 모odal리티와 네트워크 아키텍처에서 Unified Focal Loss의 성능 향상은 일관된가?

주요 결과

• Unified Focal Loss는 CVC-ClinicDB, DRIVE, BUS2017, BraTS20, KiTS19의 다섯 데이터셋 전반에서 딱딱한 손실, 교차 엔트로피 손실, Combo loss, Tversky loss를 포함한 여섯 가지 기준 손실 함수를 항상 능가했다.
• 매우 불균형한 KiTS19 데이터셋에서 Unified Focal Loss는 신장과 신장 종양 분할 모두에서 가장 높은 DSC(Dice Similarity Coefficient)를 기록했으며, 이는 Combo loss조차도 능가했다.
• BraTS20 데이터셋에서 딱딱한 손실과 Tversky 손실이 불안정한 기울기로 인해 실패한 상황에서도 Unified Focal Loss는 안정적인 학습을 유지하고 뛰어난 분할 품질을 달성했다.
• 정성적 결과에서는 Unified Focal Loss로 생성된 분할 결과가 다른 손실 함수에 비해 더 선명한 경계와 더 적은 가짜 양성/음성 결과를 보였다.
• 기타 손실 함수와 비교해 학습 시간에 유의미한 차이가 없었으며, 이는 효율적인 계산과 실용적인 구현 가능성을 시사한다.
• 모든 클래스에 동일한 하나의 γ 값을 사용했음에도 불구하고, Unified Focal Loss는 다른 복합 손실 함수보다 더 뛰어난 성능을 달성했으며, 이는 초모수 설정에 대해 강건함을 시사한다.