R EDUNDANCY -F REE C OMPUTATION G RAPHS FOR G RAPH N EURAL N ETWORKS
Anonymous authors
Paper under double-blind review
A BSTRACT
Graph Neural Networks (GNNs) are based on repeated aggregations of information across nodes’ neighbors in a graph.
However, because common neighbors are shared between different nodes, this leads to repeated and inefﬁcient computations.
We propose Hierarchically Aggregated computation Graphs (HAGs), a new GNN graph representation that explicitly avoids redundancy by managing intermediate aggregation results hierarchically, and eliminating repeated computations and unnecessary data transfers in GNN training and inference.
We introduce an accurate cost function to quantitatively evaluate the runtime performance of different HAGs and use a novel search algorithm to ﬁnd optimized HAGs.
Experiments show that the HAG representation signiﬁcantly outperforms the standard GNN graph representation by increasing the end-to-end training throughput by up to 2.8 × and reducing the aggregations and data transfers in GNN training by up to 6.3 × and 5.6 × .
Meanwhile, HAGs improve runtime performance by preserving GNN computation, and maintain the original model accuracy for arbitrary GNNs.
I NTRODUCTION
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 ﬁnal GNN layer are used for downstream prediction tasks, such as node classiﬁcation, graph classiﬁcation, 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 deﬁne 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 signiﬁcant 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
A
B
C
D
C D C D
B
A
A B A B
A
B
C
D
A B C
D
A C
D
B
C
D A B
A B ! "
#$% ! &#$% ! '
#$% !
[Figure]
Because nodes in the input graph share common neighbors, the GNN-graph performs redundant computation (e.g., both { A, B } and { C, D } are aggregated twice)
) By identifying common computational patterns, the HAG avoids repeated computation.
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 ﬁnd optimized HAGs.
Theoretically, we prove that the search algorithm ﬁnds HAGs with strong performance guarantees: (1) for GNN models whose neighborhood aggregations require a speciﬁc ordering on a node’s neighbors, the algorithm ﬁnds a globally optimal HAG under the cost function; and (2) for other GNN models, the algorithm ﬁnds 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 ﬁnds 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 ﬁve 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.
R ELATED W ORK
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 efﬁciency 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 ﬁnd 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 ﬁnd optimal HAGs.
[Table]
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.
H IERARCHICALLY A GGREGATED C OMPUTATION G RAPHS (HAG S )
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 = x v , ∀ v ∈ V 2: for k = 1 to K do 3: for v ∈ V do 4: a ( k ) v ← A GGREGATE ( { h ( k − 1) u | u ∈ N ( v ) } ) 5: h ( k ) v ← U PDATE ( 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 x v 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 x v .
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 ﬁnal 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 ﬁnal 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 deﬁne 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 efﬁcient 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. HAG R EPRESENTATION
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 (cid:98) G =
( (cid:98) V , (cid:98) E ) has nodes (cid:98) V
= V ∪ V A and edges (cid:98) E , where V is the set of nodes in the original graph, and V A is a new set of aggregation nodes.
Each aggregation node in V A 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., V A = ∅ ).
Our HAG abstraction is general and applicable to many existing GNN models.
Table 1 shows how to use our abstraction to deﬁne existing GNNs, which can be further divided into two categories.
• Set A GGREGATE .
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 A GGREGATE .
Another class of GNNs require a speciﬁc 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 preﬁxes of the sequence of aggregated nodes, which can then be reused among nodes.
We further deﬁne two properties for the aggregation nodes V A .
First, for each v ∈ V A , (cid:98) a v denotes its intermediate aggregation result, and (cid:98) N ( v ) denotes the in-neighbors of node v .
To capture the aggregation hierarchy in a HAG, we use a recursive function to deﬁne (cid:98) a v .
[Equation]
Second, (cid:98) C ( v ) denotes the set of input activations h ( k − 1) u used to compute (cid:98) a v in the recursive procedure.
[Equation]
Intuitively, (cid:98) C ( v ) deﬁnes the coverage of node v in a HAG.
For the HAG example in Figure 1c, (cid:98) 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 A GGREGATE , (cid:98) C ( · ) is an unordered set:
[Equation]
For a sequential A GGREGATE , (cid:98) C ( · ) is an ordered list:
[Equation]
where u 1 , ..., u m are the ordered in-neighbors of v .
3. GNN S WITH HAG S
Algorithm 2 A GNN abstraction with HAGs.
We exclude layer index superscripts in (cid:98) a v to denote that (cid:98) a v does not need to be memoized for back propagation, and its memory can be reused across all layers.
[Equation] 2: for [Equation] do 3: for
[Equation] do 4: compute (cid:98) a v using Equation (1) 5: for v ∈ V do 6: [Equation] A
[Equation] [Equation] u [Equation]
u ∈ (cid:98) N v
[Equation]
7: h
[Equation]
k
[Equation]
← U PDATE ( a ( k ) v
[Equation]
h ( k −
[Equation] ) v ) We extend the GNN abstraction to make it also applicable to HAGs, as shown in Algorithm 2. The extension does not require any modiﬁcation 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 ﬁrst compute (cid:98) a v for all aggregation nodes V A (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 (cid:98) a v .
Memory overhead.
Although Algorithm 2 includes new intermediate variables (cid:98) a v , the memory overhead for storing (cid:98) a v is negligible since (cid:98) 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 (cid:98) 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 deﬁne 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 (cid:98) 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 efﬁcient 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 , V A ) are equivalent if and only if N ( v )
= (cid:98) C ( v ) for all v ∈ V , where N ( v ) is v ’s neighbors in the input graph and (cid:98) C ( · ) is deﬁned in
Equation 3 and 4.
HAG S EARCH A LGORITHM
For a GNN model and an input GNN-graph, there exists a large space of equivalent HAGs with the same model accuracy but various runtime performance.
Our goal is to explore the search space to discover a HAG with optimized runtime performance.
First, we deﬁne a realistic cost function to quantitatively evaluate the runtime performance of different HAGs.
Second, we introduce an efﬁcient search algorithm that ﬁnds an optimized HAG with the following guarantees:
• For GNNs with sequential A GGREGATE , the HAG search algorithm can ﬁnd globally optimal HAGs under the cost function.
• For GNNs with set A GGREGATE , ﬁnding an optimal HAG is NP-hard by a reduction from the NP-hard maximum coverage problem (see Appendix for the proof).
The search algorithm ﬁnds at least a (1 − 1 /e ) -approximation of globally optimal HAGs based on the submodularity property.
4. C OST F UNCTION
We introduce a cost function that estimates the runtime performance of a HAG by measuring the computation cost to perform one epoch of GNN training on the HAG.
The computation cost of a GNN model includes aggregating the neighbors of each node by calling
A GGREGATE and updating the activations of each node via U PDATE , as shown in Algorithm 2.
For a
GNN model M , we assume the cost of performing A GGREGATE on two elements is α M , and the cost of computing an U PDATE is β M .
In Algorithm 2, computing (cid:98) a v with | (cid:98) N v | neighbors requires performing ( | (cid:98) N v | − 1) binary aggregations, whose cost is α M × ( | (cid:98) N v | − 1) .
Therefore, the total computation cost of training a GNN model M on a HAG (cid:98) G is
[Equation]
Since |V| is determined by the input graph, our goal is to minimize (cid:0) | (cid:98) E| − | (cid:98)
V A | (cid:1) .
4. S EARCH A LGORITHM
We present a HAG search algorithm that ﬁnds a globally optimal HAG for GNNs with sequential A GGREGATE and a (1 − 1 /e ) -approximation of globally optimal HAGs for GNNs with set A GGREGATE .
In addition to an input GNN-graph and a GNN model, the algorithm also takes a hyper-parameter capacity , deﬁning an upper limit on the number of intermediate aggregation nodes (i.e., |V A | ).
Algorithm 3 shows the pseudocode of the HAG search algorithm.
We start with an input GNN-graph, and iteratively insert aggregation nodes into the current HAG to merge highly redundant aggregations and remove unnecessary computation and data transfers.
In each iteration, we ﬁnd a binary aggregation with the highest redundancy and insert a new aggregation node w in V A to represent the binary aggregation results (line 12-15).
All nodes containing this binary aggregation can directly use the output of w without recomputing the aggregation (line 16-18).
The search algorithm iteratively reduces the computation cost of the HAG by eliminating the most redundant aggregation in each iteration.
The redundancy scores are maintained in a heap structure.
For a GNN model with a sequential A GGREGATE , Theorem 2 shows that our search algorithm ﬁnds an equivalent HAG with globally optimal computation cost.
We prove the theorem in Appendix.
Theorem 2.
For any GNN-graph G =
( V , E ) and any GNN model M with a sequential A GGREGATE ,
Algorithm 3 returns an equivalent HAG with globally minimum cost as long as capacity ≥ |E| .
For a GNN model with a set A GGREGATE , Theorem 3 shows that our search algorithm ﬁnds a HAG that is at least a (1 − 1 /e ) -approximation of the globally optimal HAGs (see proof in Appendix).
Theorem 3.
For any GNN-graph G and GNN model M with a set A GGREGATE , Algorithm 3 gives a (1 − 1 /e ) -approximation of globally optimal HAGs under the cost function.
Formally, let (cid:98) G be the
HAG returned by Algorithm 3, and
(cid:98) G o is a globally optimal HAG under the capacity constraint,
[Equation]
Time and space complexity.
Our HAG algorithm achieves low theoretical complexity and has negligible runtime overhead.
In particular, the overall time complexity of Algorithm 3 is O ( capacity × |V| + |E| × log |V| ) , and the space complexity is O ( capaccity × |V| + |E| )
(see Appendix for the proof).
One key optimization is a heap data structure for maintaining the redundancy scores of the highest O ( | V | ) node pairs.
Finding the most redundant node pair over the entire graph thus only takes
O (1) time, and updating the redundancy scores (i.e., line 15 and 18) each takes O (log | V | ) time.
E XPERIMENTS
The HAG representation maintains the predictive performance of GNNs but has much better runtime performance.
This section evaluates the runtime performance of HAGs on ﬁve real-world graph datasets.
We evaluate HAGs along three dimensions: (a
5. I MPLEMENTATION
Existing deep learning frameworks such as TensorFlow and PyTorch are designed for spatial data structures (e.g., images and text), and have limited support for irregular data structures such as graphs.
As a result, GNN models in existing frameworks translate graph structures to sparse adjacent matrices and use matrix operations to perform GNN training.
We implemented the following operations in TensorFlow r1.14 to support GNN training with HAGs.
First, graph to hag automatically transforms an input GNN-graph to an equivalent HAG with optimized performance.
Second, hag aggregate takes a HAG and nodes’ activations as inputs, and computes the aggregated activations of all nodes.
Finally, hag aggregate grad computes the gradients of hag aggregate for back propagation.
Our implementation minimizes changes to existing GNN programs: a GNN application can directly use all HAG optimizations by only modifying a few lines of code.
5. E XPERIMENTAL S ETUP
[Table]
COLLAB
[Table]
Datasets.
Table 2 summarizes the public datasets used in our experiments.
BZR is a chemical compound dataset, where each node is an atom and an edge is a chemical bond between two atoms (Kriege and Mutzel, 2012).
PPI contains a number of protein-protein interaction graphs, each of which corresponds to a different human tissue (Zitnik and Leskovec, 2017).
REDDIT is an online discussion forum dataset, with each node being a Reddit post and each edge being commenting relations.
For both
PPI and REDDIT, we directly use prepossessed data from Hamilton et al.
IMDB and COLLAB are two collaboration datasets for graph classiﬁcation (Yanardag and Vishwanathan, 2015).
IMDB is a movie collaboration dataset, with each node representing an actor/actress, while COLLAB is a scientiﬁc collaboration dataset, with each node representing a researcher.
All experiments were performed running TensorFlow r1.14 on NVIDIA Tesla V100 GPUs.
Following previous work (Kipf and Welling, 2016; Hamilton et al., 2017), each GNN model has two GNN layers and one SoftMax layer.
For graph classiﬁcation datasets, each GNN model also includes a mean-pooling layer to gather graph-level activations.
For all experiments, we set the maximum capacity of |V A | in a HAG to be |V| / 4 , which achieves high performance on real-world graphs.
5. E ND - TO -E ND P ERFORMANCE
[Figure]
We ﬁrst measure the per-epoch training time and inference latency to run a 2-layer GCN model on different graph datasets.
We follow previous work (Hamilton et al., 2017; Kriege and Mutzel, 2012; Yanardag and Vishwanathan, 2015) to split the datasets into training/validation/testing sets, and use the testing sets to measure the average inference latency.
Figure 2 compares the per-epoch training time and inference latency between GNN-graphs and
HAGs.
Compared to GNN-graphs, HAGs can improve the training and inference performance by up to 2.8 × and 2.9 × , respectively, while maintaining the same network accuracy.
We note this improvement is achieved completely automatically, and computing a HAG is inexpensive.
Thus, because the improvement is essentially for free, we believe there is no reason not to use HAGs in preference to GNN-graphs.
R e l a t i v e N u m b e r o f
A gg r e g a t i o n s 16% 65% 55% 68% 30% 41%
BZR
PPI REDDIT IMDB COLLAB Mean 0.
R e l a t i v e
D a t a
T r a n s f e r s 18% 71% 80% 29% 45% 43%
GNN-graph
HAG (a) Set Aggregations.
R e l a t i v e N u m b e r o f
A gg r e g a t i o n s 57% 97% 82% 63% 54% 69%
BZR
PPI REDDIT IMDB COLLAB Mean 0.
R e l a t i v e
D a t a
T r a n s f e r s 59% 97% 83% 67% 55% 71%
GNN-graph
HAG (b) Sequential Aggregations.
5. A GGREGATION P ERFORMANCE
We further compare the aggregation performance of GNN-graphs and HAGs on the following two metrics: (1) the number of binary aggregations performed in each GNN layer; and (2) the size of data transfers between GPU threads to perform the aggregations.
Note that aggregating a neighbor’s activations requires transferring the activations from GPU global memory to a thread’s local memory.
Figure 3 shows the comparison results.
For GNNs with set aggregations, HAGs reduce the number of aggregations by 1.5-6.3 × and the size of data transfers by 1.3-5.6 × .
For GNNs with sequential aggregations, HAGs reduce aggregations and data transfers by up to 1.8 × and 1.9 × , respectively.
Although the search algorithm ﬁnds a globally optimal HAG for sequential aggregations (Theorem 2) and a (1 − 1 /e ) -approximation of globally optimal HAGs for set aggregations (Theorem 3), we observe the performance improvement is more signiﬁcant for set aggregations.
Optimality for HAGs with set aggregation involves more potential redundancy compared to sequential aggregations, due to permutation invariance of set aggregation.
Thus higher performance can be achieved with HAGs for set aggregations, though optimal solutions are more difﬁcult to compute.
It is also worth noting that the HAG search algorithm can ﬁnd highly optimized HAGs even on very sparse graphs.
For example, on the COLLAB dataset with a graph density of 0.01%, our algorithm reduces the number of aggregations and data transfers by 3.3 × and 2.2 × , respectively.
5. HAG S EARCH A LGORITHM
[Figure]
on the COLLAB dataset
HAGs with different capacities.
We train GCN for a maximum of 350 epochs by following prior work and
We evaluate the performance of the HAG search algorithm.
Recall that the search algorithm uses a hyper-parameter capacity to control the number of aggregation nodes in a
HAG.
A larger capacity allows the algorithm to eliminate more redundant aggregations and achieves lower cost.
Figure 4 shows the end-to-end GCN training time on the
COLLAB dataset using HAGs with different capacities.
A larger value of capacity can consistently improve the training performance, which indicates that the cost function is an appropriate metric to evaluate and compare the performance of different HAGs.
By gradually increasing the capacity, the search algorithm eventually ﬁnds a HAG with ∼ 100K aggregation nodes, which consume 6MB of memory (0.1% memory overhead) while improving the training performance by 2.8 × .
In addition, the HAG search time is negligible compared to the end-to-end training time. Limitation.
One limitation of the HAG graph representation is that it does not apply to attentionbased GNN models, such as GAT (Veliˇckovi´c et al., 2017), which uses the attention scores (Vaswani et al., 2017) between nodes to guide activation updates and results in limited redundancy that can be optimized by HAG.
Conclusion.
We introduce HAG, a new graph representation to eliminate redundancy in GNNs.
We propose a cost function to estimate the performance of different HAGs and use a search algorithm to ﬁnd optimized HAGs.
We show that HAGs outperform existing GNN-graphs by improving the end-to-end training performance and reducing the aggregations and data transfers in GNN training.
A P ROOF OF T HEOREM
Proof.
It is sufﬁcient to prove that if N ( v ) = (cid:98) C ( v ) for all v ∈ V , then the GNN-graph G and the
HAG (cid:98) G generate the same outputs (i.e., h ( k ) v ) for every GNN layer.
We prove this by induction.
Assume a GNN-graph G and a HAG (cid:98) G generate the same outputs for the ( k -1)-th layer, we prove the two graphs produce the same outputs for the k -th GNN layer.
In Algorithm 2, (cid:98) a v is the aggregation results of node v , which is deﬁned as
[Equation]
This proves that Algorithm 1 and Algorithm 2 compute the same a ( k ) v .
In addition, both algorithms use the same U PDATE function that takes a ( k ) v and h ( k − 1) v as inputs and computes h ( k ) v , which applies that the two algorithms compute the same h ( k ) v .
B P ROOF OF T HEOREM
Proof.
Sequential aggregations require a speciﬁc ordering of a node’s neighbors.
Let N v denote the ordered list of node v ’s neighbors and L ( i ) v denote a list of the ﬁrst i elements in N v :
[Equation]
N v (2) , ..., N v ( i ) (cid:1) where N v ( i ) is the i -th neighbor of node v .
L ( i ) v represents a necessary intermediate aggregation step for computing a ( k ) v (since sequential aggregations are not commutative), and therefore any HAG must compute L ( i ) v as an intermediate aggregation.
Counting the number of distinct L ( i ) v (where v ∈ V and 1 < i ≤ |N v | ) provides a lower bound on the number of aggregations any equivalent HAG must perform.
Assuming (cid:98) G o is a globally optimal HAG under the cost model, we have:
[Equation]
where lb is the number of distinct L ( i ) v that must be computed by any equivalent HAG.
Assuming (cid:98) G is the output HAG of Algorithm 3, we prove that cost ( M , (cid:98) G )
= cost ( M , (cid:98) G o ) by using contradiction.
In the case cost ( M , (cid:98) G ) > cost ( M , (cid:98) G o ) , (cid:98) G must perform more than lb aggregations.
Case 1.
One possible case is that (cid:98) G computes at least one aggregation that is not a preﬁx of any
N v , indicating that (cid:98) G performs some useless aggregations, which contradicts with the fact that all intermediate aggregations added to (cid:98) G must be used at least once.
Case 2.
The other possible case is that (cid:98) G computes the aggregation of some L ( i ) v multiple times.
However, in Algorithm 3, each iteration reduces the number of aggregations by at least 1, and there are |E| aggregations initially.
This implies there cannot be redundant aggregations after |E| iterations, which contradicts with the precondition of Case 2.
C P ROOF OF T HEOREM
Proof.
The idea of the proof is to build a monotone submodular function Cormen et al.
based on the cost model.
For any GNN-graph G and an equivalent (cid:98) G , we deﬁne
[Equation]
where V A is the set of aggregation nodes in (cid:98) G , and E and (cid:98) E are the set of edges in G and (cid:98) G , respectively.
f ( (cid:98) G ) measures the number of aggregations that can be saved by using (cid:98) G for GNN training.
We begin by deﬁning the subset relations between different HAGs.
For two HAGs (cid:98) G and (cid:98) G (cid:48) , we deﬁne (cid:98) G ⊆ (cid:98) G (cid:48)
iff V A is a subset of V (cid:48)
A , where V A and V (cid:48)
A are the aggregation nodes in (cid:98) G and (cid:98) G (cid:48) , respectively.
Prove that f ( (cid:98) G ) is monotone.
We show that for all (cid:98) G ⊆ (cid:98) G (cid:48) , f ( (cid:98) G ) ≤ f ( (cid:98) G (cid:48) ) .
This is true since (cid:98) G ⊆ (cid:98) G (cid:48) indicates that (cid:98) G (cid:48) contains all aggregation nodes in (cid:98) G , which applies that (cid:98) G (cid:48) can at least save the same number of aggregations as (cid:98) G .
Prove that f ( (cid:98) G ) is submodular.
We show that for all (cid:98) G ⊆ (cid:98) G (cid:48) and any aggregation node n , f ( (cid:98) G + { n } ) − f ( (cid:98) G ) ≥ f
( (cid:98) G (cid:48)
+ { n } ) − f ( (cid:98) G [Equation]
This inequality holds because f ( (cid:98) G + { n } ) − f ( (cid:98) G ) measures the number of aggregations we can further save by adding aggregation n to the existing HAG, which monotonically decreases as we add more aggregation nodes to the HAG.
Let (cid:98) G ( i ) denote the result HAG after the i -th iteration of Algorithm 3.
(cid:98) G ( i ) includes exactly i aggregation nodes.
Let (cid:98) G
o denote the optimal HAG under the cost model with k aggregation nodes.
We claim via induction that for 0 ≤ i ≤ k ,
[Equation]
The base case is trivially true.
In the i -th step, Algorithm 3 selects an aggregation node a i by maximizing the marginal gain f
( (cid:98) G ( i )
+ a i ) − f
( (cid:98) G
( i ) ) .
Observe that the remaining aggregation nodes includes (cid:98) G
o \ (cid:98) G ( i − 1) , a set of at most k elements.
The submodularity applies that
[Equation]
and this implies that the aggregation node a i has marginal value
[Equation]
Assuming that Inequality 7 holds for (cid:98) G ( i − 1) , we have
[Equation]
which proves Inequality 7.
Therefore, we have
[Equation]
By taking in the deﬁnition of f ( · ) , we have
[Equation]
T r a i n i n g S p ee d u p s 1.2x 1.9x 1.7x 1.2x 2.6x 1.7x BZR
PPI REDDIT IMDBCOLLABMean
I n f e r e n c e S p ee d u p s 1.2x 1.7x 1.8x 1.1x 2.5x 1.6x GNN-graph HAG (a) GIN.
T r a i n i n g
S p ee d u p s 1.4x 2.1x 1.8x 1.2x 3.1x 1.8x
BZR
PPI REDDIT IMDB COLLABMean
I n f e r e n c e S p ee d u p s 1.3x 2.3x 2.4x 1.2x 3.3x 2.0x GNN-graph
HAG
[Figure]
D C OMPLEXITY OF THE HAG S EARCH A LGORITHM
Theorem 4.
The overall time complexity of Algorithm 3 is O ( capacity × |V| + |E| × log |V| ) , where capacity is the upper bound on the number of aggregation nodes in a HAG.
Proof.
We use a heap to maintain the redundancy score of each potential node pair and only update the heap when we add and remove edges in (cid:98) E .
Since the depth of the heap is at most O (log |V| ) 1 , querying the most redundant binary aggregation and modifying (cid:98) E each takes O (log |V| ) time.
First, we calculate the number of queries and updates to the heap structure:
• The algorithm iteratively pull the most redundant binary aggregation from the heap and add it to V A .
Since the number of vertices in V A is smaller than capacity , the total number of queries is O ( capacity ) .
• The algorithm inserts two new edges into (cid:98) E in line 16 and removes one edge from (cid:98) E in line 19.
Since line 16 can be invoked at most O ( capacity ) times, the total number of invocations to line 19 is O ( |E| + 2 × capacity ) .
Therefore, the overall number of updates is O ( |E| + capacity ) .
Second, the enumeration over all vertices in V (line 17) involves time complexity of O ( capacity ×|V| ) .
Therefore, the overall time complexity of Algorithm 3 is
[Equation]
E P ERFORMANCE E VALUATION
We further evaluate the effectiveness of the HAG optimizations on two additional GNN models:
GIN Xu et al.
and SGC Wu et al.
The experimental setup is described in Section 5.2.
We measure the per-epoch training time and inference latency for GIN and SGC.
We follow previous work Xu et al.
; Wu et al.
(2019) to split the datasets into training/validation/testing sets, and use the testing sets to measure the average inference latency.
Figure 5 compares the performance between GNN-graphs and HAGs.
Compared to GNN-graphs,
HAGs improve the training performance by up to 2.6 × and 3.1 × on GIN and SGC, while maintaining the same network accuracy.
This shows that the HAG optimization is general and applicable to various GNN models.
[Figure]
F T IME - TO - ACCURACY E VALUATION
We compare the time-to-accuracy performance between HAG and GNN-graph.
We train a 2-layer
GCN model (with 64 hidden dimensions in each layer) on the Reddit dataset until the test accuracy exceeds 95%.
We follow previous work Kipf and Welling (2016) to set all hyper-parameters and split the dataset.
Figure 6 shows the results.
Each dot indicates the training time and test accuracy of each epoch.
Training the GCN model using the HAG representation achieves the same training and test accuracy at the end of each epoch.
This is because the HAG optimization maintains the same results as the original GNN-graph representation and therefore preserves the model accuracy.
It takes 55 training epochs to achieve a test accuracy of 95% for both HAG and GNN-graph, and HAG improves the end-to-end training time by 1.8 × .