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: 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.