PartialFed: Cross-Domain Personalized Federated Learning via Partial Initialization

1 Introduction . Endless collection of texts , images and videos for training data-hungry models increases potential threats to the safety systems . Any attacks to these data centers could cause billions of information leakage . A safer way is to keep the user data purely local on their devices . But this contradicts with the widely adopted stochastic gradient descent ( SGD ) training procedure , which usually requires data communication for random batch sampling . Federated Learning [ 12 ] develops a paradigm for training large models under such situation . Locally trained models are aggregated using FedAvg [ 19 ] and then served as initialization model for the next local iteration . The training data confidentiality is achieved by only allowing transferring the model ’ s parameters ( rather than the data ) under different cryptograph algorithms , which include differential privacy [ 7 ] , homomorphic encryption [ 37 ] , block chain [ 22 ] , etc . However , recent studies [ 39 , 26 , 25 ] have shown that FedAvg does not provide satisfactory results in the presence of data heterogeneity . A major problem is that all models are designed to fit an “ average client ” [ 29 ] , which is difficult when local and global distributions are deviated from each other . This phenomenon is not uniquely observed . Researchers in multi-domain learning have also discovered that directly fitting non-identical domains into a single feature extractor is suboptimal [ 3 ] . The leading solution to this problem is to reconfigure the network into domain-agnostic and 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . domain-specific layers [ 23 , 33 ] . Similar conclusion has also been drawn simultaneously in the field of multi-task learning [ 27 ] . The performance is boosted by letting each task choosing task-agnostic and task-specific layers on the fly . Based on such generous observations , we propose a Personalized Federated Learning ( PFL ) method from the idea of client-agnostic and client-specific initialization . Initialization is considered because it plays the essential role of transferring global knowledge in FL . Figure 1 dipicts the overall architecture . At the heart of our algorithm is a mixed initialization strategy . Instead of fully utilizing the averaged global parameters for initialization , clients will only select a fraction of them , and load the remaining parameters from previous local models . The selection process is decided by a customized loading strategy , which might vary from client to client and time to time . To validate PartialFed thoroughly , we first propose PartialFed-Fix , where the loading strategies are inspired by human priors in functionality and classification of different parts of the network . Then we propose an automatic strategy which is learned jointly with the network parameters by gradient descent . We name this algorithm PartialFed-Adaptive . In the case of PartialFed-Fix , our algorithm can be viewed as a federated version of multi-path networks in the multi-domain learning [ 24 , 33 ] . The shared parameters of different clients are jointly learned by FedAvg , while the client-specific parameters are only learned locally . Treating the network as a combination of global and local blocks makes it possible to learn the knowledge from other clients while keeping the local knowledge stored safely . Although PartialFed-Fix can achieve improved performance , we argue that the fixed loading strategy is a suboptimal solution for PFL since different clients may have different dependencies on the global model . For example , a client might prefer the classifier of the global model at the start of training , which often acquires better generalization ability . In the end , the client is likely to train a local classifier , which brings superior personalized performance . We show that our dynamic algorithm PartialFed-Adaptive is able to capture this change of parameter loading behaviors during the training process . Despite its simplicity , PartialFed gains surprising performance on many non i.i.d . FL experiments . We construct real world FL experiments by introducing cross-domain classification and detection dataset [ 30 , 33 ] . For example , on Office-Home dataset , our PartialFed-Fix and PartialFed-Adaptive surpass FedAvg by 4.88 % and 5.43 % on average accuracy , respectively . Similarly , on UODB dataset , PartialFed-Fix and PartialFed-Adaptive outperform FedAvg by 2.65 % and 2.68 % on AP50 , respectively . More interestingly , the adaptive loading strategy greatly reduce the possibility of clients getting inferior performance caused by the distribution deviation . 2 Related Work . Personalized Federated Learning The pioneering work of FedAvg [ 19 ] aims to train a global model for all clients . But the goal is hard to achieve when data distributions are non i.i.d .. Wang et al . [ 31 ] directly finetune the global model on local dataset for personalized performance . Mansour et al . [ 18 ] provide theoretical study of clustering , data interpolation and model interpolation in personalized learning . FedProx [ 13 ] introduces a regularization based algorithm . FedCurv [ 26 ] and FedCL [ 35 ] further consider using Elastic Weight Consolidation ( EWC ) for parameters importance estimation . T. Dinh et al . [ 28 ] theoretically improve the vanilla regularization with Moreau Envelopes . [ 5 , 8 ] consider using model interpolation to address the PFL problem . Meta learning is also utilized to learn a better initialization for personalized finetuning [ 6 ] . The most related works to this paper is parameter decoupling strategies which split networks into global and private layers . Arivazhagan et al . [ 2 ] consider global base network with personalized classifiers while FedBN [ 14 ] and SiloBN [ 1 ] consider private Batch Normalization ( BN ) layers . The special cases of PartialFed-Fix are equivalent to these methods , and we further examined their combinations and more strategies . Multi-Domain Learning The problem of training on heterogeneous distributions has been widely studied in multi-domain learning . The difference is that multi-domain learning aggregate data together while FL store data in silos . Our method is related to the recent work [ 3 , 23 , 24 ] , which aims to learn a single network across all domains with minimal number of task specific parameters . Bilen and Vedaldi [ 3 ] propose to use domain-specific batch and instance normalization while Rebuffi et al . [ 23 , 24 ] further take series and parallel residual adapters as the domain-specific parts . Wang et al . [ 33 ] extend the idea to object detection with SE-ResNet [ 10 ] and introduce a multi-domain detection benchmark . Data-driven architectures are also explored in [ 20 , 38 , 32 ] where network pruning or neural architecture search are employed . 3 Partially-Loaded Federated Learning . 3.1 Preliminary . Federated Learning We formally describe the problem setup in this section . Assume there are N ∈ N clients c ∈ C who are willing to participate in the joint training . Each client owns a private dataset Dc = { ( xci , yci ) : i ∈ { 1 , ... , nc } } , where nc is the cardinality of the dataset . We use n = ∑ c∈C nc to denote the total amount of data in all clients . The data distribution P c ( x , y ) varies from client to client . To be more specific , we do not put any i.i.d . assumptions on either P c ( x ) or P c ( y|x ) , which indicates domain gap and label difference are allowed in the framework . At the start of each federated iteration t , a central server will send a global model Wt to all clients as parameter initialization . Clients will then use their private data Dc to train Wt and get an updated Wct . All W c t , c ∈ C are then sent to server and fused by FedAvg : Wt+1 ← ∑ c∈C nc n Wct ( 1 ) The resulting model is used as initialization for the next federated iteration . When networks are structurally heterogeneous in each client , Wt denotes the maximum set of shareable parameters . The algorithm will stop on its convergence . Note that we do not incorporate client selection as in McMahan et al . [ 19 ] , which can be easily added if required . PartialFed Compared to vanilla FL , PartialFed changes the client ’ s parameter initialization process . After receiving global model Wt , each client c will additionally load the local model Wct−1 from last iteration ( See Alg . 1 ) . Instead of purely initializing training with global model Wt , PartialFed uses a mixed initialization , where each parameter is either loaded from global model Wt or local model Wct−1 . We hope that by partially loading global parameters , client can benefit from sharing knowledge while avoiding local knowledge forgetting caused by FedAvg . The choice of mixing strategy Act will be discussed in section 3.2 and 3.3 . All the models are finetuned locally in the last training iteration . A major advantage of PartialFed is the minimal changes to the original framework . Without sharing gradients , subset of data and customized attributes , clients can enjoy themselves from a highest parameter privacy protocol . In addition , since the server do not need extra operations like knowledge distillation [ 16 ] except for FedAvg , the uploaded models are not required to do any forward computation . This enables homomorphic encryption [ 37 ] and protects client from information leakage caused by model attacks [ 36 ] . Algorithm 1 PartialFed : Partially-Loaded Federated Learning ClientIteration ( c , Wt ) : Load Wct−1 from local storage Get partial loading strategy Act . partial initialization Get partially loaded initialization Wct , init by : Wct , init [ i ] = { Wct−1 [ i ] if A c t [ i ] = 0 Wt [ i ] if Act [ i ] = 1 , i denotes the index of parameter ( 2 ) Set Wct = W c t , init for each local epoch from 1 to E do . regular training for batch b ∈ Dc do w ← w − ηO ` ( w ; b ) , for w ∈Wct end for end for Save Wct on local storage Return Wct Server : Initialize W0 = [ w0,0 , w1,0 , ... , wk,0 ] at random for each round t = 1 , 2 , . . . do for client c ∈ C in parallel do Wct ← ClientIteration ( c , Wt ) . device training end for Wt+1 ← ∑ c∈C nc n W c t . execute FedAvg end for 3.2 PartialFed-Fix . We first discuss the fixed manual strategies in this section . Deep networks are often intuitively divided into different functional parts : feature extractor , classifier , etc . These intuitions can be transformed into corresponding loading strategies . An example strategy is shown in Equ.3 , where parameters except fully connected layer ( fc ) are loaded from global model . We empirically validate the efficiency by transferring commonly used assumptions to our algorithm in this section . Act [ i ] = { 1 if layer i is fc 0 else ( 3 ) '' Bottom-up v.s . Top-down '' One of the commonly accepted assumptions in deep network is that blocks close to input image are more related to low-level features while top blocks are associated with high-level concepts . Sharing bottom layers assumes that different clients share similar low-level texture features while the opposite assumes high-level visual concepts are composed analogously . We examine whether PFL training benefits more from sharing global `` feature extractor '' or global `` classifier '' . To be more specific , we use ResNet18 [ 9 ] for these experiments , which has four stages ( each with two ResBlock ) and a fully-connected layer ( Figure 2 ) . A detailed description of the used dataset can be found in section 4.1 . The results are listed in Table 1 . Firstly , all experiments with partial parameters loading outperform the full parameter loading , which proves the superiority of the proposed PartialFed . Secondly , there are no significant gap with additional parameters loading from global s1 to s1-s4 , but loading global fc in this case severely degrades the performance . Thirdly , we find there ’ s consistent accuracy drop from strategy fc to s2-fc . We hypothesis that sharing too much high-level concepts will hurt the personalized federated learning . Batch Normalization Batch statistics are considered as a key factor when heterogeneous distributions occur and has been widely studied in domain adaptation [ 15 ] , multi-domain learning [ 3 ] , and federated learning [ 1 , 14 ] . We study how the loading strategies of BN will affect the PFL training . Table 2 displays the results on Office-Home dataset . Loading BN locally ( i.e . w/o BN ) gains improvements on all four domains . Given that the parameters of BN layers are only a very small fraction of ResNet , the improvement is non-trivial . On the other hand , since the BN layer is such brittle in non i.i.d . settings , we hope that a multi-distribution normalization strategy can be proposed in the future . Recalling that previous section shows loading fc and s1-s4 simultaneously could harm the performance . We hope this rule can be combined with the BN rule . This leads to the last row of Table 2 , where both global BN and global fc are not loaded . Surprisingly , this simple modification significantly improves FedAvg from 75.44 % to 80.32 % , which is 4.88 % improvements . Skip Loading Except for the `` shared feature extractor - private classifier '' network partition , another widely adopted routine in multi-domain and multi-task learning is a loop of `` share-private-shareprivate- ... '' sub-networks [ 33 , 27 ] . This paradigm is found to be helpful when tasks interference happens . This paragraph examines the performance of skip sharing with PartialFed . For simplicity , we label each ResBlock in ResNet18 with letters as in Figure 2 . The same number of blocks ( 4 blocks ) is loaded globally for all experiments . The input convolution and fc layer are loaded as default . Experimental results are revealed in Table 3 . Methods that load adjacent layers ( e.g . AaBb ) get similar mean accuracy except for the CcDd strategy . Since CcDd is a kind of top-down strategies mentioned above , which does not perform well in general , we infer this is the reason for its inferior performance . Interestingly , the two strategies with skip loading , ABCD and abcd , get higher performance than all the other strategies , which justify the idea of skip sharing . We attribute the efficiency of skip loading to the fact that clients have the demand of learning client-specific knowledge at all levels of the network instead of only a part of them .

Summary: This paper explores a simple but novel modification to FL approaches to better deal with settings involving non-homogeneous clients. They find that by allowing clients to mix local and global parameters for initializing local training, the added flexibility can notably improve the performance of the respective local end-models. They perform both an exploration of natural fixed mixing strategies and also propose an adaptive mixing strategy that allows for further improvements on average.