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 proposed a novel idea PartialFed, which loads a subset of the global model’s parameters rather than loading the entire model. Specifically, they validate the algorithm with manually decided loading strategies inspired by various expert priors, named PartialFed-Fix. Then they develop PartialFed-Adaptive, which automatically selects a personalized loading strategy for each client.