Redundancy-Free Computation Graphs for Graph Neural Networks

1 INTRODUCTION . Graph neural networks ( GNNs ) have shown state-of-the-art performance across a number of tasks with graph-structured data , such as social networks , molecule networks , and webpage graphs ( Kipf and Welling , 2016 ; Hamilton et al. , 2017 ; Ying et al. , 2018 ; Xu et al. , 2019 ) . GNNs use a recursive neighborhood aggregation scheme — in a GNN layer , each node aggregates its neighbors ’ activations from the previous GNN layer and uses the aggregated value to update its own activations . The activations of the final GNN layer are used for downstream prediction tasks , such as node classification , graph classification , or link prediction . Due to the clustering nature of real-world graphs , different nodes in a graph may share a number of common neighbors . For example , in webpage graphs , different websites under the same domain generally have a number of common links ( i.e. , neighbors ) . As another example , in recommender systems , users in the same group may have interests in common items . However , existing GNN representations do not capture these common neighbors in real-world graphs , leading to redundant and unnecessary computation in both GNN training and inference . In particular , most existing GNN models ( Kipf and Welling , 2016 ; Xu et al. , 2019 ; Wu et al. , 2019 ) use the full-batch training method that computes the activations for all nodes in each layer . Existing GNN representations use a computation graph ( referred to as GNN-graph ) to define computation in a GNN layer . The GNN-graph includes a tree structure for each node v in the input graph describing how to compute v ’ s activations by aggregating the previous-layer activations of v ’ s neighbors . Figure 1b shows the GNN-graph of the input graph in Figure 1a ; for example , for node A , its neighbor ’ s activations h ( k−1 ) B , h ( k−1 ) C and h ( k−1 ) D from the layer k − 1 are aggregated to compute new activations h ( k ) A for the layer k ( see the top portion of Figure 1b ) . The new activations of the other nodes are computed similarly using the previous activations of their neighbors . Notice that this representation results in redundant computation and data transfers . In this small example , both { A , B } and { C , D } are aggregated twice . In wider and multi-layer GNNs , the redundancies in existing GNN representations account for a significant fraction of all computation . For example , our experiments show that in modern GNNs up to 84 % of the aggregations are redundant and avoidable . In this paper , we propose a new GNN representation called Hierarchically Aggregated computation Graphs ( HAGs ) . Figure 1c shows one possible HAG for the input graph in Figure 1a . HAGs are functionally equivalent to standard GNN-graphs ( produce the same output ) , but represent common neighbors across different nodes using aggregation hierarchies , which eliminates redundant computation and unnecessary data transfers in both GNN training and inference . In addition , a HAG is agnostic to any particular GNN model , and can be used to eliminate redundancy for arbitrary GNNs . For a GNN-graph , there exist numerous equivalent HAGs with different aggregation hierarchies and runtime performance . We introduce an accurate cost function to quantitatively estimate the performance of different HAGs and develop a novel HAG search algorithm to automatically find optimized HAGs . Theoretically , we prove that the search algorithm finds HAGs with strong performance guarantees : ( 1 ) for GNN models whose neighborhood aggregations require a specific ordering on a node ’ s neighbors , the algorithm finds a globally optimal HAG under the cost function ; and ( 2 ) for other GNN models , the algorithm finds HAGs whose runtime performance is at least a ( 1− 1/e ) approximation ( ≈ 63 % ) of globally optimal HAGs using the submodularity property ( Mossel and Roch , 2007 ) . Empirically , the algorithm finds HAGs that lead to major improvements , reducing the runtime by up to 2.9× . Our HAG abstraction maintains the predictive performance of GNNs but leads to much faster training and inference . We evaluate the performance of HAGs on five real-world datasets and along three dimensions : ( a ) end-to-end training and inference runtime ; ( b ) number of aggregations ; and ( c ) size of data transfers . Experiments show that HAGs increase the end-to-end training and inference performance by up to 2.8× and 2.9× , respectively . In addition , compared to GNN-graphs , HAGs reduce the number of aggregations and the size of data transfers by up to 6.3× and 5.6× , respectively . 2 RELATED WORK . Graph neural networks have been used to solve various real-world tasks with relational structures ( Kipf and Welling , 2016 ; Hamilton et al. , 2017 ; Ying et al. , 2018 ; Xu et al. , 2019 ) . FastGCN ( Chen et al. , 2018 ) and SGC ( Wu et al. , 2019 ) accelerate GNN training using importance sampling and removing nonlinearilities . This paper solves an orthogonal problem : how to optimize GNN efficiency while maintaining network accuracy . Our methods can be applied automatically to eliminate redundancy for arbitrary GNN models . Join-trees are a tree decomposition technique that maps a graph into a corresponding tree structure to solve optimization problems on the graph , such as query optimization ( Flum et al. , 2002 ) . Although a join-tree provides a possible way to find optimal HAGs for a GNN-graph , its time complexity is exponential in the treewidth of a graph ( Arnborg et al. , 1987 ) , and real graphs tend to have very large treewidths . For example , Adcock et al . ( 2016 ) shows that the treewidth of real-world social networks grows linearly with the network size , making it infeasible to use join-trees to find optimal HAGs . Computation reduction in DNNs . Several techniques have been proposed to reduce computation in DNNs , including pruning weights ( Han et al. , 2015 ) and quantization ( Han et al. , 2016 ) . These techniques reduce computation at the cost of modifying networks , resulting in decreased accuracy ( as reported in these papers ) . By contrast , we propose a new GNN representation that accelerates GNN training by eliminating redundancy in GNN-graphs while maintaining the original network accuracy . 3 HIERARCHICALLY AGGREGATED COMPUTATION GRAPHS ( HAGS ) . Algorithm 1 An abstraction for GNNs . V is the set of nodes in an input graph , and N ( v ) denotes the set of neighbors for node v. 1 : h ( 0 ) v = xv , ∀v ∈ V 2 : for k = 1 to K do 3 : for v ∈ V do 4 : a ( k ) v ← AGGREGATE ( { h ( k−1 ) u |u ∈ N ( v ) } ) 5 : h ( k ) v ← UPDATE ( a ( k ) v , h ( k−1 ) v ) 6 : Goal : minimize L ( { h ( K ) v |v ∈ V } ) Existing GNN-graph representation . An input graph G = ( V , E ) has nodes V and edges E . For each node v ∈ V , N ( v ) denotes the set of neighbors of v , and xv denotes the input node features . A GNN iteratively learns representations for individual nodes over the entire graph through a number of GNN layers , as shown in Algorithm 1 . The learned activations of node v at layer k is h ( k ) v , and we initialize h ( 0 ) v with xv . At the k-th layer , a ( k ) v denotes the aggregated activations of v ’ s neighbors , which is combined with h ( k−1 ) v to compute an updated activation h ( k ) v . The learned node activations of the final layer ( i.e. , h ( K ) v ) are used for downstream learning tasks , and a GNN model generally minimizes a loss function L that takes the final node activations as inputs ( line 6 ) . Existing GNN models use a GNN computation graph ( GNN-graph ) to describe the computation in each GNN layer , as shown in Figure 1b . For each node v in the input graph , the GNN-graph includes an individual tree structure to define how to compute the activations h ( k ) v of node v by aggregating the previous-layer activations of v ’ s neighbors ( i.e. , { h ( k−1 ) u , u ∈ N ( v ) } ) . GNN-graphs are efficient at expressing direct neighborhood relations between nodes , but are not capable of capturing common neighbors across multiple nodes , leading to redundant computation in GNN training and inference . 3.1 HAG REPRESENTATION . We propose a new graph representation called Hierarchically Aggregated computation Graphs ( HAGs ) for GNNs . HAGs eliminate redundancy in the GNN-graph representation by hierarchically managing and reusing intermediate aggregation results . A HAG Ĝ = ( V̂ , Ê ) has nodes V̂ = V ∪ VA and edges Ê , where V is the set of nodes in the original graph , and VA is a new set of aggregation nodes . Each aggregation node in VA represents the intermediate aggregations result for a subset of nodes ( i.e. , aggregation on a subset of h ( k−1 ) v ) . For the HAG example in Figure 1c , the new nodes AB and CD denote the aggregation results of { A , B } and { C , D } , respectively . A HAG can contain a multi-level aggregation hierarchy . For example , Figure 1c can also have a third aggregation node ABCD that depends on AB and CD . Similar to edges in GNN-graphs , an edge ( u , v ) in a HAG denotes an aggregation relation — computing v ’ s activations requires aggregating u ’ s activations . The standard GNN-graph representation can be considered as a special case in the HAG representation with no intermediate aggregation nodes ( i.e. , VA = ∅ ) . Our HAG abstraction is general and applicable to many existing GNN models . Table 1 shows how to use our abstraction to define existing GNNs , which can be further divided into two categories . • Set AGGREGATE . Most GNNs assume the neighbors of a node have no ordering , and the aggregations are associative and commutative operations that are invariant to the order in which the aggregations are performed . Examples include GCN with summation aggregations and GraphSAGE-P with element-wise pooling aggregations ( Table 1 ) . Note that set aggregations in GNNs are designed to be order invariant and thus can be performed in a hierarchical fashion as we do in HAGs . • Sequential AGGREGATE . Another class of GNNs require a specific ordering of a node ’ s neighbors and the aggregations are not commutative . Examples include N -ary Tree-LSTM ( Tai et al. , 2015 ) and the LSTM variant of GraphSAGE ( Hamilton et al. , 2017 ) . HAGs can be applied in the case of sequential aggregations as well . Rather than identifying common subsets of neighbors , we identify the common prefixes of the sequence of aggregated nodes , which can then be reused among nodes . We further define two properties for the aggregation nodes VA. First , for each v ∈ VA , âv denotes its intermediate aggregation result , and N̂ ( v ) denotes the in-neighbors of node v. To capture the aggregation hierarchy in a HAG , we use a recursive function to define âv . âv = AGGREGATE ( { h ( k−1 ) u u ∈ V âu u ∈ VA ∣∣∣∣∣u ∈ N̂ ( v ) ) ( 1 ) Second , Ĉ ( v ) denotes the set of input activations h ( k−1 ) u used to compute âv in the recursive procedure . âv = AGGREGATE ( { h ( k−1 ) u |u ∈ Ĉ ( v ) } ) ( 2 ) Intuitively , Ĉ ( v ) defines the coverage of node v in a HAG . For the HAG example in Figure 1c , Ĉ ( A ) = { B , C , D } because h ( k−1 ) B , h ( k−1 ) C , and h ( k−1 ) D are used as inputs to compute h ( k ) A . For a set AGGREGATE , Ĉ ( · ) is an unordered set : Ĉ ( v ) = ⋃ u∈N̂ ( v ) { u u ∈ V Ĉ ( u ) u ∈ VA ( 3 ) For a sequential AGGREGATE , Ĉ ( · ) is an ordered list : Ĉ ( v ) = ( Ĉ ( u1 ) , ... , Ĉ ( um ) ) ( 4 ) where u1 , ... , um are the ordered in-neighbors of v. 3.2 GNNS WITH HAGS Algorithm 2 A GNN abstraction with HAGs . We exclude layer index superscripts in âv to denote that âv does not need to be memoized for back propagation , and its memory can be reused across all layers . 1 : h ( 0 ) v = xv , ∀v ∈ V 2 : for k = 1 to K do 3 : for v ∈ VA do 4 : compute âv using Equation ( 1 ) 5 : for v ∈ V do 6 : a ( k ) v ← AGGREGATE ( { âu|u ∈ N̂v } ) 7 : h ( k ) v ← UPDATE ( a ( k ) v , h ( k−1 ) v ) We extend the GNN abstraction to make it also applicable to HAGs , as shown in Algorithm 2 . The extension does not require any modification to a GNN model , and the only difference is how to compute neighborhood aggregations ( i.e. , a ( k ) v ) in each GNN layer . Before the original aggregations , we add one step that precomputes commonly used intermediate aggregation results to reduce redundant computation . In Algorithm 2 , we first compute âv for all aggregation nodes VA ( line 3-4 ) . This is performed recursively following the aggregation hierarchy of the HAG . We then compute the neighborhood aggregations a ( k ) v ( line 5-6 ) using the precomputed intermediate aggregations âv . Memory overhead . Although Algorithm 2 includes new intermediate variables âv , the memory overhead for storing âv is negligible since âv is not used for back propagation and can be saved in a constant memory across all GNN layers . In the experiments , we show that HAGs can increase the training throughput by 2.8× , while maintaining the original model accuracy at the cost of 0.1 % memory overhead to save âv . Meanwhile , storing a HAG representation requires less memory than the original GNN-graph , as HAGs generally contain much fewer edges . Equivalence between GNN-graphs and HAGs . We define a GNN-graph and a HAG to be equivalent if they produce the exact same outputs and gradients for a GNN model . Formally , a GNN-graph G and a HAG Ĝ are equivalent for a GNN model if ( 1 ) the GNN model outputs the same activations ( i.e. , h ( k ) v ) at each GNN layer , and ( 2 ) the GNN model computes the same gradients for all trainable parameters in back propagation . Equivalent graphs guarantee the same predictive performance , and therefore can be used interchangeably for both GNN training and inference . Theorem 1 provides a simple yet efficient condition to check graph equivalence . We prove the theorem in the Appendix . Theorem 1 . A GNN-graph with nodes V and a HAG with nodes ( V , VA ) are equivalent if and only if N ( v ) = Ĉ ( v ) for all v ∈ V , where N ( v ) is v ’ s neighbors in the input graph and Ĉ ( · ) is defined in Equation 3 and 4 .

Summary: In this paper, authors propose a way to speed up the computation of GNN. More specifically, the hierarchically aggregate computation graphs are proposed to aggregate the intermediate node and utilize this to speed up a GNN computation. Authors proof that the computation based on HAGs are equivalent to the vanilla computation of GNN (Theorem1). Moreover, for sequential aggregation, it can find a HAG that is at least (1-1/e)-approximation of the globally optimal HAGs (Theorem 3). These theoretical results are nice. Through experiments, the authors demonstrate that the proposed method can get faster computation than vanilla algorithms.