[논문 리뷰] On the influence of Dice loss function in multi-class organ segmentation of abdominal CT using 3D fully convolutional networks

Chen Shen, Holger R. Roth외 5명|arXiv (Cornell University)|2018. 01. 18.

Radiomics and Machine Learning in Medical Imaging참고 문헌 6인용 수 31

한 줄 요약

이 연구는 3D 완전 컨volution 네트워크(FCNs)를 사용한 다중 분류 복부 장기 분할에서 세 가지 Dice 손실 가중 전략—균일, 단순, 제곱—의 영향을 조사한다. 연구에서는 클래스 균형 가중치와 초기 학습률의 선택이 분할 정확도에 상당한 영향을 미친다는 것을 입증하며, 특히 단순 가중 전략이 학습률 0.01일 때 평균 Dice 스코어 81.0%를 기록하여 가장 높은 성능을 보였다.

ABSTRACT

Deep learning-based methods achieved impressive results for the segmentation of medical images. With the development of 3D fully convolutional networks (FCNs), it has become feasible to produce improved results for multi-organ segmentation of 3D computed tomography (CT) images. The results of multi-organ segmentation using deep learning-based methods not only depend on the choice of networks architecture, but also strongly rely on the choice of loss function. In this paper, we present a discussion on the influence of Dice-based loss functions for multi-class organ segmentation using a dataset of abdominal CT volumes. We investigated three different types of weighting the Dice loss functions based on class label frequencies (uniform, simple and square) and evaluate their influence on segmentation accuracies. Furthermore, we compared the influence of different initial learning rates. We achieved average Dice scores of 81.3%, 59.5% and 31.7% for uniform, simple and square types of weighting when the learning rate is 0.001, and 78.2%, 81.0% and 58.5% for each weighting when the learning rate is 0.01. Our experiments indicated a strong relationship between class balancing weights and initial learning rate in training.

연구 동기 및 목표

• 3D 복부 CT 볼륨에서의 다장기 분할 성능에 대해 다양한 Dice 손실 가중 전략의 영향을 분석하기 위해.
• Dice 기반 손실 함수를 사용할 때 초기 학습률이 학습 수렴과 분할 정확도에 미치는 영향을 조사하기 위해.
• 다양한 병변 빈도를 가진 복부 장기의 분할 향상을 위한 최적의 손실 가중치와 학습률 조합을 규명하기 위해.
• 데이터 증강 없이 3D U-Net 모델의 일반화 능력과 안정성을 다양한 손실 함수 구성에서 평가하기 위해.
• 다중 분류 의료 영상 분할 작업에서 3D FCNs의 하이퍼파라미터 선택에 대한 실증적 지침을 제공하기 위해.

제안 방법

• 복강 CT 볼륨에서의 엔드 투 엔드 3D 다장기 분할을 위해 백본 네트워크로 3D U-Net 아키텍처를 사용하였다.
• 세 가지 Dice 손실 가중 전략을 구현: 균일(Wu=1), 단순(Ws = N / (L|Rl| + ε)), 제곱(Wq = N / (L|Rl|² + ε)), 여기서 |Rl|은 클래스 l의 볼록수이다.
• 확률적 경사 하강법을 사용하여 두 가지 초기 학습률 μ = 0.001 및 μ = 0.01로 모델을 학습시켰다.
• 전체 CT 볼륨에서 64×64×64 서브볼륨을 무작위로 자르는 방식으로 미니배치 학습과 전체 볼륨 추론을 가능하게 하였다.
• 팔 개의 복부 장기에서 분할 성능을 평가하기 위해 Dice 유사도 계수(DSC)를 사용하였다.
• 가중치 수식에서 0으로 나누는 것을 방지하기 위해 ε = 1을 적용하였다.

실험 결과

연구 질문

• RQ1 Dice 손실 함수의 균일, 단순, 제곱 가중 전략이 여러 복부 장기의 분할 정확도에 어떻게 영향을 미치는가?
• RQ2 초기 학습률을 0.001과 0.01로 변경했을 때, Dice 기반 손실 함수의 수렴 및 성능에 어떤 영향을 미치는가?
• RQ3 손실 가중치와 학습률의 어떤 조합이 모든 장기에서 가장 높은 평균 Dice 스코어를 도출하는가?
• RQ4 클래스 불균형과 볼록 빈도는 다양한 Dice 손실 가중 전략의 효과성에 어떻게 영향을 미치는가?
• RQ5 가중 전략 간에 훈련 동역학(예: 손실 곡선 행동)이 유의미하게 다를까? 이는 서로 다른 수렴 패턴을 시사하는가?

주요 결과

• 학습률 0.001일 때 균일 가중 전략이 모두 8개의 장기에서 평균 Dice 스코어 81.3%로 가장 높은 성능을 기록하였다.
• 학습률 0.01일 때 단순 가중 전략이 다른 전략들을 압도하며 평균 Dice 스코어 81.0%를 기록했으며, 동맥, 문맥정맥, 췌장에서의 성능 향상이 뚜렷했다.
• 제곱 가중 전략은 μ = 0.01일 때 뚜렷한 향상을 보였으며, 평균 Dice 스코어 58.5%를 기록했고, 간, 비장, 위의 분할 성능 향상이 두드러졌다.
• 단순 가중 전략은 학습률을 높임에 따라 가장 큰 성능 향상을 보였으며, μ = 0.001일 때 59.5%에서 μ = 0.01일 때 81.0%로 상승하였다.
• 훈련 곡선 분석 결과 균일 가중 전략 모델은 10,000 반복 이후에 수렴한 반면, 단순 및 제곱 전략은 10,000 반복 이후에도 향상 잠재력이 있었음을 보였다.
• 데이터 증강 없이도 강한 과적합 현상이 관찰되지 않아, 훈련 데이터셋이 충분히 크고 다양하다는 것을 시사하였다.