Personalized federated learning with dynamic clustering and model distillation

Abstract

Federated learning is a distributed machine learning technique that allows various data sources to work together to train models while keeping their raw data private. However, federated learning faces many challenges when dealing with non-independent and identically distributed (Non-IID) data, especially the problem of data heterogeneity, which can significantly degrade model performance. To address this challenge, we propose a new algorithm for personalized federated learning, known as pfedCluster. The core of the pfedCluster algorithm is to dynamically cluster clients using hierarchical tree clustering, which ensures minimal intra-cluster distance and maximal inter-cluster distance, thus optimizing the clustering effect. Additionally, the algorithm facilitates knowledge transfer between clusters through knowledge distillation, further enhancing model performance. This method improves model personalization by dynamically adjusting the clustering structure to suit varying data distributions. Experimental results show that pfedCluster effectively improves model performance on MNIST and CIFAR-10 datasets, demonstrating significant advantages in dealing with data heterogeneity compared to traditional federated learning algorithms. Our code is at https://github.com/NtuEEEJackie/pFedCluster.