On Heterogeneously Distributed Data, Sparsity Matters

1 INTRODUCTION Data privacy raises increasingly intensive concerns , and governments have enacted legislation to regulate the privacy intrusion behavior of mobile users , e.g. , the General Data Protection Regulation ( Voigt & Von dem Bussche , 2017 ) . Traditional distributed learning approaches , requiring massive users ’ data to be collected and transmitted to a central server for training model , soon may no longer be realistic under the increasingly stringent regulations on users ’ private data . On this ground , federated learning ( FL ) , a distributed training paradigm emerges as a successful solution to cope with privacy concerns , which allows multiple clients to perform model training within the local device without the necessity to exchange the data to other entities . In this way , the data privacy leakage problem could be potentially relieved . Despite the promising prospect , several notorious issues are afflicting practical performance of FL : • The global model produced by weight average ( or FedAvg and its non-personalized variants ) exhibits unsatisfactory performance in a Non-IID data distribution setting . To alleviate this problem , the most popular idea is to integrate personalized features into the global model , and produce dedicated model for each local distribution . However , how to make this integration is an open problem that remains unresolved . Prior works on personalized FL ( PFL ) zero in this issue , but the existing methods either demonstrate weak generalization towards different model architectures ( Arivazhagan et al. , 2019 ) , or require extra computation and storage ( Li et al. , 2021 ) . • The communication and training overhead is prohibitively high for both the FL and PFL . Clients in FL/PFL responsible for model training are mostly edge-devices with limited computation capacity and low bandwidth , and may not be powerful enough to fulfill a modern machine learning task with large deep neural networks . Existing studies ( Li et al. , 2020 ; Vahidian et al. , 2021 ) integrate model compression into FL/PFL to save communication and computation overhead . However , both methods embrace the technique of dense-to-sparse training , which still requires a large amount of communication at the beginning of training . In addition , how to effectively aggregate the dynamic sparse models is another challenging problem that remains unresolved . In this work , we propose FedSpa ( see Figure 2 ) , which has two key features to counter the above two challenges : ( i ) FedSpa does not deploy a single global model , but allows each client to own its unique sparse model masked by a personalized mask , which successfully alleviates the Non-IID challenge . ( ii ) FedSpa allows each client to train over an evolutionary sparse model with constant sparsity1 throughout the whole federated training process , which consistently alleviates the computation overhead of clients . Besides , all the local models in FedSpa are sparse models , which requires a smaller amount of communication cost in each communication round . Theoretically , we conclude that with the rise of Non-IID extent , setting a higher sparsity may result in a better convergence on the personalized models of FedSpa . Empirically , in the Non-IID setting , we demonstrate that FedSpa accelerates the convergence ( respectively 76.2 % and 38.1 % less communication rounds to reach the best accuracy of FedAvg ( McMahan et al. , 2016 ) and Ditto ( Li et al. , 2021 ) ) , in- creases the final accuracy ( up to 21.9 % and 4.4 % higher accuracy than FedAvg and Ditto , respectively ) , reduces the communication overhead ( 50 % less parameters communicated than the dense solutions ) , and lowers the computation ( 15.3 % lower floating-point operations ( FLOPs ) than algorithms trained with fully dense model ) . To the end , we summarize our contribution as : • We present a novel formulation of the sparse personalized FL ( SPFL ) problem , which can be applied to various network architectures by enforcing personalized sparse masks to a global model . • We propose a solution dubbed as FedSpa to solve the SPFL problem . By our novel design , FedSpa reduces the communication and computation overhead of the general FL solution . • Two sparse-to-sparse mask searching techniques are integrated as plugins of our solution . To adapt our PFL training context , we modify the DST-based mask searching technique to enable a warm-start of the searching process , which achieves superior performance . • We theoretically present the convergence analysis of the personalized models . Experimental results demonstrate the superiority of FedSpa and also coincides with the theoretical conclusion – with the rise of data heterogeneity , setting a higher sparsity of FedSpa may potentially result in a better convergence on its personalized models . 2 RELATED WORKS . Federated learning ( FL ) ( McMahan et al. , 2016 ) is seriously afflicted by the issue of heterogeneously distributed ( or Non-IID ) data . Personalized FL ( PFL ) , initiated by recent literature ( Li et al. , 2021 ; Arivazhagan et al. , 2019 ) , is shown to be effective to counter this issue of FL . In this work , we propose an alternative yet effective way to enhance PFL with personalized sparse models . 2.1 PERSONALIZED FEDERATED LEARNING . We categorize PFL into five genres . Firstly , PFL via layer partition , e.g. , FedPer ( Arivazhagan et al. , 2019 ) , LG-FedAvg ( Liang et al. , 2020 ) , FedRep ( Collins et al. , 2021 ) , is to divide the global model layers into shared layers and personalized layers . For the shared layers , weights average as in FedAvg is adopted , while for personalized layers , models are trained only locally and will not 1Sparsity specifies the ratio of parameters that are set to 0 ( or inactive ) in a model . be exchanged with others . Secondly , PFL via regularization , e.g. , Ditto ( Li et al. , 2021 ) , L2GD ( Hanzely & Richtárik , 2020 ) is to add a proximal term on the local model to force the local model and global model closely in the local model fine-tuning stage . Similarly , Sarcheshmehpour et al . ( 2021 ) ; SarcheshmehPour et al . ( 2021 ) propose a total variation ( TV ) regularization to form the network lasso ( nLasso ) problem , and primal-dual methods adapted from ( Jung , 2020 ) are proposed to solve the nLasso problems . Thirdly , PFL via model interpolation , e.g. , MAPPER ( Mansour et al. , 2020 ) , APFL ( Deng et al. , 2020 ) achieves personalization by linearly interpolating the weights of the cluster ( global ) model and local model as the personalized model . Fourthly , PFL via transfer learning , e.g. , FedMD ( Li & Wang , 2019 ) , FedSteg ( Yang et al. , 2020 ) , and Fedhealth ( Chen et al. , 2020 ) , is to either use model and domain-specific local fine-tuning or knowledge distillation to adapt the global model into the personalized model . Finally , PFL via model compression , e.g. , LotteryFL ( Li et al. , 2020 ) and Sub-FedAvg ( Vahidian et al. , 2021 ) , achieves personalization via employing principle model compression techniques , such as weight pruning and channel pruning , over the shared global model . 2.2 SPARSE DEEP NEURAL NETWORKS . Methods to Sparsify neural networks can be classified into two genres : dense-to-sparse methods and sparse-to-sparse methods . Dense-to-sparse methods train from a dense model , and compress the model along the training process . Iterative pruning , first proposed by Frankle & Carbin ( 2018 ) , shows promising performance in dynamically searching for a sparse yet accurate network . Recently , sparse-to-sparse methods have been proposed to pursue training efficiency . Among them , dynamic sparse training ( DST ) ( Bellec et al. , 2018 ; Evci et al. , 2020 ; Liu et al. , 2021 ) is the most successful technique that allows sparse networks , trained from scratch , to match the performance of their dense equivalents . Stemming from the first work – sparse evolutionary training ( Mocanu et al. , 2018 ; Liu et al. , 2020 ) , DST has evolved as a class of sparse training methods absorbing many advanced techniques , e.g. , weight redistribution ( Mostafa & Wang , 2019 ; Dettmers & Zettlemoyer , 2019 ) , gradient-based regrowth ( Dettmers & Zettlemoyer , 2019 ; Evci et al. , 2020 ) , and extra weight exploration ( Jayakumar et al. , 2020 ; Liu et al. , 2021 ) . Our work also achieves personalization via model compression . We emphasize that three main progresses are made towards SOTA compression-based PFL : ( i ) We rigorously formulate the sparse personalized FL problem , filling the gap left by the prior works . ( ii ) While prior works either vaguely describe their model aggregation as ” aggregating the Lottery Ticket Network via FedAvg ” ( Li et al. , 2020 ) , or ” taking the average on the intersection of unpruned parameters in the network ” ( Vahidian et al. , 2021 ) , we explicitly formulate the aggregation as averaging the sparse update from clients . ( iii ) Both the two prominent prior works use the idea of iterative pruning to prune the network from dense to sparse . We instead provide two sparse-to-sparse training alternatives to plug in our solution , which largely reduces the costs of communication at the beginning of the training process , and exhibits remarkable performance . 3 PROBLEM FORMULATION . We assume a total number of K clients within our FL system , and we consistently use k to index a specific client . First , we give a preliminary introduction on the general FL and PFL problem . General FL problem . Letw ∈ Rd be the global weight . General FL takes the formulation as below ( P1 ) min w f̃ ( w ) = 1 K K∑ k=1 { F̃k ( w ) : = E [ L ( x , y ) ∼Dk ( w ; ( x , y ) ) ] } , whereD = D1∪· · ·∪DK is the joint distribution of k local heterogeneous distributions , data ( x , y ) is uniformly sampled according to distribution Dk wrt . client k , and L ( · ; · ) is the loss function . Ultimate PFL Problem . Let { wk } be the personalized models . The ultimate PFL is defined as : ( P2 ) min { w1 , ··· , wK } f̂ ( w1 , · · · , wK ) = 1 K K∑ k=1 { F̂k ( wk ) : = E [ L ( x , y ) ∼Dk ( wk ; ( x , y ) ) ] } , according to ( Zhang et al. , 2020 ; Hanzely et al. , 2021 ) . The problem could be separately solved by individual client with no communication . However , if the local data is insufficient , poor performance could be observed by this direct solution , since the local models can not be boosted by other clients . Regularized PFL problem . Regularized PFL can ease the heterogeneous challenges encountered by general FL , while escaping the curse of insufficient samples encountered by the ultimate PFL problem . Inspired by ( Chen & Chao , 2021 ) , we formulate the Regularized PFL problem as follows : ( P3 ) min { w1 , ··· , wK } f̄ ( w1 , · · · , wK ) = 1 K K∑ k=1 { F̄k ( wk ) : = E [ L ( x , y ) ∼Dk ( wk ; ( x , y ) ) ] } +R ( · ) , where wk denote the personalized models and R ( · ) is the regularizer that enables information exchange between clients , making the problem tractable even the local samples are insufficient . However , it remains controversial about how to define the regularizer . Also , the gap between regularized PFL problem ( P3 ) and the ultimate PFL problem ( P2 ) still remains unspecified in existing PFL study . In this work , our ultimate goal is to solve problem ( P2 ) , which requires information exchange between clients to ensure effective solution . Below , instead of utilizing regularizer as in ( P3 ) , we alternatively propose a novel Sparse PFL ( SPFL ) problem to reach the same goal . Sparse PFL problem . By introducing personalized masks into FL , we derive the SPFL problem as ( P4 ) min w f ( w ) = 1 K K∑ k=1 { Fk ( m ∗ k w ) : = E [ L ( x , y ) ∼Dk ( m ∗ k w ; ( x , y ) ) ] } , where m∗k ∈ { 0 , 1 } d is a personalized sparse binary mask for k-th client . denotes the Hadamard product for two given vectors . Our ultimate goal is to find a global model w , such that the personalized model for k-th client can be extracted from the global model by personalized mask m∗k , i.e. , m∗k w. The element of m∗k being 1 means that the weight in the global model is active for k-th personalized model , otherwise , remains dormant . Thus , the information exchange between all personalized models is enforced by a shared global model w. Compared with existing PFL algorithms , solving our SPFL problem ( P4 ) does not sacrifice additional computation and storage overhead of clients , since we do not maintain both personalized local models and global model in clients as in ( Li et al. , 2021 ; Mansour et al. , 2020 ) . On contrary , the solution to our problem could potentially lower the communication and computation overhead . Moreover , SPFL problem ( P4 ) can be applied to most of the model architectures without modelspecific hyper-parameter tuning , since we do not make model-specific separation of the public and personalized layer in ( Arivazhagan et al. , 2019 ; Liang et al. , 2020 ; Collins et al. , 2021 ) , or domainspecific fine-tuning in ( Chen et al. , 2020 ; Yang et al. , 2020 ) .

Summary: This work studies personalized federated learning through sparse local masks. For each client, the local model is subset of the global model. The paper proposes and analysis a new method, named FedSpa, which trains sparse local models for all clients reducing communication size, the amount of computation and memory cost.