[논문 리뷰] Local contrastive loss with pseudo-label based self-training for semi-supervised medical image segmentation

Krishna Chaitanya, Ertunç Erdil외 2명|arXiv (Cornell University)|2021. 12. 17.

Advanced Neural Network Applications인용 수 7

한 줄 요약

이 논문은 레이블이 있는 데이터와 레이블이 없는 데이터를 모두 활용하여, 무작위 증강이 아닌 의미적 가짜 레이블을 기반으로 하여 양성 쌍을 정의함으로써, 반독립적인 의료 영상 분할을 위한 새로운 局소 대비 손실을 제안한다. 이 방법은 단 한 개 또는 두 개의 3D 레이블이 있는 볼륨만을 사용하여 심장 및 전립선 MRI 데이터셋에서 최신 기술 수준의 성능을 달성하며, 기존의 반독립적 및 대비 학습 방법들을 크게 능가한다.

ABSTRACT

Supervised deep learning-based methods yield accurate results for medical image segmentation. However, they require large labeled datasets for this, and obtaining them is a laborious task that requires clinical expertise. Semi/self-supervised learning-based approaches address this limitation by exploiting unlabeled data along with limited annotated data. Recent self-supervised learning methods use contrastive loss to learn good global level representations from unlabeled images and achieve high performance in classification tasks on popular natural image datasets like ImageNet. In pixel-level prediction tasks such as segmentation, it is crucial to also learn good local level representations along with global representations to achieve better accuracy. However, the impact of the existing local contrastive loss-based methods remains limited for learning good local representations because similar and dissimilar local regions are defined based on random augmentations and spatial proximity; not based on the semantic label of local regions due to lack of large-scale expert annotations in the semi/self-supervised setting. In this paper, we propose a local contrastive loss to learn good pixel level features useful for segmentation by exploiting semantic label information obtained from pseudo-labels of unlabeled images alongside limited annotated images. In particular, we define the proposed loss to encourage similar representations for the pixels that have the same pseudo-label/ label while being dissimilar to the representation of pixels with different pseudo-label/label in the dataset. We perform pseudo-label based self-training and train the network by jointly optimizing the proposed contrastive loss on both labeled and unlabeled sets and segmentation loss on only the limited labeled set. We evaluated on three public cardiac and prostate datasets, and obtain high segmentation performance.

연구 동기 및 목표

• 큰 레이블 데이터셋에 대한 의존도를 줄임으로써 의료 영상 분할의 높은 레이블링 비용 문제를 해결한다.
• 가짜 레이블에서 유래한 의미적 레이블 정보를 통합하여 반독립적 분할에서 국소 수준의 표현 학습을 향상시킨다.
• 기존의 국소 대비 방법이 비지도 학습 대체 레이블이나 공간적 근접성에 의존하는 한계를 극복한다.
• 제한된 레이블 데이터에 대한 분할 손실과 레이블이 있는 데이터와 없는 데이터에 대한 대비 손실을 동시에 최적화하는 공동 학습 프레임워크를 개발한다.
• 반복적인 가짜 레이블 정밀화와 일致성 정규화를 통해 모델의 일반화 능력과 강인성을 향상시킨다.

제안 방법

• 같은 가짜 레이블을 가진 픽셀을 양성 쌍으로 간주하고, 다른 가짜 레이블을 가진 픽셀을 음성 쌍으로 간주하는 국소 대비 손실을 정의한다.
• 학생 네트워크를 사용하여 레이블이 없는 이미지의 가짜 레이블을 생성하고, 학습 도중 반복적으로 정밀화한다.
• 레이블이 있는 데이터셋에 대해서만 표준 분할 손실을 최적화하고, 레이블이 있는 데이터셋과 없는 데이터셋에 모두 적용되는 제안된 대비 손실을 함께 최적화함으로써 자기 학습을 수행한다.
• 동일한 이미지의 증강된 시각 간의 특징 비교를 위해 내부 이미지 및 외부 이미지 표현 매칭 기법을 사용한다.
• 변환된 시각 간의 겹침 임계값을 기반으로 가짜 레이블을 필터링하여 일치성 정규화를 적용함으로써 예측 품질을 향상시킨다.
• 무작위 가중치 초기화 또는 자기지도 사전학습에서 확보한 사전학습된 가중치를 사용하여 모델을 초기화한다.

실험 결과

연구 질문

• RQ1 무작위 증강이나 공간적 근접성 대비로 가짜 레이블 기반의 의미적 정보가 반독립적 의료 영상 분할에서 국소 표현 학습을 향상시키는가?
• RQ2 레이블이 있는 데이터와 없는 데이터에 모두 적용되는 대비 손실을 공동 최적화하면 기존 자기 학습 방식보다 더 나은 분할 성능을 낼 수 있는가?
• RQ3 정답 레이블이 없는 상황에서 가짜 레이블의 품질이 대비 손실 성능에 미치는 영향은 어떠한가?
• RQ4 내부 이미지 대비와 외부 이미지 대비와 같은 다양한 표현 매칭 전략이 모델 수렴과 정확도에 미치는 영향은 무엇인가?
• RQ5 레이블이 없는 데이터에서 사전학습을 수행하면 제안된 대비 자기 학습 프레임워크의 성능을 추가로 향상시킬 수 있는가?

주요 결과

• 제안된 방법은 단 한 개의 3D 레이블 볼륨만을 사용하여 ACDC 데이터셋에서 평균 DSC 0.881을 달성하며, 최신 기술 수준의 반독립적 방법들을 능가한다.
• 여덟 개의 레이블 볼륨을 사용할 경우, 내부 이미지 매칭에서는 DSC 0.885, 외부 이미지 매칭에서는 DSC 0.897를 기록하여 강건성과 확장성을 입증한다.
• 가짜 레이블을 딱딱한 임계값(0.7 또는 0.8)으로 필터링하는 것은 모든 가짜 레이블을 사용하는 경우(임계값 = 0)보다 성능 향상이 없음을 시사하며, 이 설정에서는 고신뢰도 필터링이 유익하지 않음을 의미한다.
• 두 개의 레이블 볼륨을 사용할 경우 외부 이미지 매칭이 내부 이미지 매칭보다 더 우수한 성능을 보이며, 이는 환자 간 정렬이 일반화 능력을 향상시킨다는 것을 시사한다.
• 특히 한두 개의 레이블 볼륨을 사용하는 저샷 설정에서 동시대의 대비 학습 및 반독립적 학습 기반 모델들보다 상당한 성능 향상을 달성한다.
• 레이블이 없는 데이터에서의 사전학습과 제안된 대비 자기 학습 프레임워크의 조합은 추가적인 성능 향상을 가져오며, 사전학습과 자기지도 학습의 상호보완성을 확인한다.