Model Pruning Enables Efficient Federated Learning on Edge Devices

Yuang Jiang; Shiqiang Wang; Victor Valls; Bong Jun Ko; Wei-Han Lee; Kin K. Leung; Leandros Tassiulas

Source-linked AI summary

Model Pruning Enables Efficient Federated Learning on Edge Devices

Yuang Jiang, Shiqiang Wang, Victor Valls, Bong Jun Ko, Wei-Han Lee, Kin K. Leung, Leandros Tassiulas

arXiv:1909.12326v5Sep 26, 2019cs.LGcs.DCstat.ML

TL;DR

Federated learning must train models on edge devices with limited computation and communication resources. PruneFL combines adaptive and distributed pruning with FL to adapt model size during training. Experiments report substantially reduced training time while achieving accuracy similar to the original model.

Problem
Edge clients have limited computation and communication resources, while existing compression methods do not fully provide efficient computation and automatic model-size adaptation.
Method
PruneFL performs initial pruning at a selected client and further adaptive pruning during FL, adapting model size using training information.
Results
PruneFL significantly reduces training time, and its automatically sized pruned model converges to accuracy very similar to the original model.
Takeaways & Limitations
PruneFL reduces FL cost on edge devices while producing a small model for efficient inference and a lottery ticket of the original model.

Abstract

from arXiv · show

Federated learning (FL) allows model training from local data collected by edge/mobile devices while preserving data privacy, which has wide applicability to image and vision applications. A challenge is that client devices in FL usually have much more limited computation and communication resources compared to servers in a datacenter. To overcome this challenge, we propose PruneFL -- a novel FL approach with adaptive and distributed parameter pruning, which adapts the model size during FL to reduce both communication and computation overhead and minimize the overall training time, while maintaining a similar accuracy as the original model. PruneFL includes initial pruning at a selected client and further pruning as part of the FL process. The model size is adapted during this process, which includes maximizing the approximate empirical risk reduction divided by the time of one FL round. Our experiments with various datasets on edge devices (e.g., Raspberry Pi) show that: (i) we significantly reduce the training time compared to conventional FL and various other pruning-based methods; (ii) the pruned model with automatically determined size converges to an accuracy that is very similar to the original model, and it is also a lottery ticket of the original model.

I. INTRODUCTION

Federated learning must train models across resource-constrained clients despite privacy, bandwidth, and computation limits. PruneFL addresses this gap with adaptive, distributed pruning that reduces training and inference costs while retaining similar accuracy.

Client devices have limited computation, communication bandwidth, memory, and storage, making DNN training potentially time- and energy-intensive.
Existing compression methods either risk reducing final accuracy or reduce communication without producing a smaller model for efficient computation and inference.
Automatic compressed-model size adaptation remains largely unexplored because training dynamics are unpredictable and efficient solutions must be found quickly with minimal overhead.
PruneFL extends federated learning with adaptive and distributed parameter pruning during the FL procedure.
Experiments on real edge devices show that PruneFL outperforms single-client and multi-client pruning alternatives, especially with heterogeneous data and computational power.
PruneFL jointly targets communication and computation efficiency during training and inference, while its two-stage pruning addresses statistical and device heterogeneity.

IV. PRUNEFL

PruneFL begins by pruning at one selected client, then performs adaptive reconfiguration during federated training using updates from participating clients. The model can grow or shrink as needed to improve training efficiency.

Two-stage Distributed Pruning: PruneFL uses initial pruning at a selected client followed by further pruning involving the server and clients during FL.
Two-stage Distributed Pruning: Initial pruning starts FL with a small model, reducing computation and communication time per round.
Two-stage Distributed Pruning: Further pruning uses all participating clients’ data and is integrated with FedAvg, allowing the model to grow or shrink for efficient training.

B. Adaptive Pruning

PruneFL repeatedly reconfigures the model during federated learning, selecting parameters to retain by balancing approximate empirical-risk reduction against round time. This adaptive process can shrink or regrow the model while preserving convergence under stated assumptions.

PruneFL reconfigures the model at intervals, removing or adding parameters before continuing training with the resulting model and mask.Reconfiguration occurs at a selected client during initial pruning and at the server between federated-learning rounds during further pruning.
The method estimates risk reduction using the sum of squared stochastic-gradient components, treating g_j^2 as parameter importance.The approximation uses a first-order expansion and ignores proportionality terms independent of relative parameter importance.
During federated learning, further pruning uses aggregated client gradients and can adapt the model size using information from all participating clients.The initial model may be suboptimal because it uses one client’s data, whereas further pruning incorporates all participating clients.
It models round time as a fixed overhead plus the summed time contributions of retained parameter components.The parameter time coefficient may depend on neural-network layer, and the linear model is supported by empirical observations.
Among parameters eligible for pruning or restoration, PruneFL maximizes approximate risk reduction per unit training time while retaining protected parameters.The protected set typically contains parameters whose removal could substantially increase empirical risk after reconfiguration.
The resulting optimization is solved by sorting g_j^2/t_j ratios, giving global optimality with O(|P| log |P|) time complexity.Adding a parameter improves the overall ratio when its individual ratio exceeds the current ratio.
Under the stated assumptions and learning-rate condition, the adaptive procedure has a convergence bound, with remaining-parameter gradients vanishing asymptotically.The decreasing number of nonzero prunable parameters supports eventual stabilization of the parameter vector.

