[논문 리뷰] Anatomical and Diagnostic Bayesian Segmentation in Prostate MRI $-$Should Different Clinical Objectives Mandate Different Loss Functions?

Anindo Saha, Joeran S. Bosma외 3명|arXiv (Cornell University)|2021. 10. 25.

Prostate Cancer Diagnosis and Treatment인용 수 5

한 줄 요약

이 연구는 프로스타트 MRI에서 해부학적 분할(기관 구조)과 진단적 분할(임상적으로 중요한 암 검출)이라는 서로 다른 임상 목적이 Bayesian U-Net 모델의 손실 함수를 별도로 필요로 하는지 조사한다. 분포 기반 손실 함수, 특히 포칼 손실이 경계 기반 손실보다 유의하게 뛰어난 성능을 보이며 병변 검출에서 모델 캘리브레이션 향상 덕분에 뛰어난 성능을 보였고, 해부학적 분할에서는 모든 손실 함수가 유사한 성능를 보였지만, 포칼 손실은 경계 근처의 높은 불확실성으로 인해 성능가 낮아졌다.

ABSTRACT

We hypothesize that probabilistic voxel-level classification of anatomy and malignancy in prostate MRI, although typically posed as near-identical segmentation tasks via U-Nets, require different loss functions for optimal performance due to inherent differences in their clinical objectives. We investigate distribution, region and boundary-based loss functions for both tasks across 200 patient exams from the publicly-available ProstateX dataset. For evaluation, we conduct a thorough comparative analysis of model predictions and calibration, measured with respect to multi-class volume segmentation of the prostate anatomy (whole-gland, transitional zone, peripheral zone), as well as, patient-level diagnosis and lesion-level detection of clinically significant prostate cancer. Notably, we find that distribution-based loss functions (in particular, focal loss) are well-suited for diagnostic or panoptic segmentation tasks such as lesion detection, primarily due to their implicit property of inducing better calibration. Meanwhile, (with the exception of focal loss) both distribution and region/boundary-based loss functions perform equally well for anatomical or semantic segmentation tasks, such as quantification of organ shape, size and boundaries.

연구 동기 및 목표

• 딥 러닝 분할에서 프로스타트 MRI의 서로 다른 임상 목표—기관 수준의 해부학적 구조와 병리 수준의 암 검출—이 별도의 손실 함수를 필요로 하는지 확인하기.
• 분포 기반(예: 교차 엔트로피, 포칼 손실)과 영역/경계 기반(예: Dice, 경계 손실) 손실 함수가 해부학적 및 진단적 작업에서 분할 성능에 미치는 영향 평가하기.
• 다양한 손실 함수 하에서 베이지안 U-Net 모델의 모델 캘리브레이션과 불확실성 정량화 평가하기.
• 병변 검출에 대해 표준 지표(DSC, AUROC)와 작업 특화 지표(FROC)를 사용해 성능 비교하기.
• 임상 작업 목표에 기반한 손실 함수 선택에 대한 증거 기반 권고 제시하기.

제안 방법

• 200건의 프로스타트 MRI 검사(ProstateX 데이터셋)를 기반으로, 전체 선립선, 전이존, 주변존 및 임상적으로 중요한 프로스타트 암에 대해 병변 수준의 주석이 있는 확률적 3D U-Net을 몬테카를로 드롭아웃을 사용해 학습.
• 네 가지 손실 함수 평가: 교차 엔트로피(CE), 균형 잡힌 교차 엔트로피(BCE), 포칼 손실(FL), 소프트 Dice-경계 손실(DB), 작업별 클래스 가중치 및 하이퍼파라미터 설정.
• 정규 분포 노이즈, 플립, 회전, 이동, 스케일링을 포함한 데이터 증강을 통해 모델의 강건성 향상.
• 기대 캘리브레이션 오차(ECE)와 신뢰도 다이어그램을 사용해 모델 캘리브레이션 평가, 불확실성 정량화에는 몬테카를로 드롭아웃 활용.
• 성능 평가 지표로 해부학적 분할에는 3D Dice 유사도 계수(DSC), 환자 수준 진단에는 AUROC, 병변 검출에는 부분 FROC 아래 면적 사용.
• 실험은 3개의 독립 실행에서 5겹 교차 검증을 사용하며, 각 이미지당 100회의 몬테카를로 추론 샘플을 통해 불확실성 추정.

실험 결과

연구 질문

• RQ1 해부학적 작업과 진단적 작업을 구분할 때, 손실 함수의 선택이 프로스타트 MRI 분할 성능에 유의미한 영향을 미치는가?
• RQ2 분포 기반 손실 함수(예: 포칼 손실)가 영역/경계 기반 손실보다 더 우수한 모델 캘리브레이션과 불확실성 추정을 제공하는가?
• RQ3 포칼 손실이 클래스 가중치나 아키텍처 설계 외에도 암 병변 검출에서 우수한 성능를 보이는 이유는 암묵적 캘리브레이션 특성 덕분인가?
• RQ4 소프트 Dice나 경계 손실과 같은 영역/경계 기반 손실 함수는 다중 병변 검출과 같은 팬옵티크 분할 작업에서는 최적화되지 않아, DSC가 높더라도 성능이 떨어지는가?
• RQ5 DSC와 같은 표준 평가 지표는 병변 감지 민감도와 캘리브레이션 중요성이 더 높은 진단 작업에서의 성능 차이를 포착하지 못하는가?

주요 결과

• 포칼 손실(FL)이 병변 검출에서 소프트 Dice-경계(DB) 손실보다 유의미하게 뛰어나, 최대 검출 민감도가 15.7% 향상되었고, 환자당 0.1–2.5개의 가짜 경고가 발생할 때 부분 FROC 아래 면적에서 0.25 향상됨.
• FL은 ECE가 낮고 더 신뢰할 수 있는 신뢰도 다이어그램을 보여, 더 나은 모델 캘리브레이션을 확보함으로써 위험 분류에 유리함.
• 해부학적 분할(WG, TZ, PZ)에서는 FL이 CE, BCE, DB보다 낮은 DSC(예: WG 70, TZ 64, PZ 50)를 보였는데, 주로 기관 경계 근처에서 높은 예측 불확실성 때문임.
• DSC가 열 劣한 것으로 나타났지만, FL의 이진화된 argmax 예측 결과는 다른 손실 함수와 비교해 다소 낮은 성능(예: WG의 경우 88–90% DSC)을 보였고, 이는 딱딱한 분할에서의 강건성을 시사함.
• 영역/경계 기반 손실 함수인 DB는 팬옵티크 분할 작업에 적합하지 않아, FROC 곡선이 평탄하고 검출 민감도가 낮아, 극단적으로 자신감 있는 예측을 유도함.
• 분포 기반 손실 함수(특히 FL)는 암묵적인 캘리브레이션 향상과 과적합 감소를 유도해, 단지 클래스 가중치나 아키텍처 설계 외에도 더 나은 진단 성능 기여함.