F EDERATED LEARNING USING A MIXTURE OF EXPERTS
Anonymous authors
Paper under double-blind review
A BSTRACT
Federated learning has received attention for its efﬁciency and privacy beneﬁts, in settings where data is distributed among devices.
Although federated learning shows signiﬁcant promise as a key approach when data cannot be shared or centralized, current incarnations show limited privacy properties and have shortcomings when applied to common real-world scenarios.
One such scenario is heterogeneous data among devices, where data may come from different generating distributions.
In this paper, we propose a federated learning framework using a mixture of experts to balance the specialist nature of a locally trained model with the generalist knowledge of a global model in a federated learning setting.
Our results show that the mixture of experts model is better suited as a personalized model for devices when data is heterogeneous, outperforming both global and local models.
Furthermore, our framework gives strict privacy guarantees, which allows clients to select parts of their data that may be excluded from the federation.
The evaluation shows that the proposed solution is robust to the setting where some users require a strict privacy setting and do not disclose their models to a central server at all, opting out from the federation partially or entirely.
The proposed framework is general enough to include any kinds of machine learning models, and can even use combinations of different kinds.
I NTRODUCTION
[Figure]
In many real-world scenarios, data is distributed over a large number of devices, due to privacy concerns or communication limitations.
Federated learning is framework that can data in a distributed learning setup.
This allows for exploiting both the compute power of all participating clients, and to beneﬁt from a large joint training data set.
Furthermore, this is beneﬁcial 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 ﬁeld 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 ﬁrm 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
[Equation]
where (cid:96) k is the loss for client k and ( x, y ) samples from the k th client’s data distribution p k .
A central server is coordinating training between the K local clients.
The most prevalent algorithm for solving this optimization is the federated averaging (F ED A VG ) algorithm (McMahan et al., 2017).
In this solution, each client has its own client model, parameterized by w k 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 w g .
In F ED A VG , the k client models are combined via layer-wise averaging of parameters, weighted by the size of their respective local datasets:
[Equation]
where n k is the size of the dataset of client k and n = (cid:80) k n k .
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 deﬁned 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 F ED A VG 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
h k ( x ) that aggregate the two models’ output depending on the input.
We show advantages of this approach over ﬁne-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 inﬂuence 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 beneﬁt 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.
R ELATED 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 (Koneˇcn`y 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 inﬂuential 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 ﬁrst train a generalist model and then ﬁne-tune it using more speciﬁc 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 ﬁne-tuned locally on each client, using its speciﬁc 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.
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 & Richt ´ arik (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 ﬁnd optimal α i they optimize α ∗ i = arg min α i ∈ [0 , 1] f i ( α i v + (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 speciﬁc strengths of a global model trained with federated averaging, and a local model trained locally on each client.
F EDERATED 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 f g be the global model with parameters w g .
We denote the index of clients by k and the local models by f kl with parameters w kl .
The gating function is denoted by h k , parameterized with w kh .
Training in the proposed framework is divided into three main parts.
First, a global model f g is trained using federated averaging using opt-in data (see Section 3.1).
Second, a local model f kl is trained using all available data on a client.
Third, f g and f k l are further trained together with a gating model h k 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. P RIVACY 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 f kl and the gating model h k .
This is a very ﬂexible 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 D k is split into two non-overlapping datasets, D k O and D k I , one of which has to be non-empty.
The local model f kl and the gating model
h k is trained using the whole dataset D k
= D k O ∪ D k
I , while the global model f g is trained with F ED A VG using only the opt-in dataset D k I .
This is visualized in Figure 1.
3. O PTIMIZATION
Step 1: Federated averaging.
We train the global model using F ED A VG .
In other words, globally we optimize
[Equation]
for the opt-in dataset D k I .
Here (cid:96) k is the loss for the global model w g on client k for the prediction f g ( x ) =
ˆ y g , and D k
I is the k th clients opt-in data distribution.
Step 2: Train local models.
The local models f kl are trained only locally, sharing no information between clients, minimizing the the local loss over w kl ∈ R d
[Equation]
Here (cid:96) k is the loss for the prediction ˆ y l = f kl ( w kl ; x ) from the local model on the input x and D k is the k th clients dataset.
Step 3: Train local mixtures.
The local mixture of experts are trained using the gating models h k , with the prediction error given by weighing the trained models f g and f kl :
[Equation]
In other words, at the end of a communication round, given f kl and f g , we optimize the mixture equation 5:
[Equation]
locally for every client k
= 1 , . .
. , n .
Here (cid:96) k is the loss from predicting ˆ y for the label y given the input x with the model from equation 5 over the data distribution D k of client k .
A summary of the method is described in Algorithm 1.
Algorithm
1: input: Models participating in F ED A VG w 1 , . .
w k , local expert models w kl , local gate w kh , learning rate lr , decay rates β 1 , β2: Randomly initialize w 1 , . . .
, w k with the same seed.
3: Randomly initialize w k l and w k h . 4: w g ← F ED A VG ( w 1 , . . . , w k ) //
Train for E local epochs and G communication rounds 5: for client k do 6: w kl ← Adam ( w kl , lr, β 1 , β 2 ) //
Train local model on client k 7: w g , w kl , w kh ← Adam ( w g , w kl , w kh , lr, β 1 , β 2 ) //
Train mixture of experts on client k 8: end for
3. E XPERIMENTAL SETUP
Datasets.
Our experiments are carried out on the datasets CIFAR-10, CIFAR-100 (Krizhevsky et al., 2009) and Fashion-MNIST (Xiao et al., 2017).
In order to simulate heterogeneous client data without labels overlapping between clients , we partition the data into 5 clients for CIFAR-10 and
Fashion-MNIST, and 50 clients for CIFAR-100.
For CIFAR-100 we used a client sampling strategy where we in each communication round in the federation randomly sampled a fraction of 0.1 of clients to participate.
Sampling non-iid.
The datasets are sampled in such a way that each client’s datasets contains two majority classes which together form p % of the client data and the remaining classes form (1 − p )% of the client data.
We perform experiments where we vary p to see what effect the degree of heterogeneity has on performance.
In the extreme case p = 1 . 0 , each client only has two labels in total.
Note that this is an extreme partitioning, where there is no overlap of majority class labels between clients, i.e. the local data distributions P i ∩ P j = ∅ for all pairs of clients i, j .
Also note that, we perform experiments with very little data in our experimental setup.
A summary of the experimental set-up can be seen in Table 1.
Opt-out factor.
Some users might want to opt-out from participating to a global model, due to privacy reasons.
These users will still receive the global model.
To simulate this scenario in the experimental evaluation, we introduce an opt-out factor denoted by q .
This is a fraction deciding the number of clients participating in the F ED A VG optimization.
The clients that participate in the federated learning optimization have all their data in D k I , while the rest of the clients that optout have all their data in D k O .
q = 0 means all clients are opt-in and participating.
We perform experiments varying q , to see how robust our algorithm is to different levels of client participation.
In Figure 1 we visualize how the opt-out factor can be used.
Models.
In our setup, both the local model f l and the global model f
g are CNNs with the same architecture.
However, they are not constrained to be the same model and could be implemented any two differentiable models.
The CNN has two convolutional layers with a kernel size of 5, and two fully-connected layers.
All layers have ReLU activations.
The gating function h k has the same architecture as f g and f l , but with a sigmoid activation in the last layer.
We use Adam (Kingma &
Ba, 2014) to optimize all models, with a learning rate of 0 . 0001 in all experiments.
Baselines.
We use three different models as baselines.
First, the locally trained model f kl for each client.
Second, F ED A VG .
Third, the ﬁnal model output from F ED A VG ﬁne-tuned for each client on its own local data.
We train f kl , the ﬁne-tuned model and the mixture using early stopping forepochs, monitoring validation loss on each client.
Evaluation .
For evaluation we have a held-out validation set for each client.
For both CIFARand CIFAR-100 we have n = 400 data points for evaluation per client, sampled with the same majority class fraction p .
We report an average accuracy over all clients.
Experimental set-up
Datasets
CIFAR-10, CIFAR-100, Fashion-MNIST
Models ( f g , f kl , h k )
CNNs: 2 conv.
layers (5x5 kernels) and two FC layers
Activation functions
ReLU
Communication rounds (FedAvg)
Local epochs (FedAvg)
Local batch size (FedAvg)
Epochs (mixture) 200 with early stopping
Batch size (mixture)
Optimizer
Adam
Learning rate
[Table]
Decay rates ( β 1 , β 2 ) (0 . 9 , 0 . 999)
Training data size per client { 5 , 10 , 20 , 38 , 50 , 62 , 75 , 88 , 100 , 500 }
Majority class fraction p { 0 . 2 , 0 . 3 , 0 . 4 , 0 . 5 , 0 . 6 , 0 . 7 , 0 . 8 , 0 . 9 , 1 . 0 }
Validation data size per client
R ESULTS
For the sake of reproducibility, all code will be made available.
In Table 2 we report accuracies and standard deviations on CIFAR-10 for all models when data is highly non-iid, i.e. for p = { 0 . 8 , 0 .
In Figure 3 we report accuracies over all majority fractions p .
By comparing Figures 3a, 3b and 3c, we see that F ED A VG performs substantially worse when the opt-out fraction increases, i.e. when more clients opt-out from the federation.
As a consequence of a weakened global model, the ﬁne-tuned baseline decreases much in performance as well.
However, our proposed mixture leverages both the global and local models and does not degrade as much in performance.
In Figures 2 and 4 we see that we see that the mixture performs just as well, or better, than the ﬁne-tuned baseline for CIFAR-100 and Fashion-MNIST, respectively.
Experiments were also carried out to see what effect training data size has on performance.
This is summarized as a heatmap in Figure 5a for CIFAR-10 and in Figure 5b for Fashion-MNIST, where the difference in validation accuracy between the mixture and the ﬁne-tuned baseline is shown for different majority class fractions p and different client train set sizes.
We see here that the mixture model outperforms the baseline most of the time, a trend which gets stronger when the training set size increases.
[Table]
[Figure] [Figure]
[Figure]
D ISCUSSION
To address the problems of learning a personalized model in a federated setting when the client data is heterogeneous, we have proposed a novel framework for federated mixtures of experts where a global model is combined with local specialist models.
We ﬁnd that with skewed non-iid data on the clients, our approach outperforms all other baselines, including F ED A VG , a locally trained model, and models trained ﬁrst with F ED A VG and then ﬁne-tuned on each local client in most settings.
The experimental evaluation for CIFAR-10 shows that our approach outperforms all other methods, including the strong ﬁne-tuning baseline (see Figure 3).
In Figure 5 we see that our proposed model outperforms the ﬁne-tuned baseline in most settings for CIFAR10 and Fashion-MNIST, especially when the number of data points per client increase.
For CIFAR-100, the proposed framework outperforms all other methods, regardless of the level of skewness (see Figure 2).
In this setting, a large part of the training data for each client comes from a very limited set of the available classes: two out of 100.
As such, very few training examples will be available from the minority classes.
This is a crucial result: the proposed framework is very robust to extremely skewed training data.
The framework also gives strong privacy guarantees, where clients are able to opt-out from federation, keeping their data private.
The experiments show that our proposed solution is robust to a high opt-out fraction of users, as seen in Figures 3 and 4, whereas the ﬁne-tuned baseline is not.
C ONCLUSIONS
In this work, we have presented a framework for federated learning that builds on mixtures of experts.
This framework allows us to learn a model that balances the generalist nature of the global federated model and the specialist nature of the local models.
Our approach is not only an intuitive approach for the generalist vs specialist balance, but also allows for varying participation of the different clients in the federation.
Clients may either opt-in and participate in the federation, or opt-out entirely by training only a local model with all its local data but still receive a global model from the opt-in participants.
This gives a ﬂexible solution for strong privacy guarantees in real-world settings.
Our experiments show that in the setting where many clients opt-out from federation, the ﬁne-tuned baseline degrades in performance whereas our proposed mixture model does not.
The proposed framework is compatible with any gradient-based machine learning models, and can incorporate combinations of these, strengthening the potential of this direction of research, and leveraging the beneﬁcial properties of ensembles of various machine learning models.
The experimental evaluation conducted in this work showed the proposed solution to achieve stateof-the-art results in two different benchmark datasets when data is highly skewed, and when parts of the clients in the federation opts out from the training.