[논문 리뷰] Contrastive learning of global and local features for medical image segmentation with limited annotations

Krishna Chaitanya, Ertunç Erdil외 2명|arXiv (Cornell University)|2020. 06. 18.

Radiomics and Machine Learning in Medical Imaging인용 수 241

한 줄 요약

논문은 의료 영상 분할을 위한 도메인- 및 문제 특화 대비 학습으로 분할 성능을 향상시키고, 로컬 대조 손실을 추가하며 부피 구조를 활용해 주석이 적은 상황에서 성능을 개선합니다.

ABSTRACT

A key requirement for the success of supervised deep learning is a large labeled dataset - a condition that is difficult to meet in medical image analysis. Self-supervised learning (SSL) can help in this regard by providing a strategy to pre-train a neural network with unlabeled data, followed by fine-tuning for a downstream task with limited annotations. Contrastive learning, a particular variant of SSL, is a powerful technique for learning image-level representations. In this work, we propose strategies for extending the contrastive learning framework for segmentation of volumetric medical images in the semi-supervised setting with limited annotations, by leveraging domain-specific and problem-specific cues. Specifically, we propose (1) novel contrasting strategies that leverage structural similarity across volumetric medical images (domain-specific cue) and (2) a local version of the contrastive loss to learn distinctive representations of local regions that are useful for per-pixel segmentation (problem-specific cue). We carry out an extensive evaluation on three Magnetic Resonance Imaging (MRI) datasets. In the limited annotation setting, the proposed method yields substantial improvements compared to other self-supervision and semi-supervised learning techniques. When combined with a simple data augmentation technique, the proposed method reaches within 8% of benchmark performance using only two labeled MRI volumes for training, corresponding to only 4% (for ACDC) of the training data used to train the benchmark. The code is made public at https://github.com/krishnabits001/domain_specific_cl.

연구 동기 및 목표

• 의료 영상 분할에서 라벨링 데이터의 필요를 줄이기 위해 자기 감독 학습을 활용하는 동기를 제시한다.
• 부피형 의료 영상에 맞춘 글로벌 및 로컬 대조 학습 확장을 도입한다.
• MRI/CT의 도메인 특화 신호 및 로컬 영역 구분을 통해 픽셀 단위 예측을 개선한다.
• 세 가지 MRI 데이터셋에서 접근 방식을 평가하고 전처리 작업, 데이터 증강 및 반지도 학습 방법과의 비교를 수행한다.

제안 방법

• 전역 대조 손실을 확장해 도메인 특화 유사성 신호로 교차 볼륨 슬라이스 정렬을 활용한다.
• 디코더 특징 맵에서 작동하는 로컬 대조 손실을 도입해 분할을 위한 영역 수준 구분을 촉진한다.
• 두 가지 도메인 정보 기반 글로벌 대조 전략(G^D-와 G^D)과 두 가지 로컬 전략(L^R 및 L^D)을 정의해 양성/음성 쌍을 구성한다.
• 글로벌 손실로 인코더를 사전 학습한 뒤 이를 고정하고 로컬 손실로 디코더 블록을 먼저 사전 학습한 다음 분할을 위한 전체 미세 조정을 수행한다.
• 공동 손실 가중치가 필요 없는 단계적 파이프라인으로 글로벌 및 로컬 사전 학습을 결합한다.

실험 결과

연구 질문

• RQ1 도메인 특화 글로벌 신호가 정렬된 볼륨에서 분할을 위한 대조 학습을 향상시킬 수 있는가?
• RQ2 디코더 특징에 로컬 대조 손실을 추가하는 것이 글로벌만 사용하는 경우보다 픽셀별 표현을 더 잘 만들 수 있는가?
• RQ3 도메인 정보 기반 대조 전략은 제한 주석 체제에서 무작위 증강이나 사전 과제와 비교해 어떤 차이가 있는가?
• RQ4 제안된 사전 학습은 데이터 증강 및 반지도 학습 방법과 상호 보완적인가?

주요 결과

• 볼륨 간 슬라이스 유사성을 활용하는 글로벌 전략(G^D)이 무작위 증강 기반 기준을 능가한다.
• 로컬 디코더 사전 학습(L^R 또는 L^D)은 인코더 전용 사전 학습보다 추가 이득을 제공한다.
• 제안된 조합(G^D + L^R)은 1–2개의 학습 볼륨으로도 ACDC, Prostate, MMWHS 데이터셋 전반에서 상당한 개선을 보인다.
• 이 방법은 두 개의 주석이 달린 볼륨만으로 벤치마크 성능에 근접하고, 데이터 증강이나 반지도 학습 방법과 결합했을 때 이득이 유지된다.