An Improved Analysis of Gradient Tracking for Decentralized Machine Learning

1 Introduction . Methods that train machine learning models on decentralized data offer many advantages over traditional centralized approaches in core aspects such as data ownership , privacy , fault tolerance and scalability [ 12 , 33 ] . Many current efforts in this direction come under the banner of federated learning [ 17 , 29 , 28 , 12 ] , where a central entity orchestrates the training and collects aggregate updates from the participating devices . Fully decentralized methods , that do not rely on a central coordinator and that communicate only with neighbors in an arbitrary communication topology , are still in their infancy [ 24 , 18 ] . The work of Lian et al . [ 24 ] on decentralized stochastic gradient descent ( D-SGD ) has spurred the research on decentralized training methods for machine learning models . This lead to improved theoretical analyses [ 16 ] and to improved practical schemes , such as support for time-varying topologies [ 32 , 3 , 16 ] and methods with communication compression [ 45 , 51 , 15 , 47 ] . One of the most challenging aspect when training over decentralized data is data-heterogeneity , i.e . training data that is in a non-IID fashion distributed over the devices ( for instance in data-center training ) or generated in non-IID fashion on client devices [ 21 , 13 , 22 , 23 ] . For example , the D-SGD method has been shown to be affected by the heterogenity [ 16 ] . In contrast , certain methods can mitigate the impact of heterogeneous data in decentralized optimization . For instance the gradient tracking ( GT ) methods developed by Lorenzo and Scutari [ 26 ] and Nedić et al . [ 34 ] , or the later D2 method by Tang et al . [ 46 ] which is designed for communication typologies that remain fixed and do not change over time . ∗Current affiliation : CISPA Helmholtz Center for Information Security . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . It is well known that GT methods do not depend on the heterogeneity of the data and that they converge linearly on distributed strongly convex problem instances without stochastic noise [ 26 , 34 ] . However , when we apply these methods in the context of machine learning , we need to understand how they are affected by stochastic noise and how they behave on non-convex tasks . In this paper , we develop a new , and improved , analysis of the gradient tracking algorithm with a novel proof technique . Along with the parallel contribution [ 55 ] that developed a tighter analysis of the D2 algorithm , we now have a more accurate understanding of in which setting GT works well and in which ones it does not , and our results allow for a more detailed comparison between the D-SGD , GT and D2 methods ( see Section 5 below ) . Our analysis improves over all existing results that analyze the GT algorithm . Specifically , we prove a weaker dependence on the connectivity of the network ( spectral gap ) which is commonly incorporated into the convergence rates via the standard parameter p. For example , in the strongly convex setting with stochastic noise we prove that GT converges at the rate Õ ( σ2 nε + 1 c · ( σ√ pε + 1 p log 1 ε ) ) where σ2 is an upper bound on the variance of the stochastic noise , and c ≥ p a new parameter ( often a constant ) . By comparing this result with the previously best known upper bound , Õ ( σ2 nε + 1 p · ( σ√ pε + 1 p log 1 ε ) ) , by Pu and Nedić [ 39 ] , we see that our upper bound improves the last two terms by a factor of cp ≥ 1 and that the first term matches with known lower bounds [ 37 ] . The D2 algorithm [ 46 ] only converges under the assumption that c is a constant2 and the recent upper bound from [ 55 ] coincides with our worst case complexity for GT on all topologies where D2 can be applied . We provide additional comparison of GT convergence rates in the Tables 1 and 2 . Contributions . Our main contributions can be summarized as : • We prove better complexity estimates for the GT algorithm than known before with a new proof technique ( which might be of independent interest ) . • In the non-asymptotic regime ( of importance in practice ) , the convergence rate depends on the network topology . By defining new graph parameters , we can give a tighter description of this dependency , explaining why the worst case behavior is rarely observed in practice ( see Section 5.1 ) . We verify this dependence in numerical experiments . • We show that in the presence of stochastic noise , the leading term in the convergence rate of GT is optimal—we are the first to derive this in the non-convex setting—and matching the unimprovable rate of all-reduce mini-batch SGD . 2 Related Work . Decentralized Optimization . Decentralized optimization methods have been studied for decades in the optimization and control community [ 48 , 30 , 52 , 6 ] . Many decentralized optimization methods 2In D2 the smallest eigenvalue of the mixing matrix W must bounded from below : mini λi ( W ) ≥ − 13 . [ 30 , 11 ] are based on gossip averaging [ 14 , 53 , 4 ] . Such methods usually also work well on nonconvex problems and can be used used for training deep neural networks [ 3 , 24 , 46 ] . There exists other methods , such as based on alternating direction method of multipliers ( ADMM ) [ 52 , 10 ] , dual averaging [ 6 , 31 , 41 ] , primal-dual methods [ 2 , 19 ] , block-coordinate methods for generalized linear models [ 8 ] or using new gradient propagation mechanisms [ 50 ] . Decentralized Optimization with Heterogeneous Objective Functions . There exists several algorithms that are agnostic to data-heterogeneity . Notably , EXTRA [ 42 ] and decentralized primal-dual gradient methods [ 2 ] do not depend on the data heterogeneity and achieve linear convergence in the strongly convex noiseless setting . However , these algorithms are not designed to be used for non-convex tasks . D2 [ 46 , 55 ] ( also known as exact diffusion [ 56 , 57 ] ) and Gradient Tracking ( GT ) [ 26 ] ( also known as NEXT [ 26 ] or DIGing [ 34 ] ) are both algorithms that are agnostic to the data heterogeneity level , can tolerate the stochastic noise , and that can be applied to non-convex objectives such as the training of deep neural networks in machine learning . A limitation of the D2 algorithm is that it is not clear how it can be applied to time-varying topologies , and that it can only be used on constant mixing topologies with negative eigenvalue bounded from below by − 13 . Other authors proposed algorithms that perform well on heterogeneous DL tasks [ 25 , 59 ] , but theoretical proofs that these algorithms are independent of the degree of heterogeneity are still pending . Gradient Tracking . There is a vast literature on the Gradient Tracking method itself . A tracking mechanism was used by Zhu and Martínez [ 61 ] as a way to track the average of a distributed continuous process . Lorenzo and Scutari [ 26 ] applied this technique to track the gradients , and analyzed its asymptotic behavior in the non-convex setting with a time-varying topologies . Nedić et al . [ 34 ] analyze GT ( named as DIGing ) in the strongly convex noiseless case with a time-varying network . Qu and Li [ 40 ] extend the GT analysis to the non-convex , weakly-convex and strongly convex case without stochastic noise . Nedić et al . [ 35 ] allow the different stepsizes on different workers . Yuan et al . [ 58 ] analyze asymptotic behavior of GT for dynamic optimization . Pu and Nedić [ 39 ] studied the GT method on stochastic problems and strongly convex objectives . Further , Xin et al . [ 54 ] analyze asymptotic behavior of GT with stochastic noise . For non-convex stochastic functions GT was analyzed by Zhang and You [ 60 ] and Lu et al . [ 27 ] . Li et al . [ 20 ] combine GT with variance reduction to achieve linear convergence in the stochastic case . Tziotis et al . [ 49 ] obtain second order guarantees for GT . 3 Setup . We consider optimization problems where the objective function is distributed across n nodes , min x∈Rd [ f ( x ) : = 1 n n∑ i=1 [ fi ( x ) = Eξ∼Di Fi ( x , ξ ) ] ] , ( 1 ) where fi : Rd → R denotes the local function available to the node i , i ∈ [ n ] : = { 1 , . . . n } . Each fi is a stochastic function fi ( x ) = Eξ∼Di Fi ( x , ξ ) with access only to stochastic gradients∇Fi ( x , ξ ) . This setting covers empirical risk minimization problems with Di being a uniform distribution over the local training dataset . It also covers deterministic optimization when Fi ( x , ξ ) = fi ( x ) , ∀ξ . We consider optimization over a decentralized network , i.e . when there is an underlying communication graph G = ( V , E ) , |V | = n , each of the nodes ( e.g . a connected device ) can communicate only along the edges E. In decentralized optimization it is convenient to parameterize communication by a mixing matrix W ∈ Rn×n , where wij = 0 if and only if nodes i and j are not communicating , ( i , j ) /∈ E. Definition 1 ( Mixing Matrix ) . A matrix with non-negative entries W ∈ [ 0 , 1 ] n×n that is symmetric ( W =W > ) and doubly stochastic ( W1=1 , 1 > W =1 > ) , where 1 denotes the all-one vector in Rn . 3.1 Notation . We use the notation x ( t ) i ∈ Rd , y ( t ) i ∈ Rd to denote the iterates and the tracking sequence , respectively , on node i at time step t. For vectors zi ∈ Rd ( zi could for instance be x ( t ) i or y ( t ) i ) defined for i ∈ [ n ] we denote by z̄ = 1n ∑n i=1 zi . We use both vector and matrix notation whenever it is more convenient . For vectors zi ∈ Rd defined for i ∈ [ n ] we denote by a capital letter the matrix with columns zi , formally Z : = [ z1 , . . . , zn ] ∈ Rd×n , Z̄ : = [ z̄ , . . . , z̄ ] ≡ Z 1n11 > , ∆Z : = Z − Z̄ . ( 2 ) We extend this definition to gradients of ( 1 ) , with∇F ( X ( t ) , ξ ( t ) ) , ∇f ( X ( t ) ) ∈ Rd×n : ∇F ( X ( t ) , ξ ( t ) ) : = [ ∇F1 ( x ( t ) 1 , ξ ( t ) 1 ) , . . . , ∇Fn ( x ( t ) n , ξ ( t ) n ) ] , ∇f ( X ( t ) ) : = [ ∇f ( x ( t ) 1 ) , . . . , ∇f ( x ( t ) n ) ] .

Summary: This paper analyzes Gradient Tracking, a very common decentralized optimization algorithm, in the stochastic setting. The authors first introduce the standard gradient tracking algorithms, along with the assumptions they use to analyze it. Then, they provide convergence results in 3 main settings: non-convex, strongly-convex and weakly-convex. In each case, the leading order term matches the rates of centralized mini-batch SGD, meaning that the methods are fast (at least when this term dominates). It is argued that many interesting extensions naturally fit in this framework and could be adapted without too much efforts.  Then, a proof sketch is given, and highlights the importance of studying the consensus difference X_t - \bar{X_t}, where \bar{X_t} is the average over all nodes, instead of studying just X_t. The other key ingredient is that a simple recursion, can be written with the consensus difference on both GT variables. Although the operator J involved in this recursion is not contractive at each step (spectral radius greater than 1), J^i with i large enough is contractive (spectral radius smaller than 1/2), which allows to prove convergence of the recursion.  Then, the rates obtained for Gradient Tracking are compared with rates previously obtained in the literature, and they compare favorably overall. Finally, some toy experiments are given to analyze the impact of the graph-dependent constants c and p that were introduced and that appear in the higher order terms. These experiments seem to confirm the theoretical rates.