On the Convergence of FedAvg on Non-IID Data

1 INTRODUCTION . Federated Learning ( FL ) , also known as federated optimization , allows multiple parties to collaboratively train a model without data sharing ( Konevcnỳ et al. , 2015 ; Shokri and Shmatikov , 2015 ; McMahan et al. , 2017 ; Konevcnỳ , 2017 ; Sahu et al. , 2018 ; Zhuo et al. , 2019 ) . Similar to the centralized parallel optimization ( Jakovetic , 2013 ; Li et al. , 2014a ; b ; Shamir et al. , 2014 ; Zhang and Lin , 2015 ; Meng et al. , 2016 ; Reddi et al. , 2016 ; Richtárik and Takác , 2016 ; Smith et al. , 2016 ; Zheng et al. , 2016 ; Shusen Wang et al. , 2018 ) , FL let the user devices ( aka worker nodes ) perform most of the computation and a central parameter server update the model parameters using the descending directions returned by the user devices . Nevertheless , FL has three unique characters that distinguish it from the standard parallel optimization Li et al . ( 2019 ) . ∗Equal contribution . First , the training data are massively distributed over an incredibly large number of devices , and the connection between the central server and a device is slow . A direct consequence is the slow communication , which motivated communication-efficient FL algorithms ( McMahan et al. , 2017 ; Smith et al. , 2017 ; Sahu et al. , 2018 ; Sattler et al. , 2019 ) . Federated averaging ( FedAvg ) is the first and perhaps the most widely used FL algorithm . It runs E steps of SGD in parallel on a small sampled subset of devices and then averages the resulting model updates via a central server once in a while.1 In comparison with SGD and its variants , FedAvg performs more local computation and less communication . Second , unlike the traditional distributed learning systems , the FL system does not have control over users ’ devices . For example , when a mobile phone is turned off or WiFi access is unavailable , the central server will lose connection to this device . When this happens during training , such a non-responding/inactive device , which is called a straggler , appears tremendously slower than the other devices . Unfortunately , since it has no control over the devices , the system can do nothing but waiting or ignoring the stragglers . Waiting for all the devices ’ response is obviously infeasible ; it is thus impractical to require all the devices be active . Third , the training data are non-iid2 , that is , a device ’ s local data can not be regarded as samples drawn from the overall distribution . The data available locally fail to represent the overall distribution . This does not only bring challenges to algorithm design but also make theoretical analysis much harder . While FedAvg actually works when the data are non-iid McMahan et al . ( 2017 ) , FedAvg on non-iid data lacks theoretical guarantee even in convex optimization setting . There have been much efforts developing convergence guarantees for FL algorithm based on the assumptions that ( 1 ) the data are iid and ( 2 ) all the devices are active . Khaled et al . ( 2019 ) ; Yu et al . ( 2019 ) ; Wang et al . ( 2019 ) made the latter assumption , while Zhou and Cong ( 2017 ) ; Stich ( 2018 ) ; Wang and Joshi ( 2018 ) ; Woodworth et al . ( 2018 ) made both assumptions . The two assumptions violates the second and third characters of FL . Previous algorithm Fedprox Sahu et al . ( 2018 ) doesn ’ t require the two mentioned assumptions and incorporates FedAvg as a special case when the added proximal term vanishes . However , their theory fails to cover FedAvg . Notation . Let N be the total number of user devices and K ( ≤ N ) be the maximal number of devices that participate in every round ’ s communication . Let T be the total number of every device ’ s SGDs , E be the number of local iterations performed in a device between two communications , and thus TE is the number of communications . Contributions . For strongly convex and smooth problems , we establish a convergence guarantee for FedAvg without making the two impractical assumptions : ( 1 ) the data are iid , and ( 2 ) all the devices are active . To the best of our knowledge , this work is the first to show the convergence rate of FedAvg without making the two assumptions . We show in Theorem 1 , 2 , and 3 that FedAvg has O ( 1T ) convergence rate . In particular , Theorem 3 shows that to attain a fixed precision , the number of communications is T E = O [ 1 ( ( 1 + 1 K ) EG2 + ∑N k=1 p 2 kσ 2 k + Γ +G 2 E ) ] . ( 1 ) Here , G , Γ , pk , and σk are problem-related constants defined in Section 3.1 . The most interesting insight is that E is a knob controlling the convergence rate : neither setting E over-small ( E = 1 makes FedAvg equivalent to SGD ) nor setting E over-large is good for the convergence . This work also makes algorithmic contributions . We summarize the existing sampling3 and averaging schemes for FedAvg ( which do not have convergence bounds before this work ) and propose a new scheme ( see Table 1 ) . We point out that a suitable sampling and averaging scheme is crucial for the convergence of FedAvg . To the best of our knowledge , we are the first to theoretically demonstrate 1In original paper ( McMahan et al. , 2017 ) , E epochs of SGD are performed in parallel . For theoretical analyses , we denote by E the times of updates rather than epochs . 2Throughout this paper , “ non-iid ” means data are not identically distributed . More precisely , the data distributions in the k-th and l-th devices , denote Dk and Dl , can be different . 3Throughout this paper , “ sampling ” refers to how the server chooses K user devices and use their outputs for updating the model parameters . “ Sampling ” does not mean how a device randomly selects training samples . that FedAvg with certain schemes ( see Table 1 ) can achieve O ( 1T ) convergence rate in non-iid federated setting . We show that heterogeneity of training data and partial device participation slow down the convergence . We empirically verify our results through numerical experiments . Our theoretical analysis requires the decay of learning rate ( which is known to hinder the convergence rate . ) Unfortunately , we show in Theorem 4 that the decay of learning rate is necessary for FedAvg with E > 1 , even if full gradient descent is used.4 If the learning rate is fixed to η throughout , FedAvg would converge to a solution at least Ω ( η ( E − 1 ) ) away from the optimal . To establish Theorem 4 , we construct a specific ` 2-norm regularized linear regression model which satisfies our strong convexity and smoothness assumptions . Paper organization . In Section 2 , we elaborate on FedAvg . In Section 3 , we present our main convergence bounds for FedAvg . In Section 4 , we construct a special example to show the necessity of learning rate decay . In Section 5 , we discuss and compare with prior work . In Section 6 , we conduct empirical study to verify our theories . All the proofs are left to the appendix . 2 FEDERATED AVERAGING ( FEDAVG ) Problem formulation . In this work , we consider the following distributed optimization model : min w { F ( w ) , N∑ k=1 pkFk ( w ) } , ( 2 ) where N is the number of devices , and pk is the weight of the k-th device such that pk ≥ 0 and∑N k=1 pk = 1 . Suppose the k-th device holds the nk training data : xk,1 , xk,2 , · · · , xk , nk . The local objective Fk ( · ) is defined by Fk ( w ) , 1 nk nk∑ j=1 ` ( w ; xk , j ) , ( 3 ) where ` ( · ; · ) is a user-specified loss function . Algorithm description . Here , we describe one around ( say the t-th ) of the standard FedAvg algorithm . First , the central server broadcasts the latest model , wt , to all the devices . Second , every device ( say the k-th ) lets wkt = wt and then performs E ( ≥ 1 ) local updates : wkt+i+1 ←− wkt+i − ηt+i∇Fk ( wkt+i , ξkt+i ) , i = 0 , 1 , · · · , E − 1 where ηt+i is the learning rate ( a.k.a . step size ) and ξkt+i is a sample uniformly chosen from the local data . Last , the server aggregates the local models , w1t+E , · · · , wNt+E , to produce the new global model , wt+E . Because of the non-iid and partial device participation issues , the aggregation step can vary . 4It is well know that the full gradient descent ( which is equivalent to FedAvg with E = 1 and full batch ) do not require the decay of learning rate . 5The sampling scheme is proposed by Sahu et al . ( 2018 ) for FedAvg as a baseline , but this convergence rate is our contribution . 6The convergence relies on the assumption that data are balanced , i.e. , n1 = n2 = · · · = nN . However , we can use a rescaling trick to get rid of this assumption . We will discuss this point later in Section 3 . IID versus non-iid . Suppose the data in the k-th device are i.i.d . sampled from the distribution Dk . Then the overall distribution is a mixture of all local data distributions : D = ∑N k=1 pkDk . The prior work Zhang et al . ( 2015a ) ; Zhou and Cong ( 2017 ) ; Stich ( 2018 ) ; Wang and Joshi ( 2018 ) ; Woodworth et al . ( 2018 ) assumes the data are iid generated by or partitioned among the N devices , that is , Dk = D for all k ∈ [ N ] . However , real-world applications do not typically satisfy the iid assumption . One of our theoretical contributions is avoiding making the iid assumption . Full device participation . The prior work Coppola ( 2015 ) ; Zhou and Cong ( 2017 ) ; Stich ( 2018 ) ; Yu et al . ( 2019 ) ; Wang and Joshi ( 2018 ) ; Wang et al . ( 2019 ) requires the full device participation in the aggregation step of FedAvg . In this case , the aggregation step performs wt+E ←− N∑ k=1 pkw k t+E . Unfortunately , the full device participation requirement suffers from serious “ straggler ’ s effect ” ( which means everyone waits for the slowest ) in real-world applications . For example , if there are thousands of users ’ devices in the FL system , there are always a small portion of devices offline . Full device participation means the central server must wait for these “ stragglers ” , which is obviously unrealistic . Partial device participation . This strategy is much more realistic because it does not require all the devices ’ output . We can set a threshold K ( 1 ≤ K < N ) and let the central server collect the outputs of the first K responded devices . After collecting K outputs , the server stops waiting for the rest ; the K + 1-th to N -th devices are regarded stragglers in this iteration . Let St ( |St| = K ) be the set of the indices of the first K responded devices in the t-th iteration . The aggregation step performs wt+E ←− N K ∑ k∈St pkw k t+E . It can be proved that NK ∑ k∈St pk equals one in expectation . Communication cost . The FedAvg requires two rounds communications— one broadcast and one aggregation— per E iterations . If T iterations are performed totally , then the number of communications is b 2TE c. During the broadcast , the central server sends wt to all the devices . During the aggregation , all or part of the N devices sends its output , say wkt+E , to the server . 3 CONVERGENCE ANALYSIS OF FEDAVG IN NON-IID SETTING In this section , we show that FedAvg converges to the global optimum at a rate of O ( 1/T ) for strongly convex and smooth functions and non-iid data . The main observation is that when the learning rate is sufficiently small , the effect of E steps of local updates is similar to one step update with a larger learning rate . This coupled with appropriate sampling and averaging schemes would make each global update behave like an SGD update . Partial device participation ( K < N ) only makes the averaged sequence { wt } have a larger variance , which , however , can be controlled by learning rates . These imply the convergence property of FedAvg should not differ too much from SGD . Next , we will first give the convergence result with full device participation ( i.e. , K = N ) and then extend this result to partial device participation ( i.e. , K < N ) .

Summary: This paper analyzes the convergence of FedAvg, the most popular algorithm for federated learning. The highlight of the paper is removing the following two assumptions: (i) the data are iid across devices, and (ii) all the devices are active. For smooth and strongly convex problems, the paper proves an O(1/T) convergence rate to global optimum for learning rate decaying like 1/t with time. It is also shown that with constant learning rate eta, the solution found can be necessarily Omega(eta) away from the optimum (for a specific problem instance), thus justifying the decaying learning rate used in the positive result.