Data Heterogeneity-Aware Personalized Federated Learning for Diagnosis

Huiyan Lin, Heng Li, Haojin Li, Xiangyang Yu, Kuai Yu, Chenhao Liang, Huazhu Fu, Jiang Liu

Abstract: Personalized federated learning is an extension of federated learning that aims to improve prediction accuracy for diverse clients by tailoring models to their individual data. However, the inherent agnosticism of the data across clients poses a challenge in the awareness of the client data characteristics, impacting the effectiveness of personalization. To overcome this challenge, we propose a data heterogeneity-aware algorithm for personalization in federated learning, which involves assessing the heterogeneity across client data using uncertainty. Specifically, a heterogeneity weight is determined based on the predictive uncertainty of the global model on client-specific data. Subsequently, an adaptive fusion of the global model and the previous client model is enabled using the heterogeneity weight to personalize the initialization of the client model training in each iteration. Experiments conducted on diagnosis in two imaging modalities, particularly under non-independent and identically distributed (non-IID) scenarios, demonstrate the superior performance of our proposed algorithm compared to state-of-the-art approaches.