Linear Convergence in Federated Learning: Tackling Client Heterogeneity and Sparse Gradients

1 Introduction . In a canonical federated learning ( FL ) architecture , a set S of clients periodically communicate with a central server to find a global statistical model that solves the following problem [ 1–5 ] : min x∈Rd f ( x ) , where f ( x ) = 1 m m∑ i=1 fi ( x ) . ( 1 ) Here , m is the number of clients , fi : Rd → R is the local objective ( loss ) function of client i , and f ( x ) is the global objective function . Some of the core distinguishing tenets of the FL paradigm are as follows [ 1–5 ] . First , due to privacy considerations , clients can not directly share their local training data with the server . Second , differences in the clients ’ data-sets may cause the clients to have nonidentical loss functions with different minima - this is known as statistical or objective heterogeneity . Third , due to variability in hardware ( CPU , memory ) and power ( battery level ) , i.e. , due to systems or device heterogeneity , the client devices may have different computation speeds ; in particular , this may lead to slow and straggling devices that affect convergence guarantees . Finally , communicationefficiency is a major concern , dictating the need to reduce the number of communication rounds , and also the size of the messages transmitted in each round . The above considerations pose unique technical challenges that we aim to address in this paper . In a typical FL setting , to reduce the number of communication rounds , clients perform multiple local training steps in isolation before communicating with the server . Due to such local steps , the popular 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . FedAvg algorithm suffers from a “ client-drift phenomenon '' under objective heterogeneity [ 6–11 ] : the local iterates of each client drift-off towards the minimum of their own local loss function , leading to slow convergence rates . For analysis on FedAvg , we refer the reader to [ 6 , 8 , 12–21 ] . Recently , several new algorithms such as FedProx [ 22 ] , SCAFFOLD [ 11 ] , FedSplit [ 10 ] , and FedNova [ 23 ] have been proposed as improvements to FedAvg . Despite these advances , there remain gaps in our understanding of the extent to which these algorithms match the guarantees of a centralized baseline.1 For instance , even for simple , deterministic settings , FedProx [ 22 ] and FedNova [ 23 ] exhibit a fundamental speed-accuracy conflict under objective heterogeneity ; see [ 8 , 9 ] and Section 2 . Specifically , with constant step-sizes , these algorithms converge linearly , but potentially to an incorrect point . Thus , convergence to the minimum of the global loss function necessitates diminishing step-sizes , which , in turn , leads to sub-linear convergence . Thus , fast convergence comes at the expense of accuracy . Although SCAFFOLD [ 11 ] and FedSplit [ 10 ] employ variance-reduction and operatorsplitting techniques , respectively , to tackle objective heterogeneity , it is not known whether the rates in these papers are tight . More importantly , neither SCAFFOLD nor FedSplit account for the effects of systems heterogeneity or compression , both of which are key challenges in FL . Indeed , due to systems heterogeneity , the number of local steps may vary across clients , causing some clients to make much less progress than others in each round [ 23 ] . Moreover , while empirical studies [ 24 , 25 ] have revealed significant benefits of biased sparsification , theoretical guarantees for such methods in a federated setting have remained elusive . In this context , our contributions are as follows . • A New Algorithm : Motivated by the above concerns , we develop a general algorithmic framework called FedLin that simultaneously accounts for objective heterogeneity , systems heterogeneity , and gradient sparsification . The key components of FedLin include a gradient correction term in the local update rule that exploits memory ; the use of client-specific learning rates ; and error-feedback mechanisms at the clients and the server . • Matching Centralized Rates : For smooth and strongly convex losses , we show that FedLin converges to the global minimum linearly in the deterministic setting , and with a O ( 1/T ) rate for a general stochastic oracle model , thereby matching centralized rates ( up to constants ) . We then present matching rates for smooth , convex and non-convex settings as well . Importantly , our results hold under arbitrary objective and systems heterogeneity . In contrast , the only other work in FL ( as far as we are aware ) that investigates both objective and systems heterogeneity [ 23 ] provides results only for the non-convex setting , under a bounded dissimilarity assumption . Moreover , the FedNova algorithm in [ 23 ] suffers from the speed-accuracy conflict , while FedLin does not . • Quantifying the Price of Multiple Local Steps : We establish a lower bound for FedLin that matches the upper-bound we obtain for smooth , strongly convex losses . In doing so , we provide the first ( as far as we are aware ) tight linear convergence rate analysis . Our lower bound highlights the price paid for performing multiple local steps , i.e. , the effect of infrequent communication on the convergence rate . In particular , our analysis reveals , perhaps surprisingly , that there exist simple instances ( involving quadratic losses ) for which performing multiple local steps does not improve the rate of convergence , indicating that even mild statistical heterogeneity can hurt . Our analysis also provides valuable insights into the limitations of gradient-tracking/variance-reduction techniques . • Analyzing the Impacts of Gradient Sparsification at Server and at Clients : While several works explore the effect of unbiased random quantization in distributed settings [ 26–31 ] , there are only a handful of papers [ 15 , 32 ] that also consider the effect of local steps in FL . Different from all these works , we explore the impacts of sparsifying gradients using a biased TOP-k operator , both at the server side and at the clients . Our results in this context ( i ) constitute the first formal study of gradient sparsification in a federated setting ; ( ii ) reveal key differences between up-link and down-link compression ; and ( iii ) quantify the effect of the compression level on the convergence rate . Notably , FedLin preserves linear convergence rates despite aggressive gradient sparsification . Basic Notation and Terminology : Referring to ( 1 ) , let x∗ ∈ argminx∈Rd f ( x ) , and x∗i ∈ argminx∈Rd fi ( x ) . Every FL algorithm mentioned in this paper operates in rounds t ∈ { 1 , . . . , T } . In each round t , every client performs a certain number of local steps in isolation , starting from a common global model x̄t . We will denote by x ( t ) i , ` client i ’ s estimate of the model at the ` -th local step of round t. In particular , x ( t ) i,0 = x̄t , ∀i ∈ S. 1By a centralized baseline , we refer to a setup where each client can communicate with every other client at all time steps via the server . 2 Motivation : Speed-Accuracy Trade-Off . To motivate our work , we first show how some recently proposed FL algorithms , namely FedProx [ 22 ] and FedNova [ 23 ] , exhibit a fundamental speed-accuracy trade-off even in simple , deterministic settings . Specifically , we show that these schemes do not , in general , guarantee convergence to the global minimum with constant step-sizes . This , in turn , necessitates diminishing step-sizes , leading to sub-linear convergence rates . Our analysis here is inspired by that in [ 8 ] for FedAvg . We consider a deterministic quadratic model where the local loss function of client i is given by fi ( x ) = 1/2‖A1/2i ( x− ci ) ‖2 , where Ai is a symmetric positive-definite matrix . We begin by assuming that all clients perform the same number of local steps H . The following is the FedProx update rule where a proximal term is added to mitigate client-drift : x ( t ) i , ` +1 = x ( t ) i , ` − η ( ∇fi ( x ( t ) i , ` ) + β ( x ( t ) i , ` − x̄t ) ) , ` = 0 , . . . , H − 1 ; x̄t+1 = 1 m ∑ i∈S x ( t ) i , H . ( 2 ) Proposition 1 . For any step-size η > 0 , T rounds of FedProx amount to performing T rounds of parallel GD on the surrogate optimization problem given by min x 1 m ∑ i∈S 1 2 ∥∥∥∥ ( H−1∑ ` =0 [ I − η ( Ai + βI ) ] ` Ai ) 1/2 ( x− ci ) ∥∥∥∥2 . ( 3 ) Proposition 1 shows that even when clients perform the same number of local updates , FedProx minimizes a surrogate objective function ( 3 ) whose minimum may not , in general , coincide with the minimum of the original problem . When β = 0 , FedProx reduces to FedAvg , and our observations continue to hold . To capture systems heterogeneity as in [ 23 ] , suppose now that client i performs τi local steps . Define τeff , 1/m ∑ i∈S τi and αi , τeff/τi , ∀i ∈ S. The update rule of FedNova relies on normalized aggregation of cumulative local gradients , and is given by x ( t ) i , ` +1 = x ( t ) i , ` − η∇fi ( x ( t ) i , ` ) ; x̄t+1 = x̄t − η m ∑ i∈S αi τi−1∑ ` =0 ∇fi ( x ( t ) i , ` ) , ( 4 ) where ` = 0 , . . . , τi − 1 , i ∈ S. Although FedNova can accommodate any local solver whose accumulated gradients are expressible as a linear combination of local gradients , we choose gradient descent , a simple solver , to isolate the impact of normalized aggregation - the essence of FedNova . Algorithm 1 FedLin 1 : Input : Client step-sizes ηi , i ∈ S , compression levels δc and δs , initial iterate x̄1 ∈ Rd , g1 = ∇f ( x̄1 ) , initial compression errors ρi,1 = 0 , ∀i ∈ S and e1 = 0 2 : for t = 1 , . . . , T do 3 : for i = 1 , . . . , m do 4 : for ` = 0 , . . . , τi − 1 do 5 : x ( t ) i , ` +1 ← x ( t ) i , ` − ηi ( ∇fi ( x ( t ) i , ` ) −∇fi ( x̄t ) + gt ) ; x ( t ) i,0 = x̄t 6 : end for 7 : Transmit x ( t ) i , τi to server 8 : end for 9 : Server transmits x̄t+1 = 1/m ∑ i∈S x ( t ) i , τi 10 : for i = 1 , . . . , m do 11 : Transmit hi , t+1 = Cδc ( ρi , t +∇fi ( x̄t+1 ) ) to server 12 : ρi , t+1 ← ρi , t +∇fi ( x̄t+1 ) − hi , t+1 13 : end for 14 : Server transmits gt+1 = Cδs ( et + 1/m ∑ i∈S hi , t+1 ) 15 : et+1 ← et + 1/m ∑ i∈S hi , t+1 − gt+1 16 : end for Proposition 2 . For any step-size η > 0 , T rounds of FedNova amount to performing T rounds of parallel GD on the surrogate optimization problem given by min x 1 m ∑ i∈S 1 2 ∥∥∥∥ ( τi−1∑ ` =0 [ I − ηAi ] ` αiAi ) 1/2 ( x− ci ) ∥∥∥∥2 . ( 5 ) For the proofs of Propositions 1 and 2 , see Appendix B . Proposition 2 shows that in the presence of both objective and systems heterogeneity , FedNova minimizes a surrogate loss function whose minimum may not coincide with x∗.2 Observe from ( 3 ) and ( 5 ) that using a larger learning rate η introduces more distortion to the original problem . In Figure 1 , we see how FedProx and FedNova both converge to incorrect minimizers , even for simple instances with two clients and deterministic , quadratic losses . In contrast , FedLin , our proposed approach that we develop in the next section , guarantees linear convergence to the global minimum . Main Takeaway : The main message we want to convey here is that even for deterministic settings , there are non-trivial challenges posed by objective and systems heterogeneity that only get amplified when one additionally considers biased compression . For such scenarios , it is not at all apparent whether ( and to what extent ) one can match even the basic centralized benchmark of achieving linear convergence for smooth , strongly convex loss functions . To focus on the above unresolved issues , we will primarily consider a deterministic model in this paper . Nonetheless , the general approach we develop applies to the stochastic setting as well , as aptly demonstrated by Theorem 4 in Section 4 . 3 Proposed Algorithm : FedLin In this section , we develop our proposed algorithm FedLin , formally described in Algorithm 1 . FedLin is initialized from a common global iterate x̄1 ∈ Rd . For simplicity , we assume that g1 = ∇f ( x̄1 ) , i.e. , every client has access to the true gradient of f ( · ) initially ; we can allow g1 to be arbitrary as well without affecting the convergence guarantees . FedLin proceeds in rounds : in each round t , starting from a common global model x̄t , each client i performs τi local training steps in parallel , as per line 5 of Algorithm 1 . The key features of our local update rule are as follows : exploiting past gradients to account for objective heterogeneity , using client-specific step-sizes to tackle systems heterogeneity , and employing error-feedback to account for gradient sparsification . We now discuss each of these features in detail . 2In a follow-up work to [ 8 ] , the authors in [ 9 ] generalize their framework to encompass proximal methods such as FedProx as well . As such , Propositions 1 and 2 in this section turn out to be special cases of the results in [ 9 ] . We were not aware of [ 9 ] at the time of submission of this paper . To gain intuition regarding the local step in line 5 , note that the ideal local update at client i is x ( t ) i , ` +1 = x ( t ) i , ` − ηi∇f ( x ( t ) i , ` ) . However , this requires client i to have access to the gradients of all other clients - which it does not , since clients do not communicate between rounds . To get around this , client i exploits memory , and uses the gradient of the global function ∇f ( x̄t ) from the beginning of round t ( when the clients last communicated ) as a guiding direction in its update rule . However , since ∇f ( x̄t ) is evaluated at a stale point x ( t ) i,0 = x̄t , client i subtracts off ∇fi ( x̄t ) from ∇f ( x̄t ) , and adds in the most recently evaluated gradient∇fi ( x ( t ) i , ` ) . This results in the update rule : x ( t ) i , ` +1 = x ( t ) i , ` − ηi ( ∇fi ( x ( t ) i , ` ) −∇fi ( x̄t ) +∇f ( x̄t ) ) . Our local update rule in line 5 is precisely of the above form , where gt is an inexact version of ∇f ( x̄t ) to account for gradient sparsification . When each client i performs τi local-steps , our analysis reveals that the bound on the drift-term ‖xi , ` − x̄t‖ scales linearly in τi ( see Lemma 9 in Appendix F ) . Accordingly , to compensate for such drift at client i , the step-size ηi needs to be chosen to vary inversely with the number of local steps τi . In fact , the requirement that ηi ∝ 1/τi also turns out to be necessary ( see Theorem 5 ) , providing further motivation for the choice of client-specific learning rates in FedLin . To explain the gradient sparsification module , let us denote by Cδ : Rd → Rd the TOP-k operator , where δ = d/k , and k ∈ { 1 , . . . , d } . Given any x ∈ Rd , let Eδ ( x ) be a set containing the indices of the k largest-magnitude components of x . Then , the TOP-k operator we consider is given by ( Cδ ( x ) ) j = ( x ) j if j ∈ Eδ ( x ) , and ( Cδ ( x ) ) j = 0 otherwise . Here , we use ( x ) j to denote the j-th component of a vector x . Clearly , a larger δ implies more aggressive compression . We employ a standard error-feedback mechanism [ 33–35 ] at both the server and the clients to account for gradient sparsification . At client i , ρi , t represents the accumulated error due to gradient sparsification . At the end of round t , instead of just compressing∇fi ( x̄t+1 ) , client i instead compresses∇fi ( x̄t+1 ) + ρi , t , to account for gradient coordinates not transmitted in the past . It then updates the aggregate error via line 12 . An analogous description applies to the error-feedback scheme at the server , where et is the aggregate error at the beginning of round t. The parameters of FedLin are the client step-sizes { ηi } i∈S , and the compression levels δc and δs at the clients and at the server , respectively . We now comment on some related algorithmic ideas . Related Algorithmic Approaches : In the related but different setting of distributed optimization , we note that the idea of exploiting past gradients has been used to design gradient-tracking algorithms [ 36–40 ] . In the context of FL , this idea is also related to the variance-reduction technique employed in SCAFFOLD [ 11 ] . A major difference of FedLin with the above works is that none of them consider the effect of systems heterogeneity or biased compression . In particular , accounting for the inexact gradient term gt in our update rule introduces new technical challenges that we address in this paper . There are some additional basic differences between FedLin and SCAFFOLD . To see this , consider the update rule of FedLin without sparsification : x ( t ) i , ` +1 = x ( t ) i , ` − ηi ( ∇fi ( x ( t ) i , ` ) −∇fi ( x̄t ) +∇f ( x̄t ) ) . Now suppose the global model x̄t at the beginning of round t has already converged to x∗ . Since x ( t ) i,0 = x̄t , ∀i ∈ S , and ∇f ( x∗ ) = 0 , it is easy to see that the iterates of the clients do not evolve any further , as one would ideally want . Thus , the global optimum x∗ can be viewed as a fixed-point of the FedLin update rule . Adapting to our notation , and considering the case when there is no noise in the gradients , the update rule of SCAFFOLD takes the form x ( t ) i , ` +1 = x ( t ) i , ` − η ( ∇fi ( x ( t ) i , ` ) − ci + c ) , where ci is a ‘ control-variate ’ maintained by client i , and c is the average of the ci ’ s . Importantly , the control variates { ci } i∈S used in round t of SCAFFOLD contain stale terms from round t− 1 . As a result , even if x̄t = x∗ , it may very well be that ( ∇fi ( x̄t ) − ci + c ) 6= 0 , causing the iterates of the clients to move away from x∗ , and requiring further rounds of communication to average out the imbalance . Thus , the fixed-point property we discussed for FedLin does not hold in general for SCAFFOLD . Our simulations in Section 7 reveal that FedLin converges much faster relative to SCAFFOLD on a simple linear regression model ; we conjecture it is precisely due to the reason described above . Keeping aside the differences due to systems heterogeneity and compression , the FedSVRG algorithm in [ 1 ] includes a similar gradient correction term as in FedLin , but makes use of certain additional diagonal scaling and pre-conditioning matrices . Although promising empirical results are reported for FedSVRG in [ 1 ] , these results come with no supporting theoretical guarantees of convergence . In contrast , we will develop rigorous complexity guarantees for FedLin in the following sections . Specifically , we will show that FedLin guarantees linear convergence rates despite the challenges of objective heterogeneity , systems heterogeneity , and aggressive gradient sparsification .

Summary: This paper proposes an optimization method for federated learning, FedLin, that is specifically designed to ensure fast convergence even in the presence of heterogeneous data. Notably, the authors show that FedLin converges to the global optimum (or just a critical point, in the case of non-convex functions) without decaying the client learning rate. This is in contrast to methods such as FedAvg, FedProx, and FedNova, which require decaying the client loss in order to circumvent "objective mismatch", where the methods are actually optimizing an altered surrogate loss.  This work also derives lower bounds matching the convergence of FedLin on quadratics, that explain the price and cost of performing multiple local update steps. They also analyze convergence using gradient sparsification at both the client and server. Finally, the paper shows the empirical benefits of FedLin over other methods on synthetic least squares and logistic regression problems.