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.

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

References in corpus (8)

• Bootstrap your own latent: A new approach to self-supervised Learning
• Zero-Shot Text-to-Image Generation
• BEiT: BERT Pre-Training of Image Transformers
• Tackling the Objective Inconsistency Problem in Heterogeneous Federated Optimization
• 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