PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Rui Yan, Liangqiong Qu, Qingyue Wei, Shih-Cheng Huang, Liyue Shen, Daniel Rubin, Lei Xing, Yuyin Zhou

arXiv:2205.08576v2cs.CVcs.LG

TL;DR

Medical federated learning is challenged by heterogeneous institutional data and insufficient labels, while centralized pre-training is impractical or mismatched to decentralized medical settings. The paper addresses this with Transformer-based federated masked image modeling on decentralized target-task data, followed by fine-tuning. Across medical imaging datasets, the approach is reported to improve robustness, label efficiency, and generalization, including gains under severe non-IID conditions.

  • Problem

    Medical FL suffers from heterogeneous data distributions and label deficiency, while existing supervised or federated self-supervised approaches have limitations under strongly non-IID or sparse-label settings.

  • Method

    The framework performs federated self-supervised pre-training with Transformer-based masked image modeling on decentralized data, then transfers the learned representations through fine-tuning.

  • Results

    The method is consistently robust across medical tasks and, under severe heterogeneity, improves test accuracy by 5.06% on Retina, 1.53% on Derm, and 4.58% on COVID-FL versus ImageNet-supervised pre-training.

  • Takeaways & Limitations

    Federated masked image modeling provides a label-efficient approach that remains robust to non-IID data and supports out-of-distribution generalization and limited-label fine-tuning.

Abstract

from arXiv · show

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.

I. INTRODUCTION

Medical federated learning must address both heterogeneous, non-IID institutional data and limited labels without centralizing sensitive medical images. The paper proposes decentralized masked-image-modeling pre-training with Transformers and evaluates its robustness across medical imaging tasks.

  • Motivation: Medical FL faces pronounced data heterogeneity, including label-distribution, quantity, and acquisition-protocol differences across institutions.These non-IID partitions contribute to performance degradation in decentralized training.
  • Limitations of existing approaches: Existing approaches can deteriorate under strongly skewed distributions, while supervised ImageNet pre-training may suffer domain discrepancy on medical images.Prior work also reports that federated contrastive learning can be sub-optimal when data are highly heterogeneous or clients have limited data.
  • Evaluation: The study implements BEiT and MAE with Transformers and evaluates them across centralized and federated settings, heterogeneous medical datasets, and varying labeled-data fractions.Experiments include diabetic retinopathy, dermatology, and chest X-ray images.
  • Proposed framework: The framework uses federated self-supervised pre-training with masked image modeling to learn visual representations directly from decentralized data without labels.The learned knowledge is subsequently transferred to target tasks through fine-tuning.
  • Reported contributions: The proposed methods are reported as more label-efficient and robust to non-IID data than ImageNet supervised baselines and existing FL algorithms.The evaluation also includes COVID-FL, a chest X-ray benchmark curated from eight medical sites for realistic federated testing.

II. RELATED WORK

Prior federated learning research addresses private decentralized training, self-supervised representation learning, and data heterogeneity, but challenges remain under strongly non-IID and limited-data conditions. This paper introduces masked-image-modeling pre-training with Transformers to improve federated learning across heterogeneous medical imaging data.

  • Federated Learning: Federated learning trains models across private decentralized client data, with FedAvg aggregating local stochastic-gradient updates at a server.Non-IID partitions can cause weight divergence and performance degradation.
  • Self-supervised Learning: Self-supervised learning uses unlabeled data through pretext tasks, then transfers learned representations to downstream tasks while reducing annotation costs.Masked image modeling reconstructs signals from corrupted images and does not heavily depend on large samples or specific augmentations.
  • Federated Self-supervised Learning: Existing federated contrastive-learning methods often rely on sufficient local data diversity, leaving performance under highly heterogeneous small-client partitions unresolved.Prior studies mainly evaluate label-distribution skew with more than 10,000 images per client.
  • Federated Self-supervised Learning: The proposed framework is the first federated self-supervised pre-training approach using masked image modeling to address data heterogeneity and label deficiency together.It couples masked image modeling with Transformers for robust representation learning under distribution shifts and limited data.
  • Problem Setting: The framework targets generalized global models when client distributions differ and local labeled datasets are relatively small.Its training procedure combines federated self-supervised pre-training with supervised federated fine-tuning.

B. Generalized Framework

