FedPE: Adaptive Model Pruning-Expanding for Federated Learning on Mobile Devices
Liping Yi, Xiaorong Shi, Nan Wang, Jinsong Zhang, Gang Wang, Xiaoguang Liu
Abstract
Recently, federated learning (FL) as a new learning paradigm allows multi-party to collaboratively train a shared global model with privacy protection. However, vanilla FL running on heterogeneous mobile edge devices still faces three crucial challenges: communication efficiency, statistical heterogeneity, and system heterogeneity. <italic xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">To tackle them simultaneously, we devise</i> <monospace xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">FedPE</monospace> <italic xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">, a communication-efficient and personalized federated learning framework, which allows each client to search for personalized optimal local subnets adaptive to system capacity in each round of FL.</i> It consists of three core components: a) <italic xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">adaptive pruning-expanding</i> controls model pruning or expanding according to the accuracy variations of local models, b) <italic xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">error compensation strategy</i> promotes the pruned or expanded subnets to be Lottery Ticket Networks (LTNs), c) the <italic xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">fair aggregation rule</i> aggregates local models with their real-time contributions as coefficients to boost the performance of the aggregated global model. The integration of the three components facilitates that only <italic xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">personalized optimal subnets with different footprints</i> interact between the server and clients, which effectively reduces communication costs and enhances the robustness of FL to statistical and system heterogeneity. We also prove the convergence of <monospace xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">FedPE</monospace> and design an optimal hyperparameter searching (OHS) algorithm based on <italic xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">Pareto optimization</i> to search for optimal hyperparameters for <monospace xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">FedPE</monospace> . Extensive experiments evaluated on five real-world datasets with IID or Non-IID distributions demonstrate that <monospace xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">FedPE</monospace> configured with found optimal hyperparameters achieves <inline-formula xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink"><tex-math notation="LaTeX">$1.86\times -121\times$</tex-math></inline-formula> communication efficiency improvement with almost no accuracy degradation, presenting the best trade-off between model accuracy and communication cost.