Distributed Training of Graph Convolutional Networks using Subgraph Approximation

1 INTRODUCTION . Graphs are used to model data in a diverse range of applications such as social networks ( Hamilton et al. , 2017a ) , biological networks ( Fout et al. , 2017 ) and e-commerce interactions ( Yang , 2019 ) . Recently , there has been an increased interest in applying machine learning techniques to graph data . A breakthrough was the development of Graph Convolutional Networks ( GCNs ) , which generalize the function of Convolutional Neural Networks ( CNNs ) , to operate on graphs ( Kipf & Welling , 2017 ) . Training GCNs , much like training CNNs , is a memory and computationally demanding task and may take days or weeks to train on large graphs . GCNs are particularly memory demanding since processing a single vertex may require accessing a large number of neighbors . If the GCN algorithm accesses neighbors which are multiple hops away , it may need to traverse disparate sections of the graph with no locality . This requirement poses challenges when the graph does not fit in memory , which is often the case for large graph representations . Previous approaches to alleviate issues of memory and performance seeks to constrain the number of dependent vertices by employing sampling-based methods ( Hamilton et al. , 2017b ; Chen et al. , 2018 ; Huang et al. , 2018 ; Chiang et al. , 2019 ) . However , sampling the graph may limit the achievable accuracy ( Jia et al. , 2020 ) . CNN training turned to data parallel distributed training to alleviate the memory size constraints and the computing capacity of a single machine . Here , a number of machines independently train the same neural network model but on an assigned chunk of the input data . To construct a single global model , the gradients or parameters are aggregated at regular intervals , such as at the completion of a batch ( Dean et al. , 2012 ) . While this approach works well for CNNs , it is more challenging for GCNs due to the nature of the input data . If the training input is a set of images , each image is independent . Hence , when the dataset is split into chunks of images , no information is directly carried between the chunks . In contrast , if the input is a graph , the equivalent of splitting the dataset into chunks is to divide the large graph into a number of subgraphs . However , GCNs train on graphs which rely not only on information inherent to a vertex but also information from its neighbors . When a GCN trains on one subgraph and needs information from a neighbor in another subgraph , the information from the edges that span across machines creates a communication bottleneck . For example , in our experimental evaluations we note that when the Reddit dataset ( Hamilton et al. , 2017b ) is partitioned into five subgraphs , there are about 1.3 million ( bidirectional ) edges that span across machines . If each edge is followed , 2.6 million messages would be sent , each containing a feature vector of 602 single-precision floats , summing up to about 6.3 GB of data transferred each epoch , where one epoch traverses the entire graph . Each training typically requires many epochs . In our experiments running distributed training on the GCN proposed by Kipf & Welling ( 2017 ) one epoch takes about one second . Thus the bandwidth required is enormous . This is a serious scalability issue : since the time it takes to finish one epoch decreases when the training is scaled out on more machines , the required bandwidth increases . Even with optimizations such as caching that might be possible ( e.g . vertex caching ( Yang , 2019 ) ) , the amount of communication is significant . Alternatively , instead of communicating between subgraphs , a GCN may ignore any information that spans multiple machines ( Jia et al. , 2020 ) . However , we observe that this approach results in accuracy loss due to suboptimal training . For example , when distributed on five compute nodes and ignoring edges that span across machines , we note in our evaluations that the accuracy of GraphSAGE trained on the Reddit dataset drops from 94.94 % to 89.95 % , a substantial loss in accuracy . To mitigate this problem , we present a distributed training approach of GCNs which reaches an accuracy comparable with training on a single-machine system , while keeping the communication overhead low . This allows for GCNs to train on very large graphs using a distributed system where each machine has limited memory , since the memory usage on each machine scales with the size of its allotted ( local ) subgraph . We achieve this by approximating the non-local subgraphs and making the approximations available in the local machine . We show that , by incurring a small vertex overhead in each local machine , a high accuracy can be achieved . Our contributions in this paper are : • A novel subgraph approximation scheme to enable distributed training of GCNs on large graphs using memory-constrained machines , while reaching an accuracy comparable to that achieved on a single machine . • An evaluation of the approach on two GCNs and two datasets , showing that it considerably increases convergence accuracy , and allows for faster training while keeping the memory overhead low . The rest of the paper is organized as follows . Section 2 walks through the necessary background and presents experimental data that motivates our approach . Section 3 describes our approach to efficiently distribute GCNs . Section 4 describes our evaluation methodology , while Section 5 presents the results of our evaluation . Section 6 discusses related work before we conclude in Section 7 . 2 BACKGROUND AND MOTIVATION . In this section we first describe GCNs . Then , we show that GCNs are resilient to approximations in the input graph , and describe how that motivates our approach in Section 3 . 2.1 GRAPH CONVOLUTIONAL NETWORKS . GCNs are neural networks that operate directly on graphs as the input data structure . A GCN aims to define an operation similar to a convolutional layer in CNNs , where a parameterized filter is applied to a pixel and its neighbors in an image ( Fig 1 ) . In the graph setting , the pixel corresponds to a vertex , and the edges from the vertex point out its neighbors ( Fig 2 ) . Early approaches to define this operation used graph convolutions from spectral graph theory ( Bruna et al. , 2014 ; Henaff et al. , 2015 ) . While theoretically sound , these approaches are difficult to scale to large graphs . ( Wu et al. , 2020 ) . Recent research ( e.g . ( Hamilton et al. , 2017b ; Chen et al. , 2018 ) ) instead takes a spatial-based approach , which enables computation in batches of vertices . The goal of a GCN is to learn the parameters of a neural network given a graph G = ( V , E ) where V is a set of vertices and E is a set of edges . The input is a feature matrix X , where each row represents the feature vector xv for every vertex v ∈ V , and an adjacency matrix A which represents the edges e ∈ E. Each layer in can be described as Hi = f ( Hi−1 , A ) , H0 = X , where i is the layer order and f is a propagation rule ( Kipf & Welling , 2017 ) . The number of layers and the propagation rule differ between GCNs . While our proposal is general to all GCNs of this type , we evaluate it on the GCN proposed by Kipf & Welling ( 2017 ) , which we call KW-GCN , and GraphSAGE proposed by Hamilton et al . ( 2017b ) . Both GCNs are described in more detail in Section 4.1 . Figure 1 : Image convolution . Figure 2 : Graph convolution . 2.2 LIMITATIONS TO INPUT APPROXIMATIONS . Our proposal builds upon the observation that GCNs may be resilient to approximations in the input graph . To investigate the extent of resilience , we trained both KW-GCN and GraphSAGE on the Reddit dataset ( Hamilton et al. , 2017b ) , described in Section 4.2 , for the task of vertex classification . To introduce approximations in the input graph , we randomly select an increasingly larger fraction of the training vertices to be deleted . We then use the approximated graph to train each GCN and compute the micro F1 score ( the micro-averaged harmonic mean of precision and recall , which we call accuracy ) on the complete test set . Figure 3 shows the results . We make the following qualitative observation : while GraphSAGE and KW-GCN are not equivalent in their resiliency towards vertex deletion , both networks demonstrate that up until a certain point it is possible to delete vertices without a large accuracy drop . While this might seem like a trivial but promising approach to train more efficiently by simply dropping vertices , in practice the threshold for deletion must be large to make a significant impact on the memory footprint : 50 % of the vertices has to be removed to reduce the memory footprints by 2X . Such a high deletion fraction incurs an unacceptable accuracy penalty . In this work , we instead leverage this observation to gain partial knowledge about non-local subgraphs in a distributed setting . While this approach increases the total workload by introducing redundancy , we show in Section 5 that only a small overhead is needed to substantially increase accuracy . In the next section we describe our proposed solution that exploits the approximation resilience property . 3 DISTRIBUTED TRAINING OF GCNS USING SUBGRAPH APPROXIMATION . A common approach to achieve distributed training is to leverage data parallelism : the training data is divided into mutually exclusive sets , and each machine is assigned one of these sets . Aside from decreasing training time , data parallelism also reduces memory footprint by lowering the amount of data each machine has to handle . If the graph does not fit in memory , the performance will be significantly impacted due to storage access time ( Ham et al. , 2016 ; Ahn et al. , 2015 ) . The main challenge of distributed subgraph training is accuracy . Each machine has only knowledge about the graph structure in its own partition . The result is that important information embedded in other parts of the graph is not captured in local subgraphs . Therefore , the local neural network may not learn global patterns . This is a unique problem to graph data structures . If the input to the neural network consists of images , no information is lost when the dataset is distributed , since each image is an independent data element . If the input is a graph , some edges that span the partitions are essentially unaccounted for during training . If the convolution operator ( Figure 2 ) targets a vertex at the edge of the subgraph , the neighbor data needed is not available in the local machine . A choice has to be made : either acknowledge the edge by trying to fetch the missing data from the non-local machine , or ignore the edge . Both approaches introduce inefficiencies . If the edge is acknowledged , there is a risk of overflowing the network with communication calls , since the machine that holds the neighboring vertex first has to be identified , and then that machine has to send out its data . The number of communication calls grows with the number of machines . However , if the edge is ignored , the execution time of each machine decreases since the communication overhead is minimized , but information is lost : the convolution will be approximate . Algorithm 1 : Breadth-first random sampling of vertices Input : Np : number of vertices in partition p ; M : number of partitions ; P : partitions ; O : overlap Output : Po : Partitions with overlap O 1 for p in P do 2 nts = O·Np M−1 ; Po [ p ] = p 3 for op in P \ p do 4 sample ( Po [ p ] , p , op , nts ) 5 Function sample ( Po [ p ] , p , op , nts ) : 6 selection = neighbors ( p , op ) 7 if size of selection > = nts then 8 sn = random ( selection , nts ) ; Po [ p ] = sn ∪ Po [ p ] 9 else 10 Po [ p ] = selection ∪ Po [ p ] 11 hop = neighbors ( selection , op ) \ selection 12 if ! hop is empty then sample ( Po [ p ] , hop , op , nts - size of selection ) ;

Summary: The paper presents a subgraph approximation method to reduce communication in distributed GCN training. The authors observe that up until a certain point it is possible to delete vertices in a graph without a large accuracy loss for GCN training.  Based on the observation, the authors propose the subgraph approximation idea. They claim that if the subgraph on each machine has additional information of 1 or 2 hop neighbors, the accuracy of GCN training can be maintained.