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: The paper investigates sparse training to accelerate on-device federated learning. It begins by taking an off-the-shelf sparse training method (SWAT) and analyzes its limitations. Then, the method proposes three strategies to improve upon SWAT. Experiments are conducted on CIFAR-10 dataset and show some improvements compared to SWAT.