Label-Efficient Self-Supervised Federated Learning for Tackling Data Heterogeneity in Medical Imaging

Abstract

The collection and curation of large-scale medical datasets from multiple institutions is essential for training accurate deep learning models, but privacy concerns often hinder data sharing. Federated learning (FL) is a promising solution that enables privacy-preserving collaborative learning among different institutions, but it generally suffers from performance deterioration due to heterogeneous data distributions and a lack of quality labeled data. In this paper, we present a robust and label-efficient self-supervised FL framework for medical image analysis. Our method introduces a novel Transformer-based self-supervised pre-training paradigm that pre-trains models directly on decentralized target task datasets using masked image modeling, to facilitate more robust representation learning on heterogeneous data and effective knowledge transfer to downstream models. Extensive empirical results on simulated and real-world medical imaging non-IID federated datasets show that masked image modeling with Transformers significantly improves the robustness of models against various degrees of data heterogeneity. Notably, under severe data heterogeneity, our method, without relying on any additional pre-training data, achieves an improvement of 5.06%, 1.53% and 4.58% in test accuracy on retinal, dermatology and chest X-ray classification compared to the supervised baseline with ImageNet pre-training. In addition, we show that our federated self-supervised pre-training methods yield models that generalize better to out-of-distribution data and perform more effectively when fine-tuning with limited labeled data, compared to existing FL algorithms. The code is available at https://github.com/rui-yan/SSL-FL.

arXiv commentCode and trained models are available at: https://github.com/rui-yan/SSL-FL

References in corpus (23)

• A Simple Framework for Contrastive Learning of Visual Representations
• Representation Learning with Contrastive Predictive Coding
• Federated Learning: Challenges, Methods, and Future Directions
• ChestX-ray8: Hospital-scale Chest X-ray Database and Benchmarks on Weakly-Supervised Classification and Localization of Common Thorax Diseases
• Bootstrap your own latent: A new approach to self-supervised Learning
• The Future of Digital Health with Federated Learning
• Federated Learning with Non-IID Data
• Zero-Shot Text-to-Image Generation
• On the Convergence of FedAvg on Non-IID Data
• BEiT: BERT Pre-Training of Image Transformers
• SCAFFOLD: Stochastic Controlled Averaging for Federated Learning
• Measuring the Effects of Non-Identical Data Distribution for Federated Visual Classification
• Tackling the Objective Inconsistency Problem in Heterogeneous Federated Optimization
• Ensemble Distillation for Robust Model Fusion in Federated Learning
• Using Self-Supervised Learning Can Improve Model Robustness and Uncertainty
• Models Genesis
• The impact of patient clinical information on automated skin cancer detection
• Group Knowledge Transfer: Federated Learning of Large CNNs at the Edge
• Think Locally, Act Globally: Federated Learning with Local and Global Representations
• BIMCV COVID-19+: a large annotated dataset of RX and CT images from COVID-19 patients
• Federated Contrastive Learning for Volumetric Medical Image Segmentation
• Federated Unsupervised Representation Learning
• Divergence-aware Federated Self-Supervised Learning

Cited by in corpus (5)

• SelfFed: Self-Supervised Federated Learning for Data Heterogeneity and Label Scarcity in Medical Images
• Privacy-preserving machine learning for healthcare: open challenges and future perspectives
• Mind the Gap: Federated Learning Broadens Domain Generalization in Diagnostic AI Models
• Boosting multi-demographic federated learning for chest radiograph analysis using general-purpose self-supervised representations
• Decentralized Personalization for Federated Medical Image Segmentation via Gossip Contrastive Mutual Learning