[논문 리뷰] Generalised Wasserstein Dice Score for Imbalanced Multi-class Segmentation using Holistic Convolutional Networks

Lucas Fidon, Wen-Zhi Li외 5명|UCL Discovery (University College London)|2017. 07. 03.

Artificial Intelligence in Healthcare인용 수 16

한 줄 요약

이 논문은 레이블 공간 트리와 다중 척도 감독을 통한 통합형 CNN을 활용하여 클래스 간 의미 관계를 통합한 새로운 일반화된 워셔스타인 디스 점수를 제안한다. 이는 불균형한 다중 클래스 뇌종양 세분화를 향상시킨다. 제안된 방법은 기준 손실 및 아키텍처보다 더 의미적으로 유의미한 세분화를 달성하며, 특히 희귀 종양 아형에 대해 더 높은 디스 점수를 기록한다.

ABSTRACT

The Dice score is widely used for binary segmentation due to its robustness to class imbalance. Soft generalisations of the Dice score allow it to be used as a loss function for training convolutional neural networks (CNN). Although CNNs trained using mean-class Dice score achieve state-of-the-art results on multi-class segmentation, this loss function does neither take advantage of inter-class relationships nor multi-scale information. We argue that an improved loss function should balance misclassifications to favour predictions that are semantically meaningful. This paper investigates these issues in the context of multi-class brain tumour segmentation. Our contribution is threefold. 1) We propose a semantically-informed generalisation of the Dice score for multi-class segmentation based on the Wasserstein distance on the probabilistic label space. 2) We propose a holistic CNN that embeds spatial information at multiple scales with deep supervision. 3) We show that the joint use of holistic CNNs and generalised Wasserstein Dice scores achieves segmentations that are more semantically meaningful for brain tumour segmentation.

연구 동기 및 목표

• 희귀 종양 아형 간 의미 관계를 활용하여 다중 클래스 뇌종양 세분화에서 클래스 불균형 문제를 해결한다.
• 새로운 통합형 CNN 아키텍처를 통해 다중 척도 공간 특징을 통합하여 세분화의 견고성을 향상시킨다.
• 계층적 레이블 구조를 인코딩하는 미분 가능한 손실 함수를 개발하여 보다 우수한 일반화 성능을 달성한다.
• 의미 인식 손실과 다중 척도 감독을 조합할 경우 더 뛰어난 세분화 성능을 달성함을 입증한다.

제안 방법

• 레이블 공간에서 사전 정의된 거리 행렬을 사용하여 의미 관계를 인코딩한 레이블 확률 벡터 간 워셔스타인 거리 기반 일반화된 디스 점수를 제안한다.
• 다중 척도 예측을 생성하고 각 척도에서 깊은 감독을 적용하여 특징 학습을 향상시키는 통합형 컨볼루션 신경망(HCNN)을 설계한다.
• 두 단계 학습 전략을 적용: 초기에 평균 소프트 디스 손실로 사전 학습한 후 제안된 워셔스타인 디스 손실로 미세 조정하여 학습 안정성을 확보한다.
• 수동으로 할당된 간선 가중치를 갖는 트리 구조 레이블 공간을 사용하여 종양 유형 간 해부학적 및 병리학적 관계를 반영한다 (예: 괴사 핵, 비강화, 강화).
• 모든 클래스에 동시에 확률적 레이블 맵을 비교하기 위해 워셔스타인 거리를 적용하여 상호 클래스 오분류 페널티의 동시 최적화를 가능하게 한다.
• 학습된 또는 고정된 융합 전략을 사용하여 고해상도 세부 정보와 넓은 수신장 영역의 맥락을 융합한다.

실험 결과

연구 질문

• RQ1 워셔스타인 기반 일반화된 디스 점수는 불균형한 다중 클래스 의료 영상 세분화에서 성능 향상에 기여하는가?
• RQ2 거리 행렬을 통한 계층적 레이블 관계 통합은 모델의 일반화 능력과 의미 일관성에 어떤 영향을 미치는가?
• RQ3 깊이 감독을 통한 다중 척도 특징 학습은 뇌종양 MRI에서 세분화 정확도 향상에 어느 정도 기여하는가?
• RQ4 두 단계 학습 전략(평균 디스 사전 학습 후 워셔스타인 디스 미세 조정)은 단일 손실로의 엔드 투 엔드 학습보다 더 나은 수렴과 성능을 이끌어내는가?

주요 결과

• 제안된 일반화된 워셔스타인 디스 손실($\mathcal{L}_{D^{M_{\text{tree}}}}$)은 특히 괴사 핵 및 비강화 종양과 같은 희귀 클래스에서 평균 소프트 디스 손실보다 더 높은 디스 점수를 기록했다.
• 다중 척도 감독과 깊은 감독을 통합한 통합형 CNN은 표준 U-Net 유사 아키텍처보다 더 일관되고 정확한 세분화를 생성했다.
• HCNN의 융합 예측은 전 종양 디스 점수 89% (±6), 핵심 종양 73% (±22), 강화 종양 74% (±23), $D^{M_{\text{tree}}}$ 87% (±4)를 기록했으며, 개별 척도 및 기준 모델을 모두 초월했다.
• 평균 디스 손실로 사전 학습한 후 워셔스타인 디스 손실로 미세 조정하는 전략($\mathcal{L}_{D^{M_{\text{tree}}}-\text{PT}}$)은 직접 $\mathcal{L}_{D^{M_{\text{tree}}}}$로 학습하는 것보다 안정적인 학습과 더 높은 참 양성률을 달성했다.
• 스케일 1은 수신장이 부족해 낮은 디스 점수를 기록했고, 스케일 3 및 4는 해상도가 낮아 얇은 종양 영역을 포착하지 못해 척도와 해상도 사이의 상충 관계를 드러냈다.
• 정성적 결과에서는 유일하게 융합 예측 및 스케일 2 예측만이 강화 종양 영역을 정확히 세분화했으며, 이는 가장 작고 가장 중요한 구조였다.