D EGREE -Q UANT : Q UANTIZATION -A WARE T RAINING FOR G RAPH N EURAL N ETWORKS
Anonymous authors
Paper under double-blind review
A BSTRACT
Graph neural networks (GNNs) have demonstrated strong performance on a wide variety of tasks due to their ability to model non-uniform structured data.
Despite their promise, there exists little research exploring methods to make them more efﬁcient at inference time.
In this work, we explore the viability of training quantized
GNNs, enabling the usage of low precision integer arithmetic during inference.
For
GNNs seemingly unimportant choices in quantization implementation cause dramatic changes in performance.
We identify the sources of error that uniquely arise when attempting to quantize GNNs, and propose an architecturally-agnostic and stable method, Degree-Quant , to improve performance over existing quantizationaware training baselines commonly used on other architectures, such as CNNs.
We validate our method on six datasets and show, unlike previous attempts, that models generalize to unseen graphs.
Models trained with Degree-Quant for INT quantization perform as well as FP32 models in most cases; for INT4 models, we obtain up to 26% gains over the baselines.
Our work enables up to 4.7 ⇥ speedups on CPU when using INT8 arithmetic.
I NTRODUCTION
GNNs have received substantial attention in recent years due to their ability to model irregularly structured data.
As a result, they are extensively used for applications as diverse as molecular interactions (Duvenaud et al., 2015; Wu et al., 2017), social networks (Hamilton et al., 2017), recommendation systems (van den Berg et al., 2017) or program understanding (Allamanis et al., 2018).
Recent advancements have centered around building more sophisticated models including new types of layers (Kipf & Welling, 2017; Velickovic et al., 2018; Xu et al., 2019) and better aggregation functions (Corso et al., 2020).
However, despite GNNs having few model parameters, the compute required for each application remains tightly coupled to the input graph size.
A 2-layer
Graph Convolutional Network (GCN) model with 32 hidden units would result in a model size of just 81KB but requires 19 GigaOPs to process the entire Reddit graph.
We illustrate this growth in ﬁg. 1.
One major challenge with graph architectures is therefore performing inference efﬁciently, which limits the applications they can be deployed for.
For example, GNNs have been combined with CNNs for SLAM feature matching (Sarlin et al., 2019), however it is not trivial to deploy this technique on smartphones, or even smaller devices, whose neural network accelerators often do not implement ﬂoating point arithmetic, and instead favour more efﬁcient integer arithmetic.
Integer quantization is one way to lower the compute, memory and energy budget required to perform inference, without requiring modiﬁcations to the model architecture; this is also useful for model serving in data centers.
Although quantization has been well studied for CNNs and language models (Jacob et al., 2017;
Wang et al., 2018; Zafrir et al., 2019; Prato et al., 2019), there remains relatively little work addressing
GNN efﬁciency (Mukkara et al., 2018; Jia et al., 2020; Zeng & Prasanna, 2020; Yan et al., 2020).
To the best of our knowledge, there is no work explicitly characterising the issues that arise when quantizing GNNs or showing latency beneﬁts of using low-precision arithmetic.
The recent work of
Wang et al.
(2020) explores only binarized embeddings of a single graph type (citation networks).
In Feng et al.
(2020) a heterogeneous quantization framework assigns different bits to embedding and attention coefﬁcients in each layer while maintaining the weights at full precision (FP32).
Due to the mismatch in operands’ bit-width the majority of the operations are performed at FP32 after data casting, making it impractical to use in general purpose hardware such as CPUs or GPUs.
In h lv h l +1 v
CNN
GNN h l +1 =
h l ⇤ K
h l +1 v =
( h lv ,
V u 2 N ( v ) [  ( h lu , h lv , e uv )]) addition they do not demonstrate how to train networks which generalize to unseen input graphs.
Our framework relies upon uniform quantization applied to all elements in the network and uses bit-widths (8-bit and 4-bit) that are supported by off-the-shelf hardware such as CPUs and GPU for which efﬁcient low-level operators for common operations found in GNNs exists.
This work considers the motivations and problems associated with quantization of graph architectures, and provides the following contributions:
• The explanation of the sources of degradation in GNNs when using lower precision arithmetic.
We show how the choice of straight-through estimator (STE) implementation, node degree, and method for tracking quantization statistics signiﬁcantly impacts performance.
• An architecture-agnostic method for quantization-aware training on graphs, Degree-Quant (DQ), which results in INT8 models often performing as well as their FP32 counterparts.
At INT4, models trained with DQ typically outperform quantized baselines by over 20%.
We show, unlike previous work, that models trained with DQ generalize to unseen graphs .
• We show that quantized networks achieve up to 4.7 ⇥ speedups on CPU with INT8 arithmetic, relative to full precision ﬂoating point, with 4-8 ⇥ reductions in runtime memory usage.
B ACKGROUND 2.
[Equation]
N EURAL N ETWORKS (MPNN S )
Many popular GNN architectures may be viewed as generalizations of CNN architectures to an irregular domain: at a high level, graph architectures attempt to build representations based on a node’s neighborhood (see ﬁg. 2).
Unlike CNNs, however, this neighborhood does not have a ﬁxed ordering or size.
This work considers GNN architectures conforming to the MPNN paradigm (Gilmer et al., 2017).
A graph G = ( V, E ) has node features X 2 R N ⇥ F , an incidence matrix I 2 N 2 ⇥ E , and optionally Ddimensional edge features E 2 R E ⇥ D .
The forward pass through an MPNN layer consists of message passing, aggregation and update phases: h ( i ) l
+1 =
( h ( i ) l ,
V j 2 N ( i )
[  ( h ( j ) l , h ( i ) l , e ij )]) .
Messages from node u to node v are calculated using function  , and are aggregated using a permutation-invariant function
V .
The features at v are subsequently updated using  .
We focus on three architectures with corresponding update rules:
Graph Convolution Network (GCN): h ( i ) l +1 =
P
[Equation]
, 2017), where d i refers to the degree of node i .
Graph Attention Network (GAT): h ( i ) l +
=
↵ i,i Wh ( i ) l + P j 2 N ( i ) ( ↵ i,j Wh ( j ) l ) , where ↵ represent attention coefﬁcients (Velickovic et al., 2018).
Graph Isomorphism Network (GIN): h ( i ) l +1 = f ⇥
[(1 +
[Equation]
) h
[Equation]
) l +
P
[Equation]
N (
i ) h
j ) l ] , where f is a learnable function (e.g. a MLP) and ✏ is a learnable constant (Xu et al., 2019).
2. Q UANTIZATION FOR N ON -G RAPH N EURAL N ETWORKS
Quantization allows for model size reduction and inference speedup without changing the model architecture.
While there exists extensive studies of the impact of quantization at different bitwidths (Courbariaux et al., 2015; Han et al., 2015; Louizos et al., 2017) and data formats (Micikevicius et al., 2017; Carmichael et al., 2018; Kalamkar et al., 2019), it is 8-bit integer (INT8) quantization that has attracted the most attention.
This is due to INT8 models reaching comparable accuracy levels to FP32 models (Krishnamoorthi, 2018; Jacob et al., 2017), offer a 4 ⇥ model compression, and result in inference speedups on off-the-shelf hardware as 8-bit arithmetic is widely supported.
Quantization-aware training (QAT) has become the de facto approach towards designing robust quantized models with low error (Wang et al., 2018; Zafrir et al., 2019; Wang et al., 2018).
In their simplest forms, QAT schemes involve exposing the numerical errors introduced by quantization by simulating it on the forward pass Jacob et al.
and make use of STE (Bengio et al., 2013) to compute the gradients—as if no quantization had been applied.
For integer QAT, the quantization of a tensor x during the forward pass is often implemented as: x q = min( q max , max( q min , b x/s + z c )) , where q min and q max are the minimum and maximum representable values at a given bit-width and signedness, s is the scaling factor making x span the [ q min , q max ] range and, z is the zero-point , which allows for the real value 0 to be representable in x q .
Both s and z are scalars obtained at training time.
hen, the tensor is dequantized as: ˆ x = ( x q  z ) s , where the resulting tensor ˆ x ⇠ x for a high enough bit-width.
This similarity degrades at lower bit-widths.
Other variants of integer QAT are presented in Jacob et al.
(2017) and Krishnamoorthi (2018).
To reach performance comparable to FP32 models, QAT schemes often rely on other techniques such as gradient clipping , to mask gradient updates based on the largest representable value at a given bit-width; noisy QAT, which stochastically applies QAT to a portion of the weights at each training step (Fan et al., 2020); or the re-ordering of layers (Sheng et al., 2018; Alizadeh et al., 2019).
Q UANTIZATION FOR GNN S
In this section, we build an intuition for why GNNs would fail with low precision arithmetic by identifying the sources of error that will disproportionately affect the accuracy of a low precision model.
Using this insight, we propose our technique for QAT with GNNs, Degree-Quant .
Our analysis focuses on three models: GCN, GAT and GIN.
This choice was made as we believe that these are among the most popular graph architectures, with strong performance on a variety of tasks (Dwivedi et al., 2020), while also being representative of different trends in the literature.
3. S OURCES OF E RROR
QAT relies upon the STE to make an estimate of the gradient despite the non-differentiable rounding operation in the forward pass.
If this approximation is inaccurate, however, then poor performance will be obtained.
In GNN layers, we identify the aggregation phase, where nodes combine messages from a varying number of neighbors in a permutation-invariant fashion, as a source of substantial numerical error, especially at nodes with high in-degree.
Outputs from aggregation have magnitudes that vary signiﬁcantly depending on a node’s in-degree: as it increases, the variance of aggregation values will increase.
Over the course of training q min and q max , the quantization range statistics, become severely distorted by infrequent outliers, reducing the resolution for the vast majority of values observed.
This reults in increased rounding error for nodes with smaller in-degrees.
Controlling q min and q max hence becomes a trade-off balancing truncation error and rounding error .
We can derive how the mean and variance of the aggregation output values vary as node in-degree, n , increases for each of the three GNN layers.
Suppose we model incoming message values for a single output dimension with random variables
X i , without making assumptions on their exact distribution
[Figure]
or independence.
Further, we use Y n as the random variable representing the value of node output after the aggregation step.
With GIN layers, we have Y n = (1 + ✏ )
X
+
P ni =
X i .
It is trivial to prove that E ( Y n ) =
O ( n ) .
The variance of the aggregation output is also O ( n ) in the case that that
[Equation]
P i Var ( X i ) .
We note that if
P i
[Equation]
= j Cov ( X i , X j ) is large then it implies that the network has learned highly redundant features, and may be a sign of over-ﬁtting.
Similar arguments can be made for GCN and GAT layers; we would expect GCN aggregation values to grow like O ( p n ) , and GAT aggregation values to remain constant ( O (1) ) due to the attention coefﬁcients.
We empirically validate these predictions on GNNs trained on Cora; results are plotted in ﬁg. 3.
We see that the aggregation values do follow the trends predicted, and that for the values of in-degree in the plot (up to 168) the covariance terms can be neglected.
As expected, the variance and mean of the aggregated output grow fastest for GIN, and are roughly constant for GAT as in-degree increases.
From this empirical evidence, it would be expected that GIN layers are most affected by quantization.
By using GIN and GCN as examples, we can see how aggregation error causes error in weight updates.
Suppose we consider a GIN layer incorporating one weight matrix in the update function i.e. h ( i ) l
+1 =
f
( Wy ( i ) GIN ) , where f is an activation function, y ( i ) GIN = (1 + ✏ )
h ( i ) l +
P j 2 N ( i ) h
( j ) l , and
N ( i ) denotes the in-neighbors of node i .
Writing y
( i ) GCN =
P k 2 N ( i ) ( 1 p d i d k Wh ( j ) l ) , we see that the derivatives of the loss with respect to the weights for GCN and GIN are:
GIN
[Equation]
The larger the error in y ( i ) GIN —caused by aggregation error—the greater the error in the weight gradients for GIN, which results in poorly performing models being obtained.
The same argument applies to GCN, with the h ( j ) >
l and y ( i ) GCN terms introducing aggregation error into the weight updates.
3. O UR M ETHOD : D EGREE -Q UANT
To address these sources of error we propose Degree-Quant (DQ), a method for QAT with GNNs.
We consider both inaccurate weight updates and unrepresentative quantization ranges .
Stochastic Protection from Quantization to Improve Weight Update Accuracy .
DQ aims to encourage more accurate weight updates by stochastically protecting nodes in the network from quantization.
At each layer a protective node mask is generated; all masked nodes have the phases of the message passing, aggregation and update performed at full precision.
This includes messages sent by protected nodes to other nodes, as shown in ﬁg. 4 (a detailed diagram is shown in ﬁg. 11).
It is also important to note that the weights used at all nodes are the same quantized weights; this is motivated by the fact that our method is used to encourage more accurate gradients to ﬂow back to the weights through high in-degree nodes.
At test time protection is disabled: all nodes operate at low precision.
To generate the mask, we pre-process each graph before training and create a vector of probabilities p with length equal to the number of nodes.
At training time, mask m is generated by sampling using the Bernoulli distribution: m ⇠ Bernoulli ( p ) .
In our scheme p i is higher if the in-degree of node i is large, as we ﬁnd empirically that high in-degree nodes contribute most towards error in weight updates.
We use a scheme with two hyperparameters, p min and p max ; nodes with the maximum in-degree are assigned p max as their masking probability, with all other nodes assigned a probability calculated by interpolating between p min and p max based on their in-degree ranking in the graph.
Percentile Tracking of Quantization Ranges .
Figure 3 demonstrates large ﬂuctuations in the variance of the aggregation output as in-degree increases.
Since these can disproportionately affect the ranges found by using min-max or momentum-based quantization, we propose using percentiles .
While percentiles have been used for post-training quantization (Wu et al., 2020), we are the ﬁrst (to the best of our knowledge) to propose making it a core part of QAT; we ﬁnd it to be a key contributor to achieving consistent results with graphs.
Using percentiles involves ordering the values in the tensor and clipping a fraction of the values at both ends of the distribution.
The fraction to clip is a hyperparameter.
We are more aggressive than existing literature on the quantity we discard: we clip the top and bottom 0.1%, rather than 0.01%, as we observe the ﬂuctuations to be a larger issue with
GNNs than with CNNs or DNNs.
Quantization ranges are more representative of the vast majority of values in this scheme, resulting in less rounding error .
We emphasize that a core contribution of DQ is that it is architecture-agnostic .
Our method enables a wide variety of architectures to use low precision arithmetic at inference time.
Our method is also orthogonal —and complementary—to other techniques for decreasing GNN computation requirements, such as sampling based methods which are used to reduce memory consumption (Zeng et al., 2020), or weight pruning (Blalock et al., 2020) approaches to achieve further model compression.
E XPERIMENTS
In this section we ﬁrst analyse how the choice of quantization implementation affects performance of
GNNs.
We subsequently evaluate Degree-Quant against the strong baselines of: FP32, INT8-QAT and, INT8-QAT with stochastic masking of weights (Fan et al., 2020).
We refer to this last approach as noisy QAT or nQAT.
To make explicit that we are quantizing both weights and activations, we use the notation W8A8.
We repeat the experiments at INT4.
Our study evaluates performance on six datasets and includes both node-level and graph-level tasks.
The datasets used were Cora, CiteSeer, ZINC,
MNIST and CIFAR10 superpixels, and REDDIT-BINARY.
Across all datasets INT8 models trained with Degree-Quant manage to recover most of the accuracy lost as a result of quantization.
In some instances, DQ-INT8 outperform the extensively tuned FP32 baselines.
For INT4, DQ outperforms all
QAT baselines and results in double digits improvements over QAT-INT4 in some settings.
Details about each dataset and our experimental setup can be found in appendix A.1.
[Table]
4. I MPACT OF Q UANTIZATION G RADIENT E STIMATOR ON C ONVERGENCE
The STE is a workaround for when the forward pass contains non-differentiable operations (e
[Table]
g. rounding in QAT) that has been widely adopted in practice.
While the choice of STE implementation generally results in marginal differences for CNNs—even for binary networks (Alizadeh et al.
, 2019)—it is unclear whether only marginal differences will also be observed for GNNs.
Motivated by this, we study the impact of four off-the-shelve quantization procedures on the three architectures evaluated for each type of dataset; the implementation details of each one is described in appendix A.3.
We perform this experiment to ensure that we have the strongest possible QAT baselines.
Results are shown in table 1.
We found the choice quantization implementation to be highly dependent on the model architecture and type of problem to be solved: we see a much larger variance than is observed with CNNs; this is an important discovery for future work building on our study.
We observe a general trend in all INT4 experiments beneﬁting from momentum as it helps smoothing out the quantization statistics for the inherently noisy training stage at low bitwidths.
This trend applies as well for the majority of INT8 experiments, while exhibiting little impact on MNIST.
For INT8 Cora-GCN, large gradient norm values in the early stages of training (see ﬁg. 5) mean that these models not beneﬁt from momentum as quantization ranges fail to keep up with the rate of changes in tensor values; higher momentum can help but also leads to instability.
In contrast,
GAT has stable initial training dynamics, and hence obtains better results with momentum.
For the molecules dataset ZINC, we consistently obtained lower regression loss when using momentum.
We note that GIN models often suffer from higher performance degradation (as was ﬁrst noted in ﬁg. 3), specially at W4A4.
This is not the case however for image datasets using superpixels.
We believe that datasets with Gaussian-like node degree distributions (see ﬁg. 9) are more tolerant of the imprecision introduced by quantization, compared to datasets with tailed distributions.
We leave more in-depth analysis of how graph topology affects quantization as future work.
4. O BTAINING Q UANTIZATION BASELINES
Our FP32 results, which we obtain after extensive hyperparameter tuning, and those from the baselines are shown at the top of table 2.
We observed large gains on MNIST, CIFAR10 and, ZINC.
For our QAT-INT8 and QAT-INT4 baselines, we use the quantization conﬁgurations informed by our analysis in section 4.1.
For Citeseer we use the best resulting setup analysed for Cora, and for CIFAR10 that from MNIST.
Then, the hyperparameters for each experiment were ﬁne tuned individually, including noise rate n 2 [0 . 5 , 0 .
95] for nQAT experiments.
QAT-INT8 and QAT-INT4 results in table 2 and QAT-INT4, with the exception of MNIST (an easy to classify dataset), corroborate our hypothesis that GIN layers are less resilient to quantization.
This was ﬁrst observed in ﬁg. 3.
In the case of ZINC, while all models results in noticeable degradation, GIN sees a more severe 16% increase of regression loss compared to our FP32 baseline.
For QAT W4A4 an accuracy drop of over 35% and 47% is observed for Cora and Citeseer respectively.
The stochasticity induced by nQAT helped in recovering some of the accuracy lost as a result of quantization for citation networks (both
INT8 and INT4) but had little impact on other datasets and harmed performance in some cases.
[Table] [Table]
[Table]
4. C OMPARISONS OF D EGREE -Q UANT WITH E XISTING Q UANTIZATION A PPROACHES
Degree-Quant provides superior quantization for all GNN datasets and architectures.
Our results with DQ are highlighted in gray in table 2 and table 3.
Citation networks trained with DQ for W8A manage to recover most of the accuracy lost as a result of QAT and outperform most of nQAT baselines.
In some instances DQ-W8A8 models outperform the reference FP32 baselines.
At 4-bits,
DQ results in even larger gains compared to W4A4 baselines.
We see DQ being more effective for GIN layers, outperforming INT4 baselines for Cora (+ 24 . 9% ), Citeseer (+ 26 . 2% ) and REDDITBINARY (+ 23 . 0% ) by large margins.
Models trained with DQ at W4A4 for graph classiﬁcation and graph regression also exhibit large performance gains (of over 10% ) in most cases.
For ZINC, all models achieve over 20% lower regression loss.
Among the top performing models using DQ, ratios of p min and p max in [0 . 0 , 0 . 2] were the most common.
Figure 10 in the appendix shows validation loss curves for GIN models trained using different DQ probabilities on the REDDIT-BINARY dataset.
D ISCUSSION
Latency and Memory Implications .
In addition to offering signiﬁcantly lower memory usage ( 4 ⇥ with INT8), quantization can reduce latency—especially on CPUs.
We found that with INT arithmetic we could accelerate inference by up to 4 . 7 ⇥ .
We note that the latency beneﬁt depends on the graph topology and feature dimension, therefore we ran benchmarks on a variety of graph datasets, including Reddit 2 , Zinc, Cora, Citeseer, and CIFAR-10; Zinc and Reddit results are shown in table 4, with further results given in the appendix.
For a GCN layer with in- and out-dimension of 128, we get speed-ups of: 4 . 3 ⇥ on Reddit, 2 . 5 ⇥ on Zinc, 1 . 3 ⇥ on Cora, 1 . 3 ⇥ on Citeseer and, 2 . 1 ⇥ on CIFAR-10.
It is also worth emphasizing that quantized networks are necessary to efﬁciently use accelerators deployed in smartphones and smaller devices as they primarily accelerate integer arithmetic, and that CPUs remain a common choice for model serving on servers.
The decrease in latency on CPUs is due to improved cache performance for the sparse operations; GPUs, however, see less beneﬁt due to their massively-parallel nature which relies on mechanisms other than caching to hide slow random memory accesses, which are unavoidable in this application.
Ablation Study:
Beneﬁts of Percentile Ranges .
Figure 5 shows the value of percentiles during training.
We see that when using absolute min/max the upper range grows to over double the range required for 99.9% of values, effectively halving the resolution of the quantized values.
DQ is more stable, and we obtained strong results with an order of magnitude less tuning relative to the baselines.
Ablation Study: Source of Degradation at INT4 .
Figure 6 assesses how INT8 GAT (without DQ) degrades as single elements are converted to INT4, in order to understand the precipitous drop in accuracy in the INT4 baselines; further plots for GCN and GIN are included in the appendix.
We observe that most elements cause only modest performance losses relative to a full INT8 model.
DQ is most important to apply to elements which are constrained by numerical precision , such as the aggregation and message elements in GAT.
Weight elements, however, are consistently unaffected.
Ablation Study: Effect of Stochastic Element in Degree-Quant .
We observe that the stochastic protective masking in DQ alone often achieves most of the performance gain over the QAT baseline; results are given in table 9 in the appendix.
The beneﬁt of the percentile-based quantization ranges is stability , although it can yield some performance gains.
The full DQ method provides consistently good results on all architectures and datasets, without requiring an extensive analysis as in 4.1.
C ONCLUSION
This work has presented Degree-Quant, an architecture-agnostic and stable method for training quantized GNN models that can be accelerated using off-the-shelf hardware.
With 4-bit weights and activations we achieve 8 ⇥ compression while surpassing strong baselines by margins regularly exceeding 20%.
At 8-bits, models trained with DQ perform on par or better than the baselines while achieving up to 4.7 ⇥ lower latency than FP32 models.
Our work offers a comprehensive foundation for future work in this area and is a ﬁrst step towards enabling GNNs to be deployed more widely, including to resource constrained devices such as smartphones.