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 paper proposes an interesting approach to FL with heterogeneous client data. The main idea is  to learn a common weight vector for all clients but allow them to individually mask this global weight vector.   **** I appreciate the effort that authors put into adressing my concerns which have been partially covered. Thus, i have slighltly raised the grade of my recommendation. However, i would ask the area chair to verify if the authors have addressed all my concerns satisfactorily. In particular, im still not sure about the relation between (P4) and regularized PFL (P3) which is quite generic and also allows for reguarlizers that are indicator functions of sets.