Mime: Mimicking Centralized Stochastic Algorithms in Federated Learning

1 INTRODUCTION . Federated learning has become an important paradigm in large-scale machine learning where the training data remains distributed over a large number of clients , which may be mobile phones or network sensors ( Konečnỳ et al. , 2016b ; a ; McMahan et al. , 2017 ; Mohri et al. , 2019 ; Kairouz et al. , 2019 ) . A centralized model , here referred to as a server model , is then trained , without ever transmitting client data over the network , thereby providing some basic levels of data privacy and security . Two important settings are distinguished in Federated learning ( Kairouz et al. , 2019 , 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 ( Holst , 2019 ) . Thus , in that setting , we may never make even a single pass over the entire clients ’ data during training . The cross-device setting is further characterized by resource-poor 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 . Here , we are interested in the cross-device setting , for which we will formalize and study stochastic optimization algorithms . The de facto standard algorithm for this setting is FEDAVG ( McMahan et al. , 2017 ) , which performs multiple SGD updates on the available clients , before communicating to the server . While this approach can reduce the total amount 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 ( Karimireddy et al. , 2020 ) . Furthermore , algorithmic innovations such as momentum ( Sutskever et al. , 2013 ; Cutkosky and Orabona , 2019 ) , adaptivity ( Kingma and Ba , 2014 ; Zaheer et al. , 2018 ; Zhang et al. , 2019 ) , and clipping ( You et al. , 2017 ; 2019 ; Zhang et al. , 2020 ) are critical to the success of deep learning applications and need to be incorporated into the client updates , replacing the SGD update of FEDAVG . Perhaps due to such deficiencies , there exists a large gap in performance between the centralized setting , where data is centrally collected on the server , and the federated setting ( Zhao et al. , 2018 ; Hsieh et al. , 2019 ; Hsu et al. , 2019 ; Karimireddy et al. , 2020 ) . To overcome such deficiencies , we propose a new framework , MIME , that mitigates client drift and adapts arbitrary centralized optimization algorithms , e.g . SGD with momentum or Adam , to the federated setting . In each local client update , MIME uses global statistics , e.g . momentum , and an SVRG-style correction to mimic the updates of the centralized algorithm run on i.i.d . data . These global statistics are computed only at the server level and kept fixed throughout the local steps , thereby avoiding a bias due to the atypical local data of any single client . Contributions . We summarize our main results below . • We formalize the cross-device federated learning problem , and propose a new framework MIME that can adapt arbitrary centralized algorithms to this setting . • We prove that incorporating server momentum into each local client update reduces client drift and leads to optimal statistical rates . • Further , we quantify the usefulness of performing multiple local updates on a single client by carefully tracking the bias ( client-drift ) introduced . This is the first analysis showing improved rates by taking additional multiple steps for general smooth functions . • Finally , we also 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 are 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 Zinkevich et al . ( 2010 ) derived an analysis with asymptotic convergence . Sharper and more refined analyses of the same method , sometimes called local SGD , were provided by Stich ( 2019 ) , and more recently by Stich and Karimireddy ( 2019 ) , Patel and Dieuleveut ( 2019 ) , Khaled et al . ( 2020 ) , and Woodworth et al . ( 2020b ) , for identical functions . Their analysis was extended to heterogeneous clients in ( Wang et al. , 2019 ; Yu et al. , 2019b ; Karimireddy et al. , 2020 ; Khaled et al. , 2020 ; Koloskova et al. , 2020 ) . Charles and Konečnỳ ( 2020 ) 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 Karimireddy et al . ( 2020 ) and Woodworth et al . ( 2020a ) 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 , Karimireddy et al . ( 2020 ) 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 . In contrast , we focus on the cross-device setting where clients may be visited only once during training and where they are stateless . This is akin to the difference between the finite-sum and stochastic settings in traditional centralized optimization . Improvements to FedAvg : Hsu et al . ( 2019 ) and Wang et al . ( 2020c ) observed that using server momentum significantly improves over vanilla FEDAVG . This idea was generalized by Reddi et al . ( 2020 ) , 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 . 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 , Li et al . ( 2018 ) proposed to add a regularizer to ensure client updates remain close . However , its usefulness is unclear ( cf . Fig . 5 , Karimireddy et al. , 2020 ; Wang et al. , 2020b ) . Other orthogonal directions which can be combined with MIME include tackling computation heterogeneity , where some clients perform many more updates than others ( Wang et al. , 2020b ) , improving fairness by modifying the objective ( Mohri et al. , 2019 ; Li et al. , 2019 ) , incorporating differential privacy ( Geyer et al. , 2017 ; Agarwal et al. , 2018 ; Thakkar et al. , 2020 ) , Byzantine adversaries ( Pillutla et al. , 2019 ; Wang et al. , 2020a ; He et al. , 2020a ) , secure aggregation ( Bonawitz et al. , 2017 ; He et al. , 2020b ) , etc . We refer the reader to the extensive survey by Kairouz et al . ( 2019 ) for additional discussion . 2 PROBLEM SETUP . This section formalizes the problem of cross-device federated learning . We first examine some key challenges of this setting ( cf . Kairouz et al. , 2019 ) to ensure our formalism captures the difficulty : 1 . Communication cost between the server and the clients is a major concern and the source of bottleneck in federated learning ; thus , a key metric for optimization in this setting is the number of communication rounds . 2 . Each client is likely to participate at most once , due to the extremely large number of clients ; furthermore , each individual client may have very little data of its own . 3 . There may be a wide heterogeneity or non-i.i.d.-ness due to the difference of data distributions for the clients . Thus , our objective will be to minimize the following quantity within the fewest number of clientserver communication rounds : f ( x ) = Ei∼D [ 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 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 . There is some inherent tension between the second and the third challenge outlined above : if there exists a client with arbitrarily different data whom we may never encounter during training , then there is no hope to actually minimize f . Thus for ( 1 ) to be tractable , it is necessary to assume bounded dissimilarity between different fi . ( A1 ) G2-BGD or bounded gradient dissimilarity : there exists G ≥ 0 such that Ei∼D [ ‖∇fi ( x ) −∇f ( x ) ‖2 ] ≤ G2 , ∀x . Next , we also characterize the variance in the Hessians . Note that if fi ( · ; ζ ) is L-smooth , ( A2 ) is always satisfied with δ ≤ 2L and hence is more of a definition rather than an assumption . Note that however , in realistic examples we expect the clients to be similar and hence that δ L. ( A2 ) δ-BHD or bounded Hessian dissimilarity : Almost surely , f is δ-weakly convex i.e . ∇2fi ( x ) −δI and the loss function of any client i satisfies ‖∇2fi ( x ; ζ ) −∇2f ( x ) ‖ ≤ δ , ∀x . In addition , we assume that f ( x ) is bounded from below by f ? and is L-smooth , as is standard . 3 USING MOMENTUM TO REDUCE CLIENT DRIFT . In this section we examine the tension between reducing communication by running multiple client updates each round , and degradation in performance due to client drift ( Karimireddy et al. , 2020 ) . To simplify the discussion , we assume a single client is sampled each round and that clients use full-batch gradients . Server-only approach . A simple way to avoid the issue of client drift is to take no local steps . We sample a client i ∼ D and run SGDm with momentum parameter β and step size η : xt = xt−1 − η ( ( 1− β ) ∇fi ( xt−1 ) + βmt−1 ) , mt = ( 1− β ) ∇fi ( xt−1 ) + βmt−1 . ( 2 ) Here , the gradient ∇fi ( xt ) is unbiased i.e . E [ ∇fi ( xt ) ] = ∇f ( xt ) and hence we are guaranteed convergence . However , this strategy can be communication-intensive and we are likely to spend all our time waiting for communication with very little time spent on computing the gradients . FedAvg approach . To reduce the overall communication rounds required , we need to make more progress in each round of communication . Starting from y0 = xt−1 , FEDAVG ( McMahan et al. , 2017 ) runs multiple SGD steps on the sampled client i ∼ D yk = yk−1 − η∇fi ( yk−1 ) for k ∈ [ K ] , ( 3 ) and then a pseudo-gradient g̃t = − ( yK −xt ) replaces∇fi ( xt−1 ) in the SGDm algorithm ( 2 ) . This is referred to as server-momentum since it is computed and applied only at the server level ( Hsu et al. , 2019 ) . However , such updates give rise to client-drift resulting in performance worse than the naive server-only strategy ( 2 ) . This is because by using multiple local updates , ( 3 ) starts overfitting to the local client data , optimizing fi ( x ) instead of the actual global objective f ( x ) . The net 2 which can be quite different from the true global optimum x ? . Server momentum only speeds up the convergence to the wrong point in this case . In contrast , MIME uses unbiased momentum and applies it locally at every update . This keeps the updates of MIME closer to the true optimum x ? . effect is that FEDAVG moves towards an incorrect point ( see Fig 1 , left ) . If K is sufficiently large , approximately yK x ? i , where x ? i : = arg min x fi ( x ) ⇒ Ei∼D [ g̃t ] ( xt − Ei∼D [ x ? i ] ) . Further , the server momentum is based on g̃t and hence is also biased . Thus , it can not correct for the client drift . We next see how a different way of using momentum can mitigate client drift . Mime approach . FEDAVG experiences client drift because both the momentum and the client updates are biased . To fix the former , we compute momentum using only global statistics as in ( 2 ) : mt = ( 1− β ) ∇fi ( xt−1 ) + βmt−1 . ( 4 ) To reduce the bias in the local updates , we will apply this unbiased momentum every step : yk = yk−1 − η ( ( 1− β ) ∇fi ( yk−1 ) + βmt−1 ) for k ∈ [ K ] . ( 5 ) Note that the momentum term is kept fixed during the local updates i.e . there is no local momentum used , only global momentum is applied locally . Since mt−1 is a moving average of unbiased gradients computed over multiple clients , it intuitively is a good approximation of the general direction of the updates . By taking a convex combination of the local gradient with mt−1 , the update ( 5 ) is potentially also less biased . In this way MIME combines the communication benefits of taking multiple local steps and prevents client-drift ( see Fig 1 , right ) . Sec . C makes this intuition precise .

Summary: This paper proposes a way to apply various variance reduction/momentum based method to the federated learning scenario, especially when there is distribution drift among the clients. The main claim of this paper is that the global statistics (momentum, control variance et al) should be update at the server side only, which is helpful to reduce the bias of these terms. The paper also provides convergence analysis for their methods and attains the best convergence result for their MimeMVR method.