Federated Mixture of Experts
Authors: Matthias Reisser
Christos Louizos
Efstratios Gavves
Max Welling
Venue:  (2021)
Keywords: Federated Learning
personalized models
non-i.i.d data

Abstract: Federated learning (FL) has emerged as the predominant approach for collaborative training of neural network models across multiple users, without the need to gather the data at a central location. One of the important challenges in this setting is data heterogeneity; different users have different data characteristics. For this reason, training and using a single global model might be suboptimal when considering the performance of each of the individual user’s data. In this work, we tackle this problem via Federated Mixture of Experts, FedMix, a framework that allows us to train an ensemble of specialized models. FedMix adaptively selects and trains a user-specific selection of the ensemble members. We show that users with similar data characteristics select the same members and therefore share statistical strength while mitigating the effect of non-i.i.d data. Empirically, we show through an extensive experimental evaluation that FedMix improves performance compared to using a single global model while requiring similar or less communication costs.

Under review as a conference paper at ICLR 2021
FEDERATED MIXTURE OF EXPERTS
Anonymous authors Paper under double-blind review
ABSTRACT
Federated learning (FL.) has emerged as the predominant approach for collabora- tive training of neural network models across multiple users, without the need to gather the data at a central location. One of the important challenges in this set- ting is data heterogeneity; different users have different data characteristics. For this reason, training and using a single global model might be suboptimal when considering the performance of each of the individual user's data. In this work, we tackle this problem via Federated Mixture of Experts, FedMi x, a framework that allows us to train ensemble of specialized models. FedMix adaptively selects and trains user-specific selection of the ensemble members. We show that users with similar data characteristics select the same members and therefore share statistical strength while mitigating the effect of non-i.i.d data. Empirically, we show through an extensive experimental evaluation that FecMix improves performance compared to using a single global model while requiring similar less communication costs.
INTRODUCTION
An ever-increasing amount of devices are being con- nected the internet, sensing their environment, and generating vast amounts of data. The term federated learning (FL) has been established to describe the scenario where we aim to learn from the data gener- ated by this "federation" of devices (McMahan et al.) 2016). Not only does the number of sensing devices increase, but also their processing power increas- ing continuously to the point that it becomes viable to perform inference and training of machine learn- ing models on device. In federated learning, the goal is to learn from these client devices' data without collecting the data centrally, which naturally allows for more private exchange of information.
Gradient Divergence
3.75
3.25
3.00
2.75
FedAvg
2.50
FedMix
2.25
500
1000
1500
2000
Communication rounds
Figure sliding window the gradient
divergence (defined in Appendix D), on Ci- farl0 in the setup of Scction 4 for FedAvg and FedMix (K 4).
Several challenges arise in the federated scenario. Federated devices are generally snoource-constraied both in their computational capacity as well as in communication bandwidth and latency. practical example, smartphone has limited heat dissipation capacity and must communicate via Wi-Fi. From global perspective, devices' pro- cessing power and network connection can be highly heterogeneous across geographical regions status of device owners, causing practical issues (Bonawitz et al.) 2019) and raising questions of fairness FL (Li et al. 2019 Mohri et al. 2019).
and
One of the key challenges in FL that we aim to address in this work is the non-i.i.d nature of the shards of data that are distributed across devices. In non-fedeeated machine learning, assuming independent and identically distributed data is generally justifiable and not detrimental to model performance. In FL however, each client performs a series of parameter updates on its own data shard to amortize the costs of communication. Over time, the direction of progress across shards with non-i.i.d data starts diverging (as shown Figure []. which can set back training progress, significantly slow down convergence and decrease model performance (Hsu et al.) 2019).
To this end, we propose Federated Mixture of Experts (FedMix), an algorithm for FL that allows for training ensemble of specialized models instead of single global model. In FedMi x, expert

---

Under review as a conference paper at ICLR 2021
models are learning to specialize in regions of the input space such that, for given expert, each client's progress on that expert is aligned. FedMix x allows each client to learn which experts are relevant for its shard and show how it be extended for inference client. FedMi x shows competitive performance against the established standard in FI., FedAvg (McMahan al. 2016) Deng al. 2020) across range of visual classification tasks. Code will be released upon publication.
previously unseen
FEDERATED MIXTURE OF EXPERTS
Federated learning McMahan et al. 2016) deals with the problem of learning server model with parameters W, e.g., neural network, from dataset D = {[(x,, (i),,.1).... ................ (xN, Ux)] of N datapoints that distributed across shards, i.e., = D, .... Ds. without accessing the shard specific datasets directly. By defining loss function (D, w) per shard, the total risk can be written as
No NN (D.w:w)
arg miim Σ
L.(D.::) = N, Σ L(Daiw).
is easy see that this objective corresponds to empirical risk minimization over the joint dataset D with loss L(-) for each datapoint. In federated learning one is interested in reducing the communi- cation costs; for this reason McMahan et al. (2016) propose to do multiple gradient updates for w in the inner optimization objective for each shard thus obtaining "local"" models with parameters w,- These multiple gradient updates are denoted as "local epochs", i.e., amount of passes through the en- tire local dataset, with an abbreviation of E. Each of the shards then communicates the local model W, to the server and the server updates the global model at "round" by averaging the parameters N w This constitutes federated averaging (F (FedAvg) (McMahan
of the local models w = Σ. N. et al. 2016), the standard in federated learning.
One of the main challenges in federated learning is the fact that usually the data are non-i.id. dis- tributed across the shards S, that is p(D(s; # p(D]s,) for * j. On the one hand, this make learning a single global model from all of the data with the classical FedAvg problematic. On the other hand, there is extreme that does not suffer from this issue; learning S individual models, ie., only optimizing W, on D., Although these individual models by definition do not suffer from non-i.i.d data, clearly should aim do better and exchange meaningful information between clients to learn more robust and expressive models.
2.1 THE FEDMIX ALGORITHM
S
With FedMix, we propose to strike balance be- tween the two aforementioned extremes; learning a single global model and learning S individual mod- els. For this reason, mulation, the Mixture of Experts (MoE). The clas- sical formulation of MoE model (Jacobs et al.) [99]]Jrrdan & Jacobs 1994) contains a set of K ex- perts and gating mechanism that is responsible for choosing an expert for a given data-point. A MoE model for data point (x,y) generally be de- scribed by
revisit an old model for-
y
Figure 2: FedMix graphical model. The generative model is depicted with solid lines and the inference model with dashed lines.
PWLK ( Σ pw. z)pe
where z is a categorical variable that denotes the expert, wy are the parameters of expert & and 0 are the parameters the selection mechanism.
The MoE was proposed model for datasets where different subsets of the data exhibit different relationships between input x and output y. Instead of training a single global model to fit this relationship everywhere, each expert performs well different subset of the input space. The gating function models the decision boundary between input regions, assigning data-points from subsets of the input region their respective experts.
2

---

