$\rm A^2Q$: Aggregation-Aware Quantization for Graph Neural Networks
As graph data size increases, the vast latency and memory consumption during
inference pose a significant challenge to the real-world deployment of Graph
Neural Networks (GNNs). While quantization is a powerful approach to reducing
GNNs complexity, most previous works on GNNs quantization fail to exploit the
unique characteristics of GNNs, suffering from severe accuracy degradation.
Through an in-depth analysis of the topology of GNNs, we observe that the
topology of the graph leads to significant differences between nodes, and most
of the nodes in a graph appear to have a small aggregation value. Motivated by
this, in this paper, we propose the Aggregation-Aware mixed-precision
Quantization ($\rm A^2Q$) for GNNs, where an appropriate bitwidth is
automatically learned and assigned to each node in the graph. To mitigate the
vanishing gradient problem caused by sparse connections between nodes, we
propose a Local Gradient method to serve the quantization error of the node
features as the supervision during training. We also develop a Nearest Neighbor
Strategy to deal with the generalization on unseen graphs. Extensive
experiments on eight public node-level and graph-level datasets demonstrate the
generality and robustness of our proposed method. Compared to the FP32 models,
our method can achieve up to a 18.6x (i.e., 1.70bit) compression ratio with
negligible accuracy degradation. Morever, compared to the state-of-the-art
quantization method, our method can achieve up to 11.4\% and 9.5\% accuracy
improvements on the node-level and graph-level tasks, respectively, and up to
2x speedup on a dedicated hardware accelerator.
Zeyu Zhu, Fanrong Li, Zitao Mo, Qinghao Hu, Gang Li, Zejian Liu, Xiaoyao Liang, Jian Cheng
# A\({}^{2}\)Q: Aggregation-Aware Quantization for Graph Neural Networks

###### Abstract

As graph data size increases, the vast latency and memory consumption during inference pose a significant challenge to the real-world deployment of Graph Neural Networks (GNNs). While quantization is a powerful approach to reducing GNNs complexity, most previous works on GNNs quantization fail to exploit the unique characteristics of GNNs, suffering from severe accuracy degradation. Through an in-depth analysis of the topology of GNNs, we observe that the topology of the graph leads to significant differences between nodes, and most of the nodes in a graph appear to have a small aggregation value. Motivated by this, in this paper, we propose the Aggregation-Aware mixed-precision Quantization (A\({}^{2}\)Q) for GNNs, where an appropriate bitwidth is automatically learned and assigned to each node in the graph. To mitigate the vanishing gradient problem caused by sparse connections between nodes, we propose a Local Gradient method to serve the quantization error of the node features as the supervision during training. We also develop a Nearest Neighbor Strategy to deal with the generalization on unseen graphs. Extensive experiments on eight public node-level and graph-level datasets demonstrate the generality and robustness of our proposed method. Compared to the FP32 models, our method can achieve up to a 18.6x (i.e., 1.70bit) compression ratio with negligible accuracy degradation. Moreover, compared to the state-of-the-art quantization method, our method can achieve up to 11.4% and 9.5% accuracy improvements on the node-level and graph-level tasks, respectively, and up to 2x speedup on a dedicated hardware accelerator.

## 1 Introduction

Recently, Graph Neural Networks (GNNs) have attracted much attention due to their superior learning and representing ability for non-Euclidean geometric data. A number of GNNs have been widely used in real-world applications, such as recommendation system (Jin et al., 2020), and social network analysis (Lerer et al., 2019), etc. Many of these tasks put forward high requirements for low-latency inference. However, the real-world graphs are often extremely large and irregular, such as Reddit with 232,965 nodes, which needs 19G floating-point operations (FLOPs) to be processed by a 2-layer Graph Convolutional Network (GCN) with only 81KB parameters (Tailor et al., 2020), while ResNet-50, a 50-layer DNN, only takes 8G FLOPs to process an image (Canziani et al., 2016). What is worse, it requires a huge amount of memory access for GNNs inference, e.g., the nodes features size of Reddit is up to 534MB, leading to high latency. Therefore, the aforementioned problems pose a challenge to realize efficient inference of GNNs.

Neural network quantization can reduce the model size and accelerate inference without modifying the model architecture, which has become a promising method to solve this problem in re
cent years. Unfortunately, there remain some issues in the existing works on GNNs quantization. Feng et al. (2020) only quantizes the node feature and keeps floating point calculations during inference. Tailor et al. (2020) proposes a degree-quant training strategy to quantize GNNs to the low-bit fixed point but causes a large accuracy drop, e.g., 11.1% accuracy drops when quantizing to 4bits. Moreover, some works (Wang et al., 2021b; Bahri et al., 2021; Wang et al., 2021a; Jing et al., 2021) quantize GNNs into 1-bit and compute with XNOR and bit count operations. However, these 1-bit quantization methods are either restricted to the node-level tasks or can not generalize well to other GNNs.

