FedPaI: Achieving Extreme Sparsity in Federated Learning via Pruning at Initialization

2025-04-01

AI Summary

To address high communication and computation overhead in federated learning (FL) on resource-constrained edge devices—and the incompatibility of existing iterative pruning methods with decentralized, non-IID data—this paper proposes FedPaI, a “pruning-at-initialization” framework. FedPaI is the first to introduce the Pruning-at-Initialization (PaI) paradigm into FL: it identifies personalized sparse architectures at initialization and freezes connection patterns throughout training. It supports both structured and unstructured pruning, and introduces sparse-aware server aggregation and non-IID-adaptive optimization. Experiments under non-IID settings demonstrate lossless accuracy at up to 98% sparsity, alongside 6.4–7.9× speedup in training time. FedPaI significantly outperforms state-of-the-art iterative pruning approaches in efficiency, accuracy, and adaptability to heterogeneous data distributions.

Technology Category

• Machine Learning: Distributed Machine Learning & Federated Learning
• Search and Optimization: Learning to Search
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Systems and Infrastructure for Web, Mobile and WoT: Federated Web and WoT systems, including distributed, federated and edge-based data processing
• Security and Privacy: Security and privacy of machine learning and AI applications

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.

Problem

• Achieving extreme sparsity in Federated Learning via pruning
• Reducing communication and computation overhead in FL
• Adapting pruning to diverse hardware and software environments

Innovation

• Pruning at Initialization for extreme sparsity
• Supports structured and unstructured pruning
• Personalized client-side pruning mechanisms