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

Hari Sowrirajan, Jingbo Yang외 2명|arXiv (Cornell University)|2020. 10. 11.

COVID-19 diagnosis using AI참고 문헌 28인용 수 52

한 줄 요약

MoCo-CXR는 Momentum Contrast 사전학습을 흉부 X-레이에 적용하여 ImageNet 사전학습보다 더 높은 품질의 표현과 제한된 라벨 데이터에서 특히 더 나은 전달력을 얻는다.

ABSTRACT

Contrastive learning is a form of self-supervision that can leverage unlabeled data to produce pretrained models. While contrastive learning has demonstrated promising results on natural image classification tasks, its application to medical imaging tasks like chest X-ray interpretation has been limited. In this work, we propose MoCo-CXR, which is an adaptation of the contrastive learning method Momentum Contrast (MoCo), to produce models with better representations and initializations for the detection of pathologies in chest X-rays. In detecting pleural effusion, we find that linear models trained on MoCo-CXR-pretrained representations outperform those without MoCo-CXR-pretrained representations, indicating that MoCo-CXR-pretrained representations are of higher-quality. End-to-end fine-tuning experiments reveal that a model initialized via MoCo-CXR-pretraining outperforms its non-MoCo-CXR-pretrained counterpart. We find that MoCo-CXR-pretraining provides the most benefit with limited labeled training data. Finally, we demonstrate similar results on a target Tuberculosis dataset unseen during pretraining, indicating that MoCo-CXR-pretraining endows models with representations and transferability that can be applied across chest X-ray datasets and tasks.

연구 동기 및 목표

• 라벨이 없는 흉부 X-레이 데이터를 활용하여 대조적 사전학습을 통해 견고한 표현을 학습한다.
• 적절한 증강 및 학습 설정으로 흉부 X-레이 특성에 맞게 MoCo를 적응시킨다.
• 다양한 라벨 데이터 구성 하에서 표현의 품질과 외부 데이터셋으로의 전달 가능성을 평가한다.
• MoCo-CXR 초기화를 사용한 엔드투엔드 미세조정의 이점을 입증한다.
• MoCo-CXR가 사전학습 도메드를 넘어 외부 데이터셋으로 전달될 수 있음을 보여준다.

제안 방법

• Chest X-ray 특성에 맞는 증강(무작위 회전 및 수평 뒤집기)을 포함한 수정된 MoCo 설정으로 CheXpert에서 사전학습한다.
• MoCo 사전학습 이전에 수렴 이점을 활용하기 위해 ImageNet 가중치에서 초기화한다.
• 대형 X-레이 이미지에 적합한 더 작은 배치 크기를 가능하게 하는 모멘텀-대조 큐를 사용한다.
• 다양한 라벨 분수와 두 백본(ResNet18, DenseNet121)을 사용하여 CheXpert와 Shenzhen 데이터로 미세조정한다.
• 고정된 백본에서 선형 분류기를 통해 표현을 평가하고 엔드투엔드 미세조정 성능을 평가한다.
• 500회의 테스트 재표본 추출을 통한 부트스트랩으로 통계적 유의성을 평가한다.

실험 결과

연구 질문

• RQ1 MoCo-CXR 표현이 흉부 X-레이 작업에서 선형 평가에서 ImageNet-사전학습 표현을 능가하는가?
• RQ2 라벨이 적게 있을 때 MoCo-CXR 사전학습이 더 큰 이점을 제공하는가?
• RQ3 CheXpert에서 학습된 표현이 다른 작업을 가진 외부 흉부 X-레이 데이터셋(Shenzhen)으로 전달될 수 있는가?
• RQ4 ImageNet 초기화와 비교하여 MoCo-CXR 초기화가 엔드투엔드 미세조정 성능을 향상시키는가?

주요 결과

• MoCo-CXR 사전학습 선형 모델은 일관되게 ImageNet 사전학습 선형 모델보다 우수하며, 매우 낮은 라벨 분수에서 유의미한 이득이 나타난다(예: 0.1% 라벨: 0.096 AUC 향상; 95% CI 0.061, 0.130).
• 엔드투엔드 MoCo-CXR 사전학습 모델은 낮은 라벨 데이터 분수에서 더 큰 이득을 보인다(예: 0.1%: 0.037 AUC 향상; 95% CI 0.015, 0.062).
• MoCo-CXR 표현은 외부 Shenzhen 결핵 데이터셋으로 전달되며, 6.25% 라벨 분수에서 선형 모델이 ImageNet보다 0.054 AUC 만큼 우수(95% CI 0.024, 0.086).
• Shenzhen에서의 엔드투엔드 향상은 선형 전달보다 작으며, 제한된 데이터에서 포화 가능성을 시사하지만 여전히 비제로 이득이 있다.
• MoCo-CXR의 이점은 CheXpert 작업(흉막삼출 등) 전반과 ResNet18 및 DenseNet121 백본을 모두 사용해도 지속된다.