[논문 리뷰] Tversky loss function for image segmentation using 3D fully convolutional deep networks

Seyed Sadegh Mohseni Salehi, Deniz Erdoğmuş외 1명|arXiv (Cornell University)|2017. 06. 18.

AI in cancer detection인용 수 27

한 줄 요약

이 논문은 다발성 경화증 병변 검출과 같은 극도로 불균형한 의료 데이터에서 이미지 분할 성능을 향상시키기 위해 3D 전결합 신경망에 Tversky 손실 함수를 제안한다. Tversky 지수의 초매개변수를 조정하여 정밀도보다 재현율을 우선시함으로써, 표준 Dice 손실에 비해 F₂ 점수, Dice 계수, 정밀도-재현율 곡선 아래 면적 등에서 유의미하게 향상된 성능을 달성하였으며, β=0.7일 때 최적의 성능을 보였다.

ABSTRACT

Fully convolutional deep neural networks carry out excellent potential for fast and accurate image segmentation. One of the main challenges in training these networks is data imbalance, which is particularly problematic in medical imaging applications such as lesion segmentation where the number of lesion voxels is often much lower than the number of non-lesion voxels. Training with unbalanced data can lead to predictions that are severely biased towards high precision but low recall (sensitivity), which is undesired especially in medical applications where false negatives are much less tolerable than false positives. Several methods have been proposed to deal with this problem including balanced sampling, two step training, sample re-weighting, and similarity loss functions. In this paper, we propose a generalized loss function based on the Tversky index to address the issue of data imbalance and achieve much better trade-off between precision and recall in training 3D fully convolutional deep neural networks. Experimental results in multiple sclerosis lesion segmentation on magnetic resonance images show improved F2 score, Dice coefficient, and the area under the precision-recall curve in test data. Based on these results we suggest Tversky loss function as a generalized framework to effectively train deep neural networks.

연구 동기 및 목표

• 의료 영상 분할에서 병변 볼륨이 비병변 볼륨에 비해 극도로 적은 불균형 데이터 문제를 해결하기 위해.
• 딥 러닝 기반 분할에서 정밀도와 재현율 간의 트레이드오���을 개선하고, 특히 거짓 음성(false negative)을 줄이기 위해 재현율을 높이는 데 초점을 맞추기 위해.
• 다발성 경화증 병변 분할과 같은 극도로 불균형한 상황에서 표준 Dice 손실보다 우수한 성능을 보이는 일반화된 손실 함수를 개발하기 위해.
• 실제 MRI 데이터에서 3D 전결합 U-net 아키텍처에 Tversky 손실의 효과를 검증하기 위해.

제안 방법

• Tversky 지수를 기반으로 한 Tversky 손실 함수를 제안하며, 이는 Dice 계수와 Fβ 점수의 일반화된 형태로 정밀도와 재현율의 균형을 맞추기 위함이다.
• Tversky 손실 레이어를 3D U-net 아키텍처에 통합하여 클래스 불균형을 고려한 엔드 투 엔드 학습을 가능하게 한다.
• 초매개변수 α와 β를 사용하여 가짜 양성(false positive)과 가짜 음성(false negative)에 대한 페널티를 조절하며, β가 높을수록 음성 오류에 대한 민감도가 증가한다.
• 고해상도 국소 특징과 저해상도 전역 특징을 결합하기 위해 스위프 커넥션을 갖춘 3D 전결합 신경망을 사용한다.
• α와 β 값을 다양하게 설정하여 Tversky 손실을 사용해 네트워크를 학습하고, 이들이 분할 성능에 미치는 영향을 평가한다.
• Dice 유사도 계수(DSC), F₂ 점수, 정밀도-재현율 곡선 아래 면적(APR), 민감도/특이도 등의 표준 지표를 사용해 성능을 평가한다.

실험 결과

연구 질문

• RQ1 표준 Dice 손실에 비해 Tversky 손실 함수가 극도로 불균형한 의료 영상 데이터에서 분할 성능을 향상시킬 수 있는가?
• RQ2 Tversky 지수의 초매개변수 α와 β는 병변 분할에서 정밀도-재현율 트레이드오프에 어떤 영향을 미치는가?
• RQ3 β=0.7으로 설정한 Tversky 손실은 작은 또는 희박한 병변을 더 잘 탐지할 수 있는 일반화 능력과 높은 민감도를 제공하는가?
• RQ4 Tversky 손실은 불균형 데이터셋에서 정밀도-재현율 곡선 아래 면적(APR)을 향상시키는 데 표준 Dice 손실보다 더 효과적인가?
• RQ5 Tversky 손실은 고도로 불균형한 클래스를 가진 의료 영상 분할 작업에서 딥 네트워크 학습을 위한 일반화된 프레임워크로 사용될 수 있는가?

주요 결과

• β=0.7일 때 Tversky 손실 함수는 F₂ 점수 57.32를 기록하여, 표준 Dice 손실(β=0.5)의 F₂ 점수 51.77보다 유의미하게 높은 성능을 보였다.
• β=0.7일 때 Tversky 손실은 Dice 계수 56.42를 기록하였으며, 이는 표준 Dice 손실의 53.42보다 더 높은 진짜값과의 겹침을 의미한다.
• β=0.7일 때 정밀도-재현율 곡선 아래 면적(APR)은 56.04에 도달하여 모든 테스트 설정 중에서 가장 높은 성능을 보였으며, 불균형 데이터에서의 우수한 성능을 확인했다.
• β=0.7일 때 민감도(재현율)는 56.85%에 도달하였으며, 이는 표준 Dice 손실에서 기록한 49.85%보다 뚜렷한 향상으로 병변 탐지 능력 향상을 보여준다.
• β=0.7로 설정한 Tversky 손실로 학습된 네트워크는 Dice 손실 모델이 완전히 실패한 저밀도 병변 케이스에서도 병변을 탐지하는 데 성공하였다.
• 결과적으로 β 조정을 통해 모델을 재현율 우선으로 세밀하게 튜닝할 수 있으며, 이는 거짓 음성이 더 심각한 임상 응용 분야에서 매우 중요하다.