Efficient Federated Learning Mechanism Based on Layer-wise Model Pruning
Zhijie Jiang, Zhe Zhang, Yanchao Zhao · 2024
As a computing paradigm tailored for resource-constrained client devices, federated learning based on model pruning compresses the model size by removing unimportant parameters in the neural network, which has shown outstanding results in improving model efficiency and reducing computing costs. However, previous works simply customized a unified static model pruning rate, ignoring the heterogeneous capabilities of clients and the impact of pruning on different layers of the model during continuous iteration. In this paper, we design a novel Federated learning framework based on Dynamic Layer-wise Pruning, named FedDLP, which is capable of pruning at the hierarchical level depending on the client’s capability and model similarity to improve model efficiency and maintain model performance. This framework consists of two parts. First, pre-training customizes the initial pruning rate: We set the initial pruning rate for each layer according to the different capabilities of heterogeneous clients during the pre-training stage. Second, adaptively optimize the pruning rate: We use cosine similarity to quantify the contribution of each layer of the client model to the global model, thereby adaptively and dynamically optimizing the model pruning rate. Experimental results verify that our proposed method improves model efficiency by 2 to 3.5× compared to the state-of-the-art baselines, while achieving an accuracy difference of no more than 2% compared to the unpruned model.