Federated Averaging as Expectation Maximization

1 INTRODUCTION . Smart devices have become ubiquitous in today ’ s world and are generating large amounts of potentially sensitive data . Traditionally , such data is transmitted and stored in a central location for training machine learning models . Such methods rightly raise privacy concerns and we seek the means for training powerful models , such as neural networks , without the need to transmit potentially sensitive data . To this end , Federated Learning ( FL ) ( McMahan et al. , 2016 ) has been proposed to train global machine learning models without the need for participating devices to transmit their data to the server . The Federated Averaging ( FedAvg ) ( McMahan et al. , 2016 ) algorithm communicates the parameters of the machine learning model instead of the data itself , which is a more private means of communication . The FedAvg algorithm was originally proposed through empirical observations . While it can be shown that it converges ( Li et al. , 2019 ) , its theoretical understanding in terms of the model assumptions as well as the underlying objective function is still not well understood . The first contribution of this work improves our understanding of FedAvg ; we show that FedAvg can be derived by applying the general Expectation-Maximization ( EM ) framework to a simple hierarchical model . This novel view has several interesting consequences : it sheds light on the algorithmic choices of FedAvg , bridges FedAvg with meta-learning , connects several extensions of FedAvg and provides fruitful ground for future extensions . Apart from theoretical grounding , the FL scenario poses several practical challenges , especially in the “ cross-device ” setting ( Kairouz et al. , 2019 ) that we consider in this work . In particular , communicating model updates over multiple rounds across a large amount of devices can incur significant communication costs . Communication via the public internet infrastructure and mobile networks is potentially slow and not for free . Equally important , training ( and inference ) takes place on-device and is therefore restricted by the edge-devices ’ hardware constraints on memory , speed and heat dissipation capabilities . Therefore , jointly addressing both of these issues is an important step towards building practical FL systems , as also discussed in Kairouz et al . ( 2019 ) . Through the novel EM view of FedAvg that we introduce , we develop our second contribution , FedSparse . FedSparse allows for learning sparse models at the client and server via a careful choice of priors within the hierarchical model . As a result , it tackles the aforementioned challenges , since it can simultaneously reduce the overall communication and computation at the client devices . Empirically , FedSparse provides better communication-accuracy trade-offs compared to both FedAvg as well as methods proposed for similar reasons ( Caldas et al. , 2018 ) . 2 FEDAVG THROUGH THE LENS OF EM The FedAvg algorithm is a simple iterative procedure realized in four simple steps . At the beginning of each round t , the server communicates the model parameters , let them be w , to a subset of the devices . The devices then proceed to optimize w , e.g. , via stochastic gradient descent , on their respective dataset via a given loss function Ls ( Ds , w ) : = 1 Ns Ns∑ i=1 L ( Dsi , w ) ( 1 ) where s indexes the device , Ds corresponds to the dataset at device s and Ns corresponds to its size . After a specific amount of epochs of optimization on Ls is performed , denoted as E , the devices communicate the current state of their parameters , let it be φs , to the server . The server then performs an update to its own model by simply averaging the client specific parameters wt = 1S ∑ s φs . 2.1 THE CONNECTION TO EM . We now ask the following question ; does the overall algorithm correspond to a specific optimization procedure on a given objective function ? Let us consider the following objective function : argmax w 1 S S∑ s=1 log p ( Ds|w ) , ( 2 ) where Ds corresponds to the shard specific dataset that has Ns datapoints , p ( Ds|w ) corresponds to the likelihood of Ds under the server parameters w. Now consider decomposing each of the shard specific likelihoods as follows : p ( Ds|w ) = ∫ p ( Ds|φs ) p ( φs|w ) dφs , p ( φs|w ) ∝ exp ( −λ 2 ‖φs −w‖2 ) , ( 3 ) where we introduced the auxiliary latent variables φs , which are the parameters of the local model at shard s. The server parameters w act as “ hyperparameters ” for the prior over the shard specific parameters and λ acts as a regularization strength that prevents φs from moving too far from w. How can we then optimize the resulting objective in the presence of these latent variables φs ? The traditional way to optimize such objectives is through Expectation-Maximization ( EM ) . EM consists of two steps , the E-step where we form the posterior distribution over these latent variables p ( φs|Ds , w ) = p ( Ds|φs ) p ( φs|w ) p ( Ds|w ) , ( 4 ) and the M-step where we maximize the probability of Ds w.r.t . the parameters of the model w by marginalizing over this posterior argmax w 1 S ∑ s Ep ( φs|Ds , wold ) [ log p ( Ds|φs ) + log p ( φs|w ) ] = argmax w 1 S ∑ s Ep ( φs|Ds , wold ) [ log p ( φs|w ) ] . ( 5 ) If we perform a single gradient step for w in the M-step , this procedure corresponds to doing gradient ascent on the original objective , a fact we show in Appendix D. When posterior inference is intractable , hard-EM is sometimes employed . In this case we make “ hard ” assignment for the latent variables φs in the E-step by approximating p ( φs|Ds , w ) with its most probable point , i.e . φ∗s = argmax φs p ( Ds|φs ) p ( φs|w ) p ( Ds|w ) = argmax φs log p ( Ds|φs ) + log p ( φs|w ) . ( 6 ) This is usually easier to do as we can use techniques such as stochastic gradient ascent . Given these hard assignments , the M-step then corresponds to another simple maximization argmaxw 1 S ∑ s log p ( φ ∗ s|w ) . As a result , hard-EM corresponds to a block coordinate ascent type of algorithm on the following objective function argmax φ1 : S , w 1 S ∑ s ( log p ( Ds|φs ) + log p ( φs|w ) ) , ( 7 ) where we alternate between optimizing φ1 : S and w while keeping the other fixed . How is this framework the same as FedAvg ? By letting λ → 0 in Eq . 3 it is clear that the hard assignments in the E-step mimic the process of optimizing a local model on the data of each shard . In fact , even by optimizing the model locally with stochastic gradient ascent for a fixed number of iterations with a given learning rate we implicitly assume a specific prior over the parameters ; for linear regression , this prior is a Gaussian centered at the initial value of the parameters ( Santos , 1996 ) whereas for non-linear models it bounds the distance from the initial point . After obtaining φ∗s the M-step then corresponds to argmaxw Lr : = 1S ∑ s− λ 2 ‖φ ∗ s −w‖2 , and we can easily find a closed form solution by setting the derivative of the objective w.r.t . w to zero and solving for w : ∂Lr ∂w = 0⇒ λ S ∑ s ( φ∗s −w ) = 0⇒ w = 1 S ∑ s φ∗s . ( 8 ) It is easy to see that the optimal solution for w given φ∗1 : S is the same as the one from FedAvg . Of course , FedAvg does not optimize the local parameters φs to convergence at each round , so one might wonder whether the correspondence to EM is still valid . It turns out that the alternating procedure of EM corresponds to block coordinate ascent on a single objective function , the variational lower bound of the marginal log-likelihood ( Neal & Hinton , 1998 ) of a given model . More specifically for our setting , we can see that the EM iterations perform block coordinate ascent to optimize : argmax w1 : S , w 1 S ∑ s Eqws ( φs ) [ log p ( Ds|φs ) + log p ( φs|w ) ] +H [ qws ( φs ) ] ( 9 ) where ws are the parameters of the variational approximation to the posterior distribution p ( φs|Ds , w ) and H [ q ] corresponds to the entropy of the q distribution . To obtain the procedure of FedAvg we can use a ( numerically ) deterministic distribution for φs , qws ( φs ) : = N ( ws , I ) . This leads us to the same objective as in Eq . 7 , since the expectation concentrates on a single term and the entropy of qws ( φs ) becomes a constant independent of the optimization . In this case , the optimized value for φs after a fixed number of steps corresponds to the ws of the variational approximation . It is interesting to contrast recent literature under the lens of this framework . Optimizing the same hierarchical model with hard-EM but with a non-trivial λ results into the same procedure that was proposed by Li et al . ( 2018 ) . Furthermore , using the difference of the local parameters to the global parameters as a “ gradient ” ( Reddi et al. , 2020 ) is equivalent to hard-EM on the same model where in the M-step we take a single gradient step . In addition , this view makes precise the idea that FedAvg is a meta-learning algorithm ( Jiang et al. , 2019 ) ; the underlying hierarchical model we optimize is similar to the ones used in meta-learning ( Grant et al. , 2018 ; Chen et al. , 2019 ) . How can we then use this novel view of FedAvg to our advantage ? The most straightforward way is to use an alternative prior which would result into different behaviours in local training and server side updating . For example , one could use a Laplace prior , which would result into the server selecting the median instead of averaging , or a mixture of Gaussians prior , which would result into training an ensemble of models at the server . In order to tackle the communication and computational costs , which is important for “ cross-device ” FL , we chose a sparsity inducing prior , namely the spike and slab prior . We describe the resulting algorithm , FedSparse , in the next section . 3 THE FEDSPARSE ALGORITHM : SPARSITY IN FEDERATED LEARNING Encouraging sparsity in FL has two main advantages ; the model becomes smaller and less resource intensive to train and furthermore , it cuts down on communication costs as the pruned parameters do not need to be communicated . The golden standard for sparsity in probabilistic models is the spike and slab ( Mitchell & Beauchamp , 1988 ) prior . It is a mixture of two components , a delta spike at zero , δ ( 0 ) , and a continuous distribution over the real line , i.e . the slab . More specifically , by adopting a Gaussian slab for each local parameter φsi we have that p ( φsi|θi , wi ) = ( 1− θi ) δ ( 0 ) + θiN ( φsi|wi , 1/λ ) , ( 10 ) or equivalently as a hierarchical model p ( φsi|θi , wi ) = ∑ zsi p ( zsi|θi ) p ( φsi|zsi , wi ) , p ( zsi ) = Bern ( θi ) , ( 11 ) p ( φsi|zsi = 1 , wi ) = N ( φsi|wi , 1/λ ) , p ( φsi|zsi = 0 ) = δ ( 0 ) , ( 12 ) where zsi plays the role of a “ gating ” variable that switches on or off the parameter φsi . Now consider using this distribution for the prior over the parameters in the federated setting . w , θ will be the server side model weights and probabilities of the binary gates . In order to stay close to the FedAvg paradigm of simple point estimation and since approximate inference for complex posteriors , such as those that arise in neural networks , is still an open problem , we will perform hard-EM in order to optimize w , θ . By using approximate distributions qws ( φs|zs ) , qπs ( zs ) , the variational lower bound for this model becomes argmax w1 : S , w , π1 : S , θ 1 S ∑ s Eqπs ( zs ) qws ( φs|zs ) [ log p ( Ds|φs ) + log p ( φs|w , zs ) + log p ( zs|θ ) − log qws ( φs|zs ) ] +H [ qπs ( zs ) ] . ( 13 ) For the shard specific weight distributions , as they are continuous , we will use qws ( φsi|zsi = 1 ) : = N ( wsi , ) , q ( φsi|zsi = 1 ) : = N ( 0 , ) with ≈ 0 which will be , numerically speaking , deterministic . For the gating variables , as they are binary , we will use qπsi ( zsi ) : = Bern ( πsi ) with πsi being the probability of activating local gate zsi . In order to do hard-EM for the binary variables , we will remove the entropy term for the qπs ( zs ) from the aforementioned bound as this will encourage the approximate distribution to move towards the most probable value for zs . Furthermore , by relaxing the spike at zero to a Gaussian with precision λ2 , i.e. , p ( φsi|zsi = 0 ) = N ( 0 , 1/λ2 ) , and by plugging in the appropriate expressions into Eq . 13 we can show that the local and global objectives will be argmax ws , πs Ls ( Ds , w , θ , ws , πs ) : = Eqπs ( zs ) [ Ns∑ i L ( Dsi , ws zs ) ] − λ 2 ∑ j πsj ( wsj −wj ) 2 − λ0 ∑ j πsj + ∑ j ( πsj log θj + ( 1− πsj ) log ( 1− θj ) ) + C , ( 14 ) argmax w , θ L : = 1 S S∑ s=1 Ls ( Ds , w , θ , ws , πs ) ( 15 ) respectively , where λ0 = 12 log λ2 λ and C is a constant independent of the variables to be optimized . The derivation can be found at Appendix E. It is interesting to see that the final objective at each shard intuitively tries to find a trade-off between four things : 1 ) explaining the local dataset Ds , 2 ) having the local weights close to the server weights ( regulated by λ ) , 3 ) having the local gate probabilities close to the server probabilities and 4 ) reducing the local gate activation probabilities so as to prune away a parameter ( regulated by λ0 ) . The latter is an L0 regularization term , akin to the one proposed by Louizos et al . ( 2017 ) . Now let us consider what happens at the server after the local shard , through some procedure , optimized ws and πs . Since the server loss for w , θ is the sum of all local losses , the gradient for each of the parameters will be ∂L ∂w = ∑ s λπs ( ws −w ) , ∂L ∂θ = ∑ s ( πs θ − 1− πs 1− θ ) . ( 16 ) Setting these derivatives to zero , we see that the stationary points are w = 1∑ j πj ∑ s πsws , θ = 1 S ∑ s πs ( 17 ) i.e. , a weighted average of the local weights and an average of the local probabilities of keeping these weights . Therefore , since the πs are being optimized to be sparse through the L0 penalty , the server probabilities θ will also become small for the weights that are used by only a small fraction of the shards . As a result , to obtain the final sparse architecture , we can prune the weights whose corresponding server inclusion probabilities θ are less than a threshold , e.g. , 0.1 . It should be noted that the sums and averages of Eq . 16 , 17 respectively can be easily approximated by subsampling a small set of clients S′ from S. Therefore we do not have to consider all of the clients at each round , which would be prohibitive for the “ cross-device ” setting of FL .

Summary: The paper proposes to re-interpret federated averaging (FedAvg) as a version of the expectation-maximization (EM) algorithm under a particular probabilistic model. Further, the authors propose to use spike-and-slab sparsity inducing priors over the local model parameters to sparsify the learned models (the corresponding method is called FedSparse), which naturally reduces the communication cost (only non-zero parameters need to be sent over the network). Improvement in communication efficiency is showcased on a few standard federated datasets.