FedPer++: Toward Improved Personalized Federated Learning on Heterogeneous and Imbalanced Data

Abstract

Federated learning is an emerging technique to collaboratively train machine learning models over multiple clients without exposing private data but suffers from heterogeneous data distributions across clients, which results in slow convergence and degraded model performance. This challenge has motivated several approaches to learn a personalized model that has better model accuracy than the global model for each participating client. On the other hand, the success of multi-task learning suggests that data in similar tasks often share a common feature representation, while the output layer (classifier) for classification is more task-correlated. Existing methods usually neglect the collaboration between locally trained classifiers, which makes each client failed to fully benefit from data in other clients. Therefore, we perform a linear combination based classifier collaboration method for achieving better personalized model performance. Moreover, training a local feature extractor based on local data is prone to over-fitting when local data is insufficient and cannot be fully corrected by global aggregation. To tackle this issue, we propose a feature-regularized training strategy to mitigate the local over-fitting risk and reduce the parameter divergence across clients, facilitating the global feature extractor aggregation. With extensive experiments performed on Fashion-MNIST and CIFAR-10 datasets, we demonstrate the performance improvement and robustness of our method over existing methods.

Topics & Concepts