A. Using Sparse Matrices

PruneFL uses actual sparse matrices for pruned models rather than masking dense parameters, addressing implementation overhead while evaluating efficiency across edge-device settings and FL tasks.

PruneFL uses dense matrices for full-sized models and sparse matrices for pruned convolutional and fully connected weights.This implementation is intended to realize computational benefits that binary masking of dense parameters does not provide.
Sparse storage uses bitmap or value-index tuples, with bitmap storage and communication overhead of 1/32 for 32-bit parameters.When sparsity patterns remain fixed, only nonzero values need to be exchanged, eliminating additional pattern cost.
Sparse computation is beneficial only below a hardware- and software-dependent density threshold, so the implementation selects dense or sparse representation accordingly.Sparse matrix multiplication is described as linear in the number of nonzero entries for the relevant operation.
The evaluation uses four image-classification tasks with Conv-2, VGG-11, ResNet-18, and MobileNetV3-Small models.FEMNIST and CelebA use non-IID client partitions, while CIFAR-10 and ImageNet-100 use IID partitions.
Experiments run on a prototype with Raspberry Pi clients and a CPU-only personal computer, plus a simulation using measured edge-device timings.The prototype includes nine Raspberry Pi 4 devices and one personal computer across ten clients.
The experiments compare test accuracy over time against conventional FL, iterative pruning, online learning, SNIP, and SynFlow.Enhanced baseline variants match PruneFL’s converged model size when the original methods cannot determine size automatically.

A. Time Measurement

PruneFL reduces federated training cost on edge devices through sparse computation, lower communication, and adaptive pruning while retaining accuracy comparable to conventional FL.

Computation time: 11.24 to 6.34 seconds per round: decreasing model density reduces computation time for five local iterations.The reduction is moderate because sparse computation only partially improves backward passes.
Communication time: 35.88 to 1.04 seconds per round: sparse models reduce communication time more noticeably than computation time.PruneFL exchanges pruned models except during periodic reconfiguration rounds, yielding average communication cost 0.02 + 1.96d.
Training time reduction: PruneFL outperforms conventional FL by a significant margin on FEMNIST in both prototype and simulation test-accuracy-versus-time results.Prototype and simulation results match closely; initial pruning takes less than 500 seconds and is negligible compared with further pruning.
Training time reduction: PruneFL consistently accelerates training across datasets, models, and baselines while converging to accuracy similar to conventional FL.SNIP can fail to reach conventional FL accuracy on CIFAR-10, while online learning and SynFlow can remain at random-guess accuracy.
Time and FLOPs to reach target accuracy: Less than one-third of conventional FL’s time is required to reach 80% accuracy on FEMNIST, with similar savings in accumulated FLOPs per client.PruneFL also saves more than 33% of time, exceeding two hours, compared with SNIP and SynFlow.
Comparing with additional baselines: 17.65 seconds per round: applying PruneTrain’s regularization increases computation time by 57% compared with PruneFL.This comparison uses FEMNIST with the Conv-2 model.

C. Finding a Lottery Ticket

PruneFL combines initial and further pruning to find a lottery ticket while adapting model size during training. Experiments show that it avoids single-stage drawbacks and discovers task-dependent densities, although manually selecting pruning density can harm convergence.

Lottery-ticket results: PruneFL with both initial and further pruning finds a lottery ticket whose reinitialized model learns comparably fast to the original FL model.The initial-pruning-only model is not a lottery ticket, whereas further pruning alone or both stages find one.
Lottery-ticket results: PruneFL’s two-stage design avoids the suboptimal accuracy of initial-only pruning and the slower learning speed of further-only pruning.Both stages are needed to avoid the respective drawbacks reported for the single-stage alternatives.
Model-size adaptation: 13.4% density is reached by overparameterized Conv-2 on FEMNIST, while underparameterized ResNet-18 on ImageNet-100 converges around 67.7%.These densities illustrate that PruneFL adapts model size differently across architectures and datasets.
Model-size adaptation: SNIP pruning Conv-2 to 30% or 1% density slows training, and the 1% model fails to converge to the original model’s accuracy.PruneFL automatically determines a proper density instead of relying on a manually selected target.
Limited and targeted sizes: With density limits dl = 10% and dt = 5%, constrained PruneFL keeps model size below 10%, reaches 5% at training end, and retains fast convergence with similar accuracy.The maximum model density is scheduled to decrease linearly toward the target.
Layer-wise pruning: PruneFL automatically assigns higher convergence densities to input and output layers than to many other layers, indicating their relative importance.This layer-wise pattern is reported for VGG-11 on CIFAR-10.

