Clustered Federated Learning Based on Client's Prototypes

Weimin Lai, Zirong Xu, Qiao Yan

Abstract: Federated learning empowers multiple parties to train machine learning models collaboratively while preserving data within local confines. However, due to the intrinsic non-independent and identically distributed (Non-IID) nature of client-side data, aggregating these diverse local models into a global model remains a significant challenge. Clustered Federated Learning (CFL) offers a compelling solution to mitigate the impact of data heterogeneity by organizing clients into clusters. Existing CFL methods often involve unstable, time-consuming cluster identity estimation during training. In this work, we propose an innovative federated learning approach that efficiently identifies similarities among client data distributions by analyzing client’s prototypes. Our method accurately identifies client cluster identities during the initialization phase. Clients employ the same randomly initialized model to compute the client’s prototypes for their local data and provide it to the server. The server performs one-shot client clustering by comparing the similarity of the client’s prototypes. We further design a unique inter-cluster aggregation strategy that adapts inter-cluster aggregation weights based on the similarity of the client’s prototypes, thereby enhancing model convergence. We evaluate two mixed datasets with three Non-IID settings, and our approach outperforms several popular baseline methods. Compared to IFCA and FlexCFL, our approach yields more reasonable clustering results. It achieves an average test accuracy improvement of over 11.51% in the feature and label shift Non-IID setting of the Digits-5 dataset.