RelaySum for Decentralized Deep Learning on Heterogeneous Data

1 Introduction . Ever-growing datasets lay at the foundation of the recent breakthroughs in machine learning . Learning algorithms therefore must be able to leverage data distributed over multiple devices , in particular for reasons of efficiency and data privacy . There are various paradigms for distributed learning , and they differ mainly in how the devices collaborate in communicating model updates with each other . In the all-reduce paradigm , workers average model updates with all other workers at every training step . In federated learning [ 24 ] , workers perform local updates before sending them to a central server that returns their global average to the workers . Finally , decentralized learning significantly generalizes the two previous scenarios . Here , workers communicate their updates with only few directly-connected neighbors in a network , without the help of a server . Decentralized learning offers strong promise for new applications , allowing any group of agents to collaboratively train a model while respecting the data locality and privacy of each contributor [ 25 ] . At the same time , it removes the single point of failure in centralized systems such as in federated learning [ 12 ] , improving robustness , security , and privacy . Even from a pure efficiency standpoint , decentralized communication patterns can speed up training in data centers [ 2 ] . In decentralized learning , workers share their local stochastic gradient updates with the others through gossip communication [ 41 ] . They send their updates to their neighbors , which iteratively ∗Equal contribution . Corresponding authors thijs.vogels @ epfl.ch and lie.he @ epfl.ch . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . propagate the updates further into the network . The workers typically use iterative gossip averaging of their models with their neighbors , using averaging weights chosen to ensure asymptotic uniform distribution of each update across the network . It will take τ rounds of communication for an update from worker i to reach a worker j that is τ hops away , and when it first arrives , the update is exponentially weakened by repeated averaging with weights < 1 . In general networks , worker j will never exactly , but only asymptotically receive its uniform share of the update . The slow distribution of updates not only slows down training , but also makes decentralized learning sensitive to heterogeneity in workers ’ data distributions . We study an alternative mechanism to gossip averaging , which we call RelaySum . RelaySum operates on spanning trees of the network , and distributes information exactly uniformly within a finite number of gossip steps equal to the diameter of the network . Rather than iteratively averaging models , each node acts as a ‘ router ’ that relays messages through the whole network without decaying their weight at every hop . While naive all-to-all routing requires n2 messages to be transmitted at each step , we show that on trees , only n messages ( one per edge ) are sufficient . This is enabled by the key observation that the routers can merge messages by summation to avoid any extra communication compared to gossip averaging . RelaySum achieves this using additional memory linear in the number of edges , and by tailoring the messages sent to different neighbors . At each time step , RelaySum workers receive a uniform average of exactly one message from each worker . Those messages just originate from different time delays depending on how many hops they travelled . The difference between gossip averaging and RelaySum is illustrated in Figure 1 . The RelaySum mechanism is structurally similar to Belief Propagation algorithms for inference in graphical models . This link was made by Zhang et al . [ 50 ] , who used the same mechanism for decentralized weighted average consensus in control . We use RelaySum in the RelaySGD learning algorithm . We theoretically show that this algorithm is not affected by differences in workers ’ data distributions . Compared to other algorithms that have this property [ 36 , 31 ] , RelaySGD does not require the selection of averaging weights , and its convergence does not depend on the spectral gap of the averaging matrix , but instead on the network diameter . While RelaySum is formulated for trees , it can be used in any decentralized network . We use the Spanning Tree Protocol [ 30 ] to construct spanning trees of any network in a decentralized fashion . RelaySGD often performs better on any such spanning tree than gossip-based methods on the original graph . When the communication network can be chosen freely , the algorithm can use double binary trees [ 33 ] . While these trees have logarithmic diameter and scale to many workers , RelaySGD in this setup uses only constant memory equivalent to two extra copies of the model parameters and sends and receives only two models per iteration . Surprisingly , in deep learning with highly heterogeneous data , prior methods that are theoretically independent of data heterogeneity [ 36 , 31 ] , perform worse than heuristic methods that do not have this property , but use cleverly designed time-varying communication topologies [ 2 ] . In extensive tests on image- and text classification , RelaySGD performs better than both kinds of baselines at equal communication budget . 2 Related work . Out of the multitude of decentralized optimization methods , first-order algorithms that interleave local gradient updates with a form of gossip averaging [ 29 , 11 ] show most promise for deep learning . Such algorithms are theoretically analyzed for convex and non-convex objectives in [ 28 , 11 , 29 ] , and [ 19 , 36 , 2 , 20 ] demonstrate that gossip-based methods can perform well in deep learning . In a gossip averaging step , workers average their local models with the models of their direct neighbors . The corresponding ‘ mixing matrix ’ is a central object of study . The matrix can be doublystochastic [ 29 , 19 , 16 ] , column-stochastic [ 38 , 26 , 39 , 2 ] , row-stochastic [ 40 , 44 ] , or a combination [ 42 , 43 , 32 ] . Column-stochastic methods use the push-sum consensus mechanism [ 13 ] and can be used on directed graphs . Our analysis borrows from the theory developed for those methods . While gossip averages in general requires an infinite number of steps to reach exact consensus , another line of work identifies mixing schemes that yield exact consensus in finite steps . For some graphs , this is possible with time-independent averaging weights [ 15 , 6 ] . One can also achieve finite-time consensus with time-varying mixing matrices . On trees , for instance , exact consensus can be achieved by routing updates to a root node and back , in exactly diameter number of steps [ 15 , 6 ] . On some graphs , tighter bounds can be established [ 8 ] . For fully-connected networks with n workers , Assran et al . [ 2 ] design a sparse time-varying communication scheme that yields exact consensus in a cycle of log n averaging steps and performs well in deep learning . The ‘ relay ’ mechanism of RelaySGD was previously used by Zhang et al . [ 50 ] in the control community for the decentralized weighted average consensus problem , but they do not use it in the context of optimization . Zhang et al . also introduce a modified algorithm for loopy graphs , but this modification makes the achieved consensus inexact . The ‘ relay ’ mechanism effectively turns a sparse graph into a fully-connected graph with communication delays . Work on delayed consensus [ 27 ] and optimization [ 37 , 1 ] analyzes such schemes for centralized distributed algorithms . Those consensus schemes are , however , not directly applicable to decentralized optimization . A fundamental challenge in decentralized learning is dealing with data that is not identically distributed among workers . Because , in this case , workers pursue different optima , workers may drift [ 29 ] and this can harm convergence . There is a large family of algorithms that introduce update corrections that provably mitigate such data heterogeneity . Examples applicable to non-convex problems are exact diffusion [ 45 ] , Gradient Tracking [ 22 , 31 , 48 ] , D2 [ 36 ] , PushPull [ 32 ] . To tackle the same challenge , Lin et al . [ 20 ] , Yuan et al . [ 46 ] propose modifications to local momentum to empirically improve performance in deep learning , but without provable guarantees . Lu and De Sa [ 23 ] propose DeTAG which overlaps multiple consecutive gossip steps and gradient computations to accelerate information diffusion . This technique could be applied to the RelaySum mechanism , too . 3 Method . Setup We consider standard decentralized optimization with data distributed over n ≥ 1 nodes : f ? : = minx∈Rd [ f ( x ) = 1n ∑n i=1 [ fi ( x ) : = Eξ∼Di Fi ( x , ξi ) ] ] . Here Di denotes the distribution of the data on node i and fi : Rd → R the local optimization objectives . Workers are connected by a network respecting a graph topology G = ( V , E ) , where V = { 1 , . . . , n } denotes the set of workers , and E the set of undirected communication links between them ( without self loops ) . Each worker i can only directly communicate with its neighbors Ni ⊂ V . Decentralized learning with gossip We consider synchronous first-order algorithms that interleave local gradient-based updates x ( t+1/2 ) i = x ( t ) i + u ( t ) i with message exchange between connected workers . For SGD with typical gossip averaging ( DPSGD [ 19 ] ) , the local updates can be written as u ( t ) i = −γ∇fi ( x ( t ) i , ξ ( t ) i ) , and the messages exchanged between pairs of connected workers ( i , j ) are m ( t ) i→j = x ( t+1/2 ) i ∈ Rd . Each timestep , the workers average their model with received messages , x ( t+1 ) i = Wiix ( t+1/2 ) i + ∑ j∈Ni Wijm ( t ) j→i , ( DP-SGD ) using averaging weights defined by a gossip matrix W ∈ Rn×n . In this scheme , an update u ( t1 ) i from any worker i will be linearly incorporated into the model x ( t2 ) j at a later timestep t2 with weight ( Wt2−t1 ) ij . The gossip matrix must be chosen such that these weights asymptotically converge to 1n , distributing all updates uniformly over the workers . This setup appears in , for example , [ 19 , 16 ] . Uniform model averaging If the graph topology is fully-connected , any worker can communicate with any other worker , and it is ideal to use ‘ all-reduce averaging ’ , x ( t+1 ) i = 1 n ∑n j=1 x ( t+1/2 ) j . Contrary to the decentralized scheme ( DP-SGD ) , this algorithm does not degrade in performance if data is distributed heterogeneously across workers . In sparsely connected networks , however , all-reduce averaging requires routing messages through the network . On arbitrary networks , such a routing protocol requires at least a number of communication steps equal to the network diameter τmax—the minimum number of hops some messages have to travel . RelaySGD In this paper , we approximate the all-reduce averaging update as x ( t+1 ) i = 1 n ∑n j=1 x ( t−τij+1/2 ) j , ( RelaySGD ) where τij is minimum number of network hops between workers i and j ( and τii = 0 ) . Since it takes τij steps to route a message from worker i to j , this scheme could be implemented using a peer-to-peer routing protocol like Ethernet . Of course , this naive implementation drastically increases the bandwidth used compared to gossip averaging . The key insight of this paper is that , on tree networks , the RelaySGD update rule can be implemented while using the same communication volume per step as gossip averaging , using additional memory linear in the number of a worker ’ s direct neighbors . RelaySum To implement RelaySGD , we require a communication mechanism that delivers sums of delayed ‘ parcels ’ s ( t ) w = ∑n j=1 p ( t−τwj ) j to each worker w in a tree network , where the parcel p ( t ) j is created by worker j at time t. To simplify the exposition , let us first consider the simplest type of tree network : a chain . In a chain , a worker w is connected to workers w− 1 and w+ 1 , if those exist , and the delays are τij = |i− j| . We can then decompose s ( t ) w = n∑ j=1 p ( t−τwj ) j = p ( t ) w + w−1∑ j=1 p ( t−τwj ) j︸ ︷︷ ︸ parcels from the ‘ left ’ + n∑ j=w+1 p ( t−τwj ) j︸ ︷︷ ︸ parcels from the ‘ right ’ . The sum of parcels from the ‘ left ’ will be sent as one message m ( w−1 ) →w from worker w − 1 to w , and the sum of data from the ‘ right ’ will be sent as one message m ( w+1 ) →w from w + 1 to w. Neighboring workers can compute these messages from the messages they received from their neighbors in the previous timestep . Compared to typical gossip averaging , RelaySum requires additional memory linear in the number of neighbors , but it uses the same volume of communication . Algorithm 1 shows how this scheme is generalized to general tree networks and incorporated into RelaySGD . Along with the model parameters , we send scalar counters that are used in the first few iterations of the algorithm t ≤ τmax to correct for messages that have not yet arrived . Spanning trees RelaySGD is formulated on tree networks , but it can be used on any communication graph by constructing a spanning tree . In a truly decentralized setting , we can use the Spanning Tree Protocol [ 30 ] used in Ethernet to find such trees in a decentralized fashion . The protocol elects a leader as the root of the tree , after which every other node finds the fastest path to this leader . On the other hand , when the decentralized paradigm is used in a data center to reduce communication , RelaySGD can run on double binary trees [ 33 ] used in MPI and NCCL [ 10 ] . The key idea of double binary trees is to use two different communication topologies for different parts of the model . We communicate odd coordinates using a balanced binary tree A , and communicate the even coordinates with a complimentary tree B . The trees A and B are chosen such that internal nodes ( with 3 edges ) Algorithm 1 RelaySGD Input : ∀ i , x ( 0 ) i = x ( 0 ) ; ∀ i , j , m ( −1 ) i→j = 0 , counts c ( −1 ) i→j = 0 , learning rate γ , tree network 1 : for t = 0 , 1 , . . . do 2 : for node i in parallel 3 : x ( t+ 1/2 ) i = x ( t ) i −γ∇fi ( x ( t ) i ) ( or Adam/momentum ) 4 : for each neighbor j ∈ Ni do 5 : Send m ( t ) i→j = x ( t+1/2 ) i + ∑ k∈Ni\jm ( t−1 ) k→i ( relay messages from other neighbors ) 6 : Send corresponding counters c ( t ) i→j = 1 + ∑ k∈Ni\j c ( t−1 ) k→i 7 : Receive ( m ( t ) j→i , c ( t ) j→i ) from node j 8 : end for 9 : n̄ ( t+1 ) i = 1 + ∑ j∈Ni c ( t ) j→i ( n̄ converges to the total number of workers ) 10 : xt+1i = 1 n̄ ( t+1 ) i ( x ( t+1/2 ) i + ∑ j∈Ni m ( t ) j→i ) ( = 1n ∑n j=1 x ( t−τij+1/2 ) j ) 11 : end for 12 : end for in one tree are leaves ( with only 1 edge ) in the other . Using the combination of two trees , RelaySGD requires only constant extra memory equivalent to at most 2 model copies ( just like the Adam optimizer [ 14 ] ) , and it sends and receives the equivalent of 2 models ( just like on a ring ) .

Summary: This paper proposes a decentralized SGD method, RelaySGD, for faster mixing in distributed deep learning, with the aim of improving performance on nodes with heterogeneous data sources. Convergence guarantees are provided for smooth non-convex functions under the weak growth condition, as well as for smooth strongly convex functions under the weak growth condition. The method relies on undirected, static, spanning-tree graphs.