[논문 리뷰] MoCo-Pretraining Improves Representations and Transferability of Chest X-ray Models

Hari Sowrirajan, Jing Yang외 2명|arXiv (Cornell University)|2021. 05. 04.

COVID-19 diagnosis using AI인용 수 10

한 줄 요약

이 연구는 체크엑스퍼트에서 MoCo 사전학습을 통해 흉부 X-ray 해석을 위한 표현과 전이 가능성 향상이 가능함을 입증한다. 레이블이 0.1%인 경우에만 MoCo 사전학습 모델이 선형 평가에서 AUC 0.096 향상, 종단간 최적화에서 AUC 0.037 향상하여 흉막 축적 검출에 있어 일관된 성능 향상을 보였으며, 레이블 비율에 관계없이 전이 성능이 향상되었고, 샤먼 흉부결핵 데이터셋으로의 전이도 성공하였다.

ABSTRACT

Self-supervised approaches such as Momentum Contrast (MoCo) can leverage unlabeled data to produce pretrained models for subsequent fine-tuning on labeled data. While MoCo has demonstrated promising results on natural image classification tasks, its application to medical imaging tasks like chest X-ray interpretation has been limited. Chest X-ray interpretation is fundamentally different from natural image classification in ways that may limit the applicability of self-supervised approaches, including that (1) classification depends on differences in a small number of pixels, (2) X-rays are large and grayscale, (3) there are far fewer unlabeled chest X-ray images than natural images. In this work, we investigate whether MoCo-pretraining leads to better representations or initializations for chest X-ray interpretation. We conduct MoCo-pretraining on CheXpert, a large labeled dataset of X-rays, followed by supervised fine-tuning experiments on the pleural effusion task. Using 0.1% of labeled training data, we find that a linear model trained on MoCo-pretrained representations outperforms one trained on representations without MoCo-pretraining by an AUC of 0.096 (95% CI 0.061, 0.130), indicating that MoCo-pretrained representations are of higher quality. Furthermore, a model fine-tuned end-to-end with MoCo-pretraining outperforms its non-MoCo-pretrained counterpart by an AUC of 0.037 (95% CI 0.015, 0.062) with the 0.1% label fraction. These AUC improvements are observed for all label fractions for both the linear model and an end-to-end fine-tuned model with the greater improvements for smaller label fractions. Finally, we observe similar results on a small, target chest X-ray dataset (Shenzhen dataset for tuberculosis) with MoCo-pretraining done on the source dataset (CheXpert), which suggests that pretraining on unlabeled X-rays can provide transfer learning benefits for a target task. Our study demonstrates that MoCo-pretraining provides high-quality representations and transferable initializations for chest X-ray interpretation.

연구 동기 및 목표

• MoCo 자기지도 학습 사전학습이 흉부 X-ray 해석을 위한 표현 품질을 향상시키는지 조사하기.
• 제한된 레이블 데이터로 MoCo 사전학습 특징의 전이 가능성 평가하기.
• 낮은 데이터 환경에서 MoCo 사전학습이 최적화를 위한 더 나은 초기화를 제공하는지 평가하기.
• CheXpert(원천 데이터셋)에서 학습한 MoCo 사전학습 특징이 Shenzhen TB 데이터셋(목표 데이터셋)으로의 전이 학습 이점을 검증하기.

제안 방법

• 대규모 레이블이 없는 흉부 X-ray 데이터셋(CheXpert)에서 MoCo(Momentum Contrast) 자기지도 학습을 적용하여 일반적인 표현 학습하기.
• 다양한 비율의 레이블 데이터를 사용하여 하류 흉막 축적 분류 작업에서 MoCo로 학습한 특징을 미세조정하기.
• 고정된 MoCo 특징에 선형 분류기 적용하는 선형 탐색 전략과 전체 모델을 미세조정하는 종단간 최적화 전략을 평가하기.
• Shenzhen 결핵 데이터셋을 목표 도메인으로 사용하여 CheXpert에서 사전학습한 MoCo 특징의 제로샷 전이 가능성 테스트하기.
• 95% 신뢰구간을 사용하여 통계적 유의성을 평가하기 위해 AUC로 성능 측정하기.
• 다양한 레이블 비율과 데이터셋에서 MoCo 사전학습 모델을 비사전학습 기반 모델과 비교하기.

실험 결과

연구 질문

• RQ1 MoCo 사전학습은 비사전학습 모델 대비 흉부 X-ray 영상에 대해 학습된 표현 품질을 향상시키는가?
• RQ2 MoCo 사전학습은 제한된 레이블 데이터로 하류 작업에서 전이 학습 성능을 얼마나 향상시키는가?
• RQ3 선형 및 종단간 최적화 설정 모두에서 다양한 레이블 비율에서 MoCo 사전학습 모델의 성능 향상은 어떻게 비교되는가?
• RQ4 CheXpert에서 사전학습한 MoCo 특징는 Shenzhen 데이터셋에서 다른 목표 작업(예: 결핵 검출)으로 효과적으로 전이될 수 있는가?

주요 결과

• 레이블 데이터가 0.1%일 경우, MoCo 사전학습 특징에 기반한 선형 모델이 비사전학습 기반 모델보다 AUC가 0.096 높게 나타났다 (95% 신뢰구간 0.061, 0.130).
• 0.1%의 데이터로 학습한 종단간 최적화 모델은 MoCo 사전학습 가중치를 사용할 경우 비사전학습 대비 AUC에서 0.037 향상되었다 (95% 신뢰구간 0.015, 0.062).
• MoCo 사전학습의 성능 향상은 모든 레이블 비율에서 일관되게 관찰되었으며, 특히 낮은 데이터 환경에서 상대적 향상이 더 컸다.
• CheXpert(원천)에서 사전학습한 특징를 Shenzhen TB 데이터셋(목표)으로 전이했을 때 유사한 성능 향상이 관찰되어 MoCo 특징의 전이 가능성 확인됨.
• MoCo 사전학습은 극도로 낮은 데이터 가용성에서도 더 높은 품질의 표현을 생성하여 하류 작업으로의 일반화 능력을 향상시켰다.
• 결과는 MoCo 사전학습이 흉부 X-ray 해석을 위한 더 나은 표현과 더 나은 모델 초기화를 제공함을 시사한다.