Federated Vision Transformer with Adaptive Focal Loss for Medical Image Classification

TL;DR

This paper tackles privacy-conscious medical image classification by integrating a Vision Transformer into a federated learning framework and addressing dual-level data imbalance. It introduces Dynamic Adaptive Focal Loss (DAFL), which couples a trainable focal loss with real-time client- and class-level imbalance signals, and a distribution-aware aggregation scheme that down-weights highly imbalanced clients. Theoretical analysis links DAFL to gradient rectification and variance reduction, while experiments on RSNA-ICH, Ocular Disease, and ISIC demonstrate improved minority-class performance, faster convergence, and strong interpretability via attention rollout. The work advances clinical FL by providing a robust, efficient method that maintains privacy, improves reliability under non-IID conditions, and yields clinically meaningful attention localization.

Abstract

While deep learning models like Vision Transformer (ViT) have achieved significant advances, they typically require large datasets. With data privacy regulations, access to many original datasets is restricted, especially medical images. Federated learning (FL) addresses this challenge by enabling global model aggregation without data exchange. However, the heterogeneity of the data and the class imbalance that exist in local clients pose challenges for the generalization of the model. This study proposes a FL framework leveraging a dynamic adaptive focal loss (DAFL) and a client-aware aggregation strategy for local training. Specifically, we design a dynamic class imbalance coefficient that adjusts based on each client's sample distribution and class data distribution, ensuring minority classes receive sufficient attention and preventing sparse data from being ignored. To address client heterogeneity, a weighted aggregation strategy is adopted, which adapts to data size and characteristics to better capture inter-client variations. The classification results on three public datasets (ISIC, Ocular Disease and RSNA-ICH) show that the proposed framework outperforms DenseNet121, ResNet50, ViT-S/16, ViT-L/32, FedCLIP, Swin Transformer, CoAtNet, and MixNet in most cases, with accuracy improvements ranging from 0.98\% to 41.69\%. Ablation studies on the imbalanced ISIC dataset validate the effectiveness of the proposed loss function and aggregation strategy compared to traditional loss functions and other FL approaches. The codes can be found at: https://github.com/AIPMLab/ViT-FLDAF.

Figures (8)

• Figure 1: Flowchart of our proposed method. The framework integrates a ViT backbone within a federated learning structure. The client-level imbalance parameter (Eq. \ref{['eq:client_imbalance']}) and the global class-level imbalance parameter (Eq. \ref{['eq:class_imbalance']}) are calculated based on data distributions to dynamically guide the federated aggregation and local training processes, respectively.
• Figure 2: The number of samples of each client in RSNA-ICH, Ocular disease and ISIC datasets.
• Figure 3: Classification accuracy for each epoch using test samples of Ocular disease, ISIC, and RSNA-ICH.
• Figure 4: The confusion matrices obtained by DAFL on ISIC (Left), Ocular Disease (Middle), and RSNA-ICH (Right).
• Figure 5: Decision Curve Analysis (DCA) comparing the net benefit of various models across different thresholds for three datasets (ISIC, Ocular Disease, and RSNA-ICH). The x-axis represents the threshold probability, while the y-axis shows the net benefit.