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 proposes MIME, a general algorithmic framework for federated learning (FL) that can adapt to (and match the performance of) arbitrary centralized algorithms like Adam, Momentum SGD, and SGD. The authors show that via the MIME framework the convergence of any centralized algorithm can be translated into the convergence of the algorithm in the FL setting. The authors also propose a momentum-based variance reduction (MVR) algorithm that the authors claim to be faster than any centralized algorithm. Finally, the authors conduct extensive numerical experiments to show that the algorithms under the MIME framework perform well on real datasets.