[논문 리뷰] Double-Uncertainty Weighted Method for Semi-supervised Learning

Yixin Wang, Yao Zhang외 5명|arXiv (Cornell University)|2020. 10. 19.

Advanced Neural Network Applications참고 문헌 13인용 수 12

한 줄 요약

이 논문은 교사-학생 프레임워크를 사용하여 반감독 의료 영상 분할을 위한 이중 불확실성 가중 방법을 제안한다. 베이지안 딥 러닝 환경에서 몬테 카를로 드롭아웃을 통해 분할 불확실성과 새로운 특징 불확실성을 추정함으로써, 일致성 손실을 동적으로 가중하여 교사 모델의 신뢰성과 학생의 학습을 향상시킨다. 두 개의 공개 데이터셋에서 최소한의 레이블 데이터로 최신 기술 수준의 성능을 달성한다.

ABSTRACT

Though deep learning has achieved advanced performance recently, it remains a challenging task in the field of medical imaging, as obtaining reliable labeled training data is time-consuming and expensive. In this paper, we propose a double-uncertainty weighted method for semi-supervised segmentation based on the teacher-student model. The teacher model provides guidance for the student model by penalizing their inconsistent prediction on both labeled and unlabeled data. We train the teacher model using Bayesian deep learning to obtain double-uncertainty, i.e. segmentation uncertainty and feature uncertainty. It is the first to extend segmentation uncertainty estimation to feature uncertainty, which reveals the capability to capture information among channels. A learnable uncertainty consistency loss is designed for the unsupervised learning process in an interactive manner between prediction and uncertainty. With no ground-truth for supervision, it can still incentivize more accurate teacher's predictions and facilitate the model to reduce uncertain estimations. Furthermore, our proposed double-uncertainty serves as a weight on each inconsistency penalty to balance and harmonize supervised and unsupervised training processes. We validate the proposed feature uncertainty and loss function through qualitative and quantitative analyses. Experimental results show that our method outperforms the state-of-the-art uncertainty-based semi-supervised methods on two public medical datasets.

연구 동기 및 목표

• 레이블이 부족한 의료 영상 데이터는 비용이 많이 들고 시간이 오래 걸리므로, 이를 해결하고자 한다.
• 훈련 과정에 불확실성 추정을 통합하여 교사-학생 반감독 학습의 신뢰성을 향상시키고자 한다.
• 합성 컨볼루션 커널 채널을 통해 불확실성을 캡처하는 새로운 특징 불확실성 지표를 제안하고 검증하고자 한다.
• 모델 예측과 상호작용하는 학습 가능한 불확실성 일致성 손실을 설계하여 지도 학습 없이도 교사 출력을 정밀하게 개선하고자 한다.
• 불확실성으로 가중치를 동적으로 조정하여 지도 학습과 비지도 학습 간의 균형을 조절하고자 한다.

제안 방법

• 교사 모델은 몬테 카를로 드롭아웃을 사용하여 베이지안 신경망을 근사함으로써, 분할 불확실성과 특징 불확실성을 추정할 수 있도록 훈련된다.
• 특징 불확실성은 컨볼루션 레이어의 채널(특히 4번째 업샘플링 레이어)을 기준으로 계산되어 각 채널의 특징 추출 신뢰도를 평가한다.
• 학습 가능한 불확실성 일치 손실이 도입되어, 불확실성 추정치에 가중된 일치성 기반으로 학생과 교사의 출력 간 불일치를 최소화함으로써 교사 예측을 반복적으로 개선한다.
• 총 손실은 레이블이 있는 데이터에 대한 지도 학습 교차 엔트로피 손실과 레이블이 있는 및 없는 데이터에 대한 불확실성 가중 일치 손실을 조합한다.
• 불확실성 추정치는 일치 손실의 가중치로 기능하여, 높은 불확실성 예측의 영향을 훈련 중 감소시킨다.
• 이 방법은 두 개의 공개 의료 영상 벤치마크인 2018 Atrial Segmentation Challenge와 MICCAI 2019 KiTS Challenge에서 평가된다.

실험 결과

연구 질문

• RQ1 합성 컨볼루션 커널 채널을 중심으로 한 특징 불확실성은 반감독 분할에서 모델의 해석 가능성과 성능을 향상시킬 수 있는가?
• RQ2 모델 예측과 상호작용하는 학습 가능한 불확실성 일치 손실은 지도 학습 없이도 교사 모델 정확도를 향상시킬 수 있는가?
• RQ3 불확실성 기반의 일치 손실 가중치는 반감독 학습에서 지도 학습과 비지도 학습 간 균형에 어떤 영향을 미치는가?
• RQ4 레이블 데이터가 부족한 상황에서 기존의 불확실성 기반 반감독 접근법보다 제안된 방법이 우수한 성능을 보일 수 있는가?
• RQ5 불확실성 추정은 신뢰할 수 없는 예측을 걸러내고 의료 영상 분할의 일반화 능력을 향상시키는 데 어느 정도 기여하는가?

주요 결과

• KiTS 2019 신장 분할 데이터셋에서, 이 방법은 단지 10장의 레이블이 있는 스캔만으로도 딱도 스코어 0.9206을 달성하여 모든 베이스라인을 앞서갔다.
• 32장의 레이블이 있는 스캔을 사용했을 때, 딱도 스코어는 0.9365에 도달하여 최신 기술 수준의 방법들보다 일관되게 향상된 성능을 보였다.
• 제안된 불확실성 일치 손실은 훈련 과정에서 교사 모델의 예측 불확실성을 감소시켰으며, 그 결과 그림 2(c)에서 평균 분할 불확실성 값이 감소함을 확인할 수 있었다.
• 특징 불확실성 분석 결과, 채널 간 불확실성 값이 훈련이 진행되면서 감소함을 확인하여 특징 추출의 신뢰도 향상을 시사했다.
• 매우 적은 레이블 샘플, 예를 들어 4장의 레이블 스캔에서도 성능 향상이 두드러졌으며, 이 경우 딱도 스코어에서 1위를 기록했다.
• 불확실성 가중 메커니즘이 UA-MT와 같은 기존의 베이스라인 방법의 성능을 향상시켜, 딱도 스코어에서 0.0087 향상된 성과를 보였다.