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: The authors propose a new approach to partially local federated learning. In contrast to previous works their approach does not require agents to maintain state (i.e. store the local parameters of the model). Instead, they reconstruct the personalized parts of the model on the fly from local data.  Such a state-less approach is crucial for any practical deployment of FL at scale since is enables new clients that have not yet trained their local parameters to participate in training and to perform inference. Their approach (FedRecon) is inspired by model-agnostic meta learning, where fast reconstruction at test-time is a key design criteria.   Empirically the performance of FedRecon is evaluated along different axis. The authors show that their approach trains global parameter that are particularly well suited for reconstruction, the approach achieves performance comparable to personalization through fine-tuning of a global model, and they study the qualitative behavior of their method wrt several hyperparameters.  In addition the authors report on a successful real-world deployment of the method across hundreds of millions of clients and discuss the practical challenge of data sparsity.   The code for reconstructing the experiments is packaged as a library to allow users to seamlessly integrate other tasks and algorithms.