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: Federated learning has emerged as a promising approach to training models at the edge devices. This paper makes an observation that most algorithms used within federated learning, including the popular FedAvg, could be cast as instances of EM methods. The paper then continues to propose FedSparse, a federated learning framework imposing sparse priors (specifically Bernoulli-Gaussian priors), and concludes with some experimental results on FedSparse.