Under review as a conference paper at ICLR 2021
In this work, show that, in the federated scenario, sub-dividing the input region through MoE can alleviate the consequences of non-i..d data by aligning gradient updates across experts (Figure D.. In Federated Mixture of Experts (FedMix) enrich this model by conditioning the gating mechanism on the shard assignment s. Whatever characteristics make shard & different from other shards can manifest in learning different, localized gating mechanism that does not need to be communicated to the server. In choosing K = 1, FedMix recovers the standard setting of federated averaging. K S in combination with fixing p(z s|x, s) recovers S independent models.
From a global perspective, we are interested in maximizing the following single objective:
s N.
N,
ΣΣ log Pwi. (Va,) X5.5 8)
log
Po. (2) (2)xx.5. P) (ys D 2)
(3)
Given the graphical model decomposition depicted in Figure 2 the objective in Eq. 3 corresponds to a federated MoE, where we have omitted the generative models p(z[s). We will briefly touch upon the role of learning generative models in Appendix but focus on the discriminative part of the model, ie.. the MoE, this paper.
While it is possible to optimize Eq. 3 directly, we have found empirically that it is hard to achieve both: avoiding collapse to single expert, thus obtaining FedAvg, and specialization of the experts. Instead, we propose to form variational lower-bound on Eq. 3 with global varia- tional approximation (=......) to the true posterior p(z|x, 8) with parameters Φb At test time, p(ylx", s) = Ex=pP(ylss ;)) =)p(=xx", s) can be readily evaluated without requiring q. This allows us condition [o ...) on any available iede-informttion at training time that might result in bet- ter specialization in the non-i.i.d federated scenario. In this paper mainly consider classification tasks whose non-i.i.d nature predominantly stems from the non-i.i.d distribution labels y. Other or additional known sources of misalignment could be included to further improve this approxima- tion, such as a manufacturer-id for a medical device in a medical scenario, a geographic identifier, general domain-speciff information. We show one such additional example in 4.2 The lower bound to be maximized in F FedMix x therefore is as follows:
S N.
logPwi.k 8, (Ja,)
(4)
K L (=(=1,..)))og llog Pwr, (Va,iXa,;;, z)po, (x,i,,;))) log (o(=/le.)))
(5)
(2 lo. loggw, pw. (x,,i x z)po, (2
(2/p0,2//x00 (2|xa,))))) H (as(s)))..)).
(6)
0.8
Conditioning only on y allows us to efficiently pa- rameterize the variational approximation, incurring only a small communication overhead. While would be possible to condition (o(2ly) having localized approximations with parameters d that do not need be communicated, we found global approximation to help align the gating mech- anisms across shards. A global (2(2))) (lyy) encourages shards that contain data with the same label to assign them the same expert.
Validation accuracy
0.7
0.6
thus
0.5
0.4
FedAvg
0.3
FedMix, K 4, with H(g)
FedMixx k= 4, without M(g)
0.2
0
500
1000
1500
2000
Communication rounds
Specialization of the experts is key ingredient for FedMix x to be successful: with specialization, the gradients for each expert become aligned across shards (see Figure []]. the hold-out accuracy im- proves (see Figure 3). and the communication costs decrease as each shard may only need to access a subset of the experts. We find that performing maximum a-posteriori (MAP) inference for gener- ally leads to better and more personalized models. By removing the entropy term from equation 4 (s(z)y) and therefore po, (z|x, s) encouraged concentrate and select only one expert for
Figure 3: The effect of specialisation (with- out Hg)) compared to an ensemble (with (g))) and FedAvg on C0 The experi- mental setup is identical to what is described in Section 4
3

---

Under review as a conference paper at ICLR 2021
given data point. In the extreme case where client's shard contains only data that is assigned to the same expert, we can reduce communication by receiving and sending updates for that single expert only. We show in Section
$ that communicating and evaluating experts based on thresholding the aggregate [=[=))) = Ey [o (=ly)] can reduce communication and computation overhead.
Figure 3] compares FedMix with and without the entropy term to standard FedAvg tion of communication steps. With the entropy term, FedMix x develops no expert specializa- tion and collapses to an ensemble of K 4 models. One drawback of the heavy specializa- tion with MAP inference is that sometimes FedMix x prematurely completely prunes experts, i.e. (=.(= k k| s) = 0 Vx, s. This can be undesirable as we lose model capacity that can be used for better modeling the data. As (2ly)) is one of the main training signals of pe, (2|x, s), we introduce the marginal entropy term the server, H H (Fp(t)) (=>(=]))]) as regularizer. Notice that this leads to different training dynamics than locally optimizing the lower bound with the entropy included and we, empirically, found that alleviates premature pruning, while still leading to specialized models. Figure 14a in Appendix H visualizes the development of go (g)) over time from initially uniform to high specialisation for the experiment depicted in Figure
func-
Server Side Updates In general federated learning algorithm, central server selects subset S" C. {I,...,S} of clients at time and transmits the current estimate of the global parameters w to them. These clients perform a series of mini-batch gradient updates with data from their shard DD on a local loss function, which can come at the price of each client moving in possibly different directions in parameter space. In generalized FedAvg (Reddi et al.) 2020). the server interprets A w' w as a single-step gradient update from client 8, averages those gradients and applies an optimizer such as Adam (Kingma & Ba 2014) receive w ++1 In light of non-i.i.d data across clients, this averaging strategy can result in slow progress since averaging updates in a highly non-convex parameter space can be sub-optimal. In FedMix, this effect is mitigated since for a given expert, the data that is used to update its parameters are aligned better across shards.
FedMix offers a second way to improve convergence speed by modifying the server-side updates. In generalized FedAvg, the individual gradients returned by the subset S of clients are averaged according to
S' Σ p(s) A p(s) No
N
A
In FedMix, x, we can speed up convergence by considering expert-specific updates 4 = w w If client pruned away expert from its local gating mechanism, Ai., will be zero. We propose normalize the effective magnitude of the resulting update A by up-weighing the updates of all other clients that do consider expert l for their local mixture:
N₂ Ns
A Σ p(s|z k)) Ai... p(s|z k) x p(2 kss)p(s), p(s)
(8)
Computing p(=]s) = Ex [po, (2|8, x)) prior to sending updates to the server involves evaluating potentially large neural network models. Therefore choose approximate p(z[s) ss gz(2))) Ey,D, [x ( ww), which involves just a single matrix multiplication. We discuss the implications of
sending 2 (2 to privacy and how these fare relative to FedAvg in Appendix C
Pruning experts Maintaining the cheap-to-comput marginal posterior per shard offers an addi- tional opportunity to increase computation speed during local shard iterations and reduce overall communication costs. We propose to "prune away" experts locally from the MoE if (=(2)) does not surpass threshold n/K. In order to still optimize valid bound, we need to re-normalize (=(=))) before evaluation of the loss function (Pal et al. 2005). We evaluate the same threshold prior to sending updates to the server in order to avoid communicating parameters that have not changed during the client's iterations. Once the server selects a client for another round, it provides the client that were updated by the client in the previous round. We empiri- cally find that the entropy of q(2(s) decreases steadily and we prune away experts & with probability {(= = k|s) < // K without significant drop in performance. We explore the consequences of prun- ing experts in the experiment section and in Appendix E Algorithm 2 shows how FedMi x can be enriched by pruning.
only those experts
4

---

Under review as a conference paper at ICLR 2021
Algorithm 1 The FedMix algorithm. a are the client and server learning rates respectively
Designing robust gates In the federated scenario, N₂ often much smaller than N and espe- cially small in relation to the com- plexity of the data we try model. Any localized parameters therefore are prone overfitting. On the other hand, the global parameters of an expert are trained using all data-points assigned to that expert across all shards, allowing to learn more robust features.
function SERVER SIDE
Initialize & and K vectors W -WW wK] for round & in .. do S + random subset of the clients
Initialize Aw = 0, 3 =
for s in S do
W "..((z + CLIENT SIDE(s, Φ. W) end for p(s|z) t p(2)s)p(ss) EeEs p(z)s)p(s)
for s in S' do
We can make use of the robust- ness of these experts' features for the gating mechanism by condi- tioning them instead of train- ing an entirely separate model for Po, (=|x, s). Let define (x) as intermediary features of expert k. Since not all experts might be used for a given shard and in order scale with K, we average over the marginal posterior of the train- ing set at that shard before applying linear transformation to compute the input to the softmax gates:
A
pss|2 =k)(w(-) w Vk N.
A
No (db-1- (5-)
end for A
Vo 1o -00(=ly=c)ppy-00000
I:K < ADAM(A, :) + ADAM(A
end for end function
function CLIENT SIDE(S, Φb W)
Get local parameters 0₂ for epoch in 1,..., do
for batch b E B do
Lg < Bea(slus) [log Pww (y6l x6: z)Do, ( (z x6. s)]
K Asl= k)s)/u(x) k=1
o+ aVaa VoL
hx(x)
Easlz
W + a V wL.
0.+ Vo, Lss
Po. (2|x, s)= SM (A. (x) +b,)
end for end for (Z]ss) + EDyD, lgs (z|y)] return WI.K, Φ, (2(s)
where 0, (A., b.) local learnable parameters and SM rep- resents the softmax function.
end function
Inference at test time We consider three variants for test-time evaluation of FedMix. In the first case, client that participated in training is presented with new data point (x", s). Predictions can then be straightforwardly done by selecting the that maximizes Σx= p(u x", <)p(2[x*, second, more challenging, scenario data set D.... Here we propose to instantiate and train the local gating mechanism by optimizing the parameters 9₂ of po, (2|x, s*) via MAP inference the local objective. Afterwards, predictions can be made in a manner similar to the first case.
, s). In the client s* is introduced together with new labelled local
Finally, we consider the case in which new client s" has no labelled dataset available. Without a local gating function, simply ensembling experts exhibits almost random behaviour since experts can be overly confident on -iitr--diitritution data (Snoek et al. 2019). We therefore propose to en- semble across local gating mechanisms to compute p(2[x") >> ==1 Pe, (z|x", ())(s(x*) ); method which works well in practice. In Appendix [ we discuss results for new shard inference as well as more principled approach which makes use the graphical model formulation in Figure
RELATED WORK
FedMix has similarities to many recent works in the topic of federated learning. Two methods closely related to ours are described in (Sattler et al. 2019 Briggs et al. 2020). The authors propose to perform hierarchical clustering on the updates returned from each shard in order to incrementally create separate models for groups of users, with a cluster assignment mechanism based on hand- crafted heuristies. FedMi x instead takes different approach; starts with fixed set of K models
5

---

Under review as a conference paper at ICLR 2021
and then optimizes with gradient descent at each shard a per-datapoint model assignment mechanism that can better fit the peculiarities of the local dataset.
Another closely related work is presented by Mansour et al. (020), where the authors propose to similarly create an ensemble of K-models and assign to each shard the model that achieves the lowest training loss the local dataset. This is closer the assignment that happens in FedMix with one main difference; FedMix x takes into account the uncertainty in the selection mechanism as well with p(z]x, x, s) instead of selecting the top performing component during training. This is beneficial early in training where the models have not fully specialized yet. Using local and global model parts has also been explored by Liang al. (2020). The authors propose to have local feature extractor at each shard and global classifier on top of those features as opposed to having K --epparate models and a local selection mechanism as in FedMix x. This setup yields improvements upon the vanilla federated averaging algorithm, however there are two potential drawbacks; first, empirically, the authors had to start their procedure from a pre-trained model with FedAvg and secondly, they have to ensemble all of the different feature extractors for predictions in new shards. In our experiments, we omit the pre-training step and show that the ensembling strategy fails.
Federated learning in the non-i.i.d setting can also be improved upon in other ways. Li et al. (2018) propose to employ a proximal regularizer at the shard level in order to prevent the local models from drifting too far from the global model, thus making federated learning more robust. Diang et al. (019) notice that FedAvg and Reptile (Nichol et al. 2018). a meta-learning algorithm, are essentially the same algorithm and thus propose fine tuning with Reptile in order improve the personalized performance of the global model. In a similar vein, there are promising new works that explore the meta-learning view of federated learning (Chen et al. 2018 Khodak et 2019 Fallah et al.) 2020). Improving the personalized performance of the global model has also been done without meta-learning in works such as by [eng et al.] (2020); Mansour al. (2020). In
general, such improvements are complementary to FedMix x and can be used to further enhance its performance. We refer the interested readers to the recent surveys by (Kairouz al. 2019 Kulkarni
et al. 2020).
4 EXPERIMENTS
We evaluate FedMi x three datasets: Cifar10, Cifar100 (Krizhevsky al. 2009) and Femnist Caldas et al., 2018), a 62-way image classification problem on hand-written digits and letters that naturally non-i.i.d due different writing styles of 3500 users. In Appendix A we detail the experimental setup and provide additional ablation studies in Appendix
4.1 PERFORMANCE VS. COMMUNICATION
We compare FedMix along several dimensions to baselines such as (generalized) FedAvg (Reddi et al., 2020), biased FedAvg, and the Local/Global approach of Liang et al. (2020). In biased FedAvg, we allow each client to learn a personalized bias vector be of its output layer. This will allow biased FedAvg model the label skew at each client but, fundamentally, cannot model any other form of non-i.i.d-ness. Similarly, Liang et al. (2020) propose to split the model into local and global components by having local feature extractors and learning the upper layers of the neural network via FedAvg. We experimented with splitting LeNet-5 at every intermediate layer and report results with the best performing split: keeping the input layer local. For ResNet-20, splicting after the first block performed best. First, we show that training with FedMix achieves higher global model accuracy. Second, we show that for a given communication budget, FedMix can be competitive against baselines. Table contains test set results for Cifar10, evaluated after 2k communication steps and Femnist averaged over the last 2k steps. Figures 4dl 4f reveal that the models have converged at that point. For Cifar100, observe overfitting and report test-set accuracies averaged over the last 10 evaluations. Figure44shove learning curves for different settings of F FedMix. x. All of these results were obtained by using the server version of the global parameters; the results with additional local fine-tuning at the client level are discussed in Appendix G
For Cifar10 we can that except for the initial training phase, FedMix improves over FedAvg for a given communication budget. For the same communication costs per communication round, FedMi x offers better performance than FedAvg, can be seen in comparing F FedMix (K 2) to

