Communication-Efficient Sampling for Distributed Training of Graph Convolutional Networks

Training Graph Convolutional Networks ( GCNs ) is expensive as it needs to aggregate data recursively from neighboring nodes . To reduce the computation overhead , previous works have proposed various neighbor sampling methods that estimate the aggregation result based on a small number of sampled neighbors . Although these methods have successfully accelerated the training , they mainly focus on the single-machine setting . As real-world graphs are large , training GCNs in distributed systems is desirable . However , we found that the existing neighbor sampling methods do not work well in a distributed setting . Specifically , a naive implementation may incur a huge amount of communication of feature vectors among different machines . To address this problem , we propose a communication-efficient neighbor sampling method in this work . Our main idea is to assign higher sampling probabilities to the local nodes so that remote nodes are accessed less frequently . We present an algorithm that determines the local sampling probabilities and makes sure our skewed neighbor sampling does not affect much to the convergence of the training . Our experiments with node classification benchmarks show that our method significantly reduces the communication overhead for distributed GCN training with little accuracy loss . 1 INTRODUCTION . Graph Convolutional Networks ( GCNs ) are powerful models for learning representations of attributed graphs . They have achieved great success in graph-based learning tasks such as node classification ( Kipf & Welling , 2017 ; Duran & Niepert , 2017 ) , link prediction ( Zhang & Chen , 2017 ; 2018 ) , and graph classification ( Ying et al. , 2018b ; Gilmer et al. , 2017 ) . Despite the success of GCNs , training a deep GCN on large-scale graphs is challenging . To compute the embedding of a node , GCN needs to recursively aggregate the embeddings of the neighboring nodes . The number of nodes needed for computing a single sample can grow exponentially with respect to the number of layers . This has made mini-batch sampling ineffective to achieve efficient training of GCNs . To alleviate the computational burden , various neighbor sampling methods have been proposed ( Hamilton et al. , 2017 ; Ying et al. , 2018a ; Chen et al. , 2018b ; Zou et al. , 2019 ; Li et al. , 2018 ; Chiang et al. , 2019 ; Zeng et al. , 2020 ) . The idea is that , instead of aggregating the embeddings of all neighbors , they compute an unbiased estimation of the result based on a sampled subset of neighbors . Although the existing neighbor sampling methods can effectively reduce the computation overhead of training GCNs , most of them assume a single-machine setting . The existing distributed GCN systems either perform neighbor sampling for each machine/GPU independently ( e.g. , PinSage ( Ying et al. , 2018a ) , AliGraph ( Zhu et al. , 2019 ) , DGL ( Wang et al. , 2019 ) ) or perform a distributed neighbor sampling for all machines/GPUs ( e.g. , AGL ( Zhang et al. , 2020 ) ) . If the sampled neighbors on a machine include nodes stored on other machines , the system needs to transfer the feature vectors of the neighboring nodes across the machines . This incurs a huge communication overhead . None of the existing sampling methods or the distributed GCN systems have taken this communication overhead into consideration . In this work , we propose a communication-efficient neighbor sampling method for distributed training of GCNs . Our main idea is to assign higher sampling probabilities for local nodes so that remote nodes will be accessed less frequently . By discounting the embeddings with the sampling probability , we make sure that the estimation is unbiased . We present an algorithm to generate the sampling probability that ensures the convergence of training . To validate our sampling method , we conduct experiments with node classification benchmarks on different graphs . The experimental results show that our method significantly reduces the communication overhead with little accuracy loss . 2 RELATED WORK . The idea of applying convolution operation to the graph domain is first proposed by Bruna et al . ( 2013 ) . Later , Kipf & Welling ( 2017 ) and Defferrard et al . ( 2016 ) simplify the convolution computation with localized filters . Most of the recent GCN models ( e.g. , GAT ( Velickovic et al. , 2018 ) , GraphSAGE ( Hamilton et al. , 2017 ) , GIN ( Xu et al. , 2019 ) ) are based on the GCN in Kipf & Welling ( 2017 ) where the information is only from 1-hop neighbors in each layer of the neural network . In Kipf & Welling ( 2017 ) , the authors only apply their GCN to small graphs and use full batch for training . This has been the major limitation of the original GCN model as full batch training is expensive and infeasible for large graphs . Mini-batch training does not help much since the number of nodes needed for computing a single sample can grow exponentially as the GCN goes deeper . To overcome this limitation , various neighbor sampling methods have been proposed to reduce the computation complexity of GCN training . Node-wise Neighbor Sampling : GraphSAGE ( Hamilton et al. , 2017 ) proposes to reduce the receptive field size of each node by sampling a fixed number of its neighbors in the previous layer . PinSAGE ( Ying et al. , 2018a ) adopts this node-wise sampling technique and enhances it by introducing an importance score to each neighbor . It leads to less information loss due to weighted aggregation . VR-GCN ( Chen et al. , 2018a ) further restricts the neighbor sampling size to two and uses the historical activation of the previous layer to reduce variance . Although it achieves comparable convergence to GraphSAGE , VR-GCN incurs additional computation overhead for convolution operations on historical activation which can outweigh the benefit of reduced number of sampled neighbors . The problem with node-wise sampling is that , due to the recursive aggregation , it may still need to gather the information of a large number of nodes to compute the embeddings of a mini-batch . Layer-wise Importance Sampling : To further reduce the sample complexity , FastGCN ( Chen et al. , 2018b ) proposes layer-wise importance sampling . Instead of fixing the number of sampled neighbors for each node , it fixes the number of sampled nodes in each layer . Since the sampling is conduced independently in each layer , it requires a large sample size to guarantee the connectivity between layers . To improve the sample density and reduce the sample size , Huang et al . ( 2018 ) and Zou et al . ( 2019 ) propose to restrict the sampling space to the neighbors of nodes sampled in the previous layer . Subgraph Sampling : Layer-wise sampling needs to maintain a list of neighbors and calculate a new sampling distribution for each layer . It incurs an overhead that can sometime deny the benefit of sampling , especially for small graphs . GraphSAINT ( Zeng et al. , 2020 ) proposes to simplify the sampling procedure by sampling a subgraph and performing full convolution on the subgraph . Similarly , ClusterGCN ( Chiang et al. , 2019 ) pre-partitions a graph into small clusters and constructs mini-batches by randomly selecting subsets of clusters during the training . All of the existing neighbor sampling methods assume a single-machine setting . As we will show in the next section , a straightforward adoption of these methods to a distributed setting can lead to a large communication overhead . 3 BACKGROUND AND MOTIVATION . In a M -layer GCN , the l-th convolution layer is defined as H ( l ) = Pσ ( H ( l−1 ) ) W ( l ) where H ( l ) represents the embeddings of all nodes at layer l before activation , H ( 0 ) = X represents the feature vectors , σ is the activation function , P is the normalized Laplacian matrix of the graph , and W ( l ) is the learnable weights at layer l. The multiple convolution layers in the GCN can be represented as H ( M ) = Pσ ( H ( l−1 ) ( ... σ ( PXW ( 1 ) ︸ ︷︷ ︸ H ( 1 ) ) ... ) ) W ( M ) . ( 1 ) The output embedding H ( M ) is given to some loss function F for downstream learning tasks such as node classification or link prediction . GCN as Multi-level Stochastic Compositional Optimization : As pointed out by Cong et al . ( 2020 ) , training a GCN with neighbor sampling can be considered as a multi-level stochastic compositional optimization ( SCO ) problem ( although their description is not accurate ) . Here , we give a more precise connection between GCN training and multi-level SCO . Since the convergence property of algorithms for multi-level SCO has been extensively studied ( Yang et al. , 2019 ; Zhang & Xiao , 2019 ; Chen et al. , 2020 ) , this connection will allow us to study the convergence of GCN training with different neighbor sampling methods . We can define the graph convolution at layer l ∈ [ 1 , M ] as a function f ( l ) = Pσ ( H ( l−1 ) ) W ( l ) . The embedding approximation with neighbor sampling can be considered as a stochastic function f ( l ) ωl = P̃ ( l ) σ ( H ( l−1 ) ) W ( l ) where P̃ ( l ) is a stochastic matrix with Eωl [ P̃ ( l ) ] = P . Therefore , we have f ( l ) = Eωl [ f ( l ) ωl ] . The loss function of the GCN can be written as L ( θ ) = Eω ( M+1 ) [ f ( M+1 ) ω ( M+1 ) ( EωM [ f ( M ) ωM ( ... Eω1 [ f ( 1 ) ( θ ) ] ... ) ] ) ] . ( 2 ) Here , θ is the set of learnable weights at all layers { W ( 1 ) , ... , W ( M ) } , f ( M+1 ) = F ( H ( M ) ) , and the stochastic function f ( M+1 ) ω ( M+1 ) corresponds to the mini-batch sampling . Distributed Training of GCN : As the real-world graphs are large and the compute/memory capacity of a single machine is limited , it is always desirable to perform distributed training of GCNs . A possible scenario would be that we train a GCN on a multi-GPU system . The global memory of a single GPU can not accommodate the feature vectors of all nodes in the graph . It will be inefficient to store the feature vectors on the CPU main memory and move the feature vectors to GPU in each iteration of the training process because the data movement incurs a large overhead . We want to split the feature vectors and store them on multiple GPUs so that each GPU can perform calculation on its local data . Another possible scenario would be that we have a large graph with rich features which can not be store on a single machine . For example , the e-commerce graphs considered in AliGraph ( Zhu et al. , 2019 ) can ‘ contain tens of billions of nodes and hundreds of billions of edges with storage cost over 10TB easily ’ . Such graphs need to be partitioned and stored on different machines in a distributed system . Figure 1 shows an example of training a two-layer GCN on four GPUs . Suppose full neighbor convolution is used and each GPU computes the embeddings of its local nodes . GPU-0 needs to compute the embeddings of node A and B and obtain a stochastic gradient g̃0 based on the loss function . GPU-1 needs to compute the embeddings of node C and D and obtain a stochastic gradient g̃1 . Similarly , GPU-2 and GPU-3 compute the embeddings of their local nodes and obtain stochastic gradient g̃2 and g̃3 . The stochastic gradients obtained on different GPUs are then averaged and used to update the model parameters . Communication of Feature Vectors : As shown in Figure 1 , the computation of a node ’ s embedding may involve reading the feature vector of a remote node . To compute the embedding of node A on GPU-0 , we need the intermediate embeddings of node B and E , which in turn need the feature vectors of node A , B , C , E and F ( Note that the feature vector of node E itself is needed to compute its intermediate embedding ; the same for node B ) . Since node C , E , F are not on GPU-0 , we need to send the feature vectors of node C from GPU-1 and node E , F from GPU-2 . Similarly , to compute the embedding of node B on GPU-0 , we need feature vectors of node B , C , D , E , F and G , which means that GPU-0 needs to obtain data from all of the other three GPUs . This apparently incurs a large communication overhead . Even with neighbor sampling , the communication of the feature vectors among the GPUs are unavoidable . In fact , in our experiments on a four-GPU workstation , the communication can take more than 60 % of the total execution time with a naive adoption of neighbor sampling . The problem is expected to be more severe on distributed systems with multiple machines . Therefore , reducing the communication overhead for feature vectors is critical to the performance of distributed training of GCNs .

Summary: The paper presents a sampling-based approach to speeding-up training of GCNs in distributed systems. The key step in this task involves exchanging and aggregating messages sent along the edges of the graph. If the nodes of the graph are partitioned between several machines, then exchanging those messages involve costly communication between the nodes. To reduce this communication, the paper introduces a variant of the node sampling approach of Chen et al. (2018b) and Zou et al. (2019), where the probabilities of nodes in other machines are scaled down by some factor s. The approach is evaluated experimentally, showing that it reduces the amount of communication while (essentially) preserving the accuracy.