FedPaI: Achieving Extreme Sparsity in Federated Learning via Pruning at Initialization
Haonan Wang, Zeli Liu, Kajimusugura Hoshino, Tuo Zhang, John Paul Walters, Stephen Crago
TL;DR
This work tackles the resource bottlenecks of federated learning on edge devices by introducing FedPaI, a framework that applies Pruning at Initialization (PaI) to fix sparse connectivity patterns at the start of training. It integrates both unstructured PaI with client-side personalization (FedPaI-U) and structured PaI on the server (FedPaI-S) with sparsity-aware aggregation, enabling extreme sparsity (up to 98%) while preserving accuracy across IID and non-IID data. Empirical results on CIFAR-10 with VGG19/ResNet18 show FedPaI outperforms existing efficient FL methods and delivers substantial speedups (6.4–7.9x) due to reduced communication and compute, with robustness to heterogeneous data. The framework’s flexibility—supporting both pruning paradigms and adapting to hardware constraints—positions it as a scalable solution for efficient FL in resource-constrained edge environments.
Abstract
Federated Learning (FL) enables distributed training on edge devices but faces significant challenges due to resource constraints in edge environments, impacting both communication and computational efficiency. Existing iterative pruning techniques improve communication efficiency but are limited by their centralized design, which struggles with FL's decentralized and data-imbalanced nature, resulting in suboptimal sparsity levels. To address these issues, we propose FedPaI, a novel efficient FL framework that leverages Pruning at Initialization (PaI) to achieve extreme sparsity. FedPaI identifies optimal sparse connections at an early stage, maximizing model capacity and significantly reducing communication and computation overhead by fixing sparsity patterns at the start of training. To adapt to diverse hardware and software environments, FedPaI supports both structured and unstructured pruning. Additionally, we introduce personalized client-side pruning mechanisms for improved learning capacity and sparsity-aware server-side aggregation for enhanced efficiency. Experimental results demonstrate that FedPaI consistently outperforms existing efficient FL that applies conventional iterative pruning with significant leading in efficiency and model accuracy. For the first time, our proposed FedPaI achieves an extreme sparsity level of up to 98% without compromising the model accuracy compared to unpruned baselines, even under challenging non-IID settings. By employing our FedPaI with joint optimization of model learning capacity and sparsity, FL applications can benefit from faster convergence and accelerate the training by 6.4 to 7.9 times.