FedMap: Iterative Magnitude-Based Pruning for Communication-Efficient Federated Learning

TL;DR

FedMap is introduced, a novel method that aims to enhance the communication efficiency of FL deployments by collaboratively learning an increasingly sparse global model through iterative, unstructured pruning through iterative, unstructured pruning.

Abstract

Federated Learning (FL) is a distributed machine learning approach that enables training on decentralized data while preserving privacy. However, FL systems often involve resource-constrained client devices with limited computational power, memory, storage, and bandwidth. This paper introduces FedMap, a novel method that aims to enhance the communication efficiency of FL deployments by collaboratively learning an increasingly sparse global model through iterative, unstructured pruning. Importantly, FedMap trains a global model from scratch, unlike other methods reported in the literature, making it ideal for privacy-critical use cases such as in the medical and finance domains, where suitable pre-training data is often limited. FedMap adapts iterative magnitude-based pruning to the FL setting, ensuring all clients prune and refine the same subset of the global model parameters, therefore gradually reducing the global model size and communication overhead. The iterative nature of FedMap, forming subsequent models as subsets of predecessors, avoids parameter reactivation issues seen in prior work, resulting in stable performance. In this paper we provide an extensive evaluation of FedMap across diverse settings, datasets, model architectures, and hyperparameters, assessing performance in both IID and non-IID environments. Comparative analysis against the baseline approach demonstrates FedMap's ability to achieve more stable client model performance. For IID scenarios, FedMap achieves over $90$\% pruning without significant performance degradation. In non-IID settings, it achieves at least $~80$\% pruning while maintaining accuracy. FedMap offers a promising solution to alleviate communication bottlenecks in FL systems while retaining model accuracy.

Figures (12)

• Figure 1: The stepwise and continuous pruning schedules exemplified. We design the schedules, s.t. the model can be trained to a good initial solution at the start. After $\mathbf{s}$ FL rounds, clients prune the global model further, s.t. $(1-p_g)\cdot d$ parameters remain. We unanimously use an exponential series, where $p_g = 0.25$, progressively removing 25% of the remaining model parameters.
• Figure 2: Comparison of pruning schedules trained on CIFAR10. Top row is with Resnet56 and the bottom row is with MobileNet.
• Figure 3: Accuracy vs. rounds across variable number of local epochs ($L$) and step-widths ($s$) for CIFAR10.
• Figure 4: Average performance per pruning step ($s$) across variable different values of local epochs ($L$) for CIFAR10.
• Figure 5: Communicated bits (Mb) per round of FL for exp. B. on CIFAR10.