---

Under review as a conference paper at ICLR 2021
FedAvg LeNet-5 with 61 channels and FedMix (K 4) to FedAvg LeNet-5 with 113 channels. The channel count has been chosen such that the model size approximates 2 and separate standard LeNet experts respectively. Against biased F FedAvg and Local/Global Liangooaaaaangggggggggga et al. (2020). FedMix can improve significantly in terms of inference on a new shard (Figure 12 in Appendix F). When endowing FedMi with the ability to learn single personalized bias vector b, across all experts, can harness some of the advantages of biased FedAvg for inference on a new shard while retaining FedMix x' advantage when performing inference on new shard.
For Cifar100, we can see that learning progresses much faster for FedMix in terms of communi- cation rounds than FedAvg. Against FedAvg, we can see that FecMix improves with both K-4 and K=10 on a per round progression, and with K=10 and expert pruning, it eventually provides better accuracy and communication trade-off than FedAvg. Biased FedAvg seems be the better inductive bias in this task, although one should be weary that its usefulness is limited to label skew and does not generalize to other sources of non-i.i.d.ness (see section (.2). This is in contrast to F FedMix which is designed to handle any type of explicit non-i.id.neess.
For Femnist we compare against FedAvg and observe that generally FedMix x more prone to overfitting due to the extra amount of model parameters and expert specialization. For this reason, we learn the parameters of the local gating networks via FedAvg across all clients. The gating network is still localized through the of ((2(s) in pooling of (x) (equation 9). We that increasing K improves the performance of FedMi x. Since FedMis x effectively requires the model to first identify the right expert, followed by evaluating that expert's prediction, compounding can lead to F FedMix x being outperformed by FedAvg, such as for K 2. Biased FedAvg performs similar standard FedAvg since the non-i.i.d-nes. does not predominantly stem from label-skew.
Validation accurac
0.8-
Valldation accuracy
0.5
0.850
0,875
0.7-
0.4-
FesAug
Feddug
0.800
0.5
Facihing Le1bee-11133)
Fechwny Lelee--1101)
0.2-
Foskiis K 4. - 0..00
Fostiis K=
0.775
Locattilobal
Riased Fedligg
0.4-
Fecitiic
0.1
0.750
Locattiooall
Locatiriball
Fedliix K
0.0
100
200
300
200
400
600
800
1000
250
500
750
1000
1250
GO communicated
GB communicated
G0 communicateed
(a) Ci
(b) Cifar100
(c) Femnist
Polddtooon accurac
0.7-
0.4
Validation
FecAug
Dlased reditg
0.3-
Velidaiton
0.80
reshug
0.5-
Feoking Lultul-11133)
Suaad huddng Fodiis K-44
LocalSiicbal
0.4-
Fochiic d
Fealisa K
6.1
LocaV/rotial
loca /Clcbal
Fediiix
6.0
500
1000
1500
2000
2000
4000
5000
2000
4000
6000
Communication rounds
Communication rounds
Communication munds
(d) Cifar10
(c) Cifar100
(f) Femmist
Figure 4: Average accuracy across all clients (y-axis) as a function of the amount of GB commu- nicated (top row) and as function of communication rounds (bottom row). Cifar 10 models trained on the standard 45k training split. x-axes have been truncated for improved visibility. Best viewed in color.
98.5
99.5
4.2 ROTATED MNIST
98.0
90.0
To show that FedMix x is not lim- ited to label-skew, we create federated rotated MNIST dataset with 100 clients. Instead of la- bel skew, each client randomly chooses multiple of 45 degrees
Valdattoon
PadAva Blased Fedlig
Fackog
07.0
07.0
Blased Festing
96.5
redMiix i staly)
96.5
Teciitt K s.gle
FedMiix sttyrtt
Fectiiie rodr
500
1000
1500
500
1000
1500
Communicaaion rounds
Communication rounds
Figure 5: Only rotation
Figure 6: Rotation and labels

