Prototype-based fine-tuning for mitigating data heterogeneity in federated learning

In federated learning with data heterogeneity, the global model often exhibits a severe imbalance in fitting data from different categories, and clients may not be able to obtain useful information from the impaired global model. To address this challenge, Federated Learning Based on Model Repair (FedMR) is proposed to repair the global model by a set of prototypes with minimal divergence. The repair step of FedMR is executed after global aggregation and before local training. Different clients first obtain similar local prototypes on the same feature extractor, and then fine-tune the global classifier with these local prototypes. The repaired classifier is aggregated at the server and broadcast to all clients, enabling them to start local training from a consensus point. This approach effectively mitigates the adverse effects of uneven sample distribution. In most experimental configurations, FedMR outperforms the state-of-the-art federated learning algorithms.