Breaking the centralized barrier for cross-device federated learning

1 Introduction . Federated learning ( FL ) is an increasingly important large-scale learning framework where the training data remains distributed over a large number of clients , which may be mobile phones or network sensors [ 38 , 37 , 43 , 44 , 28 ] . A server then orchestrates the clients to train a single model , here referred to as a server model , without ever transmitting client data over the network , thereby providing some basic levels of data privacy and security . Two important settings are distinguished in FL [ 28 , Table 1 ] : the cross-device and the cross-silo settings . The cross-silo setting corresponds to a relatively small number of reliable clients , typically organizations , such as medical or financial institutions . In contrast , in the cross-device federated learning setting , the number of clients may be extremely large and include , for example , all 3.5 billion active android phones [ 25 ] . Thus , in that setting , we may never make even a single pass over ∗This work was also appears under the alternative title “ Mime : Mimicking Centralized Stochastic Algorithms in Federated Learning ” [ 31 ] . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . the entire clients ’ data during training . The cross-device setting is further characterized by resourcepoor clients communicating over a highly unreliable network . Together , the essential features of this setting give rise to unique challenges not present in the cross-silo setting . In this work , we are interested in the more challenging cross-device setting , for which we will formalize and study stochastic optimization algorithms . Importantly , recent advances in FL optimization , such as SCAFFOLD [ 32 ] or FedDyn [ 1 ] , are not anymore applicable since they are designed for the cross-silo setting . The problem . The de facto standard algorithm for the cross-device setting is FEDAVG [ 43 ] , which performs multiple SGD updates on the available clients before communicating to the server . While this approach can reduce the frequency of communication required , performing multiple steps on the same client can lead to ‘ over-fitting ’ to its atypical local data , a phenomenon known as client drift [ 32 ] . This in turn leads to slower convergence and can , somewhat counter-intuitively , require larger total communication [ 69 ] . Despite significant attention received from the optimization community , the communication complexity of heterogeneous cross-device has not improved upon that of simple centralized methods , which take no local steps ( aka SERVER-ONLY methods ) . Furthermore , algorithmic innovations such as momentum [ 59 , 14 ] , adaptivity [ 35 , 75 , 77 ] , and clipping [ 71 , 72 , 76 ] are critical to the success of deep learning applications . The lack of a theoretical understanding of the impact of multiple client steps has also hindered adapting these techniques in a principled manner into the client updates , in order to replace the vanilla SGD update of FEDAVG . To overcome such deficiencies , we propose a new framework , MIME , that mitigates client drift and can adapt an arbitrary centralized optimization algorithm , e.g . SGD with momentum or Adam , to the federated setting . In each local client update , MIME uses global optimizer state , e.g . momentum or adaptive learning rates , and an SVRG-style correction to mimic the updates of the centralized algorithm run on i.i.d . data . This optimizer state is computed only at the server level and kept fixed throughout the local steps , thereby avoiding overfitting to the atypical local data of any single client . Contributions . We summarize our main results below . • MIME framework . We formalize the cross-device federated learning problem , and propose a new framework MIME that can adapt arbitrary centralized algorithms to this setting . • Convergence result . We prove a result showing that MIME successfully reduces client drift . We also prove that the convergence of any generic algorithm in the centralized setting translates convergence of its MIME version in the federated setting . • Speed-up over centralized methods . By carefully tracking the bias introduced due to multiple local steps , we prove that MIME with momentum-based variance reduction ( MVR ) can beat a lower bound for centralized methods , thus breaking a fundamental barrier . This is the first such result in FL , and also the first general result showing asymptotic speed-up due to local steps . • Empirical validation . We propose a simpler variant , MIMELITE , with an empirical performance similar to MIME . We report the results of thorough experimental analysis demonstrating that both MIME and MIMELITE indeed converge faster than FEDAVG . Related work . Analysis of FEDAVG : Much of the recent work in federated learning has focused on analyzing FEDAVG . For identical clients , FEDAVG coincides with parallel SGD , for which [ 78 ] derived an analysis with asymptotic convergence . Sharper and more refined analyses of the same method , sometimes called local SGD , were provided by [ 56 ] , and more recently by [ 57 ] , [ 47 ] , [ 34 ] , and [ 70 ] , for identical functions . Their analysis was extended to heterogeneous clients in [ 68 , 74 , 32 , 34 , 36 ] . [ 11 ] derived a tight characterization of FedAvg with quadratic functions and demonstrated the sensitivity of the algorithm to both client and server step sizes . Matching upper and lower bounds were recently given by [ 32 ] and [ 69 ] for general functions , proving that FEDAVG can be slower than even SGD for heterogeneous data , due to the client-drift . Comparison to SCAFFOLD : For the cross-silo setting where the number of clients is relatively low , [ 32 ] proposed the SCAFFOLD algorithm , which uses control-variates ( similar to SVRG ) to correct for client drift . However , their algorithm crucially relies on stateful clients which repeatedly participate in the training process . FedDyn [ 1 ] reduces the communication requirements , but also requires persistent stateful clients . In contrast , we focus on the cross-device setting where clients may be visited only once during training and where they are stateless ( and thus SCAFFOLD and FedDyn are inapplicable ) . This is akin to the difference between the finite-sum ( corresponding to cross-silo ) and stochastic ( cross-device ) settings in traditional centralized optimization [ 39 ] . Comparison to FedAvg and variants : [ 26 ] and [ 67 ] observed that using server momentum significantly improves over vanilla FEDAVG . This idea was generalized by [ 49 ] , who replaced the server update with an arbitrary optimizer , e.g . Adam . However , these methods only modify the server update while using SGD for the client updates . We henceforth refer to this meta algorithm as FedAvg . FedAvgSGD , FedAvgMom , FedAvgAdam denote specific instantiations of the server optimizer in FedAvg with SGD , Momentum or Adam . MIME , on the other hand , ensures that every local client update resembles the optimizer e.g . MIME would apply momentum in every client update and not just at the server level . Beyond this , [ 40 ] proposed to add a regularizer to ensure client updates remain close . However , this may slow down convergence ( cf . Fig . 5 and [ 32 , 66 ] ) . Other orthogonal directions which can be combined with MIME include tackling computation heterogeneity , where some clients perform many more updates than others [ 66 ] , improving fairness by modifying the objective [ 44 , 41 ] , incorporating differential privacy [ 20 , 2 , 61 ] , Byzantine adversaries [ 48 , 65 , 30 ] , secure aggregation [ 8 , 24 ] , etc . We defer additional discussion to the extensive survey by [ 28 ] . Momentum based variance reduction . Initial optimal methods for stochastic non-convex optimization like SPIDER [ 17 ] and SARAH [ 46 ] required intermittently computing very large batch gradients . Subsequently , it was shown that momentum based variance reduction ( MVR ) methods obtained a similar optimal rate without needing such large batch gradient computations [ 62 , 14 ] . Momentum is an exponential moving average of many stochastic gradients and so it has much smaller variance than the stochastic gradients themselves . However , because these gradients are computed at different parameters it also has a bias . MVR adds a small additional correction term which significantly reduces this bias and provides improved rates . 2 Problem setup . This section formalizes the problem of cross-device federated learning [ 28 ] . Cross-device FL is characterized by a large number of client devices like mobile phones which may potentially connect to the server at most once . Due to their transient nature , it is not possible to store any state on the clients , precluding an algorithm like SCAFFOLD . Furthermore , each client has only a few samples , and there is wide heterogeneity in the samples across clients . Finally , communication is a major bottleneck and a key metric for optimization in this setting is the number of communication rounds . Thus , our objective will be to minimize the following quantity within the fewest number of clientserver communication rounds : f ( x ) = Ei∼C [ fi ( x ) : = 1 ni ni∑ ν=1 fi ( x ; ζi , ν ) ] . ( 1 ) Here , fi denotes the loss function of client i and { ζi,1 , . . . , ζi , ni } its local data . Since the number of clients is extremely large , while the size of each local data is rather modest , we represent the former as an expectation and the latter as a finite sum . In each round , the algorithm samples a subset of clients ( of size S ) and performs some updates to the server model . Due to the transient and heterogeneous nature of the clients , it is easy to see that the problem becomes intractable with arbitrarily dissimilar clients . Thus , it is necessary to assume bounded dissimilarity across clients . ( A1 ) G2-BGV or bounded inter-client gradient variance : there exists G ≥ 0 such that Ei∼C [ ‖∇fi ( x ) −∇f ( x ) ‖2 ] ≤ G2 , ∀x . Next , we also characterize the variance in the Hessians . ( A2 ) δ-BHV or bounded Hessian variance : Almost surely , the loss function of any client i satisfies ‖∇2fi ( x ; ζ ) −∇2f ( x ) ‖ ≤ δ , ∀x . This is in contrast to the usual smoothness assumption that can be stated as : ( A2 * ) L-smooth : ‖∇2fi ( x ; ζ ) ‖ ≤ L , ∀x , a.s. for any i . Note that if fi ( x ; ζ ) is L-smooth then ( A2 ) is satisfied with δ ≤ 2L , and hence ( A2 ) is weaker than ( A2 * ) . In realistic examples we expect the clients to be similar and hence that δ L. In addition , we assume that f ( x ) is bounded from below by f ? and is L-smooth , as is standard . 3 Mime framework . In this section we describe how to adapt an arbitrary centralized optimizer ( referred to as the “ base ” optimizer ) which may have internal state ( e.g . momentum ) to the federated learning problem ( 1 ) while ensuring there is no client-drift . Algorithm 4 describes our framework . We develop two variants , MIME and MIMELITE , which consist of three components i ) a base optimizer we are seeking to mimic , ii ) the global ( server ) optimizer state computation , and iii ) the local client updates . Algorithm 1 Mime and MimeLite input : initial x and s , learning rate η and base optimizer B = ( U , V ) for each round t = 1 , · · · , T do sample subset S of clients communicate ( x , s ) to all clients i ∈ S communicate c← 1|S| ∑ j∈S ∇fj ( x ) ( only Mime ) on client i ∈ S in parallel do initialize local model yi ← x for k = 1 , · · · , K do sample mini-batch ζ from local data gi ← ∇fi ( yi ; ζ ) −∇fi ( x ; ζ ) + c ( Mime ) gi ← ∇fi ( yi ; ζ ) ( MimeLite ) update yi ← yi − ηU ( gi , s ) end for compute full local-batch gradient∇fi ( x ) communicate ( yi , ∇fi ( x ) ) end on client s ← V ( 1 |S| ∑ i∈S ∇fi ( x ) , s ) ( update optimizer state ) x← 1|S| ∑ i∈S yi ( update server parameters ) end for Base optimizer . We assume the centralized base optimizer we are imitating can be decomposed into two steps : an update step U which updates the parameters x , and a optimizer state update step V ( · ) which keeps track of global optimizer state s. Each step of the base optimizer B = ( U , V ) uses a gradient g to update the parameter x and the optimizer state s as follows : x← x− η U ( g , s ) , s← V ( g , s ) . ( BASEOPT ) As an example , consider SGD with momentum . The state here is the momentum mt and uses the following update steps : xt = xt−1 − η ( ( 1− β ) ∇fi ( xt−1 ) + βmt−1 ) , mt = ( 1− β ) ∇fi ( xt−1 ) + βmt−1 . Thus , SGD with momentum can be represented in the above generic form with U ( g , s ) = ( 1 − β ) g + βs and V ( g , s ) = ( 1 − β ) g + βs . Table 5 in Appendix shows how other algo- rithms like Adam , Adagrad , etc . can be represented in this manner . We keep the update U to be linear in the gradient g , whereas V can be more complicated . This implies that while the parameter update step U is relatively resilient to receiving a biased gradient g while V can be much more sensitive . Compute optimizer state globally , apply locally . When updating the optimizer state of the base algorithm , we use only the gradient computed at the server parameters . Further , they remain fixed throughout the local updates of the clients . This ensures that these optimizer state remain unbiased and representative of the global function f ( · ) . At the end of the round , the server performs s← V ( 1 |S| ∑ i∈S ∇fi ( x ) , s ) , ∇fi ( x ) = 1ni ∑ni ν=1∇fi ( x ; ζi , ν ) . ( OPTSTATE ) Note that we use full-batch gradients computed at the server parameters x , not client parameters yi . Local client updates . Each client i ∈ S performs K updates using U of the base algorithm and a minibatch gradient . There are two variants possible corresponding to MIME and MIMELITE differentiated using colored boxes . Starting from yi ← x , repeat the following K times yi ← yi − ηU ( gi , s ) ( CLIENTSTEP ) where gi ← ∇fi ( yi ; ζ ) for MIMELITE , and gi ← ∇fi ( yi ; ζ ) −∇fi ( x ; ζ ) + 1|S| ∑ j∈S ∇fj ( x ) for MIME . MIMELITE simply uses the local minibatch gradient whereas MIME uses an SVRG style correction [ 27 ] . This is done to reduce the noise from sampling a local mini-batch . While this correction yields faster rates in theory ( and in practice for convex problems ) , in deep learning applications we found that MIMELITE closely matches the performance of MIME . Finally , there are two modifications made in practical FL : we weight all averages across the clients by the number of datapoints ni [ 43 ] , and we perform K epochs instead of K steps [ 66 ] .

Summary: The paper introduces a new generalized algorithmic optimization framework named MIME for adapting and analyzing the convergence of any optimization algorithm applied in centralized settings into cross-device federated learning environments. The algorithmic reduction is accomplished by decoupling the centralized optimization algorithm into optimizer (server-side) and parameter state (client-side) updates in federated settings. The convergence of the proposed adaptation is shown both theoretically and empirically over a range of challenging real-world datasets for different optimization algorithms. Interestingly, when applying the proposed federated adaptation for a specific momentum-based optimization algorithm, the adaptation is proven to be asymptotically faster than its centralized counterpart.