---

Under review as a conference paper at ICLR 2021
Table 1: Average test-set accuracies across clients and communication costs (rounds and GB) for Cifar10, Cifar100 and Femnist. For Cifarl0 we report the performance after 2k rounds, for Cifar100 we report the average over the last 10 evaluations at the end of training and for Femnist we report the average over the last 2k communication rounds.
Method
Cifar 10
Cifar 100
Femmist
Acc.
Comm
Acc.
Comm
Acc.
Comm
KedAvg
69.36%
2k, 137.3GB
49.58%
30k, 608G
85.73%
6k, 272GB
biased FedAvg
83.9%
2k, 130.92GB
56.14%
6k; 123.1GB
85.82%
6k, 272GB
Liang et al. (020)
83.16%
2k, 129.93GBB
41.09%
6k; 116.7GB
86.75%
6k, 271, 8GB
FedMix K=2
72.25%
2k, 261.4GB
48.29%
30k, 1130GB
84.02%
6k, 544.4GB
FedKix K=4
79.88%
2k, 522.5GB
48.55%
24k, 1972GB
86.11%
6k, 1089GB
FedMix K=10
80.00%
2k, 1304.1GB
52.59%
6k, 1130GB
87.21%
6k, 2724GB
from a different probability distribution over 8 possible rotation angles to rotate digit with. Each client's distribution is drawn from Dir(a: 1.0). At test time, each data-point is randomly rotated according to the client's distribution. Additionally, we create a dataset where instead of uniform sampling of labels, replicate the non-i.i.d label skew described for Cifarl0 above and combine it with the rotation non-i..d-ness We compare F FedMix x where q is conditioned on y or on the degrees of rotation for data- data-point against baselines. Figures andd 6 that biased FedAvg can improve in the presence of label skew but collapses to standard F FedAvg for rotation non-i.i.d-nes. FedMix on the other hand has advantage in both cases. Furthermore, in presence of both sources of non- ii..dness, FedMix x with ((2lg) can outperform biased FedAvg. In the presence of only rotation (Figure 5), conditioning F FeaMix on the true source of non-i.i. is advantageous.
Perm.
10
4.3 LABEL PERMUTATIONS
Apart from non-i.i.d-ness in p(y) and p(r), we can expect the mapping p(ylz) itself to be different between clients. We replicate the experimental setup of (Sattler et al.] 2019) with 20 clients for Cifar10, C' = 1.0 and E = 3, albeit with LeNet-5. Each client is randomly assigned one of 4 different permutations of its labels, determining the cluster assignment (2[ss) that FeaMix has to learn. Although FedMix is designed to dis- tinguish different regions of the input space, we show that it can perform user clustering. The gating function p(2[x, jr, s) learns to correctly identify, for each datapoint, the expert corresponding to the permutation of thus illustrates this effect (please note that the ordering of columns is arbitrary). It should be mentioned, that FedMix converges very quickly compared to the results in (Sattler this result was obtained after only 10 rounds. After this
2
4
6
Client
8
12
14
16
recovering the original clustering; Figure 7
18
0123
0123
al. 2019), point, FedMix essentially trains 4 independent models on the respective i.i.d subsets of the data that have the same label permutation. Appendix contains additional discussions for K A 4 and experimental details.
Figure 7: Ground truth and ((2(s) after 10 rounds. Greyscale rep- resents probabilities (white: 0; black: 1)
DISCUSSION
With FedMix we have introduced a federated learning algorithm that explicitly takes the non-i.i.d characteristics federated dataset into account. Clients can learn to align specialized experts on sub-regions of the data space and achieve higher performance compared to FedAvg, especially in situations where the source of the non-i.i.d nature is known. This assumption is very strong in real-world federated scenarios and we expect more flexible alignment process than global g==ly) to be the most interesting avenue for future research. We show encouraging results on experiments with non-i.i.d-nest in p(y (z)) (permutation) and p(z) (rotation). In the future, will explore ways to perform automatic selection of as well as automatic selection of architecture elements to share between experts, trading-off gradient alignment and communication budgets.
8

---

Under review as a conference paper at ICLR 2021
REFERENCES
Keith Bonawitz, Hubert Fichner, Wolfgang Grieskamp, Dzmitry Huba, Alex Ingerman, Vladimir Ivanov, Chloe Kiddon, Jakub Konecny, Stefano Mazzocchi, H Brendan McMahan, al. Towards federated learning at scale: System design. arXiv preprint Setiv:1902.010000 2019.
Christopher Briggs, Zhong Fan, and Peter Andras. Federated learning with hierarchical clustering of local updates to improve training on non-iid data. arXiv preprint arXiv:2004.11991, 2020.
Sebastian Caldas, Peter Wu, Tian Li, Jakub Konečný, H Brendan McMahan, Virginia Smith, and Ameet Talwalkar. Leaf: A benchmark for federated settings. arXiv preprint arXiv:18/2.0109)
2018.
Fei Chen, Mi Luo, Zhenhua Dong, Zhenguo Li, and Xiugiang He. Federated meta-learning with fast convergence and efficient communication. arXiv preprint arXiv:/802.07876, 2018.
Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. Adaptive personalized federated learning. arXiv preprint WeXiv:200.134611 2020.
Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning: A meta- learning approach. arXiv preprint arXiv:2002.07948, 2020.
Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Measuring the effects of non-identical data
distribution for federated visual classification. arXiv preprint arXiv:1909.000 06335, 2019.
Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. arXiv preprint arXiv:1502.03166, 2015.
Robert A Jacobs, Michael Jordan, Steven Nowlan, and Geoffrey E Hinton. Adaptive mixtures of local experts. Neural computation, 3(1):79-87, 1991.
Yihan Jiang, Jakub Konečný, Keith Rush, and Sreeram Kannan. Improving federated learning per- sonalization via model agnostic meta learning, arXiv preprint arXiv:1909, 12488, 2019.
Michael Jordan and Robert A Jacobs. Hierarchical mixtures of experts and the em algorithm.
Neural computation, 6(2):181-214, 1994.
Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Keith Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Advances and open problems in federated learning. arXiv preprint arXiv:1912.04977, 2019.
Mikhail Khodak, Maria-Florina F Balcan, and Ameet S Talwalkar. Adaptive gradient-based meta- learning methods. In Advances in Neural Information Processing Systems, pp. 5915-5926, 2019.
Heejae Kim, Taewoo Kim, and Chan-Hyun Youn. On federated learning of deep networks from non-(iid) data: Parameter divergence and the effects of hyperparametric methods, 2020. URL wwtp:/nweeeettteeeeeeeeeeeeeeeee
ht
net forum? SesrummmdddSJeAJJS--.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Alex Krizhevsky et al. Learning multiple layers of features from tiny images. 2009.
Viraj Kulkarni, Milind Kulkamn, and Aniruddha Pant. Survey of personalization techniques for federated learning. arXiv preprint arXiv:2003.0866 2020.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEFE, 86(11):228 36(11):2278-2344 1998.
Yann LeCun, Corinna Cortes, and CJ Burges. Mnist handwritten digit database. 2010.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization heterogeneous networks. arXiv preprint arXiv: 18/2.0//7 2018.
Tian Li, Maziar Sanjabi, and Virginia Smith. Fair resource allocation in federated learning. arXiv preprint arXiv: /905. 10497, 2019.

