[논문 리뷰] Semi-Supervised Brain Lesion Segmentation with an Adapted Mean Teacher Model

Wenhui Cui, Yi Liu외 7명|arXiv (Cornell University)|2019. 03. 04.

Brain Tumor Detection and Classification참고 문헌 16인용 수 11

한 줄 요약

이 논문은 정규화된 MRI 데이터와 비정규화된 MRI 데이터를 모두 활용하는 개선된 Mean Teacher 모델을 사용한 준지도 학습 뇌 병변 세그멘테이션 방법을 제안한다. 동일한 비정규화된 스캔에 대한 노이즈가 섞인 증강 데이터 간의 세그멘테이션 일관성과 지수 이동 평균을 통한 학생-교수 모델 업데이트를 통한 자체 앙상블 전략을 적용함으로써, 최신 기술 수준의 성능을 달성하였으며, 기존의 준지도 학습 방법을 뛰어넘고, 단지 10개의 정규화된 스캔만을 사용할 경우 완전히 지도 학습 기반의 베이스라인을 초월하는 결과를 보였다.

ABSTRACT

Automated brain lesion segmentation provides valuable information for the analysis and intervention of patients. In particular, methods based on convolutional neural networks (CNNs) have achieved state-of-the-art segmentation performance. However, CNNs usually require a decent amount of annotated data, which may be costly and time-consuming to obtain. Since unannotated data is generally abundant, it is desirable to use unannotated data to improve the segmentation performance for CNNs when limited annotated data is available. In this work, we propose a semi-supervised learning (SSL) approach to brain lesion segmentation, where unannotated data is incorporated into the training of CNNs. We adapt the mean teacher model, which is originally developed for SSL-based image classification, for brain lesion segmentation. Assuming that the network should produce consistent outputs for similar inputs, a loss of segmentation consistency is designed and integrated into a self-ensembling framework. Specifically, we build a student model and a teacher model, which share the same CNN architecture for segmentation. The student and teacher models are updated alternately. At each step, the student model learns from the teacher model by minimizing the weighted sum of the segmentation loss computed from annotated data and the segmentation consistency loss between the teacher and student models computed from unannotated data. Then, the teacher model is updated by combining the updated student model with the historical information of teacher models using an exponential moving average strategy. For demonstration, the proposed approach was evaluated on ischemic stroke lesion segmentation, where it improves stroke lesion segmentation with the incorporation of unannotated data.

연구 동기 및 목표

• 의료 영상 세그멘테이션 분야에서 정규화된 뇌 MRI 데이터의 높은 비용과 부족함을 해결하기 위해.
• 제한된 정규화된 데이터가 존재할 경우 세그멘테이션 성능을 향상시키기 위해.
• 원래 이미지 분류를 위한 목적을 가진 Mean Teacher 프레임워크를 준지도 학습 뇌 병변 세그멘테이션에 적응시키기 위해.
• 일관성 정규화와 자체 앙상블 전략을 통해 비정규화된 데이터를 활용하여 모델의 강건성과 일반화 능력을 향상시키기 위해.
• 실제 임상 MRI 데이터를 사용하여 뇌 경색 병변 세그멘테이션에 대해 제안된 방법의 효과성을 입증하기 위해.

제안 방법

• 세그멘테이션을 위한 DeepMedic CNN 아키텍처를 기반으로 한 학생-교수 모델 아키텍처를 채택하였다.
• 학습 중인 학생 모델은 정규화된 데이터에 대한 교차 엔트로피 손실과 비정규화된 데이터에 대한 Dice 기반 일관성 손실의 가중합을 최소화하도록 한다.
• 비정규화된 샘플은 가우시안 노이즈를 통해 증강되어 두 개의 유사한 입력을 생성하며, 모델은 이들 간의 예측 일관성을 강제한다.
• 교수 모델은 학생 모델의 가중치에 대한 지수 이동 평균(EMA)을 사용하여 업데이트되며, 이는 이전 정보를 유지하는 데 기여한다.
• 학생 및 교수 모델은 교대로 학습되며, 교수 모델은 학생 모델의 안정적인 타겟을 제공한다.
• 최종 세그멘테이션은 학습된 교수 모델을 사용하여 수행되며, 이는 여러 학습 단계 동안의 지식을 통합한다.

실험 결과

연구 질문

• RQ1 Mean Teacher 프레임워크는 MRI 영상에서 준지도 학습 뇌 병변 세그멘테이션에 효과적으로 적응될 수 있는가?
• RQ2 비정규화된 데이터의 노이즈가 섞인 증강 데이터에 대한 예측 일관성 강화가 세그멘테이션 성능 향상에 기여하는가?
• RQ3 학생-교수 업데이트 동역학을 통한 자체 앙상블 전략은 기존 준지도 학습 방법 대비 모델의 일반화 능력을 향상시키는가?
• RQ4 정규화된 스캔 수가 매우 적을 경우 제안된 방법의 성능는 어떠한가?
• RQ5 자원이 제한된 조건에서, 제안된 방법이 완전히 지도 학습 기반의 베이스라인과 기존의 준지도 학습 방법을 모두 능가하는가?

주요 결과

• 제안된 방법은 테스트 세트에서 평균 Dice 스코어 0.6518 ± 0.2484를 기록하여, 20개의 정규화된 스캔으로 학습된 베이스라인 DeepMedic 모델보다 유의미하게 뛰어난 성능을 보였다.
• 정규화된 스캔이 단지 10개일 경우에도, 제안된 방법은 30개의 정규화된 스캔으로 학습된 베이스라인 모델과 유사한 Dice 스코어를 달성하였다.
• 모든 데이터 설정에서 경쟁 기반 준지도 학습 방법인 DeepMedic-EN 및 DeepMedic-UDA를 모두 능가하였으며, 특히 학습 데이터가 제한적일 경우 두드러진 성능 향상을 보였다.
• 동일한 비정규화된 스캔에 대한 노이즈가 섞인 증강 데이터 간의 Dice 계수 기반 일관성 손실이 효과적인 정규화를 제공하였다.
• EMA 기반의 교수 모델 업데이트 전략은 학습을 안정화시키고, 특히 자료가 적은 환경에서 모델의 강건성을 향상시켰다.
• 비정규화된 데이터를 통합함으로써 모델의 일반화 능력이 뚜렷이 향상되었으며, 정규화된 예제가 최소한일지라도 Dice 스코어가 크게 향상되었다.