Partial Model Pruning and Personalization for Wireless Federated Learning

To address the heterogeneity of devices' data and guarantee high computation and communication efficiency of federated learning (FL), we consider an FL framework with partial model pruning and personalization. This framework splits the learning model into a global part with model pruning shared with all devices to learn data representations and a personalized part to be fine-tuned for a specific device. It can adapt the model size during FL to reduce both computation and communication latency and increases the learning accuracy for devices with non-independent and identically distributed data. The computation and communication latency and convergence of the proposed FL framework are analyzed. To maximize the convergence rate and guarantee learning accuracy, Karush-Kuhn-Tucker (KKT) conditions are deployed to jointly optimize the pruning ratio and bandwidth allocation. Finally, experimental results demonstrate that the proposed FL framework achieves a remarkable reduction of approximately 50% computation and communication latency compared with FL with partial model personalization.