ZeroFL: Efficient On-Device Training for  Federated Learning with Local Sparsity

1 INTRODUCTION . Despite it being a relatively new subfield of machine learning ( ML ) , Federated Learning ( FL ) ( McMahan et al. , 2017 ; Reddi et al. , 2021 ; Horvath et al. , 2021 ) has become an indispensable tool to enable privacy-preserving collaboratively learning , as well as to deliver personalised models tailored to the end-user ’ s local data and context ( Arivazhagan et al. , 2019 ; Hilmkil et al. , 2021 ; Cheng et al. , 2021 ) . For example : next-word prediction ( Hard et al. , 2018 ) , physical activity detection ( Doherty et al. , 2017 ) , keyword spotting ( Hard et al. , 2020 ) , among others . Unlike standard centralised training , which normally takes place on the Cloud and makes use of powerful hardware ( Hazelwood et al. , 2018 ) , FL is envisioned to run on commodity devices such as smartphones or IoT devices often running of batteries , which are orders of magnitude more restricted in terms of compute , memory and power consumption . This triplet of factors drastically limits the complexity of the ML models that can be trained on-device in a federated manner , ceiling their usefulness for the aforementioned applications as a result . In order to adjust the memory and compute footprints of complex ML model to the FL setting , the research community has presented a number of approaches including : the use of distillation ( Hinton et al. , 2015 ) to enable the aggregation on the server side of heterogeneous model architectures ( e.g . based on the compute capabilities of each device ) that collaboratively train a single global model ( Lin et al. , 2020 ; Zhu et al. , 2021 ) ; group knowledge transfer algorithm ( He et al. , 2020 ) ; federated dropout , by which clients perform local training on a sub-model of the global model ( Caldas et al. , 2019 ) , translates into lower overall communication costs and , enables better support for heterogeneous pools of clients regardless of their compute capabilities ( Horvath et al. , 2021 ) ; and , more generally , better aggregation strategies that enable faster convergence ( Li et al. , 2018 ; Wang et al. , 2020 ; Reddi et al. , 2021 ) , reducing in this way overall device utilization ( e.g . fewer local epochs ) and number of communication rounds . Other optimization techniques such as quantization and sparsity have been used in the context of FL but mostly as a way to reduce communication costs ( Liu et al. , 2021 ; Amiri et al. , 2020 ; Shahid et al. , 2021 ) but not to accelerate on-device training . The use of sparse operations ( e.g . convolutions ) at training time has recently been shown to be an effective technique to accelerate training in centralised settings ( Sun et al. , 2017 ; Goli & Aamodt , 2020 ; Raihan & Aamodt , 2020 ) . The resulting models are as good or close to their densely-trained counterparts despite reducing by up to 90 % their FLOPs budget and , resulting in an overall up to 3.3× training speedup . Acceleration is achieved by performing sparse convolutions during the forward and/or backward pass , which requires at least one of the operands ( i.e . inputs , weights , gradients ) to be sufficiently sparse and , software and hardware support for such operations . However , it is unclear how the different FL-specific challenges ( i.e . data imbalance , stateless clients , periodic aggregation ) will restrict the quality of the global model . This work considers the challenges and opportunities of inducing high levels of sparsity to accelerate training on-device for FL workloads , and provides the following contributions : • The first framework for Federated Learning that leverages sparsity as a mechanism to accelerate on-device training by inducing up to 95 % sparse weights and activations.This work considers three popular datasets : CIFAR-10 and FEMNIST for image classification and , SpeechCommands for audio classification . • A study on the unique aspects that arise when introducing sparsity at training time in FL : the degree of overlap between non-zero values decreases with layer-depth index and , the locations of zero-valued weights in the global model remain constant throughout most of the training rounds . Our discussion sets the foundations for future research in this area . • A technique that alleviates the accuracy degradation when applying a state-of-the-art offthe-shelf sparsification method to the FL domain . ZeroFL achieves +2.1 % and +3.7 % higher accuracy than baselines when inducing 90 % and 95 % sparsity respectively . In addition , ZeroFL also leverages sparsity when transferring the local models to the central server reducing communication costs by 3.0× while still outperforming competitive baselines . 2 RELATED WORK . Pruning neural networks involves discarding parts of the model ( e.g . individual weights or entire channels ) that are irrelevant for solving the task at hand . This procedure generally produces a lightweight model representation more suitable for deployment on constrained devices with limited memory and compute budgets . In this section we detail how different forms of pruning or sparsification have been used to accelerate inference and , to a lesser extent , training . We also discuss how these have been introduced to reduce communication costs in distributed and federated learning . Unstructured pruning . Frameworks relying on unstructured pruning ( Han et al. , 2015a ; b ; Guo et al. , 2016 ; Molchanov et al. , 2017 ) often achieve higher compression ratios at the expense of inference stages being as compute intensive in practice as those of the original model . This is because , assuming pruning has been homogeneously applied on the model , sparse operations can only be efficiently accelerated on supported hardware , such as modern GPUs ( Wang , 2020 ; Zachariadis et al. , 2020 ; Hong et al. , 2018 ) or custom accelerators ( Zhang et al. , 2016 ; Lu et al. , 2019 ; Srivastava et al. , 2020 ) , for a sufficiently high sparsity ratio . The lower the ratio , the less likely sparse operations would translate into measurable speedups . In the case of CPUs , speedups due to sparse operations where one operand is unstructurally sparse are often only feasible at 90 % sparsity ratios or higher ( Hong et al. , 2019 ; Wang , 2021 ) . Structured pruning . Methods that apply structured pruning ( He et al. , 2018 ; 2017 ; Jian-Hao Luo & Lin , 2017 ; Yu et al. , 2018 ; Molchanov et al. , 2019 ; Wang et al. , 2017 ) , on the other hand , trade compression for acceleration potential . These approaches modify the underlying computational graph by discarding entire channels , resulting in smaller but still dense convolution operations , or by removing the nodes all together if an entire layer is set to be removed by the chosen pruning strategy . As a result , structured pruning frameworks are the preferred option when aiming to accelerate inference on general purpose hardware . A body of work across structured and unstructured pruning methods , attempts to induce structure in otherwise randomly sparse networks S. Gray & Kingma ( 2017 ) ; Ren et al . ( 2018 ) ; Wen et al . ( 2020 ) ; Verelst & Tuytelaars ( 2020 ) . This is often referred to as block sparsity and consists in subdividing the matrix representations of inputs or weights into tiles ( e.g . 16×16 tiles ) , and restrict the training in such a way that some tiles contain only zeros while the rest remain dense and real-valued . Matrix-matrix multiplications following such a pattern can be accelerated at lower global sparsity ratios compared to those following unstructured sparsity Hoefler et al . ( 2021 ) . Other forms of constraining how sparsity occurs have been proposed , for example , a cache-aware reordering on the sparsity pattern of the weights Elsen et al . ( 2020 ) . This can be used to ensure high cache reuse on Cortex-A mobile CPUs , resulting in 2.4× acceleration of MobileNets . Sparse training . The majority of works making use of sparsity are envisioned for either model compression or to accelerate inference . Only recently , sparse operations have been considered to accelerate training . The work of Sun et al . ( 2017 ) presented a mechanism to induce high levels of sparsity in the gradients during backpropagation and , demonstrated large speedups when training MLP-only models . More recently , Goli & Aamodt ( 2020 ) build upon the observation that gradients from consecutive batches are near identical . They present a framework to reuse a random sample of previously computed gradients and their thresholded difference w.r.t gradients from the current batch , resulting in a sparse tensor . Their framework accelerates training of CNNs by performing sparse convolutions during the backward pass at the cost of pre-computing partial gradients during forward pass . Closer to our work is SWAT ( Raihan & Aamodt , 2020 ) , a framework that relies on sparsified weights during inference and sparsified weights and activations for backward propagation . Compared to the previous two frameworks , SWAT achieves superior performance in terms of training acceleration and accuracy retention at high sparsity ratios . Compression on communication . Konečnỳ et al . ( 2016 ) proposed to restricts the updates of weight matrices to have a pre-specified structure in order to reduce the total communication cost . The structure can either be random or low-rank structure . ATOMO ( Wang et al. , 2018 ) introduced a generalised gradient decomposition and sparsification technique , aiming to reduce the gradient sizes communicated upstream . Han et al . ( 2020 ) proposed a different way of aggregation in the server , which instead of aggregating model weights , it aggregates the sparsified gradients after every local update step . However , since the method requires to aggregate sparsified gradient after every step , it can not benefit from multiple local updates . Hence it might require extra communication rounds to reach the target performance . PruneFL Jiang et al . ( 2019 ) reduced both computation and communication overhead to minimize the overall training time by including an initial pruning at one selected client and further pruning as a part of FL process . Nevertheless , none of the aforementioned works explored the challenges of extending state-of-theart sparsification methods to federated learning as a way to accelerate on-device training . With ZeroFL , a framework specifically tailored to the FL setting , achieves better accuracy retention than with existing methods that remain exclusive to the centralised training paradigm . 3 BACKGROUND . This section describes the state-of-the-art sparse training method SWAT ( Raihan & Aamodt , 2020 ) ; the way we adapt it to the FL contexts ; and the related challenges that would need be addressed . 3.1 SPARSE WEIGHTS AND ACTIVATIONS TRAINING . The SWAT framework embodies two strategies in the training process . During each forward pass , the weights are partitioned into active weights and non-active weights by a top-K ( in magnitude ) operator and only the active weights are used . For the lth layer in the model , the layer maps the input activations al−1 onto feature maps ol using function fl : ol = fl ( al−1 , wl ) . In this work we consider fl being the 3× 3 convolution in the l-th layer . In the backward pass , the gradient of input activations ( 5al−1 ) and the gradient of weights ( 5wl ) are calculated represented by functions Gl and Hl , as shown below : 5al−1 = Gl ( 5al , wl ) ( 1 ) 5wl = Hl ( 5al , al−1 ) ( 2 ) Then in the backward pass , the retained layer inputs al−1 are also partitioned into active and nonactive by using the same top-K procedure . This results in full gradients and active weights being used in Eq . 1 , while full gradients and active activations are used in Eq . 2 . It is worth noticing that even weights and activations and sparsified in the forward and backward pass , the gradients generated through the training process are dense . Therefore , the resulting model is a dense . The compute cost of updating weights wl given a dense5wl tensor is negligible compare to the savings due to performing the underlying convolutions in Eq.1 & 2 , as this is essentially a weighted sum .