---

Under review as a conference paper at ICLR 2021
Paul Pu Liang, Terrance Liu, Liu Ziyin, Ruslan Salakhutdinov, and Louis- Phiis-Philippe Morency. Think locally, act globally: Federated learning with local and global representations. arXiv preprint arXfv:2001.01523, 2020.
Yishay Mansour, Mehryar Mohri, Jae Ro, and Ananda Theertha Suresh. Three approaches for personalization with applications to federated learning. arXiv preprini arXiv:2002.10611 2020.
H Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, et al. Communication-effcieeett learning deep networks from decentralized data. arXiv preprint arXiv:1602.05622, 2016.
Mehryar Mohri, Gary Sivek, and Ananda Theertha Suresh. preprint arXiv: 1902:00990000 2019.
Agnostic federated learning. arXiv
Eric Nalisnick, Akihiro Matsukawa, Yee Whye Teh, Dilan Gorur, and Balaji Lakshminarayanan Do deep generative models know what they don't 't know? arXiv preprint arXiv:1810.009136, 2018.
Alex Nichol, Joshua Achiam, and John Schulman. On first-order meta-learning algorithms. arXiv preprint 1803 02099,, 2018
Chris Pal, Charles Sutton, and Andrew McCallum. Fast inference and learning with sparse belief' propagation. Advances in Neural Information Processing Systems (NIPS), 2005.
Sashank Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and H Brendan McMahan. Adaptive federated optimization. arXiv preprint arXiv:2003.
00295, 002202 2020.
Felix Sattler, Klaus-Robert Müller, and Wojeiech Samek. Clustered federated learning: Model-agnostic distributed multi-task optimization under privacy constraints. arXiv preprint 1rXiv:1910.01991, 2019.
Jasper Snock, Yaniv Ovadia, Emily Fertig, Balaji Lakshminarayanan, Scbastian Nowozin, Scul- ley, Joshua Dillon, Jie Ren, and Zachary Nado. Can you trust your model's uncertainty? eval- uating predictive uncertainty under dataset shift. In Advances in Neural Information Processing Systems, pp. 13969-19910 13980, 2019.
Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European Conference on Computer Vision (ECCV), pp. 3-19, 2018.
10

---

Under review as a conference paper at ICLR 2021
A EXPERIMENTAL SETUP
For Cifar10, we replicate the federated data split of Hsu et al. (2019). The dataset split across 100 clients, whose data-points drawn according their label from Dir(a 1.0) distribution without replacement. For the base model, we use a LeNet-5 architecture (eCunn et al.) 1998). We use a SGD optimizer with learning rate of 0.05 and batch size B 64 locally and the Adam (Kingma & Ba (2114) optimizer with its default hyperparameters at the server by interpreting the difference of the local from the global model 10 clients without replacement on each round (but with replacement across rounds) and train for E local epochs. For Ci we replicate the data split of Reddi et al. (2020). The dataset is split into 500 clients by using hierarchical model over the coarse and fine labels, with the same hyperparameters as the ones provided by Reddi et al. (2020). The other hyperparameters are the same 10 with the exception of the batch size, where architecture, where we use a ResNet-20 with group normalization (Wu He| 2018) layers instead of batch normalization (loffe & Szegedy 2013). We augment the data by random cropping from pixel padded image and horizontal flipping. Finally, for the Femnist dataset, we similarly followed the setup of Reddi et al. (2020) with the same LeNet-5 architecture and hyperparameters of Cifar10 with the exception of the batch size where we used = 20. h4 (x) is defined as the input to an expert's output layer.
gradient (Reddi et al. 2020). We sample
B 20, well as the
B PERMUTATION EXPERIMENTS
For the experiments in Section we train FedMix with K 4 and base LeNet-5 model. The optimization hyperparameters are chosen identically to what is described in A except for the optimization of (=(ss). The local objective in this case is
K
Lg =
1g (2|s) plu,,, p( - Bq (2|s) log Or (=]s)]
(10)
In order to speed up convergence, we found it beneficial to not perform gradient ascent and instead solve it directly using Lagrange multiplies, i.e., Vn,ALo(Ds, w, 9x) A Th 1) 0.
The solution for म = q(2 = k|s) at client takes the form of
Th II p(u+.,, = k(x,,)) Nin /
N p(y,,, k(o.,,)) :) Ni0
ΣI
(11)
For these permutation experiments, FecMix is tasked recover cluster-assinmmen,, as opposed to partitioning of the input space. We have found it unnecessary to change & from its default value of 1. For an approximation of Ta& with a mini-batch of size M, we approximate
N.
K
log n
log togp(U..i,22 k(x,,)) - log
exp
log p(us,i, = 2 k(xx,,))
(12)
N₂
N
T=1
==1
i=1
1
logp(ys,,, klar,,)) log
exp
logplus.s, k(
(13)
M
M
and perform dampening nf a at1--1 a) R. We found our initial pick of a 0.75 to perform well.
B. MODIFYING K
In Section 4.3 we choose K = a-priori the number of experts that fits the data generating process with clusters. Figure we plot q(>l)) (y) in case we miss-specify K and choose and K 5. In this plot, we can consider K 4 as the ground truth for the correct cluster assignment. Please note that ordering of the columns is arbitrary and can be different across the three experiments. Comparing K with K same cluster, while still correctly disentangling all clients. For K 3, FedMix correctly identifies two clusters while being forced to mix together the other clusters into expert.
that F FedMix assigns two experts to the
11

---

