[논문 리뷰] Semi-supervised Left Atrium Segmentation with Mutual Consistency Training

Yicheng Wu, Minfeng Xu외 3명|arXiv (Cornell University)|2021. 03. 04.

Cardiac Valve Diseases and TreatmentsMedicine참고 문헌 18인용 수 20

한 줄 요약

이 논문은 상호 일致성 훈련을 통해 레이블이 없는 과도한 난이도의 영역을 활용함으로써 모델의 일반화 능력을 향상시키는 반감독 3D 좌심방 분할 방법인 MC-Net을 제안한다. 두 개의 약간 다른 디코더와 순환식 가짜 라벨 체계를 사용하여 저엔트로피이자 일관된 예측을 강제함으로써, MC-Net은 LA 데이터베이스에서 20%의 레이블 데이터로도 기존의 6개의 SOTA 방법을 능가하는 최고 성능을 달성한다.

ABSTRACT

Semi-supervised learning has attracted great attention in the field of machine learning, especially for medical image segmentation tasks, since it alleviates the heavy burden of collecting abundant densely annotated data for training. However, most of existing methods underestimate the importance of challenging regions (e.g. small branches or blurred edges) during training. We believe that these unlabeled regions may contain more crucial information to minimize the uncertainty prediction for the model and should be emphasized in the training process. Therefore, in this paper, we propose a novel Mutual Consistency Network (MC-Net) for semi-supervised left atrium segmentation from 3D MR images. Particularly, our MC-Net consists of one encoder and two slightly different decoders, and the prediction discrepancies of two decoders are transformed as an unsupervised loss by our designed cycled pseudo label scheme to encourage mutual consistency. Such mutual consistency encourages the two decoders to have consistent and low-entropy predictions and enables the model to gradually capture generalized features from these unlabeled challenging regions. We evaluate our MC-Net on the public Left Atrium (LA) database and it obtains impressive performance gains by exploiting the unlabeled data effectively. Our MC-Net outperforms six recent semi-supervised methods for left atrium segmentation, and sets the new state-of-the-art performance on the LA database.

연구 동기 및 목표

• 좌심방 분할에서 레이블이 부족한 의료 데이터 문제를 해결하기 위해 레이블이 없는 데이터를 효과적으로 활용한다.
• 레이블이 없는 상황에서 어려운, 모호한 영역(예: 작은 분지, 흐린 경계)을 강조함으로써 모델의 일반화 능력을 향상시킨다.
• 지식적 불확실성( epistemic uncertainty )을 감독 신호로 활용하여 과도한 난이도의 영역에서의 모델 불확실성을 감소시킨다.
• 기존 방법을 능가하는 반감독 프레임워크를 개발하여 공개된 LA 데이터베이스에서 성능을 향상시킨다.
• 두 디코더 간의 상호 일치성이 레이블이 없는 고불확실성 영역에서 특징 학습을 향상시킬 수 있는가를 입증한다.

제안 방법

• MC-Net은 동일한 입력에 대해 다양한 예측을 생성하기 위해 공유 인코더와 두 개의 약간 다른 디코더를 사용한다.
• 두 디코더 간의 예측 차이를 이용해 불확실성을 추정하고 훈련을 이끌어낸다.
• 순환식 가짜 라벨 체계는 예측 불일치를 비지도 손실로 변환하여 상호 일치성을 강제한다.
• 예측 불확실성을 최소화하기 위해 날카기 함수를 통해 소프트 가짜 라벨을 생성하고 엔트로피 정규화를 적용한다.
• 두 디코더의 출력 간 격차를 최소화함으로써 일致성 정규화를 적용하여 저엔트로피이자 안정적인 예측을 촉진한다.
• 프레임워크는 레이블이 있는 데이터와 레이블이 없는 데이터를 동시에 사용하여 엔드 투 엔드로 훈련되며, 순환식 가짜 라벨 메커니즘은 예측을 동적으로 개선한다.

실험 결과

연구 질문

• RQ1 두 디코더 간의 상호 일치성 훈련이 반감독식 좌심방 분할에서 일반화 능력을 향상시키는가?
• RQ2 이전에는 간과되었던 레이블이 없는 난이도 높은 영역을 효과적으로 활용하여 모델의 불확실성을 줄일 수 있는가?
• RQ3 디코더 간 일치성을 강제하는 순환식 가짜 라벨 체계가 기존의 일치성 정규화보다 더 나은 분할 성능을 내는가?
• RQ4 MC-Net이 LA 데이터베이스에서 레이블이 20%뿐인 경우에도 최고 성능을 달성할 수 있는가?
• RQ5 동일한 디코더 대비 서로 다른 디코더를 사용한 모델 다양성이 성능에 어떤 영향을 미치는가?

주요 결과

• MC-Net은 LA 데이터베이스에서 레이블이 20%뿐인 경우에도 90.34%의 Dice 스코어를 달성하여 100% 레이블 데이터를 사용한 완전히 감독된 V-Net 베이스라인(91.14%)을 능가한다.
• 20% 레이블 데이터로 MC-Net은 95% Hausdorff 거리(95HD) 6.00 박자와 평균 표면 거리(ASD) 1.77 박자를 기록하여, 기존의 6개의 SOTA 방법을 모두 능가한다.
• 절단 실험을 통해 두 개의 다른 디코더(V2d-Net)를 사용할 경우 동일한 디코더(V2-Net)를 사용할 경우보다 성능 향상이 확인되었으며, 레이블이 10%일 때 Dice 스코어가 1.18% 향상되었다.
• 순환식 가짜 라벨(CPL) 구성 요소를 추가함으로써 20% 레이블 데이터 조건에서 95HD는 9.36에서 6.00 박자로, ASD는 2.18에서 1.77 박자로 감소하여 그 효과가 입증되었다.
• 특히 어려운 영역(예: 작은 분지)에서 후처리 없이도 더 완전한 좌심방 분할 결과를 도출하며 고립된 영역이 적은 편이다.
• 이 방법은 다른 형태 제약 모델과 호환되며, 다른 반감독 의료 영상 분할 작업의 성능 향상에도 확장 가능하다.