Achieving Linear Speedup with Partial Worker Participation in Non-IID Federated Learning

mKT + 1T ) for full worker participation and a convergence rate O ( √ K√ nT + 1T ) for partial worker participation , where K is the number of local steps , T is the number of total communication rounds , m is the total worker number and n is the worker number in one communication round if for partial worker participation . Our results also reveal that the local steps in FL could help the convergence and show that the maximum number of local steps can be improved to T/m in full worker participation . We conduct extensive experiments on MNIST and CIFAR-10 to verify our theoretical results . 1 INTRODUCTION . Federated Learning ( FL ) is a distributed machine learning paradigm that leverages a large number of workers to collaboratively learn a model with decentralized data under the coordination of a centralized server . Formally , the goal of FL is to solve an optimization problem , which can be decomposed as : min x∈Rd f ( x ) : = 1 m m∑ i=1 Fi ( x ) , where Fi ( x ) , Eξi∼Di [ Fi ( x , ξi ) ] is the local ( non-convex ) loss function associated with a local data distribution Di and m is the number of workers . FL allows a large number of workers ( such as edge devices ) to participate flexibly without sharing data , which helps protect data privacy . However , it also introduces two unique challenges unseen in traditional distributed learning algorithms that are used typically for large data centers : • Non-independent-identically-distributed ( non-i.i.d . ) datasets across workers ( data heterogeneity ) : In conventional distributed learning in data centers , the distribution for each worker ’ s local dataset can usually be assumed to be i.i.d. , i.e. , Di = D , ∀i ∈ { 1 , ... , m } . Unfortunately , this assumption rarely holds for FL since data are generated locally at the workers based on their circumstances , i.e. , Di 6= Dj , for i 6= j . It will be seen later that the non-i.i.d assumption imposes significant challenges in algorithm design for FL and their performance analysis . • Time-varying partial worker participation ( systems non-stationarity ) : With the flexibility for workers ’ participation in many scenarios ( particularly in mobile edge computing ) , workers may randomly join or leave the FL system at will , thus rendering the active worker set stochastic and time-varying across communication rounds . Hence , it is often infeasible to wait for all workers ’ responses as in traditional distributed learning , since inactive workers or stragglers will significantly slow down the whole training process . As a result , only a subset of the workers may be chosen by the server in each communication round , i.e. , partial worker participation . In recent years , the Federated Averaging method ( FedAvg ) and its variants ( McMahan et al. , 2016 ; Li et al. , 2018 ; Hsu et al. , 2019 ; Karimireddy et al. , 2019 ; Wang et al. , 2019a ) have emerged as a prevailing approach for FL . Similar to the traditional distributed learning , FedAvg leverages local computation at each worker and employs a centralized parameter server to aggregate and update the model parameters . The unique feature of FedAvg is that each worker runs multiple local stochastic gradient descent ( SGD ) steps rather than just one step as in traditional distributed learning between two consecutive communication rounds . For i.i.d . datasets and the full worker participation setting , Stich ( 2018 ) and Yu et al . ( 2019b ) proposed two variants of FedAvg that achieve a convergence rate of O ( mKT + 1√ mKT ) with a bounded gradient assumption for both strongly convex and nonconvex problems , where m is the number of workers , K is the local update steps , and T is the total communication rounds . Wang & Joshi ( 2018 ) and Stich & Karimireddy ( 2019 ) further proposed improved FedAvg algorithms to achieve an O ( mT + 1√ mKT ) convergence rate without bounded gradient assumption . Notably , for a sufficiently large T , the above rates become O ( 1√ mKT ) 1 , which implies a linear speedup with respect to the number of workers.2 This linear speedup is highly desirable for an FL algorithm because the algorithm is able to effectively leverage the massive parallelism in a large FL system . However , with non-i.i.d . datasets and partial worker participation in FL , a fundamental open question arises : Can we still achieve the same linear speedup for convergence , i.e. , O ( 1√ mKT ) , with non-i.i.d . datasets and under either full or partial worker participation ? In this paper , we show the answer to the above question is affirmative . Specifically , we show that a generalized FedAvg with two-sided learning rates achieves linear convergence speedup with non-i.i.d . datasets and under full/partial worker participation . We highlight our contributions as follows : • For non-convex problems , we show that the convergence rate of the FedAvg algorithm on non-i.i.d . dataset areO ( 1√ mKT + 1T ) andO ( √ K√ nT + 1T ) for full and partial worker participation , respectively , where n is the size of the partially participating worker set . This indicates that our proposed algorithm achieves a linear speedup for convergence rate for a sufficiently large T . When reduced to the i.i.d . case , our convergence rate is O ( 1TK + 1√ mKT ) , which is also better than previous works . We summarize the convergence rate comparisons for both i.i.d . and non-i.i.d . cases in Table 1 . It is worth noting that our proof does not require the bounded gradient assumption . We note that the SCAFFOLD algorithm ( Karimireddy et al. , 2019 ) also achieves the linear speedup but extra variance reduction operations are required , which lead to higher communication costs and implementation complexity . By contrast , we do not have such extra requirements in this paper . • In order to achieve a linear speedup , i.e. , a convergence rate O ( 1√ mKT ) , we show that the number of local updatesK can be as large as T/m , which improves the T 1/3/m result previously shown in Yu et al . ( 2019a ) and Karimireddy et al . ( 2019 ) . As shown later in the communication complexity comparison in Table 1 , a larger number of local steps implies relatively fewer communication rounds , thus less communication overhead . Interestingly , our results also indicate that the number of local updates K does not hurt but rather help the convergence with a proper learning rates choice in full worker participation . This overcomes the limitation as suggested in Li et al . ( 2019b ) that local SGD steps might slow down the convergence ( O ( KT ) for strongly convex case ) . This result also reveals new insights on the relationship between the number of local steps and learning rate . 1This rate also matches the convergence rate order of parallel SGD in conventional distributed learning . 2To attain accuracy for an algorithm , it needs to take O ( 1 2 ) steps with a convergence rate O ( 1√ T ) , while needing O ( 1 m 2 ) steps if the convergence rate is O ( 1√ mT ) ( the hidden constant in Big-O is the same ) . In this sense , one achieves a linear speedup with respect to the number of workers . Notation . In this paper , we let m be the total number of workers and St be the set of active workers for the t-th communication round with size |St| = n for some n ∈ ( 0 , m ] . 3 We use K to denote the number of local steps per communication round at each worker . We let T be the number of total communication rounds . In addition , we use boldface to denote matrices/vectors . We let [ · ] it , k represent the parameter of k-th local step in the i-th worker after the t-th communication . We use ‖·‖2 to denote the ` 2-norm . For a natural number m , we use [ m ] to represent the set { 1 , · · · , m } . The rest of the paper is organized as follows . In Section 2 , we review the literature to put our work in comparative perspectives . Section 3 presents the convergence analysis for our proposed algorithm . Section 4 discusses the implication of the convergence rate analysis . Section 5 presents numerical results and Section 6 concludes this paper . Due to space limitation , the details of all proofs and some experiments are provided in the supplementary material . 2 RELATED WORK . The federated averaging ( FedAvg ) algorithm was first proposed by McMahan et al . ( 2016 ) for FL as a heuristic to improve communication efficiency and data privacy . Since then , this work has sparked many follow-ups that focus on FL with i.i.d . datasets and full worker participation ( also known as LocalSGD ( Stich , 2018 ; Yu et al. , 2019b ; Wang & Joshi , 2018 ; Stich & Karimireddy , 2019 ; Lin et al. , 2018 ; Khaled et al. , 2019a ; Zhou & Cong , 2017 ) ) . Under these two assumptions , most of the theoretical works can achieve a linear speedup for convergence , i.e. , O ( 1√ mKT ) for a sufficiently large T , matching the rate of the parallel SGD . In addition , LocalSGD is empirically shown to be communication-efficient and enjoys better generalization performance ( Lin et al. , 2018 ) . For a comprehensive introduction to FL , we refer readers to Li et al . ( 2019a ) and Kairouz et al . ( 2019 ) . 3For simplicity and ease of presentation in this paper , we let |St| = n. We note that this is not a restrictive condition and our proofs and results still hold for |St| ≥ n , which can be easily satisfied in practice . Algorithm 1 A Generalized FedAvg Algorithm with Two-Sided Learning Rates . Initialize x0 for t = 0 , · · · , T − 1 do The server samples a subset St of workers with |St| = n. for each worker i ∈ St in parallel do xit,0 = xt for k = 0 , · · · , K − 1 do Compute an unbiased estimate git , k = ∇Fi ( xit , k , ξit , k ) of∇Fi ( xit , k ) . Local worker update : xit , k+1 = x i t , k − ηLgit , k . end for Let ∆it = x i t , K − xit,0 = −ηL ∑K−1 k=0 g i t , k . Send ∆ i t to the server . end for At Server : Receive ∆it , i ∈ S. Let ∆t = 1|S| ∑ i∈S ∆ i t. Server Update : xt+1 = xt + η∆t . Broadcasting xt+1 to workers . end for For non-i.i.d . datasets , many works ( Sattler et al. , 2019 ; Zhao et al. , 2018 ; Li et al. , 2018 ; Wang et al. , 2019a ; Karimireddy et al. , 2019 ; Huang et al. , 2018 ; Jeong et al. , 2018 ) heuristically demonstrated the performance of FedAvg and its variants . On convergence rate with full worker participation , many works ( Stich et al. , 2018 ; Yu et al. , 2019a ; Wang & Joshi , 2018 ; Karimireddy et al. , 2019 ; Reddi et al. , 2020 ) can achieve linear speedup , but their convergence rate bounds could be improved as shown in this paper . On convergence rate with partial worker participation , Li et al . ( 2019b ) showed that the original FedAvg can achieve O ( K/T ) for strongly convex functions , which suggests that local SGD steps slow down the convergence in the original FedAvg . Karimireddy et al . ( 2019 ) analyzed a generalized FedAvg with two-sided learning rates under strongly convex , convex and non-convex cases . However , as shown in Table 1 , none of them indicates that linear speedup is achievable with non-i.i.d . datasets under partial worker participation . Note that the SCAFFOLD algorithm ( Karimireddy et al. , 2019 ) can achieve linear speedup but extra variance reduction operations are required , which lead to higher communication costs and implementation complexity . In this paper , we show that this linear speedup can be achieved without any extra requirements . For more detailed comparisons and other algorithmic variants in FL and decentralized settings , we refer readers to Kairouz et al . ( 2019 ) .

Summary: This paper provides a new analysis for the FedAvg algorithm, which assumes the data on different workers are non-IID and the objective functions are non-convex. The new analysis improved the existing bounds of FedAvg. Besides, the analysis is also extended to the non-stationary network, where the number of workers participating in the optimization may vary.