Few-Round Learning for Federated Learning

1 Introduction . Today , valuable data are being collected increasingly at distributed edge nodes such as mobile phones , wearable client devices and smart vehicles/drones . Directly sending these local data to the central server for model training raises significant privacy concerns . To address this issue , an emerging trend known as federated learning ( FL ) [ 13 , 9 , 1 , 11 , 20 , 16 , 15 ] , where server uploading of local data is not necessary , has been actively researched . In FL , a large group of distributed clients interested in solving the same task ( e.g. , classification on given categories of images ) collaborate in training a single global model without sharing their data . While standard supervised learning uses some dataset D to find the model φ that would minimize a loss function f ( φ , D ) , FL in comparison seeks the model φ that minimizes the averaged version of the local losses f ( φ , Dk ) , computed at each node k using local data Dk . The learning process typically starts from a randomly initialized or some pretrained model and is carried out through iterative aggregation of the local model updates . 1.1 Backgrounds and Main Contributions . Motivation . Unfortunately , FL generally requires a large number of communication rounds between the server and the clients for model exchange , to achieve a desired level of performance . This makes ∗Equal contribution . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . the implementation of FL a significant challenge in bandwidth-limited or time-sensitive applications . Especially in real-time applications ( e.g. , connected vehicles or drones ) where the model should quickly adapt to dynamically evolving environments , the requirement on many communication rounds becomes a major bottleneck . Goal and challenge . To tackle this problem from the service provider ’ s perspective , we aim to prepare an initial model that can quickly adapt to any group ( focusing on its own task ) of clients within only a few rounds of FL . The key challenge here is that the task of the group conducting FL ( i.e. , the downstream task for which the prepared model will be used ) is generally not known when the service provider prepares the initial model . In the context of classification , different tasks mean classification involving different sets of classes . For example , classifying diseases A , B , C ( task 1 ) is a different task compared to classifying diseases D , E , F ( task 2 ) . Since the group conducting FL for the downstream task can include classes that are unseen during preparation , existing FL approaches can not tackle this problem . Key idea . Our key idea is to adopt meta-learning ( which enables reliable prediction even when the task at inference is unseen when the model was meta-trained ) to prepare the initial model that enables few-round FL . In other words , we aim to meta-train an initial model for few-round downstream FL . Once meta-training is over , the service provider would offer the trained model to some clients who want to solve a common task after collaborating through a quick few rounds of FL . These clients may or may not be the participants of the earlier meta-training phase , and their classification task is generally considered unseen during meta-training . A high-level description of our idea is depicted in Fig . 1 ( b ) . Given a small target value R , we take an episodic training approach to enable R-round FL for any group of clients . In essence , we find the initial model φ that would minimize the average of local losses f ( θR ( φ ) , Dk ) , where θR ( φ ) is the model to be updated from φ through R rounds of FL among future clients in the deployment stage . Despite the high practical significance of this problem formulation , to the best of our knowledge , this is the first work to propose a meta-learning strategy geared to few-round FL . It is also worth mentioning that model preparation is not a real-time requirement and can often be done when bandwidth demands are sparse . Comparison with personalized FL . We stress that our idea has a different purpose and approach relative to the recent line of works on federated meta-learning [ 12 , 4 ] , which initiate a model for personalized optimizations at local clients ( see Fig . 1 ( a ) ) . The goal of these approaches is to obtain a personalized local model at each client within a few steps of gradient descents , in the deployment stage . To achieve this goal , in the preparation stage , a few steps of local updates and meta-update are first performed at each participant independently ( with its own local data ) , and FL ( or aggregation ) is adopted just to take advantage of data of various participants : these approaches seek φ that minimizes the average of local losses f ( θk ( φ ) , Dk ) , where θk ( φ ) is the local model updated from φ through a number of gradient steps using local data Dk . In contrast to personalized FL that focuses on local client models in the deployment stage , our few-round learning inherits the ability of FL at deployment to obtain a global model . Hence , for our scheme , it is natural to adopt FL in the preparation stage to mimic the R-round FL scenario at deployment ; in the preparation stage , meta-update is performed at each participant after the collaborative R FL rounds . To sum , our approach aims to prepare an initial model that leads to a global model within a “ few rounds of FL ” , while personalized FL aims for an initial model leading to personalized models within “ few steps of local updates ” based only on the local data . These are obviously two completely different problems with distinct solutions . Main contributions . Technically , we utilize a model-agnostic meta-learning ( MAML ) approach to prepare the initial model via an episodic training strategy . While directly applying MAML independently to each local model leads to existing solutions on personalized FL [ 12 , 4 ] , in our approach , R rounds of local updates and aggregations are first performed in each episode before the meta-update process . This unique episode construction compared to personalized FL methods mimics the deployment stage where actual inference is preceded by an R-round FL procedure . Another key ingredient in our solution is to adopt prototype aggregation in each FL round to construct global prototypes that serve as better class representatives compared to the locally computed prototypes , in learning embedding space . This strategy is especially effective when a non-IID ( independent , identically distributed ) data distribution across clients tends to induce a significantly biased model after performing local updates . The global prototypes serve as prior knowledge , a form of regularization , and prevent local models from overfitting to the local data . Moreover , the global prototypes ( reflecting all classes across clients ) can assist the local models to learn a more general embedding space . We call this approach a global prototype-assisted learning ( GPAL ) strategy . Our main contributions are summarized as follows : • We formulate a new problem of high practical significance , namely , few-round learning , where the goal is to prepare an initial model that can quickly adapt to any group of clients within only a few rounds of FL . • We propose a meta-training algorithm specifically geared to R rounds of FL followed by inference , to be performed by a group of clients on a possibly unseen task . • We guarantee convergence of our meta-training algorithm via theoretical analysis . • We show via experiments that our scheme outperforms existing pretraining approaches including fine-tuning via FedAvg and personalized FL in both IID and non-IID scenarios . 1.2 Related Works . Few-shot learning . Few-shot learning is an instantiation of meta-learning . In the context of image classification , few-shot learning typically involves episodic training where each episode of training data is arranged into a few training ( support ) sample images and validation ( query ) samples to mimic inference that uses only a few examples [ 19 ] . Through a repetitive exposure to a series of varying episodes with different sets of image classes , the model learns to handle new tasks ( classification against unseen classes ) each time . Two widely-known few-shot learning methods with different philosophical twists , which are also conceptually relevant to the present work , are MAML [ 5 ] and Prototypical Networks [ 18 ] . MAML attempts to generate an initial model from which different models targeting different tasks can be obtained quickly via just a few gradient updates . The idea is that the initial model is learned via meta-training to develop an internal representation that is close in some sense to a variety of unseen tasks . Prototypical Networks , on the other hand , learn embedding space such that model outputs cluster around class prototypes , the class-specific centroids of the embedder outputs . With episodic training , simple Prototypical Networks are surprisingly effective in learning inductive bias for successful generalization to new tasks . We stress that our few-round learning scheme ( that targets a few global rounds of FL ) has different purpose and technical approach compared to the existing works on few-shot learning ( that targets few shots of data sample ) . Nevertheless , we take advantage of both concepts on MAML and Prototypical Networks to achieve our own goal : we adopt MAML in updating the initial model specifically geared to R-round FL , and adopt both prototype aggregation and prototype-assisted learning strategies to learn a general embedding space and successfully handle the non-IID issue in FL . Federated meta-learning . Recent research activity has focused on improving model personalization via federated meta-learning [ 12 , 3 , 4 , 7 ] . The common goal of these works is to generate an initial model based on which each new client can find its own optimized model via a few local gradient steps and using only its own data . In these works , meta-learning employed during federated learning intends to enable each client to handle previously unseen tasks , in the spirit of MAML of [ 5 ] . User-specific next-word prediction at individual smartphones , for example , is a possible application . Compared to this line of work , we focus on creating an initial model that leads to a high-accuracy global model , rather than personalized models . In this way , we seek to take advantage of a higher variety of data as well as the larger data volume that would be made available through collaborative learning of a group of distributed nodes . A clear example is the diagnosis of a broader class of diseases that would be possible through collaborative training across more examples contributed by a larger group of individuals . Personalized FL methods ( e.g. , [ 12 , 4 ] ) especially have disadvantage in non-IID settings where each client necessarily lacks a sufficient variety of data . The results are reported in Section 4 . One-shot FL . Another line of work recently focused on one-shot FL , where the goal is to train a global model with just one communication round between the server and the clients . The authors of [ 6 ] proposed an ensemble method to choose reliable client-specific models from given clients . In the work of [ 17 ] , local clients send XOR-encoded MNIST image data to the server , and the server decodes it to train the global model . While the server would need certain data in advance to decode the received results , XOR operation can serve as data augmentation while preserving privacy . In the fusion learning of [ 8 ] , each local client uploads both the model parameters and the distribution parameters to the server . The server generates artificial data samples from the distribution parameters to train a global model . When the data gets complex , however , it is not clear whether conversion into a simple distribution would be reliable . Compared to the existing works on one-shot FL that employ some randomly initialized model , the key difference of our method is the use of meta-learning to prepare an initial model which can adapt to unseen tasks of individual groups ’ of clients within R rounds of FL . The advantage of our scheme compared to these methods is shown in Section 4 .

Summary: This paper presents a method for using meta-learning to train an initial model, which is intended to achieve optimal performance on an arbitrary downstream tasks when trained for R rounds of federated learning (for some small R). The paper presents aa meta-learning-flavored algorithm for conducting the pretraining step, and provides theoretical bounds on the gradient of the resulting loss (Theorem 1). The paper includes some empirical results on CIFAR100, miniImageNet, and FEMNIST, with comparison to some other baselines (although no baselines exist from other works for this particular task, according to the authors).