STEM: A Stochastic Two-Sided Momentum Algorithm Achieving Near-Optimal Sample and Communication Complexities for Federated Learning

1 Introduction . In Federated Learning ( FL ) , multiple worker nodes ( WNs ) collaborate with the goal of learning a joint model , by only using local data . Therefore it has become popular for machine learning problems where datasets are massively distributed [ 1 ] . In FL , the data is often collected at or off-loaded to multiple WNs which in collaboration with a server node ( SN ) jointly aim to learn a centralized model ⇤Corresponding Author : Mingyi Hong . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . [ 2 , 3 ] . The local WNs share the computational load and since the data is local to each WN , FL also provides some level of data privacy [ 4 ] . A classical distributed optimization problem that K WNs aim to solve : min x2Rd ⇢ f ( x ) : = 1 K KX k=1 f ( k ) ( x ) : = 1 K KX k=1 E⇠ ( k ) ⇠D ( k ) [ f ( k ) ( x ; ⇠ ( k ) ) ] . ( 1 ) where f ( k ) : Rd ! R denotes the smooth ( possibly non-convex ) objective function and ⇠ ( k ) ⇠ D ( k ) represents the sample/s drawn from distribution D ( k ) at the kth WN with k 2 [ K ] . When the distributions D ( k ) are different across the WNs , it is referred to as the heterogeneous data setting . The optimization performance of non-convex FL algorithms is typically measured by the total number of samples accessed ( cf . Definition 2.2 ) and the total rounds of communication ( cf . Definition 2.3 ) required by each WN to achieve an ✏-stationary solution ( cf . Definition 2.1 ) . To minimize the sample and the communication complexities , FL algorithms rely on the following four key design elements : ( i ) the WNs ’ local model update directions , ( ii ) Minibatch size to compute each local direction , ( iii ) the number of local updates before WNs share their parameters , and ( iv ) the SN ’ s update direction . How to find effective FL algorithms by ( optimally ) designing these parameters has received significant research interest recently . Contributions . The main contributions of this work are listed below : 1 ) We propose the Stochastic Two-Sided Momentum ( STEM ) algorithm , that utilizes certain momentum-assisted stochastic gradient directions for both the WNs and SN updates . We show that there exists an optimal trade off between the minibatch sizes and number of local updates , such that on the trade-off curve STEM requires Õ ( ✏ 3/2 ) 2 samples and Õ ( ✏ 1 ) communication rounds to reach an ✏-stationary solution ; see Figure 1 for an illustration . These complexity results are the best achievable for first-order stochastic FL algorithms ( under certain assumptions , cf . Assumption 1 ) ; see [ 5–8 ] and [ 9 , 10 ] , as well as Remark 1 of this paper for discussions regarding optimality . To the best of our knowledge , STEM is the first algorithm which – ( i ) simultaneously achieves the optimal sample and communication complexities for FL and ( ii ) can optimally trade off the minibatch sizes and the number of local updates . 2 ) A momentum-less special case of our STEM result further reveals some interesting insights of the classical FedAvg algorithm ( a.k.a . the Local SGD ) [ 11–13 ] . Specifically , we show that for FedAvg , there also exists a trade-off between the minibatch sizes and the number of local updates , such that it requires O ( ✏ 2 ) samples and O ( ✏ 3/2 ) communication rounds to achieve an ✏-stationary solution . 2The notation Õ ( · ) hides the logarithmic factors . Collectively , our insights on the trade-offs provide practical guidelines for choosing different design elements for FL algorithms . Related Works . FL algorithms were first proposed in the form of FedAvg [ 11 ] , where the local update directions at each WN were chosen to be the SGD updates . Earlier works analyzed these algorithms in the homogeneous data setting [ 19–25 ] , while many recent studies have focused on designing new algorithms to deal with heterogeneous data settings , as well as problems where the local loss functions are non-convex [ 9 , 10 , 12–16 , 18 , 26–32 ] . In [ 12 ] , the authors showed that Parallel Restarted SGD ( Local SGD or FedAvg [ 11 ] ) achieves linear speed up while requiring O ( ✏ 2 ) samples and O ( ✏ 3/2 ) rounds of communication to reach an ✏-stationary solution . In [ 14 ] , a Momentum SGD was proposed , which achieved the same sample and communication complexities as Parallel Restarted SGD [ 12 ] , without requiring that the second moments of the gradients be bounded . Further , it was shown that under the homogeneous data setting , the communication complexity can be improved to O ( ✏ 1 ) while maintaining the same sample complexity . The works in [ 15 , 16 ] conducted tighter analysis for FedAvg with partial WN participation with O ( 1 ) local updates and batch sizes . Their analysis showed that FedAvg ’ s sample and communication complexities are both O ( ✏ 2 ) . Additionally , SCAFFOLD was proposed in [ 15 ] , which utilized variance reduction based local update directions [ 33 ] to achieve the same sample and communication complexities as FedAvg . Similarly , VRL-SGD proposed in [ 29 ] also utilized variance reduction and showed improved communication complexity of O ( ✏ 1 ) , while requiring the same computations as FedAvg . Importantly , both SCAFFOLD and VRL-SGD ’ s guarantees were independent of the data heterogeneity . The FedProx proposed in [ 10 ] used a penalty based method to improve the communication complexity of FedAvg ( i.e. , the Parallel Restarted and Momentum SGD [ 14 , 12 ] ) to O ( ✏ 1 ) . FedProx used a gradient similarity assumption to model data heterogeneity which can be stringent for many practical applications . This assumption was relaxed by FedPD proposed in [ 9 ] . Recently , the works [ 17 , 18 ] proposed to utilize hybrid momentum gradient estimators [ 7 , 8 ] . The MIME algorithm [ 17 ] matched the optimal sample complexity ( under certain smoothness assumptions ) of O ( ✏ 3/2 ) of the centralized non-convex stochastic optimization algorithms [ 5–8 ] . Similarly , Fed-GLOMO [ 18 ] achieved the same sample complexity while employing compression to further reduce communication . Both MIME and Fed-GLOMO required O ( ✏ 3/2 ) communication rounds to achieve an ✏-stationary solution . Please see Table 1 for a summary of the above discussion . The comparison of Local SGD ( FedAvg ) to Minibatch SGD for convex and strongly convex problems with homogeneous data setting was first conducted in [ 19 ] and later extended to heterogeneous setting in [ 13 ] . It was shown that Minibatch SGD almost always dominates the Local SGD . In contrast , it was shown in [ 24 ] that Local SGD dominates Minibatch SGD in terms of generalization performance . Although existing FL results are rich , but they are somehow ad hoc and there is a lack of principled understanding of the algorithms . We note that the proposed STEM algorithmic framework provides a theoretical framework that unifies all existing FL results on sample and communication complexities . Notations . The expected value of a random variable X is denoted by E [ X ] and its expectation conditioned on an Event A is denoted as E [ X|Event A ] . We denote by R ( and Rd ) the real line ( and the d-dimensional Euclidean space ) . The set of natural numbers is denoted by N. Given a positive integer K 2 N , we denote [ K ] , { 1 , 2 , . . . , K } . Notation k · k denotes the ` 2-norm and h· , ·i the Euclidean inner product . For a discrete set B , |B| denotes the cardinality of the set . Uniform distribution over a discrete set { 1 , . . . , T } is denoted as U { 1 , . . . , T } . 2 Preliminaries . Before we proceed to the algorithms , we make the following assumptions about problem ( 1 ) . Assumption 1 ( Sample Gradient Lipschitz Smoothness ) . The stochastic functions f ( k ) ( · , ⇠ ( k ) ) with ⇠ ( k ) ⇠ D ( k ) for all k 2 [ K ] , satisfy the mean squared smoothness property , i.e , we have Ekrf ( k ) ( x ; ⇠ ( k ) ) rf ( k ) ( y ; ⇠ ( k ) ) k2 L2kx yk2 for all x , y 2 Rd . Assumption 2 ( Unbiased gradient and Variance Bounds ) . ( i ) Unbiased Gradient . The stochastic gradients computed at each WN are unbiased E [ rf ( k ) ( x ; ⇠ ( k ) ) ] = rf ( k ) ( x ) , 8 ⇠ ( k ) ⇠ D ( k ) , 8 k 2 [ K ] . ( ii ) Intra- and inter- node Variance Bound . The following bounds hold : Ekrf ( k ) ( x ; ⇠ ( k ) ) rf ( k ) ( x ) k2 2 , krf ( k ) ( x ) rf ( ` ) ( x ) k2 ⇣2 , 8 ⇠ ( k ) ⇠ D ( k ) , 8k , ` 2 [ K ] . Note that Assumption 1 is stronger than directly assuming f ( k ) ’ s are Lipschitz smooth ( which we will refer to as the averaged gradient Lipschitz smooth condition ) , but it is still a rather standard assumption in SGD analysis . For example it has been used in analyzing centralized SGD algorithms such as SPIDER [ 5 ] , SNVRG [ 6 ] , STORM [ 7 ] ( and many others ) as well as in FL algorithms such as MIME [ 17 ] and Fed-GLOMO [ 18 ] . The second relation in Assumption 2- ( ii ) quantifies the data heterogeneity , and we call ⇣ > 0 as the heterogeneity parameter . This is a typical assumption required to evaluate the performance of FL algorithms . If data distributions across individual WNs are identical , i.e. , D ( k ) = D ( ` ) for all k , ` 2 [ K ] then we have ⇣ = 0 . Next , we define the ✏-stationary solution for non-convex optimization problems , as well as quantify the computation and communication complexities to achieve an ✏-stationary point . Definition 2.1 ( ✏-Stationary Point ) . A point x is called ✏-stationary if krf ( x ) k2 ✏ . Moreover , a stochastic algorithm is said to achieve an ✏-stationary point in t iterations if E [ krf ( xt ) k2 ] ✏ , where the expectation is over the stochasticity of the algorithm until time instant t. Definition 2.2 ( Sample complexity ) . We assume an Incremental First-order Oracle ( IFO ) framework [ 34 ] , where , given a sample ⇠ ( k ) ⇠ D ( k ) at the kth node and iterate x , the oracle returns ( f ( k ) ( x ; ⇠ ( k ) ) , rf ( k ) ( x ; ⇠ ( k ) ) ) . Each access to the oracle is counted as a single IFO operation . We measure the sample ( and computational ) complexity in terms of the total number of calls to the IFO by all WNs to achieve an ✏-stationary point given in Definition 2.1 . Definition 2.3 ( Communication complexity ) . We define a communication round as a one back-andforth sharing of parameters between the WNs and the SN . Then the communication complexity is defined to be the total number of communication rounds between any WN and the SN required to achieve an ✏-stationary point given in Definition 2.1 . 3 The STEM algorithm and the trade-off analysis In this section , we discuss the proposed algorithm and present the main results . The key in the algorithm design is to carefully balance all the four design elements mentioned in Sec . 1 , so that sufficient and useful progress can be made between two rounds of communication . Let us discuss the key steps of STEM , listed in Algorithm 1 . In Step 10 , each node locally updates its model parameters using the local direction dkt , computed by using b stochastic gradients at two Algorithm 1 The Stochastic Two-Sided Momemtum ( STEM ) Algorithm 1 : Input : Parameters : c > 0 , the number of local updates I , batch size b , stepsizes { ⌘t } . 2 : Initialize : Iterate x ( k ) 1 = x̄1 = 1K PK k=1 x ( k ) 1 , descent direction d ( k ) 1 = d̄1 = 1 K PK k=1 d ( k ) 1 with d ( k ) 1 = 1 B P ⇠ ( k ) 1 2B ( k ) 1 rf ( k ) ( x ( k ) 1 ; ⇠ ( k ) 1 ) and |B ( k ) 1 | = B for k 2 [ K ] . 3 : Perform : x ( k ) 2 = xk1 ⌘1d ( k ) 1 , 8 k 2 [ K ] 4 : for t = 1 to T do 5 : for k = 1 to K do # at the WN 6 : d ( k ) t+1 = 1 b X ⇠ ( k ) t+12B ( k ) t+1 rf ( k ) ( x ( k ) t+1 ; ⇠ ( k ) t+1 ) + ( 1 at+1 ) ⇣ d ( k ) t 1 b X ⇠ ( k ) t+12B ( k ) t+1 rf ( k ) ( x ( k ) t ; ⇠ ( k ) t+1 ) ⌘ where we choose |B ( k ) t+1| = b , and at+1 = c · ⌘2t ; 7 : if t mod I = 0 then # at the SN 8 : d ( k ) t+1 = d̄t+1 : = 1 K PK k=1 d ( k ) t+1 9 : x ( k ) t+2 : = x̄t+1 ⌘t+1d̄t+1 = 1 K PK k=1 x ( k ) t+1 ⌘t+1d̄t+1 # server-side momentum 10 : else x ( k ) t+2 = x ( k ) t+1 ⌘t+1d ( k ) t+1 # worker-side momentum 11 : end if 12 : end for 13 : end for 14 : Return : x̄a where a ⇠ U { 1 , ... , T } . consecutive iterates x ( k ) t+1 and x ( k ) t . After every I local steps , the WNs share their current local models { x ( k ) t+1 } K k=1 and directions { d ( k ) t+1 } K k=1 with the SN . The SN aggregates these quantities , and performs a server-side momentum step , before returning x̄t+1 and d̄t+1 to all the WNs . Because both the WNs and the SN perform momentum based updates , we call the algorithm a stochastic two-sided momentum algorithm . The key parameters are : b the minibatch size , I the local update steps between two communication rounds , ⌘t the stepsizes , and at the momentum parameters . One key technical innovation of our algorithm design is to identify the most suitable way to incorporate momentum based directions in FL algorithms . Although the momentum-based gradient estimator itself is not new and has been used in the literature before ( see e.g. , in [ 7 , 8 ] and [ 17 , 18 ] to improve the sample complexities of centralized and decentralized stochastic optimization problems , respectively ) , it is by no means clear if and how it can contribute to improve the communication complexity of FL algorithms . We show that in the FL setting , the local directions together with the local models have to be aggregated by the SN so to avoid being influenced too much by the local data . More importantly , besides the WNs , the SN also needs to perform updates using the ( aggregated ) momentum directions . Finally , such two-sided momentum updates have to be done carefully with the correct choice of minibatch size b , and the number of local updates I . Overall , it is the judicious choice of all these design elements that results in the optimal sample and communication complexities . Next , we present the convergence guarantees of the STEM algorithm . 3.1 Main results : convergence guarantees for STEM In this section , we analyze the performance of STEM . We first present our main result , and then provide discussions about a few parameter choices . In the next subsection , we discuss a special case of STEM related to the classical FedAvg and minibatch SGD algorithms . Theorem 3.1 . Under the Assumptions 1 and 2 , suppose the stepsize sequence is chosen as : ⌘t = ̄ ( wt + 2t ) 1/3 , ( 2 ) where we define : ̄ = ( bK ) 2/3 2/3 L , wt = max ⇢ 2 2 , 4096L3I3̄3 2t , c3̄3 4096L3I3 . Further , let us set c = 64L 2 bK + 2 24̄3LI = L 2 ✓ 64 bK + 1 24 ( bK ) 2I ◆ , and set the initial batch size as B = bI ; set the local updates I and minibatch size b as follows : I = O ( T/K2 ) ⌫/3 , b = O ( T/K2 ) 1/2 ⌫/2 ( 3 ) where ⌫ satisfies ⌫ 2 [ 0 , 1 ] . Then for STEM the following holds : ( i ) For x̄a chosen according to Algorithm 1 , we have : Ekrf ( x̄a ) k2 = O ✓ f ( x̄1 ) f⇤ K2⌫/3T 1 ⌫/3 ◆ + Õ ✓ 2 K2⌫/3T 1 ⌫/3 ◆ + Õ ✓ ⇣2 K2⌫/3T 1 ⌫/3 ◆ . ( 4 ) ( ii ) For any ⌫ 2 [ 0 , 1 ] , we have Sample Complexity : The sample complexity of STEM is Õ ( ✏ 3/2 ) . This implies that each WN requires at most Õ ( K 1✏ 3/2 ) gradient computations , thereby achieving linear speedup with the number of WNs present in the network . Communication Complexity : The communication complexity of STEM is Õ ( ✏ 1 ) . The proof of this result is relegated to the Supplemental Material . A few remarks are in order . Remark 1 ( Near-Optimal sample and communication complexities ) . Theorem 3.1 suggests that when I and b are selected appropriately , then STEM achieves Õ ( ✏ 3/2 ) and Õ ( ✏ 1 ) sample and communication complexities . Taking them separately , these complexity bounds are the best achievable by the existing FL algorithms ( upto logarithmic factors regardless of sample or batch Lipschitz smooth assumption ) [ 35 ] ; see Table 1 . We note that the O ( ✏ 3/2 ) complexity is the best possible that can be achieved by centralized SGD with the sample Lipschitz gradient assumption ; see [ 5 ] . On the other hand , the O ( ✏ 1 ) complexity bound is also likely to be the optimal , since in [ 9 ] the authors showed that even when the local steps use a class of ( deterministic ) first-order algorithms , O ( ✏ 1 ) is the best achievable communication complexity . The only difference is that [ 9 ] does not explicitly assume the inter-node variance bound ( i.e. , the second relation in Assumption 2- ( ii ) ) . We leave the precise characterization of the communication lower bound with inter-node variance as future work . Remark 2 ( Large Batch Sizes and/or Local Updates ) . At first glance , it may seem that the requirement of STEM to compute large mini-batches and/or local updates ( cf . Table 1 ) to achieve this ( near ) optimal performance is a drawback , however , we note that it is in fact an advantage of STEM that it allows the WNs to perform larger number of local updates ( or compute large minibatches ) without communicating often . This follows from the fact that irrespective of the number of local updates ( or batch sizes ) STEM achieves near-optimal communication complexity while attaining optimal overall sample complexity . Moreover , note that even with b = I = O ( 1 ) ( i.e. , b and I are chosen as constants ) , STEM achieves the same ( optimal ) sample and communication complexities as achieved by FedGLOMO [ 18 ] and MIME [ 17 ] . We further note that to the best of our knowledge the algorithms that achieve the communication complexity of O ( ✏ 1 ) either require the number of local updates or the batch-sizes that depend on the solution accuracy ✏ . For example , FedProx [ 10 ] , FedPD [ 9 ] , and FedDyn [ 36 ] rely on solving the “ local problems '' to achieve an ✏-accuracy , which implies that the number of local updates ( or the batch sizes ) implicitly depends on the desired solution accuracy ✏ , as is the case for STEM . Similarly , as shown in [ 12 ] and [ 14 ] the communication complexity of FedAvg and its momentum version can be improved from O ( ✏ 2 ) to O ( ✏ 3/2 ) when the number of local updates ( or batch size ) is chosen as O ( ✏ 1/2 ) ( cf . Section 3.2 for a more detailed discussion ) . Remark 3 ( The Optimal Batch Sizes and Local Updates Trade-off ) . The parameter ⌫ 2 [ 0 , 1 ] is used to balance the local minibatch sizes b , and the number of local updates I . Eqs . in ( 3 ) suggest that when ⌫ increases from 0 to 1 , b decreases and I increases . Specifically , if ⌫ = 1 , then b is a constant but I = O ( T 1/3/K2/3 ) . In this case , each WN chooses a small minibatch while executing multiple local updates , and STEM resembles a FedAvg ( a.k.a . Local SGD ) algorithm but with double-sided momentum update directions , and is referred to as Fed STEM . In contrast , if ⌫ = 0 , then b = O ( T 1/2/K ) but I is a constant . In this case , each WN chooses a large batch size while executing only a few , or even one , local updates , and STEM resembles the Minibatch SGD , but again with different update directions , and is referred to as Minibatch STEM . Such a trade-off can be seen in Fig . 1b . Due to space limitation , these two special cases will be precisely stated in the supplementary materials as corollaries of Theorem 3.1 . Algorithm 2 The FedAvg Algorithm 1 : Input : { ⌘t } Tt=0 ; I , the # of local updates per communication round ; b , the minibatch sizes . 2 : for t = 1 to T do 3 : for k = 1 to K do 4 : d ( k ) t = 1 b P ⇠ ( k ) t 2B ( k ) t rf ( k ) ( x ( k ) t ; ⇠ ( k ) t ) with |B ( k ) t | = b 5 : x ( k ) t+1 = x ( k ) t ⌘td ( k ) t 6 : if t mod I = 0 then 7 : x ( k ) t+1 = x̄t+1 = 1 K PK k=1 x ( k ) t+1 8 : end if 9 : end for 10 : end for 11 : Return : x̄a where a ⇠ U { 1 , ... , T } . Remark 4 ( The Sub-Optimal Batch Sizes and Local Updates Trade-off ) . From our proof ( Theorem C.10 included in the supplemental material ) , we can see that STEM requires Õ max ( b · I ) ✏ 1 , K 1✏ 3/2 samples and Õ max ✏ 1 , ( b · I ) 1K 1✏ 3/2 and communication rounds . According to the above expressions , if b · I increases beyond O ( K 1✏ 1/2 ) , then the sample complexity will increase from the optimal Õ ( ✏ 3/2 ) ; otherwise , the optimal sample complexity Õ ( ✏ 3/2 ) is maintained . On the other hand , if b · I decreases beyond O ( K 1✏ 1/2 ) , the communication complexity increases from Õ ( ✏ 1 ) . For instance , if we choose b = O ( 1 ) and I = O ( 1 ) the communication complexity becomes Õ ( ✏ 3/2 ) while the optimal sample complexity Õ ( ✏ 3/2 ) is maintained . This trade-off is illustrated in Figure 1a , where we maintain the optimal sample complexity , while changing b and I to generate the trade-off surface . Remark 5 ( Data Heterogeneity ) . The term Õ ⇣ ⇣2 K2⌫/3T 1 ⌫/3 ⌘ in the gradient bound ( 4 ) captures the effect of the heterogeneity of data across WNs , where ⇣ is the parameter characterizing the intra-node variance and has been defined in Assumption 2- ( ii ) . Highly heterogeneous data with large ⇣2 can adversely impact the performance of STEM . Note that such a dependency on ⇣ also appears in other existing FL algorithms , such as [ 9 , 14 , 18 ] . However , there is one special case of STEM that does not depend on the parameter ⇣ . This is the case where I = 1 , i.e. , the minibatch SGD counterpart of STEM where only a single local iteration is performed between two communication rounds . We have the following corollary . Corollary 1 ( Minibatch STEM ) . Under Assumptions 1 and 2 , and choose the algorithm parameters as in Theorem 3.1 . At each WN , choose I = 1 , b = ( T/K2 ) 1/2 , and the initial batch size B = b · I . Then STEM satisfies : ( i ) For x̄a chosen according to Algorithm 1 , we have Ekrf ( x̄a ) k2 = O ⇣f ( x̄1 ) f⇤ T ⌘ + Õ ⇣ 2 T ⌘ . ( ii ) Minibatch STEM achieves Õ ( ✏ 3/2 ) sample and Õ ( ✏ 1 ) communication complexity . Next , we show that FedAvg also exhibits a trade-off similar to that of STEM but with worse sample and communication complexities .

Summary: The paper proposes a new method to solve the finite sum problem in federated learning where each local model is an expectation function. The STEM algorithm uses momentum variance-reduced estimator in local updates (similar to STORM estimator in stochastic optimization) and also send these estimators to the server when performing server-side update. Therefore, STEM includes momentum updates for both local and server updates. Convergence analysis reveals that STEM nearly match the best known sample complexity in centralized setting and it also matches the lower bound communication complexity up to a log factor. Numerical experiments on neural network training not only show advantage of STEM compared with FedAvg and SCAFFOLD but also presenting the tradeoff between communication frequency and minibatch size.