Federated-EM with heterogeneity mitigation and variance reduction
Authors: Aymeric Dieuleveut
gersende fort
Eric Moulines
Geneviève Robin
Venue: NeurIPS 2021 Poster (2021)
Keywords: Expectation-Maximization
Federated Learning
Latent variables models

Abstract:  The Expectation Maximization (EM) algorithm is the default algorithm for inference in latent variable models. As in any other field of machine learning, applications of latent variable models to very large datasets make the use of advanced parallel and distributed architecture mandatory. This paper introduces FedEM, which is the first extension of the EM algorithm to the federated learning context. FedEM is  a new communication efficient method, which handles partial participation of local devices, and is robust to  heterogeneous distribution of the datasets. To alleviate the communication bottleneck, FedEM compresses appropriately defined complete data sufficient statistics. We also develop and analyze an extension of FedEM to further incorporate a variance reduction scheme. In all cases, we derive finite-time complexity bounds for smooth non-convex problems.  Numerical results are presented to support our theoretical findings, as well as an application to federated missing values imputation for biodiversity monitoring.


Federated Expectation Maximization with heterogeneity mitigation and variance reduction
Aymeric Dieuleveut
Gersende Fort Institut de Mathématiques de Toulouse
Centre de Mathématiques Appliquées
Ecole Polytechnique. France
Université Toulouse; CNRS
Institut Polytechnique de Paris
UPS, Toulouse, France
aymeric. www.lleeeut@olllyeecceeeee
edu
gersende. ferreende.fort@matt. forttmathh....ooooooommomoc..... univ-toul ouse. fr
Eric Moulines
Geneviève Robin
Centre de Mathématiques Appliquées Ecole Polytechnique, France CS Dpt, HSE University, Russian Federation
Laboratoire de Mathématiques et Modélisation d d'Évry
Université d' d'Evry Val d'Essonne; CNRS
eric. wtt..oooolieespllleeceniiiiieeee edu
Evry-Courcouronnes, France generieve.
Abstract
The Expectation Maximization (EM) algorithm is the default algorithm for infer- ence in latent variable models. As in any other field of machine learning, appli- cations of latent variable models to very large datasets makes the use of advanced parallel and distributed architectures mandatory. This paper introduces FedEM, which is the first extension of the EM algorithm to the federated learning context. FedEM a new communication efficient method, which handles partial participa- tion of local devices, and is robust to heterogeneous distributions of the datasets. To alleviate the communication bottleneck, FedEM compresses appropriately de- fined complete data sufficient statisties. We also develop and analyze extension FedEM to further incorporate a variance reduction scheme. all cases, we de- rive finite-time complexity bounds for smooth non-convex problems. Numerical results are presented to support our theoretical findings, as well as an application to federated missing values imputation for biodiversity monitoring.
Introduction
The Expectation Maximization (EM) algorithm the most popular approach for inference in latent variable models. The EM algorithm, special instance of the Majorize/Minimize algorithm [24], was formalized by [8] and is without doubt one of the fundamental algorithms in machine learning- Applications include among many others finite mixture analysis. latent factor models inference, and missing data imputation; see [38, 29, 26, 13] and references therein. As in any other field of machine learning, training latent variable models on very large datasets make the use of advanced parallel and distributed architectures mandatory. Federated Learning (FL) [22, 39], which exploits the computation power of large number of edge devices to perform distributed machine learning, a powerful framework to achieve this goal.
The conventional EM algorithm is not suitable for FL settings. We propose several new distributed versions of the EM algorithm supporting compressed communication. More precisely, our objective
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

