Personalized Federated Learning by Domain-Aware Network Pruning and Re-growth

Federated learning (FL) is a machine learning paradigm where multiple clients train their local machine learning models collaboratively (without sharing private data). One of the main challenges in FL is statistical heterogeneity of the data distributions across clients. Personalized FL (PFL) mitigates statistical heterogeneity by collaborative model training across homogeneous clients. In this paper, we propose a novel personalized federated learning by domain-aware network pruning and re-growth, called FedDNPR, that is more accurate as compared to existing PFL methods while maintaining high efficiency. This is achieved by 1) introducing a regularization term capturing heterogeneity of weights in iterative network pruning in order to reduce network sharing among unrelated clients, and 2) iterative network re-growing only from weights of related clients to increase network sharing among related clients. With FedDNPR, model clustering is performed considering the similarity between gradient updates in the last layers of the networks, with cosine as the similarity measure to achieve both accuracy and efficiency in model personalization. With extensive experimental evaluation, we show that FedDNPR significantly outperforms the state-of-the-art PFL approaches, while maintaining comparable efficiency.