Under review as a conference paper at ICLR 2021
K 3
K 4
K
0
4
6
Client
10
12
14
16
18
012
0123
01234
Figure 8: ((2[s) convergence for different number of experts K. Greyscale corresponds to prob-
abilities (white: 0; black: 1)
C PRIVACY IMPLICATIONS
Privacy is one of the key motivations for research and deployment of Federated Learning. Even though privacy is not focus of this paper, briefly discuss some implications of making explicit use of ((y) () FedMix x. The update rule described equation requires access to the marginal {<]ss) Ey p(yl ( (2(v)) the server. the same time, the server has access the parameters that were used in computing p(z(s) before being sent to the server. Therefore, in principle, could solve g==|s) = Σy (1
with respect to p(y)s) and thus obtain the marginal label distribution at the client. In practice this is not as straightforward to do as (a) the probability matrix (2(g) not always invertible and (b) solutions that use the pseudo-inverse, empirically, are not very accurate in capturing the entire distribution. With the additional constraints that the marginal needs to sum to one contains only positive elements and that N. p(y]s) E E, in some cases, a reconstruction can become possible. As the number of classes exceeds the number of experts, this reconstruction becomes more unlikely. We leave a thorough characterization of these properties to future work.
Irrespective of FedMix, we want to shine light on the possibility to reconstruct p(y(s) at the server- side in standard FedAvg. Assume randomly initialized model being sent to client where the client performs a single full data set update step on the output layer's 's bias vector by. Assuming a softmax cross-entropy loss Ls, the average gradient with respect the k-th entry by takes the form of
0L. Obg
N. Σ
1liik = Vi,,mee] Ti.
(14)
N
where 1 is the indicator function and Ti. is the softmax probability of class k of datapoint i. With a randomly initialized model, these softmax probabilities can be assumed be uniform, leading to an average gradient of
0L. 0b
No.
1 No
1
p(yls)
(15)
No
0L2 to the where Nº is the number of classes. Upon sending the updated bias vector by a abn
server, can easily reconstruct the marginal label distribution.
Figure 9 shows, for every client in our Cifar10 setup (Appendix A). in red the true marginal p(y]s) and in blue the reconstructed marginal based on (the same) randomly initialized model being sent to each client. Clearly, have high congruence. In Figure I0, we investigate the more realistic with mini-bater stochastic gradient descent at the client level. order to avoid reconstructing the multi-step update, simply normalize the difference (b b.) and interpret it
setup of
12

---

Under review as a conference paper at ICLR 2021
marginal label distribution. We that multiple updates (on average: 8) reduce the congruence between the true and reconstructed marginal, however the information leakage is still remarkable.
Figure Histogram representation pwls) as well its server-side reconstruction for every client for the Cifar10 setup described in Appendix A The x-axis per sub-plot enumerates the 10 classes. For every class c, the left bar in red represents p(y = cls) and the right bar blue represents its reconstruction. Each client performed single full data-set update step.
13

---

Under review as a conference paper at ICLR 2021
Histogram representation of p(yls) well as its server-side reconstruction for every client for the Cifarl0 setup described Appendix A The x-axis per sub-plot enumerates the 10 classes. For every class c, the left bar in red represents ply cls) and the right bar in blue represents its reconstruction. Each client performed multiple mini- update steps (on average 8).
Figure 10:
D GRADIENT DIVERGENCE
We aim to track the divergence of updates for a subset 5' of shards at the server time step t for FedMix and FedAvg. Therefore we define a metric inspired by Kim al. (2020); /Sattler et al.) 2019) to define divergence of gradients Ai. w ++1 for some subset wa of the parameters of expert k as
5|= k) 0.5 (1
D A
GD(A)) Σ p(s 0| k))(s
(16)
||
||
||
For FedAvg, the above metrie collapses to
A A
GD(A)) Σ
Σ p(s -) 0.5
(17)
|| 1- -A111
i=1
In Figure in the main text, we plot the sum of GD(A)) across all parameters we (i.e., convolutional kernels, weights and biases) of the LeNet-5 experts in comparison for FedAvg.
14

---

