[논문 리뷰] Bootstrapping Semi-supervised Medical Image Segmentation with Anatomical-aware Contrastive Distillation

Chenyu You, Weicheng Dai외 3명|arXiv (Cornell University)|2022. 06. 06.

Medical Imaging and Analysis인용 수 7

한 줄 요약

ACTION는 부족한 라벨 데이터에서의 클래스 불균형 문제를 해결하기 위해 부드러운 음성 샘플 라벨링과 활성 샘플링을 통한 경계선 음성 히트를 활용한 해부학적 인지 대비 분산 프레임워크를 제안한다. 이는 ACDC 및 LiTS 데이터셋에서 최신 기술을 초월하며, 라벨이 3%인 경우 ACDC에서 87.5%의 Dice 스코어를 달성한다.

ABSTRACT

Contrastive learning has shown great promise over annotation scarcity problems in the context of medical image segmentation. Existing approaches typically assume a balanced class distribution for both labeled and unlabeled medical images. However, medical image data in reality is commonly imbalanced (i.e., multi-class label imbalance), which naturally yields blurry contours and usually incorrectly labels rare objects. Moreover, it remains unclear whether all negative samples are equally negative. In this work, we present ACTION, an Anatomical-aware ConTrastive dIstillatiON framework, for semi-supervised medical image segmentation. Specifically, we first develop an iterative contrastive distillation algorithm by softly labeling the negatives rather than binary supervision between positive and negative pairs. We also capture more semantically similar features from the randomly chosen negative set compared to the positives to enforce the diversity of the sampled data. Second, we raise a more important question: Can we really handle imbalanced samples to yield better performance? Hence, the key innovation in ACTION is to learn global semantic relationship across the entire dataset and local anatomical features among the neighbouring pixels with minimal additional memory footprint. During the training, we introduce anatomical contrast by actively sampling a sparse set of hard negative pixels, which can generate smoother segmentation boundaries and more accurate predictions. Extensive experiments across two benchmark datasets and different unlabeled settings show that ACTION significantly outperforms the current state-of-the-art semi-supervised methods.

연구 동기 및 목표

• 왜곡된 데이터 분포로 인해 희귀한 해부학적 구조가 자주 잘못 분류되는 의료 영상 분할의 클래스 불균형 문제를 해결한다.
• 모든 음성 샘플을 동일하게 취급하는 표준 대비 학습의 한계를 극복한다. 이는 잘못된 음성 샘플을 유발할 수 있다.
• 지속적인 전역 표현과 지역 해부학적 특징을 유지하면서 균형 잡힌 표현을 학습함으로써 분할 정확도를 향상시킨다.
• 학습 중에 희소한 쿼리-키 쌍을 사용하는 활성 샘플링을 통해 메모리 및 계산 오버헤드를 줄인다.
• 낮은 감독 설정 하에서 기존 반감성 방법을 능가하는 강력하고 확장 가능한 프레임워크를 개발한다.

제안 방법

• 메모리 백에 저장된 해부학적으로 유사한 앵커를 기반으로 교사 모델의 확률 분포를 사용해 음성 샘플을 소프트 라벨링하는 반복적 대비 분산 파이프라인을 도입한다.
• 스테이지 II에서 국소 대비 분산을 적용하여 의미적으로 유사하지만 동일하지 않은 음성 예제를 샘플링함으로써 특징의 다양성을 향상시킨다.
• 희소하고 어려운 음성 픽셀을 활성 샘플링하여, 쿼리 특징을 해당 클래스의 평균(양성 키)에 가깝게 하고 다른 클래스(음성 키)에서 멀어지게 하는 AnCo(해부학적 대비) 손실을 설계한다.
• 훈련 안정성과 일반화 성능 향상을 위해 교사 모델에는 약한 증강을, 학생 모델에는 강한 증강을 적용한다.
• 스테이지 III에서 신뢰도 임계값을 적용한 가짜 라벨에 대해 교차 엔트로피 손실을 통합하여 예측을 정밀하게 보정하고 성능을 향상시킨다.
• 전체 특징 저장을 피하고 동적 샘플링 및 모멘터리 기반 타겟 네트워크 업데이트에 의존함으로써 메모리 사용량을 최소화한다.

실험 결과

연구 질문

• RQ1 음성 샘플의 소프트 라벨링이 반감성 의료 영상 분할에서 표현 학습에 기여하는가?
• RQ2 해부학적으로 어려운 음성 픽셀을 활성 샘플링하면 불균형 데이터셋에서 더 선명한 분할 경계와 더 높은 성능을 달성하는가?
• RQ3 전역적 의미 관계와 국소 해부학적 특징을 모두 모델링하는 대비 분산 프레임워크가 최소한의 메모리 오버헤드로 뛰어난 성능을 낼 수 있는가?
• RQ4 제안된 해부학적 대비 손실(AnCo)이 의료 영상에서 클래스 불균형 문제를 다룰 때 표준 대비 학습보다 우월한가?
• RQ5 소프트 라벨링, 국소 분산, 해부학적 대비의 조합이 낮은 라벨 비율 환경에서 분할 정확도를 얼마나 향상시키는가?

주요 결과

• ACTION는 라벨이 3%인 ACDC 데이터셋에서 87.5%의 Dice 스코어를 달성하여, 감독 기반 베이스라인(89.2% Dice)과 모든 최신 반감성 SSL 방법을 뛰어넘는다.
• LiTS 데이터셋에서는 5% 및 10% 라벨 설정 모두에서 최고 성능을 기록하며, 이는 이전 방법들보다 크게 앞서는 성능을 보인다.
• 제거 실험 결과, 무작위 샘플링, 국소 분산, 해부학적 대비 등의 모듈을 제거할 경우 성능 저하가 심각하게 발생함을 확인하여 각 모듈의 필수성을 입증한다.
• 교사에 약한 증강, 학생에 강한 증강을 적용한 조합이 가장 높은 성능을 보이며, 안정성과 다양성 사이의 최적의 균형을 이루는 것으로 나타났다.
• ACTION는 RV 및 Myo와 같은 소형 및 희귀 구조의 경우 특히 선명하고 정확한 분할 경계를 생성하며, 정성적 결과에서 이를 시각적으로 확인할 수 있다.
• ACTION는 ACDC에서 낮은 ASD(1.12 voxel)를 기록하여 경계선 정위치의 높은 정밀도를 보이며, 다양한 비율의 비라벨 데이터에 대해서도 강건성을 유지한다.