[논문 리뷰] Self-Ensembling Contrastive Learning for Semi-Supervised Medical Image Segmentation

Jinxi Xiang, Zhuowei Li외 3명|arXiv (Cornell University)|2021. 05. 27.

Domain Adaptation and Few-Shot Learning참고 문헌 19인용 수 11

한 줄 요약

이 논문은 자기-앙상블링 대비 학습(SeCL)을 제안하며, 이는 레이블이 없는 데이터를 사용하여 이미지 수준의 감독 손실과 특징 수준의 대비 손실을 결합하는 준지도 학습 의료 영상 분할 방법이다. 학생-교수 네트워크와 지수 이동 평균 업데이트를 활용하고, 해부학적 인식 대비 샘플링을 적용함으로써, 제한된 레이블을 가진 MRI 및 CT 분할 벤치마크에서 최신 기술 수준(SOTA) 성능을 달성하며, 특징의 밀도와 분리도를 크게 향상시킨다.

ABSTRACT

Deep learning has demonstrated significant improvements in medical image segmentation using a sufficiently large amount of training data with manual labels. Acquiring well-representative labels requires expert knowledge and exhaustive labors. In this paper, we aim to boost the performance of semi-supervised learning for medical image segmentation with limited labels using a self-ensembling contrastive learning technique. To this end, we propose to train an encoder-decoder network at image-level with small amounts of labeled images, and more importantly, we learn latent representations directly at feature-level by imposing contrastive loss on unlabeled images. This method strengthens intra-class compactness and inter-class separability, so as to get a better pixel classifier. Moreover, we devise a student encoder for online learning and an exponential moving average version of it, called teacher encoder, to improve the performance iteratively in a self-ensembling manner. To construct contrastive samples with unlabeled images, two sampling strategies that exploit structure similarity across medical images and utilize pseudo-labels for construction, termed region-aware and anatomical-aware contrastive sampling, are investigated. We conduct extensive experiments on an MRI and a CT segmentation dataset and demonstrate that in a limited label setting, the proposed method achieves state-of-the-art performance. Moreover, the anatomical-aware strategy that prepares contrastive samples on-the-fly using pseudo-labels realizes better contrastive regularization on feature representations.

연구 동기 및 목표

• 제한된 레이블 데이터를 가진 준지도 학습 의료 영상 분할 성능을 향상시키기 위해.
• 레이블이 없는 이미지에 대비 손실을 적용하여 특징 수준의 표현 학습을 향상시키기 위해.
• 해부학적 및 구조적 사전 지식을 활용하여 의료 영상 분할에서 의미 있는 양성 및 음성 샘플을 구성하는 과제를 해결하기 위해.
• 학생-교수 프레임워크를 통한 자기-앙상블링을 통합하여 표현 학습의 안정성과 향상을 높이기 위해.
• 실제 의료 영상 환경에서 지역 인식 및 해부학적 인식 대비 샘플링 전략의 효과를 검증하기 위해.

제안 방법

• 이중으로 최적화하는 자기-앙상블링 대비 학습 프레임워크를 제안하며, 이미지 수준의 감독 손실과 특징 수준의 대비 손실을 함께 최적화한다.
• 학생-교수 인코더 아키텍처를 사용하며, 교수 네트워크는 학생의 가중치를 지수 이동 평균(EMA)으로 업데이트하여 학습 안정성을 높인다.
• 레이블이 없는 이미지의 증강된 시각을 기반으로 대비 손실을 적용하여 특징 공간에서 동일 클래스의 밀도와 이질 클래스의 분리도를 강화한다.
• 두 가지 대비 샘플링 전략을 도입한다: 구조 유사성 기반의 지역 인식 전략과 공간 가이드를 위한 의사 레이블을 사용하는 해부학적 인식 전략.
• t-SNE 시각화를 통해 해부학적 인식 전략 하에서 특징 표현의 군집화가 향상되었음을 확인한다.
• 대비 손실 계수 λ에 대해 램프업 스케줄링을 적용하여 학습 중에 점차 영향력을 증가시킨다.

실험 결과

연구 질문

• RQ1 이미지 수준의 감독과 결합된 특징 수준의 대비 학습이 준지도 학습 의료 영상 분할 성능을 향상시킬 수 있는가?
• RQ2 학생-교수 프레임워크를 통한 자기-앙상블링이 낮은 감독 환경에서 표현 학습을 어떻게 향상시키는가?
• RQ3 지역 인식 전략과 해부학적 인식 전략 중 어느 것이 의료 영상에서 더 나은 특징 분리도를 제공하는가?
• RQ4 실시간 의사 레이블 기반 샘플링이 구조 기반 샘플링보다 더 강건하고 일반화된 표현을 생성하는가?
• RQ5 대비 정규화가 기존의 일致성 기반 준지도 학습 방법을 보완하여 분할 정확도를 향상시킬 수 있는가?

주요 결과

• SeCL은 제한된 레이블 조건 하에서 MRI 및 CT 분할 데이터셋에서 최신 기술 수준 성능을 달성하며, 해부학적 인식 전략을 사용한 MMWHS-2017 데이터셋에서 LV, RV, AO의 Dice 스코어는 각각 0.8535, 0.7708, 0.8667을 기록한다.
• 해부학적 인식 대비 샘플링 전략이 지역 인식 샘플링 전략보다 우수하며, LA 구조에서 3.5% 높은 Dice 스코어를, WHS 전체에서 1.5% 높은 스코어를 기록한다.
• 자기-앙상블링 메커니즘에서 α=0.9가 최적의 균형을 이룬다: α=0.99는 적응 속도가 느리며, α=0은 학습이 불안정하고 성능이 열악하다.
• t-SNE 시각화 결과, 해부학적 인식 샘플링을 사용한 SeCL이 해부학적 구조에 걸쳐 더 밀도 있고 잘 분리된 특징 군집을 생성함을 확인한다.
• 대비 손실은 특징 표현 품질을 크게 향상시키며, 더 나은 군집화와 낮은 HD(Hausdorff 거리) 값으로 증명된다 — SeCL(AACS)는 LV에서 3.708 mm의 HD를 기록하여 Mean Teacher(6.467 mm)를 초월한다.
• 제거 실험 결과, 특징 수준의 대비 학습이 필수적임을 입증한다: 이를 제거하면 성능 저하가 발생하며, 학생-교수 메커니즘은 학습 안정성과 수렴 성능을 향상시킨다.