Reducing the Communication Cost of Federated Learning through Multistage Optimization

1 INTRODUCTION . In federated learning ( FL ) ( McMahan et al. , 2017 ; Kairouz et al. , 2019 ; Li et al. , 2020 ) , distributed clients interact with a central server to learn a model without directly sharing their data with the server . The training objective is to solve the following minimization problem : min x [ F ( x ) = 1 N ∑ i∈ [ N ] Fi ( x ) ] ( 1 ) where variable x is the model parameter , i indexes the clients ( or devices ) , N is the number of clients , and Fi ( x ) is a ( strongly convex ) loss that only depends on that client ’ s data . Typical FL deployments have two properties that make optimizing Eq . ( 1 ) challenging : ( i ) Data heterogeneity : Let Fi ( x ) = Ezi∼Pi [ f ( x ; zi ) ] for some loss function f , where zi is the client ’ s data and Pi ’ s capture the heterogeneity of the clients ’ data distributions . Formally , we define client heterogeneity ζ2 and heterogeneity at the optimum ζ2∗ as follows , where the optimum is defined as x ∗ : = arg minx F ( x ) : ζ2 : = sup x 1 N ∑ i∈ [ N ] ‖∇F ( x ) −∇Fi ( x ) ‖2 , ζ2∗ : = ( 1/N ) ∑ i∈ [ N ] ‖∇Fi ( x∗ ) ‖2 , ( 2 ) These definitions capture how different the local client gradients are from the overall global gradient ( worst-case ) ( Woodworth et al. , 2020a ) . ( ii ) Communication cost : In many deployments of FL , clients have limited bandwidths ( e.g. , mobile devices ) and hence high communication cost . Due to these two challenges , most federated optimization algorithms alternate between local rounds of computation , where clients process only their own data to save communication , and global rounds , where clients synchronize with the central server to resolve disagreements due to heterogeneity . We assume clients communicate with the server in each global round ( full participation ) . This commonly-studied setting ( Woodworth et al. , 2020a ; Gorbunov et al. , 2020 ; Yuan et al. , 2021 ; Mitra et al. , 2021 ; Karimireddy et al. , 2020b ; Reddi et al. , 2020 ) applies to cross-silo FL ( Kairouz et al. , 2019 ) . Several federated algorithms navigate this trade-off between reducing communication and resolving data heterogeneity by modifying the amount and the nature of local and global computation ( McMahan et al. , 2017 ; Li et al. , 2018 ; Wang et al. , 2019b ; a ; Li et al. , 2019 ; Karimireddy et al. , 2020b ; a ; Al-Shedivat et al. , 2020 ; Reddi et al. , 2020 ; Charles & Konečnỳ , 2020 ; Mitra et al. , 2021 ) . Recently , Woodworth et al . ( 2020a ) showed a lower bound on the convergence rate ( with respect to number of communication rounds R ) of any first-order federated optimization algorithms for strongly-convex , smooth objectives as a function of client data heterogeneity ζ∗ ( Section 2 , Eq . ( 6 ) ) . This lower bound highlights two gaps in the literature : ( 1 ) Today , no single optimization algorithm has been shown to match this lower bound for all values of ζ∗ . ( 2 ) Although different baseline optimization algorithms match the lower bound in different heterogeneity regimes , it remains unknown if any algorithm can match the lower bound when 0 < ζ2∗ < β 3/2∆/µ1/2 , where ∆ is the initial function suboptimality ( see Assumption 6 ) , and the functions Fi are β-smooth and µ-strongly convex ( see App . C for definitions of smoothness and strong convexity ) . ∆ can be sizable if one starts optimization far away from the optimum , and the gap between the upper and the lower bound in this regime is significant ( § 2 ) . To explain these gaps more precisely , we first examine several baseline optimization algorithms . To begin , consider the simple baseline of Minibatch Stochastic Gradient Descent ( SGD ) ( Woodworth et al. , 2020a , shortened as SGD ) . Here , in each global communication round , the central server sends each device the current model . Then each device calculates a single local stochastic gradient at this model over a minibatch of size K and sends it to the server . The server averages these to make a global minibatch gradient over NK samples and performs a model update ( Algorithm 2 ) . Further , this approach can be accelerated in the Nesterov sense ( Ghadimi & Lan , 2012 , AC-SA ) . Woodworth et al . ( 2020a ) recently showed that Accelerated Minibatch SGD ( shortened as AC-SA ) is nearly optimal ( up to exponentially shrinking condition numbers factors in communication rounds R ) in the highly-heterogeneous regimes when ζ2∗ > β 3/2∆/µ1/2 . However , in other regimes , AC-SA may be suboptimal as it does not exploit local computation to reduce communication . A popular alternative exploiting local computation is Federated Averaging ( FedAvg ) ( McMahan et al. , 2017 ) 1 ( Algorithm 6 ) . Here each client i independently runs SGD on its own objective Fi for K iterations ( local updates ) , and periodically communicates their respective most recent iterate to the server , which then averages these iterates . Compared to ( Accelerated ) SGD , this approach allows clients to learn from local data in between communication rounds . However , on strongly convex objectives , FedAvg has only been shown to achieve the lower bound when the data is homogeneous , i.e . ζ2∗ = ζ 2 = 0 . In fact , FedAvg ’ s known rates can only outperform AC-SA ’ s rates when heterogeneity is almost insignificant : ζ2 ≤ µ , where is the target sub-optimality gap for Eq . ( 1 ) ( Woodworth et al. , 2020a ) ( for more details on see Definition 1 ) . Further , an algorithm specific lower bound shows that FedAvg is provably sub-optimal ; improving the analysis will not bring FedAvg meaningfully closer to optimal rates or even allow FedAvg to beat AC-SA when heterogeneity is non-negligible Woodworth et al . ( 2020a ) . Several approaches have been proposed to improve FedAvg via acceleration in the homogeneous setting Yuan & Ma ( 2020 ) and cross-device variance reduction ( Karimireddy et al. , 2020b ; Gorbunov et al. , 2020 ) and adaptivity and momentum ( Reddi et al. , 2020 ) in the heterogeneous setting . Empirically , these variants outperform baselines like ( Accelerated ) SGD , though gains in the heterogeneous setting are not explained by existing analyses . In summary , currently there are only two baseline optimization algorithms that ( nearly , up to exponentially shrinking condition numbers ) theoretically match the lower bound of Woodworth et al . ( 2020a ) in some heterogeneity regime : FedAvg when ζ2 = 0 ( i.e. , ζ2∗ = 0 ) and AC-SA when ζ2∗ > β 3/2∆/µ1/2 . Several natural questions follow . ( 1 ) Does there exist an algorithm that achieves the lower bound when 0 < ζ2∗ < β 3/2∆/µ1/2 ? In fact , Woodworth et al . ( 2020a ) posed a weaker open question , asking whether it was possible to improve on AC-SA in the regime where data heterogeneity was bounded but not insignificant . ( 2 ) If so , can the same algorithm simultaneously achieve optimal convergence rates across all values of ζ∗ ? This Work . We answer both questions affirmatively when ζ ' ζ∗ by drawing on multistage algorithms , in which learning rates decay during optimization . Such algorithms have achieved fast convergence rates in convex optimization ( Aybat et al. , 2019 ; Fallah et al. , 2020 ) . In the federated setting , they have improved practical performance ( Reddi et al. , 2020 ; Charles & Konečnỳ , 2020 ) . 1This approach is also referred to as Local Stochastic Gradient Descent ( Stich , 2018 ; Khaled et al. , 2020 ; Woodworth et al. , 2020b ) , LocalUpdate ( Charles & Konečnỳ , 2020 ) . Our work exploits the following insight : far from the optimum , local computation can reduce communication ; near the optimum , it is important to communicate frequently ( per model update ) to resolve heterogeneity . This intuition has been explored in the homogeneous setting ( Wang & Joshi , 2019 ) , where the theoretical benefits do not appear to give order-wise convergence improvements . Contributions . In this paper , we first show that multistage optimization , when applied to baseline federated algorithms , achieves comparable or better worst-case convergence rates than known baselines ( including SGD and FedAvg ) in all heterogeneity regimes . In doing so , multistage algorithms resolve an open problem from ( Woodworth et al. , 2020b ) , which asked if one can design an algorithm that combines the advantages of both FedAvg and SGD and enjoys guarantees that dominate both . We propose multistage algorithms that ( nearly ) match the lower bound of Woodworth et al . ( 2020a ) on convergence rates for first-order federated algorithms . We show these theoretical gains by first analyzing an algorithm that ( nearly ) achieves the lower bound : a simple two-stage combination of ASG ( a variant of AC-SA in Theorem 6 ) and FedAvg . Loosely , we first run FedAvg up to an “ optimization error floor ” that depends on client heterogeneity ; then , we use ASG to complete the optimization ( Thm . 2 ) . Further , we show that similar gains can be obtained by a more general multistage algorithm which gradually moves from the FedAvg update to ASG update . This multistage variant ( Thm . 3 ) has the added benefit of not requiring the knowledge of the noise bound σ2 , the heterogeneity ζ2 , or the initial distance to the optimum ∆ . This intuition is not specific to the combination of ASG and FedAvg . We demonstrate the generality of multistaging by applying it to other pairs of FL optimization algorithms to obtain the same nearly-optimal worst-case convergence ( but possibly much better average-case rates , § 4 ) . Finally , we demonstrate the practical value of these theoretical insights through both convex and nonconvex experiments in § 4 . 2 MODEL AND RELATED WORK . We assume that each Fi is µ-strongly convex ( Assumption 1 ) and β-smooth ( Assumption 2 ) . We let κ = β/µ denote the condition number . N is the number of clients ; all clients participate in each communication round . A discussion of when a subset of clients participate per round is in App . L. Heterogeneity ζ and heterogeneity at optimum ζ∗ are defined as in equation 2 ( Assumptions 3 and 4 , respectively ) . We assume unless otherwise specified that ζ2 > 0 , that the problem is heterogeneous . We assume the client gradient variance is upper bounded as σ2 , formally defined in Assumption 5 . We bound the initial suboptimality gap as ∆ ( see Assumption 6 ) . Optimization proceeds in rounds ; K is the number of client iterations per communication round , and R is the number of communication rounds . In some cases , clients take minibatches of size B in their local steps ; unless specified otherwise , B is 1 . Our goal is to minimize the sub-optimality error EF ( x̂ ) − F ( x∗ ) ( where x̂ is the solution returned by an optimization algorithm ) after R rounds ( or to minimize communication round complexity , the number of rounds R to a given error ) . R is a standard proxy for the communication cost . Since communication costs in FL are typically much larger than computation cost ( Kairouz et al. , 2019 ) , our goal is to minimize round complexity R for sufficiently large K. We use the notation Õ to hide polylogarithmic factors .

Summary: It is known that if the level of heterogeneity is sufficiently high, then accelerated minibatch SGD is optimal for federated optimization matching the known lower bound of (Woodworth et al., 2020a). On the other hand, when the level of heterogeneity is very low, then FedAvg/LocalSGD outperforms the former in terms of communication complexity and needs only a few communication rounds given enough local computation. This paper *proposes* a multi-stage optimization procedure and *claims* that it nearly matches the lower bound for all heterogeneity levels.