Federated Reconstruction: Partially Local Federated Learning

1 Introduction . Federated learning is a machine learning setting in which distributed clients solve a learning objective on sensitive data via communication with a coordinating server [ 44 ] . Typically , clients collaborate to train a single global model under an objective that combines heterogeneous local client objectives . For example , clients may collaborate to train a next word prediction model for a mobile keyboard application without sharing sensitive typing data with other clients or a centralized server [ 28 ] . This paradigm has been scaled to production and deployed in cross-device settings [ 3 , 28 , 56 ] and cross-silo settings [ 11 , 13 ] . However , training a fully global federated model may not always be ideal due to heterogeneity in clients ’ data distributions . Yu et al . [ 58 ] show that global models can perform worse than purely local ( non-federated ) models for many clients ( e.g. , those with many training examples ) . Moreover , in some settings privacy constraints completely prohibit fully global federated training . For instance , for models with user-specific embeddings , such as matrix factorization models for collaborative filtering [ 37 ] , naively training a global federated model involves sending updates to user embeddings on the server , directly revealing potentially sensitive individual preferences [ 21 , 47 ] . To address this , we explore partially local federated learning . In this setting , models are partitioned into global g and local parameters l such that local parameters never leave client devices . This enables training on sensitive user-specific parameters as in the collaborative filtering setting , and we show it can also improve robustness to client data heterogeneity and communication cost for other settings , 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . since we are effectively interpolating between local and federated training . Previous works have looked at similar settings [ 4 , 41 ] . Importantly , these approaches can not realistically be applied at scale in cross-device settings because they assume clients are stateful or always-available : in practice , clients are sampled from an enormous population with unreliable availability , so approaches that rely on repeated sampling of the same stateful clients are impractical ( Kairouz et al . [ 34 ] [ Table 1 ] ) . Other work has demonstrated that stateful federated algorithms in partial participation regimes can perform worse than stateless algorithms due to the state becoming `` stale '' [ 48 ] . Previous methods also do not enable inference on new clients unseen during training , preventing real-world deployment . These limitations motivate a new method for partially local federated learning , balancing the benefits of federated aggregation and local training . This approach should be : 1 . Model-agnostic : works with any model . 2 . Scalable : compatible with large-scale cross-device training with partial participation . 3 . Practical for inference : new clients can perform inference . 4 . Fast : clients can quickly adapt local parameters to their personal data . In this work , we propose combining federated training of global parameters with reconstruction of local parameters ( see Figure 1 ) . We show that our method relaxes the statefulness requirement of previous work and enables fast personalization for unseen clients without additional communication , even for models without user-specific embeddings . Our contributions : We make the following key contributions : • Introduce a model-agnostic framework for training partially local and partially global models , satisfying the above criteria . We propose a practical algorithm instantiating this framework ( FEDRECON ) . • Justify the algorithm via a connection to model-agnostic meta learning ( see Section 4.2 ) , showing that FEDRECON naturally leads to fast reconstruction at test time ( see Table 1 ) . • Demonstrate FEDRECON ’ s empirical performance over existing approaches for applications in collaborative filtering and next word prediction , showing that our method outperforms standard centralized and federated training in performance on unseen clients ( see Table 1 ) , enables fast adaptation to clients ’ personal data ( see Figure 3 ) , and matches the performance of other federated personalization techniques with less communication ( see Figure 2 ) . • Release an open-source library for evaluating algorithms across tasks in this setting.1 • Describe the successful deployment of this approach at scale for collaborative filtering in a real-world mobile keyboard application ( see Section 7 ) . 2 Related Work . Previous works have explored personalization of federated models via finetuning [ 52 , 58 ] , meta learning / bi-level optimization [ 10 , 16 , 18 , 33 ] , and model interpolation [ 14 , 27 , 43 ] . Some works aim to improve training convergence with heterogeneous client gradient updates [ 35 , 39 ] , while 1https : //git.io/federated_reconstruction others address client resource heterogeneity [ 15 , 49 ] . All of these approaches require communicating all client parameters during training , which can be unreasonable due to privacy and communication constraints for some models ( discussed further in Section 3 ) , which motivates methods that aggregate only part of a model as in our work . Arivazhagan et al . [ 4 ] and Liang et al . [ 41 ] aggregate part of a model , but these approaches do not meet the criteria from Section 1 . Similar to other works proposing local parameters [ 22 , 31 , 40 ] , both approaches require clients to maintain local models across rounds , which is problematic when sampling clients from large populations ( criterion 2 ) . Arivazhagan et al . [ 4 ] assumes that all clients are available for training at all times and do not propose a method for performing inference on new clients ( criterion 3 ) . Liang et al . [ 41 ] requires new inference clients to be able to ensemble the outputs of all other clients ’ local models to evaluate on new data , which is unrealistic in practice due to communication and privacy constraints ( criterion 3 ) . These constraints are crucial : with previous methods most clients do not have a practical way to perform inference . Previous methods were also proposed for specific model types ( criterion 1 ) : Arivazhagan et al . [ 4 ] explores personalization layers after shared base layers and Liang et al . [ 41 ] learns personal representations of local data . Finally , as we discuss in Section 4.2 , our method optimizes a meta learning objective for training global parameters that lead to fast reconstruction ( criterion 4 ) . Federated Collaborative Filtering : We evaluate our approach on collaborative filtering [ 37 ] in Section 5.1.1 . Prior work has explored federated matrix factorization : Ammad-Ud-Din et al . [ 2 ] avoids sending the user matrix to the server by storing it locally , aggregating only the item matrix globally . Chai et al . [ 9 ] applies homomorphic encryption to aggregation of the item matrix . Flanagan et al . [ 20 ] studies federated collaborative filtering as a multi-view learning problem . Each approach requires clients to maintain state , unlike our method . Ammad-Ud-Din et al . [ 2 ] and Chai et al . [ 9 ] also do not address the problem of inference on unseen users . Federated Meta Learning : Our approach is motivated by a connection to meta learning , described in Section 4.2 . Other federated learning works have also established connections to meta learning : Jiang et al . [ 33 ] observed that training a global federated model that can be easily personalized via finetuning can be studied in the model-agnostic meta learning ( MAML ) framework [ 19 ] , and FEDAVG is performing the distributed version of the REPTILE meta learning algorithm presented by Nichol et al . [ 46 ] . Chen et al . [ 10 ] , Fallah et al . [ 18 ] , and Lin et al . [ 42 ] apply the MAML algorithm and variants in federated settings . Khodak et al . [ 36 ] aims to improve upon these methods by learning client similarities adaptively . These methods do not address the partially local federated learning setting , where some parameters are not aggregated globally . 3 Partially Local Federated Learning . Typically , federated learning of a global model optimizes : min x2Rd F ( x ) = Ei⇠P [ fi ( x ) ] ( 1 ) where fi ( x ) = E⇠2Di [ fi ( x , ⇠ ) ] is the local objective for client i , x is the d-dimensional model parameter vector , P is the distribution of clients , and ⇠ is a data sample drawn from client i ’ s data Di . In practical cross-device settings , fi ( x ) may be highly heterogeneous for different i , and the number of available clients may be large and constantly changing due to partial availability . Only a relatively small fraction of clients may be sampled for training . To motivate partially local federated learning , we begin by considering models that can be partitioned into user-specific parameters and non-user-specific parameters . An example is matrix factorization in the collaborative filtering setting [ 30 , 37 ] : in this scenario , a ratings matrix R 2 RU⇥I representing user preferences is factorized into a user matrix P 2 RU⇥K and an items matrix Q 2 RI⇥K such that R ⇡ PQ > , where U is the number of users and I is the number of items . For each user u , this approach yields a K-dimensional user-specific embedding Pu . To train this type of model in the federated setting , we can not naively use the popular FEDAVG algorithm [ 44 ] or other ( personalized ) algorithms that involving aggregation of all model parameters . A simple application of global learning algorithms might require every client to be sent every other client ’ s personal parameters , which is clearly unreasonable for both privacy and communication . A more sophisticated approach might be to have each client communicate only their own personal parameters with the server . In this case , the server still has access to individual user parameters , Algorithm 1 Federated Reconstruction Training Input : set of global parameters G , set of local parameters L , dataset split function S , reconstruction algorithm R , client update algorithm U Server executes : g ( 0 ) ( initialize G ) for each round t do S ( t ) ( randomly sample m clients ) for each client i 2 S ( t ) in parallel do ( ( t ) i , ni ) ClientUpdate ( i , g ( t ) ) end for n = P i2S ( t ) ni g ( t+1 ) g ( t ) + ⌘s P i2S ( t ) ni n ( t ) i end for ClientUpdate : ( Di , s , Di , q ) S ( Di ) l ( t ) i R ( Di , s , L , g ( t ) ) g ( t ) i U ( Di , q , l ( t ) i , g ( t ) ) ( t ) i g ( t ) i g ( t ) ni |Di , q| return ( t ) i , ni to the server which in this setting can be trivially used to recover sensitive user-item affinities , negating the privacy benefit of not centralizing the data ( again unreasonable ) . Thus a practical federated learning algorithm for this setting should be partially local : it should enable clients to train a subset of parameters entirely on-device . However , approaches that involve stateful clients storing their local parameters across rounds are undesirable in large-scale cross-device settings since clients are unlikely to be sampled repeatedly , causing state to be infrequently available and become stale , degrading performance ( Reddi et al . [ 48 ] [ Sec . 5.1 ] ) . Additionally , since only a fraction of clients participate in training , all other clients will be left without trained local parameters , preventing them from performing inference using the model . In a large population setting with hundreds of millions of clients as described in Section 7 , this can mean 99 % + of clients do not have a complete model , preventing practical deployment . Thus an algorithm for this setting ideally should not depend on stateful clients and should provide a way to perform inference on unseen clients . Though we have motivated partially local federated learning via a setting that contains privacysensitive user-specific parameters , we will later show that this paradigm can also improve robustness to heterogeneity in fi ( x ) and reduce communication cost , even for models without user-specific parameters . In this case , the partition between local and global parameters is determined by the use-case and communication limitations . As an example , in Section 5.1.2 we motivate a next word prediction use-case , where having a partially local model can be useful for handling diverse client inputs while reducing communication . Achieving partially local federated learning in a practical cross-device setting with large , changing client distribution P and stateless clients is one of the key contributions of our work .

Summary: This paper introduces a partially local federated learning method named FEDRECON based on federated reconstruction.  The experimental results demonstrate the effectiveness of FEDRECON. The authors also deploy this approach in a real-world mobile keyboard application.