[논문 리뷰] Self-Paced Contrastive Learning for Semi-supervised Medical Image Segmentation with Meta-labels

Jizong Peng, Ping Wang외 2명|arXiv (Cornell University)|2021. 07. 29.

Domain Adaptation and Few-Shot Learning인용 수 29

한 줄 요약

이 논문은 의료 영상의 메타-레이블을 이용한 자기 주도형 대조 학습 프레임워크를 도입하여 인코더를 사전 학습하고 반지도 학습(segmentation)을 규제하며, 라벨 달린 스캔이 극히 적은 경우에도 성능을 향상시킵니다.

ABSTRACT

Pre-training a recognition model with contrastive learning on a large dataset of unlabeled data has shown great potential to boost the performance of a downstream task, e.g., image classification. However, in domains such as medical imaging, collecting unlabeled data can be challenging and expensive. In this work, we propose to adapt contrastive learning to work with meta-label annotations, for improving the model's performance in medical image segmentation even when no additional unlabeled data is available. Meta-labels such as the location of a 2D slice in a 3D MRI scan or the type of device used, often come for free during the acquisition process. We use the meta-labels for pre-training the image encoder as well as to regularize a semi-supervised training, in which a reduced set of annotated data is used for training. Finally, to fully exploit the weak annotations, a self-paced learning approach is used to help the learning and discriminate useful labels from noise. Results on three different medical image segmentation datasets show that our approach: i) highly boosts the performance of a model trained on a few scans, ii) outperforms previous contrastive and semi-supervised approaches, and iii) reaches close to the performance of a model trained on the full data.

연구 동기 및 목표

• 라벨 데이터가 희소하고 비라벨 데이터가 제한적이거나 이용 불가능할 때 의료 영상 분할 성능 개선을 목표로 한다.
• 자유롭게 이용 가능한 메타-레이블(예: 슬라이스 위치, 환자 ID, 심장 단계)을 활용하여 표현 학습을 향상시킨다.
• 잡음이 있는 메타-레이블을 처리하고 인코더 표현과 라벨 기반 가중치를 함께 최적화하는 자기 주도형 대조 손실을 제안한다.
• 메타-레이블 사전 학습과 반지도 학습 및 Mean Teacher를 결합하면 분할 정확도가 더 향상됨을 보여준다.
• 제안된 접근 방식이 여러 의료 영상 데이터셋에서 기존의 대조 학습 및 반지도 방법보다 우수함을 보인다.

제안 방법

• 메타-레이블을 동일한 메타-레이블을 가진 증강 샘플 또는 같은 원본 이미지에서 나온 샘플을 묶어 약한 감독으로 사용하는 대조 손실을 정의한다.
• 레이블 노이즈를 완화하고 더 어려운 샘플을 점진적으로 포함하기 위해 양수 쌍에 적응 가중치를 부여하는 자기 주도 정규화기를 도입한다.
• 하드형 및 선형 SPL 정규화기에 대해 가중 양수를 위한 닫힌 형태 해를 제공하여 가중치 업데이트를 안내한다.
• 여러 메타-레이블 대조 손실을 하나의 목적 함수로 가중치를 두고 합치고, 프리트레이닝을 거친 후 제한된 라벨 데이터로 반지도 학습을 수행하도록 프레임워크를 확장하며(Mean Teacher를 이용한 일반화 손실 포함, SP Con 손실을 정규화로 사용).
• 대조 학습 임베딩을 위한 인코더-디코더(U-Net)와 프로젝션 헤드를 사용한다; 인코더를 사전 학습한 뒤 반지도 설정에서 제한된 라벨 데이터로 미세 조정한다.
• 고려를 위한 self-paced 속도 gamma를 에폭에 걸쳐 스케줄링하여 높은 신뢰도 쌍과 낮은 신뢰도 쌍의 포함을 균형 있게 조절한다.

실험 결과

연구 질문

• RQ1 메타-레이블이 의료 영상에서 활용되어 소수의 라벨로 분할 성능을 향상시키는 자기 주도형 대조 프레임워크가 효과적인가?
• RQ2 자기 주도 가중치가 대조 양수 쌍의 노이즈를 완화하고 하류 분할 표현을 향상시키는가?
• RQ3 사전 학습에 여러 메타-레이블을 결합하는 것이 최종 분할 성능에 어떤 영향을 미치는가?
• RQ4 대조 메타-레이블 학습을 반지도 학습(예: Mean Teacher)의 규제로 통합하면 전체 학습 수준에 근접한 성능 향상을 가져오는가?

주요 결과

• 메타-레이블이 있는 자기 주도 대조 학습이 한정된 라벨의 감독 학습보다 분할 성능을 개선한다.
• 메타-레이블 대조 손실로의 프리트레이닝이 데이터셋 전반에서 비지도 대조 프리트레이닝보다 우수하다.
• 여러 메타-레이블을 결합하면 단일 메타-레이블을 사용할 때보다 더 좋은 성능이 나온다.
• SP-Con 손실을 프리트레이닝 및 반지도 규제로 모두 도입하고, 필요 시 Mean Teacher를 적용하면 매우 적은 주석 스캔으로도 성능이 거의 전면 감독 학습 수준에 근접한다.
• 선형 자기 주도 정규화가 일반적으로 hard 정규화보다 더 강한 개선을 이끈다.
• 세 가지 데이터셋에서 제안한 방법이 낮은 라벨 설정에서 기존의 대조 및 반지도 방법보다 성능이 우수하다.