[Paper Review] SPCXR: Self-supervised Pretraining using Chest X-rays Towards a Domain Specific Foundation Model

Syed Muhammad Anwar, Abhijeet Parida+5 more|arXiv (Cornell University)|Nov 23, 2022

COVID-19 diagnosis using AI4 citations

TL;DR

This paper proposes SPCXR, a self-supervised vision transformer pretraining framework using unlabeled chest X-rays to learn domain-specific representations. By applying group-masked self-supervised learning on large-scale CXR data, the method achieves state-of-the-art performance in downstream tasks like COVID-19 and pneumonia classification, as well as lung segmentation, with a 25% accuracy gain on pediatric data compared to supervised baselines.

ABSTRACT

Chest X-rays (CXRs) are a widely used imaging modality for the diagnosis and prognosis of lung disease. The image analysis tasks vary. Examples include pathology detection and lung segmentation. There is a large body of work where machine learning algorithms are developed for specific tasks. A significant recent example is Coronavirus disease (covid-19) detection using CXR data. However, the traditional diagnostic tool design methods based on supervised learning are burdened by the need to provide training data annotation, which should be of good quality for better clinical outcomes. Here, we propose an alternative solution, a new self-supervised paradigm, where a general representation from CXRs is learned using a group-masked self-supervised framework. The pre-trained model is then fine-tuned for domain-specific tasks such as covid-19, pneumonia detection, and general health screening. We show that the same pre-training can be used for the lung segmentation task. Our proposed paradigm shows robust performance in multiple downstream tasks which demonstrates the success of the pre-training. Moreover, the performance of the pre-trained models on data with significant drift during test time proves the learning of a better generic representation. The methods are further validated by covid-19 detection in a unique small-scale pediatric data set. The performance gain in accuracy (~25%) is significant when compared to a supervised transformer-based method. This adds credence to the strength and reliability of our proposed framework and pre-training strategy.

Motivation & Objective

• Address the challenge of limited annotated medical imaging data for rare or pediatric lung diseases by leveraging large-scale unlabeled CXR data for pretraining.
• Develop a domain-specific foundation model for chest radiology that generalizes across diverse clinical data distributions and imaging devices.
• Improve performance on low-resource downstream tasks—particularly in pediatric populations—through self-supervised pretraining on adult CXR data.
• Demonstrate transferability of learned representations to both classification and segmentation tasks using a unified pre-trained model.
• Validate the robustness of the pre-trained model under data drift and distribution shifts typical in real-world clinical deployment.

Proposed method

• Employ a group-masked self-supervised learning (GMML) strategy to pretrain a vision transformer (ViT-S) on 100,000 unlabeled pediatric and adult CXR images from multiple sources.
• Use masked image modeling where patches of the input CXR are masked and the model is trained to reconstruct them, encouraging learning of global and local anatomical features.
• Fine-tune the pre-trained ViT-S within the UNETR architecture for semantic lung segmentation using the Shenzhen and Belarus datasets.
• Apply the same pre-trained weights to downstream classification tasks (e.g., COVID-19, pneumonia, healthy) on RSNA, RICORD, and pediatric datasets without additional annotation.
• Use AdamW optimizer with weight decay of 1e-5 and a 5e-4 initial LR for 400 epochs during fine-tuning on 256×256 image resolution.
• Evaluate performance using standard metrics: accuracy, AUC-ROC, TPR, TNR, mean IoU, Dice coefficient, and Hausdorff surface distance for segmentation.

Experimental results

Research questions

• RQ1 Can self-supervised pretraining on large-scale unlabeled CXRs improve performance on low-resource downstream tasks such as pediatric COVID-19 detection?
• RQ2 To what extent does pretraining on adult CXRs transfer to unseen pediatric datasets with limited data and significant distribution shifts?
• RQ3 How does the proposed GMML-based self-supervised pretraining compare to fully supervised fine-tuning in terms of accuracy and robustness across diverse CXR data distributions?
• RQ4 Can a single pre-trained vision transformer model effectively serve multiple downstream tasks, including both classification and segmentation, in chest radiology?
• RQ5 Does the self-supervised representation learning lead to better generalization and reduced overfitting on small datasets, especially when labeled data is scarce?

Key findings

• The self-supervised pre-trained SPCXR model achieved a 94.73% accuracy on a small pediatric CXR dataset for binary classification (healthy vs. COVID-19), a 25% improvement over a supervised transformer baseline.
• On the RSNA and RICORD test sets, the model reached 98.25% accuracy for adult CXR classification, demonstrating strong performance on standard benchmarks.
• For lung segmentation, the pre-trained model reduced the Hausdorff surface distance by approximately 46% compared to a fully supervised UNETR, indicating superior boundary localization.
• The model outperformed a thoracic radiologist in identifying COVID-19 from CXRs, with a higher true positive rate and lower false positive rate, improving clinical triage reliability.
• The model achieved 99.67% accuracy in distinguishing bacterial pneumonia from viral COVID-19, highlighting its ability to learn subtle pathological patterns.
• Attention heatmaps confirmed that the model attended to clinically relevant regions in most cases, though some attention was observed outside the lungs, indicating a need for improved interpretability.