APPENDIX

For general positive monotone time functions, the appendix analyzes a greedy solution to the pruning objective. The resulting guarantee is local rather than globally optimal, and multiple local optima may exist.

General T(M): Algorithm A.1 solves the pruning objective for a general monotone positive T(M) with complexity O(|P|^2).The algorithm is presented as a local optimization procedure for the general case.
General T(M): The general-T(M) result is weaker than the linear case because multiple local optimal solutions may exist.The appendix attributes the weaker guarantee to difficulty expressing cost changes when sets differ by more than one element.

2) Proof of Theorem 2:

The proof of Theorem 2 derives convergence for federated learning with masked parameters under smoothness, variance, independence, and pruning-rate assumptions. The resulting bound includes a client-count acceleration term and depends on exponentially decreasing prunable parameters.

Proof setup: The analysis extends a prior convergence theorem to masked gradients or function values, treating masking as replacing selected parameter entries with zeros.The proof introduces pruned-value notation for the masked quantities.
Proof setup: The local and averaged parameter update rules are combined with smoothness, Lipschitzness, unbiasedness, variance, and client-independence assumptions.The proof repeatedly bounds terms using these stated assumptions and Jensen’s inequality.
Pruning perturbation: The proof’s pruning perturbation decreases geometrically because the fraction of non-zero prunable parameters halves every h iterations.Added-back parameters are assigned zero values, so the difference between successive reconfigured vectors is attributed to removed parameters.
Convergence result: A convergence bound is obtained under additional conditions including K ≥ 24NI^2β^2, bounded iterates, and exponentially decreasing non-zero prunable parameters.The bound is stated in (B.22) and is dominated by an asymptotic term.
Convergence result: Using more clients can accelerate convergence by a factor involving 1/√N under the theorem’s assumptions.This interpretation follows the discussion immediately after the convergence bound.

1) Model architecture details:

The appendix points readers to a table containing the model architectures used in the experiments. The architecture details are provided separately from the proof material.

Model architecture details: Model architecture specifications are listed in Table C.1.The passage directs the reader to the appendix table for these details.
Model architecture details: Table C.1 serves as the appendix reference for the architectures used in the study.The supplied passage identifies the table but does not enumerate its layers or hyperparameters.
Model architecture details: The architecture details are separated from the theorem-proof discussion in the appendix.Only the location of the specifications is stated here.

2) Gradient Computation:

Sparse matrix operations support PruneFL’s forward and backward computations, but only part of backpropagation gains computational speed from sparsity.

The forward pass multiplies dense input x by sparse weight u to produce dense output y.For a fully connected layer, y = x·u.
The weight and input gradients are gu = xTgy and gx = gyuT, respectively.x and gy are dense backpropagation quantities, while u is sparse.
Sparse matrix implementation accelerates computation of gx when the weight u is sparse.
Weight-gradient computation does not improve in speed because current implementations first form a dense gradient and then apply u’s sparse pattern.The dense computation includes zero-valued entries before selecting the retained values.

3) FLOPs Computation:

PruneFL reconfigures the model during federated learning, accumulates parameter importance across rounds, and evaluates efficiency, accuracy, lottery-ticket behavior, and model-size changes across datasets.

FLOPs Computation:: Reconfiguration begins when selected-client training accuracy exceeds 1.5 times random-guess accuracy.This enables earlier pruning for easy tasks while avoiding pruning during near-random initialization.
FLOPs Computation:: Importance measures are summed locally after each update and sent to the server in the next reconfiguration round.Non-reconfiguration rounds compute and exchange only the remaining fixed parameter set.
FLOPs Computation:: Training time within each layer is approximately linear in the number of parameters when sparse implementation is used.The experiments vary parameters in selected fully connected and convolutional layers while fixing other layers.
FLOPs Computation:: The random-client-selection evaluation covers CIFAR-10, ImageNet-100, FEMNIST, and CelebA under dataset-specific client partitions.The figures report test accuracy versus time, lottery-ticket results, and parameter count versus round for four datasets.
FLOPs Computation:: Convergence accuracy for PruneFL is similar to conventional FL, although conventional FL can have a slight advantage when accuracy is compared after the same number of rounds.The reported accuracies average the last five measured evaluations.