Under review as a conference paper at ICLR 2021
E ABLATION STUDIES
We investigate the characteristics of FedMix by varying the number of experts K for Cifar10, as learning curves
TaaeessTTaannde well as by selecting different levels 1) for pruning experts. Figures] I TTc [nC of several values K. We can see that a higher number of experts has the potential to achieve higher accuracies the cost of more required communication. With increasing K, the modeling task for single expert becomes progressively easier and the challenge is moved towards identifying the right expert for given data-point. This could subsequently lead overfitting in the federated setting not for the shared feature extractor h(x) (Eq. equation 9). In case we choose K to be equal to the number of classes, we observe that FedMix assigns one cluster each class and the p(z|x, s) is reminiscent of [iang al. (2020). In case we further increase K, we observe marginally improved performance by multiple experts being assigned to the same label.
Figures ll and 11d depict the influence of different pruning levels 1 the communication effi ciency and final performance FedMix with 4. We can see that higher levels of pruning do reduce communication significantly, however they cause some loss of performance.
0.8
0.8
Validation accuracy
Validation accuracy
0.7
FedAvg
0.6
Biased FedAug
FedMix 2
0.6
FedMix
0.4
FedMix 6
0.5
No pruning
FedMix
0.8
FedMix
FedMix =
0.99
0.2
FedMix
0.995
500
1000
1500
100
200
300
communicated
GB communicated
(a)
(b)
Validation accuracy
0.7
Validation accuracy
0.7
0.6
FedAvg
Blased FedAvg
FedMixx =
FedMix K =
0.4
FedMixx
No pruning
FedMix =
n 0.8
FedMix 10
0.9
0.2
FedMix
0.3
n 0.99
FedMix K 14
0.995
0.2
1000
200 GB communicated
GB communicated
(d)
Ablation studies the effect of K the average local accuracy (a) and shard accuracy (c) as well as on the influence on local (b) and new shard (d) accuracy of the threshold for pruning experts before communication for K 4.
Figure 11:
F NEW SHARD INFERENCE
In the main text, we discussed performance of the individual algorithms when evaluating new data on a shard that took part in training and therefore has access to trained local model parameters. This is the case for all considered methods except FedAvg, for which there exist no localized parameters. For inference on a new shard, we use the sum of all local validation sets as a proxy. [iang et al. (2020) propose ensemble the representations of the local feature extractors across clients before evaluating the global part of the network. We find this approach to work quite poorly in practice. Presumably pretraining the local feature extractors with FedAvg ameliorates this behaviour. In
15

---

Under review as a conference paper at ICLR 2021
biased FedAvg, ensemble the individual local biases across clients receive single global bias. For FedMix, x, we marginalize the local gating predictions to achieve a global gating prediction p(z|x*) = Σ Pe, (z|x", s))(ss)))pss) Predictions can then be made by marginalizing across experts using this global gating function: p(y)x") Σ Pw, (y x"; =){(=(x*).
Interestingly, Cifar10 and C0far100 display different behaviours when comparing FedMix to the baselines. For Cifarl0, biased FedAvg performs much worse compared to FedAvg while FedMix performs well. For Cifar100, biased FedAvg is indistinguishable from FedAvg, outperforming FedMix. We leave the investigation of these discrepeancies to future work.
0.5
Validation accuracy
Validation accuracy
0.6
FedAvg
Blased FedAvg
0.3
FedAvg
0.4
FedAvg, LeNet-5(61)
Biased FedAvg
FedAvg. LeNet-5(113)
FedMix K
FedMix K
FedMix n 0.99
FedMix K
FedMix K 10
0.2
Biased FedMix K K-4
0.1
FedMix K 10, 0.9
Local/Global
Local/Global
0.0
200 GB communicated
300
200
400
600
800
1000
GB communicated
(a) Cifar10
(b) Cifar100
0.4
Validation accuracy
Validation accuracy
0.6
0.3
FedAvg Biased FedAvg
FedAvg
Blased FedAvg
0.4
FedAvg, (66)
FedAvg, LeNet-5(1133
FedMix K
FedMix K = 4
FedMix K= 2
Local/Giobal
FedMix K 4
0.1
0.2
Biased FedMix K
Local/Globall
0.0
0
500
1000
1500
2000
0
2000
4000
6000
Communication rounds
Communication rounds
(c) Cifar10
(d) Cifar100
client (y-axis) a function of the amount of GB communicated (top row) and as a function of communication rounds (bottom row). Cifar 10 models are trained on the standard 45k training split. x-axes have been truncated for improved visibility. Best viewed in color.
Figure 12: Accuracy
As alternative marginalizing the local gate predictions using p(s), investigating the graphical model Figure reveals the possibility for marginalization with p(s(x*):
5 p(2|x') Σ Po. (2)x"
s)p(s(x"), (e)x)s)sss))p((xx
p(s(x") & p(x"|s))p(ps
(18)
It is for this third evaluation case that training local generative models p(x(s) for each client becomes interesting, as they allow to compute the responsibilities p(s]x*) at test time. In practice, however, the success of this approach depends heavily on the correciness of p(s(x*), which in turn depends on the ability the local generative models p(x[s) 3) to assign high probability to data that resembles D, and low probability to out-of-distrrbution data. Training and calibrating generative models for this task is in itself an active area of research (Nalisnick et al.) 2018) and investigating how clients in federated setting might exchange information facilitate this process is not yet explored. We
16

---

Under review as a conference paper at ICLR 2021
therefore leave thorough evaluation of this to future work and only present here MNIST (LeCun et al. 2010 experiment.
We train FeaMix with K 4 and experts of two hidden layer ReLU MLP with 200 hidden units on MNIST. We split the dataset into S 100 clients according to the procedure described in Liang et al. (2020). FedMix x achieves 97.7% average validation accuracy compared to 97.0% with FedAvg after 600 communication steps. Independently for each client, we train a small vari- ational autoencoder with 32-dimensional latent space using two-layer MLP with 512 and 256 hidden units respectively as encoder and mirrored decoder structure. We optimize the VAEs using Adam with standard hyper parameters, B 10 and perform early stopping on the local validation sets after no improvement for three epochs. After training, each client communicates their VAE to the server, where we evaluate pl p(s[x*) to marginalize over the local gating functions according x", ,( With this procedure, FedMix achieves 95.9% test set ac-
to p(z[x) = - curacy, compared to FedAvg with 96.9%. Marginalizing with p(s) instead of p(s[x*) achieves 96.63% showing the limitation of the approach in that any error in p(s[x*) propagates into the ex- pert assignment. Reliable -fistii-iittttti detection capabilities the individual estimators for p(x(s) are therefore necessary.
G LOCAL FINE-TUNING
In the main text we have compared algorithms by evaluating local data-sets on the perver-provided global parameters. For biased FedAvg, Local/Global and FedMix those parameters complete with the client-specific local parameters. In Appendix FF] we have discussed the perfor- mance of FedMix when evaluating on new shard. Here, we discuss the evaluation when using locally fine-tuned models. Figure 13 shows the learning curves of Cifarl0 when evaluating the validation set on parameters just before communication fine-tuning for E = 1 epochs on the local data-set. This form of model personalization is an option for all federated learning methods. Due to the relatively small size of the local data-sets, this form of personalization can cause significant levels of overfitting, as is evident from Figure [3] Since regularization with respect to this form of personalization adds another level of complexity to the experimental design and hyperparameter tuning, we leave thorough evaluation to future work.
made
the server. This corresponds local
0.9
0.9
Validation accuracy
Validation accuracy
0.8
0.8
0.7
FedAvg
0.7
FedAvg
Blased FedAvg
Blased FedAvg
FedAvg, LeNet-5(61)
FedAvg, LeNet-5(61)
0.6
FedAvg, LeNet-5(1133
0.6
FedAvg, LeNet-5(113)
FedMix K-2
FedMix
FedMix K= 4
0.5
FedMix K = 4
Biased FedMix 4
Biased FedMix K
Local/Global
Local/Global
0.4
0.4
100
200
300
500 Communication rounds
1000
1500
2000
GB communicated
Figure 13: Average accuracy of models evaluated after E = 1 local epochs of fine-tuning
17

---

Under review as a conference paper at ICLR 2021
H SPECIALIZATION
In Figure 14a we how (=(g))) changes over time. The entropy of qo(zly c), ie. the weight with which each expert is assigned to specific label с decreases over time until each label as- The original MoE formulation, however, quickly
signed exactly one expert with probability collapses to single expert, as can be see in Figure [4b] In the federated learning seenario, this corresponds to performing standard FecAvg on the single surviving expert.
MoE
t= 200
== 1000
2000
4000
9
Cifar10 label
Cifar10 label
1 3
1 2 3
0 3
0
1 2
3
0
0123
(b) Visualization of Σ p(2(z,)
for original mixture of experts formulation (equation 3 with S 1) after the first epoch of training. The mixture collapses single expert.
(a) Visualization of (=lg) at different communication rounds for FedMix x with K = 4 on Cifar10. Greyscale corresponds to prohabili- ties; white corresponds to zero and black corresponds to one. Probabil- ities across experts (horizontally).
Figure 14: Expert specialisation
FEDMIX WITH EXPERT PRUNING
The FedMix algorithm (Algorithm [ can be extended to ignore the communication of experts between server and client if, for this given client, that particular expert is not updated during training. If qs(2 k(s) = Ep(pls) (2 = k y)] 0 for an expert A on client s, then parameters of that expert observe no gradient and the expert effectively pruned from the local library of available experts. Since the parameterization of (=(=)s) does not allow for exact values of zero, we introduce such that (2 = s) < T/K, we consider expert & to be pruned. Since the value ( k|s) during local optimization at the client is subject to change, it is possible that gø(2 = k(s) briefly lies above the threshold even though it had just been pruned away. Therefore the server considers expert k pruned for given client s' only if go(2 k|s s') < 0.9 n/K. Algorithm 2 details this approach. Note that we write p(z(s) = Ex-D, [Po, (z|x, s)), i.c. the true marginal, however in practice make of the cheap compute marginal approximation (=(s)) EEy-D, (g v)).
18

---

Under review as a conference paper at ICLR 2021
Algorithm The FedMix algorithm. 3 are the client and server learning rates respectively, 7) E [0, 1] the pruning threshold
function SERVER SIDE
Initialize & and K vectors W [W1, W wk]
for round in I....T do
S" f- random subset of the clients
Initialize A = 0, A =0
Initialize p(z(s) = Kvs
for s in S" do
W' + WR |p(z = k|s) 2 0.9 n
W Φ p(z|s) CLIENT SIDE(s, Φ, W))
Store p(2(s)
end for
p(s|z) + p(zls)p(s)) Eses ((2)))p(8)
for in S' do
A W4 = p(s|z = k)(w w Vk
A +
Ng. (0=-1 Φ5:)
N₂
end for A
V (E c)ply = -))
(+1
ADAM(A ! :9 9)
1:K
ADAM(A (2,)) 3)
end for end function
function CLIENT SIDE(s, Φb, W)
Get local parameters 0g for epoch in 1,..., E do
for batch E B do
% (2)ss) (=== k|s) 0f go(= k(s) n/ K
Ey-DD, [9>(2
(=(s)) + (z|s) Σx 1=======a k(s)
(s(y))) - (o(2)))
d/(2 = kly)= go = k|s) < n/K K =
d6 + (6(()))// Σn do = kly)
L, - (elu))) flg Pw: (ybll Xb: +)}+
E
(zlvs) [log po, (2 x2. s)]
++ = a Ls
W V L.
0- = a Vo. Ls
end for
end for
{(2s)) + E Ey [26(2]
W' t [ = k(s) 2 n/K]
(=(s)) < (2(8))
d= kls) 0 g(= = k(s) n/ K
(2(s)) < (z(ss)) d(= k(s)
Renormalization
return (z)))((2())
end function
19


Reviews:
[{"review_id": "u6K_L3Z-e0D", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "-- Summary --\nThe paper proposes a method for federated learning of a mixture of experts model (FedMix). The approach allows training an ensemble of models each of which specializes to a subset of clients with similar data characteristics. The authors argue that this way of training an ensemble reduces the gradient divergence/interference, improves the overall performance, and sometimes reduces the communication overhead. The new method is evaluated a few federated image classification datasets.\n\n\n-- Overall evaluation --\nI find the idea of training a model ensemble instead of a single global model extremely appealing in the context of federated learning. The fact that the proposed approach allows models in the mixture (i.e., experts) to automatically specialize to different subsets of clients throughout training without having to run any clustering on the server makes the approach particularly appealing. Having said that, I also believe that the paper in the current form has many weaknesses, including:\n- a significant lack of clarity throughout section 2 (see comments and questions below),\n- in terms of methods, fairly ad-hoc changes introduced into the evidence lower-bound objective given in eq. 5 (removal of the entropy term and addition of a different regularize, which is justified in a very handwavy way),\n- weak experimental results, which indicate that FedMix is, in fact, worse than the baselines both in terms of performance and the communication cost (unless I'm misreading Table 1 and Figure 4).\n\nI believe that the paper has interesting and novel ideas, but falls short on both presenting them as well as getting them to work empirically. I would not recommend accepting the paper in the current form.\n\n\n-- Comments and questions --\n\n- Using MoE in federated learning makes a lot of sense intuitively. The authors further specialize MoE to FL by introducing a local gating network and then tying everything together through a global q distribution that approximates the posterior. These design choices are justified by stating that running federated optimization on the objective given in eq. 3 does not work well. I wonder, however, why not use a shared gating network instead of having to approximate the objective?\n\n- I find the way lower bound in eqs. 5-6 is introduced and then the entropy term is dropped and substituted with some regularize (see 2nd paragraph on page 4) very ad-hoc. Typically, one would use a graphical model such as the one given in Figure 2 to derive a loss function in a principled way, without additional \"hacks.\" Why use the graphical model formalism in the first place, if the final loss is being designed using some additional intuitive heuristics? Relatedly, perhaps there's a way to change the graphical model somehow and show that the final loss corresponds to the evidence lower bound in a slightly different probabilistic model? (e.g., with a different prior)\n\n- Since $q_\\phi(z \\mid y)$ is a distribution over the experts in the mixture conditional on the class of the data point, what happens if K > # classes? It seems like q will select a separate model for each output class, so each expert will specialize to a single class only. How would that affect performance? Is K < # classes a requirement?\n\n\n-- Experiments --\n\n- If I'm not misreading the results in Table 1, FedMix adds a drastic communication overhead (due to multiple experts in the mixture) and does not perform as well as biased FedAvg on the CIFAR 10/100 datasets. Results for two baselines for EMNIST are missing.\n\n- Similarly, results presented on Figure 4 for CIFAR datasets demonstrate that the baselines are better.\n\n- For FedAvg, was the accuracy computed for the global model or for the global model after fine-tuning locally? Since FedMix targets to improve personalized FL, it should be compared with FedAvg with a simple fine-tuning-based personalization.\n\n\n-- Minor points --\n\n- Figure 1: are the bold curves some kind of running averages? Also, if the figure is given in the main text, it's strange that the definition of gradient divergence or alignment is provided in the appendix. I would recommend giving it somewhere in the main text.\n- Figure 3: which dataset and task are these results for?\n- Section 4, paragraph 1: I believe the EMNIST dataset has 3400 users if the version from tensorflow federated is used.\n- Very minor: In Eq. 3, the notation for the sum over S is strange; maybe change to $\\sum_{s=1}^S$?", "extra_scores": {}, "extra_text": {}}, {"review_id": "oCNPe4bmfVK", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors propose using an ensemble of experts, but where the mixing proportions depend on both the data and the client (providing the personalization in the method). The proposed method is very interesting, but there are several issues which need to be addressed.\n\nThe method requires significantly higher communication cost to train, since multiple versions of the model parameters need to be communicated, as well as parameters for other models used for the mixing. The authors propose suggestions for reducing the number of models that need to be communicated, but they rely on very special behaviour of the different models (namely that they separate well so that clients do not actually need all of them). \n\nThe clients need to transmit mixing parameters $q(z|s)$ as well as the $\\phi$ which generate them. Given access to $N_s$ and the discrete nature of $y$, this would allow the server to reconstruct the label distribution of each user's data, which is a huge reduction in the privacy achieved by federated learning. \n\nThe algorithm is quite complicated, and appears fairly fragile, as is demonstrated by the different functions for $h_k(x)$ required for just 3 data sets. Selection of tuning parameters in federated learning is a challenging and task, and having one which is a function makes the method quite challenging to use in practice.\n\nExperimental results are not very compelling, with a simple personalized bias term outperforming the proposed model in terms of both number of communication rounds and total volume of data transmitted, and does not produce the privacy violation the proposed method produces. Although the simple personalized bias does not generalize beyond label distribution, there are many many similar personalization proposals for federated learning which do, and it is not clear whether the proposed method is competitive with those. Of course comparison with all existing methods is not feasible, but given the method must be entirely judged by empirical results, comparison with at least some should be done.\n\nOverall it is not clear that the extra costs induced by the method are justified by the improvement in convergence speed, especially compared to more simple personalization methods which do not incur any of these extra costs (either in communication or privacy). ", "extra_scores": {}, "extra_text": {}}, {"review_id": "VMYQO5L2wZo", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Strength:\n* The idea of training an ensemble of specialized models in FL is interesting, although not novel.\n*  The evaluation is solid. I would be more impressed if the authors could also evaluate some SOTA language models.\n\nWeakness:\n*  The technical contribution is limited and the novelty is not sufficient. I am not against \"A+B\" research if it does bring drastic practical improvement. However, the proposed method is worse than some prior works on CIFAR 10 even using ten-time bandwidth (correct me if I interpret the results wrongly).\n*  There are many other works with similar technical consideration although their intention might be different. Please check the line of works about personalized FL and compare with them if possible.\n*  The writing is not clear given that the idea is actually simple.", "extra_scores": {}, "extra_text": {}}, {"review_id": "7oDIoYkhxl", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper proposes a novel algorithm, which is a federated form on mixture of experts, called Federated Mixture of Experts (FedMix). In FedMix, an ensemble of specialized models is trained instead of a single global model. This strikes a compromise between training a single global model and one model per client. A gating mechanism is employed, to choose an expert model that is responsible for a given data point, thus aligning the gradient updates across experts and alleviating the consequences of non-i.i.d. data.\n\nPros:\n- the paper is well written and easy to follow\n- it tackles an important problem, namely how to cope with non-iid data in FL\n- experiments are performed on different datasets\n\nCons:\n- A major drawback of the paper is the lack of  in-depth experimental evaluation and theoretical insights. On Cifar-10 and Cifar-100 the proposed method performs worse than related state-of-the-art methods. On the rotated MNIST experiments in 4.2 it performs better, however, this section is very short and the insights from the presented results are limited.\n\n- The STC approach presented in \"Robust and Communication-Efficient Federated Learning from Non-IID Data\" performs even better than biased FedAve on Cifar-10. It may well be that it also performs well in the rotated MNIST experiments. Anyways I am not convinced that the proposed method is superior to state-of-the-art techniques in non-iid settings (on Cifar-10 and -100 it is clearly not).\n\n- The paper does not provide any new theoretical insights, it only reports empirical results. For instance, can you give guarantees that the proposed approach will result in K specialised models if the client data was generated from K different distributions (see  proof in Sattler et al. 2020).\n\n- The authors mention the conceptual similarity of the proposed approach to Clustered Federated Learning (Sattler et al. 2020, Briggs et al., 2020), that allow to jointly train specialised models. Why do you not compare the proposed approach to these methods? I assume that Clustered FL will solve the rotated MNIST problem (although it may be worse in terms of communication-efficiency). More insightful experiments would be great.\n\nOverall, I am not convinced by the proposed method and be the reported results.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents an approach, FedMix, for federated learning using mixture of experts (MoE). The basic idea is to learn an ensemble of models and user-specific combination weights (mixing proportions).

The reviewers appreciated the MoE formulation for federated learning. However, there were multiple concerns from the reviewers, which include lack of clarity (regarding the variational lower bound that is being used), significant communication cost and privacy concerns (the server can infer the users' label distributions), weak experimental results, and lack of any theoretical support (which isn't that big an issue if the paper were stronger on other aspects). The author feedback was considered and the reviewers engaged in some discussions (also with the authors). In the end, however, the reviewer were still not convinced that the paper is ready to be published in its current state. Based on my own reading of the paper, the reviews, and the author response, I agree with this assessment.

The authors are advised to take into account the feedback from the reviewers and resubmit to another venue.
Decision: Reject