Adaptive Personalized Federated Learning

1 INTRODUCTION . With the massive amount of data generated by the proliferation of mobile devices and the internet of things ( IoT ) , coupled with concerns over sharing private information , collaborative machine learning and the use of federated optimization ( FO ) is often crucial for the deployment of large-scale machine learning ( McMahan et al. , 2017 ; Kairouz et al. , 2019 ; Li et al. , 2020b ) . In FO , the ultimate goal is to learn a global model that achieves uniformly good performance over almost all participating clients without sharing raw data . To achieve this goal , most of the existing methods pursue the following procedure to learn a global model : ( i ) a subset of clients participating in the training is chosen at each round and receive the current copy of the global model ; ( ii ) each chosen client updates the local version of the global model using its own local data , ( iii ) the server aggregates over the obtained local models to update the global model , and this process continues until convergence ( McMahan et al. , 2017 ; Mohri et al. , 2019 ; Karimireddy et al. , 2019 ; Pillutla et al. , 2019 ) . Most notably , FedAvg by McMahan et al . ( 2017 ) uses averaging as its aggregation method over local models . Due to inherent diversity among local data shards and highly non-IID distribution of the data across clients , FedAvg is hugely sensitive to its hyperparameters , and as a result , does not benefit from a favorable convergence guarantee ( Li et al. , 2020c ) . In Karimireddy et al . ( 2019 ) , authors argue that if these hyperparameters are not carefully tuned , it will result in the divergence of FedAvg , as local models may drift significantly from each other . Therefore , in the presence of statistical data heterogeneity , the global model might not generalize well on the local data of each client individually ( Jiang et al. , 2019 ) . This is even more crucial in fairness-critical systems such as medical diagnosis ( Li & Wang , 2019 ) , where poor performance on local clients could result in damaging consequences . This problem is exacerbated even further as the diversity among local data of different clients is growing . To better illustrate this fact , we ran a simple experiment on MNIST dataset where each client ’ s local training data is sampled from a subset of classes to simulate heterogeneity . Obviously , when each client has samples from less number of classes of training data , the heterogeneity among them will be high and if each of them has samples from all classes , the distribution of their local training data becomes almost identical , and thus heterogeneity will be low . The results of this experiment are depicted in Figure 1 , where the generalization and training losses of the global models of the FedAvg ( McMahan et al. , 2017 ) and SCAFFOLD ( Karimireddy et al. , 2019 ) on local data diverge when the diversity among different clients ’ data increases . This observation illustrates that solely optimizing for the global model ’ s accuracy leads to a poor generalization of local clients . To embrace statistical heterogeneity and mitigate the effect of negative transfer , it is necessary to integrate the personalization into learning instead of finding a single consensus predictor . This pluralistic solution for FO has recently resulted in significant research in personalized learning schemes ( Eichner et al. , 2019 ; Smith et al. , 2017 ; Dinh et al. , 2020 ; Mansour et al. , 2020 ; Fallah et al. , 2020 ; Li et al. , 2020a ) . To balance the trade-off between the benefit from collaboration with other users and the disadvantage from the statistical heterogeneity among different users ’ domains , in this paper , we propose an adaptive personalized federated learning ( APFL ) algorithm which aims to learn a personalized model for each device that is a mixture of optimal local and global models . We theoretically analyze the generalization ability of the personalized model on local distributions , with dependency on mixing parameter , the divergence between local and global distributions , as well as the number of local and global training data . To learn the personalized model , we propose a communication efficient optimization algorithm that adaptively learns the model by leveraging the relatedness between local and global models as learning proceeds . As it is shown in Figure 1 , by progressively increasing the diversity , the personalized model found by the proposed algorithm demonstrates a better generalization compared to the global models learned by FedAvg and SCAFFOLD . We supplement our theoretical findings with extensive corroborating experimental results that demonstrate the superiority of the proposed personalization schema over the global and localized models of commonly used federated learning algorithms . 2 PERSONALIZED FEDERATED LEARNING . In this section , we propose a personalization approach for federated learning and analyze its statistical properties . Following the statistical learning theory , in a federated learning setting each client has access to its own data distributionDi on domain Ξ : = X×Y , whereX ∈ Rd is the input domain and Y is the label domain . For any hypothesis h ∈ H the loss function is defined as ` : H×Ξ→ R+ . The true risk at local distribution is denoted by LDi ( h ) = E ( x , y ) ∼Di [ ` ( h ( x ) , y ) ] . We use L̂Di ( h ) to denote the empirical risk of h on distribution Di . We use D̄ = ( 1/n ) ∑n i=1Di to denote the average distribution over all clients . 2.1 PERSONALIZED MODEL . In a standard federated learning scenario , where the goal is to learn a global model for all devices cooperatively , the learned global model obtained by minimizing the joint empirical distribution D̄ , i.e. , minh∈H L̂D̄ ( h ) by proper weighting . However , as alluded to before , a single consensus predictor may not perfectly generalize on local distributions when the heterogeneity among local data shards is high ( i.e. , the global and local optimal models drift significantly ) . Meanwhile , from the local user perspective , the key incentive to participate in “ federated ” learning is the desire to seek a reduction in the local generalization error with the help of other users ’ data . In this case , the ideal situation would be that the user can utilize the information from the global model to compensate for the small number of local training data while minimizing the negative transfer induced by heterogeneity among distributions . This motivates us to mix the global model and local model with a controllable weight as a joint prediction model , namely , the personalized model . Here we formally introduce our proposed adaptive personalized learning schema , where the goal is to find the optimal combination of the global and the local models , in order to achieve a better client-specific model . In this setting , global server still tries to train the global model by minimizing the empirical risk on the aggregated domain D̄ , i.e. , h̄∗ = arg minh∈H L̂D̄ ( h ) , while each user trains a local model while partially incorporating the global model , with some mixing weight αi , i.e. , ĥ∗loc , i = arg minh∈H L̂Di ( αih + ( 1 − αi ) h̄∗ ) . Finally , the personalized model for ith client is a convex combination of h̄∗ and ĥ∗loc , i : hαi = αiĥ ∗ loc , i + ( 1− αi ) h̄∗ , ( 1 ) It is worth mentioning that , hαi is not necessarily the minimizer of empirical risk L̂Di ( · ) , because we optimize ĥ∗loc , i with partially incorporating the global model . Example 1 . Let us illustrate a simple situation where mixed model does not necessarily coincide with local ERM model . To this end , consider a setting where the hypothesis class H is the set of all vectors in R2 , lying in ` 2 unit ball : H = { h ∈ R2 : ‖h‖2 ≤ 1 } . Assume the local empirical minimizer is known to be [ 1 , 0 ] > , and h̄∗ = [ −1 , 0 ] > , and α is set to be 0.5 . Now , if we wish to find a ĥ ∗ loc , i , such that hαi = α ∗ ĥ ∗ loc , i + ( 1 − α ) ∗ h̄ ∗ coincides with local empirical minimizer , we have to solve : 0.5∗h+0.5∗ [ −1 , 0 ] > = [ 1 , 0 ] > , subject to ‖h‖2 ≤ 1 . This equation has no feasible solution , implying that it is not necessarily true that hαi coincides with local empirical minimizer . In fact , in most cases , as we will show in the convergence of the proposed algorithm , hαi will incur a residual risk if evaluated on the training set drawn from Di . 2.2 GENERALIZATION GUARANTEES . We now characterize the generalization of the mixed model . We present the learning bounds for classification and regression tasks . For classification , we consider a binary classification task , with squared hinge loss ` ( h ( x ) , y ) = ( max { 0 , 1− yh ( x ) } ) 2 . In the regression task , we consider the MSE loss ` ( h ( x ) , y ) = ( h ( x ) − y ) 2 . Even though we present learning bounds under these two loss functions , our analysis can be generalized to any convex smooth loss . Before formally presenting the generalization bound , we introduce the following quantity to measure the empirical complexity of a hypothesis classH over a training set S. Definition 1 . Let S be a fixed set of samples and consider a hypothesis class H. The worst case disagreement between a pair of models measured by absolute loss is quantified by : λH ( S ) = suph , h′∈H 1 |S| ∑ ( x , y ) ∈S |h ( x ) − h′ ( x ) | . The empirical discrepancy characterizes the complexity of hypothesis class over some finite set . The similar concepts are also employed in the related multiple source PAC learning or domain adaption ( Kifer et al. , 2004 ; Mansour et al. , 2009 ; Ben-David et al. , 2010 ; Konstantinov et al. , 2020 ; Zhang et al. , 2020 ) . We now state the main result on the generalization of the proposed personalization schema . The proof of the theorem is provided in Appendix D. Theorem 1 . Let hypothesis classH be compact closed set with finite VC dimension d. Assume loss function ` is Lipschitz continuous with constant G , and bounded in [ 0 , B ] . Then with probability at least 1−δ , there exists a constantC , such that the risk of the mixed model hαi = αiĥ∗loc , i+ ( 1−αi ) h̄∗ on the ith local distribution Di is bounded by : LDi ( hαi ) ≤ 2α2i LDi ( h∗i ) + 2C √ d+ log ( 1/δ ) mi +GλH ( Si ) + 2 ( 1− αi ) 2 ( L̂D̄ ( h̄∗ ) +B‖D̄ − Di‖1 + C √ d+ log ( 1/δ ) m ) , ( 2 ) wheremi , i = 1 , 2 , . . . , n is the number of training data at ith user , m = m1 + . . .+mn is the total number of all data , Si to be the local training set drawn from Di , ‖D̄ − Di‖1 = ∫ Ξ |P ( x , y ) ∼D̄ − P ( x , y ) ∼Di |dxdy , is the difference between distributions D̄ and Di , and h∗i = arg minh∈H LDi ( h ) . Remark 1 . We note that a very analogous work to ours is Mansour et al . ( 2020 ) , where a generalization bound is provided for mixing global and local models . However , their bound does not depend on αi , and hence we can not see how it impacts the generalization ability . In Theorem 1 , by omitting constant terms , we observe that the generalization risk of hαi on Di mainly depends on three key quantities : i ) m : the number of global data drawn from D̄ , ii ) divergence between distributions D̄ andDi , and iii ) mi : the amount of local data drawn fromDi . Usually , the first quantity m , the amount of global data is fairly large compared to individual users , so global model usually has a better generalization . The second quantity characterizes the data heterogeneity between the average distribution and ith local distribution . If this divergence is too high , then the global model may hurt the local generalization . For the third quantity , as amount of local data mi is often small , the generalization performance of local model can be poor . Optimal mixing parameter . We can also find the optimal mixing parameter α∗i that minimizes generalization bound in Theorem 1 . Notice that the RHS of ( 2 ) is quadratic in αi , so it admits a minimum value at α∗i = ( L̂D̄ ( h̄∗ ) +B‖D̄ − Di‖1 + C √ d+log ( 1/δ ) m ) ( L̂D̄ ( h̄∗ ) +B‖D̄ − Di‖1 + C √ d+log ( 1/δ ) m ) + ( LDi ( h∗i ) + 2C √ d+log ( 1/δ ) mi +GλH ( Si ) ) . The optimal mixture parameter is strictly bounded in [ 0 , 1 ] , which matches our intuition . If the divergence term is large , then the value becomes close to 1 , which implies if local distribution drifts too much from average distribution , it is preferable to take more local models . If mi is small , this value will be negligible , indicating that we need to mix more of the global model into the personalized model . Conversely , if mi is large , then this term will be again roughly 1 , which means taking the majority of local model will give the desired generalization performance .

Summary: This paper studies the personalization aspect of the federated learning problem. The authors propose a new framework in which they replace the common global model in the original federated learning formulation with a convex combination of the global model and a local model. They later introduce an adaptive optimization algorithm for their formulation and provide generalization bounds as well as convergence guarantees for both strongly-convex and nonconvex settings.  Finally, numerical experiments and comparison with other methods in the literature are provided to support the theoretical results.