Summary: This work proposes ZeroFL; a method that allows for sparse neural network training at the edge along with reduced upload communication, both of which being important aspects in federated learning.  ZeroFL essentially follows a prior work on sparse neural network training, namely SWAT, and adapts it appropriately for the federated setting. The SWAT method works by setting a target weight sparsity percentage (denoted by $sp$) which is enforced in both the forward and backward pass by a top-k operation on the weights. This allows for using sparse convolutions and thus offering training time speed-ups. The authors note that it is important to use the same set of sparse weights for inference as well, in order to not affect performance.   ZeroFL then applies this idea in the federated setting by letting each client perform SWAT type of training with a target sparsity of $sp$. Instead of the clients then communicating to the server the exact set of sparse weights obtained at the end of their local training procedure, ZeroFL proposes to add an extra “slack” $r_{mask}$ and communicate $1 - sp + r_{mask}$ percentage of weights (instead of $1 - sp$). The authors motivate this change via the results of an ablation study which showed significantly lower performance when directly applying SWAT to the federated setting. They argue that with the extra $r_{mask}$ weights, a “cleaner” training signal from the clients can be given as these weights are not “corrupted” by the weights of the other clients when averaging at the server. The authors further propose three separate weights of updating the server model, i.e., 1) using the client top-k weights, 2) using the difference between the server and client top-k weights and 3) using the top-k differences between the server and client weights.    The authors experimentally validate the ZeroFL method on cifar 10 using 100 clients and both iid and non-iid splits.