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: A method to reduce communication costs when performing distributed training in networks where communication between nodes requires more than one hop. The trade off is increased memory usage and delayed communication (parameters arrive at a later iteration).  === Review Update ===  Thank you for the clarifications. In particular, it is true that all virtual topologies used by modern distributed ML systems are sparse, so I retract my comment regarding "comparison on fully-connected" topologies as it certainly does not make sense.  After reading both the authors' response and the other reviewers' comments, I've decided to keep my score as-is.