[논문 리뷰] Combo Loss: Handling Input and Output Imbalance in Multi-Organ Segmentation

Saeid Asgari Taghanaki, Yefeng Zheng외 6명|arXiv (Cornell University)|2018. 05. 07.

Radiomics and Machine Learning in Medical Imaging인용 수 4

한 줄 요약

이 논문은 교육 과정 학습 프레임워크 내에서 딱딱한 손실 함수와 가중치가 부여된 교차 엔트로피를 조합하여, 다중 기관 의료 영상 분할에서 입력 및 출력 클래스 불균형을 동시에 해결하는 새로운 손실 함수인 Combo Loss를 제안한다. 이는 훈련 중에 잘못된 양성/음성 결과 간의 트레이드오프를 명시적으로 제어할 수 있게 하여 수렴 속도를 빠르게 하고 PET, MRI, 초음파 데이터셋에서 최신 기술 수준의 성능을 달성하며, Focal Loss 대비 DSC 점수에서 최대 10.8% 향상된다.

ABSTRACT

Simultaneous segmentation of multiple organs from different medical imaging modalities is a crucial task as it can be utilized for computer-aided diagnosis, computer-assisted surgery, and therapy planning. Thanks to the recent advances in deep learning, several deep neural networks for medical image segmentation have been introduced successfully for this purpose. In this paper, we focus on learning a deep multi-organ segmentation network that labels voxels. In particular, we examine the critical choice of a loss function in order to handle the notorious imbalance problem that plagues both the input and output of a learning model. The input imbalance refers to the class-imbalance in the input training samples (i.e., small foreground objects embedded in an abundance of background voxels, as well as organs of varying sizes). The output imbalance refers to the imbalance between the false positives and false negatives of the inference model. In order to tackle both types of imbalance during training and inference, we introduce a new curriculum learning based loss function. Specifically, we leverage Dice similarity coefficient to deter model parameters from being held at bad local minima and at the same time gradually learn better model parameters by penalizing for false positives/negatives using a cross entropy term. We evaluated the proposed loss function on three datasets: whole body positron emission tomography (PET) scans with 5 target organs, magnetic resonance imaging (MRI) prostate scans, and ultrasound echocardigraphy images with a single target organ i.e., left ventricular. We show that a simple network architecture with the proposed integrative loss function can outperform state-of-the-art methods and results of the competing methods can be improved when our proposed loss is used.

연구 동기 및 목표

• 딥 러닝 기반 다중 기관 분할에서 소수의 전경 볼륨 픽셀(입력 불균형)과 잘못된 양성/음성 결과(출력 불균형)의 이중적 과제를 해결하기 위해.
• 훈련 중 잘못된 양성과 음성 결과 간의 트레이드오프를 명시적이고 학습 가능한 방식으로 제어할 수 있는 손실 함수를 개발하기 위해.
• PET, MRI, 초음파를 포함한 다양한 영상 모odal리티에서 극도로 불균형한 의료 영상 데이터에 대해 분할 성능을 향상시키기 위해.
• 구조적 변경 없이 기존 아키텍처를 개선할 수 있는 플러그인 손실 함수를 제공하여 일반화 능력과 수렴 성능을 향상시키기 위해.
• 후처리에 의존도를 줄이기 위해, 미분 가능하고 적응적인 손실 함수를 통해 잘못된 양성 및 음성 결과를 직접 처벌하기 위해.

제안 방법

• Combo Loss는 분할 안정성 향상과 악성 국소 최소값 회피를 위해 딱딱한 유사도 계수(DSC)와 가중치가 부여된 교차 엔트로피(WCE)를 하나의 미분 가능한 손실 함수로 조합한다.
• 손실 함수는 다음과 같이 정의된다: L_combo = α × (1 - DSC) + (1 - α) × WCE이며, 여기서 α는 DSC와 WCE 항 간의 균형을 조절한다.
• WCE 항에 두 번째 하이퍼파라미터 β를 도입하여 잘못된 양성과 음성 결과에 대한 처벌 간의 트레이드오프를 명시적으로 제어한다.
• 최적화 안정성과 수렴 속도 향상을 위해 점진적으로 WCE 항의 영향력을 증가시키는 교육 과정 학습 방식을 적용한다.
• 구조적 변경 없이 인코더-디코더 네트워크(예: 3D U-Net, V-Net)에서 엔드 투 엔드로 손실 함수를 적용한다.
• 하이퍼파라미터 α와 β는 각 데이터셋에서 성능을 최적화하기 위해 교차 검증을 통해 튜닝된다.

실험 결과

연구 질문

• RQ1 통합된 손실 함수가 다중 기관 분할에서 입력 및 출력 불균형을 효과적으로 다룰 수 있는가?
• RQ2 딱딱한 손실과 가중치가 부여된 교차 엔트로피를 조합함으로써 불균형한 의료 영상 데이터에서 모델의 수렴과 성능 향상이 어떻게 향상되는가?
• RQ3 손실 함수의 학습 가능한 하이퍼파라미터를 통해 잘못된 양성과 음성 결과 간의 트레이드오프를 얼마나 정밀하게 제어할 수 있는가?
• RQ4 제안된 Combo Loss가 기존의 손실 함수들인 교차 엔트로피, 딱딱한 손실, Focal Loss보다 다양한 영상 모달리티에서 우수한 성능을 보이는가?
• RQ5 Combo Loss는 아키텍처 변경 없이 기존의 분할 네트워크를 대체하는 플러그인 손실 함수로 사용될 수 있는가?

주요 결과

• PET 데이터셋에서 Combo Loss는 표준 교차 엔트로피 손실 대비 잘못된 양성률을 34.6% 감소시키고 잘못된 음성률을 27.8% 감소시켰다.
• MRI 데이터셋에서 3D V-Net를 사용할 경우, Combo Loss는 DSC 점수를 2.3% 향상시키고 Focal Loss 대비 잘못된 음성률을 47.4% 감소시켰다.
• 초음파 데이터셋에서 3D U-Net에 적용했을 때, Combo Loss는 Focal Loss 대비 DSC 점수에서 10.8% 향상된 성과를 기록했다.
• 잘못된 양성 및 음성 결과를 최소화하기 위한 최적의 β 값은 초음파에서 0.7, MRI에서 0.6로 확인되어 하이퍼파라미터 제어의 효과를 입증했다.
• U-Net과 V-Net의 원래 손실 함수를 대체했을 때, Combo Loss는 MRI에서 DSC 점수를 0.87에서 0.91로, 초음파에서 0.85에서 0.92로 향상시켜 일관된 성능 향상을 보였다.
• 표준 교차 엔트로피보다 더 빠른 수렴 속도를 달성했고 하우스도르프 거리에서 이상치 성능도 감소시켜 더 견고한 분할 성능을 입증했다.