[논문 리뷰] 3D Semi-Supervised Learning with Uncertainty-Aware Multi-View Co-Training

Yingda Xia, Fengze Liu외 7명|arXiv (Cornell University)|2018. 11. 29.

AI in cancer detection인용 수 24

한 줄 요약

이 논문은 공간 변환을 통해 3D 의료 영상의 다양한 시야를 생성하고, 2D 사전 훈련된 모델에서 초기화된 비대칭 3D 커널을 사용하여 특징의 다양성을 유도하는 불확실성 인식 다중 시야 공동 훈련(UMCT)을 제안한다. 베이지안 딥 러닝을 통해 드롭아웃을 활용해 지식적 불확실성(epistemic uncertainty)을 추정하고, 이를 바탕으로 불확실성 가중 레이블 융합을 수행함으로써 가짜 레이블의 신뢰도를 향상시킨다. 이로 인해 레이블이 제한된 NIH 췌장 및 LiTS 간암 종양 세그먼테이션에서 최신 기술 수준의 성능을 달성한다.

ABSTRACT

While making a tremendous impact in various fields, deep neural networks usually require large amounts of labeled data for training which are expensive to collect in many applications, especially in the medical domain. Unlabeled data, on the other hand, is much more abundant. Semi-supervised learning techniques, such as co-training, could provide a powerful tool to leverage unlabeled data. In this paper, we propose a novel framework, uncertainty-aware multi-view co-training (UMCT), to address semi-supervised learning on 3D data, such as volumetric data from medical imaging. In our work, co-training is achieved by exploiting multi-viewpoint consistency of 3D data. We generate different views by rotating or permuting the 3D data and utilize asymmetrical 3D kernels to encourage diversified features in different sub-networks. In addition, we propose an uncertainty-weighted label fusion mechanism to estimate the reliability of each view's prediction with Bayesian deep learning. As one view requires the supervision from other views in co-training, our self-adaptive approach computes a confidence score for the prediction of each unlabeled sample in order to assign a reliable pseudo label. Thus, our approach can take advantage of unlabeled data during training. We show the effectiveness of our proposed semi-supervised method on several public datasets from medical image segmentation tasks (NIH pancreas & LiTS liver tumor dataset). Meanwhile, a fully-supervised method based on our approach achieved state-of-the-art performances on both the LiTS liver tumor segmentation and the Medical Segmentation Decathlon (MSD) challenge, demonstrating the robustness and value of our framework, even when fully supervised training is feasible.

연구 동기 및 목표

• 3D 의료 부피 데이터의 레이블링 비용이 높다는 점을 해결하기 위해, 풍부한 비라벨 데이터를 활용한다.
• 공간 변환을 통한 시야 다양성과 비대칭 3D 커널 설계를 도입하여 준지도 학습 3D 세그먼테이션 성능을 향상시킨다.
• 베이지안 딥 러닝을 활용해 예측 불확실성을 추정함으로써 공동 훈련에서 가짜 레이블 품질을 향상시킨다.
• 다중 시야 훈련을 통해 전체 레이블이 있는 설정에서도 성능 향상이 가능함을 보여준다.
• 2D 사전 훈련된 가중치가 부피 세그먼테이션 작업에서 3D 네트워크 초기화에 미치는 영향을 조사한다.

제안 방법

• 공간 변환(예: 회전, 순열)을 3D 데이터에 적용하여 공동 훈련을 위한 고유한 시야를 생성한다.
• 서브넷에서 비대칭 3D 컨볼루션 커널(예: 3×3×1)을 사용하고, 사전 훈련된 2D 모델에서 초기화하여 모델 수준의 시야 차이를 유도한다.
• 드롭아웃을 활용한 베이지안 딥 러닝을 적용하여 지식적 불확실성을 추정하고 각 시야의 예측에 대한 신뢰도 점수를 산출한다.
• 불확실성 가중 레이블 융합(ULF) 모듈은 예측의 추정된 불확실성에 따라 가중치를 적용하여 비라벨 데이터의 가짜 레이블을 계산한다.
• 다른 시야의 가짜 레이블과 진짜 레이블을 함께 사용하여 다중 시야 네트워크의 엔드 투 엔드 훈련 중 지도 신호를 제공한다.
• 이 프레임워크는 준지도 학습과 완전히 지도 학습 훈련 모두를 지원하며, 모든 데이터가 레이블이 있을 경우에도 공동 훈련 목표를 적용한다.

실험 결과

연구 질문

• RQ1 데이터 수준과 모델 수준의 시야 다양성을 도입한 다중 시야 공동 훈련이 3D 준지도 학습 세그먼테이션 성능을 향상시킬 수 있는가?
• RQ2 베이지안 딥 러닝을 통한 불확실성 추정이 3D 의료 영상 세그먼테이션의 공동 훈련에서 가짜 레이블 신뢰도를 향상시키는 데 어떤 영향을 미치는가?
• RQ3 2D 사전 훈련된 가중치가 부피 세그먼테이션 작업에서 3D 네트워크 초기화에 어떤 영향을 미치는가?
• RQ4 불확실성 가중 레이블 융합이 표준 평균화나 무게 없는 융합보다 다중 시야 공동 훈련에서 성능을 더 잘 향상시키는가?
• RQ5 제안된 공동 훈련 프레임워크는 레이블이 제한된 완전히 지도 학습 설정에서도 성능 향상을 이끌 수 있는가?

주요 결과

• NIH 췌장 데이터셋에서 제안된 UMCT 프레임워크는 라벨이 10%뿐인 조건에서 78.70%의 Dice 스코어를 기록하여 다른 준지도 학습 방법보다 뛰어난 성능을 보였다.
• LiTS 간암 종양 데이터셋에서 UMCT는 준지도 학습에서 최신 기술 수준의 성능을 달성하여 레이블이 제한된 조건에서도 강력한 일반화 능력을 입증했다.
• UMCT의 완전히 지도 학습 버전은 NIH 데이터셋에서 78.70%의 Dice 스코어를 기록하여 표준 3D ResNet-50 및 V-Net보다 뛰어난 성능을 보였으며, 이는 다중 시야 공동 훈련의 이점임을 시사한다.
• 제거 분석 결과, 2D 초기화된 비대칭 3D 커널은 대칭 아키텍처 대비 성능을 1% 이상 향상시켰다.
• 불확실성 가중 레이블 융합은 세 개 이상의 시야에서 성능 향상이著명하게 향상되었으며, 이는 신뢰도 추정이 신뢰할 수 없는 가짜 레이블의 노이즈를 줄이는 데 기여함을 보여준다.
• 초기화에 대한 제거 분석 결과, 2D 사전 훈련된 가중치는 훈련 안정성과 최종 성능을 크게 향상시켜 3D 세그먼테이션에서 그 중요성을 입증했다.