Lightweight and Personalized Federated Learning via Distribution Similarity

• 1、School of Computer and Information Technology, Beijing Jiaotong University
• 2、 Engineering Research Center of Network Management Technology for High-Speed Railway of Ministry of Education, Beijing Jiaotong University

Abstract：Traditional approaches in federated learning (FL) attempt to collaboratively learn a shared global model under the orchestration of a central server. However, clients usually only have limited communication bandwidth, and in the case of heterogeneous data, learning a single global model may not be suitable for all clients. Such statistical heterogeneity and communication bottlenecks are two key issues that hinder the development of federated learning. To solve the above problems, this paper proposes a lightweight personalized federated learning framework based on distribution similarity. Without sacrificing data privacy, clients with similar distribution were grouped together by analyzing the cosine distance between client data subspaces, so that clients in the same cluster could benefit from each other. By applying unstructured pruning to find a small subnetwork for each cluster, the communication cost can be significantly reduced due to the compact size of the sparse subnetwork. Experimental results show that the proposed method outperforms the existing methods, achieving up to 33.73% improvement in inference accuracy and 2.33 times reduction in communication cost.

Keywords： Federated learning Personalization Clustering Distribution similarity Communication efficiency Data heterogeneity