Federated Mixture of Experts

1 INTRODUCTION An ever-increasing amount of devices are being connected to 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 generated by this “ federation ” of devices ( McMahan et al. , 2016 ) . Not only does the number of sensing devices increase , but also their processing power is increasing continuously to the point that it becomes viable to perform inference and training of machine learning 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 . Several challenges arise in the federated scenario . Federated devices are generally resource-constrained , both in their computational capacity as well as in communication bandwidth and latency . In a practical example , a smartphone has limited heat dissipation capacity and must communicate via Wi-Fi . From a global perspective , devices ’ processing power and network connection can be highly heterogeneous across geographical regions and socio-economical status of device owners , causing practical issues ( Bonawitz et al. , 2019 ) and raising questions of fairness in FL ( Li et al. , 2019 ; Mohri et al. , 2019 ) . 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-federated 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 in Figure 1 ) , 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 an ensemble of specialized models instead of a single global model . In FedMix , expert models are learning to specialize in regions of the input space such that , for a given expert , each client ’ s progress on that expert is aligned . FedMix allows each client to learn which experts are relevant for its shard and we show how it can be extended for inference on a previously unseen client . FedMix shows competitive performance against the established standard in FL , FedAvg ( McMahan et al. , 2016 ; Deng et al. , 2020 ) across a range of visual classification tasks . Code will be released upon publication . 2 FEDERATED MIXTURE OF EXPERTS . Federated learning ( McMahan et al. , 2016 ) deals with the problem of learning a server model with parameters w , e.g. , a neural network , from a datasetD = { ( x1 , y1 ) , . . . , ( xN , yN ) } ofN datapoints that is distributed across S shards , i.e. , D = D1 ∪ · · · ∪ DS , without accessing the shard specific datasets directly . By defining a loss function Ls ( Ds ; w ) per shard , the total risk can be written as arg min w S∑ s=1 Ns N Ls ( Ds ; w ) , Ls ( Ds ; w ) : = 1 Ns Ns∑ i=1 L ( Dsi ; w ) . ( 1 ) It is easy to see that this objective corresponds to empirical risk minimization over the joint datasetD with a loss L ( · ) for each datapoint . In federated learning one is interested in reducing the communication costs ; for this reason McMahan et al . ( 2016 ) propose to do multiple gradient updates for w in the inner optimization objective for each shard s , thus obtaining “ local ” models with parameters ws . These multiple gradient updates are denoted as “ local epochs ” , i.e. , amount of passes through the entire local dataset , with an abbreviation of E. Each of the shards then communicates the local model ws to the server and the server updates the global model at “ round ” t by averaging the parameters of the local models wt = ∑ s Ns N w t s. This constitutes federated averaging ( FedAvg ) ( McMahan 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.i.d . distributed across the shards S , that is p ( D|si ) 6= p ( D|sj ) for i 6= j . On the one hand , this can make learning a single global model from all of the data with the classical FedAvg problematic . On the other hand , there is one extreme that does not suffer from this issue ; learning S individual models , i.e. , only optimizing ws on Ds . Although these individual models by definition do not suffer from non-i.i.d data , clearly we should aim to do better and exchange meaningful information between clients to learn more robust and expressive models . 2.1 THE FEDMIX ALGORITHM With FedMix , we propose to strike a balance between the two aforementioned extremes ; learning a single global model and learning S individual models . For this reason , we revisit an old model formulation , the Mixture of Experts ( MoE ) . The classical formulation of a MoE model ( Jacobs et al. , 1991 ; Jordan & Jacobs , 1994 ) contains a set ofK experts and a gating mechanism that is responsible for choosing an expert for a given data-point . A MoE model for a data point ( x , y ) can generally be described by pw1 : K , θ ( y|x ) = K∑ z=1 pwz ( y|x , z ) pθ ( z|x ) , ( 2 ) where z is a categorical variable that denotes the expert , wk are the parameters of expert k and θ are the parameters of the selection mechanism . The MoE was proposed as a 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 on a 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 to their respective experts . In this work , we show that , in the federated scenario , sub-dividing the input region through a MoE can alleviate the consequences of non-i.i.d data by aligning gradient updates across experts ( Figure 1 ) . In Federated Mixture of Experts ( FedMix ) we enrich this model by conditioning the gating mechanism on the shard assignment s. Whatever characteristics make shard s different from other shards can manifest in learning a different , localized gating mechanism that does not need to be communicated to the server . In choosingK = 1 , FedMix recovers the standard setting of federated averaging . K = S in combination with fixing p ( z = s|x , s ) = 1 recovers S independent models . From a global perspective , we are interested in maximizing the following single objective : S∑ s=1 Ns∑ i=1 log pw1 : K , θs ( ys , i|xs , i , s ) = S∑ s=1 Ns∑ i=1 log [ K∑ z=1 pθs ( z|xs , i , s ) pwz ( ys , i|xs , i , z ) ] ( 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 ( x|s ) . We will briefly touch upon the role of learning generative models in Appendix F but focus on the discriminative part of the model , i.e. , the MoE , in 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 a single expert , thus obtaining FedAvg , and specialization of the experts . Instead , we propose to form a variational lower-bound on Eq . 3 with a global variational approximation qφ ( z| . . . ) to the true posterior p ( z|x , y , s ) with parameters φ . At test time , p ( y|x∗ , s ) = ∑K k=1 p ( y|x∗ , z ) p ( z|x∗ , s ) can be readily evaluated without requiring q . This allows us to condition qφ ( z| . . . ) on any available side-information at training time that might result in better specialization in the non-i.i.d federated scenario . In this paper we mainly consider classification tasks whose non-i.i.d nature predominantly stems from the non-i.i.d distribution of labels y . Other or additional known sources of misalignment could be included to further improve this approximation , such as a manufacturer-id for a medical device in a medical scenario , a geographic identifier , or general domain-specific information . We show one such additional example in 4.2 . The lower bound to be maximized in FedMix therefore is as follows : S∑ s=1 Ns∑ i=1 logpw1 : K , θs ( ys , i|xs , i , s ) ( 4 ) ≥ S∑ s=1 Ns∑ i=1 K∑ z=1 qφ ( z|ys , i ) [ log pwz ( ys , i|xs , i , z ) pθs ( z|xs , i , s ) − log qφ ( z|ys , i ) ] ( 5 ) = S∑ s=1 Ns∑ i=1 ( K∑ z=1 qφ ( z|ys , i ) [ log pwz ( ys , i|xs , i , z ) pθs ( z|xs , i , s ) ] ) +H ( qφ ( z|ys , i ) ) . ( 6 ) Figure 3 : The effect of specialisation ( without H ( q ) ) compared to an ensemble ( with H ( q ) ) and FedAvg on Cifar10 . The experimental setup is identical to what is described in Section 4 . Conditioning only on y allows us to efficiently parameterize the variational approximation , incurring only a small communication overhead . While it would be possible to condition qφ ( z|y ) on s , thus having localized approximations with parameters φs that do not need to be communicated , we found a global approximation to help align the gating mechanisms across shards . A global qφ ( z|y ) encourages shards that contain data with the same label to assign them to the same expert . Specialization of the experts is a key ingredient for FedMix to be successful ; with specialization , the gradients for each expert become aligned across shards ( see Figure 1 ) , the hold-out accuracy improves ( 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 z generally leads to better and more personalized models . By removing the entropy term from equation 4 , qφ ( z|y ) and therefore pθs ( z|x , s ) are encouraged to concentrate and select only one expert for a given data point . In the extreme case where a 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 4 that communicating and evaluating experts based on thresholding the aggregate qφ ( z|s ) = Ey∼Ds [ qφ ( z|y ) ] can reduce communication and computation overhead . Figure 3 compares FedMix with and without the entropy term to standard FedAvg as a function of communication steps . With the entropy term , FedMix develops no expert specialization and collapses to an ensemble of K = 4 models . One drawback of the heavy specialization with MAP inference is that sometimes FedMix prematurely completely prunes experts , i.e. , pθs ( z = k|x , s ) ≈ 0 ∀x , s. This can be undesirable as we lose model capacity that can be used for better modeling the data . As qφ ( z|y ) is one of the main training signals of pθs ( z|x , s ) , we introduce the marginal entropy term in the server , H ( Ep ( y ) [ qφ ( z|y ) ] ) , as a regularizer . Notice that this leads to different training dynamics than locally optimizing the lower bound with the entropy included and we , empirically , found that it alleviates premature pruning , while still leading to specialized models . Figure 14a in Appendix H visualizes the development of qφ ( z|y ) over time from initially uniform to high specialisation for the experiment depicted in Figure 3 . Server Side Updates In a general federated learning algorithm , a central server selects a subset S′ ⊂ { 1 , . . . , S } of clients at time t and transmits the current estimate of the global parameters wt to them . These clients perform a series of mini-batch gradient updates with data from their shard Ds 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 ∆ts = w t − wt+1s as a single-step gradient update from client s , averages those gradients and applies an optimizer such as Adam ( Kingma & Ba , 2014 ) to receive wt+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 ∆t = S′∑ s=1 p ( s ) ·∆ts , p ( s ) = Ns NS′ . ( 7 ) In FedMix , we can speed up convergence by considering expert-specific updates ∆tk , s = w t k − wt+1k , s . If a client s pruned away expert k from its local gating mechanism , ∆ t k , s will be zero . We propose to normalize the effective magnitude of the resulting update ∆k by up-weighing the updates of all other clients that do consider expert k for their local mixture : ∆tk = S′∑ s=1 p ( s|z = k ) ·∆tk , s , p ( s|z = k ) ∝ p ( z = k|s ) p ( s ) , p ( s ) = Ns NS′ . ( 8 ) Computing p ( z|s ) = Ex∼Ds [ pθs ( z|s , x ) ] prior to sending updates to the server involves evaluating potentially large neural network models . Therefore we choose to approximate p ( z|s ) ≈ qφ ( z|s ) = Ey∼Ds [ qφ ( z|y ) ] , which involves just a single matrix multiplication . We discuss the implications of sending qφ ( z|s ) to privacy and how these fare relative to FedAvg in Appendix C. Pruning experts Maintaining the cheap-to-compute marginal posterior per shard offers an additional 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 qφ ( z|s ) does not surpass a threshold η/K . In order to still optimize a valid bound , we need to re-normalize qφ ( z|y ) 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 only those experts to the client that were updated by the client in the previous round . We empirically find that the entropy of q ( z|s ) decreases steadily and we prune away experts k with probability q ( z = k|s ) < η/K without significant drop in performance . We explore the consequences of pruning experts in the experiment section and in Appendix E. Algorithm 2 shows how FedMix can be enriched by pruning . Algorithm 1 The FedMix algorithm . α , β are the client and server learning rates respectively function SERVER SIDE Initialize φ and K vectors W = [ w1 , . . . , wK ] for round t in 1 , . . . T do S′ ← random subset of the clients Initialize ∆tW = 0 , ∆ t φ = 0 for s in S′ do Wts , φ t s , p ( z|s ) ← CLIENT SIDE ( s , φ , W ) end for p ( s|z ) ← p ( z|s ) p ( s ) / ∑ s∈S′ p ( z|s ) p ( s ) for s in S′ do ∆twk+ = p ( s|z = k ) ( w t−1 k −wts , k ) ∀k ∆tφ+ = Ns NS′ ( φt−1 − φts ) end for ∆tφ− = ∇φH ( ∑ cqφ ( z|y=c ) p ( y=c ) ) wt+11 : K ← ADAM ( ∆tw1 : K , β ) φt+1 ← ADAM ( ∆tφ , β ) end for end function function CLIENT SIDE ( s , φ , W ) Get local parameters θs for epoch e in 1 , . . . , E do for batch b ∈ B do Ls ← Eqφ ( z|yb ) [ log pwz ( yb|xb , z ) pθs ( z|xb , s ) ] φ+ = α∇φLs W+ = α∇WLs θs+ = α∇θsLs end for end for q ( z|s ) ← Ey∼Ds [ qφ ( z|y ) ] return w1 : K , φ , q ( z|s ) end function Designing robust gates In the federated scenario , Ns is often much smaller than N and especially small in relation to the complexity of the data we try to model . Any localized parameters therefore are prone to 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 . We can make use of the robustness of these experts ’ features for the gating mechanism by conditioning on them instead of training an entirely separate model for pθs ( z|x , s ) . Let us define hk ( x ) as intermediary features of expert k. Since not all experts might be used for a given shard and in order to scale with K , we average over the marginal posterior of the training set at that shard before applying a linear transformation to compute the input to the softmax gates : hs ( x ) = K∑ k=1 qφ ( z = k|s ) hk ( x ) pθs ( z|x , s ) = SM ( ATs hs ( x ) +bs ) ( 9 ) where θs = ( As , bs ) are local learnable parameters and SM represents the softmax function . Inference at test time We consider three variants for test-time evaluation of FedMix . In the first case , a client s that participated in training is presented with a new data point ( x∗ , s ) . Predictions can then be straightforwardly done by selecting the y that maximizes ∑K z=1 p ( y|x∗ , z ) p ( z|x∗ , s ) . In the second , more challenging , scenario a new client s∗ is introduced together with a new labelled local data set Ds∗ . Here we propose to instantiate and train the local gating mechanism by optimizing the parameters θs of pθs ( z|x , s∗ ) via MAP inference at the local objective . Afterwards , predictions can be made in a manner similar to the first case . Finally , we consider the case in which a 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 out-of-distribution data ( Snoek et al. , 2019 ) . We therefore propose to ensemble across local gating mechanisms to compute p ( z|x∗ ) = ∑S s=1 pθs ( z|x∗ , s ) p ( s|x∗ ) ; a method which works well in practice . In Appendix F we discuss results for new shard inference as well as a more principled approach which makes use of the graphical model formulation in Figure 2 .

Summary: The 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.