[논문 리뷰] A Novel Focal Tversky loss function with improved Attention U-Net for lesion segmentation

Nabila Abraham, Naimul Khan|arXiv (Cornell University)|2018. 10. 18.

Radiomics and Machine Learning in Medical Imaging인용 수 14

한 줄 요약

이 논문은 소규모 병변의 의료 영상 분할에서 정밀도와 재현율의 균형을 향상시키는 새로운 포커스 Tversky 손실 함수를 제안하며, 다중 척도 영상 피라미드를 통합한 개선된 어텐션 U-Net을 함께 사용한다. 이 방법은 BUS 2017 데이터셋에서 DICE 스코어에 25.7%의 상대적 향상을 기록하고, ISIC 2018에서는 3.6% 향상을 기록하여 표준 U-Net에 비해 DICE 손실을 사용하는 모델를 능가하며, 극도로 불균형한 데이터에서 안정적이고 균형 잡힌 성능을 보인다.

ABSTRACT

We propose a generalized focal loss function based on the Tversky index to address the issue of data imbalance in medical image segmentation. Compared to the commonly used Dice loss, our loss function achieves a better trade off between precision and recall when training on small structures such as lesions. To evaluate our loss function, we improve the attention U-Net model by incorporating an image pyramid to preserve contextual features. We experiment on the BUS 2017 dataset and ISIC 2018 dataset where lesions occupy 4.84% and 21.4% of the images area and improve segmentation accuracy when compared to the standard U-Net by 25.7% and 3.6%, respectively.

연구 동기 및 목표

• 소규모 병변이 영상 면적의 극히 소수를 차지하는 경우에 발생하는 클래스 불균형 문제를 해결하기 위해.
• 소규모 병리적 영역의 의미적 분할에서 정밀도와 재현율 간의 트레이드오���을 개선하기 위해.
• 다중 척도 영상 피라미드를 어텐션 U-Net 아키텍처에 통합하여 소규모 병변에 대한 특징 표현을 향상시키기 위해.
• 훈련 과정에서 어려운 예외(잘 분류되지 않은 소규모 병변)에 더 집중하고 쉬운 음성 예외를 억제하기 위해 손실 함수를 설계하기 위해.
• 극도로 불균형한 데이터셋에서 정밀도와 재현율의 변동성이 낮고 안정적이고 균형 잡힌 분할 성능을 달성하기 위해.

제안 방법

• Tversky 지수의 역수를 하이퍼파rameter γ의 거듭제곱으로 승격시켜 쉬운 예외를 억제하고 어려운, 잘못 분류된 픽셀에 집중하는 포커스 Tversky 손실(FTL)을 제안한다.
• 클래스별 가중치 α와 β를 사용해 Tversky 지수를 조정하며, α > β로 설정하여 가짜 음성 결과를 줄이고 소규모 병변의 재현율을 향상시킨다.
• 각 인코더 레이어에 다중 척도 영상 피라미드 입력을 도입하여 맥락적 특징을 유지하고, 다운샘플링 과정에서 소규모 병변의 세부 정보 손실을 방지한다.
• 스킵 연결과 어텐션 게이트를 통합하여 어텐션 U-Net을 개선하여 주목할 만한 병변 영역을 동적으로 강조하고 배경 노이즈를 억제한다.
• 모델이 수렴에 가까워졌을 때도 최종 레이어를 Tversky 손실로 감독하는 하이브리드 훈련 전략을 사용하여 오버억제를 방지하고 오류 신호 전파를 유지한다.
• 학습률 감소를 적용한 확률적 경사 하강법을 사용하고, BUS 2017 및 ISIC 2018 데이터셋에서 하이퍼파rameter(α, β, γ) 튜닝을 위해 그리드 서치를 수행한다.

실험 결과

연구 질문

• RQ1 Tversky 지수를 기반으로 한 수정된 포커스 손실이 소규모 병변 분할에서 정밀도와 재현율의 균형을 향상시킬 수 있는가?
• RQ2 다중 척도 영상 피라미드를 통합할 경우, 극도로 불균형한 데이터셋에서 소규모 병변의 특징 표현과 분할 정확도에 어떤 영향을 미치는가?
• RQ3 포커스 Tversky 손실과 어텐션 U-Net 아키텍처를 결합하면 표준 U-Net에 비해 더 안정적이고 정확한 분할 성능을 달성할 수 있는가?
• RQ4 대규모 내부 클래스 변동성과 소규모 병변 크기를 가진 데이터셋에서 훈련할 경우, 손실 함수 설계가 수렴 행동과 성능 안정성에 어떤 영향을 미치는가?
• RQ5 제안된 방법은 병변 크기 분포와 영상 복잡도가 다른 데이터셋 간에 일관된 성능 향상을 달성할 수 있는가?

주요 결과

• α=0.7, β=0.3, γ=4/3로 설정된 제안된 포커스 Tversky 손실은 표준 DICE 및 Tversky 손실에 비해 정밀도-재현율 균형을 크게 향상시킨다.
• BUS 2017 데이터셋에서 개선된 어텐션 U-Net에 포커스 Tversky 손실을 적용한 결과, DICE 스코어가 0.804를 기록하여 기준 U-Net에 비해 25.7%의 상대적 향상을 기록했다.
• ISIC 2018 데이터셋에서 기준 U-Net 대비 DICE 스코어가 3.6% 향상되었으며, 표준편차 0.7%로 낮아 안정적인 성능을 보였다.
• 다중 척도 영상 피라미드로 인해 특징 표현이 향상되었으며, 특히 소규모 병변에서 다운샘플링 과정에서의 정보 손실 감소와 함께 분할 정확도 향상에 기여했다.
• 제거 분석 결과, 포커스 Tversky 손실과 개선된 어텐션 U-Net을 결합할 경우 두 데이터셋에서 가장 높은 DICE 스코어와 가장 균형 잡힌 정밀도-재현율 스코어를 확보했다.
• 수렴에 가까워졌을 때도 최종 레이어에서 Tversky 손실을 감독하여 오류 신호 전파를 유지함으로써, ISIC 2018에서 높은 재현율(0.926)을 유지하면서 정밀도가 과도하게 떨어지는 것을 방지했다.