Personalized Federated Learning with First Order Model Optimization

1 INTRODUCTION . Federated learning ( FL ) has shown great promise in recent years for training a single global model over decentralized data . While seminally motivated by effective inference on a general test set similar in distribution to the decentralized data in aggregate ( McMahan et al. , 2016 ; Bonawitz et al. , 2019 ) , here we focus on federated learning from a client-centric or personalized perspective . We aim to enable stronger performance on personalized target distributions for each participating client . Such settings can be motivated by cross-silo FL , where clients are autonomous data vendors ( e.g . hospitals managing patient data , or corporations carrying customer information ) that wish to collaborate without sharing private data ( Kairouz et al. , 2019 ) . Instead of merely being a source of data and model training for the global server , clients can then take on a more active role : their federated participation may be contingent on satisfying client-specific target tasks and distributions . A strong FL framework in practice would then flexibly accommodate these objectives , allowing clients to optimize for arbitrary distributions simultaneously in a single federation . In this setting , FL ’ s realistic lack of an independent and identically distributed ( IID ) data assumption across clients may be both a burden and a blessing . Learning a single global model across non-IID data batches can pose challenges such as non-guaranteed convergence and model parameter divergence ( Hsieh et al. , 2019 ; Zhao et al. , 2018 ; Li et al. , 2020 ) . Furthermore , trying to fine-tune these global models may result in poor adaptation to local client test sets ( Jiang et al. , 2019 ) . However , the non-IID nature of each client ’ s local data can also provide useful signal for distinguishing their underlying local data distributions , without sharing any data . We leverage this signal to propose a new framework for personalized FL . Instead of giving all clients the same global model average weighted by local training size as in prior work ( McMahan et al. , 2016 ) , for each client we compute ∗Corresponding author ; work done while interning at NVIDIA a weighted combination of the available models to best align with that client ’ s interests , modeled by evaluation on a personalized target test distribution . Key here is that after each federating round , we maintain the client-uploaded parameters individually , allowing clients in the next round to download these copies independently of each other . Each federated update is then a two-step process : given a local objective , clients ( 1 ) evaluate how well their received models perform on their target task and ( 2 ) use these respective performances to weight each model ’ s parameters in a personalized update . We show that this intuitive process can be thought of as a particularly coarse version of popular iterative optimization algorithms such as SGD , where instead of directly accessing other clients ’ data points and iteratively training our model with the granularity of gradient decent , we limit ourselves to working with their uploaded models . We hence propose an efficient method to calculate these optimal combinations for each client , calling it FedFomo , as ( 1 ) each client ’ s federated update is calculated with a simple first-order model optimization approximating a personalized gradient step , and ( 2 ) it draws inspiration from the “ fear of missing out ” , every client no longer necessarily factoring in contributions from all active clients during each federation round . In other words , curiosity can kill the cat . Each model ’ s personalized performance can be saved however by restricting unhelpful models from each federated update . We evaluate our method on federated image classification and show that it outperforms other methods in various non-IID scenarios . Furthermore , we show that because we compute federated updates directly with respect to client-specified local objectives , our framework can also optimize for outof-distribution performance , where client ’ s target distributions are different from their local training ones . In contrast , prior work that personalized based on similarity to a client ’ s own model parameters ( Mansour et al. , 2020 ; Sattler et al. , 2020 ) restricts this optimization to the local data distribution . We thus enable new features in personalized FL , and empirically demonstrate up to 70 % improvement in some settings , with larger gains as the number of clients or level of non-IIDness increases . Our contributions . 1 . We propose a flexible federated learning framework that allows clients to personalize to specific target data distributions irrespective of their available local training data . 2 . Within this framework , we introduce a method to efficiently calculate the optimal weighted combination of uploaded models as a personalized federated update 3 . Our method strongly outperforms other methods in non-IID federated learning settings . 2 RELATED WORK . Federated Learning with Non-IID Data While fine-tuning a global model on a client ’ s local data is a natural strategy to personalize ( Mansour et al. , 2020 ; Wang et al. , 2019 ) , prior work has shown that non-IID decentralized data can introduce challenges such as parameter divergence ( Zhao et al. , 2018 ) , data distribution biases ( Hsieh et al. , 2019 ) , and unguaranteed convergence Li et al . ( 2020 ) . Several recent methods then try to improve the robustness of global models under heavily non-IID datasets . FedProx ( Li et al. , 2020 ) adds a proximal term to the local training objective to keep updated parameter close to the original downloaded model . This serves to reduce potential weight divergence defined in Zhao et al . ( 2018 ) , who instead allow clients to share small subsets of their data among each other . This effectively makes each client ’ s local training set closer in distribution to the global test set . More recently , Hsu et al . ( 2019 ) propose to add momentum to the global model update in FedAvgM to reduce the possibly harmful oscillations associated with averaging local models after several rounds of stochastic gradient descent for non-identically distributed data . While these advances may make a global model more robust across non-IID local data , they do not directly address local-level data distribution performance relevant to individual clients . Jiang et al . ( 2019 ) argue this latter task may be more important in non-IID FL settings , as local training data differences may suggest that only a subset of all potential features are relevant to each client . Their target distributions may be fairly different from the global aggregate in highly personalized scenarios , with the resulting dataset shift difficult to handle with a single model . Personalized Federated Learning Given the challenges above , other approaches train multiple models or personalizing components to tackle multiple target distributions . Smith et al . ( 2017 ) propose multi-task learning for FL with MOCHA , a distributed MTL framework that frames clients as tasks and learns one model per client . Mixture methods ( Deng et al. , 2020 ; Hanzely & Richtárik , 2020 ; Mansour et al. , 2020 ) compute personalized combinations of model parameters from training both local models and the global model , while Peterson et al . ( 2019 ) ensure that this is done with local privacy guarantees . Liang et al . ( 2020 ) apply this mixing across network layers , with lower layers acting as local encoders that map a client ’ s observed data to input for a globally shared classifier . Rather than only mix with a shared global model , our work allows for greater control and distinct mixing parameters with multiple local models . Fallah et al . ( 2020 ) instead optimize the global model for fast personalization through meta-learning , while T Dinh et al . ( 2020 ) train global and local models under regularization with Moreau envelopes . Alternatively , Clustered FL ( Sattler et al. , 2020 ; Ghosh et al. , 2020 ; Briggs et al. , 2020 ; Mansour et al. , 2020 ) assumes that inherent partitions or data distributions exist behind clients ’ local data , and aim to cluster these partitions to federate within each cluster . Our work does not restrict which models are computed together , allowing clients to download suitable models independently . We also compute client-specific weighted averages for greater personalization . Finally , unlike prior work , we allow clients to receive personalized updates for target distributions different from their local training data . 3 FEDERATED FIRST ORDER MODEL OPTIMIZATION . We now present FedFomo , a personalized FL framework to efficiently compute client-optimizing federated updates . We adopt the general structure of most FL methods , where we iteratively cycle between downloading model parameters from server to client , training the models locally on each client ’ s data , and sending back the updated models for future rounds . However , as we do not compute a single global model , each federated download introduces two new steps : ( 1 ) figuring out which models to send to which clients , and ( 2 ) computing their personalized weighted combinations . We define our problem and describe how we accomplish ( 1 ) and ( 2 ) in the following sections . Problem Definition and Notation Our work most naturally applies to heterogeneous federated settings where participating clients are critically not restricted to single local training or target test distribution , and apriori we do not know anything about these distributions . To model this , let C be a population with |C| = K total clients , where each client ci ∈ C carries local data Di sampled from some distribution D and local model parameters θ ` ( t ) i during any round t. Each ci also maintains some personalized objective or task Ti motivating their participation in the federation . We focus on supervised classification as a universal task setting . Each client and task are then associated with a test dataset Dtesti ∼ D∗ . We define each Ti : = minL ( θ ` ( t ) i ; D test i ) , where L ( θ ; D ) : Θ 7→ R is the loss function associated with dataset D , and Θ denotes the space of models possible with our presumed network architecture . We assume no knowledge regarding clients and their data distributions , nor that test and local data belong to the same distribution . We aim to obtain the optimal set of model parameters { θ∗1 , . . . , θ∗K } = arg min ∑ i∈ [ K ] LTi ( θi ) . 3.1 COMPUTING FEDERATED UPDATES WITH FOMO . Unlike previous work in federated learning , FedFomo learns optimal combinations of the available server models for each participating client . To do so , we leverage information from clients in two different ways . First , we aim to directly optimize for each client ’ s target objective . We assume that clients can distinguish between good and bad models on their target tasks , through the use of a labeled validation data split Dvali ⊂ Di in the client ’ s local data . Dvali should be similar in distribution to the target test dataset Dtesti . The client can then evaluate any arbitrary model θj on this validation set , and quantify the performance through the computed loss , denoted by Li ( θj ) . Second , we directly leverage the potential heterogeneity among client models . Zhao et al . ( 2018 ) explore this phenomenon as a failure mode for traditional single model FL , where they show that diverging model weights come directly from local data heterogeneity . However , instead of combining these parameters into a single global model , we maintain the uploaded models individually as a means to preserve a model ’ s potential contribution to another client . Critically , these two ideas together not only allow us to compute more personal model updates within non-IID local data distributions , but also enable clients to optimize for data distributions different from their own local data ’ s . Federated learning as an iterative local model update The central premise of our work stems from viewing each federated model download−and subsequent changing of local model parameters−as an optimization step towards some objective . In traditional FL , this objective involves performing well on the global population distribution , similar in representation to the union of all local datasets . Assuming N federating clients , we compute each global model θG at time t as : θG ( t ) = ∑N n=1 wn · θ ` ( t ) n , where wn = |Dtrainn |/ ∑N j=1 |Dtrainj | . If client ci downloads this model , we can view this change to their local model as an update : θ ` ( t+1 ) i ← θ ` ( t ) i + ∑N n=1 wn· ( θ ` ( t ) n −θ ` ( t ) i ) since ∑ n wn = 1 . This then updates a client ’ s current local model parameters in directions specified by the weights w and models { θn } in the federation . A natural choice to optimize for the global target distribution sets wn as above and in McMahan et al . ( 2017 ) , e.g . as an unbiased estimate of global model parameters . However , in our personalized scenario , we are more interested in computing the update uniquely with respect to each client ’ s target task . We then wish to find the optimal weights w = 〈w1 , . . . , wN 〉 that optimize for the client ’ s objective , minimizing Li ( θ ` i ) . Efficient personalization with FedFomo Intuitively , we wish to find models { θ ` ( t ) m : m ∈ [ N ] \i } such that moving towards their parameters leads to better performance on our target distribution , and accordingly weight these θ higher in a model average . If a client carries a satisfying number of local data points associated with their target objective Li , then they could obtain a reasonable model through local training alone , e.g . directly updating their model parameters through SGD : θ ` ( t+1 ) i ← θ ` ( t ) i − α∇θLi ( θ ` ( t ) i ) ( 1 ) However , without this data , clients are more motivated to federate . In doing so they obtain useful updates , albeit in the more restricted form of fixed model parameters { θn : n ∈ N } . Then for personalized or non-IID target distributions , we can iteratively solve for the optimal combination of client models w∗ = arg minLi ( θ ) by computing : θ ` ( t+1 ) i ← θ ` ( t ) i − α1 > ∇wLi ( θ ` ( t ) i ) ( 2 ) where 1 is a size-N vector of ones . Unfortunately , as the larger federated learning algorithm is already an iterative process with many rounds of communication , computing w∗ through Eq . 2 may be cumbersome . Worse , if the model averages are only computed server-side as in traditional FL , Eq . 2 becomes prohibitively expensive in communication rounds ( McMahan et al. , 2017 ) . Following this line of reasoning however , we thus derive an approximation of w∗ for any client : Given previous local model parameters θ ` ( t−1 ) i , set of fellow federating models available to download { θ ` ( t ) n } and local client objective captured by Li , we propose weights of the form : wn = Li ( θ ` ( t−1 ) i ) − Li ( θ ` ( t ) n ) ‖θ ` ( t ) n − θ ` ( t−1 ) i ‖ ( 3 ) where the resulting federated update θ ` ( t ) i ← θ ` ( t−1 ) i + ∑ n∈ [ N ] wn ( θ ` ( t ) n −θ ` ( t−1 ) i ) directly optimizes for client ci ’ s objective up to a first-order approximation of the optimal w∗ . We default to the original parameters θ ` ( t−1 ) i if wn < 0 above , i.e . wn = max ( wn , 0 ) , and among positive wn normalize to get final weights w∗n = max ( wn,0 ) ∑ n max ( wn,0 ) to maintain w∗ ∈ [ 0 , 1 ] and ∑ n=1 w ∗ n ∈ { 0 , 1 } . We derive Eq . 3 as a first-order approximation of w∗ in Appendix A.1 . Here we note that our formulation captures the intuition of federating with client models that perform better than our own model , e.g . have a smaller loss on Li . Moreso , we weigh models more heavily as this positive loss delta increases , or the distance between our current parameters and theirs decreases , in essence most heavily weighing the models that most efficiently improve our performance . We use local parameters at t-1 to directly compute how much we should factor in current parameters θ ` ( t ) i , which also helps prevent overfitting as Li ( θ ` ( t−1 ) i ) − Li ( θ ` ( t ) i ) < 0 causes “ early-stopping ” at θ ` ( t−1 ) i . Communication and bandwidth overhead Because the server can send multiple requested models in one download to any client , we still maintain one round of communication for model downloads and one round for uploads in between E local training epochs . Furthermore , because w in Eq . 3 is simple to calculate , the actual model update can also happen client-side , keeping the total number of communications with T total training epochs at b 2TE c , as in FedAvg . However FedFomo also needs to consider the additional bandwidth from downloading multiple models . While quantization and distillation ( Chen et al. , 2017 ; Hinton et al. , 2015 ; Xu et al. , 2018 ) can alleviate this , we also avoid worst case N2 overhead with respect to the number of active clients N by restricting the number of models downloaded M . Whether we can achieve good personalization here involves figuring out which models benefit which clients , and our goal is then to send as many helpful models as possible given limited bandwidth . To do so , we invoke a sampling scheme where the likelihood of sending model θj to client ci relies on how well θj performed regarding client ci ’ s target objective in previous rounds . Accordingly , we maintain an affinity matrix P composed of vectors pi = 〈pi,1 , . . . , pi , K〉 , where pi , j measures the likelihood of sending θj to client ci , and at each round send the available uploaded models corresponding to the top M values according to each participating client ’ s p. Initially we set P = diag ( 1 , . . . , 1 ) , i.e . each model has an equal chance of being downloaded . Then during each federated update , we update p← p + w from Eq . 3 , where w can now be negative . If N K , we may benefit from additional exploration , and employ an ε-greedy sampling strategy where instead of picking strictly in order of p , we have ε chance to send a random model to the client . We investigate the robustness of FedFomo to these parameters through ablations of ε and M in the next section .

Summary: The paper proposes a new FL method that computes in every communication round for each client a personalized model as starting point for the next round of federation. The paper defines the client-specific objective as some loss function of the weighted combination of all (or subset) models on a client-specific validation set. This personalized weighted combination of the models especially fits situations where not all clients have congruent objectives such as in non-IID settings. The paper evaluates the proposed FL algorithm on standard datasets for image classification by comparing to alternative FL methods.