Personalized federated learning with prototype representation and robust aggregation
摘要
Federated learning allows multiple clients to train models locally while ensuring client data privacy. However, the non-independent and identically distributed (non-IID) data and the model heterogeneity distributed on the clients bring challenges. To address this challenge, the deep neural network is decoupled into a feature extractor and classifier respectively, to train a personalized model for each client and facilitate more efficient parameter training. However, feature bias occurs when extracting features from the non-IID data, resulting in poor feature quality and affecting the model convergence speed. In addition, in the process of aggregating the classifier, only the number of data samples is considered, and the differences between local and global classifier model parameters are ignored, resulting in poor model performance. To solve these problems, this paper designs a personalized federated learning framework with prototype representation and robust aggregation for heterogeneous clients. Firstly, the representation of the prototype features of each class is calculated, the prototype loss between the prototype features and the sample features is constructed, and the joint loss function is further constructed to improve the discrimination ability of the model. Secondly, the Euclidean distance is used to evaluate the difference between the parameters of the local and global classifier models, and the aggregate weights are adjusted for all local classifiers to generate a global classifier. Finally, the personalized federated learning algorithm will be constructed based on the prototype feature representation and robust aggregation strategy adjustment mechanism. Experiments on three datasets show that our algorithm performs excellently on heterogeneous personalized federated learning tasks.