[Paper Review] MoCo-CXR: MoCo Pretraining Improves Representation and Transferability of Chest X-ray Models

Hari Sowrirajan, Jingbo Yang+2 more|arXiv (Cornell University)|Oct 11, 2020

COVID-19 diagnosis using AI28 references52 citations

TL;DR

MoCo-CXR adapts Momentum Contrast pretraining to chest X-rays, yielding higher-quality representations than ImageNet pretraining and better transfer, especially with limited labeled data.

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.

Motivation & Objective

• Leverage unlabeled chest X-ray data to learn robust representations via contrastive pretraining.
• Adapt MoCo for chest X-ray characteristics with suitable augmentations and training settings.
• Evaluate representation quality and transferability to external datasets under varying labeled-data regimes.
• Demonstrate end-to-end fine-tuning benefits using MoCo-CXR initializations.
• Show that MoCo-CXR transfers to external datasets beyond the pretraining domain.

Proposed method

• Pretraining on CheXpert using a modified MoCo setup with chest X-ray specific augmentations (random rotation and horizontal flip).
• Initialize from ImageNet weights to leverage convergence benefits before MoCo pretraining.
• Use a momentum-contrast queue to enable smaller batch sizes suitable for large X-ray images.
• Fine-tune with CheXpert and Shenzhen data using various label fractions and two backbones (ResNet18, DenseNet121).
• Evaluate representations via linear classifiers on frozen backbones and assess end-to-end fine-tuning performance.
• Assess statistical significance with bootstrap over 500 test replications.

Experimental results

Research questions

• RQ1 Do MoCo-CXR representations outperform ImageNet-pretrained representations in linear evaluation on chest X-ray tasks?
• RQ2 Does MoCo-CXR pretraining provide greater benefits when labeled data is scarce?
• RQ3 Can representations learned from CheXpert transfer to an external chest X-ray dataset (Shenzhen) with a different task?
• RQ4 Do MoCo-CXR initializations improve end-to-end fine-tuning performance compared to ImageNet initializations?

Key findings

• MoCo-CXR pretrained linear models consistently outperform ImageNet-pretrained linear models, with significant gains at very low label fractions (e.g., 0.1% labels: 0.096 AUC improvement; 95% CI 0.061, 0.130).
• End-to-end MoCo-CXR pretrained models show larger gains at low labeled data fractions (e.g., 0.1%: 0.037 AUC improvement; 95% CI 0.015, 0.062).
• MoCo-CXR representations transfer to an external Shenzhen tuberculosis dataset, with linear models outperforming ImageNet at 6.25% label fraction by 0.054 AUC (95% CI 0.024, 0.086).
• End-to-end improvements on Shenzhen are smaller than linear transfer, suggesting potential saturation with limited data but still non-zero gains.
• MoCo-CXR benefits persist across CheXpert tasks (pleural effusion and others) and with both ResNet18 and DenseNet121 backbones.