[논문 리뷰] Marginal loss and exclusion loss for partially supervised multi-organ segmentation

Gonglei Shi, Xiao Li외 2명|arXiv (Cornell University)|2020. 07. 08.

Advanced Neural Network Applications참고 문헌 46인용 수 6

한 줄 요약

이 논문은 부분적으로 레이블이 붙은 데이터셋에서 이미지당 일부 기관만 레이블이 붙어 있는 경우에 단일 다기관 분할 네트워크를 훈련하기 위해 마진 손실(marginal loss)과 배제 손실(exclusion loss)을 제안한다. 마진 손실은 레이블이 없는 기관과 진짜 배경을 병합된 레이블로 모델링하여 마진 확률을 통해 효과적인 훈련을 가능하게 하고, 배제 손실은 기관 간 상호 배제성을 강제한다. 이 방법은 추가 계산 없이 다섯 가지 벤치마크 데이터셋에서 최신 기술 수준의 성능을 달성하며, DICE 스코어를 크게 향상시키고 하우스도르프 거리를 감소시킨다.

ABSTRACT

Annotating multiple organs in medical images is both costly and time-consuming; therefore, existing multi-organ datasets with labels are often low in sample size and mostly partially labeled, that is, a dataset has a few organs labeled but not all organs. In this paper, we investigate how to learn a single multi-organ segmentation network from a union of such datasets. To this end, we propose two types of novel loss function, particularly designed for this scenario: (i) marginal loss and (ii) exclusion loss. Because the background label for a partially labeled image is, in fact, a `merged' label of all unlabelled organs and `true' background (in the sense of full labels), the probability of this `merged' background label is a marginal probability, summing the relevant probabilities before merging. This marginal probability can be plugged into any existing loss function (such as cross entropy loss, Dice loss, etc.) to form a marginal loss. Leveraging the fact that the organs are non-overlapping, we propose the exclusion loss to gauge the dissimilarity between labeled organs and the estimated segmentation of unlabelled organs. Experiments on a union of five benchmark datasets in multi-organ segmentation of liver, spleen, left and right kidneys, and pancreas demonstrate that using our newly proposed loss functions brings a conspicuous performance improvement for state-of-the-art methods without introducing any extra computation.

연구 동기 및 목표

• 여러 데이터셋에서 부분적으로만 레이블이 붙은 경우에 강력한 다기관 분할 모델을 훈련하는 데 도전하는 것.
• 부분적으로 레이블이 붙은 의료 영상에서 제한된 레이블 정보를 효과적으로 활용할 수 있는 손실 함수를 개발하는 것.
• 계산 비용을 증가시키지 않고도 췌장과 신장과 같이 작고 복잡한 기관의 분할 정확도와 내성 강도를 향상시키는 것.
• 레이블링 완전도가 다른 다섯 가지 벤치마크 데이터셋의 결합에서 공동 훈련을 가능하게 하는 것.

제안 방법

• 마진 손실은 배경 레이블을 마진 확률로 모델링하여, 레이블이 없는 모든 기관의 예측 확률과 진짜 배경의 확률을 합산한 후 손실을 계산함으로써 유도된다.
• 마진 손실은 교차 엔트로피나 DICE 손실과 같은 표준 손실 함수와 통합되어, 마진 교차 엔트로피 손실 또는 마진 DICE 손실을 형성할 수 있다.
• 배제 손실은 레이블이 붙은 기관과 레이블이 없는 기관의 예측 세그먼테이션 간의 이질성 측정을 통해 기관 간 상호 배제성을 강제한다.
• 배제 손실은 예측의 겹침을 방지하는 정규화 요소로 작용하며, 해부학적 구조의 본질적인 비중첩 성질을 활용한다.
• 두 손실 함수는 결합되어 표준 U-Net 기반 백본(예: nnUNet)과 함께 훈련 중에 적용되며, 부분적으로 레이블이 붙은 데이터에 대해 엔드 투 엔드 학습을 가능하게 한다.
• 이 방법은 추가 추론 계산이 필요로 하지 않으며, 기존 딥러닝 프레임워크와 호환된다.

실험 결과

연구 질문

• RQ1 마진 손실은 부분적으로 레이블이 붙은 의료 영상에서 병합된 배경 레이블이 유도하는 모호성을 효과적으로 다룰 수 있는가?
• RQ2 배제 손실을 통해 상호 배제성을 강제하면 다기관 환경에서 분할 성능이 어떻게 향상되는가?
• RQ3 마진 손실과 배제 손실을 결합함으로써 췌장과 신장과 같이 작고 분할하기 어려운 기관의 성능 향상 정도는 어느 정도인가?
• RQ4 이러한 새로운 손실 함수만을 사용하여 다섯 가지 부분적으로 레이블이 붙은 벤치마크 데이터셋의 결합에서 단일 분할 네트워크를 효과적으로 훈련시킬 수 있는가?

주요 결과

• 제안된 방법은 모든 기관에서 평균 DICE 계수 0.931을 달성하여 기존 방법들인 PaNN(0.906)과 PIPO(0.907)를 크게 앞서나간다.
• PaNN 대비 평균 하우스도르프 거리는 24.0% 감소하고, PIPO 대비 40.0% 감소하여 경계 정확도가 뛰어나다는 것을 시사한다.
• 배제 손실만으로는 췌장과 신장에서 최소한의 향상이 나타나, 마진 손실과 함께 보조 손실로써 가장 효과적으로 작용한다는 것을 시사한다.
• 레이블 수가 줄어들수록 성능이 일관되게 떨어지므로, 모델 성능이 레이블 양에 민감함을 확인한다.
• 시각화 결과는 배제 손실이 이웃한 구조로 침범하는 잘못된 기관 예측을 효과적으로 줄임을 보여주며, 특히 작은 기관에서 두드러진다.
• 실패 사례 분석은 간과 유사 배경 영역를 구분하는 데 어려움이 있으며, 비장과 췌장을 잘못 분류하는 문제를 드러내어 형태 우선 정보를 도입함으로써 향상 가능성을 시사한다.