Parameterized Knowledge Transfer for Personalized Federated Learning

1 Introduction . Federated Learning ( FL ) [ 1 ] has emerged as an efficient paradigm to collaboratively train a shared machine learning model among multiple clients without directly accessing their private data . By periodically aggregating parameters from the clients for global model updating , it can converge to high accuracy and strong generalization . FL has shown its capability to protect user privacy while there remains a crucial challenge that significantly degrades the learning performance , i.e. , statistic heterogeneity in users ’ local datasets . Given the Non-Independent and Identically Distributed ( Non-IID ) user data , the trained global model often can not be generalized well over each client [ 2–5 ] . To deal with the above issues , employing personalized models appears to be an effective solution in FL , i.e. , personalized federated learning ( pFL ) . Recent works regarding pFL include regularizationbased methods [ 6–8 ] ( i.e. , pFedMe [ 6 ] , L2SGD [ 7 ] , FedAMP [ 8 ] ) , meta-learning-based Per-FedAvg ∗Corresponding author 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . [ 9 ] and cluster-based IFCA [ 10 , 11 ] . However , in order to aggregate the parameters from all clients , it is inevitable for them to have identical model structure and size . Such constraints would prevent status quo pFL methods from further application in practical scenarios , where clients are often willing to own unique models , i.e. , with customized neural architectures to adapt to heterogeneous capacities in computation , communication and storage space , etc . Motivated by the paradigm of Knowledge Distillation ( KD ) [ 12–16 ] that knowledge can be transferred from a neural network to another via exchanging soft predictions instead of using the whole model parameters , KD-based FL training methods have been studied [ 12 , 14–19 ] to collaboratively train heterogeneous models in a privacy-preserving way . However , these works have neglected the further personalization requirement of FL clients , which can be well satisfied via the personalized knowledge transfer for heterogeneous FL users . In this paper , we seek to develop a novel training framework that can accommodate heterogeneous model structures for each client and achieve personalized knowledge transfer in each FL training round . To this end , we formulate the aggregation phase in FL to a personalized group knowledge transfer training algorithm dubbed KT-pFL , whose main idea is to allow each client to maintain a personalized soft prediction at the server that can be updated by a linear combination of all clients ’ local soft predictions using a knowledge coefficient matrix . The principle of doing so is to reinforce the collaboration between clients with similar data distributions . Furthermore , to quantify the contribution of each client to other ’ s personalized soft prediction , we parameterize the knowledge coefficient matrix so that it can be trained simultaneously with the models following an alternating way in each iteration round . We show that KT-pFL not only breaks down the barriers of homogeneous model restriction , which requires to transfer the entire parameters set in each round , whose data volume is much larger than that of the soft prediction , but also improves the training efficiency by using a parameterized update mechanism . Experimental results on different datasets and models show that our method can significantly improve the training efficiency and reduce the communication overhead . Our contributions are : • To the best of our knowledge , this paper is the first to study the personalized knowledge transfer in FL . We propose a novel training framework , namely , KT-pFL , that maintains a personalized soft prediction for each client in the server to transfer knowledge among all clients . • To encourage clients with similar data distribution to collaborate with each other during the training process , we propose the ‘ knowledge coefficient matrix ’ to identify the contribution from one client to others ’ local training . To show the efficiency of the parameterized method , we compared KT-pFL with two non-parameterized learning methods , i.e. , TopK-pFL and Sim-pFL , which calculate the knowledge coefficient matrix on the cosine similarity between different model parameters . •We provide theoretical performance guarantee for KT-pFL and conduct extensive experiments over various deep learning models and datasets . The efficiency superiority of KT-pFL is demonstrated by comparing our proposed training framework with traditional pFL methods . 2 Related Work . Personalized FL with Homogeneous Models . Recently , various approaches have been proposed to realize personalized FL with homogeneous local model structure , which can be categorized into three types according to the number of global models applied in the server , i.e. , single global model , multiple global models and no global model . single global model type is a close variety of conventional FL , e.g. , FedAvg [ 1 ] , that combine global model optimization process with additional local model customization , and consist of four different kinds of approaches : local fine-tuning [ 20–23 ] , regularization ( e.g. , pFedMe [ 6 ] , L2SGD [ 7 , 24 ] , Ditto [ 25 ] ) , hybrid local and global models [ 11 , 26 , 27 ] and meta learning [ 9 , 28 ] . All of these pFL methods apply a single global model , and thus limit the customized level of the local model at the client side . Therefore , some researchers [ 8 , 10 , 11 ] propose to train multiple global models at the server , where clients are clustered into several groups according to their similarity and different models are trained for each group . FedAMP [ 8 ] and FedFomo [ 29 ] can be regarded as special cases of the clustered-based method that each client owns a personalized global model at the server side . As a contrast , some literature waive the global model to deal with the heterogeneity problem [ 30 , 31 ] , such as multi-task learning based ( i.e. , MOCHA [ 31 ] ) and hypernetwork-based framework ( i.e. , FedHN [ 30 ] ) . However , all these methods require aggregating the model parameters from the clients , who have to apply identical model structure and size , which hinders further personalization , e.g. , employing personalized model architectures for heterogeneous clients is not feasible . Heterogeneous FL and Knowledge Distillation . To enable heterogeneous model architectures in FL , Diao et al . [ 32 ] propose to upload a different subset of global model to the server for aggregation with the objective to produce a single global inference model . Another way of personalization is to use Knowledge Distillation ( KD ) in heterogeneous FL systems ( [ 12 , 14 , 15 , 17–19 , 33 , 34 ] ) . The principle is to aggregate local soft-predictions instead of local model parameters in the server , whereby each client can update the local model to approach the averaged global predictions . As KD is independent with model structure , some literature [ 13 , 16 ] are proposed to take advantage of such independence to implement personalized FL with heterogeneous models at client sides . For example , Li et al . [ 13 ] propose FedMD to perform ensemble distillation for each client to learn well-personalized models . Different from FedMD that exchanging soft-predictions between the clients and the server , FedDF [ 16 ] first aggregates local model parameters for model averaging at the server side . Then , the averaged global models can be updated by performing knowledge transfer from all received ( heterogeneous ) client models . To sum up , most of these schemes construct an ensembling teacher by simply averaging the teachers ’ soft predictions , or by heuristically combining the output of the teacher models , which are far away from producing optimal combination of teachers . In our framework , KD is used in a more efficient way that the weights of the clients ’ soft predictions are updated together with the model parameters during every FL training iteration . 3 Problem Formulation . We aim to collaboratively train personalized models for a set of clients applying different model structures in FL . Consider supervised learning whose goal is to learn a function that maps every input data to the correct class out of C possible options . We assume that there are N clients , and each client n can only access to his private dataset Dn : = { xni , yi } , where xi is the i-th input data sample , yi is the corresponding label of xi , yi ∈ { 1 , 2 , · · · , C } . The number of data samples in dataset Dn is denoted by Dn . D = { D1 , D2 , · · · , DN } , D = ∑N n=1Dn . In conventional FL , the goal of the learning system is to learn a global model w that minimizes the total empirical loss over the entire dataset D : min w L ( w ) : = N∑ n=1 Dn D Ln ( w ) , where Ln ( w ) = 1 Dn Dn∑ i=1 LCE ( w ; xi , yi ) , ( 1 ) where Ln ( w ) is the n-th client ’ s local loss function that measures the local empirical risk over the private dataset Dn and LCE is the cross-entropy loss function that measures the difference between the predicted values and the ground truth labels . However , this formulation requires all local clients to have a unified model structure , which can not be extended to more general cases where each client applies a unique model . Therefore , we need to reformulate the above optimization problem to break the barrier of homogeneous model structure . Besides , given the Non-IID clients ’ datasets , it is inappropriate to just minimize the total empirical loss ( i.e. , minw1 , ··· , wN L ( w1 , · · · , wN ) : = ∑N n=1 Dn D Ln ( w n ) ) . To that end , we propose the following training framework : Definition 3.1 . Let s ( wn , x̂ ) denote the collaborative knowledge from client n , and x̂ denote a data sample from a public dataset Dr that all clients can access to . Define the personalized loss function of client n as Lper , n ( wn ) : = Ln ( wn ) + λ ∑ x̂∈Dr LKL ( N∑ m=1 cmn · s ( wm , x̂ ) , s ( wn , x̂ ) ) , ( 2 ) where λ > 0 is a hyper-parameters , LKL stands for Kullback–Leibler ( KL ) Divergence function and is added to the loss function to transfer personalized knowledge from a teacher to another . cmn is the knowledge coefficient which is used to estimate the contribution from client m to n. The second term in ( 2 ) allows each client to build his own personalized aggregated knowledge in the server and enhance the collaboration effect between clients with large c. The concept of collaborative knowledge can either refer to soft predictions or model parameters depending on the definition in different situations . For example , s ( wn , x̂ ) can be deemed to be a soft prediction of the client n , which are calculated with the softmax of logits zn , i.e. , s ( wn , x̂ ) = exp ( z n c /T ) ∑C c=1 exp ( z n c /T ) , logits zn is the output of the last fully connected layer on client n ’ s model , T is the temperature hyperparameter of the softmax function . Definition 3.2 . We define c ∈ RN×N as the knowledge coefficient matrix : c = c11 c12 · · · c1N c21 c22 · · · c2N ... ... . . . ... cN1 cN2 · · · cNN . ( 3 ) Our objective is to minimize min w , c L ( w , c ) : = N∑ n=1 Dn D Lper , n ( wn ) + ρ‖c− 1 N ‖2 , ( 4 ) where w = [ w1 , · · · , wN ] ∈ R ∑N n=1 dn is the concatenated vector of all weights , dn represents the dimensions of model parameter wn . 1 ∈ Rn2 is the identity matrix whose elements are all equal to 1 . The second term in ( 4 ) is a regularization term that ensures generalization ability of the whole learning system . Without the regularization term , a client with a completely different data distribution tends to set large values of the knowledge coefficient ( i.e. , equal to 1 ) , and no collaboration will be conducted during training in this case . ρ is a regularization parameter that is larger than 0 .

Summary: The work aims to address the problem of creating adequate personalized models in a federated learning regime. The problems stems from a general observation that global models do no generalize well in the Federated Learning (FL) scenarios. The authors propose to use the idea of knowledge distillation to properly transfer knowledge to different clients via soft prediction mechanics. It is based on the knowledge about predictions that can be obtained on a public dataset (also referred as collaborative knowledge). At the server side, the proposed model learns the corresponding knowledge transfer coefficients that align inputs from multiple clients and hence distill the knowledge. The training is performed in an alternating manner by switching between knowledge coefficients (updated at the server side) and the task-specific parameters of the model (updated at the client side). Advantages of the proposed approach are empirically demonstrated in comparison against other FL algorithms w.r.t. both accuracy and communication efficiency on several public datasets.