Learn Locally, Correct Globally: A Distributed Algorithm for Training Graph Neural Networks

1 INTRODUCTION . In recent years , Graph Neural Networks ( GNNs ) have achieved impressive results across numerous graph-based applications , including social networks ( Hamilton et al. , 2017 ; Deng et al. , 2019 ) , recommendation systems ( Ying et al. , 2018 ; Wang et al. , 2018 ) , and drug discovery ( Fout et al. , 2017 ; Do et al. , 2019 ) . Despite their recent success , effective training of GNNs on large-scale real-world graphs , such as Facebook social network ( Boldi & Vigna , 2004 ) , remains challenging . Although several attempts have been made to scale GNN training by sampling techniques ( Hamilton et al. , 2017 ; Zou et al. , 2019 ; Zeng et al. , 2020 ; Chiang et al. , 2019 ; Chen et al. , 2018 ; Zhang et al. , 2021 ) , they are still inefficient for training on extremely large graphs , due to the unique structure of GNNs and the limited memory capacity/bandwidth of current servers . One potential solution to tackle these limitations is employing distributed training with data parallelism , which have become almost a de facto standard for fast and accurate training for natural language processing ( Lin et al. , 2021 ; Hard et al. , 2018 ) and computer vision ( Bonawitz et al. , 2019 ; Konečnỳ et al. , 2018 ) . For ∗Equal Contribution . example , as shown in Figure 1 , moving from single machine to multiple machines reduces the training time and alleviates the memory burden on each machine . Besides , scaling the training of GNNs with sampling techniques can result in privacy concerns : existing sampling-based methods require centralized data storage and model learning , which could result in privacy concerns in real-world scenarios ( Shin et al. , 2018 ; Wu et al. , 2021 ) . Fortunately , the privacy in distributed learning can be preserved by avoiding mutual access to data between different local machines , and using only a trusted third party server to access the entire data . Nonetheless , generalizing the existing data parallelism techniques of classical distributed training to the graph domain is non-trivial , which is mainly due to the dependency between nodes in a graph . For example , unlike solving image classification problems where images are mutually independent , such that we can divide the image dataset into several partitions without worrying about the dependency between images ; GNNs are heavily relying on the information inherent to a node and its neighboring nodes . As a result , partitioning graph leads to subgraphs with edges spanning subgraphs ( cut-edges ) , which will cause information loss and hinder the performance of the model ( Angerd et al. , 2020 ) . To cope with this problem , ( Md et al. , 2021 ; Jiang & Rumi , 2021 ; Angerd et al. , 2020 ) propose to transfer node features and ( Zheng et al. , 2020 ; Tripathy et al. , 2020 ; Scardapane et al. , 2020 ) propose to transfer both the node feature and its hidden embeddings between local machines , both of which can cause significant storage/communication overhead and privacy concerns ( Shin et al. , 2018 ; Wu et al. , 2021 ) . To better understand the challenge of distributed GNN training , we compare the validation F1score in Figure 2 ( a ) and the average data communicated per round in Figure 2 ( b ) for two different distributed GNN training methods on the Reddit dataset . On the one hand , we can observe that when ignoring the cut-edges , Parallel SGD with Periodic Averaging ( PSGD-PA ( Dean et al. , 2012 ; Li et al. , 2020b ) ) suffers from significant accuracy drop and can not achieve the same accuracy as the single machine training , even by increasing the number of communication . However , Global Graph Sampling ( GGS ) can successfully reach the baseline by considering the cut-edges and allowing feature transfer , at the cost of significant communication overhead , and potential violation of privacy . In this paper , we propose a communication-efficient distributed GNN training method , called Learn Locally , Correct Globally ( LLCG ) . To reduce the communication overhead , inspired by the recent success of the distributed optimization with periodic averaging ( Stich , 2019 ; Yu et al. , 2019 ) , we propose Local Training with Periodic Averaging : where each local machine first locally trains a GNN model by ignoring the cut-edges , then sends the trained model to the server for periodic model averaging , and receive the averaged model from server to continue the training . By doing so we eliminate the features exchange phase between server and local machines , but it can result in a significant performance degradation due to the lack of the global graph structure and the dependency between nodes among different machines . To compensate for this error , we propose a Global Server Correction scheme to take advantage of the available global graph structure on the server and refine the averaged locally learned models before sending it back to each local machine . Notice that without Global Server Correction , LLCG is similar to PSGD-PA as introduced in Figure 2 . To get a deeper understanding on the necessity of Global Server Correction , we provide the first theoretical analysis on the convergence of distributed training for GNNs with periodic averaging . In particular , we show that solely averaging the local machine models and ignoring the global graph structure will suffer from an irreducible residual error , which provides sufficient explanation on why Parallel SGD with Periodic Averaging can never achieve the same performance as the model trained on a single machine in Figure 2 ( a ) . Then , we theoretically analyze the convergence of our proposal LLCG . We show that by carefully choosing the number of global correction steps , LLCG can overcome the aforementioned residual error and enjoys O ( 1/ √ PT ) convergence rate with P local machines and T iterations of gradient updates , which matches the rate of ( Yu et al. , 2019 ) on a general ( not specific for GNN training ) non-convex optimization setting . Finally , we conduct comprehensive evaluations on real-world graph datasets with ablation study to validate the effectiveness of LLCG and its improvements over the existing distributed methods . Related works . Recently , several attempts have been made on distributed GNN training . According to how they deal with the input/hidden feature of nodes that are associated with the cut-edges ( i.e. , the edges spanning subgraphs of each local machine ) , existing methods can be classified into two main categories : ( 1 ) Input feature only communication-based methods : In these methods , each local machine receives the input features of all nodes required for the gradient computation from other machines , and trains individually . However , since the number of required nodes grows exponentially with the number of layers , these methods suffer from a significant communication and storage overhead . To alleviate these issues , ( Md et al. , 2021 ) proposes to split the original graph using a min-cut graph partition algorithm that can minimize the number of cut-edges . ( Jiang & Rumi , 2021 ) proposes to use importance sampling to assign nodes on the local machine with a higher probability . ( Angerd et al. , 2020 ) proposes to sample and save a small subgraph from other local machines as an approximation of the original graph structure . Nonetheless , these methods are limited to a very shallow GNN structure and suffer from significant performance degradation when the original graph is dense . ( 2 ) Input and hidden feature communication-based methods : These methods propose to communicate hidden features in addition to the input node features . Although these methods reduce the number of transferred bytes during each communication round ( due to the smaller size of hidden embedding and less required nodes features ) , the number of communication rounds grows linearly as the number of layers , and are prone to more communication delay . To address these issues , in addition to optimal partitioning of the graph , ( Zheng et al. , 2020 ) proposes to use sparse embedding to reduce the number of bytes to communicate and ( Tripathy et al. , 2020 ) proposes several graph partitioning techniques to diminish the communication overhead . 2 BACKGROUND AND PROBLEM FORMULATION . In this section , we start by describing Graph Convolutional Network ( GCN ) and its training algorithm on a single machine , then formulate the problem of distributed GCN training . Note that we use GCN with mean aggregation for simplicity , however , our discussion is also applicable to other GNN architectures , such as SAGE ( Hamilton et al. , 2017 ) , GAT ( Velickovic et al. , 2018 ) , ResGCN ( Li et al. , 2019 ) and APPNP ( Klicpera et al. , 2019 ) . Training GCN on a single machine . Here , we consider the semi-supervised node classification in an undirected graph G ( V , E ) with N = |V| nodes and |E| edges . Each node vi ∈ V is associated with a pair ( xi , yi ) , where xi ∈ Rd is the input feature vector , yi ∈ R|C| is the ground truth label , and C is the candidate labels in the multi-class classifications . Besides , let X = [ x1 , . . . , xN ] ∈ RN×d denote the input node feature matrix . Our goal is to find a set of parameters θ = { W ( ` ) } L ` =1 by minimizing the empirical loss L ( θ ) over all nodes in the training set , i.e. , L ( θ ) = 1 N ∑ i∈V φ ( h ( L ) i , yi ) , h ( ` ) i = σ ( 1 |N ( vi ) | ∑ j∈N ( vi ) h ( ` −1 ) j W ( ` ) ) , ( 1 ) where φ ( · , · ) is the loss function ( e.g. , cross entropy loss ) , σ ( · ) is the activation function ( e.g. , ReLU ) , and N ( vi ) is the neighborhood of node vi . In practice , we can update the model parameters by the stochastic gradient computed on a sampled mini-batch ( using full-neighbors ) by ∇̃L ( θ , ξ ) = 1 B ∑ i∈ξ ∇φ ( h ( L ) i , yi ) , ( 2 ) where ξ denotes an i.i.d . sampled mini-batch of size B and we have E [ ∇̃L ( θ , ξ ) ] = ∇L ( θ ) . Distributed GCN training with periodic averaging . In this paper , we consider the distributed learning setting with P local machines and a single parameter server . The original input graph G is partitioned into P subgraphs , where Gp ( Vp , Ep ) denotes the subgraph on the p-th local machine with Np = |Vp| nodes , and Xp ∈ RNp×d as the input feature of all nodes in Vp located on the p-th machine . Then , the full-batch local gradient∇Llocalp ( θp ) is computed as ∇Llocalp ( θp ) = 1 Np ∑ i∈Vp ∇φ ( h ( L ) i , yi ) , h ( ` ) i = σ ( 1 |Np ( vi ) | ∑ j∈Np ( vi ) h ( ` −1 ) j W ( ` ) p ) , ( 3 ) where θp = { W ( ` ) p } L ` =1 is the model parameters on the p-th local machine , Np ( vi ) = { vj | ( vi , vj ) ∈ Ep } is the local neighbors of node vi on the p-th local machine . When the graph is large , the Algorithm 1 Distributed GCN training with “ Parallel SGD with Periodic Averaging ” Input : Global parameters θ̄0 , local parameters θ0p = θ̄0 , time-step t = 0 , learning rate η . 1 : for r← 1 to R do 2 : for p← 1 to P do in parallel . Parallel training on local machines 3 : Local machine p receives the global parameters θtp ← θ̄t. . Communication 4 : for k← 1 to K do 5 : t← t+ 1 . 6 : Local machine p constructs the mini-batch ξtp with neighbor sampling . 7 : Local machine p computes the stochastic gradients ∇̃Llocalp ( θtp , ξtp ) . 8 : Local machine p updates the local parameter by θt+1p = θtp− η∇̃Llocalp ( θtp , ξtp ) . 9 : end for 10 : Local machine p sends the local parameters θt+1p to the server. . Communication 11 : end for 12 : Server updates the global parameters by parameter averaging θ̄t+1 = 1 P ∑P p=1 θ t+1 p . 13 : end for Output : Server returns trained GCN model with mint E [ ‖∇L ( θ̄t ) ‖2 ] . computational complexity of forward and backward propagation could be very high . One practical solution is to compute the stochastic gradient on a sampled mini-batch with neighbor sampling , i.e. , ∇̃Llocalp ( θp , ξp ) = 1 Bp ∑ i∈ξp ∇φ ( h̃ ( L ) i , yi ) , h̃ ( ` ) i = σ ( 1 |Ñp ( vi ) | ∑ j∈Ñp ( vi ) h̃ ( ` −1 ) j W ( ` ) p ) , ( 4 ) where ξp is an i.i.d . sampled mini-batch of Bp nodes , Ñp ( vi ) ⊂ N ( vi ) is the sampled neighbors . An illustration of distributed GCN training with Parallel SGD with Periodic Averaging ( PSGD-PA ) is summarized in Algorithm 1 . Before training , the server maintains a global model θ̄0 and each local machine keeps a local copy of the same model θ0p . During training , the local machine first updates the local model θtp using the stochastic gradient ∇̃Llocalp ( θtp , ξtp ) computed by Eq . 4 for K iterations ( line 8 ) , then sends the local model θtp to the server ( line 10 ) . At each communication step , the server collects and averages the model parameters from the local machines ( line 12 ) and send the averaged model θt+1p back to each local machine . Limitations . Although PSGD-PA can significantly reduce the communication overhead by transferring the locally trained models instead of node feature/embeddings ( refer to Figure 2 ( b ) ) , it suffers from performance degeneration due to ignorance of the cut-edges ( refer to Figure 2 ( a ) ) . In the next section , we introduce a communication-efficient algorithm LLCG that does not suffer from this issue , and can achieve almost the same performance as training the model on a single machine .

Summary: The paper considers the problem of distributed training for graph learning tasks, under a setting where data privacy is significant for each individual machine and communication to/from a central parameter server is expensive. To preserve privacy each machine has only access to a distinct partition of the overall graph. The central server has access to the full graph. In the LLCG algorithm that the paper proposes, each machine trains on its local graph partition for some time before sending the parameters to the server. The server averages the received parameter, but additionally also does its own training using the full graph available to it. Theoretically the authors show that the proposed method avoids an error gap in the gradient norm that would exist if server correction is not performed. Experimental results show the proposed scheme performing similarly to GGS albeit with much lower communication costs.