The generalized framework first pre-trains encoders and decoders federatively with masked image modeling on decentralized unlabeled data, then transfers the learned encoder to supervised federated fine-tuning. It supports BEiT and MAE implementations within this two-stage protocol.

  • Two-stage framework: The framework has federated self-supervised pre-training followed by supervised federated fine-tuning for transferring knowledge to target tasks.The pre-training stage uses masked image modeling on decentralized data.
  • Masked image modeling: BEiT and MAE are integrated as masked image modeling methods and are named Fed-BEiT and Fed-MAE in the federated framework.Both methods are used for pre-training and subsequent fine-tuning.
  • Federated pre-training: During each pre-training round, clients train local autoencoders, upload their weights, and receive a server-averaged global autoencoder.The process repeats until the maximum communication round T.
  • Federated fine-tuning: The final global encoder initializes each local encoder, after which a linear classifier is appended and end-to-end federated fine-tuning uses labeled client images.This transfers the pre-trained representation into the target supervised task.
  • Masked-image reconstruction: Each local pre-training model is an encoder-decoder autoencoder trained by masking image patches and reconstructing the original masked signals.The input image is divided into patches before masking.

1) Masking:

The masking stage partitions image patches into masked and visible subsets before encoding. BEiT and MAE differ in their masking strategy and in whether masked-patch embeddings are supplied to the Transformer encoder.

  • Masking notation: The masking formulation separates patches into masked positions M and visible positions V, with masked-patch count determined by the masking ratio.The visible and masked sets together contain all P image patches.
  • Masking strategies: BEiT uses block-wise n-gram masking, whereas MAE uses random masking.These are the two masked image modeling variants implemented in the framework.
  • Patch encoding: The ViT encoder processes sequences of image patches within the federated masked-image-modeling framework.The encoder is applied to patchified images rather than whole images.
  • BEiT: BEiT provides visible-patch inputs together with learnable embeddings for masked patches, and its encoder outputs representations for all patches.Each encoded patch representation has dimension D.
  • MAE: MAE supplies the encoder only with linearly projected visible patches plus position embeddings, producing representations for visible patches.The output representation for each visible patch has dimension D.

3) Decoder:

The decoder reconstructs masked image content from encoder representations, with BEiT predicting visual tokens and MAE regressing pixels. Local encoder-decoder models minimize client objectives before federated aggregation across communication rounds.

  • The decoder performs signal reconstruction from encoded representations of input patches.
  • BEiT: BEiT uses a linear decoder to predict visual tokens at masked positions from all encoded patches.
  • MAE: MAE uses visible-patch representations, mask vectors, and position embeddings in a lightweight ViT decoder to regress masked pixels.
  • Each client trains its local encoder and decoder on local data by minimizing an expected loss objective.
  • BEiT uses cross-entropy for masked visual-token prediction, whereas MAE uses mean squared error for masked pixel prediction.
  • Federated pre-training alternates local gradient updates with weighted server averaging for up to T communication rounds before retaining the global encoder.

D. Supervised Federated Fine-tuning

Supervised federated fine-tuning initializes each client encoder with the globally pre-trained encoder and adds a linear classifier. Average-pooled representations are classified through a softmax layer using local labeled data.

  • Fine-tuning initializes each local encoder with the pre-trained global encoder and appends a linear classifier.
  • Average pooling extracts encoder representations that a softmax classifier uses to minimize cross-entropy for image classification.

IV. EXPERIMENTS

Experiments evaluate the framework across simulated and real-world federated medical imaging datasets, spanning retinal, dermatology, and chest X-ray tasks. The evaluation examines heterogeneity robustness, out-of-distribution generalization, and label efficiency.

  • The experiments compare methods on simulated and real-world federated medical datasets.
  • Tasks include diabetic retinopathy detection, skin-lesion diagnosis, and pneumonia or COVID-19 identification from chest X-rays.
  • COVID-FL represents institutions with site-specific class omissions and acquisition differences that create label and intensity heterogeneity.
  • COVID-FL: COVID-FL contains 16,044 training images and 3,974 test images under an 80%-20% split.
  • Skin-FL: Skin-FL uses 33,126 ISIC20 images as an out-of-distribution test set containing classes absent from training.
  • Skin-FL: The Skin-FL evaluation is challenging because of severe class imbalance.

B. Experiment Setup

The setup simulates varying heterogeneity on Retina and Derm and uses real-world heterogeneous COVID-FL and Skin-FL partitions. Training applies image augmentation, while evaluation metrics differ by dataset.

  • Dirichlet concentration values α = {100, 1.0, 0.5} create IID, moderate non-IID, and severe non-IID Retina and Derm partitions across five clients.
  • Smaller Dirichlet α values produce greater data heterogeneity by changing class proportions assigned to clients.
  • COVID-FL and Skin-FL contain both label-distribution and feature-distribution skew, partitioned into 12 and 10 clients respectively.
  • Data augmentation: Pre-training uses random scaling, 224 × 224 crops, color jittering, and horizontal flipping, with dataset-specific scaling ranges.
  • Data augmentation: Fine-tuning adds random rotation and uses different scaling ranges from pre-training.

3) Self-supervised FL pre-training setup:

The framework pre-trains ViT models on decentralized target-task data using federated masked image modeling, then fine-tunes them for medical image classification.

  • Self-supervised FL pre-training setup: Fed-BEiT and Fed-MAE use masked image modeling with Transformer backbones for federated self-supervised pre-training.The implementation uses ViT-B and dataset-specific image patching and mask ratios.
  • Self-supervised FL pre-training setup: Pre-training runs for 1000 communication rounds with Fed-BEiT and 1600 rounds with Fed-MAE, using cosine learning-rate decay.Both schedules include warmup periods and use the same hyperparameters between centralized and federated learning within each task.
  • Evaluation setup: The experiments evaluate accuracy on Retina, Derm, and COVID-FL, and F1-score on the imbalanced Skin-FL dataset.The setup compares proposed federated pre-training with scratch and ImageNet-based baselines.
  • Evaluation setup: Proposed models are pre-trained directly on decentralized target-task data, whereas most baselines use centralized ImageNet-22K pre-training.COVID-FL additionally includes centralized CXR14 pre-training baselines.
  • Evaluation setup: All comparison methods use ViT-B backbones, while federated fine-tuning uses 1000 rounds for random initialization and 100 rounds for other models.This keeps the backbone fixed while allowing different training durations for scratch and pre-trained models.

1) More robust to data heterogeneity:

Across medical imaging tasks, federated masked image modeling is most robust under severe non-IID conditions and remains effective with limited labels and out-of-distribution data.

  • Robustness to data heterogeneity: 5.06%, 1.53%, and 4.58% test-accuracy improvements occur on Retina, Derm, and COVID-FL under severe non-IID data versus ViT ImageNet.Fed-BEiT and Fed-MAE show the smallest accuracy discrepancies across data partitions and outperform baselines in strongly skewed settings.
  • Scope and comparison: Centralized in-domain CXR14 pre-training can be a good alternative on COVID-FL, although such large datasets are uncommon because of privacy and ownership concerns.Fed-BEiT and Fed-MAE outperform supervised CXR14 pre-training, while their gains over CXR14 self-supervised baselines are not significant.
  • Label efficiency: With approximately 70%, 30%, and 10% of Retina labels, the proposed methods consistently improve performance over supervised ImageNet pre-training in IID and non-IID settings.Training from scratch performs poorly when labeled images are limited.
  • Out-of-distribution generalization: Fed-BEiT and Fed-MAE perform slightly better than supervised ImageNet pre-training and notably better than other methods on out-of-distribution Skin-FL data.This evaluation uses Skin-FL to assess generalization beyond the training distribution.

4) Comparison with prior FL methods:

Against federated self-supervised, optimization-based, and semi-supervised baselines, the proposed methods provide stronger robustness to heterogeneous data and limited annotations.

  • Federated self-supervised baselines: Fed-MAE and Fed-BEiT surpass FedEMA by 3.47% and 2.81% test accuracy on Retina’s severe non-IID Split-3.The proposed methods outperform all compared federated self-supervised baselines on Retina Split-2 and Split-3.
  • Federated self-supervised baselines: All self-supervised methods improve robustness to heterogeneity over random initialization, while the proposed methods lead on Retina’s heterogeneous partitions.The comparison includes FedEMA, FedBYOL, FedMoCo, and FedMoCov3.
  • Limited-label learning: The proposed methods outperform all baselines with limited annotations, with larger gains when fewer labeled samples are used.This comparison evaluates Retina Split-1 using different fractions of labeled training data.
  • Communication-round comparison: The ablation compares pre-training and scratch models across total communication rounds, separating pre-training and fine-tuning rounds.Fed-BEiT and Fed-MAE use pre-training rounds plus 100 fine-tuning rounds, while ViT Scratch varies fine-tuning rounds.

1) Number of communication rounds:

Increasing federated pre-training rounds improves accuracy, but the proposed methods reach useful performance thresholds with fewer total rounds than training from scratch.

  • Communication cost: Communication cost depends on both the number of rounds and the communicated model size, which is larger during pre-training than fine-tuning.The communicated pre-training models contain 111.7M parameters for Fed-MAE and 92M for Fed-BEiT, compared with 85.8M during fine-tuning.
  • Training-data size: Reducing Retina Split-3 training images from 9000 to 6000 and 3000 narrows the performance gap between the proposed pre-training and ImageNet initialization.This result indicates that the relative advantage changes with total training-data size.
  • Mask ratio: Optimal mask ratios are 40% for BEiT and 60% for MAE on Retina and Derm, while MAE uses 30% on COVID-FL.The paper relates the COVID-FL choice to differences between chest X-ray images and retina or skin images.
  • Data augmentation: Adding grayscale scaling and color jittering improves centralized Retina MAE pre-training accuracy by 0.7%.The ablation suggests task-specific augmentations may benefit medical-image pre-training.
Loading 2205.08576v2…