Most of the above methods do not make full use of the property of GNNs and graph data, resulting in severe accuracy degradation or poor generalization. As presented in MPNN framework (Gilmer et al., 2017), GNNs processing is divided into two phase: First, in the aggregation phase, a node collects information from neighboring nodes and uses the aggregation function to generate hidden features; second, in the update phase, the hidden features are transformed into new features by an update function. We analyze the nodes features after aggregation in Figure 1 and find that the higher the in-degree is, the larger the node features tend to be after aggregation. And the features vary significantly between nodes with different in-degrees, which represent the topology of a graph. Moreover, according to Xie et al. (2014); Aiello et al. (2001), the degrees of nodes in most real-world graph data often follow the power-law distribution, i.e., nodes with a low degree account for the majority of graph data. Therefore, specially quantizing the nodes features according to the topology of the graphs will be beneficial to reduce the quantization error while achieving a higher compression ratio.

In this paper, we propose the **Aggregation-Aware Quantization** (\(\rm A^{2}Q\)) method, which quantizes different nodes features with different learnable quantization parameters, including bitwidth and step size. These parameters can be adaptively learned during training and are constrained by a penalty on memory size to improve the compression ratio. However, when quantizing the model in semi-supervised tasks, the gradients for most quantization parameters are zero due to the sparse connections between nodes, which makes the training non-trivial. We propose the **Local Gradient** method to solve this problem by introducing quantization error as supervised information. Finally, to generalize our method to unseen graphs in which the number of the nodes varies, we develop the **Nearest Neighbor Strategy** which assigns the learned quantization parameters to the unseen graph nodes. To the best of our knowledge, we are the first to introduce the mixed-precision quantization to the GNNs. Compared with the previous works, our proposed methods can significantly compress GNNs with negligible accuracy drop.

In summary, the key contributions of this paper are as follows:

1. We propose the Aggregation-Aware mixed-precision Quantization (\(\rm A^{2}Q\)) method to enable an adaptive learning of quantization parameters. Our learning method is powerful by fully

