Federated Learning for Medical Image Classification: A Comprehensive Benchmark

Abstract

Federated Learning (FL) is a promising paradigm for collaborative medical AI, yet its performance often degrades when faced with the heterogeneity of real-world clinical data. This paper provides a comprehensive benchmark of state-of-the-art FL algorithms across nine medical imaging datasets. We identify "feature shift"—variations in pixel distributions across centers—as a primary obstacle and propose a novel framework combining local Denoising Diffusion Probabilistic Models (DDPM) with label smoothing to enhance model robustness and accuracy.

Introduction: The Gap Between Theory and Practice

Federated Learning is frequently cited as the solution for training medical AI models without compromising patient privacy. By keeping data localized at hospitals and only sharing model updates, FL bypasses the hurdles of data centralization. However, most FL research relies on "natural" image datasets (e.g., MNIST, CIFAR-10) that lack the complexity of medical imaging.

In clinical settings, data is rarely independent and identically distributed (IID). A chest X-ray from Shenzhen may differ significantly from one taken in Montgomery due to different imaging hardware, patient demographics, and clinical protocols. This feature shift causes standard FL algorithms to diverge, leading to suboptimal performance in actual medical applications.

Methodology: Generative Alignment and Label Smoothing

To address these challenges, we introduce a framework that leverages Denoising Diffusion Probabilistic Models (DDPM) and Label Smoothing.

In this approach, each client trains a local diffusion model on their private data. This model is then used to generate high-quality synthetic images that augment the local training set, effectively balancing classes and smoothing the data distribution. To mitigate the risk of the model over-fitting to synthetic noise, we apply label smoothing, converting "hard" categorical labels into "soft" probabilities.

Expand

Source, p. 6

Crucially, this method maintains the privacy benefits of FL. Only the classification model parameters are transmitted to the central server; the diffusion models and the synthetic data they produce remain strictly local to each client.

Experimental Setup and the Challenge of Feature Shift

The study evaluated FL algorithms across nine diverse medical datasets, including pathology (ColonPath, CRC), dermatology (NeoJaundice, HAM10K), ophthalmology (Retino, DR), and radiology (COVID-QU-Ex, Breast).

A central component of our evaluation is a real-world multi-center Tuberculosis (TB) dataset compiled from Shenzhen, India, and Montgomery. This dataset highlights the extreme variability in medical data. As shown in the pixel frequency curves below, the same disease can manifest with vastly different image characteristics across different geographic regions.

Expand

Source, p. 7

Key Insights from the Benchmark

• Algorithm Inconsistency: No single existing FL algorithm (e.g., FedProx, MOON, FedBN) consistently outperformed others across all medical modalities. In some instances, basic FedAvg remained the most robust baseline.
• The Impact of Feature Shift: The high variability in pixel distributions is a more significant hurdle in medical FL than in natural image FL.
• Generative Benefits: By generating synthetic data locally, clients can "align" their distributions. This helps the global model converge toward a representation that generalizes better across heterogeneous sites.

Results and Performance

The proposed DDPM + Label Smoothing method emerged as the top performer, achieving the highest accuracy in 7 out of 9 experimental scenarios.

• Accuracy Gains: On the Retino dataset, the method reached 85.14% accuracy, a significant improvement over FedAvg’s 76.13%.
• Handling Imbalance: The framework successfully stabilized training in scenarios with extreme sample size imbalances between clients.
• Communication Efficiency: Because only the classification model is shared, the communication overhead remains identical to standard FedAvg, making it practical for bandwidth-constrained environments.

Limitations and Future Directions

While effective, the local training of diffusion models requires substantial computational resources (GPUs) at each client site. Additionally, this benchmark focused on classification; future research should extend these generative alignment techniques to more complex tasks like medical image segmentation and object detection.