FSL: Federated Supermask Learning

Federated learning ( FL ) allows multiple clients with ( private ) data to collaboratively train a common machine learning model without sharing their private training data . In-the-wild deployment of FL faces two major hurdles : robustness to poisoning attacks and communication efficiency . To address these concurrently , we propose Federated Supermask Learning ( FSL ) . FSL server trains a global subnetwork within a randomly initialized neural network by aggregating local subnetworks of all collaborating clients . FSL clients share local subnetworks in the form of rankings of network edges ; more useful edges have higher ranks . By sharing integer rankings , instead of float weights , FSL restricts the space available to craft effective poisoning updates , and by sharing subnetworks , FSL reduces the communication cost of training . We show theoretically and empirically that FSL is robust by design and also significantly communication efficient ; all this without compromising clients ’ privacy . Our experiments demonstrate the superiority of FSL in real-world FL settings ; in particular , ( 1 ) FSL achieves similar performances as state-of-the-art FedAvg with significantly lower communication costs : for CIFAR10 , FSL achieves same performance as Federated Averaging while reducing communication cost by ∼ 35 % . ( 2 ) FSL is substantially more robust to poisoning attacks than state-of-the-art robust aggregation algorithms . 1 INTRODUCTION . Federated Learning ( FL ) is an emerging AI technology , where mutually untrusted clients ( e.g. , Android devices ) collaborate to train a shared model , called the global model , without explicitly sharing their local training data . FL training involves a server ( e.g. , Google server ) which collects model updates from selected FL clients in each round of training , and uses them to update the global model . FL , although highly promising , faces multiple challenges ( Kairouz et al. , 2019 ; Li et al. , 2020b ) to its practical deployment , in particular , communication efficiency and robustness , which are the focus of our work . Privacy preservation is another major challenge to FL , but is orthogonal to our work . We present Federated Supermask Learning ( FSL ) , a novel approach to perform FL while concurrently achieving the two goals of robustness and communication efficiency . FSL is built on a novel learning paradigm called supermasks ( Zhou et al. , 2019 ; Ramanujan et al. , 2020 ) , which allows it to reduce communication costs while achieving significantly higher robustness . Specifically , in FSL , clients collaborate to find a subnetwork within a randomly initialized neural network which we call the supernetwork ( this is in contrast to conventional FL where clients collaborate to train a neural network ) . The goal of training in FSL is to collaboratively identify a supermask , which is a binary mask of 1 ’ s and 0 ’ s , that is superimposed on the random neural network ( the supernetwork ) to obtain the final subnetwork . The subnetwork is then used for downstream tasks , e.g. , image classification , hence it is equivalent to the global model in conventional FL . Note that in entire FSL training , weights of the supernetwork do not change . More specifically , each FSL client computes the importance of the edges of the supernetwork based on their local data . The importance of the edges is represented as a ranking vector . Each FSL client will use the edge popup algorithm ( Ramanujan et al. , 2020 ) and their data to compute their local rankings ( the edge popup algorithm aims at learning which edges in a supernetwork are more important over the other edges by minimizing the loss of the subnetwork on their local data ) . Each client then will send their local edge ranking to the server . Finally , the FSL server uses a novel voting mechanism to aggregate client rankings into a supermask , which represents which edges of the random neural network ( the supernetwork ) will form the global subnetwork . Intuitions on FSL ’ s robustness : In traditional FL algorithms , clients send large-dimension model updates ∈ Rd ( real numbers ) to the server , providing malicious clients a significant flexibility in fabricating malicious updates . By contrast , FSL clients merely share the rankings of the edges of the supernetwork , i.e. , integers ∈ [ 1 , d ] , where d is the size of the supernetwork . Therefore , as we will show both theoretically and empirically , FSL provides robustness by design and reduces the impact of untargeted poisoning attacks . Furthermore , unlike most existing robust FL frameworks , FSL does not require any knowledge about the percentages of malicious clients . Intuitions on FSL ’ s communication efficiency : In FSL , the clients and the server communicate just the rankings of the edges in the supernetwork , i.e. , a permutation of indices in [ 1 , d ] . Ranking vectors are generally significantly smaller than the global model . This , as we will show , significantly reduces the upload and download communication in FSL compared to Federated Averaging ( FedAvg ) ( McMahan et al. , 2017 ) , where clients communicate model parameters , each of 32/64 bits . Empirical results : We experiment with three datasets in real-world heterogeneous FL settings and show that : ( 1 ) FSL achieves similar performance ( e.g. , model accuracy ) as state-of-the-art FedAvg but with significantly reduced communication costs : for CIFAR10 , the accuracy and communication cost per client are 85.4 % and 40.2MB for FedAvg , while 85.3 % and 26.2MB for FSL . ( 2 ) FSL is highly robust to poisoning attacks as compared to state-of-the-art robust aggregation algorithms : from 85.4 % in the benign setting , 10 % malicious clients reduce the accuracy of FL to 56.3 % and 58.8 % with Trimmed-Mean ( Xie et al. , 2018 ; Yin et al. , 2018 ) and Multi-Krum ( Blanchard et al. , 2017 ) , respectively , while FSL ’ s performance only decreases to 79.0 % . We also compare FSL with two communication reduction methods , SignSGD ( Bernstein et al. , 2019 ) and TopK ( Alistarh et al. , 2018a ) and show that FSL produces comparable communication costs and model accuracies . For instance , on CIFAR10 , FSL , SignSGD , and TopK achieve 85.3 % , 79.1 % , and 82.1 % test accuracy , respectively , when the corresponding communication costs ( download and upload ) are 26.2MB , 20.73MB , and 30.79MB . On the other hand , FSL offers a significantly superior robustness . For instance , on CIFAR10 , 10 % ( 20 % ) malicious clients reduce the accuracy of SignSGD to 39.7 % ( 10.0 % ) , but FSL ’ s accuracy decreases to only 79.0 % ( 69.5 % ) . TopK is incompatible with existing robust aggregation algorithms , hence uses Average aggregation and is as vulnerable as FedAvg , especially in the real-world heterogeneous settings . 2 RELATED WORKS . Supermask Learning : Modern neural networks have a very large number of parameters . These networks are generally overparameterized ( Dauphin & Bengio , 2013 ; Denil et al. , 2013 ) , i.e. , they have more parameters than they need to perform a particular task , e.g. , classification . The lottery ticket hypothesis ( Frankle & Carbin , 2019 ) states that a fully-trained neural network , i.e. , supernetwork , contains sparse subnetworks , i.e. , subsets of all neurons in supernetwork , which can be trained from scratch ( i.e. , by training same initialized weights of the subnetwork ) and achieve performances close to the fully trained supernetwork . The lottery ticket hypothesis allows for massive reductions in the sizes of neural networks . ( Ramanujan et al. , 2020 ) offer a complementary conjecture that an overparameterized neural network with randomly initialized weights contains subnetworks which perform as good as the fully trained network . Poisoning Attacks and Defenses for Federated Learning ( FL ) : FL involves mutually untrusting clients . Hence , a poisoning adversary may own or compromise some of the FL clients , called malicious clients , with the goal of mounting a targeted or untargeted poisoning attack . In a targeted attack , the goal is to reduce the utility of the model on specific test inputs , while in the untargeted attack , the goal is to reduce the utility for all ( or most ) test inputs . It is shown ( Blanchard et al. , 2017 ) that even a single malicious client can mount an effective untargeted attack on FedAvg . In order to make FL robust to the presence of such malicious clients , the literature has designed various robust aggregation rules ( AGR ) ( Blanchard et al. , 2017 ; Mhamdi et al. , 2018 ; Yin et al. , 2018 ; Chang et al. , 2019 ) , which aim to remove or attenuate the updates that are more likely to be malicious according to some criterion . For instance , Multi-krum ( Blanchard et al. , 2017 ) repeatedly removes updates that are far from the geometric median of all the updates , and Trimmed-mean ( Xie et al. , 2018 ; Yin et al. , 2018 ) removes the largest and smallest values of each update dimension and calculates the mean of the remaining values . Unfortunately , these robust AGRs are not very effective in non-convex FL settings and multiple works have demonstrated strong targeted ( Wang et al. , 2020 ; Bhagoji et al. , 2019 ) and untargeted attacks ( Shejwalkar & Houmansadr , 2021 ; Fang et al. , 2020 ) on them . Communication Cost of FL : In many real-world applications of FL , it is essential to minimize the communication between FL server and clients . Especially in cross-device FL , the clients ( e.g. , mobile phones and wearable devices ) have limited resources and communication can be a major bottleneck . There are two major types of communication reduction methods : ( 1 ) Qunatization methods reduce the resolution of ( i.e. , number of bits used to represent ) each dimension of a client update . For instance , SignSGD ( Bernstein et al. , 2019 ) uses the sign ( 1 bit ) of each dimension of model updates . ( 2 ) Sparsification methods propose to use only a subset of all the update dimensions . For instance , in TopK ( Aji & Heafield , 2017 ; Alistarh et al. , 2018a ) , only the largest K % update dimensions are sent to the server in each FL round . We note that , communication reduction methods primarily focus on and succeed at reducing upload communication ( client → server ) , but they use the entire model in download communication ( server→ client ) . 3 PRELIMINARIES . 3.1 FEDERATED LEARNING . In FL ( McMahan et al. , 2017 ; Kairouz et al. , 2019 ; Konečnỳ et al. , 2016 ) , N clients collaborate to train a global model without directly sharing their data . In round t , the service provider ( server ) selects n out of N total clients and sends them the most recent global model θt . Each client trains a local model for E local epochs on their data starting from the θt using stochastic gradient descent ( SGD ) . Then the client send back the calculated gradients ( Ok for kth client ) to the server . The server then aggregates the collected gradients and updates the global model for the next round . FL can be either cross-device or cross-silo ( Kairouz et al. , 2019 ) . In cross-device FL , N is large ( from few thousands to billions ) and only a small fraction of clients is chosen in each FL training round , i.e. , n N . By contrast , in cross-silo FL , N is moderate ( up to 100 ) and all clients are chosen in each round , i.e. , n = N . In this work , we evaluate the performance of FSL and other FL baselines for cross-device FL under realistic production FL settings . 3.2 EDGE-POPUP ALGORITHM Algorithm 1 Edge-popup ( EP ) algorithm : it finds a subnetwork of size k % of the entire network θ 1 : Input : number of local epochs E , training data D , initial weights θw and scores θs , subnetwork size k % , learning rate η 2 : for e ∈ [ E ] do 3 : B ← Split D in B batches 4 : for batch b ∈ [ B ] do 5 : EP FORWARD ( θw , θs , k , b ) 6 : θs = θs − η∇ ` ( θs ; b ) 7 : end for 8 : end for 9 : return θs 10 : function EP FORWARD ( θw , θs , k , b ) 11 : m← sort ( θs ) 12 : t← int ( ( 1− k ) ∗ len ( m ) ) 13 : m [ : t ] = 0 14 : m [ t : ] = 1 15 : θp = θw m 16 : return θp ( b ) 17 : end function The edge-popup ( EP ) algorithm ( Ramanujan et al. , 2020 ) is a novel optimization method to find supermasks within a large , randomly initialized neural network , i.e. , a supernetwork , with performances close to the fully trained supernetwork . EP algorithm does not train the weights of the network , instead only decides the set of edges to keep and removes the rest of the edges ( i.e. , pop ) . Specifically , EP algorithm assigns a positive score to each of the edges in the supernetwork . On forward pass , it selects top k % edges with highest scores , where k is the percentage of the total number of edges in the supernetwork that will remain in the final subnetwork . On the backward pass , it updates the scores with the straight-through gradient estimator ( Bengio et al. , 2013 ) . Algorithm 1 presents EP algorithm ; we defer further details to Appendix D .

Summary: This paper presents a pruning-before-training algorithm for improving the communication efficiency of federated learning. In each iteration, all local devices with the same initial weights perform the pruning on their own local datasets. The edge scores produced by local devices are aggregated by the server. The order of the edge scores will be used in the next iteration of score initialization process. After the pruning convergence, the edges are pruned correspondingly and the federated learning model is trained on the pruned network.