Federated Learning's Blessing: FedAvg has Linear Speedup

1 INTRODUCTION . Federated learning ( FL ) is a machine learning paradigm where many clients ( e.g. , mobile devices or organizations ) collaboratively train a model under the orchestration of a central server ( e.g. , service provider ) , while keeping the training data decentralized ( Smith et al . ( 2017 ) ; Kairouz et al . ( 2019 ) ) . In recent years , FL has swiftly emerged as an important learning paradigm ( McMahan et al . ( 2017 ) ; Li et al . ( 2020a ) ) –one that enjoys widespread success in applications such as personalized recommendation ( Chen et al . ( 2018 ) ) , virtual assistant ( Lam et al . ( 2019 ) ) , and keyboard prediction ( Hard et al . ( 2018 ) ) , to name a few–for at least three reasons : First , the rapid proliferation of smart devices that are equipped with both computing power and data-capturing capabilities provided the infrastructure core for FL . Second , the rising awareness of privacy and the explosive growth of computational power in mobile devices have made it increasingly attractive to push the computation to the edge . Third , the empirical success of communication-efficient FL algorithms has enabled increasingly larger-scale parallel computing and learning with less communication overhead . Despite its promise and broad applicability in our current era , the potential value FL delivers is coupled with the unique challenges it brings forth . In particular , when FL learns a single statistical model using data from across all the devices while keeping each individual device ’ s data isolated ( Kairouz et al . ( 2019 ) ) , it faces two challenges that are absent in centralized optimization and distributed ( stochastic ) optimization ( Zhou & Cong ( 2018 ) ; Stich ( 2019 ) ; Khaled et al . ( 2019 ) ; Liang et al . ( 2019 ) ; Wang & Joshi ( 2018 ) ; Woodworth et al . ( 2018 ) ; Wang et al . ( 2019 ) ; Jiang & Agrawal ( 2018 ) ; Yu et al . ( 2019b ; a ) ; Khaled et al . ( 2020 ) ; Koloskova et al . ( 2020 ) ; Woodworth et al . ( 2020b ; a ) ) : 1 ) Data ( statistical ) heterogeneity : data distributions in devices are different ( and data can not be shared ) ; κ1κ̃ are condition numbers defined in Section G. Since κ1 ≥ κ̃ , this implies a speedup factor of√ κ1 κ̃ for accelerated FedAvg . 2 ) System heterogeneity : only a subset of devices may access the central server at each time both because the communications bandwidth profiles vary across devices and because there is no central server that has control over when a device is active ( the presence of “ stragglers ” ) . To address these challenges , Federated Averaging ( FedAvg ) McMahan et al . ( 2017 ) was proposed as a particularly effective heuristic , which has enjoyed great empirical success . This success has since motivated a growing line of research efforts into understanding its theoretical convergence guarantees in various settings . For instance , Haddadpour & Mahdavi ( 2019 ) analyzed FedAvg ( for non-convex smooth problems satisfying PL conditions ) under the assumption that each local device ’ s minimizer is the same as the minimizer of the joint problem ( if all devices ’ data is aggregated together ) , an overly restrictive assumption that restricts the extent of data heterogeneity . Very recently , Li et al . ( 2020b ) furthered the progress and established an O ( 1T ) convergence rate for FedAvg for strongly convex smooth Federated learning problems with both data and system heterogeneity . A similar result in the same setting Karimireddy et al . ( 2019 ) also established an O ( 1T ) result that allows for a linear speedup when the number of participating devices is large . At the same time , Huo et al . ( 2020 ) studied the Nesterov accelerated FedAvg for non-convex smooth problems and established anO ( 1√ T ) convergence rate to stationary points . However , despite these very recent fruitful pioneering efforts into understanding the theoretical convergence properties of FedAvg , it remains open as to how the number of devices–particularly the number of devices that participate in the computation–affects the convergence speed . In particular , is linear speedup of FedAvg a universal phenomenon across different settings and for any number of devices ? What about when FedAvg is accelerated with momentum updates ? Does the presence of both data and system heterogeneity in FL imply different communication complexities and require technical novelties over results in distributed and decentralized optimization ? These aspects are currently unexplored or underexplored in FL . We fill in the gaps here by providing affirmative answers . Our Contributions We provide a comprehensive and unified convergence analysis of FedAvg and its accelerated variants in the presence of both data and system heterogeneity . Our contributions are threefold . First , we establish an O ( 1/KT ) convergence rate under FedAvg for strongly convex and smooth problems and an O ( 1/ √ KT ) convergence rate for convex and smooth problems ( where K is the number of participating devices ) , thereby establishing that FedAvg enjoys the desirable linear speedup property in the FL setup . Prior to our work here , the best and the most related convergence analysis is given by Li et al . ( 2020b ) and Karimireddy et al . ( 2019 ) , which established an O ( 1T ) convergence rate for strongly convex smooth problems under FedAvg . Our rate matches the same ( and optimal ) dependence on T , but also completes the picture by establishing the linear dependence on K , for any K ≤ N , where N is the total number of devices , whereas Li et al . ( 2020b ) does not have linear speedup analysis , and Karimireddy et al . ( 2019 ) only allows linear speedup close to full participation ( K = O ( N ) ) . As for convex and smooth problems , there was no prior work that established the O ( 1√ T ) rate under both system and data heterogeneity . Our unified analysis highlights the common elements and distinctions between the strongly and convex settings . Second , we establish the same convergence rates–O ( 1/KT ) for strongly convex and smooth problems and O ( 1/ √ KT ) for convex and smooth problems–for Nesterov accelerated FedAvg . We analyze the accelerated version of FedAvg here because empirically it tends to perform better ; yet , its theoretical convergence guarantee is unknown . To the best of our knowledge , these are the first results that provide a linear speedup characterization of Nesterov accelerated FedAvg in those two problem classes ( that FedAvg and Nesterov accelerated FedAvg share the same convergence rate is to be expected : this is the case even for centralized stochastic optimization ) . Prior to our results here , the most relevant results Yu et al . ( 2019a ) ; Li et al . ( 2020a ) ; Huo et al . ( 2020 ) only concern the non-convex setting , where convergence is measured with respect to stationary points ( vanishing of gradient norms , rather than optimality gaps ) . Our unified analysis of Nesterov FedAvg also illustrates the technical similarities and distinctions compared to the original FedAvg algorithm , whereas prior works ( in the non-convex setting ) were scattered and used different notations . Third , we study a subclass of strongly convex smooth problems where the objective is overparameterized and establish a faster O ( exp ( −KTκ ) ) convergence rate for FedAvg , in contrast to the O ( exp ( −Tκ ) ) rate for individual solvers Ma et al . ( 2018 ) . Within this class , we further consider the linear regression problem and establish an even sharper rate under FedAvg . In addition , we propose a new variant of accelerated FedAvg based on a momentum update of Liu & Belkin ( 2020 ) –MaSS accelerated FedAvg–and establish a faster convergence rate ( compared to if no acceleration is used ) . This stands in contrast to generic ( strongly ) convex stochastic problems where theoretically no rate improvement is obtained when one accelerates FedAvg . The detailed convergence results are summarized in Table 1 . Connections with Distributed and Decentralized Optimization Federated learning is closely related to distributed and decentralized optimization , and as such it is important to discuss connections and distinctions between our work and related results from that literature . First , when there is neither system heterogeneity , i.e . all devices participate in parameter averaging during a communication round , nor statistical heterogeneity , i.e . all devices have access to a common set of stochastic gradients , FedAvg coincides with the “ Local SGD ” of Stich ( 2019 ) , which showed the linear speedup rate O ( 1/NT ) for strongly convex and smooth functions . Woodworth et al . ( 2020b ) and Woodworth et al . ( 2020a ) further improved the communication complexity that guarantees the linear speedup rate . When there is only data heterogeneity , some works have continued to use the term Local SGD to refer to FedAvg , while others subsume it in more general frameworks that include decentralized model averaging based on a network topology or a mixing matrix . They have provided linear speedup analyses for strongly convex and convex problems , e.g . Khaled et al . ( 2020 ) ; Koloskova et al . ( 2020 ) as well as non-convex problems , e.g . Jiang & Agrawal ( 2018 ) ; Yu et al . ( 2019b ) ; Wang & Joshi ( 2018 ) . However , these results do not consider system heterogeneity , i.e . the presence of stragglers in the device network . Even with decentralized model averaging , the assumptions usually imply that model averages over all devices is the same as decentralized model averages based on network topology ( e.g . Koloskova et al . ( 2020 ) Proposition 1 ) , which precludes system heterogeneity as defined in this paper and prevalent in FL problems . For momentum accelerated FedAvg , Yu et al . ( 2019a ) provided linear speedup analysis for non-convex problems , while results for strongly convex and convex settings are entirely lacking , even without system heterogeneity . Karimireddy et al . ( 2019 ) considers both types of heterogeneities for FedAvg , but their rate implies a linear speedup only when the number of stragglers is negligible . In contrast , our linear speedup analyses consider both types of heterogeneity present in the full federated learning setting , and are valid for any number of participating devices . We also highlight a distinction in communication efficiency when system heterogeneity is present . Moreover , our results for Nesterov accelerated FedAvg completes the picture for strongly convex and convex problems . For a detailed comparison with related works , please refer to Table 2 in Appendix Section B . 2 SETUP . In this paper , we study the following federated learning problem : min w { F ( w ) , ∑N k=1 pkFk ( w ) } , ( 1 ) where N is the number of local devices ( users/nodes/workers ) and pk is the k-th device ’ s weight satisfying pk ≥ 0 and ∑N k=1 pk = 1 . In the k-th local device , there are nk data points : x 1 k , x 2 k , . . . , x nk k . The local objective Fk ( · ) is defined as : Fk ( w ) , 1nk ∑nk j=1 ` ( w ; xjk ) , where ` denotes a userspecified loss function . Each device only has access to its local data , which gives rise to its own local objective Fk . Note that we do not make any assumptions on the data distributions of each local device . The local minimum F ∗k = minw Fk ( w ) can be far from the global minimum of Eq ( 1 ) ( data heterogeneity ) .

Summary: This paper gives convergence analysis for FedAvg and its accelerated version under data heterogeneity and system heterogeneity. The main improvement comes from a more careful analysis for one-step descent where the authors make use of the term $\alpha_{t} \sum_{k=1}^{N} p_{k}\left[F_{k}\left(\mathbf{w}^{*}\right)-F_{k}\left(\overline{\mathbf{w}}_{t}\right)\right]$ that is ignored by previous work Li (2020b). The paper focuses on how FedAvg’s convergence scales with the number of participating devices. It improves previous analysis for FedAvg under more federated settings and shows that FedAvg has linear speedup for any number of participating devices.