Personalized Heterogeneous Federated Learning with Gradient Similarity

1 INTRODUCTION . With the popularity of smartphones , personal computers , and other devices , the data stored on them has increased dramatically . These data are related to each other but exist independently on different devices . Also , their owners are often not willing to share their private data , which prompted Federated Learning ( FL ) ( McMahan et al. , 2017 ) . Conventional FL is a distributed machine learning framework in which multiple clients train a global model together under a central server ’ s coordination without sharing their local data . It ’ s common process is as follows : ( 1 ) At each global iteration , the central server broadcasts its current global model to the client ; ( 2 ) With the help of the global model , the client updates its local model which trained on its local data , and sends it back to the central server ; ( 3 ) The central server aggregates the local models to get a new global model . The FL process will repeat these three steps until convergence . Under the premise of privacy protection , it can obtain a shared global model with higher generality . Nonetheless , the wide application of FL still is facing the following two major challenges : • Statistical heterogeneity . As the user preference of each device may be different , the data distribution on each device may also be inconsistent . Consequently , the data as a whole on all the clients may be unbalanced and non-IID . This may cause undesirable performance of the FL participants . • Systems heterogeneity . Due to the differences of users ’ hardware and network bandwidth , conventional FL models are prone to generate stragglers whose models are easy to be discarded after server ’ s time-out . Firstly , to address the statistical heterogeneity , the personalization of FL is widely studied . The majority of the available personalized FL ( PFL ) methods focus on one shared global model ’ s parameters updating by client ’ s fine-tune . FedAvg ( McMahan et al. , 2017 ) is treated by ( Jiang et al. , 2019a ) as a client ’ s meta-learning-like process fine-tuning on one shared global model with higher accuracy . FedMeta ( Chen et al. , 2018 ) , also a personalized FL based on meta-learning , generates a local and personalized model for each client using a common global model . Different from the aforementioned methods , we propose a Subclass Personalized Federated Learning ( SPFL ) algorithm , which utilizes the Softmax Normalized Gradient Similarity ( SNGS ) to generate global models with each specific to one client for its parameter fine-tune . Secondly , to tackle the systems heterogeneity , the server of an asynchronous FL algorithm named FedAsync ( Xie et al. , 2019 ) , can immediately perform aggregation after receiving a local model , reducing the waiting time for stragglers . In TWAFL ( Chen et al. , 2019 ) , a time-weighted aggregation strategy is used by server to handle stragglers according to the staleness of the received model parameters . Although these asynchronous FL algorithms can overcome some difficulties encountered in synchronous aggregation , such as waiting for the response from client , they are still inefficient for directly aggregating the received local model on the server , without taking the iteration gap of stragglers ’ local models at global communication round into consideration . Therefore , we propose the Leap Gradient Approximation ( LGA ) algorithm using SNGS to predict the local model of the stragglers in each global communication round . These local models estimated by server are also aggregated with non-stragglers ’ local models on the server to attained one shared global model . Further , to consider the real scenarios where system heterogeneity and statistical heterogeneity often co-exist , personalization is employed in our asynchronous FL . Before each individual personalized model for client is dispatched from server , it is computed by using the shared global model and the SNGS value between server and that client . We refer this personalized asynchronous FL approach as Personalized Leap Gradient Approximation ( PLGA ) algorithm . Contribution The main contributions of this paper towards personalized heterogeneous FL can be summarized as follows : ( 1 ) We propose the SPFL algorithm for both IID and non-IID data in personalized synchronous FL , where the server uses the SNGS to weight the relationship between clients , and delivers the personalized global model to each client ; ( 2 ) We propose the PLGA algorithm for the personalized asynchronous FL , in which the server also applies the SNGS to weight the relationship between client and itself , and uses the first-order Taylor expansion of gradient to approximate the model of the delayed clients ; ( 3 ) The stage strategy of ResNet is applied to improve the performance of both SPFL and PLGA algorithms . 2 RELATED WORK . In general , the personalized process in FL can be achieved by a variety of ways . The work ( Hard et al. , 2018 ) shows that by adding user-related features into the local inference process to get a personalized model may be better than the global one . The completely decentralized algorithms for each client by smoothing some similar tasks or model parameters between similar data distributions are studied in ( Bellet et al. , 2018 ; Vanhaesebrouck et al. , 2017 ) . Algorithms are also proposed to achieve personalization by jointly learning the model on a similar graph structure ( Lalitha et al. , 2019 ; He et al. , 2021 ) . FedFOMO ( Zhang et al. , 2020 ) is another approach to PFL , where each client only federates with other relevant clients for model aggregation . The pFedHN ( Shamsian et al. , 2021 ) algorithm is proposed using hypernetwork to generate different personalized model for each client . Furthermore , PFL can also be studied by its relationship with other machine learning paradigms . The close relationship between FL and meta-learning is analyzed by ( Khodak et al. , 2019 ; Jiang et al. , 2019b ) . Following the classic meta-learning algorithm named MAML ( Finn et al. , 2017 ) , Reptile algorithm ( Nichol et al. , 2018 ) is proposed , which has a close relationship with the Federated Averaging algorithm . When each individualized model learning is treated as a task , such a scenario can be viewed as a multi-task learning ( Zhang & Yang , 2017 ) , in which each task learns a taskrelated model . MOCHA ( Smith et al. , 2017 ) is designed by regarding FL as a multi-task , which solves communication efficiency and fault tolerance . To provide a good global model , there are PFL applying transfer learning and domain transfer ( Mansour et al. , 2009 ; Cortes & Mohri , 2014 ; BenDavid et al. , 2010 ) . Model mixing is another strategy to approach PFL ( Deng et al. , 2020 ; Hanzely & Richtárik , 2020 ; Liang et al. , 2020 ) . Alternatively , Other PFL methods apply clustering ( Mansour et al. , 2020 ; Ghosh et al. , 2020 ; Duan et al. , 2020 ) . Finally , pFedMe ( Dinh et al. , 2020 ) applies Moreau Envelopes as clients ’ regularized loss functions to achieve personalization . There are just few studies on asynchronous FL . TWAFL ( Chen et al. , 2019 ) aggregates a temporally weighted local models on the server . FedAsync ( Xie et al. , 2019 ) introduced an algorithm to balance the parameters of the updated model and the last global model . Meanwhile , the distillation technique has also been applied in ( Li & Wang , 2019 ; Bistritz et al. , 2020 ) . 3 METHOD . Our method generates a personalized model for each client during the server aggregation process . Therefore , the initialization model of each client is personalized . At the same time , we also consider the contribution of different neural network layers to the personalized model . We will introduce the application of this personalized algorithm to synchronous FL and asynchronous FL in the following subsections . Algorithm 1 Subclass Personalized Federated Learning ( SPFL ) Algorithm Input : initialized global model parameters w0 ; initialized SNGS matrix S̃s ∈ RN×N ; initialized personalized local model of each client ’ s { w0i } = w0 ; update frequency Γ and the number S of SNGS matrix ; client learning rate β ; server learning rate α ; local epoch E ; local batch size B of each client Output : Model Parameters { wi } Def MainLoop : for t = 0 , 1 , . . . do perform the following steps in parallel for each client Ci ∈ C if t % Γ == 0 then wtgi = 1 N ∑N i w t i { gti } = ClientUpdate ( wtgi , Ci , β ) S̃s = UpdateS̃Matrix ( { gti } ) { gti } = ClientUpdate ( wtgi , Ci , β ) { wt+1si } = ServerUpdate ( { wtgi } , { gti } , S̃s , α ) Def ClientUpdate ( wtgi , Ci , β ) : wti = w t gi for e = 0 ; e < E ; e++ do for batch size b in B do wti ← wti − 2β∇Li ( f ( wti , xb ) , yb ) return { wtgi − wti } Def UpdateS̃Matrix ( { gi } ) : for s ∈ S do perform ( 3 ) and ( 6 ) to return S̃s ( i , j ) Def ServerUpdate ( { wtsi } , { gti } , S̃s , α ) : for s ∈ S do perform ( 7 ) to get wt+1si return { wt+1gi } Definition of personalized FL algorithm In FL , there is a set of N clients C = { Ci } , i = 1 , . . . , N . Define the dataset on all clients as D = { Di } , i = 1 , . . . , N , with Di as the corresponding dataset of Ci . The sample size of each datasets Di is |Di| . ( xj , yj ) is a sample of Di , with xj ∈ Rd as the corresponding input feature and yj ∈ Rt as the corresponding input label . The prediction model of client Ci defined as ŷj = f ( wi , xj ) , where wi ∈ Rz is the model parameter . The dataset on each client Ci is assumed to obey a distribution of Pi . And the loss function on each client Ci is Li ( wi ) . The optimization goal for each client is : argmin wi Li ( wi ) = EDi∼Pi [ Li ( f ( wi , xj ) , yj ) ] ≈ 1 |Di| |Di|∑ j Li ( f ( wi , xj ) , yj ) . ( 1 ) 3.1 SUBCLASS PERSONALIZED FL ( SPFL ) ALGORITHM . We compared the performance of each client ’ s independent training ( i.e. , the model in each client performs SGD only on local data available , and model averaging is not performed ) and conventional global FL algorithms , such as FedAvg . The results in Appendix B. show that FL usually outperforms each client ’ s independent training on IID data , demonstrating the global model trained with multiple client data improves the model generalizability ability of each client . On the other hand , for some clients , independent training has better performance on non-IID data . Our experiment in Appendix B illustrates that the global FL model can not adapt well to the heterogeneous data . To tackle this problem , we propose a personalized FL algorithm that improved the performance of conventional global FL algorithms on non-IID data . At the same time , it still maintains a high performance on IID data . Firstly , we introduce a similarity matrix S ∈ RN×N to model the relationship between different clients , with S ( i , j ) represents the similarity between Ci and Cj . If the data distributions Pi and Pj are similar , then the value of S ( i , j ) is relatively high , and vice versa . For any two clients , Ci and Cj , supposing they receive the global model parameters wt at the same time t. After the local update of both clients using wt , their local model parameters are wti and w t j , respectively . Then the SGD updates of the corresponding gradients of wti and w t j are : gti = w t − wti , gtj = w t − wtj . ( 2 ) The correlation between Ci and Cj based on the global model wt is measured by the cosine similarity of the gradient updates gti and g t j , as calculated in ( 3 ) : S ( i , j ) = gti T gtj ‖gti‖‖gtj‖ , ( 3 ) where ‖ · ‖ is the vector L2 normal . And the aggregation strategy on the server for Ci is as follows : wt+1i = w t i − α N∑ j |Dj | |D| S ( i , j ) g t j , ( 4 ) where α is the learning rate of the server . As the gradient itself is extremely directionality sensitive , when S ( i , j ) is negative , multiplying it with the original gradient gtj will reverse the gradient direction . Therefore , softmax function is applied to make S ( i , j ) nonnegative and normalized as S̃ ( i , j ) . S̃ ( i , j ) , the softmax version of similarity relationship between Ci and Cj , is defined as : S̃ ( i , j ) = eS ( i , j ) ∑N j e S ( i , j ) . ( 5 ) The experimental results in Appendix B show that different layers influence personalization differently in the convolutional neural network ( CNN ) . Inspired by ResNet , we use stage strategy that contains a set of network layers to measure the similarity between clients . Therefore , the similarity matrix S̃ is calculated as S̃s based on a stage as in ( 6 ) , where its subscript s indicates a stage : S̃s ( i , j ) = eSs ( i , j ) ∑N j e Ss ( i , j ) . ( 6 ) Finally , the aggregation strategy on the server for Ci base on stage is defined as : wt+1si = w t si − α N∑ j |Dj | |D| S̃s ( i , j ) g t j , ( 7 ) where α is the learning rate of the server , gtsj and w t si are the updated gradient and model parameters of the stage s on the client i at iteration round t , respectively . As the model changes with iteration round t , so does the gradient of the model . Consequently , the similarity matrix S̃ obtained based on the gradient will also change . Therefore , the similarity matrix is updated regularly during the training process . During each update , the model parameters of all clients are aggregated with equal weight to get wt of ( 2 ) , and the similarity matrix is calculated based on wt . In the view of each aggregation , it is still similar to the federated average algorithm . However , each client obtains a personalized model based on its similarity matrix in the training rounds in each aggregation . In this process , each client is individually regarded as a subclass , and the algorithm is called Subclass Personalized Federated Algorithm ( SPFL ) . The algorithm is shown in Alg . 1 .

Summary: This paper proposes two methods for personalized federated learning, one synchronous and one asynchronous. The general approach taken in both cases is to adapt the weights when averaging information from different clients, so that clients with more similar gradients are given more weight in the update for each client. The two approaches are called SPFL (synchronous) and PLGA (asynchronous). The approaches are illustrated on small image classification tasks.