L OTTERY A WARE S PARSITY H UNTING : E NABLING F EDERATED L EARNING ON R ESOURCE -L IMITED E DGE
Anonymous authors
Paper under double-blind review
A BSTRACT
Limited computation and communication capabilities of clients pose significant challenges in federated learning (FL) over resource-limited edge nodes.
A potential solution to this problem is to deploy off-the-shelf sparse learning algorithms that train a binary sparse mask on each client with the expectation of training a consistent sparse server mask yielding sparse weight tensors.
However, as we investigate in this paper, such naive deployments result in a significant drop in accuracy compared to FL with dense models, especially for clients with limited resource budgets.
In particular, our investigations reveal a serious lack of consensus among the trained sparsity masks on clients, which prevents convergence for the server mask and potentially leads to a substantial drop in model performance.
Based on such key observations, we propose federated lottery aware sparsity hunting (FLASH), a unified sparse learning framework to make the server win a lottery in terms of yielding a sparse sub-model, able to maintain classification performance under highly resource-limited client settings.
Moreover, to support FL on different devices requiring different parameter density, we leverage our findings to present heteroFLASH , where clients can have different target sparsity budgets based on their device resource limits.
Experimental evaluations with multiple models on various datasets (both IID and non-IID) show superiority of our models in closing the gap with unpruned baseline while yielding up to ∼ 10 .
1% improved accuracy with ∼ 10 .
26 × fewer communication costs, compared to existing alternatives, at similar hyperparameter settings.
Code is released as Supplementary .
I NTRODUCTION
Federated learning (FL) McMahan et al.
(2017) is a popular form of distributed training, which has gained significant traction due to its ability to allow multiple clients to learn a shared global model without the requirement to transfer their private data.
However, clients’ heterogeneity and resource limitations pose significant challenges for FL deployment over edge nodes, including mobile phones and IoT devices.
To resolve these issues, various methods have been proposed over the past few years including efficient learning for heterogeneous collaborative training Lin et al.
Zhu et al.
(2021), distillation He et al.
federated dropout techniques Horvath et al.
; Caldas et al.
(2018b), efficient aggregation for faster convergence and reduced communication Reddi et al.
Li et al.
(2020b).
However, these methods do not necessarily address the growing concerns of highly computation and communication limited edge.
Meanwhile, reducing the memory, compute, and latency costs for deep neural networks (DNNs) in centralized training for their efficient edge deployment has also become an active area of research.
In particular, recently proposed sparse learning (SL) strategies Evci et al.
Kundu et al.
(2021b);
Mocanu et al.
(2018); Dettmers & Zettlemoyer (2019); Raihan & Aamodt (2020) effectively train weights and associated binary sparse masks to allow only a fraction of model parameters to be updated during training, potentially enabling the lucrative reduction in both the training time and compute cost Qiu et al.
; Raihan & Aamodt (2020), while creating a model to meet a target parameter density denoted as d, and is able to yield accuracy close to that of the unpruned baseline .
However, the challenges and opportunities of sparse learning in FL is yet to be fully unveiled.
Only very recently, few works Bibikar et al.
; Huang et al.
(2022) have tried to leverage sparse learning in FL primarily to show their efficacy in non-IID settings.
Nevertheless, these works primarily used sparsity for non-aggressive model compression, limiting the actual benefits of sparse learning, and required multiple local epochs, that may further increase the training time for stragglers making the overall FL process inefficient Zhang et al.
Moreover, the server-side pruning used in these methods may not necessarily adhere to the layers’ pruning sensitivity 1 that often plays a crucial role in sparse model performance Kundu et al. (2021b)
; Zhang et al.
Another recent work, ZeroFL Qiu et al.
(2021), has explored deploying sparse learning in FL settings.
However,
Qiu et al.
(2021) could not leverage any advantage of model sparsity in the clients’ communication cost and had to keep significantly more parameters active compared to a target d to yield good accuracy.
Moreover, as shown in Fig. 1(b), for d = 0 . 05 , ZeroFL still suffers from substantial accuracy drop of ∼ 14% compared to the baseline.
Our Contributions.
Our contribution is fourfold.
In view of the above limitations, we first identify crucial differences between a centralized and the corresponding FL model, in learning the sparse masks for each layer.
In particular, we observe that in FL, the server model fails to yield convergent sparse masks, primarily due to the lack of consensus among clients’ later layers’ masks.
In contrast, the centralized model show significantly higher convergence trend in learning sparse masks for all layers.
We then experimentally demonstrate the utility of pruning sensitivity and mask convergence in achieving good accuracy setting the platform to close the performance gap in sparse FL.
We then leverage our findings and present federated lottery aware sparsity hunting (FLASH), a sparse FL methodology addressing the aforementioned limitations in a unified manner.
At the core, FLASH leverages a two-stage FL, a robust and low-cost layer sensitivity evaluation stage and a FL training stage. In particular, the disentangling of the layer sensitivity evaluation from sparse weight training allows us to either choose to train a sparse mask or freeze a sensitivity driven pre-defined mask.
This can further translate to a proportional communication saving.
To deal with the heterogeneity in clients’ compute-budget, we further extend our methodologies to hetero -FLASH, where individual clients can support different density based on their resources.
Here, to deal with the unique problem of the server selecting different sparse models for clients, we present server-side gradual mask sub-sampling, that identifies sparse masks via a form of layer sensitivity re-calibration, starting for models with highest to that with lowest density support.
We conduct experiments on MNIST, FEMNIST, and CIFAR-10 with different models for both IID and non-IID client data partitioning.
Experimental results show that, compared to the existing alternative Qiu et al.
(2021), at iso-hyperparameter settings, FLASH can yield up to ∼ 8 .
9% and ∼ 10 .
1% , on IID and non-IID data settings, respectively, with reduced communication of up to ∼ 10 .
2 × .
R ELATED W ORKS
Model Pruning.
Over the past few years, a plethora of research has been done to perform efficient model compression via pruning, particularly in centralized training Ma et al.
; Frankle &
Carbin (2018); Liu et al. (2021); You et al.
; He et al.
Pruning essentially identifies and removes the unimportant parameters to yield compute-efficient inference models.
More recently, sparse learning Evci et al.
Kundu et al.
(2021b) ; Dettmers & Zettlemoyer (2019); Raihan &
Aamodt (2020); Kundu et al. (2020), a popular form of model pruning, has gained significant traction as it can yield FLOPs advantage even during training.
In particular, it ensures only d % of the model parameters remain non-zero during the training for a target parameter density d ( d < 1 . 0 and sparsity is 100 − d %), potentially enabling training compute and comm. cost if deployed for FL.
Dynamic network rewiring (DNR).
We leverage DNR Kundu et al.
(2021b) , to sparsely learn the sparsity mask of each client.
In DNR, a model starts with randomly initiated mask following the target parameter density d .
After an epoch, the client evenly prunes the lowest p r % weights from each layer based on absolute magnitude, where p r is prune rate.
Note, this p r % pruning happens on top of the sparse model with density d , allowing p r % weights to be regrown.
DNR then ranks each layer based on the normalized contribution of the summed non-zero weight magnitudes.
Finally, the client regrows total p r % weights in a non-uniform way, allowing more regrowth to the layers having higher rank.
This process iteratively repeats over epochs to finally learn the mask.
Federated learning for resource and communication limited edge.
To address device heterogeneity, existing works have explored the idea of heterogeneous training Horvath et al.
Diao et al.
(2020); Yao et al.
(2021) allowing different clients to train on different fractions of full-model based on their compute-budget.
On a parallel track, various optimizations are proposed in FL training framework to accelerate convergence, thus requiring fewer communication rounds Han et al.
Gorbunov et al.
; Zhang et al.
; Li et al.
; Reddi et al.
; Islamov et al.
Albasyoni et al.
To address the issue of client resource limitations, a few research have leveraged pruning in FL Li et al.
(2020a)
; Jiang et al.
; Li et al.
In particular, LotteryFL Li et al.
(2020a) trained each client to have their personalized mask with which they are able to perform well only on their own data.
Moreover, the clients often need to send full model costing bandwidth.
PruneFL
Jiang et al.
(2022) also asks for significant communication costs as it demands participating clients to send all the gradient values to the server while updating the masks.
Only a few contemporary works Huang et al.
; Bibikar et al.
; Qiu et al.
(2021) tried to leverage the benefits of sparse learning in federated settings.
In particular, Huang et al.
(2022) relied on a randomly initialized sparse mask, and recommended keeping it frozen throughout the training, yet failed to provide any supporting intuition.
FedDST Bibikar et al.
(2021), on the other hand, leveraged the idea of RigL Evci et al.
(2020) to perform sparse learning of the clients, relied on a large number of local epochs to avoid gradient noise, and focused primarily on only highly non-IID data without targeting ultra-low density d .
More importantly, neither of these works investigated the key differences between centralized and FL sparse learning.
With similar philosophy as ours, ZeroFL Qiu et al.
(2021) first identified a key aspect of sparse learning in FL in terms of all clients’ masks to be within 30% of the total model weights to yield good accuracy at high compression.
However, ZeroFL suffered significantly in exploiting a proportional advantage in communication saving as even for low parameter density d , all clients had to download the dense model and send back a 3 × denser model.
Furthermore, these algorithms sacrifice significant accuracy at ultra-low d .
R EVISITING S PARSE L EARNING : W HY D OES IT M ISS THE M ARK IN FL?
[Table] [Figure] [Table] [Figure]
[Table] [Figure] [Table] [Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
Sparse learning uses proxies, including normalized momentum and normalized values of the non-zero weights Kundu et al.
(2021b) ; Dettmers & Zettlemoyer (2019), to decide the layers and weights that are more sensitive towards pruning and update the binary sparse mask accordingly.
Note, centralized training has shown significant benefits with sparse learning with FLOPs reduction during forward operations
Evci et al.
(2020), and potential training speed-up of up to 3 .
3 × Qiu et al.
(2021) while maintaining close to the baseline accuracy, even at d ≤ 0 .
We now use a sparse learning, namely
Kundu et al.
(2021b), in FL settings (refer to Table 1 for details) on CIFAR-10, where each client separately performs Kundu et al.
(2021b) to train a sparse server-side ResNet18 and meet a fixed parameter density d , starting from a random sparse mask.
After sending the updates to server, it aggregates them using FedAvg.
We term this as naive sparse training (NST).
Observation 1.
At high compression d ≤ 0 . 1 , the collaboratively learned FL model significantly sacrifices performance, while the centralized sparse learning yields close to baseline performance.
As shown in Fig.2(a), naive deployment of sparse learning significantly sacrifices accuracy in FL.
In particular, for d = 0 . 1 , the trained server-side model suffers an accuracy drop of 3 .
At even lower d = 0 . 05 , this drop significantly increases to 12 .
03% , hinting at serious limitations of sparse learning in FL.
However, an off-the-shelf centralized sparse learning can yield model having close to the baseline accuracy, even at d = 0 . 05 .
Observation 2.
As the training progresses, the sparse masks in centralized training tend to agree across epochs, showing convergence, while the server mask in FL does lack agreement across rounds.
Definition 1.
Sparse mask mismatch.
For a model at round t , we define the sparse mask mismatch (SM) sm t as the Jaccard distance that is measured as follows.
[Equation]
where M tl represents the sparse mask tensor for layer l at the end of round t .
Interestingly, as depicted in Fig. 2(a), the SM for centralized learning tends to zero as the training progresses.
In contrast, with the same model, dataset and d values, in FL, the SM remains > 0 .
4 indicating a substantial distinction in the sparse mask learning between centralized and federated learning.
Observation 3.
At low target density, in federated sparse learning, throughout the training rounds, the disagreement on the later layer’s masks remains more severe than the earlier ones.
As the training progresses, in centralized learning, mask for each layer shows significant convergence
[Table]
as measured by SM for the layer (Fig. 3(a)).
However, Fig. 3(b) shows in FL, the later layers’ masks differ significantly and continue to disagree over rounds with SM value as high as ∼ 0 .
This may be attributed to many possible mask choices due to the later layers’ significantly fewer non-zero parameter allocation, compared to the initial layers, driven by their respective pruning sensitivities.
For example, layer 1 requires 90% parameters to be present, compared to only 5% for layer 14 , with the later costing an SM of ∼ 0 .
To further investigate the impact of higher SM and layer sensitivity on a model’s accuracy, we performed five different training in centralized as described in Table 2.
In particular, for the training in row1 we randomly generate sparse masks with uniform density for all the layers.
For all other training, we first randomly create each layer’s mask by following its pruning sensitivity 2 and then decide to keep the layer mask frozen for some or all the layers.
For training described in rows 3-5, we allow a fraction of the mentioned layers’ masks to differ between consecutive epochs such that they meet the target SM value, creating the situation of non-convergent masks.
As Table 2 clearly shows that large SM for the layers can degrade the accuracy by up to 9 .
34% , we can safely conclude that disagreement of masks across epochs can significantly affect the model’s final performance .
Moreover, the model trained via sparse learning with sensitivity-driven pre-defined masks yields better performance than the one trained with uniform density sparse mask.
FLASH: M ETHODOLOGY
To win a lottery of having a sparse network yielding high accuracy at reduced parameters, we identify two key characteristics of sparse learning, namely, the pruning sensitivity and mask learnability towards convergence.
To explicitly adhere to these two important aspects, in FLASH, we present a In , the continue to have higher SM contrary to centralized, where every layer tend to reduce the SM as the training matures. two-stage sparse FL method, stage 1 : targeting sensitivity analysis to identify good initial sparse mask for each layer, stage 2 : targeting training to learn masks and weights.
In particular, to evaluate layer sensitivity in stage 1 , the server randomly selects a small fraction of clients ([ C d ]), each locally sparse learning Kundu et al.
(2021b) for few warm-up epochs ( E d ) ( L 4-8 in Algo. 1).
Upon collection of layer-wise sensitivity from the clients, for each layer l , the server estimates average density 3 ˆ d l as (cid:80) cdi =1 d li c d , where d li is the density at layer l in i th client.
As these averaged layer-wise density values may not necessarily yield to the target density d , for a model with K parameters we follow the following density re-calibration
[Equation]
k l is dense model’s parameter size for layer l .
For each layer l of the model, the server then creates a binary sparse mask tensor that is randomly initialized, with a fraction of 1s ∝ d lc ( L 10).
In stage, the server begins the training rounds starting with a sparse model initialization following the sparse mask computed in stage 1 ( L 11).
Particularly, at each round, the clients perform sparse learning for E epochs ( L 21-26) with the choice to either train the mask or keep training on the weights with the mask frozen ( L 25).
The later allows the clients to intermittently share masks saving up-link cost.
However, in FL settings, the masks often show poor convergence (section 3, Obs. 2).
To address this, in stage 2 , we present the idea of sparse FL with mask learnability either disabled or enabled ( L 9).
The disabled scenario ( m freez = 1 ) essentially translates to sparse learning with pre-defined layer masks at initialization, allowing all the clients to have a forcefully convergent mask ( sm t =for all t ).
Moreover, as FLASH disentangles the sensitivity evaluation stage from the training, the pre-defined mask in this scenario benefits from the notion of layer sensitivity.
We thus aptly name this scenario sensitivity-driven pre-defined sparse training (SPDST).
Interestingly, earlier research Bibikar et al.
(2021) hinted at poor model performance with pre-defined masks, contrasting ours where we see significantly improved model performance, implying the importance of stage 1 (as will be elaborated in section 5).
In the enabled mask learning scenario ( m freez = 0 ), model masks and weights are jointly trained during clients’ local learning, thus termed as joint mask weight sparse training (JMWST).
However, as highlighted earlier, clients’ naive sparse mask selection at the beginning of each round costs a considerable accuracy drop (section 3 Obs. 1).
JMWST allows the server to select a sparse model for the clients at round t .
For clients’ target density d , the aggregated server model ( L 17) at the end of round t − 1 , generally has density d S > d .
To enable efficient sampling of sparse model, we leverage the density re-calibration strategy (Eq. 2) by taking the ( t − 1) th round’s clients’ sensitivity in to consideration.
We then perform magnitude pruning to retain the top- d lc fraction of parameters for l th layer at the server and send the pruned model to clients at round t ( L 18).
The clients then perform local sparse learning, yielding another set of sparse models and so on.
Note, the aggregation and sampling is simpler in SPDST, as the server model always remains at density d .
In terms of yielding convergent masks, we indeed observed a lower SM for JMWST by ∼ 85% compared to that in NST, evaluated after 300 rounds on CIFAR-10.
Algorithm 1 details the FLASH training methods.
It is noteworthy that, the clients are only allowed to update mask after an interval of r int , rounds, which for JMWST is set to 1 by default, allowing the server to evaluate masks at the end of every round.
Extension to support heterogeneous parameter density.
To support different density budgets for different clients, we now present hetero-FLASH.
Let us assume a total of N support densities d set =
[ d 1 , .., d M ] , where d i <
d i +1 .
Now, for hetero-SPDST, we perform a sensitivity warm-up, to create the masks for the clients’ with the highest density d N .
For any other density d i , we sample a which Data: Training rounds T , local epochs E , client set [ C N ], clients per rounds c r , target density d , sensitivity warm-up epochs E d , density warm up client count c d , initial value of freeze masks m freez = 0 , training algorithm A and Aggregation type Agr . 1 M init ← createRandomMask ( d ) 2 Θ init ← initMaskedWeight ( M init ) 3 serverExecute: 4 # Calculate the layer-wise sensitivity in stage 1 5 Randomly sample c
d clients [ C d ] ⊂ [ C N ] 6 for each client c ∈
[ C d ] in parallel doΘ c ← clientExecute (Θ init ,E d , 0) # m _ freeze [Figure] S c ← computeSensitivity
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
sparse mask from that with density d i +1 .
Note, while creating the mask from d i +1 to d i , we follow the layer-wise density re-calibration approach as mentioned earlier.
For hetero-JMWST, at the beginning of each round, the server performs magnitude pruning to yield N sub models meeting N different density levels, contrasting to the creation of one model in JMWST.
Participating clients of different densities use the corresponding sub models to start their local sparse training.
In hetero-FLASH, server performs aggregation by following a form of weighted fed averaging (WFA).
In particular, with similar inspiration as Diao et al.
(2020), to give equal importance to each parameter update in such heterogeneous settings, WFA averages the values by their number of non-zero occurrences among the participating clients.
For brevity, we have provided the algorithm for hetero-FLASH in the Appendix.
E XPERIMENTS
Datasets and Models.
We evaluated the performance of FLASH on MNISTLeCun & Cortes (2010), Federated EMNIST (FEMNIST) Caldas et al. (2018a), and CIFAR-10 Krizhevsky et al.
(2009) datasets with the CNN models described in McMahan et al.
Caldas et al.
(2018a), and ResNet18, respectively.
Further model details are provided in the Appendix.
For data partitioning of MNIST and CIFAR-10, we use Latent Dirichlet Allocation (LDA)Reddi et al.
(2020) with three different α ( α = 1000 for IID and α = 1 and 0 .
1 for non-IID).
For FEMNIST, we employ the same setting as in Han et al.
(2020), which partitions the data based on the writer into 3400 clients, making it inherently non-IID.
Training Hyperparameters.
We use Clients’ starting learning rate ( η init ) as 0 .
1 that is exponentially decayed to 0.001 ( η end ) at the end of training.
Specifically, learning rate for participants at round t is η t = η init ( exp ( t T log ( η init η end ))) .
In all the sparse learning experiments, prune rate is set to 0.25 4 .
Summary of the rest of the training hyperparameters can be found in 1.
Furthermore, all the experiments were performed with three different seeds.
We report the final results as the averaged accuracy with corresponding std deviation in the [Table] [Table]
5. E XPERIMENTAL R ESULTS WITH FLASH
To understand the importance of stage 1 in FLASH methodology, we identify a baseline training with uniform layer sensitivity driven pre-defined sparse training (PDST) in FL.
Table 3 details the performance of FLASH at different levels of d , for various choices of sparse learning methods.
In particular, as we can see in Table 3 column 5 and 6, the performance of both NST and PDST produced models cost heavy accuracy drop at ultra low parameter density d = 0 .
For example, on CIFAR-10 ( α = 0 . 1 ), models from NST and PDST sacrifice an accuracy of 16 .
35% and 5 .
32% , respectively.
However, at comparatively higher density ( d = 0 . 1 ), both can yield models with a lower accuracy difference from the baseline by around 6 .
46% and 2 .
SPDST, on the other hand, can maintain close to the baseline accuracy at even ultra-low density for all data partitions.
Interestingly, for majority of the cases, it even outperforms JMWST yielded models.
These results clearly highlight the efficacy of both sensitivity driven sparse learning (as SPDST > PDST) and early mask convergence (as SPDST ≈ JMWST) in FL settings .
Importantly, for increased r int in JMWST, we observe a consistent improvement in accuracy .
The inferior accuracy at r int = 1 can be attributed to the mask divergence caused by frequent noisy gradient dependent update.
We thus believe efficient hyperparameter search including r int is essential for sparse FL model’s improved performance, particularly for JMWST.
Moreover, JMWST requires additional communication of non-zero weight indices, contrasting SPDST, where clients do not need to send the mask at all, allowing us to yield proportional communication saving as the model density .
Fig.
4 shows the acc.
vs. round comparison among the two proposed methods on different data distributions.
Comparison with ZeroFL.
Despite leveraging a form of sparse learning Raihan & Aamodt (2020),
ZeroFL required significantly higher up-link/down-link communication cost compared to the target density d .
This enables FLASH to gain a significant advantage in communication saving over ZeroFL, particularly for SPDST, as it only asks for the reduced size parameters to be communicated between the server and clients.
In particular, we evaluate the communication saving as the ratio of the dense model size and corresponding sparse model size with the tensors represented in compressed sparse
[Table]
row (CSR) format Tinney & Walker (1967).
As depicted in Table 4 5 , FLASH can yield an accuracy improvement of up to 10 .
1% at a reduced communication cost of up to 10 .
26 × (computed at up-link when both send sparse models).
5. E XPERIMENTAL R ESULTS WITH H ETERO -FLASH
Table 5 shows the performance of hetero-FLASH where the clients can have three possible density budgets as defined by the d set .
We assume the maximum capacity clients’ density budget of 0 . 2 .
To train on all the density values, we first create three sets, each having 40% , 30% , and 30% of total clients, and corresponds to density 0 .
2 , 0 . 15 , and 0 .
1 , respectively.
Now, during every round, we sample 10% from each set with corresponding target density.
Similar to the trend in FLASH, hetero-SPDST outperforms the JMWST counter-parts by up to 3 .
93% evaluated on the three datasets.
Also, following similar trend as with homogeneous density clients, with increased mask update interval ( r int ), the performance of hetero-JMWST gets a significant boost in accuracy of up to 3 . [Figure]
Dependence of initial sensitivity warm-up of participating clients.
To understand the importance of the clients’ participation in the warm-up, we experimented with six different scenarios.
In particular, we used two different values of participating clients ([10, 20]) each corresponding to three different local epoch choices ([10, 20, 40]).
As shown in Fig.
5(a), the yielded pruning sensitivity follows a similar trend.
Moreover, an SPDST training with mask chosen from any of these sensitivity lists finally yield FL models with similar performances (Fig. 5(b)), clearly demonstrating the robustness of our warm-up based sensitivity evaluation stage 1 .
Comparison with ERK+ initialization. We now compare our SPDST mask initialization, with that of parameter density distribution evaluated via ERK+ Huang et al. ; Evci et al. Notably, contrary to uniform density, ERK+ scheme keeps more weights for the layers having fewer parameters.
Note here, we use SPDST, ERK+, or uniform (PDST) as the initial mask for stage, and keep the mask frozen throughout the training of stage 2 . As shown in Fig. 6(a-b), the mask evaluation stage 1 to initialize mask allows SPDST to consistently provide superior results over the other two . We hypothesize this to the better layer sensitivity evaluation scheme of SPDST, particularly at the earlier layers, allowing it to retain more information at these layers.
Importance of parameters’ weighted fed averaging at the server.
Earlier literature Diao et al.
(2020) suggested a form of weighted fed averaging, for clients with different model sizes.
Inspired by that, we now investigate the necessity of WFA in FLASH.
In particular, we performed experiments on CIFAR-10 ( α = 1 . 0 ), both with and without WFA, during server aggregation.
As shown in
Fig. 6(c), WFA model performs inferior to the fed averaged model in FLASH.
On the contrary, hetero-FLASH, enjoys consistent superior performance with WFA 6(c-d).
The inferior performance of WFA in FLASH may hint at the fact that if a weight is non-zero only for fewer clients, as compared to other non-zero weights, giving it equal weight as the others nullifies its lower importance, that may be necessary to preserve for mask convergence. On the other hand, having WFA in hetero FLASH is necessary, as a weight’s less frequent non-zero occurrence can be due to fewer number of high-parameter density clients in a round.
Further investigations on the utility and use case of such weighted averaging is an interesting future research direction.
Time and communication overhead for stage 1 . Mask evaluation stage 1 uses one round with E d local epochs (for us E d = 10 ) per client. A normal FL stage in our settings trains the clients for T rounds, 1 epoch per client/round. Therefore, stage 1 increases the time by a factor of ( E d T + 1).
Usually,
E d << T , making the pre-training time overhead negligible.
The communication overhead of stage 1 is also negligible compared to that in each round for the stage 2 FL training. Each participant only needs to send L values for an L -layer model. So, c d clients will have a total communication overhead of ( L × c d × 32 ) bits, assuming 32 -bit number representation. Computation saving for FLASH.
The training FLOPs for a layer l ( F llayer ) can be partitioned into forward operation FLOPs ( F lfwd ), backward input ( F lback _ in ) and weight gradient ( F lback _ wt ) compute FLOPs. With the assumption of no-compute cost associated to the zero-valued weights via zero-gating logic Kundu et al. (2021a), the F llayer for FLASH with parameter density d ( d << 1 . 0 ) is given by [Equation]
where s a is d and 1 for SPDST and JMWST, respectively. Fu lx represents the corresponding FLOPs associated with an unpruned layer. Thus SPDST provides improved computation benefits along with the communication savings. Further details on FLOPs computation is provided in the Appendix.
Performance at limited communication budget.
Fig. 7(a) and (b) show the performance of FL models when the clients are communication limited.
In particular, we see both PDST and SPDST can significantly outperform other approaches in yielding a significantly well-trained model at low comm.
budget.
This can be attributed to their significantly smaller model sizes, helping them to run for higher number of rounds than others, on a limited bandwidth scenario.
C ONCLUSIONS
This paper presented federated lottery aware sparsity hunting methodologies to yield sparse server models with low parameter density while costing insignificant accuracy drop compared to the unpruned counterparts.
In particular, we demonstrated two efficient sparse learning solutions specifically tailored for FL, enabling better computation and communication benefits over existing sparse learning alternatives.
We experimentally demonstrated the superiority of our models in yielding up to ∼ 10 .
1% improved accuracy with ∼ 10 .
26 × fewer communication costs, compared to the existing alternatives
Qiu et al.
(2021), at similar hyperparameter settings.
A A PPENDIX A. M ODEL A RCHITECTURES
Table 6 shows the model architectures used for MNIST and FEMNIST datasets.
For CIFAR-10 we used ResNet
He et al.
with the first CONV layer kernel size as 3 × 3 instead of original 7 × 7 .
[Table]
A. H ETERO -FLASH A LGORITHM Algorithm details
the training algorithm in hetero-FLASH.
Note that and aggrParamUpdateMask and subSampleServerModel are the two functions that play key role in supoorting heterogeniety in sparsity ratios for different clients.
The details of these two functions are elaborated in Algortihm 3 and 4, respectively.
We plan to open-source our code upon acceptance of the paper.
Algorithm 2:
FLASH Training.
(Θ c
E m
) : 21 Θ c
Θ c 22 for i
E do Θ c i doSparseLearning ( c
)m
← checkUpdateMask
) 25 end 26 return c E then Θ tS ← ds tc
3 else
FLASH do
end
] 12 end Data: Current round id t , client set [ C r ] , aggregated Weight tS L layers, support density set d set =
[ d 1 , ..., d M ] where d i <
d i +1 , model layer-wise parameter count [ k ] =
[ k 1 , ..., k L ].
1 if size ( d set ) is 1 then //JMWST
subsampling in FLASH
M ← initMaskWithZeros () [ ˆ d 1 , ..., ˆ d L ] ← avgLayerWiseDensity ([ C r ]) r f ← d 1 × K (cid:80) Ll =1 ˆ d l .k l for layer l ← 1 to L do idx ← getSortedWeightIndeices (Θ tS , l ) n z ← int ( r f × ˆ d l × k l ) //number of non-zeros
M l [ idx [: n z ]] ← end 11 else //For hetero-FLASH for d i ∈
d set do
M i ← initMaskWithZeros () end
D ts ← getCurrentDensity (Θ tS )
[ ˆ d 1 , ..., ˆ d L ] ← getLayerWiseDensity (Θ tS ) for layer
l ← 1 to L do idx ← getSortedWeightIndeices (Θ tS , l ) for d i ∈
d set do r fi ← d i D ts n z ← int ( r fi × ˆ d l × k l )
M il [ idx [: n z ]] ← end end 26 end
A. M ORE Q UANTITATIVE A NALYSIS
Ablation with the mask update interval rounds ( r int ).
As mentioned in the original manuscript, in the case of JMWST, to save communication energy we often can choose not to update the mask every round.
We thus performed ablation with an increased frequency of mask update interval round r int from the default value of 1 (similar to Qiu et al.
Table 7 shows the results with different r int . In particular, as we can see in the table, less frequent update intervals does not degrade the final model performance. Moreover, a less frequent update can provide additional savings in terms of up-link cost for the models as the masks do not change every round. [Table]
Impact of Number of participating clients per round.
Fig.
8 (a), shows that JMWST and
SPDST follow the same pattern at the baseline model ( d = 1 . 0 ) with FedAvg.
In other words, similar to FedAvg, as the c r increases, the performance enhances.
Also, for a specific c r , JMWST and
SPDST perform better than PDST and NST.
Impact of Batch-Normalization layer statistics.
Fig 8 (b), shows the performance comparison between batch bormalization (BN) and static batch normalization (static BN, as suggested in Diao et al.
In particular in our setting, using BN layer statistics always outperform the static BN.
Revisiting sparse mask mismatch for FLASH.
As demonstrated in Figs 9, the sparse mask mismatch in the case of JMWST significantly reduces, helping the mask train in a convergent way, significantly faster than that in NST.
Fig.
10 shows the layer-wise SM, for centralized trained model (Fig. 10a) and FL trained model with sparsity (Fig. 10b-c).
In particular, the SM at later layer can significantly reduce in the case of JMWST as compared to NST, further demonstrating the convergence ability even at the later layers.
Convergence trend of proposed algorithms.
Fig.
11 shows the test accuracy vs FL rounds for
NST, PDST, SPDST, and JMWST algorithms on CIFAR-10 dataset with non-IID data distribution ( α = 1 ).
As shown in the plots, for d = 0 . 05 and d = 0 .
1 , NST has slower convergence with lower final accuracy.
Introducing consensus among the clients for the sparse mask accelerates the convergence and enhances the final performance.
Communication saving and FLOPs evaluations.
Employing sparse learning in FL helps participating clients reduce both the communication and compute costs (FLOPs) for training.
Without the loss of generality, we now evaluate the convolutional layer training FLOPs for FLASH, and demonstrate the relation of parameter density d with FLOPs and communication saving.
Let us assume a layer l has a weight tensor θ l ∈ R C o × C
i × h × w , where h and w are the height and width of the convolutional kernel, and C o and C i represent the number of filters and channels per filter, respectively.
It takes an input tensor I ∈ R C
i × H × W to produce an output tensor O ∈ R C o × R × S .
Let us also assume d to be the density of non-zero in the weight tensor for all the layers.
During training, FLOPs associated to each weight tensor update can be partitioned in to three component, namely, forward pass FLOPs ( F fwd ), backward input grad FLOPs ( F back _ in ), and backward weight grad FLOPs ( F back _ wt ).
The weight sparsity helps both F fwd and F back _ in to reduce proportionally as given below.
[Equation]
Finally, if the zero weights’ gradients flow is computed for the purpose of mask learning, then
F back _ wt can’t leverage the advantage of low parameter density.
Thus during mask training of JMWST, as the gradient needs to be dense, F back _ wt is same as that in dense computation.
However, for SPDST, zero weights remain as zero, allowing us to safely skip the associated gradient computation.
This essentially helps SPDST to extract benefits of sparsity during all the three stages of FLOPs computation.
Following Eq.
shows the F back _ wt in FLASH.
[Equation]
) where, s a = (cid:26) 1 , for JMWST d, for SPDST (7)
For similar density, clients in ZeroFL also enjoys similar benefits in F fwd and F back _ in .
However, F back _ wt can be reduced only via introduction of activation sparsity, a .
It is well surveyed in literature that having sparse activation density as low as parameter density may significantly impact model performance.
Thus generally, a > d that allows SPDST to enjoy a FLOPs benefit of ad for F back _ wt .
These computations can be easily extended to linear layers.
Also, we can safely ignore the FLOPs associated to the BN layers due to their negligible contribution to the total FLOPs.
Discussion on compute benefits at the edge.
To extract FLOPs benefits for irregular pruning in FLASH, we assume that the compute energy for the sparse network can be avoided via the means of clock-gating Yang & Kim (2018) of the zerovalued weights.
Moreover, there has been recent development of sparsity-friendly DNN accelerators Qin et al.
(2020) that can efficiently reduce the compute cost by a significant margin.
Such accelerators can leverage the yielded sparse FL models to deploy at compute constrained edges.
FLOPs vs. communication cost for different density budget.
To reach a target accuracy value, we now plot the FLOPs to uplink communication cost for different density budget in Fig. Impact of r int . As depicted in Fig.
13, the test accuracy improves with the increase in JMWST mask update interval r int .
In particular, for both α = 0 . 1 and 1 . 0 , the accuracy with increased r int can be up to ∼ 1 .
6 % and ∼ 0 .
98 %, respectively.
Interestingly, the improvement tend to saturate after certain r int .
Thus, we consider important sparse learning hyperparameter search as an interesting future research direction.