Federated-EM with heterogeneity mitigation and variance reduction

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 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 is a new communication efficient method , which handles partial participation of local devices , and is robust to heterogeneous distributions 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 . 1 Introduction . The Expectation Maximization ( EM ) algorithm is the most popular approach for inference in latent variable models . The EM algorithm , a 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 the 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 a large number of edge devices to perform distributed machine learning , is 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 ( NeurIPS 2021 ) . is to minimize a non-convex finite-sum smooth objective function Argmin✓2⇥ F ( ✓ ) , F ( ✓ ) : = 1 n nX i=1 Li ( ✓ ) + R ( ✓ ) , ⇥ ✓ Rd , ( 1 ) 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 promote sparsity , regularity , etc . In latent variable models , Li ( ✓ ) = m 1 Pm j=1 log p ( yij ; ✓ ) , where { yij } mj=1 are the m observations available for worker # i , and p ( y ; ✓ ) is the incomplete likelihood . p ( y ; ✓ ) is defined by marginalizing the complete-data likelihood p ( y , z ; ✓ ) defined as the joint probability density function of the observation y and a non-observed latent variable z 2 Z , i.e . p ( y ; ✓ ) = R Z p ( y , z ; ✓ ) µ ( dz ) where Z is the latent space and µ is a measure on Z . We focus in this paper on the case where p ( y , z ; ✓ ) belongs to a curved exponential family , given by p ( y , z ; ✓ ) : = ⇢ ( y , z ) exp hs ( y , z ) , ( ✓ ) i ( ✓ ) ; ( 2 ) where s ( y , z ) 2 Rq is the complete-data sufficient statistics , : ⇥ ! Rq and : ⇥ ! R , ⇢ : Y ⇥ Z ! R+ are vector/scalar functions . 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 ✓curr , it computes a majorizing function ✓ 7 ! Q ( ✓ , ✓curr ) given up to an additive constant by Q ( ✓ , ✓curr ) : = hs̄ ( ✓curr ) , ( ✓ ) i+ ( ✓ ) + R ( ✓ ) where s̄ ( ✓ ) : = 1 n nX i=1 s̄i ( ✓ ) ; ( 3 ) and s̄i ( ✓ ) is the ith device conditional expectation of the complete-data sufficient statistics : s̄i ( ✓ ) : = 1 m mX j=1 s̄ij ( ✓ ) , s̄ij ( ✓ ) : = Z Z s ( yij , z ) p ( z|yij ; ✓ ) µ ( dz ) , ( 4 ) where p ( z|yij ; ✓ ) : = p ( yij , z ; ✓ ) /p ( yij ; ✓ ) . As for the M step , an updated value of ✓curr is computed as a minimizer of ✓ 7 ! Q ( ✓ , ✓curr ) . The majorizing function is then updated with the new ✓curr ; this process is iterated until convergence . The EM algorithm is most useful when for any ✓curr 2 ⇥ , the function ✓ 7 ! Q ( ✓ , ✓curr ) is a convex function of the parameter ✓ 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-parametrization [ 23 ] . In the FL context , the vanilla EM algorithm is affected by three major problems : ( 1 ) the communication 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 classical 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 are not appropriate to 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 an 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 2 { 1 , . . . , n } and the observations yij to be independent but not necessarily identically distributed . In addition , our theoretical results deal with specific behaviors for each worker i 2 { 1 , . . . , n } , see e.g. , A5 , 7 and 8 . In the FL-EM setting , heterogeneity manifests itself by the non-equality of the local conditional expectations of the complete-data sufficient statistics s̄i ’ s ; modifications to the algorithms must be performed to ensure convergence at the central server . Finally , a subset of users are potentially inactive in each learning round , being unavailable or unwilling to participate . Thus , taking into account PP of the workers and its impact on the convergence of algorithms , is a major issue . • FedEM . The main contribution of our paper is a new method called FedEM , supporting communication compression , partial participation and data heterogeneity . In this algorithm , the workers compute an estimate of the local complete-data sufficient statistics s̄i using a minibatch of data , apply an unbiased compression operator to a noise compensated version ( using a 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 a variance reduction method inspired by the SPIDER framework [ 9 ] which has recently been extended to 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 additionally 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 a function h ( see e.g . Section 2 ) . We introduce a unified theoretical framework which covers the convergence of FedEM and VR-FedEM algorithms in the non-convex case and establish convergence guarantees for finding an ✏-stationary point ( see Theorem 1 and Theorem 3 ) . In both cases , we provide the number Kopt ( ✏ ) of optimization steps and the number KCE ( ✏ ) of computed conditional expectations to reach ✏-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 KCE ( ✏ ) improves on earlier results . In this respect , VR-FedEM has the same advantages as SPIDER [ 9 ] compared to SVRG [ 18 ] and SAGA [ 6 ] , or as 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 ✏ . Notations . For vectors a , b in Rq , ha , bi is the Euclidean scalar product , and k·k denotes the associated norm . For r 1 , kakr is the ` r-norm of a vector a . 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 kAkF is its Frobenius norm ; for two matrices A , B , hA , Bi : = Trace ( B > A ) . For a positive integer n , set [ n ] ? : = { 1 , · · · , n } and [ n ] : = { 0 , · · · , n } . The set of non-negative integers ( resp . positive ) is denoted by N ( resp . N ? ) . The minimum ( resp . maximum ) of two real numbers a , b is denoted by a^b ( resp . a_b ) . We will use the Bachmann-Landau notation a ( x ) = O ( b ( x ) ) to characterize an upper bound of the growth rate of a ( x ) as being b ( x ) . 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 A1 to A3 , which define the model at hand . A1 . The parameter set ⇥ ✓ Rd is a convex open set . The functions R : ⇥ ! R , : ⇥ ! Rq , : ⇥ ! R , and ⇢ ( yij , · ) : Z ! R+ , s ( yij , · ) : Z ! Rq for i 2 [ n ] ? and j 2 [ m ] ? are measurable functions . For any ✓ 2 ⇥ and i 2 [ n ] ? , the log-likelihood is finite : 1 < Li ( ✓ ) < 1 . A2 . For all ✓ 2 ⇥ and i 2 [ n ] ? , the conditional expectation s̄i ( ✓ ) is well-defined . A3 . For any s 2 Rq , the map s 7 ! Argmin✓2⇥ { ( ✓ ) + R ( ✓ ) hs , ( ✓ ) i } exists and is unique ; the singleton is denoted by { T ( s ) } . EM defines a sequence { ✓k , k 0 } that can be computed recursively as ✓k+1 = T s̄ ( ✓k ) , where the map T is defined in A3 and s̄ is defined in ( 3 ) . On the other hand , the EM algorithm can be defined through a mapping in the complete-data sufficient statistics , referred to as the expectation space . In this setting , the EM iteration defines a Rq-valued sequence { bSk , k 0 } given by bSk+1 = s̄ T ( bSk ) . Thus , we observe that the EM algorithm admits two equivalent representations : ( Parameter space ) ✓k+1 = T s̄ ( ✓k ) ; ( Expectation space ) bSk+1 = s̄ T ( bSk ) . ( 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 a fixed point to the EM algorithm in the expectation space , then ✓ ? : = T ( s ? ) is a fixed point of the EM algorithm in the parameter space , i.e. , ✓ ? = T s̄ ( ✓ ? ) ; note that the converse is also true . Define the functions hi and h from Rq to Rq by h ( s ) : = 1n Pn i=1 hi ( s ) with hi ( s ) : = s̄i T ( s ) s . h ( s ) : = 1 n nX i=1 hi ( s ) , hi ( s ) : = s̄i 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 s 7 ! h ( s ) ( see ( 3 ) for the definition of s̄ ) . Therefore , convergence of EM-based algorithms is evaluated in terms of ✏-stationarity ( see [ 14 , 10 ] ) : for all ✏ > 0 , there exists a ( possibly random ) termination time K s.t . : E h kh ( bSK ) k2 i ✏ . Another key property of EM is that it is a monotonic algorithm : each iteration leads to a decrease of the negative penalized log-likelihood i.e . F ( ✓k+1 ) F ( ✓k ) or , equivalently in the expectation space F T ( bSk+1 ) F T ( bSk ) ( for sequences { ✓k , k 0 } and { bSk , k 0 } given by ( 5 ) ) . A4 assumes that the roots of the mean field h are the roots of the gradient of F T ( see [ 7 ] for the same assumption when studying Stochastic EM ) . A5 assumes global Lipschitz properties of the functions hi ’ s . A4 . The function W : = F T : Rq ! R is continuously differentiable on Rq and its gradient is globally Lipschitz with constant L Ẇ . Furthermore , for any s 2 Rq , rW ( s ) = B ( s ) h ( s ) where B ( s ) is a q ⇥ q positive definite matrix . In addition , there exist 0 < vmin vmax such that for any s 2 Rq , the spectrum of B ( s ) is in [ vmin , vmax ] . A5 . For any i 2 [ n ] ? , there exists Li > 0 such that for any s , s0 2 Rq , khi ( s ) hi ( s0 ) k = k ( s̄i T ( s ) s ) ( s̄i T ( s0 ) s0 ) k Liks s0k . A Federated EM algorithm . Algorithm 1 : FedEM with partial participation Data : kmax 2 N ? ; for i 2 [ n ] ? , V0 , i 2 Rq ; bS0 2 Rq ; a positive sequence { k+1 , k 2 [ kmax 1 ] } ; ↵ > 0 ; a coefficient p = EA⇠PPP [ card ( A ) ] /n . Result : The FedEM-PP sequence : { bSk , k 2 [ kmax ] } 1 Set V0 = n 1 Pn i=1 V0 , i 2 for k = 0 , . . . , kmax 1 do 3 Sample Ak+1 ⇠ PPP 4 for i 2 Ak+1 do 5 ( worker # i ) 6 Sample Sk+1 , i , an approximation of s̄i T ( bSk ) 7 Set k+1 , i = Sk+1 , i Vk , i bSk 8 Set Vk+1 , i = Vk , i + ↵Quant ( k+1 , i ) . 9 Send Quant ( k+1 , i ) to the central server 10 for i /2 Ak+1 do 11 ( worker # i ) 12 Set Vk+1 , i = Vk , i ( no update ) 13 ( the central server ) 14 Set Hk+1 = Vk + ( np ) 1 P i2Ak+1 Quant ( k+1 , i ) 15 Set bSk+1 = bSk + k+1Hk+1 Set Vk+1 = Vk + ↵n 1 P i2Ak+1 Quant ( k+1 , i ) Send bSk+1 and T ( bSk+1 ) to the n workers Our first contribution , the novel algorithm FedEM is described by algorithm 1 . The algorithm encompasses partial participation of the workers : at iteration # ( k + 1 ) , only a subset Ak+1 of active workers participate to the training , see line 3 . The averaged fraction of participating workers is denoted p. Each of the active workers # i computes an unbiased approximation Sk+1 , i ( line 6 ) of s̄i T ( bSk ) ; conditionally to the past ( see Appendix D.2 for a rigorous definition ) , these approximations are independent . The workers then transmit to the central server a compressed information about the new sufficient statistics . A naive solution would be to compress and transmit Sk+1 , i bSk , but data heterogeneity between servers often prevents these local differences from vanishing at the optimum , leading to large compression errors and impairing convergence of the algorithm . Following [ 28 ] , a memory Vk , i ( initialized to hi ( bS0 ) at k = 0 ) is introduced ; and the differences k+1 , i : = Sk+1 , i Ŝk Vk , i are compressed for i 2 Ak+1 ( line 7 and line 9 ) . These memories are updated locally : Vk+1 , i = Vk , i + ↵Quant ( k+1 , i ) , at line 8 , with ↵ > 0 ( typically set to 1/ ( 1+ ! ) where ! is defined in A6 ) . On its side , the central server releases an aggregated estimate bSk+1 of the complete-data sufficient statistics by averaging the quantized difference ( np ) 1 P i2Ak+1 Quant ( k+1 , i ) and by adding Vk ( line 14 and line 15 ) . Then , it updates Vk+1 = Vk +↵n 1 Pn i=1 Quant ( k+1 , i ) , see line 15 . The final step consists in solving the M-step of the EM algorithm , i.e . in computing T ( bSk+1 ) ( see A3 ) . We finally state our assumption on the compression process . We consider a large class of unbiased compression operators Quant satisfying a variance bound : A6 . There exists ! 0 s.t . for any s 2 Rq : E [ Quant ( s ) ] = s , and E ⇥ kQuant ( s ) k2 ⇤ ( 1+ ! ) ksk2 . Intuitively , the stronger the compression is , the larger ! will be . Remark that if no compression is used , or equivalently for all s 2 Rq , Quant ( s ) = s , then A6 is satisfied with ! = 0 . An example of quantization operator satisfying A6 is the random dithering that can be described as the random operator Quant : Rq ! Rq , Quant ( x ) = ( 1/squant ) kxkr sign ( x ) bsquant ( |x|/kxkr ) + ⇠c where r 1 is user-defined , ⇠ is a uniform random variable on [ 0 , 1 ] q and squant 2 N ? is the number of levels of roundings ; see [ 17 , 2 ] . This operator satisfies A6 with ! = s 1 quant O ( q1/r + q1/2 ) ; see [ 17 , Example 1 ] . Another example , namely the block-p-quantization , 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 ] . The convergence analysis is under the following assumptions on the oracle Sk+1 , i : for any i 2 [ n ] ? , the approximations Sk+1 , i are unbiased and their conditional variances are uniformly bounded in k. For each k 2 N , denote by Fk the -algebra generated by { S ` , i , A ` ; i 2 [ n ] ? , ` 2 [ k ] } and including the randomness inherited from the quantization operator Quant up to iteration # k. A 7 . For all k 2 N , conditional to Fk , { Sk+1 , i } ni=1 are independent . Moreover , for any i 2 [ n ] ? , E [ Sk+1 , i|Fk ] = s̄i T ( bSk ) and there exists 2i > 0 such that for any k 0 E h kSk+1 , i s̄i T ( bSk ) k2 Fk i 2i . A7 covers both the finite-sum setting described in the introduction , and the online setting . In the finite-sum setting , s̄i is of the form m 1 Pm j=1 s̄ij . In that case , Sk+1 , i can be the sum over a minibatch Bk+1 , i of size b sampled at random in [ m ] ? , with or without replacement and independently of the history of the algorithm : we have Sk+1 , i = b 1 P j2Bk+1 , i s̄ij T ( bSk ) . In the online setting , the oracles Sk+1 , i come from an online processing of streaming informations ; in that case Sk+1 , i can be computed from a minibatch of independent examples so that the conditional variance 2i , which will be inversely proportional to 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 techniques 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 communication rounds , ( iii ) using compression , ( iv ) sampling clients at each step . Here , we provide a tight analysis of strategies ( i ) , ( iii ) and ( iv ) ( sampling client is part of PP ) . Regarding the interest of performing multiple iterations ( ii ) , as analyzed for example in [ 21 , 27 ] for the classical gradient settings , we note that : first , from a 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 a fundamental specificity of the EM framework ( which is not shared by the gradient framework ) . In applications , we usually do not want T to be available at each local node . However , our work allows to perform multiple local iterations of the E step before communicating with the central server . In algorithm 1 , the local statistics Sk+1 , i 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 clientdrift challenge ( in the presence of heterogeneity ) . Yet , we stress that combining compression and heterogeneity results in other challenges : it is 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 , i ’ s memory-variables . Convergence analysis , full participation regime . In this paragraph , we focus on the fullparticipation regime ( p = 1 ) : for all k 2 [ kmax ] ? , Ak = [ n ] ⇤ . We now present in Theorem 1 our key result , from which complexity expressions are derived . The proof is postponed to Appendix C. Theorem 1 . Assume A1 to A7 and set L2 : = n 1 Pn i=1 L 2 i , 2 : = n 1 Pn i=1 2 i . Let { bSk , k 2 [ kmax ] } be given by algorithm 1 , with ! > 0 , ↵ : = ( 1 + ! ) 1 and k = 2 ( 0 , max ] where max : = vmin 2L Ẇ ^ p n 2 p 2L ( 1 + ! ) p ! . ( 7 ) Denote by K the uniform random variable on [ kmax 1 ] . Then , taking V0 , i = hi ( bS0 ) for all i 2 [ n ] ? : vmin ✓ 1 L Ẇ vmin ◆ E h kh ( bSK ) k2 i 1 kmax ⇣ W ( bS0 ) minW ⌘ + L Ẇ 1 + 5 ! n 2 . ( 8 ) When there is no compression ( ! = 0 so that Quant ( s ) = s ) , we prove that the introduction of the random variables Vk , i ’ s play no role whatever ↵ > 0 and the choice of the V0 , i ’ s , and we have for any 2 ( 0 , 2vmin/LẆ ) ( see ( 29 ) in the supplemental ) ⇣ 1 L Ẇ 2vmin ⌘ E h kh ( bSK ) k2 i 1 kmax ⇣ W ( bS0 ) minW ⌘ + L Ẇ 2 n . ( 9 ) Optimizing the learning rate , we derive the following corollary ( see the proof in Appendix C ) . Corollary 2 ( of Theorem 1 ) . Choose : = ( W ( bS0 ) minW ) n kmaxLẆ ( 1+5 ! ) 2 1/2 ^ max . We get E h kh ( bSK ) k2 i 4 vmin ✓s W ( bS0 ) minW L Ẇ ( 1 + 5 ! ) 2 nkmax _ W ( bS0 ) minW maxkmax ◆ . Theorem 1 and Corollary 2 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 a “ naive ” implementation of compressed distributed EM ( i.e . an implementation without the variables Vk , i ’ s ) would not converge . Let us comment the complexity to reach an ✏-stationary point , and more precisely how the complexity evaluated in terms of the number of optimization steps depend on ! , n , 2 and ✏ . Since KOpt ( ✏ ) = kmax , from Corollary 2 we have that : Kopt ( ✏ ) = O ⇣ ( 1+ ! ) 2 n✏2 ⌘ _O ⇣ 1 max✏ ⌘ . Maximal learning rate and compression . The comparison of Theorem 1 with the no compression case ( see ( 9 ) ) shows that compression impacts max by a factor proportional to p n/ ! 3/2 as ! increases ( similar constraints were observed in the risk optimization literature , e.g . in [ 17 , 32 ] ) . This highlights two different regimes depending on the ratio p n/ ! 3/2 : if the number of workers n scales at least as ! 3 , the maximal learning rate is not impacted by compression ; on the other hand , for smaller numbers of workers n ⌧ ! 3 , compression can degrade the maximal learning rate . We highlight this conclusion with a small example in the case of scalar quantization for which ! ⇠ pq/squant : for q = 102 and squant = 4 ( obtaining a compression rate of a factor 16 ) , the maximal learning rate is almost unchanged if n 16 . Dependency on ✏ . The complexity Kopt ( ✏ ) is decomposed into two terms scaling respectively as 2✏ 2 and 1 max ✏ 1 , the first term being dominant when ✏ ! 0 . This observation highlights two different regimes : a high noise regime corresponding to max ( 1 + ! ) 2/ ( n✏ 1 ) 1 where the complexity is of order 2✏ 2 , and a low noise regime where max ( 1 + ! ) 2/ ( n✏ 1 ) 1 and the complexity is of order 1 max ✏ 1 . An extreme example of the low noise case is 2 = 0 , occurring for example in the finite-sum case ( i.e. , when s̄i = m 1 Pm j=1 s̄ij ) with the oracle Sk+1 , i = s̄i T ( bSk ) . Impact of compression for ✏-stationarity . As mentioned above , the compression simultaneously impacts the maximal learning rate ( as in ( 7 ) ) and the complexity Kopt ( ✏ ) . Consequently , the impact of the compression depends on the balance between ! , n , 2 and ✏ , and we can distinguish four different “ main ” regimes . In the following tabular , for each of the four situations , we summarize the increase in complexity Kopt ( ✏ ) resulting from compression . Complexity regime : ( Dominating term in Kopt ( ✏ ) ) ( 1+ ! ) 2 n✏2 1 max✏ max regime : ( Dominating term in ( 7 ) ) Example situation High noise 2 , small ✏ Low 2 ( e.g. , large minibatch ) larger ✏ vmin 2LẆ large ratio n/ ! 3 ⇥ ! ⇥1 p n 2 p 2L ( 1+ ! ) p ! low ratio n/ ! 3 ⇥ ! ⇥ ! 3/2/ p n Depending on the situation , the complexity can be multiplied by a factor ranging from 1 to ! _ ( ! 3/2/ p n ) . Remark that the communication cost of each iteration is typically reduced by compression of a factor at least ! . Moreover , the benefit of compression is most significant in the low noise regime and when the maximal learning rate is vmin/ ( 2LẆ ) ( e.g. , when n large enough ) . We then improve the communication cost of each iteration without increasing the optimization complexity , 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 to additionally incorporate a variance reduction scheme in FedEM . It is described in the finite-sum setting when for all i 2 [ n ] ? , s̄i : = m 1 Pm j=1 s̄ij : at each iteration # ( t , k + 1 ) , the oracle on s̄i T ( bSt , k ) will use a minibatch Bt , k+1 , i of examples sampled at random ( with or without replacement ) in [ m ] ? . Algorithm 2 : VR-FedEM Data : kout , kin , b 2 N ? ; for i 2 [ n ] ? , V1,0 , i 2 Rq ; bSinit 2 Rq ; a positive sequence { t , k+1 , t 2 [ kout ] ? , k 2 [ kin 1 ] } ; ↵ > 0 Result : sequence : { bSt , k , t 2 [ kout ] ? , k 2 [ kin ] } 1 bS1,0 = bS1 , 1 = bSinit , V1,0 = n 1 Pn i=1 V1,0 , i 2 for i = 1 , . . . , n do 3 S1,0 , i = 1 m Pm j=1 s̄ij T ( bSinit ) 4 for t = 1 , . . . , kout do 5 for k = 0 , . . . , kin 1 do 6 for i = 1 , . . . , n ( worker # i , locally ) do 7 Sample at random a batch Bt , k+1 , i of size b in [ m ] ? 8 Set St , k+1 , i = St , k , i + b 1 P j2Bt , k+1 , i ⇣ s̄ij T ( bSt , k ) s̄ij T ( bSt , k 1 ) ⌘ 9 Set t , k+1 , i = St , k+1 , i bSt , k Vt , k , i 10 Set Vt , k+1 , i = Vt , k , i + ↵Quant ( t , k+1 , i ) . 11 Send Quant ( t , k+1 , i ) to the central server 12 ( the central server ) 13 Set Ht , k+1 = Vt , k + n 1 Pn i=1 Quant ( t , k+1 , i ) 14 Set bSt , k+1 = bSt , k + t , k+1Ht , k+1 15 Set Vt , k+1 = Vt , k +↵n 1 Pn i=1 Quant ( t , k+1 , i ) 16 Send bSt , k+1 and T ( bSt , k+1 ) to the n workers 17 bSt+1,0 = bSt+1 , 1 = bSt , kin 18 Vt+1,0 = Vt , kin 19 for i = 1 , . . . , n do 20 St+1,0 , i = 1 m Pm j=1 s̄ij T ( bSt+1,0 ) 21 Vt+1,0 , i = Vt , kin , i The algorithm is decomposed into kout outer loops ( indexed by t ) , each of them having kin inner loops ( indexed by k ) . At iteration # ( k + 1 ) of the inner loops , each worker # i updates a local statistic St , k+1 , i based on a minibatch Bt , k+1 , i of its own examples { s̄ij , j 2 Bt , k+1 , i } ( see Line 8 ) : starting from bSt,0 , i : = m 1 Pm j=1 s̄ij T ( bSt , 1 ) , bSt , k+1 , i is defined in such a way that it approximates m 1 Pm j=1 s̄ij T ( bSt , k ) ( see Corollary 18 ) . Then , the worker # i sends to the central server a quantization of t , k+1 , i ( see Line 12 ) which can be seen as an approximation of ↵ 1 { hi ( bSt , k ) hi ( bSt , k 1 ) } upon noting that the variable Vt , k+1 , i defined by Line 10 approximates hi ( bSt , k ) ( see Proposition 26 ) . The central server learns the mean value Vt , k+1 = n 1 Pn i=1 Vt , k+1 , i ( see Line 15 and Lemma 21 ) and , by adding the quantized quantities , defines a field Ht , k+1 which approximates n 1 Pn i=1 hi ( bSt , k ) ( see Proposition 24 ) . Line 14 can be seen as a Stochastic Approximation update , with learning rate t , k+1 and mean field s 7 ! n 1 Pn i=1 hi ( s ) ( see ( 6 ) for the definition of hi ) . The variance reduction is encoded in the definition of St , k+1 , i , Line 8 . We have St , k+1 , i = b 1 P j2Bt , k+1 , i s̄ij T ( bSt , k ) +⌥t , k+1 , i . The first term is the natural approximation of s̄i T ( bSt , k ) based on a minibatch Bt , k+1 , i . Conditionally to the past , ⌥t , k+1 , i is correlated to the first term and biased , but its bias is canceled at the beginning of each outer loop ( see Line 20 and Appendix E.3.2 ) : ⌥t , k+1 , i defines a control variate . Such a 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 St+1,0 , i are initialized to the full sum m 1 Pm j=1 s̄ij T ( bSt , kin ) ( see Line 20 ) thus canceling the bias of S· , i ( see Proposition 17 ) . When there is a single worker and no compression is used ( n = 1 , ! = 0 ) , VR-FedEM reduces to SPIDER-EM , which has been shown to be rate optimal for smooth , non-convex finite-sum optimization [ 10 ] . Theorem 3 studies the FL setting ( n 1 and ! 0 ) : it establishes a finite time control of convergence in expectation for VR-FedEM . Assumptions A5 and A7 are replaced with A8 . A8 . For any i 2 [ n ] ? and j 2 [ m ] ? , the conditional expectations s̄ij ( ✓ ) are well defined for any ✓ 2 ⇥ , and there exists Lij such that for any s , s0 2 Rq , k ( s̄ij T ( s ) s ) ( s̄ij T ( s0 ) s0 ) k Lijks s0k . Theorem 3 . Assume A1 to 3 , A4 , A6 and A8 . Set L2 : = n 1m 1 Pn i=1 Pm j=1 L 2 ij . Let { bSt , k , t 2 [ kout ] ? , k 2 [ kin 1 ] } be given by algorithm 2 run with ↵ : = 1/ ( 1 + ! ) , V1,0 , i : = hi ( bS1,0 ) for any i 2 [ n ] ? , b : = d kin ( 1+ ! ) 2 e and t , k = : = vmin L Ẇ ✓ 1 + 4 p 2 vmax L Ẇ Lp n ( 1 + ! ) ⇣ ! + 1 + 10 ! 8 ⌘1/2◆ 1 . ( 10 ) Let ( ⌧ , K ) be the uniform random variable on [ kout ] ? ⇥ [ kin 1 ] , independent of { bSt , k , t 2 [ kout ] ? , k 2 [ kin ] } . Then , it holds E ⇥ kH⌧ , K+1k2 ⇤ 2 E ⇥ W ( bS1,0 ) ⇤ minW vmin kinkout , ( 11 ) E h kh ( bS⌧ , K ) k2 i 2 ⇣ 1 + 2 L2 ( 1 + ! ) 2 n ⌘ E ⇥ kH⌧ , K+1k2 ⇤ . ( 12 ) The proof is postponed to Appendix E. This result is a consequence of the more general Proposition 25 . We make the following comments : 1 . Eq . ( 11 ) provides the convergence of E ⇥ kH⌧ , K+1k2 ⇤ , and Eq . ( 12 ) ensures that the quantity of interest E [ kh ( bS⌧ , K ) k2 ] is controlled by E [ kH⌧ , K+1k2 ] . We observe that 2 ( 1 + 2 L 2 ( 1+ ! ) 2 n ) is uniformly bounded w.r.t. ! as , by ( 10 ) , 2 = O ! ! 1 ( ! 3 ) . 2 . Up to our knowledge , this is 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 ( ! = 0 ) and in the single-worker case ( n = 1 ) , Fort et al . [ 10 ] use kin = b : we recover this result as a particular case . When n > 1 and ! > 0 , the recommended batch size b decreases as 1/ ( 1 + ! ) 2 . Convergence rate and optimization complexity . Our step-size is chosen constant and independent of kin , kout . Indeed , contrary to Theorem 1 , there is no Bias-Variance trade-off ( as typically observed with variance reduced methods ) , and the optimal choice of is the largest one to ensure convergence . Consequently , since the number of optimization steps is koutkin , we have Kopt ( ✏ ) = O ( 1 ✏ ) . Impact of compression on the learning rate and ✏-stationarity . The compression constant ! does not directly appear in ( 11 ) , but impacts the value of . Two different regimes appear : 1. if 4 p 2 vmaxLẆ Lp n ( 1 + ! ) ! + 1+10 ! 8 1/2 ⌧ 1 ( i.e . we focus on the large ! , n asymptotics when ! 3 ⌧ n ) , then ' vminLẆ has nearly the same value as without compression [ 10 ] . The complexity is then similar to the one of SPIDER-EM [ 10 ] , with a smaller communication cost . The gain from compression is maximal in this regime . 2. if 4 p 2 vmaxLẆ Lp n ( 1 + ! ) ! + 1+10 ! 8 1/2 1 ( i.e . we focus on the large ! , n asymptotics when ! 3 n ) , then = O ⇣ vmin p n vmaxL ! 3/2 ⌘ is strictly smaller than without compression . The optimization complexity is then higher to the one of SPIDER-EM1 ( by a factor proportional to ! 3/2/ p n ) with a smaller communication cost ( typically at least ! times less bits exchanged per iteration ) . The overall trade-off thus depends on the comparison between ! and n. Complexity : 1/ ( ✏ ) regime : ( Dominating term in ( 10 ) ) Example situation vmin/LẆ large ratio n/ ! 3 ⇥1 vmin p n/ ( vmaxL ! 3/2 ) low ratio n/ ! 3 ⇥ ! 3/2/ p n We summarize these two regimes in this tabular , focusing on the large n , large ! asymptotic regimes . For the two regimes , we indicate the increase in complexity Kopt ( ✏ ) resulting from compression . We provide a discussion on computed conditional expectations complexity KCE in Appendix E.2 . 1As a corollary of [ 10 , Theorem 2 ] , the optimization complexity of SPIDER-EM is kout + kinkout that is ✏ 1 in order to reach ✏-stationarity .

Summary: This paper investigates how to federate the Expectation Maximization (EM) algorithm in the case of a curved exponential family. The 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. The authors prove the convergence of FedEM (Corollary 2). A variance-reduced version of FedEM is also proposed (VR-FedEM), also with convergence guarantees (Th. 4). Experiments are performed on toy synthetic data, MNIST and real data (missing value imputation from a tabular dataset, bird).