A DAPTIVE F ILTERS FOR L OW -L ATENCY AND M EMORY -E FFICIENT G RAPH N EURAL N ETWORKS
Anonymous authors
Paper under double-blind review
A BSTRACT
Scaling and deploying graph neural networks (GNNs) remains difﬁcult due to their high memory consumption and inference latency.
In this work we present a new type of GNN architecture that achieves superior performance to competing GNN approaches, with lower memory consumption and latency, along with characteristics suited to accelerator implementation.
Our proposal uses memory proportional to the number of vertices in the graph ( O ( V ) ), in contrast to competing methods which require memory proportional to the number of edges ( O ( E ) ); surprisingly, we ﬁnd our efﬁcient approach actually achieves higher accuracy than competing approaches across 6 large and varied datasets against strong baselines.
We achieve our results by using a novel adaptive ﬁltering approach, which can be interpreted as enabling each vertex to have its own weight matrix, and is not directly related to attention.
Following our focus on efﬁcient hardware usage, we demonstrate that our method achieves lower latency and memory consumption for the same accuracy when compared to competing approaches.
I NTRODUCTION
Graph Neural Networks (GNNs) have emerged as an effective way to build models over arbitrarily structured data.
For example, they have successfully been applied to computer vision tasks: GNNs can deliver high performance on point cloud data (Qi et al., 2017) and for feature matching across images (Sarlin et al., 2020).
Recent work has also shown that they can be applied to physical simulations (Pfaff et al., 2020; Sanchez-Gonzalez et al., 2020).
Code analysis is another application domain where GNNs have found success (Guo et al., 2020; Allamanis et al., 2017).
Our work aims to make a step towards enabling these applications, and many more, by investigating methods to cheaply boost GNN representational power.
Succinctly, we investigate whether it is possible to design GNN architectures that achieve high performance without compromising on efﬁciency.
Importantly, our method delivers beneﬁts at both training and inference time, and can be successfully applied to tasks requiring generalization to unseen graphs.
The development of hardware-efﬁcient techniques is key to the scaling and deployment of deep learning.
We have seen the deployment of convolutional neural networks (CNNs) to enable novel applications at the edge, due to innovation at the hardware and algorithmic level.
Investigating methods to improve the efﬁciency of GNNs will enable us to train larger models, on larger datasets, along with scaling down our models so that we can deploy them for applications at the edge.
Research into efﬁciency often focuses on hardware-software co-design: the development of techniques to enable the software to take better advantage of the hardware, and vice-versa.
There are several facets to this ﬁeld (Sze et al., 2020), with common examples including quantization and pruning.
Another key area is designing neural network architectures with efﬁciency as a design goal .
To date, research on GNN architectures has typically focused on boosting accuracy, with efﬁciency being a secondary consideration; many works focusing on efﬁciency tend not to address the underlying architecture design, and instead consider orthogonal approaches, such as quantization or graph sampling.
We have seen memory consumption and latency for state-of-the-art GNN architectures increase in recent years, predominantly due to models incorporating anisotropy —treating neighbors differently using learnt functions—to boost accuracy.
Surprisingly, we ﬁnd that we can achieve superior performance to existing GNNs without explicitly incorporating anisotropy into our architecture, challenging conventional wisdom in the community.
x
Messages are functions of source and target nodes and hence must be materialized b f ( x , b ) a f ( x , a ) c f ( x , c ) d f ( x , d ) x
Messages are functions of source only ; propagation can be implemented using matrix multiplication-style approaches b f ( b ) a f ( a ) c f ( c ) d f ( d )
Contributions .
We propose a new GNN architecture, Efﬁcient Graph Convolution ( EGC ).
Our proposal’s memory consumption is linear in the number of vertices in the graph ( O ( V ) )—not the number of edges ( O ( E ) ).
Our architecture’s latency is also lower than competing approaches, and is a drop-in replacement on a variety of tasks.
We provide a rigorous evaluation of our architecture across 6 large graph datasets covering both transductive and inductive use-cases.
Application domains ranging from citation graphs to code analysis and molecular property prediction are covered, and demonstrate that our approach consistently achieves better results than strong baselines, and that our approach can handle heterogeneous graphs.
We also demonstrate that our method can achieve lower memory consumption and latency for a given accuracy.
B ACKGROUND 2. H ARDWARE -S OFTWARE C O -D ESIGN FOR D EEP L EARNING
Several of the popular approaches for co-design have already been described in the introduction: quantization, pruning, and careful architecture design are all common for CNNs and Transformers (Vaswani et al., 2017).
In addition to enabling better performance to be obtained from general purpose processors such as CPUs and GPUs, these techniques are also essential for maximizing the return from specialized accelerators; while it may be possible to improve performance over time due to improvements in CMOS technology, further improvements plateau without innovation at the algorithmic level (Fuchs & Wentzlaff, 2019).
As neural network architecture designers, we cannot simply rely on improvements in hardware to make our proposals viable for real-world deployment.
2. G RAPH N EURAL N ETWORKS
Many GNN architectures can be viewed as a generalization of CNN architectures to the irregular domain: as in CNNs, representations at each node are built based on the local neighborhood using parameters that are shared across the graph.
GNNs differ as we cannot make assumptions about the the size of the neighborhood, or the ordering.
One common framework used to deﬁne GNNs is the message passing neural network (MPNN) paradigm (Gilmer et al., 2017).
A graph G =
( V, E ) has node features X ∈ R N × F , adjacency matrix A ∈ R N × N and optionally D -dimensional edge features E ∈ R E × D .
We deﬁne a function φ that calculates messages from node u to node v , a differentiable and permutation-invariant aggregator ⊕ , and an update function γ to calculate representations at layer l
+1 : h ( i ) l +1 = γ
( h ( i ) l , ⊕ j ∈N ( i )
[ φ ( h ( i ) l , h ( j ) l , e ij )]) .
Propagation rules for baseline architecture are provided in Table 1, with further details supplied in Table 11 in the Appendix.
This work builds upon the concept of using basis matrices.
The idea of combining basis matrices was proposed in R-GCN (Schlichtkrull et al., 2018) to handle multiple edge types; Xu et al.
(2021) demonstrated that this concept can also be applied to point cloud analysis.
In this work we are solv-
ing a fundamentally different problem: we are interested in designing efﬁcient GNN architectures , not merely boosting expressivity.
Scaling and Deploying GNNs While GNNs have seen success across a range of domains, there remain challenges associated with scaling and deploying them.
Graph sampling is one approach to scaling training for large graphs or models which will not ﬁt in memory.
Rather than training over the full graph, each iteration is run over a sampled sub-graph; approaches vary in whether they sample node-wise (Hamilton et al., 2017), layer-wise (Chen et al., 2018a;
Huang et al., 2018), or subgraphs (Zeng et al., 2019; Chiang et al., 2019).
Alternatively, systems for distributed GNN training have been proposed (Jia et al., 2020) to scale training beyond the limits of a single accelerator.
Some works have proposed architectures that are designed to accommodate scaling: graph-augmented
MLPs, such as SIGN (Rossi et al., 2020), are explicitly designed as a shallow architecture, as all the graph operations are done as a pre-processing step.
Other work includes applying neural architecture search (NAS) to arrange existing GNN layers (Zhao et al., 2020), or building quantization techniques for GNNs (Tailor et al., 2021).
Finally, a recent work has shown that using memoryefﬁcient reversible residuals (Gomez et al., 2017) for GNNs (Li et al., 2021) enables us to train far deeper and larger GNN models than before, thereby progressing the state-of-the-art accuracy.
Why Are Existing Approaches Not Sufﬁcient?
It is worth noting that many of these approaches have signiﬁcant limitations that we hope to address with our work.
Sampling methods are often ineffective when applied to many problems which involve model generalization to unseen graphs—a common use-case for GNNs.
We evaluated a variety of sampling approaches and observed that even modest sampling levels, which provide little beneﬁt to memory or latency, cause model performance to decline noticeably.
In addition, these methods do not accelerate the underlying GNN, hence they may not provide any overall beneﬁt to inference latency.
There is also no evidence that we are aware of that graph-augmented MLPs perform adequately when generalizing to unseen graphs; indeed, they are known to be theoretically less expressive than standard GNNs (Chen et al., 2021).
We also investigated this setup, and found that these approaches do not offer competitive accuracy with state-of-the-art approaches.
Experiment details and results, along with further discussion of the limitations of existing work, is provided in Appendix B.
In summary, efﬁcient GNN architecture design is of interest to the community as it enables us to scale our models further, potentially yielding improvements in accuracy.
In addition, for tasks where we need to generalize to unseen graphs, such as code analysis or point cloud processing, we reduce memory consumption and latency, thereby enabling us to deploy our models to more resource-constrained devices than before.
We note that efﬁcient architecture design can be usefully combined with other approaches (e.g. sampling) where appropriate.
In this work, we propose an architecture that is not only far superior to competing architectures in terms of memory consumption and inference latency, but which is more parameter efﬁcient.
In contrast to competing architectures, which require O ( E ) memory, our proposal requires only O ( V ) memory.
O UR A RCHITECTURE : E FFICIENT G RAPH C ONVOLUTION (EGC)
In this section we describe our approach, and delay theoretical analysis to the next section.
We present two versions: EGC-S (ingle), using a single aggregator, and EGC-M (ulti) which generalizes our approach by incorporating multiple aggregators.
Our approach is visualized in Figure 2.
3. A RCHITECTURE D ESCRIPTION
For a layer with in-dimension of F and out-dimension of F (cid:48) we use B basis weights Θ b ∈ R F (cid:48)
× F .
We compute the output for node i by calculating combination weighting coefﬁcients w ( i ) ∈ R B per node , and weighting the results of each aggregation using the different basis weights Θ b .
We calculate w ( i ) =
Φx ( i ) + b , where Φ ∈ R B × F and b ∈ R B are weight and bias parameters associated with calculating the combination weighting coefﬁcients.
The output for node i is computed in three steps.
First, we perform the aggregation with each set of basis weights Θ b .
Second, we compute the weighting coefﬁcients w ( i ) =
Φx ( i ) + b ∈ R B for each node i , where Φ ∈ R B × F and b ∈ R B are weight and bias parameters for calculating the combination weighting coefﬁcients.
Third, the layer x
( i ) (cid:80) w ( i ) =
Φx ( i ) + b ΘΘΘ
[Figure]
×
Run separate graph ﬁlters parameterised by Θ b over graph
Apply per-node weightings w ( i ) to each basis ﬁlter
h
( i )h ( i )h
( i )y
( i ) output for node i is the weighted combination of aggregation outputs:
[Equation]
where α ( i, j ) is some function of nodes i and j , and N ( i ) denotes the in-neighbours of i .
A popular method pioneered by GAT (Veli ˇ ckovi ´
c et al., 2018) to boost representational power is to represent α using a learned function of the two nodes’ representations.
While this enables anisotropic treatment of neighbors, and can boost performance, it necessarily results in memory consumption of O ( | E | ) due to messages needing to be explicitly materialized, and complicates hardware implementation for accelerators.
If we choose a representation for α that is not a function of the node representations— such as α ( i, j )
= 1 to recover the add aggregator used by GIN (Xu et al., 2019), or α ( i, j )
= 1 / (cid:112) deg ( i ) deg ( j ) to recover symmetric normalization used by GCN (Kipf & Welling, 2017)— then we can implement our message propagation phase using sparse matrix multiplication (SpMM), and avoid explicitly materializing each message, even for the backwards pass.
In this work, we assume α ( i, j ) to be symmetric normalization as used by GCN unless otherwise stated; we use this normalization as it is known to offer strong results across a variety of tasks; more formal justiﬁcation is provided in section 4.2.
Adding Heads as a Regularizer We can extend our layer through the addition of heads, as used in architectures such as GAT or Transformers (Vaswani et al., 2017).
These heads share the basis weights, but apply different weighting coefﬁcients per head.
We ﬁnd that adding this degree of freedom aids regularization when the number of heads ( H ) is larger than B , as bases are discouraged from specializing (see section 5.3), without requiring the integration of additional loss terms into the optimization—hence requiring no changes to code for downstream users.
To normalize the output dimension, we change the basis weight matrices dimensions to F (cid:48)
H × F .
Using (cid:107) as the concatenation operator, and making the use of symmetric normalization explicit, we obtain the EGC-S layer:
[Equation]
3. B OOSTING R EPRESENTATIONAL C APACITY
Recent work by Corso et al.
(2020) has shown that using only a single aggregator is sub-optimal: instead, it is better to combine several different aggregators.
In Equation (2) we deﬁned our layer to use only symmetric normalization.
To improve performance, we propose applying different aggregators to the representations calculated by Θ b x ( j ) .
The choice of aggregators could include different variants of summation aggregators e.g. mean or unweighted addition, as opposed to symmetric normalization that was proposed in the previous section.
Alternatively, we can use aggregators such as stddev, min or max which are not based on summation.
It is also possible to use directional aggregators proposed by Beaini et al.
(2021), however this enhancement is orthogonal to this work.
If we have a set of aggregators A , we can extend Equation (2) to obtain our EGC-M layer:
[Equation]
where ⊕ is an aggregator.
With this formulation, we are reusing the same messages we have calculated as before—but we are applying several aggregation functions to them at the same time.
Aggregator Fusion It would appear that adding more aggregators would cause latency and memory consumption to grow linearly.
However, this is not true in practice.
Firstly, since sparse operations are typically memory bound in practice, we can apply extra aggregators to data that has already arrived from memory with little latency penalty.
EGC can also efﬁciently inline the nodewise weighting operation at inference time, thereby resulting in relatively little memory consumption overhead.
The equivalent optimization is more difﬁcult to apply successfully to PNA due to the larger number of operations per-node that must be performed during aggregation, caused by scaling functions being applied to each aggregation, before all results are concatenated and a transformation applied.
More details, including proﬁling and latency measurements, can be found in Appendix D.
I NTERPRETATION AND B ENEFITS
This section will explain our design choices, and why they are better suited to the hardware.
We emphasize that our approach does not directly correspond to attention.
4. S PATIAL I NTERPRETATION : N ODE -W ISE W EIGHT M ATRICES
In our approach, each node effectively has its own weight matrix.
We can derive this by re-arranging eq.
(2) by factorizing the Θ b terms out of inner sum:
[Equation]
In contrast
[Equation]
GAT shares weights, and pushes complexity into the message calculation phase by calculating per-message weightings.
MPNN (Gilmer et al., 2017) and PNA (Corso et al., 2020) further increase complexity by explicitly calculating each message—resulting in substantial latency overhead due to the number of dense operations increasing by roughly | E | | V | .
Speciﬁcally
[Equation]
we have:
[Equation]
From an efﬁciency perspective we observe that our approach of using SpMM has better characteristics due to it requiring only O ( V ) memory consumption—no messages must be explicitly materialized to use SpMM.
We note that although it is possible to propagate messages for GAT with
SpMM, there is no way to avoid storing the weightings during training as they are needed for backpropagation, resulting in O ( E ) memory consumption.
We also note that fusing the message and aggregation steps for certain architectures may be possible at inference time, but this is a difﬁcult pattern for hardware accelerators to optimize for.
Relation To Attention Our method is not directly related to attention , which relies upon pairwise similarity mechanisms, and hence results in a O ( E ) cost when using the common formulations.
Alternatives to attention-based Transformers proposed by Wu et al. (2019a) are a closer analogue to our technique, but rely upon explicit prediction of the per-timestep weight matrix.
This approach is not viable for graphs, as the neighborhood size is not constant.
[Table]
4. S PECTRAL I NTERPRETATION : L OCALISED S PECTRAL F ILTERING
We can also interpret our EGC-S layer through the lens of graph signal processing (Sandryhaila &
Moura, 2013).
Many modern graph neural networks build on the observation that the convolution operation for the Euclidean domain when generalised to the graph domain has strong inductive biases: it respects the structure of the domain and preserves the locality of features by being an operation localised in space.
Our method can be viewed as a method of building adaptive ﬁlters for the graph domain.
Adaptive ﬁlters are a common approach when signal or noise characteristics vary with time or space; for example, they are commonly applied in adaptive noise cancellation.
Our approach can be viewed as constructing adaptive ﬁlters by linearly combining learnable ﬁlter banks with spatially varying coefﬁcients.
The graph convolution operation is typically deﬁned on the spectral domain as ﬁltering the input signal x ∈ R N on a graph with N nodes with a ﬁlter g θ parameterized by θ .
This requires translating between the spectral and spatial domain using the Fourier transform.
As on the Euclidean domain, the Fourier transform on the graph-domain is deﬁned as the basis decomposition with the orthogonal eigenbasis of the Laplace operator, which for a graph with adjacency matrix A ∈ R N × N is deﬁned as L = D − A , where D is the diagonal degree matrix with D ii =
(cid:80) Nj =
1 A ij .
The Fourier transform of a signal x ∈ R N then is F ( x ) =
U (cid:62) x , where L = UΛU (cid:62) , with orthogonal eigenvector-matrix
U ∈ R N × N and diagonal eigenvalue-matrix Λ ∈ R N × N .
The result of a signal x ﬁltered by g θ is y
=
g θ ( L ) x
= U g θ ( Λ ) U (cid:62) x where the second equality holds if the Taylor expansion of g θ exists.
Our approach corresponds to learning multiple ﬁlters and computing a linear combination of the resulting ﬁlters with weights depending on the attributes of each node locally.
The model therefore allows applying multiple ﬁlters for each node, enabling us to obtain a spatially-varying frequency response, while staying far below O ( E ) in computational complexity.
Using a linear combination of ﬁlters, the ﬁltered signal becomes y = (cid:80) B b =1 w b (cid:12) g θ b ( L ) x , where w b ∈ R N are the weights of ﬁlter b for each of the N nodes in the graph.
If we parameterize our ﬁlter using ﬁrst-order Chebyshev polynomials as used by Kipf & Welling (2017) our ﬁnal expression for the ﬁltered signal becomes
[Equation]
˜ D − 12 ˜
A ˜ D − 12 ) XΘ b , where
˜ A =
A + I N is the adjacency matrix with added self-loops and ˜ D is the diagonal degree matrix of ˜ A as deﬁned earlier.
This justiﬁes the symmetric normalization aggregator we chose in Equation (2).
Cheng et al.
(2021) proposed an approach for localized ﬁltering.
However, their approach does not generalize to unseen topologies or scale to large graphs as it requires learning the coefﬁcients of several ﬁlter matrices S k of size N × N .
Our approach does not suffer from these constraints.
E VALUATION 5. P ROTOCOL
We primarily evaluate our approach on 5 datasets taken from recent works on GNN benchmarking.
We use ZINC and CIFAR-10 Superpixels from Dwivedi et al.
and Arxiv, MolHIV and Code
[Table] [Table]
from Open Graph Benchmark (Hu et al., 2020).
These datasets cover a wide range of domains, both transductive and inductive tasks, and are larger than datasets which are typically used in GNN works.
We use evaluation metrics and splits speciﬁed by these papers.
Baseline architectures chosen reﬂect popular general-purpose choices (Kipf & Welling, 2017; Xu et al., 2019; Hamilton et al., 2017; Veli ˇ ckovi ´ c et al., 2018; Gilmer et al., 2017), along with the recent state-of-the-art PNA (Corso et al., 2020).
Code and pre-trained models are supplied in the supplementary material.
In order to provide a fair comparison we standardize all parameter counts, architectures and optimizers in our experiments.
All experiments were run using Adam (Kingma & Ba, 2014).
Further details on how we ensured a fair evaluation can be found in the appendix.
We do not use edge features in our experiments as for most baseline architectures there exist no standard method to incorporate them.
We do not use sampling, which, as explained in Section 2.2, is ineffective for 4 datasets; for the remaining dataset, Arxiv, we believe it is not in the scientiﬁc interest to introduce an additional variable.
This also applies to GraphSAGE, where we do not use the commonly applied neighborhood sampling.
All experiments were run 10 times.
5. M AIN R ESULTS
Our results across the 5 tasks are shown in Table 1.
We draw attention to the following observations:
• EGC-S is competitive with anisotropic approaches .
We outperform GAT and MPNNSum on all benchmarks, despite our resource efﬁciency.
The exception is MPNN-Max on CIFAR & Code, where the max aggregator provides a stronger inductive bias.
• EGC-M outperforms PNA .
The addition of multiple aggregator functions improves performance of EGC to beyond that obtained by PNA.
We hypothesize that our improved performance over PNA is related to PNA’s reliance on multiple degree-scaling transforms.
While this approach can boost the representational power of the architecture, we believe that it can result in a tendency to overﬁt to the training set.
• EGC performs strongly without running out of memory.
We observe that EGC is one of only three architectures that did not exhaust the VRAM of the popular Nvidia 1080/2080Ti GPUs, with 11GB VRAM , when applied to Arxiv: we had to use an RTX 8000 GPU with 48GB VRAM to run these experiments.
PNA, our closest competing technique accuracywise, exhausted memory on the Code benchmark as well.
Detailed memory consumption ﬁgures are provided in Table 2.
Overall, EGC obtains the best performance on 4 out of the 5 main datasets; on the remaining dataset (MolHIV), EGC is the second best architecture.
This represents a signiﬁcant achievement: our architecture demonstrates that we do not need to choose between efﬁciency and accuracy.
5. A DDITIONAL S TUDIES
How Should Heads and Bases be Chosen?
To understand the trade-off between the number of heads ( H ) and bases ( B ), we ran an ablation study on ZINC using EGC-S; this in shown in Figure 3.
[Figure] [Figure]
the hidden dimension.
Each experiment was tuned individually.
Setting B > H does not necessarily improve performance due to the risk of overﬁtting, and forces the usage of a smaller hidden dimension to retain a constant parameter count.
The relationship between these parameters is non-trivial.
There are several aspects to consider: (1) increasing H and B means that we spend more of our parameter budget to create the combinations, which reduces hidden dimension—as shown in Figure 3.
This is exacerbated if we use multiple aggregators: our combination dimension must be HB |A| .
(2) Increasing B means we must reduce the hidden size substantially, since it corresponds to adding more weights of size F (cid:48)
H × F .
Increasing H allows us to increase hidden size, since each basis weight becomes smaller.
We see in Figure 3 that increasing B beyond H does not yield signiﬁcant performance improvements: we conjecture that bases begin specializing for individual heads; by sharing, there is a regularizing effect, like observed in Schlichtkrull et al.
This regularization stabilizes the optimization and we observe lower trial variance for smaller B .
We advise B = H or B = H 2 .
We ﬁnd H = 8 to be effective with EGC-S; for EGC-M, where more parameters are spent on combination weights, we advise setting H = 4 .
This convention is applied consistently for Table 1; full details are supplied in the appendix and code.
[Table]
Should The Combination Weightings ( w )
Be Activated?
Any activation function will shrink the space the per-node weights Θ ( i ) h can lie in, hence we would expect it to harm performance; this is veriﬁed in Table 4.
Activating w may improve training stability, but we did not observe to be an issue in our experiments.
Another issue is that different aggregators result in outputs with different means and variances (Tailor et al., 2021), hence they need to be scaled by different factors to be combined.
Applying EGC to Large-Scale Heterogeneous Graphs We evaluated EGC on the heterogeneous OGB-MAG dataset, containing 2M nodes and 21M edges.
On a simpliﬁed, homogeneous, version of the graph, we exceed the baselines’ performance by 1.5-2%; we experiment with both symmetric normalization (EGC-S) and the mean aggregators to demonstrate that the mechanism utilized by EGC is effective, regardless of which aggregator provides the stronger inductive bias for the dataset.
Our architecture can be expanded to handle different edge types, yielding the R-EGC architecture which improves performance over R-GCN by 0.9%.
We expect that accuracy can be further improved by using sampling techniques to regularize the optimization, or using pretrained embeddings; however, adding these techniques makes comparing the results more difﬁcult as it is known that sampling techniques can affect each architecture’s performance differently (Liu et al., 2020).
[Table]
5. M EMORY AND L ATENCY B ENCHMARKS
We now assess our model’s resource efﬁciency.
For CPU measurements we used an Intel Xeon Gold 5218 and for GPU we used an Nvidia RTX 8000.
Aggregator Fusion We evaluated aggregator fusion across several topologies on both CPU and
GPU.
For space reasons, we leave full details to Appendix D. In summary, we observe that usingaggregators over standard SpMM incurs an additional overhead of only 14% , enabling us to improve model performance without excessive computational overheads at inference time.
End-to-End Latency
We provide latency and memory statistics for the parameter-normalized models in Table 2.
We draw attention to how slow and memory-intensive the O ( E ) models are for both training and inference.
The reader should note that the inference latency and memory consumption for MPNN and PNA rises by approximately 6-7 × relative to EGC, corresponding to the large | E | | V | ratio for Arxiv.
EGC-M offers substantially lower latency and memory consumption than its nearest competitor, PNA, as | E | is often ≥ 2 | V | due to the graphs being directed, however the precise beneﬁt will be dataset-dependent.
We report results on other datasets in Appendix E.
So far we have not demonstrated that our approach is more efﬁcient than the baselines on Arxiv, which is the only dataset where EGC-M and PNA are not the best performing.
To demonstrate that we are more efﬁcient we must show that we achieve lower memory consumption and latency for a given accuracy-level.
We evaluate increasing the parameter count for baseline models until they achieve the same accuracy as EGC-S; we gave these models an extra advantage over our method by increasing the hyperparameter search budget.
The results are shown in Table 5, where we observe that EGC-S is more efﬁcient once we are comparing models achieving the same accuracy.
D ISCUSSION AND C ONCLUSION
How Surprising Are Our Results?
Our results contradict common wisdom in the GNN community: they directly challenge the observations regarding anisotropy improving performance described in the widely cited benchmarking work by Dwivedi et al.
Our observations raise important questions regarding our understanding of GNN optimization: we believe that although certain baselines are more expressive, they are more difﬁcult to train.
Related observations to ours have been made in the NLP community: it has been demonstrated that convolutional-style networks can be competitive with Transformers in many cases (Tay et al., 2021).
This work has made a step towards improving the runtime efﬁciency of GNNs.
Our proposed layer can be used as a drop-in replacement for existing GNN layers, and achieves better results acrossbenchmark datasets compared to strong baselines.
We have demonstrated that our approach can be used to replace anisotropic mechanisms currently pervasive in GNN architecture design to improve parameter efﬁciency, memory consumption and latency.
In summary, our work raises important questions for the research community, while offering signiﬁcant practical beneﬁts with regard to resource consumption.
We believe the next step for our work is incorporation of edge features e.g. through line graphs (Chen et al., 2020b) or topological message passing (Bodnar et al., 2021).
E THICS S TATEMENT
The method described in this paper is generic enough that it can be applied to any problem GNNs are applied to.
The ethical concerns associated with this work are related to enabling more efﬁcient training and deployment of GNNs.
This may be positive (drug discovery) or negative (surveillance), but these concerns are inherent to any work investigating efﬁciency for machine learning systems.
R EPRODUCIBILITY S TATEMENT
We have supplied the code required to regenerate our results, along with the hyperparameters required.
In addition, we supply pre-trained models.
Please refer to the supplementary material for more details.
A F URTHER E XPERIMENT D ETAILS A. E NSURING F AIRNESS
We expand on our experimental protocol, with a particular focus on describing measures that we took to ensure that the results we report are not unfairly biased towards EGC.
For EGC-S, we use H = 8 and B = 4 , as implied by our ablation for all experiments, with the single exception of OGB-Code, where we use H = B = 8 .
The beneﬁt of using a smaller set of bases is that we can increase the hidden dimension, but this is not viable in this case since most of the 11M parameters in the model correspond to the token read-out layers, which quickly increases as the model hidden dimension grows.
As shown by section 5.3, if we cannot increase the hidden dimension, it is better to increase the bases.
This is the only exception we make.
For EGC-M, we use H = B = 4 for all experiments.
The main challenge is aggregator selection, and this remains a major challenge for our work.
We were unable to ﬁnd a satisfactory technique for automated discovery of aggregator choices, hence we rely on heuristics to ﬁnd them.
We restrict ourselves to use 3 aggregators for each model (yielding 35 possible choices).
In order to determine the aggregators, we use two heuristics: (1) aggregators should be “diverse” and (2) aggregators should be chosen based on inductive bias for task.
Using these two rules, we try up to 3 possible choices of aggregators; all choices considered are shown in Table 6.
We note that while some choices do improve performance, our conclusions are not invalidated; it is also worth noting that it is likely that better aggregator choices can be found.
A. C LUSTER D ETAILS
Most of our experiments were run on several machines in our SLURM cluster using Intel CPUs and
NVIDIA GPUs.
Each machine was running Ubuntu 18.04.
The GPU models in our cluster were
RTX 2080Ti and GTX 1080Ti.
High-memory experiments were run on V100s in our cluster and an
RTX 8000 virtual machine we had access to.
B L IMITATIONS OF E XISTING A PPROACHES
As explained in the related work, existing approaches to improving GNN efﬁciency have severe limitations.
In this section we elaborate upon them, and provide experimental evidence where necessary.
[Table] [Table]
We ﬁrst examine graph-augmented MLPs (GA-MLPs); to our knowledge there are few experimental results assessing the performance of these models when they are applied to problems requiring generalization to unseen graphs.
In the literature they are generally applied to large scale node classiﬁcation benchmarks, such as those found in the OGB benchmark suite (Wu et al., 2019b;
Rossi et al., 2020).
It is known that these models are theoretically less expressive than standard
GNNs, however their performance when applied to node classiﬁcation datasets has been acceptable.
We consider models of the form:
[Equation]
We use up to the 4th power of the diffusion operator S to emulate the depth of the corresponding
GNNs, which all use 4 layers.
We set S to use symmetric normalization, as used by GCN and EGCS; we also consider setting S to the adjacency matrix on ZINC, to emulate the operations used by
GIN.
The results are provided in Table 7.
We see that the GA-MLP models offer similar (but often worse) performance than the corresponding GNN baselines.
The models are not competitive with approaches such as GAT or MPNN, and is outperformed by a wide margin by EGC-S. Achieving state-of-the-art performance with GA-MLPs does not appear to be possible, at least with our current understanding of these models.
We now proceed to investigate sampling, in which each training step does not run over the entire graph, but over some sampled subgraph.
These methods have seen great popularity when applied to tasks such as node classiﬁcation, and they are able to deliver sampling ratios in excess of 20 × , hence yielding noticeable improvements to memory consumption (the primary limitation with large scale training).
However, we note that the beneﬁt of these methods has not been examined carefully for many graph topologies; while they have been shown to be effective for many “small-world” graph topologies—which arise in many graphs—it is not the case that all graphs fall into this category.
For example, molecule graphs would not ﬁt this category.
We ﬁrst assess the sampling strategy from GraphSAINT (Zeng et al., 2019) by applying it to EGC-S models.
For our experiments, we use the node-centric sampler.
The results are presented in Table 8; we disable sampling at test time.
Even with a relatively low dropping probability of 10% (i.e. 90% of nodes are retained), the model performance degradation is severe.
We note that the computational savings achieved are modest when using this drop probablility.
We also observed similar results when using the edge-centric sampler from GraphSAINT.
We also attempted a different approach for sampling proposed by DropEdge (Rong et al., 2020); as before, we apply it to EGC-S models.
In this simple scheme, elements of the adjacency matrix are dropped; this scheme was shown to be effective for training deep graph networks, but it is also useful for reducing the computational footprint of models, since it effectively reduces the number of messages that have to be computed.
We also provide results in Table 8.
The results are signiﬁcantly better than we observe when using GraphSAINT’s sampling strategies: at 10% drop probabilities we even see an improvement in some cases, due to the regularization effect.
However, once we increase the drop probability to levels where we would observe a noticeable reduction in computational demand, we observe that model performance declines.
In summary, while DropEdge is a more effective strategy for sampling on many inductive tasks, it is not beneﬁcial as a method to reduce the computational burden.
Finally, we discuss the limitations of other approaches proposed in the literature.
Quantization is an approach that is typically applied at inference time; mixed-precision approaches can be applied at training time, however care must be taken for GNNs to avoid biasing the gradients (Tailor et al., 2021).
Additionally, while neural architecture search (NAS) may be useful to ﬁnd memory-efﬁcient models if the search objective is set appropriately, they suffer from some limitations—primarily search time and memory consumption.
Finally, approaches such as reversible residuals are useful to architecture design, they do not tackle issues such as high peak memory usage induced by the message passing step.
C A PPROACHES FOR H ARDWARE A CCELERATION OF GNN S
The reader should note that most existing work for GNN hardware acceleration focuses on supporting only a subset of GNNs: speciﬁcally, they tend to only support models that can be implemented using SpMM. Approaches in the literature falling into this area include Chen et al. (2020a), Yan et al.
You et al. (2021) and Zhang et al. It is possible to add greater ﬂexibility to the accelerator to support more expressive message passing schemes, however this necessarily implies greater complexity. As Amdahl’s law (Amdahl, 1967) implies, increasing ﬂexibility is likely to reduce peak performance, while increasing silicon area requirements.
Therefore, aiming for the simplest primitive (as we do with EGC) is the most sensible approach to obtain hardware acceleration. D A GGREGATOR F USION The naive approach of performing each aggregation sequentially would cause a linear increase in latency with respect to |A| . However, a key observation to note is that we are memory-bound : the bottleneck with sparse operations is waiting for the data to arrive from memory.
This observation applies to both GPUs and CPUs, and justiﬁed through proﬁling. Using a proﬁler on a GTX 1080Ti we observed that SpMM using the Reddit graph Hamilton et al. (2017) with feature sizes ofachieved just 1.2% of the GPU’s peak FLOPS, with 88.5% of stalls being caused by unmet memory dependencies. The fastest processing order performs as much work as possible with data that has already been fetched from memory, rather than fetching it multiple times.
This concept is illustrated in Algorithm 1 in the appendix.
We can perform all aggregations as a lightweight modiﬁcation to the standard compressed sparse row (CSR) SpMM algorithm. The second observation we make is that storing the results of all aggregations is unnecessary at inference time. Note that the CSR SpMM algorithm processes each row in the output matrix sequentially: rather than storing the aggregations for every row, instead we should store only the weighted results.
This approach not only reduces memory consumption, but also latency as we improve the effectiveness of our cache and reduce memory system contention. In practice, this optimization is especially important when performing inference on topologies which have more frequent periods
Input: CSR A ∈ R N × N , Dense B ∈ R N × F , Combination weightings w ∈ R N ×|A|
Output: Dense C ∈ R N × F for i = 0 to A .
rows − 1 do for jj = A .
row pointer [ i ] to A .
row pointer [ i + 1] do j = A .
column index [ jj ]
Init temp arrays of length F per aggregator a ij = A .
values [ jj ] { May be faster to interleave these calls: } for ⊕ ∈
A do process row ⊕ ( a ij , B [ i, :] , temp ⊕ ) end for end for { Can be generalized to H, B > 1 : }
C
[ i, :] =
(cid:80)
⊕∈A w [ i, ⊕ ] · temp ⊕ [:] end for
[Table]
where the processing has become compute-bound, since we reduce contention between load and store units Corporation.
This is also demonstrated in Algorithm 1.
We evaluated aggregator fusion across four different topologies, on both CPU and GPU; our results can be found in Table 9.
We assumed all operations are 32-bit ﬂoating point, and that we were using three aggregators: summation-based, max, and min; these aggregators match those used for
EGC-M Code.
Our benchmarks were conducted on a batch of 10k graphs from the ZINC and Code datasets, Arxiv, and the popular Reddit dataset (Hamilton et al., 2017), which is one of the largest graph datasets commonly evaluated on in the GNN literature.
Our SpMM implementation on GPU is based on Yang et al.
Code for the kernels are provided in our repo.
As expected, our technique optimizing for input re-use achieves signiﬁcantly lower inference latency than the naive approach to applying multiple aggregators.
While the naive approach results in a mean increase in latency of 331%, our approach incurs a mean increase of only 14% relative to ordinary
SpMM, used by GCN and GIN.
The increase is topology dependent, with larger increases in latency being observed for topologies which are less memory-bound.
We also provide timings for dense matrix multiplication (i.e. XΘ ) to justify our focus on optimizing sparse operations in this work:
CSR SpMM operation is 4.7 × slower (geomean) than the corresponding weight multiplication.
We believe further optimizations of the operations used by architecture are achievable through the use of auto-tuning frameworks e.g. TVM (Chen et al., 2018b), but this lies beyond the scope of this work. [Table]
E L ATENCY AND M EMORY C ONSUMPTION ON O THER D ATASETS
In the evaluation [Table] we assessed the memory consumption and latency for the parameter-normalized models on Arxiv.
In this section, we consider a similar exercise for OGB Code models.
The results are provided in Table 10.
Our conclusions remain broadly similar, with EGC-M offering clear improvements to memory consumption, latency [Table] and parameter efﬁciency relative to PNA.
EGC-S is superior to GAT, with similar inference latency, better model performance [Table] and noticeably lower memory consumption.
We train models using batch size 128; the reader should note that the memory consumption ﬁgures can vary between runs (even for the same model), since graphs vary in the number of nodes and edges.
F G ENERALIZING TO H ETEROGENEOUS G RAPHS
Our R-EGC model is similar to the baseline R-GCN model included in the OGB repository.
The
OGB model deviates from the standard deﬁnition of R-GCN (Schlichtkrull et al., 2018) since it handles different node types, not just edge types.
The baseline model has weights to generate messages for each relation-type, and a weight matrix to update each individual node type.
This corresponds to:
[Equation]
where η corresponds to the node type of node i , and R represents the set of relation types.
Note that the mean aggregator is used.
We deviate from the baseline by using a single set of basis weights.
Instead, we use a different weighting calculation layers ( w ( i ) =
Φx + b ) per node and relation type.
[Equation] [Table]