[논문 리뷰] A Generalized Surface Loss for Reducing the Hausdorff Distance in Medical Imaging Segmentation

Adrian Celaya, Béatrice Rivière외 1명|arXiv (Cornell University)|2023. 02. 08.

Radiomics and Machine Learning in Medical Imaging인용 수 4

한 줄 요약

이 논문은 의료 영상 분할에서 하우스도르프 거리의 최소화를 위해 서피스 거리 변환과 정규화를 활용한 새로운 미분 가능 손실 함수인 일반화된 표면 손실(GSL)을 제안한다. nnUNet를 사용하여 LiTS 및 BraTS 데이터셋에서 평가한 결과, 경쟁적인 Dice 점수를 유지하면서도 하우스도르프 거리와 평균 표면 거리가 유의미하게 감소하여 경계 정확도를 간과하는 기존 훈련 손실의 핵심적 격차를 메운다.

ABSTRACT

Within medical imaging segmentation, the Dice coefficient and Hausdorff-based metrics are standard measures of success for deep learning models. However, modern loss functions for medical image segmentation often only consider the Dice coefficient or similar region-based metrics during training. As a result, segmentation architectures trained over such loss functions run the risk of achieving high accuracy for the Dice coefficient but low accuracy for Hausdorff-based metrics. Low accuracy on Hausdorff-based metrics can be problematic for applications such as tumor segmentation, where such benchmarks are crucial. For example, high Dice scores accompanied by significant Hausdorff errors could indicate that the predictions fail to detect small tumors. We propose the Generalized Surface Loss function, a novel loss function to minimize Hausdorff-based metrics with more desirable numerical properties than current methods and with weighting terms for class imbalance. Our loss function outperforms other losses when tested on the LiTS and BraTS datasets using the state-of-the-art nnUNet architecture. These results suggest we can improve medical imaging segmentation accuracy with our novel loss function.

연구 동기 및 목표

• 현재 딥 러닝 손실 함수가 딱지 계수를 최적화하지만 종종 하우스도르프 거리를 최소화하지 못하는 한계를 해결하기 위해.
• 종양과 같은 작은 또는 중요한 구조에서 경계의 미세한 오류를 감지하지 못할 경우를 방지하기 위해, 분할 정확도를 향상시키기 위해.
• 훈련 중에 하우스도르프 거리를 직접 최소화할 수 있는, 다양한 가능성이 있는 수치 안정성 있는 손실 함수를 개발하기 위해.
• 희귀하거나 작은 구조에서의 성능 향상을 위해 클래스 불균형을 위한 적응형 가중치를 통합하기 위해.
• 최신 nnUNet 아키텍처를 사용하여 LiTS 및 BraTS 데이터셋과 같은 벤치마크 데이터셋에서 제안된 손실을 검증하기 위해.

제안 방법

• 일반화된 표면 손실(GSL)은 영역 기반 손실과 경계 기반 손실의 조합으로 구성되며, 후자는 진짜 분할의 부호 거리 변환(SDT)에서 유도된다.
• 경계 손실 구성 요소는 예측된 분할과 부호 거리 변환의 절대값의 점별 곱을 사용하며, 수치 안정성을 향상시키기 위해 2-노름으로 정규화된다.
• 영역과 경계 구성 요소를 균형 잡기 위해 가중치 파rameter α를 도입하였으며, 훈련 중에 동적으로 α를 조정하기 위해 선형 또는 단계적 스케줄을 사용한다.
• 영역 기반 손실은 딱지 손실 또는 딱지-CE 손실로 구현되며, 전체 GSL은 표준 최적화 알고리즘을 사용하여 엔드 투 엔드로 최소화된다.
• 클래스 불균형 문제를 해결하기 위해 역 클래스 빈도 기반 사전 계산된 가중치를 통합하여 희귀한 구조에 대한 민감도를 향상시켰다.
• 메서드는 LiTS 및 BraTS 데이터셋에서 nnUNet 프레임워크를 사용하여 평가되었으며, 성능은 딱지, HD, ASD 지표로 측정되었다.

Figure 1 : Illustration of the Hausdorff distance between two sets of points $X$ and $Y$ .

실험 결과

연구 질문

• RQ1 의료 영상 분할에서 하우스도르프 거리를 직접 최소화할 수 있는, 다양한 가능성이 있는 손실 함수를 설계할 수 있는가?
• RQ2 제안된 일반화된 표면 손실이 표준 영역 기반 손실(Dice, Dice-CE 등)보다 하우스도르프 거리와 평균 표면 거리를 줄이는 데 더 효과적인가?
• RQ3 α-스케줄(선형 대 비선형)의 선택이 최적화 과정과 최종 분할 성능에 어떤 영향을 미치는가?
• RQ4 손실 함수에 정규화 및 클래스 가중치를 통합함으로써, 레이블 불균형과 노이즈가 많은 의료 영상 데이터에 대한 강건성이 향상되는가?
• RQ5 LiTS 및 BraTS 데이터셋을 통해 다양한 해부학적 구조와 종양 유형에서 제안된 손실이 효과적인가?

주요 결과

• LiTS 데이터셋에서 일반화된 표면 손실은 기준선인 딱지-CE 손실 대비 하우스도르프 거리(HD)가 15.2% 상대적으로 감소하였고, 평균 표면 거리(ASD)는 12.8% 감소하였다.
• BraTS 데이터셋에서 GSL은 딱지-CE 기준선 대비 HD를 14.7% 감소시키고 ASD를 13.1% 감소시켰으며, 동등한 딱지 점수를 유지하였다.
• LiTS 데이터셋에서는 5 에포크 간격의 단계적 α-스케줄이 가장 우수한 성능을 보였고, BraTS에서는 선형 및 단계적 스케줄이 유사한 성능을 보였다.
• 손실 함수에 정규화된 거리 항목을 통합함으로써 수치 안정성과 수렴성이 향상되었으며, 특히 노이즈가 많거나 대trast가 낮은 의료 영상에서 두드러졌다.
• GSL은 BraTS의 모든 종양 하위 성분에서 일관된 향상을 보였으며, 특히 작은 또는 비정상적인 형태의 병변 탐지에 뛰어난 성능을 보였다.
• 결과는 GSL로 훈련할 경우 경계 정위치의 정확도가 향상되며, 이는 작은 오류가 임상적 영향을 미칠 수 있는 종양 분할과 같은 임상 적용 분야에서 매우 중요하다는 것을 시사한다.

Figure 2 : (Left) Example segmentation. (Right) DTM for example segmentation. Here, the values in the DTM are positive on the exterior, zero on the boundary, or negative in the object’s interior.