[논문 리뷰] Boundary loss for highly unbalanced segmentation

Hoel Kervadec, Jihene Bouchtiba외 4명|Espace ÉTS (ETS)|2018. 12. 17.

Advanced Neural Network Applications인용 수 6

한 줄 요약

이 논문은 고도로 불균형한 의료 영상 분할을 위한 경계 손실을 제안하며, 영역 내 픽셀 합계가 아닌 윤곽 거리 기반으로 작동하여 학습 안정성과 성능을 향상시킨다. 비대칭 L2 거리 측정을 소프트맥스 출력을 이용한 영역 적분으로 표현함으로써, 기존 네트워크와의 간편한 통합이 가능해지고, WMH 및 ISLES와 같은 불균형 데이터셋에서 DSC 점수와 안정성에 있어 뚜렷한 향상을 이룬다.

ABSTRACT

Widely used loss functions for CNN segmentation, e.g., Dice or cross-entropy, are based on integrals over the segmentation regions. Unfortunately, for highly unbalanced segmentations, such regional summations have values that differ by several orders of magnitude across classes, which affects training performance and stability. We propose a boundary loss, which takes the form of a distance metric on the space of contours, not regions. This can mitigate the difficulties of highly unbalanced problems because it uses integrals over the interface between regions instead of unbalanced integrals over the regions. Furthermore, a boundary loss complements regional information. Inspired by graph-based optimization techniques for computing active-contour flows, we express a non-symmetric $L_2$ distance on the space of contours as a regional integral, which avoids completely local differential computations involving contour points. This yields a boundary loss expressed with the regional softmax probability outputs of the network, which can be easily combined with standard regional losses and implemented with any existing deep network architecture for N-D segmentation. We report comprehensive evaluations and comparisons on different unbalanced problems, showing that our boundary loss can yield significant increases in performances while improving training stability. Our code is publicly available: https://github.com/LIVIAETS/surface-loss .

연구 동기 및 목표

• 의료 영상에서 소규모 병변 탐지와 같은 고도로 불균형한 분할 작업에서 표준 영역 기반 손실 함수(예: 교차 엔트로피, DSC)의 불안정성과 열악한 성능을 해결한다.
• 클래스 불균형으로 인해 큰 배경 영역이 우세해지면서 최적화가 주로 다수 클래스에 치우치는 문제를 해결한다.
• 분할 오류가 가장 중요한 영역인 경계 영역에 집중하기 위해 윤곽 수준의 거리 지표를 활용한 손실 함수를 개발한다.
• 기존 딥 러닝 아키텍처 및 N차원 분할 작업과의 호환성을 확보하며, 아키텍처 수정 없이도 적용 가능하도록 한다.
• 특히 백색질 고강도 병변과 같은 소규모 희귀 구조에서 분할 정확도와 학습 안정성을 향상시킨다.

제안 방법

• 윤곽의 공간에서 비대칭 L2 거리 기반의 경계 손실을 제안하며, 윤곽 점에 직접 계산을 수행하지 않도록 한다.
• 네트워크의 소프트맥스 확률 출력을 이용해 이 윤곽 기반 거리를 영역 적분으로 표현함으로써, 표준 백프로파게이션을 사용한 미분 가능한 학습이 가능해진다.
• 예측 및 진짜 값 확률 맵에만 의존하는 영역 기반 손실로 경계 손실을 제안함으로써, 어떤 기존 CNN 아키텍처와도 호환 가능하다.
• 일반화된 DSC 손실과 같은 표준 영역 기반 손실과 가중치 합을 통해 조합하며, 학습 가능한 또는 스케줄링 가능한 하이퍼파rameter α를 사용한다.
• 학습 중에 α에 대한 재균형 전략을 사용하여 경계 손실의 기여도를 동적으로 조정함으로써 일반화 성능을 향상시킨다.
• 특수 레이어 유형이나 복잡한 후처리가 필요 없이 표준 딥 러닝 프레임워크를 사용해 효율적으로 구현된다.

실험 결과

연구 질문

• RQ1 표준 영역 기반 손실과 비교해, 윤곽 기하학에 기반한 손실 함수가 고도로 불균형한 의료 영상 데이터셋에서 분할 성능을 향상시킬 수 있는가?
• RQ2 경계 정보까지의 거리에 기반한 경계 손실이 소구조 분할에서 학습 안정성과 수렴에 어떤 영향을 미치는가?
• RQ3 경계 손실을 영역 기반 손실(예: 일반화된 DSC)과 조합할 때 최적의 하이퍼파ram터 스케줄링 및 가중치 설정은 무엇인가?
• RQ4 제안된 경계 손실이 부드럽고 해부학적으로 타당한 분할 경계를 유도하는 공간 정규화 효과를 제공하는가?
• RQ5 제안된 경계 손실은 이진 분할을 초월해 다중 클래스 또는 다중 구조 분할 문제에 효과적으로 적용될 수 있는가?

주요 결과

• WMH 데이터셋에서 재균형 α 전략을 사용해 GDL과 조합한 경우, DSC가 GDL 전용일 때 0.210에서 0.392로 향상되어 뚜렷한 성능 향상을 보였다.
• ISLES 데이터셋에서 재균형 α를 사용한 경계 손실은 DSC 0.644(±0.026)를 기록했으며, GDL 전용(0.511) 및 모든 일정 α 설정보다 뛰어난 성능을 보였다.
• 학습 DSC가 유사한 경우에도 검증 DSC가 크게 향상되어, 더 나은 일반화 능력과 과적합 감소를 나타냈다.
• 경계 손실의 학습 시간은 표준 손실과 유사했으며, 평균 배치 처리 시간이 0.210초로, 최소한의 계산 오버헤드를 보였다.
• 경계 손실은 특히 소규모 구조에서 더 부드럽고 정확한 윤곽을 유도했으며, 그림 1에서 시각적으로 확인되었으며, 놓친 병변이 복구된 것으로 나타났다.
• 재균형 α 스케줄링 전략은 α=1 및 α=2를 포함한 모든 일정 α 값보다 뛰어난 성능을 보였으며, 광범위한 하이퍼파ram터 튜닝 없이도 일관된 향상을 제공했다.