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 proposed a personalized federated learning algorithm which takes into account the similarity of gradient of different users to update the model. More formally, the authors define $\tilde{S}(i,j)$ as a measure of similarity between the gradients of two user $i$ and $j$, and then update the model of user $i$by weighting the gradient of user $j$ by $\tilde{S}(i,j)$. The authors study their method in various numerical settings.