[논문 리뷰] SimCVD: Simple Contrastive Voxel-Wise Representation Distillation for Semi-Supervised Medical Image Segmentation

Chenyu You, Yuan Zhou외 3명|arXiv (Cornell University)|2021. 08. 13.

Advanced Neural Network Applications인용 수 16

한 줄 요약

SimCVD는 단지 두 개의 드롭아웃 마스크를 최소한의 데이터 증강 수단으로 사용하여 대조적 볼륨 단위 표현 정규화 프레임워크를 제안한다. 이는 경계 인식 표현을 대조 학습을 통해 학습한다. 이는 최신 기술 수준의 성능을 달성하며, 라벨이 20%인 경우 LA 데이터셋에서 90.85%의 Dice 스코어를 기록하고, 10%일 경우 89.03%를 기록하여 이전 방법들보다 최대 2.22% 향상된다.

ABSTRACT

Automated segmentation in medical image analysis is a challenging task that requires a large amount of manually labeled data. However, most existing learning-based approaches usually suffer from limited manually annotated medical data, which poses a major practical problem for accurate and robust medical image segmentation. In addition, most existing semi-supervised approaches are usually not robust compared with the supervised counterparts, and also lack explicit modeling of geometric structure and semantic information, both of which limit the segmentation accuracy. In this work, we present SimCVD, a simple contrastive distillation framework that significantly advances state-of-the-art voxel-wise representation learning. We first describe an unsupervised training strategy, which takes two views of an input volume and predicts their signed distance maps of object boundaries in a contrastive objective, with only two independent dropout as mask. This simple approach works surprisingly well, performing on the same level as previous fully supervised methods with much less labeled data. We hypothesize that dropout can be viewed as a minimal form of data augmentation and makes the network robust to representation collapse. Then, we propose to perform structural distillation by distilling pair-wise similarities. We evaluate SimCVD on two popular datasets: the Left Atrial Segmentation Challenge (LA) and the NIH pancreas CT dataset. The results on the LA dataset demonstrate that, in two types of labeled ratios (i.e., 20% and 10%), SimCVD achieves an average Dice score of 90.85% and 89.03% respectively, a 0.91% and 2.22% improvement compared to previous best results. Our method can be trained in an end-to-end fashion, showing the promise of utilizing SimCVD as a general framework for downstream tasks, such as medical image synthesis, enhancement, and registration.

연구 동기 및 목표

• 딥 러닝 기반 분할에서(annotation이 부족한) 의료 영상 문제를 해결한다.
• 완전히 지도 학습된 모델들에 비해 준지도 모델의 강건성과 정확도를 향상시킨다.
• 기하학적 구조와 의미 정보를 명시적으로 모델링하여 경계 검출 성능을 향상시킨다.
• 제한된 라벨과 풍부한 비라벨 데이터로부터 강력한 볼륨 단위 표현을 학습할 수 있는 단순하면서도 효과적인 프레임워크를 개발한다.
• 합성, 향상, 정렬과 같은 후속 의료 영상 작업을 위한 엔드 투 엔드 학습을 가능하게 한다.

제안 방법

• 두 개의 독립적인 드롭아웃 마스크를 통해 생성된 동일한 3D 볼륨의 두 뷰를 대조 학습을 위한 양성 쌍으로 사용한다.
• 학생 네트워크를 통해 객체 경계의 세그멘테이션 맵과 서명 거리 맵(SDM)을 예측하도록 훈련한다.
• 두 뷰의 잠재 표현에 대해 대조 손실을 적용하여 경계 인식 특징에서의 일치를 유도한다.
• 교사 네트워크에서 학생 네트워크로 쌍별 유사성 패턴을 전달하는 구조적 정규화를 구현한다.
• 평균 교사 원칙에 따라 학생 네트워크 가중치의 지수 이동 평균을 사용해 교사 네트워크 가중치를 업데이트한다.
• 학생 네트워크를 감독 손실(세그멘테이션 및 SDM), 비감독 손실(대조, 쌍별 정규화, 일관성)의 조합으로 최적화한다.

실험 결과

연구 질문

• RQ1 단지 드롭아웃만으로도 의료 영상 분할을 위한 대조적 볼륨 단위 표현 학습에서 경쟁 가능한 성능을 달성할 수 있는가?
• RQ2 표준 데이터 증강 기법들(예: 인페인팅, 강성 변형)과 비교해 드롭아웃은 표현 품질과 일반화 능력 측면에서 어떻게 다른가?
• RQ3 서명 거리 맵을 통한 경계 인식 표현 학습이 저샷 설정에서 분할 정확도를 얼마나 향상시키는가?
• RQ4 볼륨 표현 간 쌍별 유사성의 구조적 정규화가 모델 성능을 추가로 향상시킬 수 있는가?
• RQ5 제안된 SimCVD 프레임워크는 다양한 의료 영상 작업과 데이터셋에 대해 잘 일반화되는가?

주요 결과

• 라벨이 20%인 경우 SimCVD는 LA 데이터셋에서 90.85%의 Dice 스코어를 기록하며, 이전 최신 기술 수준보다 0.91% 향상되었다.
• 라벨이 10%인 경우 SimCVD는 89.03%의 Dice 스코어를 기록하여 이전 최고 성능 방법보다 2.22% 향상되었다.
• 훈련 중 드롭아웃을 제거하면 Dice 스코어가 1.34% 유의미하게 감소하여, 드롭아웃이 표현 학습에서 핵심적인 역할을 한다는 것을 입증했다.
• 표준 데이터 증강 기법들(예: 인페인팅, 비선형 변형)을 추가하면 드롭아웃만 사용했을 때보다 성능이 떨어지며, 이는 해로운 노이즈를 유도하기 때문임을 시사한다.
• 풀링 크기가 128인 적응형 평균 풀링을 사용할 경우 최적의 성능을 기록하며, 더 큰 크기(예: 256)는 성능 향상에 기여하지 않았다.
• 제거 실험 결과 드롭아웃이 표현 붕괴를 방지하고 별개의 양성 쌍을 생성하는 효과적인 최소한의 데이터 증강 형태임을 확인했다.