is to minimize non-convex finite-sum smooth objective function
Argmingge F(O), F(0)=
n1 Σ ()) +
+ C TR
where n is the number of workers/devices which are connected to a central server, and the worker #i only has access to its local data; finally R is a penalty term which may be introduced to pro- mote sparsity, regularity, etc. In latent variable models, ()) = m وド= log p(y;;:)) where {ukg 1y=1 are the m observations available for worker #i, and p 8) the incomplete likeli- hood. p(y; 0) is defined by marginalizing the complete-datu likelihood ply, =; defined as the joint probability density function of the observation and non-observed latent variable E Z, i.e. p(y; 0) = fz p(y, 2; (dz) where Z is the latent space and μ a measure on Z. We focus this paper on the case where p(y, 2; 8) belongs to a curved exponential family, given by p 2; 0) := p(y, 2) exp { (s(y,, 2), 6(0) (0))
(2)
where s(y, 2) E R9 is the complete-date sufficient statistics, θ R and 1 p: Y x Z + R- are vector/scalar functions.
O
R.
In absence of communication constraints, the EM algorithm is a popular method to solve (1). It alternates between two steps: in the Expectation (E) step, using the current value of the iterate Ocurrs it computes a majorizing function 0 - Q Oarr)) given up to an additive constant by
(0, .urr)) == - (5( ), <(0) + (0) + R(0) where 5(0) :
50)0;
(3)
and (0) is the ith device conditional expectation of the complete-data sufficient statistics:
(0) = 1 m
(0) (0) = S/- s(yis) x) =)p(z 0)u(dz)
(4)
where p(z Wiy; 0) :- p 2 0)) (p:; 0). As for the M step, an updated value of Oeurr is computed as a minimizer of 0 + Q(, Deurr). The majorizing function is then updated with the new Ocurri this process is iterated until convergence. The EM algorithm is most useful when for any Ocurr O, the function I + Q feurr) is convex function of the parameter 0 which is solvable in @ either explicitly or with little computational effort. A major advantage of the EM algorithm stems from its invariance under homeomorphisms, contrary to classical first-order methods: the EM updates are the same for any continuous invertible re- pe-paaametrization [23].
In the FL context, the vanilla EM algorithm is affected by three major problems: (1) the communi- cation bottleneck, (2) data heterogeneity, and (3) partial participation (PP) of the workers.
When the number of workers is large, the cost of communication becomes overwhelming. A classi- cal technique to alleviate this problem is to use communication compression. Most FL algorithms are first order methods and compression is typically applied to stochastic gradients. Yet, these methods not appropriate solve (1) since (i) they do not preserve the desirable homeomorphic invariance property, and (ii) the full EM iteration is not distributed since the M step is performed by the central server only. This calls for extension of the EM algorithm to the FL setting.
Since workers are often user personal devices, the issue of data heterogeneity naturally arises. Our model in Equations (1), (3) and (4) allows the local loss functions to depend on the worker i E {1.....nz n and the observations Vij to be independent but not necessarily identically distributed. In addition, our theoretical results deal with specific behaviors for each worker i E {l..... ..). see e.g., A5, and 8. the FL-EM setting, heterogeneity manifests itself by the non-equality of the local conditional expectations of the complete-data sufficient statistics s, 's; modifications to the algorithms must be performed to ensure convergence at the central server.
Finally, subset of users are potentially inactive in each learning round, being unavailable or unwill- ing to participate. Thus, taking into account PP of the workers and its impact on the convergence of algorithms, is major issue.
FedEM. The main contribution of our paper is a new method called FedEM, supporting commu- nication compression, partial participation and data heterogeneity. In this algorithm, the workers compute an estimate of the local complete-data sufficient statistics s, using minibatch of data, apply an unbiased compression operator to a noise compensated version (using technique inspired by [17, 15)) and send the result to the central server, which performs aggregation and the M-step (i.e. the parameter update).

---

VR-FedEM. We improve FedEM by adding variance reduction method inspired by the SPIDER framework [9] which has recently been extended the EM framework [10]. For both FedEM and VR-FedEM, the central server updates the expectations of the global complete-data sufficient statistics through a Stochastic Approximation procedure [3, 4]. When compared to FedEM, VR-FedEM addi- tionally performs variance reduction for each worker, progressively alleviating the variance brought by the random oracles which provide approximations of the local complete-data sufficient statistics. Theoretical analysis. EM in the curved exponential family setting converges to the roots of function h (see e.g. Section 2). We introduce a unified theoretical framework which covers the convergence of FedEM and VR-FedEM algorithmss in the non-convex case and establish convergence guarantees for finding an e-stationary point (see Theorem 1 and Theorem 3). In both cases, we provide the number (o)t(e) of optimization steps and the number (c of computed conditional expectations to reach c-stationarity. These results show that in the Stochastic Approximation steps of VR-FedEM, the step sizes are independent of m, the number of observations per server. Furthermore, the computational cost in terms of (CE((c) improves on earlier results. In this respect, VR-FedEM has the same advantages as SP SDIDR [9] compared to SVRG [18] and SAGA [6],
SPIDER-EM [10] compared to sEM-vr [5] and FIEM [20, 11]. Lastly, our bounds demonstrate the robustness of FedEM and VR-FedEM to data heterogeneity.
Finally, seen as a root finding algorithm in a quantized FL setting, VR-FedEM can be compared to VR-DIANA [17]: we show that VR-FedEM does not require the step sizes to decrease with m and provides state of the art iteration complexity to reach a precision e.
Notations. For vectors a, 6 in R2, b) the Euclidean scalar product, and ||-I| denotes the asso- ciated norm. For r 2 1, |||]]r is the f,-norm of vector The Hadamard product a ① b denotes the entrywise product of the two vectors a, b. By convention, vectors are column- vectors. For a matrix A, A is its transpose and || is its Frobenius norm; for two matrices A, B, (A, ==
Trace B A). For positive integer n. set [n]* := {1, ..} and [n] := {0. m2.. The set of non non-negative integers (resp. positive) denoted by IN (resp. N The minimum (resp. maximum) of two real numbers a, is denoted by a/b (resp. a' b)) We will use the Bachmann- Landau notation
a(z) = Ob(()) characterize an upper bound of the growth rate of a(x) as being b(z).
2 FedEM: Expectation Maximization algorithms for federated learning
Recall the definition of the negative penalized (normalized) log-likelihood F()) from (1). Along the entire paper, we make the following assumptions Al to A3,which define the model at hand.
A1. The parameter set θ C Rd is a convex open set. The functions R - + R, o : : - R²,
θ + R. and p(y,g, :) Z + R,,, s(ytj..) Z - R. for E [n]* and [m]] measurable
functions. For any 0 E θ and i E [m]', the log- tiglliieldoodd is finite: -oo < (0) xc.
A2. For all 0 E ② and i [n]],, the conditional expectation (0) is well-defined.
A3. For any s E IR?, the map 8 + Argmingee {v)(0) + R(0) - (s, ((0)) exists and is unique; the singleton is denoted by {s))
EM defines sequence (0x, 2 0} that can be computed recursively as 0x+1 = To5(O2), where the map T is defined in A3 and 5 is defined in (3). On the other hand, the EM algorithm can be defined through a mapping in the complete-data sufficient statistics, referred as the expectation space. In this setting, the EM iteration defines a R8-valued sequence (S., k 2 0} given by S = So T(aa). Thus, we observe that the EM algorithm admits two equivalent representations:
(Parameter space) 0x+1 = T o s0x)); (Expectation space) 8 = 5o T(S.).
(5)
In this paper, we focus on the expectation space representation; see [23] for an interesting discussion on the connection of EM and mirror descent. It has been shown in [7] that if s, is fixed point to the EM algorithm in the expectation space, then 0., == T(*+) is a fixed point of the EM algorithm in the parameter space, i.e., 0, T o s(0_)) :) note that the converse is also true. Define the functions h; and h from R. to RR. by h(s) := 1
x >i h;(s) with (((s) == 5, T(s) -.
h(s) = 1 M hr(s),
n
(2(s) = 5, o T(s) -s
(6)
A key property is that the fixed points of EM in the expectation space are the roots of the mean field + h(s) (see (3) for the definition of 5). Therefore, convergence of EM-based algorithms is
3

---

evaluated in terms of c-stationarity (see [14, 10]:: for all >> O, there exists a (possibly random)
termination time K s.t.: [Ib(Sk)) S c. Another key property of EM is that is a monotonic
algorithm: each iteration leads a decrease of the negative penalized log-likelihone i.e. F(0x++1) Σ F(Oz) or, equivalently in the expectation space ( S FoT(Se) (for sequences (0x, A 2 0} and {Sk, k 2 0} given by (5)). A4 assumes that the roots of the mean field h are the roots of the gradient of F a T (see [7] for the same assumption when studying Stochastic EM). A5 assumes global Lipschitz properties of the functions h,'s. s.
A4. The function W = F IS + I is continuously differentiable on R and its gradient is globally Lipschitz with constant Ly- Furthermore, for any s E R9, V W(()) -Bss))((3) where B(s) is q x q positive definite matrix. In addition, there exist 0 < Umin S Umaa such that for any
8 E R2, the spectrum of (ss) is in [Umin; Umax]-
A5. For any E [n]". there exists L, > 0 such that for any s, s R2, [l(((s) h₂(s')ll
(15 o T(s) - s) (5, ((5')=s)) & Lills - s'll.
Federated EM algorithm.
Algorithm 1: FedEM with partial participation
Our first contribution, the novel al- gorithm FedEM is described by algo- rithm 1. The algorithm encompasses partial participation of the workers: at iteration (k + 1), only subset A&+1 of active workers participate the training, see line 3. The averaged fraction of participating workers is de- noted Each of the active workers #i computes an unbiased approximation
Data: kmax E N*; for E [n]*. Vo, E R";
So E IR;; a positive sequence
{74+1, k 1kmax - []] > coefficient
p = A-p [card(A)]/n.
Result: The FedEM-PP sequence:
{S., h E [kmax)] }
Set Vo = n
2i=1 i- Vo,;
n
for 0...., kmax 1 do
Sample AR-1 ~ IPP
1 i (line 6) (( condition-
4 for i Ax+1 do
ally the past (see Appendix D.2 for a rigorous definition), these approxi- mations are independent. The work- ers then transmit to the central server compressed information about the new sufficient statistics. A naive so- lution would be to compress and trans-
(worker #i)
6 Sample S4+1.5, an approximation of o T(Sa)
Set Ak+1, Sp+1. VR,, - S
8
Set VK+1, = Vi, + a Quant 1.+1.).
Send Quaan(AA41,,) to the central server
10 for i & Ak+1 do
(worker #i) mit 1.. Ss. but data heterogeneity 11 between servers often prevents these local differences from vanishing at the optimum, leading to large compres- sion errors and impairing convergence of the algorithm. Following [28], a memory Ve,, (initialized to (2,)
12 Set V = Va, (no update)
13 (the central server)
14
Set
H 1 = Vv (np) EieAur Quant(AR41,)
Set 8R41 = =S, + 7k+ 1H++1 7 Set
V*+1 = V2x + an EieAess Quant(A
k = 0) is introduced; and the differ-
ences A := Sk+1. S. - Va, are compressed for E Au+1 (line 7
Send SK+1 and T(Sx+1) to the n workers
and line 9). These memories are updated locally: Vi++.; = Ve.: + a Quant (A+11.) at line 8, with > 0 (typically set to 1/(1+w) where w is defined A6). On its side, the central server releases an aggregated estimate Se41 of the complete-data sufficient statistics by averaging the quantized dif- ference (mp) (aap)>'EeAAe
EiEAL Quant(Ax-1,, and by adding V₂ (line 14 and line 15). Then, it updates Vk+1 = V4_ + am Σ (uant(Ax+1,i) see line 15. The final step consists in solving the M-step of the EM algorithm, i.c. in computing T(Sk+1) (see A3).
We finally state our assumption on the compression process. We consider large class of unbiased compression operators Quant satisfying a variance bound:
A6. There exists w 2 0 s.t. for any E R2: E [Quant(s)] s, and E ((QQuant(s)]3 < (1+w)]|+}13
Intuitively, the stronger the compression is, the larger w will be. Remark that if no compression used, equivalently for all s E 129, Quant(s) = then A6 is satisfied with w 0. An example

---

of quantization operator satisfying A6 is the random dithering that can be described as the random operator Quant : R' + R2, Quant(2))
(qqunn)]|z)] sign(z) (z) [ (| /| [=|[-)))
+ {] where
r 2 is user-defined, ş is uniform random variable on [0, 1|9 and Squant E N is the number of
0 Example IJ. Another example, namely the -uont---qunttzzation is provided in the supplemental (see Appendix B). More generally, this assumption is valid for many compression operators, for example resulting in sparsification [see. e.g. 28].
levels of roundings; see [17, 2]. This operator satisfies A6 with w Squant O(g2/r_00/// + g²/2); +
see [17,
The convergence analysis is under the following assumptions on the oracle 1 for any i E [n]", the approximations 1,+ are unbiased and their conditional variances are uniformly bounded in k. For each & E N, denote by F& the o-algebra generated by (S,,,, E [m]*, e E [k]) and including the randomness inherited from the quantization operator Quant up to iteration #k.
7. For all k N. conditional Fk. {Sk+1s) independent. Moreover, for any i E [n]*. E 11
5, c T(Sx) and there exists of such that for any k 2 S of
|| - 5, T(
A7 covers both the finite-sum setting described in the introduction, and the online setting. In the finite-sum setting, 5, is of the form "ーー Sty. In that case, S can be the sum over a mini- batch I,+ of size sampled at random in [m]]". with or without replacement and independently of the history of the algorithm: we have S = b-1 D- E)eB,4.. 5. o T(Sk). In the online setting, the oracles 1 i come from online processing of streaming informations; in that case Sk+1.5 can be computed from a minibatch of independent examples so that the conditional variance o² which will be inversely proportional the size of the minibatch, can be made arbitrarily small.
Reduction of communication complexity for FL. Reducing the communication cost between workers is a crucial aspect of the FL approach [19]. In gradient based optimization, four tech- niques have been used to reduce the amount of communication: (i) increasing the minibatch size and reducing the number of iterations, (ii) increasing the number of local steps between two com- munication rounds, (iii) using compression, (iv) sampling clients at each step. Here, we provide a tight analysis of strategies (i), (iii) and (iv) (sampling client part of PP).
Regarding the interest of performing multiple iterations (ii), as analyzed for example in [21, 27] for the classical gradient settings, note that: first, from theoretical standpoint, tradeoffs between larger minibatch and more local iterations are unclear [37]. Secondly, performing local iterations is not possible in the EM setting: one iteration of EM is the combination of two steps E and M and the M step, which required the use of the map T, is only performed by the central server; this remark is fundamental specificity of the EM framework (which is not shared by the gradient framework). In applications, we usually do not want 7 to be available at each local node. However, our work allows perform multiple local iterations of the step before communicating with the central server. In algorithm 1, the local statistics Sk+1, are general enough to cover this case; see the comment above on A7.
Finally, as we do not perform local full EM iterations, we do not face the well-identified client- drift challenge (in the presence of heterogeneity). Yet, we stress that combining compression and heterogeneity results in other challenges: it known in the Gradient Descent setting (see e.g. [28, 31)), that heterogeneity strongly hinders convergence in the presence of compression. To alleviate the impact of heterogeneity, we introduce the Vk,2's memory-variables.
Convergence analysis, full participation regime. In this paragraph, we focus on the full- participation regime (p- = 1): for all k E [kmaa]* AA = [n]*. We now present in Theorem our key result, from which complexity expressions are derived. The proof is postponed to Appendix C.
Theorem 1. Assume Al to A7 and set L = n Ei L. o² = n-1 Di o7 Let {Sx, k e [kmaa]] be given by algorithm with > 0, (1 (1+w) w) and nh = E (0, 7max] where
Umia
Vn
Ymax
A 2/2L(1+
2L
Denote by K the uniform random variable on [Kima- - 11. Then, laking Vo.: h:( for all [m]"
Umin
(1 11-
Lix
[1h(3k)1P]]
5 7kimax
1
(W(S)) - min w) + +
+ 5w
0². (8)
Umin
n
5

---

n When there is no compression (w = 0 so that Quant(s) = s), we prove that the introduction of the random variables Vi..'s 's play no role whatever a and the choice of the Vo, 's and we have for any 7 E (0, 20min/ (Tw) (see (29) in the supplemental)
(1 (1-720m)) 20min
Tw
E
[I((Sk)1P]
(W(S)) min w) + + o²
ykmax
Optimizing the learning rate 7, we derive the following corollary (see the proof in Appendix C).
Corollary 2 (of Theorem 1), Choose == (WIS.) min W)n 1/2 A Ymnax- We ger
( b (I+bw]o2
Amax
(W(S.) -
min W) W + 5w)o2
(W(S)) - min W)
E [1((k))1
Umin
nkkmax
Tmax Kmax
Theorem and Corollary do not require any assumption regarding the distributional heterogeneity of workers. These results remain thus valid when workers have access to data resulting from different distributions - a widespread situation in FL. frameworks. Crucially, without assumptions on the heterogeneity of workers, the convergence of "naive" implementation of compressed distributed EM (i.e. an implementation without the variables Vi,a's ')) would not converge.
Let us comment the complexity to reach e-stationary point, and more preciscly how the com- plexity evaluated in terms of the number of optimization steps depend on w, n, o² and e. Since
Ko)a(e) kmax, from Corollary 2 we have that: Kope(e) = = [1+w]o"
ne?
Maximal learning rate and compression. The comparison of Theorem with the no compression case (see (9)) shows that compression impacts Tmax by a factor proportional to vT/ww// as w increases (similar constraints were observed in the risk optimization literature, e.g. in [17, 32])- This highlights two different regimes depending on the ratio vn/w3/2 if the number of workers n scales at least as w3, the maximal learning rate is not impacted by compression; on the other hand, for smaller numbers of workers < wa, compression rate. We highlight this conclusion with a small example in the case of scalar quantization for which v Squant: for q 102 and Squant = (obtaining compression rate of a factor 16), the maximal learning rate is almost unchanged n 2 16.
degrade the maximal learning
Dependency on The complexity (c) is decomposed into two terms scaling respectively as d²c-2 and e the first term being dominant when - O. This observation highlights two different regimes: a high noise regime corresponding to Ymax (1 + w) (ne 1) > where the complexity is of order o²c-2 2, and a low noise regime where Tmax + w)o (ne S 1 and the complexity is of order 7max An extreme example of the low noise case is 0² occurring for
example in the finite-sum case (i.e., when 5; = m Σ/= Sty) with the oracle SH+1,5 = 5; o T(Sa).
Impact of compression for e-stationarity. As mentioned above, the compression simulta- neously impacts the maximal learning rate (as in (7)) and the complexity Ke)t(e). Conse- quently, the impact of the compression depends on the balance between w, n, and c, and can distinguish four different "main" regimes. In the following tabular, for each of the four situations, we summarize the increase in complexity Kepe(e) resulting from compression.
Complexity regime:
Dl mC
(Dominating term in (c)
Timax regime:
Exumple situation
High noise o², small
Low o² (e-g., large minibatch)
(Dominating term in (7))
larger
large ratio n/w
Xwi
x1
VR
low ratio 12 n
X@
X/ //m
Depending on the situation, the complexity can be multiplied by a factor ranging from to w V (w3/2/ //T)) Remark that the communication cost of each iteration is typically reduced by compression of a factor least Moreover, the benefit of compression is most significant in the low noise regime and when the maximal learning rate is Umin/ (2L,/) (e.g., when n large enough). We then improve the communication cost of each iteration without increasing the optimization com- plexity, effectively reducing the communication budget "for free".
Because of space constraints, the results in the PP regime are postponed to Appendix A.

---

3 VR-FedEM: Federated EM algorithm with variance reduction
A novel algorithm, called VR-FedEM and described by algorithm 2, is derived additionally incor- porate variance reduction scheme in FedEM. It is described in the finite-sum setting when for all i E [n]', S; =- m Σド Siy: at each iteration # (1, A + 1), the oracle on S; o T(St,k) will use minibatch 1,, of examples sampled random (with or without replacement) in [m]".
The algorithm is decomposed into kout. outer loops (indexed by i), each of them having kin inner loops (in- dexed by k). At iteration # + 1) of the inner loops, each worker #i updates local statistic St,41.. based on a minibatch B of its own examples {5.g.] E B (see Line 8): starting from Sc,.. = m- Σy= Sty T(S S is defined in such a way that it approx- imates m > m Sij T(S.,) (see 2 1=1
Algorithm 2: VR-FedEM
Data: Kout: kins b No; for E [n]*. Vi,0, RR: Binit E IR;; a positive sequence
{7,,--11 L E [Rouu]", k E [Kin 1]]; or >> 0
Result: sequence: {Se,k, t E [hout]] [o, h E [mm]]}
$1.0 S Sinit. V1,0 = n Ei= V1.0,. for L, ...., do
1
2
S1,0,i - 1 >m Siz T(Sini))
3
m
4
for = 1, 1.... kou do
for k = 0, Km do
for i == 1, n (worker #i, locally) do
Corollary 18). Then, the worker #i sends to the central server a quan- tization of A1, (see Line 12) which can be seen as an approxima- tion of a -1 (h, - h, (S.,,,)))
Sample random batch B1,,+ of size b [m]]
Set kkt+ = Sy, +
Σ JEBI,AIL.. Big o T(S..) 5iz T(Sek-1))
-
upon noting that the variable Vi,++1.
defined by Line 10 approximates
Set Ar,t1,, = St,t11i Se,t Venu Set V V1,+1,1 = VV, + a (U, Send Quant (A.,)1,,)) to the central server
(.,...)) (see Proposition 26). The t0 central server learns the mean value I1 Vi,++1 = n=' Ei=1 V (see 12 Line 15 and Lemma 21) and, by 13 adding the quantized quantities, de- fines a field Hr,k+1 which ap- proximates n- E'= (St,k))) (see i5
(the central server)
Set He,k+1 = Vt,x + n2 -. Quant (ArR+1,s)
Set S S1 + Nr,++ H H1
Set V1,4,1 Ve, aan Eim1 Quant(Ar
Send St,4+1 and T(S to the n workers
Proposition 24). Line 14 can be seen 16 as Stochastic Approximation up- date, with learning rate 7r,4+1 and mean field + n Σ>= (s) (see
S S Sr, V Vi,kin
19 18
for i = l...., do
(6) for the definition of h.).
S m Ei= Siz T(S
The variance reduction is encoded in the definition of S1,i Line 8. We
Vi+1,0.5 = Vi.kingi
have S,,441.1 = b-1 EjeB..... St ((S.,)) Y G4,+1, The first term is the natural approximation
of 5, T(Sr,)) based on minibatch B,,+1.i. Conditionally to the past, Tr Tr,++, is correlated to the first term and biased, but its bias canceled at the beginning of each outer loop (see Line 20 and Appendix E.3.2): T Tr,++1,; defines control variate. Such variance reduction technique was first proposed in the stochastic gradient setting [30, 9, 36] and then extended to the EM setting [10, 12). At the end of each outer loop, the local approximations S1+1,0,5 are initialized to the full
sum m 1 E,T 5ij ((S...... (see Line 20) thus canceling the bias S.,, (see Proposition 17).
When there is single worker and no compression is used (n 1, w = 0), VR-FedEM reduces to SPIDER-EM, which has been shown to be rate optimal for smooth, non-convex finite-sum optimiza- tion [10]. Theorem studies the FL setting (n and > 0): it establishes finite time control of convergence in expectation for VR-FedEM. Assumptions A5 and A7 are replaced with A8.
A8. For any i |n]] and j E [m]*, the conditional expectations (iy are well defined for any 0 E θ, and there exists Lej such that for any s, s' E IR9, [| o T(s) - s) - (Sy o T(s') - s')ll م Lijll - s'll
7

---

Theorem 3. Assume Al to 3, A4, A6 and A8. Set L == n n-'m"mm Ei=1 Σ Lij. Let (8
[Kout]] h E [kin
1]}
be given by algorithm 2 run with a = 1/(1 + w), V1,0. =- -h:(S,,) for any
i E [n] [n]"bb: b= [ Amm
11-
and
Lw vn (1 w) (w+ 1 10w 1/2 -1
Ye.k y
Umin
1 + 4v2
Umax
(10)
Lw
8
Let (T, K) be the uniform random variable [hout]* x Ikin 1). independent of ( [Kout]* k E [Kin]}- Then, it holds
E 11 (E W(SLw)) min W)
(11)
Unia U ykintiout
E [I((s,
s $2(1+
[ + w)2
|||
(12)
The proof is postponed to Appendix E. This result is consequence of the more general Proposi- tion 25, We make the following comments:
1. Eq. (11) provides the convergence of E of interest E[]h((,,.k )ll²] is controlled by BI|
and Eq. (12) ensures that the quantity
We observe that 2(1 + y2 (1(++)) is uniformly bounded w.r.t. w as, by (10), 72 Ow- 3w-
Up our knowledge, this the first result on Federated EM, that leverages advanced variance reduction techniques, while being robust to distribution heterogeneity (the theorem is valid without any assumption on heterogeneity) and while reducing the communication cost.
3. Without compression (w = 0) and in the single-worked case (n = 1). Fort et al. [10] use kin b:
we recover this result as particular case. When n and w > 0, the recommended batch size decreases as 1/(1 + w
Convergence rate and optimization complexity. Our step-size y chosen constant and inde- pendent of ki kout. Indeed, contrary to Theorem I, there is no Bias- Variance trade-off (as typi- cally observed with variance reduced methods), and the optimal choice of y the largest one to ensure convergence. Consequently, since the number of optimization steps is hout Kin, we have (c)=
Impact compression on the learning rate and e-stationarity. The compression constant w does not directly appear in (11), but impacts the value of y. Two different regimes appear:
(i.e. we focus the large w, n asymptotics when 3 < n)), then = Ymui has nearly the same value as without compression [10]. The complexity then similar to the one of SPIDER-EM [10], with a smaller communication cost. The gain from compression is maximal in this regime.
1. if 4 Yumes Jn (1 + w) (w 1+10m1122 Low
<
Za
if 2f Venax L (1 + w) (w + 1+ 1/2 > (i.e. we focus on the large w, n asymptotics when Vn!
3 n)) then 7 O Vniin T
is strictly smaller than without compression. The optimization
Venax 3/-
complexity is then higher to the one of SPIDER- (by factor proportional to 3/2 /T) with a smaller communication cost (typically least times less bits exchanged per iteration). The overall trade-off thus depends on the comparison between w and n..
Complexity
1/(gr)
We summarize these two regimes in this tabular, fo-
rogime: (Deminating (10))
cusing on the large n, large w asymptotic regimes.
Example situstion
For the two regimes, we indicate the increase in
Vaina/ Lag Vainy La
large ratio nyu" low ratio n/w
complexity Kopt(e resulting from compression.
xu? ///
We provide a discussion on computed conditional expectations complexity ICEE in Appendix E.2.
As corollary of [10. Theorem 2). the optimization complexity of SP SPIDER-EN EM is Kout kinkout that is in order reach e-stationarity.
8

---

4 Numerical illustrations
In this section, we illustrate the performance of FedEM and VR-FedEM applied to inference in Gaus- sian Mixture Models (GMM), on a synthetic data set and on the MNIST data set. We also present an application to Federated missing data imputation with the analysis of the eBird data set [34, 1J.
Synthetic data. The synthetic data are from the following GMM model: for all & E [N]* and E {0, 1]. PaZz- = 9) = Tgi and conditionally to Ze = 9. Yi ~ N2(Ugg E). The 2 x 2 covariance matrix Σ is known, and the parameters be fitted are the weights (To, TT) and the expectations M1). The total number of examples is N = 10 the number of agents is n 10 and the probability of participation of servers is p 0.75. FedEM and VR-FedEM are run with y 10 2 = and a = 10 2 For FedEM, we consider the finite-sum setting when 5 = m Ej= Sy with 102; the oracle Sk+1,5 obtained by sum over a minibatch of b = 20 examples. For VR-FedEM, we set b = 5 and kin = 20. We run the two algorithms for 500 epochs (one epoch corresponds to N conditional expectation evaluations Siy). Figure/ I shows trajectory of || 1|² | given by VR-FedEM), along with the theoretical value of the mean field
given by FedEM (and [lh((ee (12 for FedEM (and [lh(S.,))l' I| for VR-FedEM). The results illustrate the variance reduction, and gives insight on the variability of the trajectories resulting from the two algorithms.
MNIST Data set. We perform similar experiment on the MNIST dataset illustrate the be- haviour of FedEM and VR-FedEM on GMM inference problem with real data. The dataset consists of N = 7 x 10 images of handwritten digits, each with 784 pixels. We pre-process the dataset by removing 67 uninformative pixels (which are always zero across all images) to obtain d = 717 pixels per image. Second, we apply principal component analysis to reduce the data dimension. We keep the dpc = 20 principal components of each observation. These N preprocessed observations are distributed at random across n = 102 servers, each containing m = 700 observations. We estimate a Rdre- multivariate GMM model with G = 10 components. Details on the multivariate Gaussian mixture model are given in the supplementary material (see Appendix F). Here again, s, is sum over the examples available at server #i; the minibatches are independent and sampled at random in [m]] with replacement; we choose = 20 and the step size is constant and set to = 10 3, The same initial value Smnt is used for all experiments: we set Smit == 3(T", , Ea), where = 1/G for all g E [G]*, the expectations 11 are sampled uniformly at random among the available examples, and E0 is the empirical covariance matrix of the N examples. Figure 3 shows the sequence of parameter estimates for the weights and the squared norm of the mean field |H l for FedEM (resp. ||H 2 for VR-FedEM) vs the number of epochs.
Federated missing values imputation for citizen science. We develop FedMissEM, a special in- stance of FedEM designed to missing values imputation in the federated setting: we apply it to the analysis of part of the eBird data base [34, 1]. citizen science smartphone application for biodiver- sity monitoring. In eBird, citizens record wildlife observations, specifying the ecological site they visited, the date, the species and the number of observed specimens. Two major challenges occur: (i) ecological sites are visited irregularly, which leads to missing values and (ii) non-professional observers have heterogeneous wildlife counting schemes.
Model and the FedMissEM algorithm. I observers participate in the programme,, there are J ccological sites and L time stamps. Each observer #i provides a / x L matrix X and a subset of indices Q2 C [J]* x [L]*. For j E [J]* and é E [L]", the variable X3, encodes the observation that would be collected by observer #i if the site #j were visited at time stamp #ł; since there are unvisited sites, we denote by Y = {Xfe, (j, f) E Qi the set of observed values and Z := {X je: (j, & 12' the set of unobserved values. The statistical model is parameterized by a matrix 0 E RJXL, where Ozz is scalar parameter characterizing the distribution of species individuals at site j and time stamp f. For instance, Oze is the log-intensity of Poisson distribution when the observations are count data the log-odd of binomial model when the observations are presence- absence data. This model could be extended to the case observers #i and #i' count different number of specimens on average at the same location and time stamp, because they do not have access to the same material or do not have the same level of expertise: heterogeneity between observers could be modeled by using different parameters for each individual #i say H J Here, we consider the case when 85 = 0,e for all (j. E) E [J]* x [L]* and i E [ We further assume that the entries (Xfe, E [4]* [J] e E [L]*} are independent with a distribution from an exponential

---

Est Esmneot brd
100 120 200
100
ES5 DESTAAOO brd ozurt 2800
1000
500
660
50
d0sematton
deservelion
Figure 2: Estimated temporal trends for Common
Figure 1: Trajectory of FedEM the number of epochs (left; blue line: [Ih((*) [?; red line: || (1-) and of VR-FedEM (right; dashed blue
Buzzard (Left) and Mallard (right). Blue crosses: es- timated monthly counts; Red triangles: number of missing values. Dotted lines: LOESS regressions for the estimated counts (blue) and the number of miss- ing values (red).
line: [h(S*)]l3 solid red line: || (11)).
03
Dauch
Opar1
Eaoch
Figure 3: [Left to right] For FedEM: Evolution of the estimates of the weights Ne for E E [G]* vs the number of epochs (first plot) and Evolution of the squared norm of the mean field | [Hg] vs the number of epochs (second plot). Then, the same things for VR-FedEM (third and fourth plots).
family with respect to some reference measure y on R of the form: x + p(z) exp(x1fe (0)e)) Algorithm 7 in Appendix F.2 provides details the model, and the pseudo-code for FedMissEM.
Application to eBird data analysis. We apply FedMissEM to the analysis of part of the eBird data base [34, 1] of field observations reported in France by 7 2, 465 observers, across J = 9, 721 sites and at = 525 monthly time points. We analyze successively two data sets corresponding to observations of two relatively common species: the Common Buzzard and the Mallard. These subsamples correspond respectively to N = 5, 980 and N 12, 185 field observations. The / field observers are randomly assigned into
10 groups (the observations of the field observers from the group E [n|]* are allocated to the server #c). For c E [n]]", server contains N. observations; in our two examples, Ne ranges between 400 and 1, 500. We run FedMissEM for 150 epochs; with = 10-4 a = 10 b = 102, a rank r = and = 0; for the distribution of the variables Xj,, we use Gaussian distribution with unknown expectation 0fO and variance 1. We recover aggregated temporal trends at the national French level for these two bird species by summing the estimated counts across ecological sites, for each time stamp; the trends are displayed in Figure 2, along with locally estimated scatterplot smoothing (LOESS).
Conclusions
We introduced FedEM which is, to the best of our knowledge, the first algorithm implementing EM in FL setting, and handles compression of exchanged information, data heterogeneity and partial par- ticipation. We further extended it to incorporate variance reduction scheme, yielding VR-FedEM. We derived complexity bounds which highlight the efficiency of the two algorithms, and illustrated our claims with numerical simulations, as well as an application to biodiversity monitoring data. In simultancously published work, Marfoq et al. [25] consider a different Federated EM algorithm, in order to address the personalization challenge by considering a mixture model. Under the as- sumption that each local data distribution is a mixture of unknown underlying distributions, their algorithm computes a model corresponding to each distribution. On the other hand, we focus on the curved exponential family, with variance reduction, partial participation and compression and on limiting the impact of heterogeneity, but do not address personalization.
Acknowledgments The work of A. Dieuleveut and E. Moulines is partially supported by ANR- 19-CH-A-0002-01 /chaire SCAI, and Hi!Paris. The work G. Fort partially supported by the Fondation Simone et Cino del Duca under the project OpSiMorE.
10

---

Broader Impact of this work This work is mostly theoretical, and we believe it does not currently present any direct societal consequence. However, the methods described this paper can be used to train machine learning models which could themselves have societal consequences. For instance, the deployment of machine learning models can suffer from gender and racial bias, or amplify existing inequalities.

---

References
[1] ebird. 2017. ebird: An online database of bird distribution and abundance [web application]- ebird, cornell lab of omithology,, ithaca, new york. available: http://www.ebird.org (accessed: 21 march 2020).
[2] D. Alistarh, T. Hoefler, M. Johansson, N. Konstantinov, S. Khirirat, and C. Renggli, The con- vergence of sparsified gradient methods. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31, pages 5973-5983. Curran Associates, Inc., 2018.
[3] A. Benveniste, M. Métivier, and P. Priouret. Adaptive Algorithms and Stochastic Approxima- tions. Springer Verlag, 1990.
[4] V. S. Borkar. Stochastic approximation. Cambridge University Press, Cambridge; Hindustan Book Agency, New Delhi, 2008, A dynamical systems viewpoint.
[5] J. Chen, J. Zhu, Y. Tch, and T. Zhang. Stochastic expectation maximization with variance re- duction. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Biance and R. Gar- nett, editors, Advances in Neural Information Processing Systems 31, pages 7967-7977. 2018.
[6] A. Defazio, F. Bach, and S. Lacoste-Julien. SAGA: A Fast Incremental Gradient Method With Support for Non-Strongly Convex Composite Objectives. In Ghahramani, M. Welling, C. Cortes, N. D. Lawrence, and K. Q. Weinberger, editors, Advances in Neural Information Processing Systems 27, pages 1646-1654. Curran Associates, Inc., 2014.
[7] B. Delyon, M. Lavielle, and E. Moulines. Convergence of a stochastic approximation version of the EM algorithm. Ann. Statist., 27(1):94-128, 1999.
[8] A. Dempster, N. Laird, and D. Rubin. Maximum Likelihood from Incomplete Data via the EM Algorithm. J. Roy. Stat. Soc. B Met., 39(1):1-13)) 1977.
[9] C. Fang, C. Li, Z. Lin, and T. Zhang. SPIDER: Near- Pear-Optimal Non-Convex Optimization via Stochastic Path-Integrated Differential Estimator. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Dasa-BBinnchi, and R. Garnett, editors, Advances in Neural Information Pro- cessing Systems 31, pages 689-699, Curran Associates, Inc., 2018.
[10] G. Fort, E. Moulines, and H.-T. Wai. A Stochastic Path Integral Differential EstimatoR Ex- pectation Maximization Algorithm. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 16972-16988 Curran Associates, Inc., 2020.
[11] G. Fort, P. Gach, and F. Moulines. Fast Incremental Expectation Maximization for finite-sum optimization: non asymptotic convergence. Statistics and Computing, 2021. Accepted for publication.
[12] G. Fort, E. Moulines, and H.-T. Wai. Ccom-SPPDER--ER Faster Variance Reduced Stochastic Expectation Maximization for Nonconvex Finite-Sum Optimization. In IEEE International Conference on Acoustics, Speech and Signal Processing, 2021.
[13) S. Frühwirth-Schnateer G. Celeux, and C. P. Robert, editors. Handbook of mixture analysis. Chapman & Hall/CRC Handbooks of Modern Statistical Methods. CRC Press, Boca Raton, FL, 2019.
[14] S. Ghadimi and G. Lan. Stochastic First- and Programming. SIAM J. Optim., 23(4):2341-2368, 2013.
Methods for Nonconvex Stochastic
[15] E. Gorbunov, F. Hanzely, and P. Richtárik. A unified theory of SGD: Variance reduction, sampling, quantization and coordinate descent. In International Conference on Artificial Intel- ligence and Statistics, pages 680-690. PMLR, 2020.
[16] Horváth and Richtarik. better alternative to error feedback for communication--effcceee distributed learning. In International Conference on Learning Representations, 2021.
12

---

[17] S. Horváth, D. Kovalev, K. Mishchenko, S. Stich, and P. Richtárik. Stochastic distributed learning with gradient quantization and variance reduction. arXiv preprint arXiv:1904.005 1904.0515,, 2019.
[18] R. Johnson and Т. Zhang. Accelerating Stochastic Gradient Descent using Predictive Variance Reduction. In C. J. C. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K. Q. Weinberger, editors, Advances in Neural Information Processing Systems 26, pages 315-323. Curran As- sociates, Inc., 2013.
[19] P. Kairouz, H. B. McMahan, B. Avent, A. Bellet, M. Bennis, A. N. Bhagoji. K. Bonawit, Z. Charles, G. Cormode, R. Cummings, R. G. L. D'Oliveira, H. Eichner, S. El Rouayheb, D. Evans, J. Gardner, Z. Garrett, A. Gascón, B. Ghazi, P. B. Gibbons, M. Gruteser, Z. Har- chaoui, C. He, L. He, Z. Huo, Hutchinson, J. Hsu, M. Jaggi, T. Javidi, G. Joshi, M. Khodak, J. Konecný. A. Korolova, F. Koushanfar, S. Koyejo, T. Lepoint, Y. Liu, P. Mittal, M. Mohri, Qi, D. Ramage, R. Raskar, M. Raykova, S. U. Stich, Z. Sun, A. Theertha Suresh, F. Tramèr, P. Vepakomma, J. Wang, L. Xiong, 7. Xu, Q. Yang, F. X. Yu, H. Yu, and S. Zhao. Advances and Open Problems in Federated Learning- Now Foundations and Trends.
R. Nock, Özgür, R. Pagh,
Song, W. Song,
[20] B. Karimi, H.-T. Wai, E. Moulines, and M. Lavielle. On the Global Convergence of (Fast) In- cremental Expectation Maximization Methods. In H. Wallach, H. Larochelle, A. Beygelzimer, d' Alché Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems 32, pages 2837-2847 Curran Associates, Inc., 2019.
[21] S. P. Karimireddy, S. Kale, M. Mohri, S. Reddi, S. Stich, and A. T. Suresh. SCAFFOLD: Stochastic controlled averaging for federated learning. In H. D. III and Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Pro- ceedings of Machine Learning Research, pages 5132-5143, PMLR, 13-18 Jul 2020. URL https: httpps:/prooeeednggg mlr. www..press/vvve http://aaa/i/rrdrrr.e.aa
html.
[22] Konečný, H. B. McMahan, F. X. Yu, P. Richtárik, A. T. Suresh, and D. Bacon. Federated learning: Strategies for improving communication efficiency. arXiv preprint arXiv:1610.05492,2000 2016.
[23] F. Kunstner, R. Kumar, and M. Schmidt. Homeomorphic-ivvaiiane of em: Non-asymptotic convergence in kl divergence for exponential families via mirror descent. In International Conference on Artificial Intelligence and Statistics, pages 3295-3303, PMLR, 2021.
[24] K. Lange. MM Optimization Algorithms. SIAM-Society for Industrial and Applied Mathemat- ics, 2016.
[25] O. Marfoq, G. Neglia, A. Bellet, L. Kameni, and R. Vidal. Federated multi-task learning under a mixture of distributions. 35th Conference on Neural Information Processing Systems (NeurIPS 2021), 2021.
[26] G. McLachlan and T. Krishnan. The EM algorithm and extensions. Wiley series in probability and statistics. Wiley, 2008
[27] B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. A. y. Arcas. Communication- Efficient Learning of Deep Networks from Decentralized Data. Artificial Intelligence and Statistics, pages 1273-1282. PMLR, 2017.
[28] K. Mishchenko, E. Gorbunov, M. Takáč, and P. Richtárik. Distributed learning with com- pressed gradient differences. arXiv preprint 19X1009900.
[29] K. Murphy and S. J. Russell. Dynamic bayesian networks: representation, inference and learn- ing. 2002.
[30] L. M. Nguyen, J. Liu, K. Scheinberg, and M. Takáč. Sarah: A novel method for machine learning problems using stochastic recursive gradient. Proceedings of the 34th International Conference Machine Learning Volume 70, ICML' 17, page 2613-2621. JMLR.org. 2017.
13

---

[31] C. Philippenko and A. Dieuleveut. Bidirectional compression in heteregenous settings for dis- tributed federated learning with partial participation: tight convergence guarantees. Techni- cal report, arXiv 2006.14591v 14591v3, 2020.
[32] C. Philippenko and A. Dieuleveut. Preserved central model for faster bidirectional compression in distributed settings. Advances in Neural Information Processing Systems (NeurlPS), 2021.
[33] Sattler, S. Wiedemann, K.-R. Müller, and W. Samek. Robust and Communication Efficcient Federated Learning From Non-i.i.. Data. IEEE Transactions on Neural Networks and Learn- ing Systems, pages 1-14, 2019.
[34] B. L. Sullivan, C. L. Wood, M. J. Iliff, R. E. Bonney, D. Fink, and S. Kelling, eBird: A citizen- based bird observation network in the biological sciences. Biological Conservation, 142(10): 2282 2292, 2009.
[35] H. Tang, C. Yu, X. Lian, T. Zhang. and J. Liu. DoubleSqueeze: Parallel Stochastic Gradient Descent with Double-pass past Error-Compern Compenatmee Compression. In International Conference on Machine Learning, pages 6155-6165. PMLR, May 2019. ISSN: 2640-3498.
[36] Z. Wang. K. Ji, Y. Zhou, Y. Liang, and V. Tarokh. SpiderBoost and Momentum: Faster Stochas- tic Variance Reduction Algorithms. In H. Wallach, H. Larochelle, A. Beygelzimer, d' Alché- Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems 32, pages 2406-24622 2416. 2019.
[37] B. Woodworth, K. K. Patel, S. Stich, Z. Dai, B. Bullins, B. Mcmahan, O. Shamir, and N. Sre- bro. Is local SGD better than minibatch SGD? H. D. III and A. Singh, editors, Pro- ceedings of the 37th International Conference on Machine Learning, volume 119 of Pro- ceedings of Machine Learning Research, pages 10334-10343, PMLR, 13-18 Jul 2020. URL https: pttps://proceedingg. mlr presss/v1$//ooooortt200. hhtmm.
[38] I.. Xu and M. I. Jordan. On convergence properties of the EM algorithm for Gaussian mixtures. Neural computation, 8(1):129-151, 1996.
[39] Q. Yang, Y. Liu, T. Chen, and Y. Tong. Federated machine learning: Concept and applications. ACM Transactions Intelligent Systems and Technology (TIST), 10(2):1-19, 2019.
14


Reviews:
[{"review_id": "yLNsaH6DxOi", "reviewer": "Reviewer_nRvM", "summary": "This paper investigates how to federate the Expectation Maximization (EM) algorithm in the case of a curved exponential family.\nThe authors propose a novel algorithm, called FedEM, which translates the EM algorithm in the FL setting, which consists in iteratively sending local sufficient statistics to the server and updating the value of the model. This algorithm takes into account the distributed aspects, communication constraints (using an unbiased compressor with an error feedback mechanism for the local statistics to send), partial participation, and data heterogeneity of the FL setting.\nThe authors prove the convergence of FedEM (Corollary 2).\nA variance-reduced version of FedEM is also proposed (VR-FedEM), also with convergence guarantees (Th. 4).\nExperiments are performed on toy synthetic data, MNIST and real data (missing value imputation from a tabular dataset, bird).", "questions": "", "limitations": "", "rating": 8, "confidence": 1, "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": "## Originality\n\nTo the best of my knowledge, this is the first work tackling the federation of the EM algorithm, which makes it novel.\n\n## Clarity \n\nOverall, the paper is very well-written and easy to follow even for non-experts. However, in the experimental section (page 8), it is not stated what figure 2 refers to - I assume it is the MNIST experiments given the number of components estimated.\n\n## Quality\n\nRegarding the theoretical results, I did not check in detail the proofs, so I cannot judge their correctness.\n\nIf any, the experimental section could be further expanded.\n- For instance, the author state that a naïve implementation of the EM algorithm in FL would not be resilient to heterogeneity, but there is no experimental demonstration of this assertion.\n- Further, there is no study of the effect of compression\n- Last, there is no comparison with the results obtained in a single node setting: does one retrieve results which are close to them? \n\n## Significance\n\nThe proposed algorithms and the convergence proofs are the main contributions of the paper. While the experimental section is a bit weak, I do think this is a minor problem given the other contributions.\n\n\n========\nUpdate following rebuttal\n========\n\nI thank the authors for clearly addressing all my criticisms in their answer, and for the time and precision they devoted to answer the other reviewers. I will raise my score to 8.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "See above", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}, {"review_id": "kT4qbzLucs", "reviewer": "Reviewer_WGsy", "summary": "The paper proposes methods to implement Expectation-Maximization (EM) algorithm under Federated Learning (FL) setting. The methods proposed incorporate gradient compression and variance reduction techniques, and also allow for partial participation during training process. Moreover, the paper gives finite-time complexity analysis of proposed methods. Finally, paper demonstrates their methods by both simulated and real data experiments.", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "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": "While the problem solved by the paper is important and interesting, the way it is solved in the paper seems straightforward. More specifically, the EM algorithm is well-studied, and both of the gradient compression and variance reduction techniques used to implement EM for FL setting are borrowed from previous literatures. The theoretical analysis, while it is appreciated, looks like adaptation from existing papers. All in all, I do not see much novelty in this paper.\n\nThe paper is technically correct. The paper gives finite-time complexity analysis of proposed methods. The paper demonstrates their methods by both simulated and real data experiments. The work is overall complete.\n\nThe submission is well-organized. The methods are clearly motivated and explained. The proof is clear and complete as well.\n\nThe problem studied by the paper will be interesting for many practitioners. However, the paper is more like a combination of existing works in to guide engineering, rather than a novel research. There is limited unique results in the work. Thus, overall, it is of limited significance. I think at least paper should make it more clear about what is new in this work.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The paper discussed about its potential negative societal impact. The paper could perhaps think more about whether EM under FL can protect privacy.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "kEQQQ50NMu0", "reviewer": "Reviewer_XoK6", "summary": "Authors propose and study a distributed variant of Expectation-Maximization for exponential families. Theoretical analysis is conducted to characterize the convergence of the proposed algorithms towards stationary points of the underlying objective function. The empirical performance of proposed algorithms is studied via numerical experiments on synthetic and open data. \n", "questions": "", "limitations": "", "rating": 3, "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": "* It is unclear how data heterogeneity is modeled and handled by the proposed algorithms. As i understood, the main model underlying the proposed algorithms are homogeneous i.i.d. draws from a single exponential famliy Eq. (2). How would your FedEM algorithm cope with completely heterogeneous data which e.g. conforms with a networked exponential family as studied in \n\nA. Jung, \"Networked Exponential Families for Big Data Over Networks,\" in IEEE Access, vol. 8, pp. 202897-202909, 2020, doi: 10.1109/ACCESS.2020.3033817.\n\n* It would help the reader to have a dedicated section or paragraph about existing work and how the proposed methods related to it. Also, Section 2 could be dedicated to the precise problem statement including the communication link constraints. \n\n* It would be good to have a dedidacted paragraph in Section 1 that summarizes the paper's contributions. This paragraph should also discuss the relation between FedEM and VR-FedEM. Is the latter a generalization of the former? Do both algorithms apply in different settings? \n\n* Pls add more discussion on how to choose the input parameters of Algorithm 1 and Algorithm 2 \n\n* The numerical experiments need to be expanded and also compare with other methods, e.g., network Lasso methods such as studied in \n\nSarcheshmehPour, Y., Tian, Y., Zhang, L., and Jung, A., “Networked Federated Multi-Task Learning”, <i>arXiv e-prints</i>, 2021. https://arxiv.org/abs/2105.12769\n\n* Pls discuss if the Assumptions 1 - 9 are satisfied for the datasets and models considered in the numerical experiments \n\n* There needs to be more discussion/analysis of the obtained numerical results from the perspective of the theoretical results. Do the plots in Figure 1 confirm the statements/insights provided by your Theorems ?  \n\n* I did not find any discussion/reference to Figure 2 in the main text. \n\n* it is mentioned that \"A major advantage of the EM algorithm stems from its invariance under homeomorphisms,...\" However, it is not clear to me how this invariance is relevant for your methods?\n\n* unclear what is meant by \"..where heterogeneity is characterized by the non-cancellation of the local gradients at the optimum,..\"\n\n* unclear what is meant by \"..conditions must therefore be derived to ensure convergence at the central server.\" which conditions are you referring to here?\n\n* \"...the complexity grows at the rate...\"  which complexity do you mean? computational complexity or sample complexity?\n\n* \"...see [16] for an interesting discussion on the connection of EM and mirror descent.\" how is this connection relevant for the following developments?\n\n* could you provide more motivation/intuition for introducing the concept of a mean-field in Eq. (6) ? does your convergence analysis make heavy use of this mean-field?\n\n* \"...that the conditional variance is (uniformly) bounded\" the conditional variance of what quantity?\n\n* \"... of the new sufficient statistics.\" not clear what the \"new sufficient statistic\" is \n\n* \"Crucially, a naive implementation of compressed distributed EM without assumptions on the heterogeneity would not converge.\" maybe add a reference to a proof for this statement \n\n* pls add more discussion on the consequences of Proposition 3. how strong does partial participation hurt the performance compared to having full participation?\n\n\n\n\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "this is theoretic work and therefore the societal impact can only be discussed on a rather high level. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "9rB5dtitBZZ", "reviewer": "Reviewer_at3L", "summary": "The authors study the problem of expectation-maximization in a federated setting. They propose two communication-efficient algorithms where the clients compress certain sufficient statistics to save communication. For each algorithm, convergence rates are presented that match their centralized counterparts in the absence of compression. ", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "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": "Originality and Pros: Both the problem studied, and the algorithms presented, are novel to the best of my knowledge. The compression strategies used are somewhat unique to the problem being studied which makes them interesting. Concrete analytical results are presented to support the proposed algorithms. \n\nClarity: The paper is mostly well-written, but certain parts could do with better explanations. In what follows, I explain these concerns in detail.\n\n(1) In the standard FL setting, statistical heterogeneity causes the local loss functions of the clients to have potentially different minima. As a result, infrequent communication leads to \"client-drift\" which, in turn, can cause convergence to incorrect points. For the federated EM problem the authors' study, I found it hard to understand what heterogeneity means (despite the explanation in lines 58-59), and what the implications of heterogeneity are for convergence. A concrete mathematical description of heterogeneity would be very helpful in this context.\n\n(2) One of the major concerns in FL is to reduce the number of communication rounds. To do so, almost all standard FL algorithms I am aware of (e.g., FedAvg, FedProx, SCAFFOLD, FedSplit, etc. ) involve periodic communication between the clients and the server. In between two communication rounds, clients perform multiple local computations on their private data. In contrast, if I understand correctly, both the algorithms proposed in the paper involve communication between the clients and the server *at every iteration*. This is quite undesirable since the number of communication rounds can be significantly higher relative to existing FL algorithms.\n\nThe authors need to discuss this potential limitation of their approach. How would periodic interactions impact the convergence guarantees of their proposed algorithms?\n\n(3) Continuing with the above point, it is precisely due to infrequent communication that data heterogeneity creates a problem in standard FL settings. Indeed, even in the presence of heterogeneity, if the clients exchanged model updates with the server at every iteration, the iterates of all clients would evolve synchronously, i.e., this would essentially be equivalent to a centralized setting. Said differently, if the communication period is 1, then the \"client-drift\" phenomenon that arises due to data heterogeneity in common FL algorithms would cease to exist. \n\nSince the authors do not consider periodic communication, how (if at all) does heterogeneity create a challenge? To make my question clearer, suppose there is no compression/quantization. Moreover, suppose all clients participate in every round of communication. Even in these scenarios, data heterogeneity leads to challenges in the analysis of common FL algorithms such as FedAvg and FedProx, precisely due to infrequent, periodic communication. In contrast, I would imagine that if there is no compression, and full client participation at all times, Fed-EM would essentially behave as a standard centralized EM algorithm. Isn't this correct? \n\nThe point I am trying to make here is that not considering infrequent communication takes away the typical challenges posed by data heterogeneity in standard FL settings. This point needs to be acknowledged and discussed.\n\n(4) The authors do not cite or discuss any of the popular FL algorithms such as FedAvg, FedProx, or SCAFFOLD (to name a few). I would encourage the authors to compare their model of heterogeneity and the communication structure of their algorithm to the ones in the following references.\n\n[R1] \"Tighter theory for local SGD on identical and heterogeneous data\", Khaled et al, AISTATS 20.\n\n[R2] \"On the convergence of fedavg on non-iid data\", Li et al., arXiv 19.\n\n[R3] \"On the convergence of federated optimization in heterogeneous networks\", Sahu et al., arXiv 18.\n\n[R4] \"Scaffold: Stochastic controlled averaging for federated learning\", Karimireddy et al., ICML 20.\n\n(5) Regarding Section 2, the authors list several assumptions without commenting on how strong these assumptions are. Are these assumptions standard in the analysis of EM algorithms? How do they compare to the typical assumptions of convexity and smoothness in standard FL optimization settings? (such as those in [R1]-[R4])\n\n(6) A couple of questions regarding Algorithm 1: (i) Is the operator T known to all clients? (ii) Lines 6-8 appear to be inspired by the \"error-feedback\" mechanism employed when one considers compression/sparsification of gradients. If so, this should be stated clearly. (iii) In line 5, what does it mean to sample a sufficient statistic S_{k+1,i}? \n\n(7) (i) Although the variance reduction method in Algorithm 2 appears to be theoretically sound, it is not adequately motivated as I explain below. Typically in FL, variance reduction or gradient tracking is used to reduce the \"client-drift\" phenomenon that arises due to data heterogeneity and infrequent communication (see [R4] above and [R6] below). Since the authors do not clearly explain what challenge heterogeneity poses for their specific problem, I found it hard to appreciate the need for Algorithm 2.\n\n(ii) Unlike Algorithm 1, Algo. 2 involves both outer and inner loops. However, no explanation is provided for the rationale behind this. \n\n(iii) How does variance reduction improve (if at all) the convergence guarantees relative to Algorithm 1? Does this come at the expense of more communication? It would be instructive if the authors clearly compared the complexity bounds of Theorem 4 with those of Theorem 1. \n\n(iv) The authors should also discuss how their variance reduction approach compares with those used in standard FL settings. In this context, the most relevant works are [R4] above and [R5], [R6] below.\n\n[R5] \"Federated Learning with Compression: Unified Analysis and Sharp Guarantees\", Haddadpour et al., AISTATS 21.\n\n[R6] \"Achieving Linear Convergence in Federated Learning under Objective and Systems Heterogeneity\", Mitra et al., arXiv 21.\n\nReference [R6] in particular employs both gradient tracking and error-feedback to tackle statistical heterogeneity and compression/sparsification, just as in VR-FedEM. Appropriate comparisons are thus warranted.\n\n========================== REVIEW UPDATE =======================================\n\nI have read the response of the authors and also the comments of the other Reviewers. The rebuttal by the authors has improved my understanding of the paper. So I am increasing my score to a 7. \n\n\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors briefly discuss some limitations and societal impacts of their work in Section 5.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4.5 hours"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers appreciate the incorporation of EM algorithm into federated learning for dealing with data heterogeneity. The paper makes good theoretical contributions but also has weaknesses in experiments. Overall, I am in favor of acceptance. Please incorporate the new results in the final version. Please revise the paper based on the reviews and rebuttals.
Decision: Accept (Poster)