Figure 1: The analysis of the average aggregated node features in different in-degrees node groups on various tasks. (a) The values at the final layer for GNNs trained on Cora. (b) The values at the 2-5 layer of GIN trained on REDDIT-BINARY. The average values are all generated from 10 runs.
utilizing the characteristic of GNNs, and the learned bitwidth is strongly related to the topology of the graph.
2. A **Local Gradient** method is proposed to train the quantization parameters in semi-supervised learning tasks. Furthermore, to generalize our method to the unseen graphs in which the number of input nodes is variable, we develop the **Nearest Neighbor Strategy** to select quantization parameters for the nodes of the unseen graphs.
3. Experiments demonstrate that we can achieve a compression ratio up to 18.6x with negligible accuracy degradation compared to the full-precision (FP32) models. Moreover, the model trained with our \(\mathrm{A^{2}Q}\) method outperforms the state-of-the-art (SOTA) method up to 11.4% with a speedup up to 2.00x in semi-supervised tasks, and obtains up to 9.5% gains with a 1.16x speedup in graph-level tasks. We provide our code at this URL: [https://github.com/weihai-98/](https://github.com/weihai-98/)\(\mathrm{A^{2}Q}\).

## 2 Related Work

_Graph Neural Networks:_ The concept of the graph neural network was first proposed in Scarselli et al. (2008), which attempted to generalize neural networks to model non-Euclidean data. In the following years, various GNN models were proposed. For example, Graph Convolution Network (GCN) (Kipf and Welling, 2016) uses a layer-wise propagation rule that is based on a first-order approximation of spectral convolutions on graphs, Graph Isomorphism Network (GIN) (Xu et al., 2018) designed a provably maximally powerful GNN under the MPNN framework, and Graph Attention Network (GAT) (Velickovic et al., 2017) introduces the attention mechanism to graph processing. Although GNNs have encouraging performance in a wide range of domains (Jin et al., 2020; Yang, 2019), the huge amount of float-point operations and memory access in process pose a challenge to efficient inference, which hinder the applications of GNNs.

_Quantized GNNs:_ As a promising method to reduce the model size and accelerate the inference process, quantization is also applied to GNNs. Some works quantize features and weights in GNNs to low bitwidths (Feng et al., 2020; Tailor et al., 2020) or even 1-bit (Wang et al., 2021; Bahri et al., 2021; Wang et al., 2021; Jing et al., 2021), i.e., use fixed-point numbers instead of floating-point numbers for computation. But when the compression ratio is high (e.g., \(<\)4bit), the performance degradation of these works is significant, and the generalization of 1-bit method is limited. There are also some works on vector quantization (VQ), which use the vectors in a codebook obtained during the training process instead of the original features (Ding et al., 2021; Huang et al., 2022). However, searching for vectors in the codebook is computationally complex.

_Mixed-Precision Quantization:_ Based on the idea that different layers have different sensitivities to quantization, mixed-precision quantization is proposed in CNNs to quantize different layers to different bitwidths for better model compression. Early works (Wang et al., 2019; Lou et al., 2019) proposed reinforcement learning (RL) based methods to search bitwidth for different layers, but they often require large computational resources, which limits the exploration of the search space. Another important class of mixed-precision method is the criteria-based method, they use the specific criteria to represent the quantization sensitivity, e.g., (Dong et al., 2019; Dong et al., 2020; Chen et al., 2021)quantize different layers with different bitwidths based on the trace of the Hessian. Recently, there are some other methods to learn the bitwidth during training (Uhlich et al., 2019; Esser et al., 2019; Jain et al., 2020). However, due to the huge difference between GNNs and CNNs, it is difficult to use these methods on GNNs directly, and our \(\mathrm{A^{2}Q}\) is the first method to introduce the mixed-precision quantization to GNNs, further improving the inference efficiency of GNNs.

## 3 Method

In this section, we describe our proposed Aggregation-Aware Quantization in detail. Firstly, we present the formulation of the mixed-precision quantization for GNNs, which fully utilizes the property of GNNs and graph data. Secondly, we introduce the Local Gradient method to address the gradient vanishing problem during training. Finally, we detail the Nearest Neighbor Strategy, which is used for generalizing our approach to the unseen graphs.
### Aggregation-Aware Quantization

We assume a graph data with \(N\) nodes and the node features are \(F\)-dimensional, i.e., the feature map is \(\mathbf{X}\in\mathbb{R}^{N\times F}\) and \(\mathbf{x}_{i}\) is the features of node \(i\). We use the learnable parameters step size \(\alpha_{i}\in\mathbb{R}_{+}\) and bitwidth \(b_{i}\in\mathbb{R}_{+}\) to quantize the features of the \(i\)-th node as:

\[\bar{\mathbf{x}}_{i}=sign(\mathbf{x}_{i})\left\{\begin{aligned} \lfloor\frac{|\mathbf{x}_{i}|}{\alpha_{i}}+0.5\rfloor,&|\mathbf{x}|< \alpha_{i}(2^{[b_{i}]-1}-1)\\ 2^{[b_{i}]-1}-1,&|\mathbf{x}_{i}|\geq\alpha_{i}(2^{[b_ {i}]-1}-1)\end{aligned}\right., \tag{1}\]

where \(\lfloor\cdot\rfloor\) is the floor function, and \([\cdot]\) is the round function to ensure the bitwidth used to quantize is an integer. The learnable parameters are \(\mathbf{s}_{X}=(\alpha_{1},\alpha_{2},...,\alpha_{N})\), and \(\mathbf{b}_{X}=(b_{1},b_{2},...,b_{N})\). Then we can obtain the fixed-point feature map \(\bar{\mathbf{X}}\), and the original feature can be represented as \(\mathbf{X}_{q}=\mathbf{S}_{X}\cdot\bar{\mathbf{X}}\), where \(\mathbf{S}_{X}=diag(\alpha_{1},\alpha_{2},...,\alpha_{N})\). Note that we use \([b]+1\) as the quantization bitwidth for the features after ReLU because the values are all non-negative.

In the update phase, the node features are often transformed with a linear mapping or an MLP in which matrix multiplication \(\mathbf{X}\mathbf{W}\) is the main computation, and the transformed node features are the input to the next layer in GNNs. In order to accelerate the update phase, we also quantize \(\mathbf{W}\). Due to the fact that \(\mathbf{W}\) in a certain layer is shared by all nodes, we quantize \(\mathbf{W}\) to the same bitwidth of 4bits for all GNNs in this paper. However, each column of \(\mathbf{W}\) has its learnable quantization step size, i.e., \(\mathbf{s}_{W}=(\beta_{1},\beta_{2},..,\beta_{F_{2}})\), where \(F_{2}\) is the output-dimension of the node features in current layer and \(\beta_{i}\) is the quantization step size for the \(i\)-th column of \(\mathbf{W}\), and we also use Eq. 1 to quantize \(\mathbf{W}\). We can obtain the integer representation \(\bar{\mathbf{W}}\) and the quantized representation \(\bar{\mathbf{W}}_{q}=\bar{\mathbf{W}}\cdot\mathbf{S}_{W}\), where \(\mathbf{S}_{W}=diag(\beta_{1},\beta_{2},...,\beta_{F_{2}})\). The float-point matrix multiplication in the update phase can be reformulated as follow:

\[\mathbf{X}\cdot\mathbf{W}\approx\mathbf{X}_{q}\cdot\mathbf{W}_{q}=(\mathbf{S}_{X}\cdot\bar{\mathbf{X} })\cdot(\bar{\mathbf{W}}\cdot\mathbf{S}_{W})=(\bar{\mathbf{X}}\cdot\bar{\mathbf{W}})\odot( \mathbf{s}_{X}\otimes\mathbf{s}_{W})\;, \tag{2}\]

where \(\odot\) denotes an element-wise multiplication, and \(\otimes\) denotes the outer product. After training, we can obtain \(\mathbf{s}_{X}\) and \(\mathbf{s}_{W}\) so that the outer product can be pre-processed before inference. An example is illustrated in Figure 2. For the aggregation phase, i.e., \(\mathbf{AX}\), \(\mathbf{A}\) is the adjacency matrix and \(\mathbf{A}\in\{0,1\}^{N\times N}\), we quantize the \(\mathbf{X}\) as the quantization way of \(\mathbf{W}\) because the nodes features involved in the aggregation process come from the update phase, in which the features lose the topology information of graphs. Then the aggregation phase can be performed by integer operations to reduce the computational overhead.

The quantization parameters \((s,b)\) are trained by the backpropagation algorithm. Since the floor and round functions used in the quantization process are not differentiable, we use the straight-through estimator (Bengio et al., 2013) to approximate the gradient through these functions, and the gradients of the quantization parameters can be calculated by:

\[\frac{\partial L}{\partial s}=\sum_{i=1}^{d}\frac{\partial L}{\partial x_{q}^ {i}}\cdot\frac{\partial x_{q}^{i}}{\partial s}\;, \tag{3}\]

where \(d\) is the dimension of the vector \(\mathbf{x}\), \((s,b)\) are the quantization parameters for \(\mathbf{x}\), and \(x_{q}^{i}\) is the value of \(i\)-th dimension in \(\mathbf{x}_{q}\). Detailed information about quantization process and the backpropagation are shown in Appendix A.1 and A.3 **Proof 2 and 3**.
In order to improve the compression ratio of the node features, we introduce a penalty term on the memory size:

\[L_{memory}=(\frac{1}{\eta}\cdot\sum_{l=1}^{L}\sum_{i=1}^{N}\mathrm{dim}^{l}\cdot b _{i}^{l}-M_{target})^{2}\;, \tag{5}\]

where \(L\) is the number of layers in the GNNs, \(N\) is the total number of nodes, \(\mathrm{dim}^{l}\) is the length of the node features in \(l\)-th layer, \(b_{i}^{l}\) is the quantization bitwidth for node \(i\) in \(l\)-th layer, \(M_{target}\) is the target memory size on the total node features memory size, and \(\eta=8*1024\), which is a constant to convert the unit of memory size to \(\mathrm{KB}\). Then the model and quantization parameters can be trained by the loss function:

\[L_{total}=L_{task}+\lambda\cdot L_{memory}\;, \tag{6}\]

where \(L_{task}\) is the task-related loss function and \(\lambda\) is a penalty factor on \(L_{memory}\).

### Local Gradient

Although the above end-to-end learning method is concise and straightforward, the gradients for the quantization parameters of nodes features, i.e., \(\frac{\partial L_{task}}{\partial s}\) and \(\frac{\partial L_{task}}{\partial b}\), are almost zero during the training process of semi-supervised tasks, which poses a significant challenge to train the quantization parameters for nodes features. We analyze the property of GNNs and graph data, and find that two reasons lead to this phenomenon: 1. The extreme sparsity of the connections between nodes in graph data. 2. Only a tiny fraction of nodes with labels are used for training in semi-supervised tasks (e.g., \(0.30\%\) in PubMed dataset). Therefore, \(\frac{\partial L_{task}}{\partial x_{q}}\) for most node features are zero (detailed proof in Appendix A.3.2), which results in that the gradients for quantization parameters of these nodes vanish according to Eq. 3 and Eq. 4. To clarify, we visualize the \(\frac{\partial L_{task}}{\partial x_{q}}\) in the second layer of GCN trained on Cora. As shown in Figure 3, most gradients for the nodes features are zero.

The gradients of the \(L_{task}\) w.r.t. quantized nodes features can be viewed as the supervised information from the labeled nodes which enable the training of the quantization parameters for nodes features. However, this supervised information is missing due to zero gradients. Considering the quantization error is related to the \(L_{task}\), we introduce the quantization error \(E=\frac{1}{d}\left|\mathbf{x}_{q}-\mathbf{x}\right|_{1}\) as the supervised information for the quantization parameters of nodes features, where \(\mathbf{x}\) is the features before quantization, \(\mathbf{x}_{q}\) is the features after quantization and \(|\cdot|_{1}\) denotes the L1 norm. We refer to this method as **Local Gradient** because the gradients are computed by the local quantization errors instead of back-propagated task-related gradients. Then the quantization parameters for node features can be trained by gradients from \(E\):

\[\frac{\partial E}{\partial s}=\frac{1}{d}\sum_{i=1}^{d}sign(x_{q}^{i}-x^{i}) \cdot\frac{\partial x_{q}^{i}}{\partial s}\;, \tag{7}\]

\[\frac{\partial E}{\partial b}=\frac{1}{d}\sum_{i=1}^{d}sign(x_{q}^{i}-x^{i}) \cdot\frac{\partial x_{q}^{i}}{\partial b}\;. \tag{8}\]

Note that the quantization parameters of \(\mathbf{W}\) are still trained by utilizing the gradients in Eq. 3.

### Nearest Neighbor Strategy

In graph-level tasks, the quantized GNNs are required to generalize to unseen graphs. In such a scenario, the number of input nodes may vary during training or inference. However, the learnable method can only train a fixed number of \((s,b)\) pairs which are the same as the number of input nodes, so it is challenging to learn the \(s\) and \(b\) for every node in graph-level tasks. To solve this problem, we propose the **Nearest Neighbor Strategy**, which allows learning of a fixed number of quantization parameters and select quantization parameters for the unseen graphs.

The proposed strategy is shown in Algorithm 1. To ensure the numerical range of \(\mathbf{x}_{q}\) is as close as to \(\mathbf{x}\) at FP32, a simple way is to keep the maximum quantization value equal to the maximum absolute value of \(\mathbf{x}\). Based on this idea, we first initialize \(m\) groups of quantization parameters, then we calculate the maximum quantization value for every group, i.e., \(q_{max}=s(2^{|b|-1}-1)\). When quantizing the features of node \(i\), the feature with the largest absolute value \(f_{i}\) in the node features \(\mathbf{x}_{i}\) is first selected, and then we find the nearest \(q_{max}\) and quantize the node features with the \((s,b)\) corresponding to this \(q_{max}\). When performing backpropagation, we first calculate the gradients of the loss function w.r.t. quantization parameters according to Eq. 3 and Eq. 4. For a specific set of quantization parameters \((s_{j},b_{j})\), we collect the gradients from the nodes that have used them
and add these gradients together. After the model has been trained, we obtain the quantization parameters \((\mathbf{s},\mathbf{b})\). Since \(q_{max}\) can be calculated and sorted in advance, searching the nearest \(q_{max}\) can be implemented by binary searching. Usually, we set \(m=1000\) for all graph-level tasks in our paper and the overhead introduced to inference time is negligible.

## 4 Experiments

### Experimental Settings

In this section, we evaluate our method on three typical GNN models, i.e., GCN, GIN, and GAT. And we compare our method with the FP32 GNN model and DQ-INT4 (Tailor et al., 2020) on eight datasets, including four node-level semi-learning tasks (Cora, CiteSeer, PubMed, ogbn-arxiv) (Hu et al., 2020; Yang et al., 2016) and four graph-level tasks (REDDIT-BINARY, MNIST, CIFAR10, ZINC) (Yanardag & Vishwanathan, 2015; Dwivedi et al., 2020), to demonstrate the generality and robustness of our method. Among these datasets, ZINC is a dataset for regression tasks, which uses regression loss as the metric of the model performance, while others are all for classification tasks.

For a fair comparison, we set the quantization bitwidth of \(\mathbf{W}\) for all GNNs to 4bits as DQ-INT4. We count the average bitwidths for nodes features in all layers of the overall model and list them in our

\begin{table}
\begin{tabular}{c l l l l l} \hline \hline Dataset & Model & Accuracy & Average bits & Compression Ratio & Speedup \\ \hline \multirow{4}{*}{**Cora**} & GCN(FP32) & 81.5\(\pm\)0.7\% & 32 & 1x & — \\  & GCN(DQ ) & 78.3\(\pm\)1.7\% & 4 & 8x & 1x \\  & GCN(ours) & **80.9\(\pm\)0.6\%** & **1.70** & **18.6x** & **2.00x** \\ \cline{2-6}  & GAT(FP32) & 83.1\(\pm\)0.4\% & 32 & 1x & — \\  & GAT(DQ ) & 71.2\(\pm\)2.9\% & 4 & 8x & 1x \\  & GAT(ours) & **82.6\(\pm\)0.6\%** & **2.03** & **15.4x** & **1.49x** \\ \hline \multirow{4}{*}{**CiteSeer**} & GCN(FP32) & 71.1\(\pm\)0.7\% & 32 & 1x & — \\  & GCN(DQ ) & 66.9\(\pm\)2.4\% & 4 & 8x & 1x \\  & GCN(ours) & **70.6\(\pm\)1.1\%** & **1.87** & **17.0x** & **1.91x** \\ \cline{2-6}  & GIN(FP32) & 66.1\(\pm\)0.9\% & 32 & 1x & — \\  & GIN(DQ ) & 60.8\(\pm\)2.1\% & 4 & 8x & 1x \\  & GIN(ours) & **65.1\(\pm\)1.7\%** & **2.54** & **12.6x** & **1.37x** \\ \hline \multirow{3}{*}{**PubMed**} & GAT(FP32) & 79.0\(\pm\)0.3\% & 32 & 1x & — \\  & GAT(DQ) & 70.6\(\pm\)12.5\% & 4 & 8x & 1x \\  & GAT(ours) & **78.8\(\pm\)0.4\%** & **2.12** & **15.1x** & **1.38x** \\ \hline \multirow{3}{*}{**ogbn-arxiv**} & GCN(FP32) & 71.7\(\pm\)0.3\% & 32 & 1x & — \\  & GCN(DQ) & 65.4\(\pm\)3.9\% & 4 & 8x & 1x \\ \cline{1-1}  & GCN(ours) & **71.1\(\pm\)0.3\%** & **2.65** & **12.1x** & **1.28x** \\ \hline \hline \end{tabular}
\end{table}
Table 1: The results comparison on node-level tasks. The average bits are counted for each task when the best results are achieved.
results, denoted by "Average bits". Since today's CPUs and GPUs can not support mixed-precision operations well, we implement a precision-scalable hardware accelerator to perform the overall inference process for GNN. The accelerator employs massive bit-serial multipliers Judd et al. (2016), therefore, the latency of the integer multiplications is determined by the bitwidth of the node features. To evaluate the performance gains of our method over DQ-INT4, we develop a cycle-accurate simulator for our accelerator. More details about accelerator architecture are shown in Appendix A.7.5. Moreover, we show the compression ratio of quantized GNNs compared to the FP32 models in terms of overall memory size. For simplicity, we use GNN(DQ) to represent the GNNs quantized by DQ-INT4 and GNN-dataset to represent the task in which we run the experiment, e.g., GCN-Cora represents the GCN model trained on Cora. Detailed information about datasets and settings is in Appendix A.5 and Appendix A.6.

### Node-Level Tasks

Table 1 shows the experimental results on three GNN architectures trained on four node-level datasets. Compared with DQ-INT4, our method can achieve significantly better accuracy on each task, even with a higher compression ratio, improving the inference performance with 1.28x to 2.00x speedups. On almost all node-level tasks, our proposed \(\mathrm{A^{2}Q}\) has negligible accuracy drop compared to the FP32 baselines while achieving 12.1x-18.6x compression ratio. Since both GIN and GAT involve more complex computations, such as the calculation of attention coefficients in GAT, it is more challenging to quantize those models, and DQ performs poorly on these two models. However, our method can overcome this problem and maintain comparable accuracy compared with the FP32 models. Our method can outperform the DQ-INT4 by 11.4% on the GAT-Cora task with a smaller bitwidth (2.03 v.s. 4). Even on ogbn-arxiv, which has a large number of nodes, \(\mathrm{A^{2}Q}\) can achieve a 12.1x compression ratio compared with FP32 baseline with comparable accuracy, which demonstrates the robustness of our method. Moreover, to demonstrate the generality of our method, we also evaluate our method on heterogeneous graphs and the inductive learning tasks and compare with more related works in Appendix A.7.1.

### Graph-Level Tasks

Table 2 presents the comparison results on the graph-level tasks. Our method can obtain better results on all tasks than DQ-INT4 with higher compression and a considerable speedup. Especially on the GIN-REDDIT-BINARY task, our method outperforms DQ-INT4 by 9.5% while achieving a 1.16x

\begin{table}
\begin{tabular}{c l l l l l} \hline \hline Dataset & Model & Accuracy (Loss\(\downarrow\)) & Average bits & Compression ratio & Speedup \\ \hline \multirow{4}{*}{**MNIST**} & GCN(FP32) & 90.1\(\pm\)0.2\% & 32 & 1x & — \\  & GCN(DQ) & 84.4\(\pm\)1.3\% & 4 & 8x & 1x \\  & GCN(ours) & **89.9\(\pm\)0.8\%** & **3.50** & **9.12x** & **1.17x** \\ \cline{2-6}  & GIN(FP32) & 96.4\(\pm\)0.4\% & 32 & 1x & — \\  & GIN(DQ) & 95.5\(\pm\)0.4\% & 4 & 8x & 1x \\  & GIN(ours) & **95.7\(\pm\)0.2\%** & **3.75** & **8.52x** & **1.07x** \\ \hline \multirow{4}{*}{**CIFAR10**} & GCN(FP32) & 55.9\(\pm\)0.4\% & 32 & 1x & — \\  & GCN(DQ) & 51.1\(\pm\)10.7\% & 4 & 8x & 1x \\  & GCN(ours) & **52.5\(\pm\)0.8\%** & **3.32** & **9.62x** & **1.25x** \\ \cline{2-6}  & GAT(FP32) & 65.4\(\pm\)0.4\% & 32 & 1x & — \\  & GAT(DQ) & 56.5\(\pm\)0.6\% & 4 & 8x & 1x \\  & GAT(ours) & **64.7\(\pm\)2.8\%** & **3.73** & **8.57x** & **1.12x** \\ \hline \multirow{2}{*}{**ZINC**} & GCN(FP32) & 0.450\(\pm\)0.008 & 32 & 1x & — \\  & GCN(DQ) & 0.536\(\pm\)0.011 & 4 & 8x & 1x \\  & GCN(ours) & **0.49\(\pm\)0.05** & **3.68** & **8.68x** & **1.08x** \\ \hline \multirow{2}{*}{**REDDIT-BINARY**} & GIN(FP32) & 92.2\(\pm\)2.3\% & 32 & 1x & — \\  & GIN(DQ) & 81.3\(\pm\)4.4\% & 4 & 8x & 1x \\ \cline{1-1}  & GIN(ours) & **90.8\(\pm\)1.8\%** & **3.50** & **9.14x** & **1.16x** \\ \hline \hline \end{tabular}
\end{table}
Table 2: The results comparison on graph-level tasks.
speedup. Even for graph datasets with similar in-degrees, such as MNIST and CIFAR10, our method also learns the appropriate bitwidths for higher compression ratio and better accuracy. Although on GIN-MINT task, the improvement of our method is relatively small due to the similarity of the in-degrees between different nodes, our method can achieve comparable accuracy with smaller bitwidth (3.75 v.s. 4).

### Analysis

To understand why our approach works, we analyze the relationship between the learned bitwidths and the topology of the graph. Figure 4(a) and 4(b) reveal that the bitwidth learned by \(\mathrm{A}^{2}\mathrm{Q}\) is strongly related to the topology of graph data in the node-level tasks. As the bitwidth increases, the average in-degrees of nodes become larger. In other words, \(\mathrm{A}^{2}\mathrm{Q}\) method tends to learn higher bitwidth for nodes with higher in-degrees. However, in GAT, as shown in Figure 4(c), the learned bits are irregular. This is because the features aggregated in GAT are topology-free. However, our method can still learn appropriate quantization bitwidths for different nodes, which improves accuracy while reducing memory usage. In addition, Figure 4 also shows the node distribution for different bitwidths and the result is consistent with power-law distribution. Since nodes in graph data mainly have low in-degrees, most of the nodes are quantized to low bitwidth (\(\leq 4\)), compressing the GNNs as much as possible. And there are also some high in-degree nodes quantized to high bitwidth, which can help to maintain the accuracy of the GNN models. As a result, the average bitwidth of the entire graph features is low, and the accuracy degradation is negligible.

For the graph-level tasks in which the number of nodes varies, our method is also aggregation-aware. We select a layer of GIN trained on REDDIT-BINARY and analyze the relationship between bitwidth and average in-degrees of nodes using the corresponding bitwidth to quantize in Figure 4(d) and 4(e). It can be seen that the bitwidth learned for nodes features input to the second layer of MLP, which is the update function in GIN for graph-level tasks, does not present a correlation with the topology of graph. We analyze the reason and find that the node features before the second layer is the result mapped by the first layer of MLP and is activated by the activation function, e.g., ReLU, which results in the node features losing the topology information. We present more experiment results in Appendix A.7. to demonstrate that our method is generally applicable.

## 5 Ablation Study

**The advantage of learning-based mixed-precision quantization:** In Figure 5, we compare our \(\mathrm{A}^{2}\mathrm{Q}\) with the manual mixed-precision method, which manually assigns high-bit to those nodes with high in-degrees and low-bit to those nodes with low in-degrees. In the figure, the postfix "learn" denotes that using \(\mathrm{A}^{2}\mathrm{Q}\) method, "manual" denotes that we assign bits to nodes and the model only learns the stepsize, and "mixed-precision" denotes that the model uses the same quantization method as DQ-INT4 but assigning different bitwidths to nodes. For the "mixed-precision", we assign 5bits to those nodes with 50% top in-degrees and assign 3bits to others. The implications are two-fold. First, compared with the DQ-INT4, which uses the same quantization bitwidth, the mixed-precision

Figure 4: The relationship between quantized bitwidth and average in-degrees of nodes. (a), (b) and (c) represent the results of three GNN models trained on CiteSeer. (d) and (e) are results about the first and the second layer of an MLP, which is the update function of GIN trained on REDDIT-BINARY. The green bars represent the average in-degrees for the certain bitwidth used by nodes and the orange polylines represent the number of the nodes that use this certain bitwidth.
method obtains 1.1% gains on GCN-Cora tasks demonstrating that the mixed-precision method is more effective. Second, the results of the learning method outperform the manual method on all tasks. Especially for the models with a high compression ratio, on GIN-CiteSeer task, learning method can achieve 21.5% higher accuracy. This demonstrates that our learning method can perform better than the assignment method according to prior knowledge for mixed-precision quantization of GNNs.

**The power of learning the quantization parameters:** Ablations of two quantization parameters \((\mathbf{s},\mathbf{b})\) on the GIN-Cora task are reported in the first row of Table 3. The "no-lr" denotes that do not use learning method, "no-lr-b" denotes that only learn the step size \(\mathbf{s}\), "no-lr-s" denotes that only learn the bitwidths \(\mathbf{b}\), and "lr-all" denotes that learn the bitwidth and step size simultaneously. We can see that learning the step size can significantly increase the accuracy and even the "no-lr-bit" model can outperform the DQ-INT4 at the same compression ratio. When learning the bitwidth and step size simultaneously, the model can achieve higher accuracy with a higher compression ratio. This is because our method learns lower bitwidths for most nodes with low in-degrees and higher bitwidths for a tiny fraction of nodes with high in-degrees, which can improve the compression ratio while achieving higher accuracy.

**Local Gradient v.s. Global Gradient:** To demonstrate the effectiveness of our Local Gradient method, we compare the models trained with and without it on the GCN-CiteSeer task in the last row of Table 3. The "Global" denotes that the model is trained with Eq. 3 and Eq. 4. The model trained with the local method outperforms the global method by 13.8% with a higher compression ratio. This is because the Local Gradient method can learn quantization parameters for all nodes, while only quantization parameters for a part of nodes can be updated with the Global Gradient method due to the extreme sparse connection in the graph on the node-level semi-supervised tasks.

**The overhead of Nearest Neighbor Strategy:** We evaluate the real inference time of the GIN model on the 2080ti GPU. On REDDIT-BINARY task, the model without the selection process requires 121.45ms, while it takes 122.60ms for the model with our Nearest Neighbor Strategy, which only introduces 0.95% overhead. But with the help of the Nearest Neighbor Strategy, our model can obtain 19.3% accuracy gains for quantized GIN on REDDIT-BINARY.

## 6 Conclusion

This paper proposes A\({}^{2}\)Q, an aggregation-aware mixed-precision quantization method for GNNs, and introduces the Local Gradient and Nearest Neighbor Strategy to generalize A\({}^{2}\)Q to the node-level and graph-level tasks, respectively. Our method can learn the quantization parameters for different nodes by fully utilizing the property of GNNs and graph data. The model quantized by our A\({}^{2}\)Q can achieve up to a 18.6x compression ratio, and the accuracy degradation is negligible compared with the FP32 baseline. Compared with the prior SOTA, DQ-INT4, our method can significantly improve 11.4% accuracy with up to a 2.00x speedup on different tasks. Our work provides a general, robust and feasible solution to speed up the inference of GNNs.