Federated learning using mixture of experts

1 INTRODUCTION In many real-world scenarios , data is distributed over a large number of devices , due to privacy concerns or communication limitations . Federated learning is a framework that can leverage this data in a distributed learning setup . This allows for exploiting both the compute power of all participating clients , and to benefit from a large joint training data set . Furthermore , this is beneficial for privacy and data security . For example , in keyboard prediction for smartphones , thousands or even millions of users produce keyboard input that can be leveraged as training data . The training can ensue directly on the devices , doing away with the need for costly data transfer , storage , and immense compute on a central server ( Hard et al. , 2018 ) . The medical field is another example area where data is extremely sensitive and may have to stay on premise , and a setting where analysis may require distributed and privacy-protecting approaches . In settings with such firm privacy requirements , standard federated learning approaches may not be enough to guarantee the needed privacy . The optimization problem that we solve in a federated learning setting is min w∈Rd L ( w ) = min w∈Rd 1 n n∑ k=1 E ( x , y ) ∼pk [ ` k ( w ; x , y ) ] ( 1 ) where ` k is the loss for client k and ( x , y ) samples from the kth client ’ s data distribution pk . A central server is coordinating training between the K local clients . The most prevalent algorithm for solving this optimization is the federated averaging ( FEDAVG ) algorithm ( McMahan et al. , 2017 ) . In this solution , each client has its own client model , parameterized bywk which is trained on a local dataset for E local epochs . When all clients have completed the training , their weights are sent to the central server where they are aggregated into a global model , parameterized by wg . In FEDAVG , the k client models are combined via layer-wise averaging of parameters , weighted by the size of their respective local datasets : wgt+1 ← ∑ k nk n wkt+1 , ( 2 ) where nk is the size of the dataset of client k and n = ∑ k nk . Finally , the new global model is sent out to each client , where it constitutes the starting point for the next round of ( local ) training . This process is repeated for a defined number of global communication rounds . The averaging of local models in parameter space generally works but requires some care to be taken in order to ensure convergence . McMahan et al . ( 2017 ) showed that all local models need to be initialized with the same random seed for FEDAVG to work . Extended phases of local training between communication rounds can similarly break training , indicating that the individual client models will over time diverge towards different local minima in the loss landscape . Similarly , different distributions between client datasets will also lead to divergence of client models ( McMahan et al. , 2017 ) . Depending on the use case , however , the existence of local datasets and the option to train models locally can be advantageous : specialized local models , optimized for the data distribution at hand may yield higher performance in the local context than a single global model . Keyboard prediction , for example , based on a global model may represent a good approximation of the population average , but could provide a better experience at the hands of a user when biased towards their individual writing style and word choices . A natural question arises : when is a global FL-trained model better than a specialized local model ? A specialist would be expected to perform better than a global generalist in a pathological non-iid setting , whereas the global generalist would be expected to perform better in an iid setting . To address the issue of specialized local models within the federated learning setting , we propose a general framework based on mixtures of experts of a local and a global model on each client . Local expert models on each client are trained in parallel to the global model , followed by training local gating functions hk ( x ) that aggregate the two models ’ output depending on the input . We show advantages of this approach over fine-tuning the global model on local data in a variety of settings , and analyze the effect that different levels of variation between the local data distributions have on performance . While standard federated learning already shows some privacy enhancing properties , it has been shown that in some settings , properties of the client and of the training data may be reconstructed from the weights communicated to the server ( Wang et al. , 2019 ) . To this end , in this paper we will work with a stronger notion of privacy . While existing solutions may be private enough for some settings , we will assume that a client that require privacy for some of its data , needs this data to not influence the training of the global model at all . Instead , our framework allows for complete opting out from the federation with all or some of the data at any given client . Clients with such preferences will still benefit from the global model and retain a high level of performance on their own , skewed data distribution . This is important when local datasets are particularly sensitive , as may be the case in medical applications . Our experimental evaluation demonstrate the robustness of our learning framework with different levels of skewness in the data , and under varying fractions of opt-out clients . 2 RELATED WORK . Distributed machine learning has been studied as a strategy to allow for training data to remain with the clients , giving it some aspects of privacy , while leveraging the power of learning from bigger data and compute ( Konečnỳ et al. , 2016 ; Shokri & Shmatikov , 2015 ; McMahan et al. , 2017 ; Vanhaesebrouck et al. , 2016 ; Bellet et al. , 2018 ) . The federated averaging technique ( McMahan et al. , 2017 ) has been influential and demonstrated that layer-wise averaging of the weights in neural network models trained separately at the clients is successful in many settings , producing a federated model that demonstrates some ability to generalize from limited subsets of data at the clients . However , it has been shown that federated averaging struggles when data is not independent and identically distributed among the clients ( the non-iid setting ) , which shows that there is a need for personalization within federated learning ( Kairouz et al. , 2019 ) . In general , addressing class imbalance with deep learning is still a relatively understudied problem ( Johnson & Khoshgoftaar , 2019 ) . A common approach for personalization is to first train a generalist model and then fine-tune it using more specific data . This approach is used in meta-learning ( Finn et al. , 2017 ) , domain adaptation ( Mansour et al. , 2009 ) , and transfer learning ( Oquab et al. , 2014 ) . This approach was proposed for the distributed setting by Wang et al . ( 2019 ) who used federated averaging to obtain a generalist model which was later fine-tuned locally on each client , using its specific training data . Some work has been inspired by the meta-learning paradigm to learn models that are specialized at the clients ( Jiang et al. , 2019 ; Fallah et al. , 2020 ) . Arivazhagan et al . ( 2019 ) combined this strategy and ideas from transfer learning with deep neural networks and presented a solution where shallow layers are frozen , and the deeper layers are retrained at every client . Zhao et al . ( 2018 ) propose a strategy to improve training on non-iid client data by creating a subset of data which is globally shared between all clients . Recent strategies have also explored knowledge distillation techniques for federated learning ( Jeong et al. , 2018 ; He et al. , 2020 ; Lin et al. , 2020 ) , which show promising results in non-iid settings . Hanzely & Richtárik ( 2020 ) proposed a solution that provides an explicit trade-off between global and local models by the introduction of an alternative learning scheme that does not take the full federation step at every round , but instead takes a step in the direction towards the federated average . Deng et al . ( 2020 ) proposed to combine a global model w trained using federated averaging , with a local model v with a weight αi . To find optimal αi they optimize α∗i = argminαi∈ [ 0,1 ] fi ( αiv + ( 1− αi ) w ) every communication round . While this weighting scheme will balance the two models , it has no way of adapting to the strengths of the different members of the mix . Mixture of experts ( Jacobs et al. , 1991 ) is the combination of several competing neural networks trained together with a gating network to solve a common task . It was presented as an ensemble method which can be trained end to end using gradient descent . In the current work , we will apply the mixture to leverage the specific strengths of a global model trained with federated averaging , and a local model trained locally on each client . 3 FEDERATED LEARNING USING A MIXTURE OF EXPERTS . In this work , we present a framework for federated learning that builds on federated averaging and mixtures of experts . Our framework includes a personalized model for each client , which is included in a mixture together with a globally trained model using federated learning . The local models never leave the clients , which gives strong privacy properties , while the global model is trained using federated averaging and leverages larger compute and data . In our framework as seen in Figure 1 , some clients can choose to opt-out from the federation , meaning that no information of their data is leaving the client , ensuring privacy of those clients . Let fg be the global model with parameters wg . We denote the index of clients by k and the local models by fkl with parameters w k l . The gating function is denoted by h k , parameterized with wkh . Training in the proposed framework is divided into three main parts . First , a global model fg is trained using federated averaging using opt-in data ( see Section 3.1 ) . Second , a local model fkl is trained using all available data on a client . Third , fg and fkl are further trained together with a gating model hk on each client locally , using all available data on the client . Steps one and two may be performed in parallel if allowed by the available resources . 3.1 PRIVACY GUARANTEES . The proposed framework allows for a strict form of privacy guarantee . Each client may choose an arbitrary part of their data which they consider being too sensitive to use for federated learning , and no information from this data will ever leave the client . The system will still leverage learning from this data by using it to train the local model fkl and the gating model h k. This is a very flexible and useful property . For example , it allows for a user to use the sensitive data in training of the private local models , while transforming it using some privatization mechanism and use the censored version to train the federated model . In general , each client dataset Dk is split into two non-overlapping datasets , DkO and DkI , one of which has to be non-empty . The local model fkl and the gating model h k is trained using the whole dataset Dk = DkO ∪ DkI , while the global model fg is trained with FEDAVG using only the opt-in dataset DkI . This is visualized in Figure 1 .

Summary: The proposed method is a federated method allowing to have a certain amount of data shared between all the learners and some data specific to each learner. The targeted field of application is classification for problems where strong privacy is crucial. The method consists in learning a global classifier (with the shared data) as well as local classifiers (one per learner, using the local data). The inference, for each learner, is done with a local expert (another neural network) trained to combine inferences from the local and global models.