FedProIn: Mitigating Client Drift for Learnable Prototypes in Federated Medical Imaging

2026-07-05

AI Summary

This work addresses the challenge of client drift caused by non-independent and identically distributed (non-IID) data in federated medical image learning, which leads to unstable convergence and degraded model performance. The authors propose a novel federated learning framework based on learnable class prototypes, uniquely decomposing client drift into feature drift and prototype drift. To mitigate these issues, they introduce a feature divergence loss and a prototype contrastive loss for regularization. Furthermore, a normalized influence-aware prototype aggregation mechanism is developed to adaptively weight local prototypes, thereby enhancing global semantic alignment. Evaluated under non-IID settings on the HAM10000 and Matek-19 datasets, the proposed method achieves classification accuracies of 81.1% and 95.8%, respectively, significantly outperforming existing baselines.

Abstract

Federated learning (FL) is severely hindered by statistical heterogeneity due to variations in scanners, acquisition protocols, and patient populations. Such non-IID data induces client drift during local optimization, leading to unstable convergence and suboptimal global models when parameter-based aggregation is applied. We propose a prototype-based, influence-aware federated learning framework (FedProIn) that uses multiple learnable class prototypes to capture shared semantic structures across heterogeneous clients. We introduce feature divergence loss and prototype contrastive loss to mitigate client drift by decomposing it into feature drift and prototype drift. In addition, we propose a normalized influence aggregation strategy that adaptively weights client prototypes according to their contribution to the global representation, reducing the impact of biased or low-quality updates. Experimental results on two publicly available medical datasets, HAM10000 and Matek-19, demonstrate that FedProIn achieves accuracies of (83.5% IID, 81.1% non-IID) on HAM10000 and (96.2% IID, 95.8% non-IID) on Matek-19, respectively, outperforming existing baselines in both conditions. Our code is available at https://github.com/harsh-kmr/FedProIn.

Problem

• federated learning
• client drift
• statistical heterogeneity
• non-IID data
• medical imaging

Innovation

• prototype-based federated learning
• client drift mitigation
• feature divergence loss
• prototype contrastive loss
• influence-aware aggregation