[논문 리뷰] Self-supervised Mean Teacher for Semi-supervised Chest X-ray Classification

Fengbei Liu, Yu Tian외 4명|arXiv (Cornell University)|2021. 03. 05.

COVID-19 diagnosis using AI참고 문헌 39인용 수 9

한 줄 요약

이 논문은 자기지도 학습(mean-teacher) 프레임워크인 S²MTS²를 제안하며, 반감기 대비 학습(JCL)과 지수 이동 평균 교사 학습을 결합하여 흉부 X선 촬영 분류를 위한 준지도 학습을 수행한다. 사전 훈련 단계에서 비라벨 데이터를 활용해 무한한 양의 양성 특징 쌍을 생성하고, 미세조정 단계에서 일致성 정규화를 적용함으로써, ChestX-ray14, CheXpert, ISIC2018에서 최신 기술(SOTA) 수준의 성능을 달성한다. 특히 라벨이 제한된 경우에도 이전 방법들보다 크게 뛰어난 성능을 보인다.

ABSTRACT

The training of deep learning models generally requires a large amount of annotated data for effective convergence and generalisation. However, obtaining high-quality annotations is a laboursome and expensive process due to the need of expert radiologists for the labelling task. The study of semi-supervised learning in medical image analysis is then of crucial importance given that it is much less expensive to obtain unlabelled images than to acquire images labelled by expert radiologists. Essentially, semi-supervised methods leverage large sets of unlabelled data to enable better training convergence and generalisation than using only the small set of labelled images. In this paper, we propose Self-supervised Mean Teacher for Semi-supervised (S$^2$MTS$^2$) learning that combines self-supervised mean-teacher pre-training with semi-supervised fine-tuning. The main innovation of S$^2$MTS$^2$ is the self-supervised mean-teacher pre-training based on the joint contrastive learning, which uses an infinite number of pairs of positive query and key features to improve the mean-teacher representation. The model is then fine-tuned using the exponential moving average teacher framework trained with semi-supervised learning. We validate S$^2$MTS$^2$ on the multi-label classification problems from Chest X-ray14 and CheXpert, and the multi-class classification from ISIC2018, where we show that it outperforms the previous SOTA semi-supervised learning methods by a large margin.

연구 동기 및 목표

• 의료 영상의 전문가 레이블링 비용이 높다는 문제를 해결하기 위해 흉부 X선 촬영 분류에서 준지도 학습(SSL)을 향상시키는 것.
• 기존의 SSL 방법들이 미세조정 단계에서 비라벨 데이터를 충분히 활용하지 못하거나, 자기지도 사전 훈련 단계에서 유한한 양의 양성 쌍에 의존하는 한계를 극복하는 것.
• 자기지도 사전 훈련과 일致성 기반의 학생-교사 학습을 통합한 새로운 프레임워크를 개발하여 표현력과 일반화 능력을 향상시키는 것.
• 다양한 라벨 데이터 비율을 가진 다중 레이블(ChestX-ray14, CheXpert) 및 다중 클래스(ISIC2018) 의료 영상 기준에서 방법의 유효성을 검증하는 것.

제안 방법

• 자기지도 사전 훈련 단계에서 비라벨 이미지로부터 유한한 수의 양성 특징 쌍을 생성하기 위해 연합 대비 학습(JCL)을 사용하여 표현 학습을 향상시킨다.
• 학생-교사 모델은 학생의 가중치에 대한 지수 이동 평균(EMA)을 사용해 견고한 교사 네트워크를 형성하며, 예측의 일致성을 강제한다.
• 자기지도 사전 훈련 단계에서는 비라벨 데이터에서 추출한 쿼리 및 키 특징 간의 무한한 양성 쌍에 대해 대비 손실을 최소화한다.
• 미세조정 단계에서는 라벨이 있는 데이터에 대한 교차 엔트로피 손실과 비라벨 데이터에서 학생 및 교사 예측 간의 일致성 손실을 병행한다.
• S4L(자기지도 사전 훈련)과 일치성 강화 SSL을 결합함으로써 기존 최신 기술(SOTA) 방법보다 더 나은 일반화 성능을 달성한다.
• 모델은 ResNet 또는 DenseNet 백본을 사용해 엔드 투 엔드로 훈련되며, 최종 분류기는 라벨이 있는 데이터와 비라벨 데이터를 모두 활용해 미세조정된다.

실험 결과

연구 질문

• RQ1 무한한 양의 양성 특징 쌍을 활용한 자기지도 사전 훈련이 준지도 흉부 X선 촬영 분류에서 표현 학습을 향상시킬 수 있는가?
• RQ2 자기지도 사전 훈련과 mean-teacher 일치성 정규화를 조합하면 기존 최신 기술(SOTA) SSL 방법보다 성능이 향상되는가?
• RQ3 자기지도 학습, JCL, mean-teacher 미세조정이라는 개별 구성 요소들이 전체 성능에 기여하는 정도는 어떠한가?
• RQ4 제안된 방법이 100%의 라벨 데이터로 미세조정되었을 경우, 완전히 지도 학습 설정에서도 최신 기술(SOTA) 성능을 달성할 수 있는가?

주요 결과

• ChestX-ray14에서 20%의 라벨 데이터를 사용할 경우, S²MTS²는 평균 AUC 81.06%를 기록하여 이전 최신 기술(SOTA) 방법보다 0.5~1.5%p 높은 성능을 보였다.
• CheXpert에서 S²MTS²는 평균 AUC 94.71%를 달성하여 이전 최신 기술(93.58%)을 1.13p 초월했다.
• ISIC2018에서 20%의 라벨 데이터를 사용할 경우, S²MTS²는 F1 스코어 62.67%를 기록하여 이전 최신 기술(59.10)보다 3.57p 향상되었다.
• 제거 실험 결과 각 구성 요소—자기지도 학습, JCL, mean-teacher—가 유의미하게 기여하며, 전체 모델은 기준 모델 대비 평균 AUC를 2.7%p 향상시켰다.
• 100%의 라벨 데이터로 훈련된 경우에도 S²MTS²는 ChestX-ray14에서 평균 AUC 71.58%를 기록하여 지도 학습 최신 기술(SOTA) 방법을 초월했다.
• 다양한 데이터 분할 및 라벨 비율에 걸쳐 뚜렷한 성능 향상을 보이며, 모든 기준에서 AUC, F1, 민감도 측면에서 일관된 개선을 보였다.