Federated Averaging as Expectation Maximization
Authors: Christos Louizos
Matthias Reisser
Joseph Soriaga
Max Welling
Venue:  (2021)
Keywords: federated
learning
sparsity
expectation
maximization
efficient
FedAvg
FedSparse
EM

Abstract: Federated averaging (FedAvg), despite its simplicity, has been the main approach in training neural networks in the federated learning setting. In this work, we show that the algorithmic choices of the FedAvg algorithm correspond to optimizing a single objective function that involves the global and all of the shard specific models using a hard version of the well known Expectation-Maximization (EM) algorithm. As a result, we gain a better understanding of the behavior and design choices of federated averaging while being able to provide interesting connections to recent literature. Based on this view, we further propose FedSparse, a version of federated averaging that employs prior distributions to promote model sparsity. In this way, we obtain a procedure that leads to reductions in both server-client and client-server communication costs as well as more efficient models. 

Under review as a conference paper at ICLR 2021
FEDERATED AVERAGING AS EXPECTACTION-MAXIIIAAA
A Anonymous authors Paper under double-blind review
ABSTRACT
Federated averaging (F (FedAvg), despite its simplicity, has been the main approach in training neural networks in the federated learning setting. In this work, we show that the algorithmic choices of the FedAvg algorithm correspond optimizing a single objective function that involves the global and all of the shard specific (EM)
models using a hard version of the well known algorithm. As result, we gain better understanding of the behavior and design choices of federated averaging while being able to provide interesting connections to recent literature. Based on this view, we further propose FedSpa rse, a version of federated averaging that employs prior distributions to promote model sparsity. In this way, we obtain procedure that leads to reductions in both server-client and client-server communication costs as well as more efficient models.
INTRODUCTION
Smart devices have become ubiquitous in today's world and generating large amounts of po- tentially 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 al. 2016) has been proposed to train global machine learning models without the need for participating devices to transmit their data the server. The Federated Averaging (F (FecAvg) (McMahan al. 2016) algorithm communicates the parameters of the machine learning model instead of the data itself, which is 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 F FedAvg can be derived by applying the general Expectation Masiiiiotaaiiiititionn (EM) framework to 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 "'rooss-devcce" 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, 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 cummunicationaaccuraa trade-offs compared to both FedAvg as well as methods proposed for similar reasons (Caldas et al., 2018).

---

Under review as a conference paper at ICLR 2021
2 FEDAVG THROUGH THE LENS OF EM
The FedAvg algorithm a simple iterative procedure realized in four simple steps. At the beginning of each round the server communicates the model parameters, them be w, a subset of the devices. The devices then proceed to optimize w, egg,, via stochastic gradient descent, on their respective dataset via given loss function
LD, w)
N. Σ L(D w)
where indexes the device, D, corresponds the dataset at device s and N, corresponds its size. After a specific amount of epochs of optimization on L, is performed, denoted as E, the devices communicate the current state of their parameters, let it be Φ., to the server. The server then performs an update to its own model by simply averaging the client specific parameters W, Ex D....
2.1 THE CONNECTION TO EM
We now ask the following question; does the overall algorithm correspond to a specific optimization procedure given objective function? Let consider the following objective function:
1
Σ log (D,,/w), S
arg max
where D. corresponds the shard specific dataset that has N, datapoints, p(D,[w) corresponds the likelihood of D under the server parameters w. Now consider decomposing each of the shard specific likelihoods follows:
p(D. w) - p(D.)o))p(( (w)d.. p(o./w) x exp (/)b.-w13) _-1)
where we introduced the auxiliary latent variables o.. 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 regularization strength that prevents Φ, from moving too far from
How can we then optimize the resulting objective in the presence of these latent variables o,? The traditional way optimize such objectives is through Expectation Expectation-Maximizttiin Maximization (EM). EM consists of two steps, the E-step where we form the posterior distribution over these latent variables
p ww)
p(o.|D., w)
(4)
p(D,/w)
and the M-step where we maximize the probability of D. w.r.t the parameters of the model w by marginalizing over this posterior
Σ Doo (log logp(d>s[w)
arg max-
arg max
D. woo
(5)
W
we perform single gradient step for w in the M-step, this procedure corresponds to doing gradient ascent on the original objective, fact we show in Appendix
When posterior inference is intractable, hard-EM is sometimes employed. In this case we make "hard" assignment for the latent variables d, in the E-step by approximating p(o., D.,, w) with its most probable point, i.e.
Φ arg max pDD.ll))p((o.lw arg max log ( Φ. p(D, (w)
Φ.
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
2

---

Under review as a conference paper at ICLR 2021
arg m E, log p(o:/w) As result, hard-EM corresponds to block coordinate ascent type of algorithm on the following objective function
E (logp((D.lo)) + log p(o)/w)).
arg max
$1.8.w S
where we alternate between optimizing Φ1.5 and w while keeping the other fixed.
How is this framework the same as FedAvg? By letting À - 0 in Eq. 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 specific prior over the parameters; for linear regression, this prior is a Gaussian centered at the initial value of the parameters (Santos) (996) whereas for non-linear models it bounds the distance from the initial point. After obtaining Ø; the M-step then corresponds to arg max L. :- _/2 Σ. žे b w wll?, and we can easily find closed form solution by setting the derivative of the objective w.r.t. w to zero and solving for w:
oh, =0 Ow
A Σ -
= w ΣΦ の5
(8)
S
It is easy to see that the optimal solution for given i is the same as the one from FedAvg. Of course, FedAvg does not optimize the local parameters Φ. 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:
Σ Bew, (6.) Pogg + logp(ó./ww). + H Hlaw, (Φ.))
arg Iax. W1:8.W W S
where , are the parameters of the variational approximation to the posterior distribution p( (Da, w) and H [4] corresponds to the entropy of the q distribution. To obtain the procedure of FedAvg we can use (numerically) deterministic distribution for Φ.. 9w, (Φ.) AWW,,, eI). This leads to the same objective as in Eq. | since the expectation concentrates on a single term and the entropy of qw. (Ø..) becomes constant independent of the optimization. In this case, the optimized value for Φs after a fixed number of steps corresponds to the w, 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 the global parameters as a "gradient" (Reddi et al.) (020) is equivalent to hard-EM on the same model where in the M-step we take single gradient step. In addition, this view makes precise the idea that FedAvg is a meta-learning algorithm (Jiang et al. 2019); the underlying hicrarchical model we optimize is similar to the ones used in meta-learning (Grant et al. 2018 Chen 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 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. is mixture of two components, a delta spike zero 5(0), and a continuous distribution over the real line, i.e. the slab. More specifically, by adopting a Gaussian slab for each local parameter Φ.. we have that
p(O.:0, w:) (1 0,)5(0) + (.Ni W:, 1/2),
(10)
3

---

Under review as a conference paper at ICLR 2021
or equivalently as hierarchical model
p( 00; w.) = - p(asi O:)p( we), p(zss) (ern(
(11)
- -l.)
1/2), p( ( - 0) 00),
(12)
where zu plays the role of a "gating" variable that switches on or off the parameter Φss. Now consider using this distribution for the prior over the parameters 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 open problem, we will perform hard-EM in order to optimize w, O. By using approximate distributions qw. (Ø, (2,), Qr, (z.), the variational lower bound for this model becomes
Σ Egg. (8.) (0b1 po, (D, ) + log w. 2.) + log p(z,
arg max W1:8. WiS.W. 1. S
log sw. (0.1z.) H HQg,, (z.)]-
(13)
For the shard specific weight distributions, as they are continuous, we will use gw, (Ø Zei = 1) N(w (Wasi e), q(@ 1) := NO, e) with se 0 which will be, numerically speaking, deterministic. For the gating variables, as they are binary, we will use qn., (Z.i) == Bern T.;) with Ti being the probability of activating local gate Zzi- In order to do hard-EM for the binary variables, we will remove the entropy term for the ęm. (z.) from the aforementioned bound as this will encourage the approximate distribution to move towards the most probable value for Z.. Furthermore, by relaxing the spike at zero to Gaussian with precision xg, i..,, p(O.á se 22 = 0) = NNO. 1/Xg), and by plugging in the appropriate expressions into Eq. 13 we can show that the local and global objectives will be
No
arg max Ls (D,, w, O, Ws = (e.)( W.
Σ
L(Dai, W. 2 25) -
Toj( (Wo( (Waj
w))
Ao Σ Taj +
Σ (Ts log 0g + (1 - Tgz) -4 log(1 - 03)) + C, (14)
1 5
arg max L =
S (D,, w, O.,,,,
Σ
(15)
w,O
respectively, where Ao = log A and C' a constant independent of the variables to be optimized. The derivation can be found Appendix D is interesting see that the final objective at each shard intuitively tries to find trade-off between four things: 1) explaining the local dataset D., 2) having the local weights close to the server weights (regulated by x), 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 Ag). The latter an Lo regularization term, akin to the one proposed by] Louizos et al. (2017).
Now let us consider what happens the server after the local shard, through some procedure, optimized w, and π,, Since the server loss for w, O is the sum of all local losses, the gradient for each the parameters will be
OL
= -EATA
AT, w),
Σ
1
(16)
Ow
00
1
Setting these derivatives to zero, we see that the stationary points are
w 1 ΣS T.W.,
0
S Σ n
(17)
i.., a weighted average of the local weights and an average of the local probabilities of keeping these weights. Therefore, since the π. are being optimized to be sparse through the Lo penalty, the server probabilities will also become small for the weights that are used by only a small fraction of the shards. As result, to obtain the final sparse architecture, corresponding server inclusion probabilities 0 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 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.
can prune the weights whose
4

---

Under review as a conference paper at ICLR 2021
3.1 REDUCING THE COMMUNICATION COST
The framework described so far allow us to learn a more efficient model. We now discuss how we can use it in order cut down both download and upload communication costs during training.
Reducing client to server communication cost In order to reduce the client to server cost we will communicate sparse samples from the local distributions instead of the distributions themselves; in this way we do not have to communicate the zero values of the parameter vector. This leads to large savings, while still keeping the server gradient unbiased. More specifically, we can express the gradients and stationary points for the server weights and probabilities as follows
dc Ow
Σ
AB (x.) [Zo(Wa w)], W = Byrns (2118)
Σ
Zo We
(18)
ac
Σ Egg (x.)
Z
0
1 S Σ
Egr, (n.) [Zo].
(19)
00
As result, can then communicate from the client only the subset of the local weights w, that non-zero in 2, 9m, (Z+). W. w, @ Z,, and the server can infer the state of z, by inspecting which parameters were omitted. Having access to those samples, the server can then form I-sample stochastic estimates of either the gradients or the stationary points for w, O. Notice that this is way to reduce communication without adding bias in the gradients of the original objective. In case that we are willing to incur extra bias, we can further use techniques such as quantization (Amiri et al., 2020) and top-k gradient selection (Lin et al., 2017) to reduce communication even further. Such approaches are left for future work.
Reducing the server client communication cost the updated distributions at each round. Unfortunately, for simple unstructured pruning, this doubles the communication cost as for each weight w; there is an associated O, that needs to be sent to the will employ structured pruning, which introduces a single additional
The server needs to communicate to the clients
client. To mitigate this effect parameter for each group of weights. For groups of moderate sizes, e.g., the set of weights of a given convolutional filter, the extra overhead is small. We can also take the communication cost reductions one step further if we allow for some bias in the optimization procedure; we can prune the global model during training after every round and thus send each of the clients only the subset of the model that has survived. Notice that this is easy to do and does not require any data at the server. The inclusion probabilities are available at the server, so we can remove the parameters that have 0 less than threshold, e.g. 0.1. This can lead to large reductions in communication costs, especially once the model becomes sufficiently sparse.
3.2 FEDSPARSE IN PRACTICE
While optimizing for w, locally is straightforward to do with gradient based optimizers, n, is more tricky, as the expectation over the binary variables z, in Eq. [4]is intractable to compute in closed form and using Monte-Carke integration does not yield reparametrizable samples. To circumvent these issues, we rewrite the objective in an equivalent form and use the hard-concrete relaxation from (Louizos et al. 2017). which can allow for the straightforward application of gradient ascent. We provide the details in Appendix IF When the client has to communicate to the server, propose to form w, by sampling from the tero-temperature relaxation, which yields exact binary samples. Furthermore, at the beginning of each round, following the practice of F FedAvg, the participating clients initialize their approximate posteriors be equal to the priors that were communicated from the server. Empirically, we found that this resulted in better global accuracy.
Local optimization
Parametrization of the probabilities Since there have been evidence that such optimization based pruning can be inferior to simple magnitude based pruning (Gale et al.) 2019), we take an approach that combines the two and reminisces the recent work of Azarian et al. (2020p. We parametrize the probabilities 0, π. as function of the model weights and magnitude based thresholds that regulate how active a parameter can be. More specifically, we use the following parametrization
W Tg
[ xg
l|2
Tag
0g : a
ss a
(20)
T
5

---

Under review as a conference paper at ICLR 2021
where the subscript denotes the group, (-)) is the sigmoid function, Tg. Tag are the global and client specific thresholds for a given group 9 and T is a temperature hyperparameter. Following Azarian et al. (2020) we also "detach" the gradient of the weights through 0, n, to avoid decreasing the probabilities by just shrinking the weights. With this parametrization we lose the ability to get a closed form solution for the server thresholds, but nonetheless can still perform gradient based optimization at the server by using the chain rule. For positive threshold, we use a parametrization in terms of a softplus function, i.e., T log(1 + exp(v)) where is the learnable parameter.
Algorithm The server side algorithm for FedSparse (assuming weight sparsity for sim- plicity). a--) is the sigmoid function, is the threshold for pruning.
Algorithm 2 The client side algorithm for FedSparse.
Get w, from the server
0 - a ((1 ((w) -/T)
Initialize v and w for round in 1 ..... do - +logg( 1 exp(v))
w.. V, f W, V
for epoch e in 1...., E. do for batch b E B do
+ o ([w) - -)
Ts <- log(1 + exp(v.))
w - 1 > </w
> prune global model
T. f o ((wal a - TT)/T)
Initialize V = 0. V V 00
L. Lolb, w. O We TT))
for in random subset of the clients do
W. t SGD(V b.)
W f- CLIENT(s, v)
Va ADAMAX(V L.)
2. f I[ f 0]
end for end for
V w + Z,
( (
w)
+ [2g (1 1) - (2-))))))(vv/
n. f- ((waal T.)
end for
Zg ~ 9T, (Z+) return Z. ② W.
w1+ vi+1 f- ADAM(V(). ADAMAX(V..
end for
RELATED WORK
Both contributions of this work share similarities with several recent works in federated learning. FedP rox (Li et al. 2018) proposed to add a proximal term to the local objective at each shard, so that it prevents the local models from drifting too far from the global model. Through the EM view of F FedAvg we show how such a local objective arises if we use a non-trivial precision for the Gaussian prior over the local parameters. Furthermore, FedAvg has been advocated be meta-learning algorithm in Jiang et al. (2019); with the EM view we make this claim precise and show that the underlying hierarchical model that FedAvg optimizes the same as the models used in several meta-learning works (Grant et al. 2018 Chen et al. 2019). Furthermore, by performing a single gradient step for the M-step in the EM view of FedAvg we see that we arrive at a procedure that has been previously explored both in a meta-learning context with the Reptile algorithm (Nichol ct al. 2018), as well as the federated learning context with the "generalized" FedAvg (Reddi et al. 2020). One important difference between meta-learning and FedAvg is that the latter maximizes the average, across shards, marginal-likelhoon in order to find the server global parameters whereas meta-learning methods usually optimize the global parameters such that the finetuned model perform well on the local validation set. Exploring such parameter estimation methods, as, e.g., in Chen et al. (201)) in the federated scenario and how these relate existing approaches that merge meta-learning with federated learning, e.g., Fallah et al. (202)), is an interesting avenue for future work. Finally, the EM view we provide here for FedAvg can also provide novel perspective for optimization works that improve model performance via model replicas (Zhang et al. 2019 Pittorino et al. 2020).
Reducing the communication costs a well known and explored topic in federated learning. FedSparse has close connections to federated dropout (Caldas et al.) 2018), as the latter can be understood via a similar hierarchical model, where gates z are global and have a fixed probability 0 for both the prior and the approximate posterior. Then at each round and for each client, the server samples these gates (due to the expectation), collects the submodels and sends those to be optimized at the shard. Compared to federated dropout, FedSparse allow us optimize the dropout rates to the data, such that they satisfy a given accuracy sparsity trade-off, dictated by the hyperparameter Ag- Another benefit we gain from the EM view of F FedAvg is that it makes it clear that the server

---

Under review as a conference paper at ICLR 2021
can perform gradient based optimization. As result, harvest the large literature efficient distributed optimization (Lin et al. 2017] Bernstein et al.) 2018; Wangni et al.) 2018 Yu et al.) 2019), which involves gradient quantization, sparsification and more general compression. On this front, there have been also other works that aim to reduce the communication cost in FL via such approaches (Sattler et al., 2019; Han al. 2020). In general, such approaches can be orthogonal to FedSparse and exploring how they can be incorporated is promising avenue for future research.
5 EXPERIMENTS
We verify in three tasks whether the FedSparse procedure leads to similar or better global models compared to FedAvg while providing reductions communication costs and efficient models. As a baseline that also reduces communication costs by sparsifying the model, consider the federated dropout procedure from Caldas et al. (2018), which we refer to as FedDrop. For each of the three tasks, present the results for FedSpa rse with regularization strengths that target three sparsity levels: low, mid and high. For the F FepDrop baseline, we experiment with multiple combinations of dropout probabilities for the convolutional and fully connected layers, For each of these tasks we report the setting that performs best in term of accuracy / communication trade-off.
The first task we consider is federated version of CIFARIO classification; we partition the data among 100 shards in non-i.i.d way by following Hsu et al. (2019). For the model we employ LeNet-5 convolutional architecture (LeCun al., 1998) with the addition of dropout(O.1) 1) for the second convolutional layer and (00pout((.3) for the first fully connected layer in order to prevent overfitting locally at the shard. We optimize the model for Ik communication rounds. For the second task we consider the 500 shard federated version of CIFAR 100 classification from Reddi et al. (2020). For the model we use a ResNet20, where we replace the batch normalization layers with group Reddi normalization, following Reddi al. (2020), and optimize for 6k rounds. For the final task
considered the non-i.i.d. Femnist classification and we use the same configuration as CIFAR10 but we optimize the model for 6k rounds. More details can be found Appendix A
We evaluate FedSparse rse and the baselines on two metrics that highlight the tradeoffs between accuracy and communication costs. On both metries the x-axis represents the total communication cost incurred up until that point and the y-axis represents two distinct model accuracies. The first one corresponds the accuracy of the global model on the union of the shard test sets, whereas the second one corresponds to the average accuracy of the shard specific "local models" on the shard specific test sets. The "local model" on each shard is the model configuration that the shard last communicated to the server, and serves as a proxy for the personalized model performance on each shard. The later metric is motivated from the meta-learning (Tiang et al. 2019) and EM view of federated averaging, and corresponds to using the local posteriors for prediction on the local test set instead of the server side priors.
5.1 EXPERIMENTAL RESULTS
The results from our experiments can be found in the following table and figures. Overall, we observed that the FedSparse models achieve their final sparsity ratios early in training, i.., after 30-50 rounds, which quickly reduces the communication costs for each round (Appendix B
Table I: Average global, local test-set accuracies across clients in %, along with total communications costs in GB and sparsity of the final model for Cifar10, Cifar100 and Femnist. We report the average over the last 10 evaluations.
Method
Femnist G.Acc. L.Acc. Comm. Spars. G.Acc. L.Acc. Comm. Spars. G.Acc. L.Acc. Comm. Spars.
Cifar 10
Cifar 100
FedAvg
69.97
86.71
65
41.11
61.61
123
85.62
90.81
272
FedDrop
71.54
86.86
35
37.65
58.02
112
85.23
89.37
174
FedSparse, low
70.30
87.65
52
20.1
40.69
61.31
113
2.2
85.03
90.82
270
1.0
FedSparse, mid
70.30
87.54
36
45.7
31.94
54.83
59
51.0
84.83
90.61
145
47.0
FedSpa rse, high 68.46
87.17
26
62.4
24.68
49.58
45
62.9
84.13
89.68
95
65.5

---

Under review as a conference paper at ICLR 2021
We can that for CIFAR 10. the FedSparse models with medium (~45%) and high (~62%) sparsity outperform all other methods for small commu- nications budgets on the global accuracy front, but eventu- ally surpassed by FedDrop on higher budgets. However, on the local accuracy front, we see that the redSparse models Pareto dominate both baselines, achiev- ing, e.g., 87% local accuracy with 43% less communication compared to FedAvg. Overall, judging the final performance only, that FedDrop reaches the best accuracy on the global model, but F redSparse reaches the best accuracy in the local models.
G:bba Aoo
3
80.0
88
Fackey
FedDrop
17.5
FecStop
56
75.0
FecSparas,
Fackpares mi
FedSparse, high
FacSparee, Nigh
G25
(a) CIFAR 10 global acc.
(b) CIFAR local acc.
On CIFAR 100, the differences are less pronounced, as the mod- els did not fully converge for the maximum number rounds we use. Nevertheless, we still observe similar patterns; for small communication budgets, the sparser models are better for both the global and local accu- racy egg. they can reach 32% global accuracy while requiring 13% less communication than
Findka
Teedkg
26
Fudlipur
FedDrop
FedDaop
FadSpprnon
Fadtipursee
16.0
Fedipparon,
C1000
FedSparsn, 11gh
45.0
08
08
CIFAR 100 global acc.
CIFARI00 local acc.
FedAvg.
Finally for Femnist we observe the most differences, as the Fem- nist task is also more communi- cation intensive due to having 3.5k shards. We see that the FedSparse algorithm Pareto dominates both FedDrop and FedAvg and, more specifically, the high sparsification setting, it can reach 84% global accuracy and 89% local accuracy while re- quiring 41% % and 51 % less com-
02
90
Gl00 Noc.
Fodlvg
Loodl
FodAng
redtlop
Fadiparne,
Fetfipamn,
FodSparse, mid
FodSparse, (il
72
sattpparee tegh
1ultparam, light
70
50
10.5
100
141%
100
200
250
50
500
150
200
250
(a) Femnist global acc.
(b) Femnist local acc.
munication compared to FedAvg respectively. Judging by the final accuracy; both FedAvg and FedSparse with the low setting reached similar global and local model performance, which is to be expected given that that particular Fe FedSparse arse setting lead to only 1% sparsity.
CONCLUSION
In this work, we showed how the FedAvg algorithm, the standard in federated learning, corresponds to applying a variant of the well known EM algorithm to a simple hierarchical model. Through this perspective, we bridge several recent works on federated learning well as connect FedAvg to meta-learning. As a straightforward extension stemming from this view, we proposed FedSparse, generalization of FedAvg with sparsity inducing priors. Empirically showed that FedSparse can learn sparse neural networks which, besides being more efficient, can also significantly reduce the communication costs without decreasing performance.
8

---

Under review as a conference paper at ICLR 2021
REFERENCES
Mohammad Mohammadi Amiri, Deniz Gunduz, Sanjeev R Kulkarni, and H Vincent Poor. Federated learning with quantized global model updates. arXiv preprint arXiv:2006.1005 2020.
Kambiz Azarian, Yash Bhalgat, Jinwon Lee, and Tijmen Blankevoort. Learned threshold pruning. arXiv preprint arXiv:2003.00073 2020.
Jeremy Bernstein, Yu-Xiang Wang, Kamyar Azizzadenesheli, and Anima Anandkumar. signsgd: Compressed optimisation for non-convex problems. arXiv preprint arXiv: 2018.
Sebastian Caldas, Jakub Konečny, Brendan McMahan, and Ameet Talwalkar. Expanding the reach of federated learning by reducing client resource requirements. arXiv preprint arXiv:18/2.072210 2018.
Yutian Chen, Abram L Friesen, Feryal Behbahani, David Budden, Matthew W Hoffmann, Ar- naud Doucet, and Nando de Freitas. Modular meta-learning with shrinkage. arXiv preprint arXiv:1909.05559, 2019.
Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning: A meta- learning approach. arXiv preprint arXiv22002.09488, 2020.
Trevor Gale, Erich Elsen, and Sara Hooker. The state of sparsity in deep neural networks. arXiv preprint (txXiv:102.0995777 2019.
Erin Grant, Chelsea Finn, Sergey Levine, Trevor Darrell, and Thomas Griffiths. Recasting gradient- 2018.
based meta-learning as hierarchical bayes. arXiv preprint
Pengchao Han, Shiqiang Wang, and Kin K Leung. Adaptive gradient sparsification for efficient federated learning: An online learning approach. arXiv preprint arXiv:2001.04756 2020.
Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Measuring the effects of non-identical data distribution for federated visual classification. arXiv preprint
2019.
Yihan Jiang, Jakub Konečný, Keith Rush, and Sreeram Kannan. Improving federated learning personalization via model agnostic meta learning. arXiv preprint arXiv:1909. 2019.
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: 2019.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1:Xx::4412.69800 2014.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278-2324 1998.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith.
Federated optimization in heterogeneous networks. arXiv preprint arXiv:1812.06122 2018.
Xiang Li, Kaixuan Huang, Wenhao Yang. Shusen Wang, and Zhihua Zhang. On the convergence of 2019.
fedavg on non-iid data. arXiv preprin! 1907.02907..22
Yujun Lin, Song Han, Huizi Mao, Yu Wang, and William J Dally. Deep gradient compression: Reducing the communication bandwidth for distributed training. arXiv preprint arXiv:1712.01887, 2017.
Christos Louizos, Max Welling, and Diederik P Kingma. Learning sparse neural networks through L_ regularization. arXiv preprint arXiv:1712.01311 2017.
H Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, et al. Communication-effiiinn learning of deep networks from decentralized data. arXiv preprint arXiv:1602.055 002.00000066,2. 2016.
Toby Mitchell and John J Beauchamp. Bayesian variable selection in linear regression. Journal of the american statistical association, 1023-1032, 1988.

---

Under review as a conference paper at ICLR 2021
Radford M Neal and Geoffrey E Hinton. A view of the em algorithm that justifies incremental, sparse, and other variants. In Learning in graphical models, pp. 355-368. Springer, 1998.
Alex Nichol, Joshua Achiam, and John Schulman. On first-order meta-learning algorithms. arXiv preprint wtXiv:/803.09999 2018.
Fabrizio Pittorino, Carlo Lucibello, Christoph Feinauer, Enrico M Malatesta, Gabriele Perugini, Carlo Baldassi, Matteo Negri, Elizaveta Demyanenko, and Riccardo Zecchina. Entropic gradient descent algorithms and wide flat minima. arXiv preprint arXiv:2006.07897 2020.
Sashank Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and Brendan McMahan. Adaptive federated optimization. arXiv preprint 00Xi::0003.00295. 2020.
Reginaldo J Santos.
Equivalence of regularization and truncated iteration for general ill-posed problems. Linear algebra and its applications, 236:25-33, 1996.
Felix Sattler, Simon Wiedemann, Klaus-Robert Müller, and Wojciech Samek. Robust and communication-effiiee federated learning from non-iid data. IEEE transactions on neural networks and learning systems, 2019.
Jianqiao Wangni, Jialei Wang, Ji Liu, and Tong Zhang. Gradient sparsification for communication- efficient distributed optimization. In Advances in Neural Information Processing Systems, pp- 1299-1309, 2018.
Yue Yu, Jiaxiang Wu, and Longho Huang. Double quantization for communication-effiiennt dis- tributed optimization. In Advances in Neural Information Processing Systems, pp. 4438 4449, 2019.
Michael Zhang, James Lucas, Jimmy Ba, and Geoffrey E Hinton. Lookahead optimizer: k steps forward, step back. In Advances in Neural Information Processing Systems, pp. 9597-9608, 2019.
APPENDIX
A EXPERIMENTAL DETAILS
For all of the three tasks we randomly select 10 clients without replacement in given round but with replacement across rounds. For the local optimizer of the weights we use stochastic gradient descent with a learning rate of 0.05, whereas for the global optimizer we use Adam (Kingma & Ba 2014) with the default hyperparameters provided in (Kingma & Ba 2014). For the pruning thresholds in FedSpars se used the Adamax (Kingma & Ba 2014) optimizer with 1e - 3 learning rate the shard level and the Adamax optimizer with le - 2 learning rate at the server. For all three of the tasks used F = 1 with batch size of 64 for CIFAR10 10 and 20 for CIFAR 100 and Femnist. It should be noted that for all the methods we performed gradient based optimization using the difference gradient for the weights (Reddi et al. 2020) instead of averaging.
For the FedDrop baseline, we used very small dropout rate of 0.01 for the input and output layer and tuned the dropout rates for convolutional and fully connected layers separately in order to optimize the accuracy communication tradeoff. For convolutional layers we considered rates in {0.1, 0.2, 0.33 whereas for the fully connected layers we considered rates {0.1, 0.2, 0.3. 0.4, 0.5). For CIFAR10 ( we did not employ the additional dropout noise at the shard level, since we found that was detrimental for the FedDrop performance. Furthermore, for Resnet20 on CIFARION we did not apply federated dropout at the output layer. For CIFARI0 the best performing dropout rates were 0.1 for the convolutional and 0.5 for the fully connected, whereas for CIFAR 100 it was 0.1 for the convolutional. For Femnist, saw that rate of 0.2 for the convolutional and rate of 0.4 for the fully connected performed better.
For Fe FedSparse, from a dense model. The temperature for the sigmoid in the parameterization of the probabilities was set to T 0.001. Furthermore, downscaled the cross-entropy term between the client side
initialized v such that the thresholds 7 lead to 0 = 0.99 initially, f.e. started
10

---

Under review as a conference paper at ICLR 2021
probabilities, Ta. and the server side probabilities, 0 by mutitiplying it with 1e 4. Since the beginning of each round we were always imitializing a = 0 and we were only optimizing for a small found that the full strength of the cross-entropy was not
number of steps before synchronizing, necessary. Furthermore, for similar reasons, i.e. we set w, = w at the beginning of each round, we also used À 0 for the drift term s w)2. The remaining hyperparameter λn dictates how sparse the final model will be. For the LeNet-5 model the An's we report are {5e - 7, 5e 6, 5e - 5} for the "low", "mid" and "high" settings respectively, which were optimized for CIFARIO and used as-is for Femnist. For CIFAR100 and Resnet20, we did not perform any pruning for the output layer and the λu for the "low", "mid" and "high" settings were (5c - 7, These were chosen so that we obtain models with comparable sparsity ratios the one CIFARIO.
5e 5} respectively.
B EVOLUTION OF SPARSITY
We show the evolution of the sparsity ratios for all tasks and configurations in the following plot. We can see that in all settings the model attains its final sparsity quite early in training.
80
60
Fetipprroo low
-40
Fediipars med
Fastparn
20
Fattquuras
FadSpanse
00
150
200
300
100
100
Slounds
Fluunds
(a) CIFARIO sparsity
(b) CIFARION 100 sparsity
(c) Femnist sparsity
C ADDITIONAL RESULTS
Convergence plots terms of communication rounds. In order to understand whether the extra noise is detrimental to the convergence speed of FedSparse, we plot the validation accuracy in terms of communication rounds for all tasks and baselines. As can be seen, there no inherent difference before FedSparse starts pruning. This happens quite early in training for CIFAR 100 thus there where we observe the most differences.
50
Gobni
16
40
Fedlparn:
FulSpenke rnd
10
Fecticasee regh
Factipane, regh
edoppaaee hop
16 Rounds 200
Rounda 10
Rounds
(a) CIFARIO
CIFAR100
Femnist
Figure 5: Evolution of the validation accuracy in terms of communication rounds.
Impact of server side pruning. order to understand whether server side pruning is harmful for convergence, plot both the global and average local validation accuracy CIFAR 10 for the "mid" setting of F FedSparse with and without server side pruning enabled. As we can see, there are noticeable differences and in fact, pruning results into slightly better overall performance.
11

---

Under review as a conference paper at ICLR 2021
87.5
85.0
Global Apc.
82.5
Local Acc.
80.0
77.5
75.0
55
FedSparse,
FedSparse, mid
FedSparse, mid, no global prune
redSparse, mid, no googlobbll-rud
50
70.0
40
60
100
20
40
80
80
100
Rounds 10
Rounds x 10
(a) CIFARIO global val. acc.
(b) CIFARI00 local val. acc.
Figure 6: Evolution of the validation accuracy in terms of communication rounds with and without
server side pruning.
D CORRESPONDENCE BETWEEN SINGLE STEP EM AND GRADIENT ASCENT
With the addition of the auxiliary variables Φ, we have that the overall objective for the server becomes
arg max
Σ log (.D,)o.)p(d) ww)d. S
1
(21)
By performing EM with single gradient step for in the M-step (instead of full maximization), we are essentially doing gradient ascent on the original objective 21 To see this, we can take the gradient of Eq. 21 w.r.t. w where Z. (().ll))p(O..w))))
/ p(Da ).)
Op(p.Iw) dØ,
(22)
5
Z.
1 Σ / (Dz,(s))((e ww Blooollp...o dØ.
(23)
Z.
dw
5 Σ / plo..D., w)
d Ologplo. dØ.
(24)
dw
where compute Eq. 24 we that we first have to obtain the posterior distribution of the local variables Φ, and then estimate the gradient for w by marginalizing over this posterior.
DERIVATION OF THE LOCAL LOSS
Let p( (:; 2si
((w 1/A), p(p.) Zsi = 0) (/0, and a(dsalz.e 1)
N(w.i, e²) 2. . Zzi = 0) = (0, 2). Furthermore, let (Zzzi) = Bern(TTs)). The local objective that stems from 13 can be rewritten
arg max Egg. (x. Juw, (b.1k.)) l..) (ogg((D.lo.)) Bg (v.) [K (Kw, (0wp:// Wiss.
E (z. [log p (0))
(25)
where we omitted from the objective the entropy of the distribution over the local gates.
One of the quantities that we are after is
Eg(,,.)) [KLgg(Os (l.00//s....))ppppp...... (
Ti: KLN (Was. ss. (w,,, 1/2)) (1 T
KIN 1/22))
(26)
The KL term for when Zyi can be written
2
1
KL(N(w..: wi. ()W,, 1/A))
- log > - log 2
+
Wai
w:)
(27)
2
12

---

Under review as a conference paper at ICLR 2021
The KL term for when zs 0 can be written
Azc² 1
KL(N(O, 2 (1),2) = log 22 - -loge+
(28)
2
2
Taking everything together we thus have
sf
dc²
L Bett..)) ( 2.:)]
(Wei - w:)² + Tas(
log A - logg e +
2 )
2
2
2
(1 1--Tsi)- 1
Agc²
1
loog A2 - -loge+
(29)
2
sti
1
A2
2
(Wet w.)2 2 Tss
log
_ () - 12)
C
2
2
(30)
Ari (Wai w:) + AO + C,
(31)
2
where A = 1 log * and (A- X2) was omitted due to 2² ~ O. In the appendix of Louizos et al. (2017), the authors argue about a hypothetical prior that results into needing À nats to transtorm that prior to the approximate posterior. Here we make this claim more precise and show that this prior is upproximately equivalent to a mixture of Gaussians prior where the precision of the non- non-zero prior component À + e (in order avoid the L2 regularization term) and the precision of the zeroth component A2 is equivalent to À eep((2 where Ån is the desired To regularization strength.
Furthermore, the cross-entropy from qT, (z,) p(z, (0) is straightforward to compute
Egg, (z.) log p Zz. 0)] = Σ (Te log 02 + (1 - Taj) log - 02))
(32)
By putting everything together we have that the local objective becomes
arg max Ey,,, (z,) W,.
Σ L(D.i, w. 2z.)
X
Σ Th (Wo) w,)² Au Σ Taj
+ Σ Tay log 0g + (1 -T)) log(1 - 0,)) + C. (33)
F LOCAL OPTIMIZATION OF THE BINARY GATES
We propose rewrite the local loss in Eq.
13 to
N,
(D,,,, w. 0 TT) = Eag., (z,)
Σ
L(Dai, W. 32x) - T T[Z.j # O[[wss - -w)2
-
0, + log(1 0,) ) (34)
M I[anj + 0] Σ [[zo] + 0] log 1-0 0,
and then replace the Bernoulli distribution QT, (z.,) with continuous relaxation, the hard-concrete distribution (Louizos et al. 2017). Let the continuous relaxation be ru, (z.). where u, are the parameters of the surrogate distribution. In this case the local objective becomes
No
LD w, A.... us) = Bru, (e.)
Σ [(Das, w. 22.22. - x Σ Ru,, (Zaj > O)(ws - w)2
Ao Σ Ru,, (Zuj > 0)+ + Σ
Run, 0) log
0, 1-0;
log(1 8,) (35)
where R,, is the cumulative distribution function (CDF) of the continuous relaxation Fin, (:). We straightforwardly optimize the surrogate objective with gradient ascent.
13


Reviews:
[{"review_id": "iWl2b98dIxg", "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**\n\nThe 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.\n\n**General comments and evaluation**\n\nIn my opinion, the paper has some interesting ideas; the effort to connect federated learning with techniques from probabilistic inference strongly resonates with me. However, the connection the authors make between federated averaging and EM is not as insightful as I anticipated. The fact that FedAvg can be seen as is a special case of EM with a particular choice of a prior distribution does not add much to our understanding of the behavior of the algorithm (e.g., convergence, issues related to different types of heterogeneity, etc.). While FedSparse is derived using Bayesian considerations, I would argue that the same algorithm can be derived using FedAvg + simple L1 regularization. Again, I do not see the benefit of the Bayesian formalism, especially given that it is essentially thrown away at the end by using the \"hard\" version of EM (where all the Gaussian distributions are collapsed). This unnecessarily complicates things without offering much insight.\n\nApart from the formalism, the authors should clarify the modeling assumptions of the FL setting they consider. Currently, there are two canonical FL settings, known as \"cross-device\" and \"cross-silo\" (https://arxiv.org/abs/1912.04977), which require FL methods to operate under drastically different assumptions -- the former assumes that the data is distributed among a very large number of devices (tens of millions) with limited memory and compute; the latter assumes a small number of clients with sufficient computational resources. It is unclear which of the settings the authors aim to address. For instance, $\\arg\\max_{w_{1:S}}$ in Eq. 13 seems meaningless in the cross-device setting since in practice S would be on the order of millions and only a subset of clients will participate in each round. I'm curious to hear which setting the authors had in mind.\n\nAll in all, I think the paper would significantly benefit from further clarification, a possibly theoretical analysis of the algorithm in heterogeneous settings (e.g., understand convergence rate and bounds on the tradeoffs between sparsity and performance) and/or a more careful and convincing empirical study.\n\n**Detailed comments and questions**\n\n- While reading the paper, a high-level question that wasn't clear to me: why use specifically \"hard\" EM for FedSparse instead of the standard EM?\n\n- In the paragraph after Eq. 13, the authors propose a bunch of approximations, starting with collapsing Gaussian distributions, then removing the entropy term from the objective, setting hyperparameters, etc. This is a lot of approximations, I would like to know the purpose and the effect of each on the final result. Currently, many of these choices are obscure and unclear.\n\n- The authors propose to use sampling of the binary variables z to compute gradients in Eqs. 18-19. The approach is indeed unbiased, but I'm worried that the increase in variance might be substantial in some practical cases. Can the authors somehow quantify this? (Ideally, provide a convergence rate for the algorithm that shows the effect of additional variance).\n\n- What is the effect of structured pruning proposed to reduce server-to-client communication cost? It seems like a heuristic and it is unclear how it affects the algorithm.\n\n**Experiments**\n\n- Could the authors define local vs. global accuracy exactly?\n- How many clients participated in each round in each benchmark? Was the server communicating with all clients or just a subset?\n- The authors show some improvement of FedSparse over FedAvg and FedDrop in terms of the total amount of communication traffic in GB. Why does that metric matter? In which use-case? This goes back to my question about modeling assumptions.\n- How many epochs per round did each client run? I would be curious to see how FedSparse compares to FedAvg if we double the amount of local computation (eg, make clients train for 2 epochs at each round instead of 1). I imagine that would improve convergence and reduce # of rounds, and as a result, reduce the total communication cost.\n- I think that the benchmark datasets considered in the study may not be the best to showcase the benefit of sparsity. I would be curious to see how FedSparse compares to FedAvg on StackOverflow dataset introduced by Reddi et al. (2020), where the sparsity benefits can be much more visible. Eg, learning extremely sparse models could enable the use of much larger vocabularies, and as a result, produce more accurate models than what is achievable with FedAvg. That would be much more compelling evidence in favor of FedSparse than the results reported in the paper.", "extra_scores": {}, "extra_text": {}}, {"review_id": "Yvd7MIKKEuC", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 2, "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 of the paper:\nThis paper brings into light a new perspective where the well known \"FedAvg\" algorithm can be modeled as expectation maximization. This reveals certain characteristics about federated averaging which can be useful in practice. Using these insights, \"Fedsparse\" algorithm is introduces that provides a trade-off between the sparsity (thus the model size) and accuracy of the models. Ample experiments are provided to validate theoretical claims.\n\nQuality:\nPresentation and the quality of writing is decent and the flow is understandable. But the major contributions are not clearly stated. I am wondering why the related work is at the end. I think it is better move the related work to an early section along with the emphasis on contributions.\n\nOriginality and significance:\nI believe the connection between FedAvg and EM and the related understandings open up new avenues for designing better federated learning algorithms. FedSparse is clearly one of such methods that allows a trade-off between accuracy and memory consumption. Overall, believe this paper adds valuable contributions to the broader ICLR community.\n\nOther comments:\nWhen communicating the parameters from server to clients, I am wondering whether there are certain heuristics that can be used to prune $\\theta$ values. Using an arbitrary threshold may not be the best given that the heterogeneous and non I.I.D. local datasets. Perhaps the authors can look into the effects of this as a future direction.\n\n==============================================================================================================\n\nAdded after reading author response:\n\n-----------------------------------------------------------------------\n\nAuthors have sufficiently addressed my concerns and I'm planning to maintain my generous score based on my initial understanding. However, other reviewers have raised many important concerns and I encourage authors to improve the paper based on those.", "extra_scores": {}, "extra_text": {}}, {"review_id": "TyO97DMLIlK", "reviewer": "AnonReviewer1", "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": "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.\n\nThe paper revolves around two disjoint stories: *1) Federated learning can be cast in the EM framework*; and *2) FedSparse can solve federated learning with sparse priors*. Unfortunately, neither of these two stories are told in a convincing manner. My personal suggestion to the authors would have been to focus on one of these messages and do a thorough work around it. Here are my questions on the different parts of the story for the authors to strengthen their story and contributions.\n\n**1-a)** The equivalences between coordinate descent and full EM, and many of the other observations made in this paper only hold when we are dealing with strongly concave convex problems. The paper lacks a rigorous statement on what conditions are needed for these methods to converge and what convergence rate guarantees could be derived. \n\n**1-b)** The observed equivalences in this paper are under the iid assumption. What happens in the presence of heterogeneous data?\n\n**1-c)** Can you add a clear discussion on the connections between (Li et al. 2018), (Reddi et al. 2020), and the proposed EM framework in terms of the pros and cons and their suitability to different scenarios inspired by the observations in this paper? And can you then propose a set of experiments to verify such observations empirically?\n\n**2-a)** How is FedSparse related to approximate message passing methods? AMP could also be cast as an instance of EM methods and is prevalent in solving sparse linear problems under the Bernoulli Gaussian prior assumption.\n\n**2-b)** Can you please explain why the global accuracy of FedSparse is capped when compared to FedDrop? It is not clear to me why a sparse prior should be a good assumption in federated scenarios.\n\n**2-c)** Can you please do a more in-depth study of the dependence of the performance of FedSparse in different scenarios? Can you also include (in the supplementary) a simulated scenario with truly sparse models and verify that FedSparse is well suited to such scenarios?\n\n==== post rebuttal ====\n\nIt is an interesting observation that federated averaging and several other variants of it may be cast as instances of expectation maximization (EM). However, unfortunately, these connections have been made in a rather informal way, not resulting in much of new insights into convergence of such federated algorithms, or other practical tweaks to them. The paper then moves on to propose FedSparse, an abrupt shift from the original observation to federated learning with sparse priors. The study of FedSparse is thin, with unclear motivations, and with not so strong experimental results. Connections with the plethora of existing literature on sparse signal processing, specifically as it is cast in the EM framework, e.g., approximate message passing, is missing. In the end, the reader is left with two interesting ideas but it is not clear what the take-home message is.\n\nMy suggestion to the authors would be to \"formalize\" the connections between federated averaging and other instances of these algorithms (e.g., draw from EM literature to provide convergence guarantees, or maybe provide other variants of these algorithms with more favorable properties) and resubmit the paper with more formal connections to the EM framework. In my opinion, FedSparse is itself a separate paper that needs a separate motivation, set of hypotheses, and experiments.", "extra_scores": {}, "extra_text": {}}, {"review_id": "5e4BHdUJndU", "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": "  \nThe paper tries to use EM to explain the optimization procedure of the training method in a federated setting, and then treats the local model as a hidden variable of EM. Then, the authors propose a new FedSparse method to reduce communication efficiency by training a sparse model.\n\nPros:\n\nIt is a novel idea to treat local mode’s parameter \\phi as a hidden variable in an EM framework. In particular, the following FedSparse algorithm is designed based on this setting. \n\nThis proposed sparsity-based model compression method is linked to federated learning in a comprehensive way. \n\nAs compared to FedAvg, the proposed method can reduce about 50% communication cost without scarifying accuracy. \n\nCons:                                                     \n\nThe FedAvg is an SGD-based optimization, and treating the local model’s parameters as hidden variable phi_s is a little bit confusing. In particular, the FedAvg has no regularization term to prevent \\phi moving too far from w, and the FedAvg’s variant version involves such a mode divergence-based regularization term – FedProx (Tian Li et. al. 2018).  Therefore, the definition of p(phi | w) in Equation 3 need more discussion.\n\nThe paper uses a classic sparsity method, spike, and slab (Mitchell & Beauchamp, 1988). The authors need to introduce a more model sparsity method for compression purposes and then justify choosing the spike and slab. \n\nIn the experiment part, the proposed method and baseline methods have no significant difference in terms of accuracy. From a communication efficiency perspective, the proposed method should be also compared to FedDrop. Moreover, reduced communication cost in FedSparse is not very high.\n\n\n\t \n ", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers agree that the EM perspective of Federated Learning is novel and interesting. However, a common criticism is that the connection made is rather shallow and not sufficiently developed. There look to be quite interesting potentials of the proposed framework and the specific FedSparse method, but I agree with the reviewers that both aspects need further development before they are in publishable form.
Decision: Reject