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 authors presented a momentum extension of the classic FedAvg method for non-convex federated learning, along with a unified framework for convergence analysis. The core result shows that the proposed method can achieve near-tight sample and communication complexities under mild conditions, whilst also revealing a trade-off between local update frequencies and local minibatch sizes. The actual performance of the proposed method was evaluated in CIFAR-10 and MNIST image classification tasks.