[논문 리뷰] Inconsistency-aware Uncertainty Estimation for Semi-supervised Medical Image Segmentation

Yinghuan Shi, Jian Zhang외 6명|arXiv (Cornell University)|2021. 10. 17.

COVID-19 diagnosis using AI인용 수 7

한 줄 요약

이 논문은 예측 불일치를 통해 다양한 오분류 비용 하에서 불확실성을 추정하는 방법을 사용하는 CoraNet을 제안한다. 이는 보수적-격렬 모듈을 통해 확실한 영역과 불확실한 영역를 식별한다. 이 방법은 CT 췌장, MR 심내막, ACDC 다중 구조 분할 작업에서 최신 기술을 초월하며, Dice 점수에서 기존 방법보다 최대 1.5% 향상되었고, 완전한 감독 학습 성능에 가까워졌다.

ABSTRACT

In semi-supervised medical image segmentation, most previous works draw on the common assumption that higher entropy means higher uncertainty. In this paper, we investigate a novel method of estimating uncertainty. We observe that, when assigned different misclassification costs in a certain degree, if the segmentation result of a pixel becomes inconsistent, this pixel shows a relative uncertainty in its segmentation. Therefore, we present a new semi-supervised segmentation model, namely, conservative-radical network (CoraNet in short) based on our uncertainty estimation and separate self-training strategy. In particular, our CoraNet model consists of three major components: a conservative-radical module (CRM), a certain region segmentation network (C-SN), and an uncertain region segmentation network (UC-SN) that could be alternatively trained in an end-to-end manner. We have extensively evaluated our method on various segmentation tasks with publicly available benchmark datasets, including CT pancreas, MR endocardium, and MR multi-structures segmentation on the ACDC dataset. Compared with the current state of the art, our CoraNet has demonstrated superior performance. In addition, we have also analyzed its connection with and difference from conventional methods of uncertainty estimation in semi-supervised medical image segmentation.

연구 동기 및 목표

• 기존의 소프트맥스 기반 신뢰도와 사전 정의된 임계값에 의존하는 준감독 의료 영상 분할에서의 전통적 불확실성 추정의 한계를 해결하기 위해.
• 비용 감수성 설정 하에서 예측이 일관되지 않은 영역를 식별함으로써 모델 초기화의 신뢰성을 향상시키기 위해.
• 확실한 영역와 불확실한 영역를 별도로 처리하는 자체 학습 전략을 개발하여 안정성을 향상시키기 위해.
• 비라벨 데이터의 확실성과 불확실성을 동시에 최적화하는 통합 모델의 엔드 투 엔드 학습을 구현하기 위해.
• 노이즈가 있는 가짜 라벨로 인한 오류 전파를 줄이며 벤치마크 데이터셋에서 뛰어난 분할 정확도를 달성하기 위해.

제안 방법

• 예측 불일치를 유도하기 위해 보수적(1:5)과 격렬적(5:1)의 서로 반대되는 오분류 비용을 가진 두 개의 하위 네트워크를 학습하는 보수적-격렬 모듈(CRM)을 도입한다.
• 두 모델 간에 일관된 예측을 보이는 영역는 확실한 영역로 분류하고, 불일치가 있는 영역는 불확실한 영역로 간주하여 동적 불확실성 맵을 구성한다.
• 모델는 세 가지 구성 요소로 이루어져 있다: 보수적-격렬 모듈(CRM), 확실한 영역 분할 네트워크(C-SN), 불확실한 영역 분할 네트워크(UC-SN)이며, 엔드 투 엔드로 학습된다.
• 별도의 자체 학습 전략을 적용한다: C-SN는 확실한 영역에서, UC-SN는 불확실한 영역에서 학습하여 오류 전파를 줄인다.
• 표준 데이터 증강(회전, 반전, 자르기)과 다중 클래스 분할 작업에서의 다수결 투표 및 신뢰도를 사용한다.
• 주요 평가 지표로 Dice 점수를 사용하며, 모든 실험에서 20%의 라벨 데이터와 80%의 비라벨 데이터를 사용한다.

실험 결과

연구 질문

• RQ1 다른 오분류 비용 하에서의 예측 불일치는 준감독 의료 영상 분할에서 불확실성 추정의 신뢰할 수 있는 대체 지표가 될 수 있는가?
• RQ2 확실한 영역와 불확실한 영역를 별도로 학습하는 전략이 통합 학습보다 분할 성능을 향상시키는가?
• RQ3 기존의 소프트맥스 기반 신뢰도와 비교해 본다면, 제안된 불일치 인식 기반의 불확실성 추정은 안정성과 정확도 면에서 어떻게 우월한가?
• RQ4 이 방법을 통해 훈련 중 노이즈가 있는 가짜 라벨로 인한 오류 전파를 어느 정도 줄일 수 있는가?
• RQ5 이러한 불확실성 인식 전략을 사용할 경우, 최소한의 라벨 데이터로 완전한 감독 학습 수준의 성능을 달성할 수 있는가?

주요 결과

• ACDC 데이터셋에서 CoraNet은 평균 Dice 점수 86.10% (±0.62)를 기록하여, 비교된 모든 베이스라인을 초월했고, 이전 최신 기술 방법보다 0.34% 높았다.
• 다중 구조 분할 작업에서 CoraNet은 심장근육에 대해 84.87%, 좌심실에 대해 90.74%, 우심실에 대해 82.69%의 Dice 점수를 달성했다.
• 불확실한 영역를 주 훈련 루프에서 분리함으로써 오류 전파를 줄여 더 안정적인 수렴을 이끌어냈다.
• 시각적 비교 결과, 특히 도전적인 영역에서 Mean Teacher보다 CoraNet이 더 선명하고 정확한 경계 예측을 생성했다.
• 모델 성능는 완전한 감독 기반 베이스라인(평균 Dice 점수 88.79%)과 일관되게 유사하여, 제한된 애너테이션 자료로도 강력한 일반화 능력을 보였다.
• 제거 분석 결과, 별도의 자체 학습 전략이 통합 학습 대비 성능 향상에 기여했음을 확인하여 설계 선택의 타당성을 입증했다.