Degree-Quant: Quantization-Aware Training for Graph Neural Networks

Shyam Anil Tailor, Javier Fernandez-Marques, Nicholas Donald Lane

Keywords: Graph neural networks, quantization, benchmark

Abstract: 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 efficient 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 quantization-aware training baselines commonly used on other architectures, such as CNNs. We validate our method on six datasets and show, unlike previous quantization attempts, that models generalize to unseen graphs. Models trained with Degree-Quant for INT8 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.7x speedups on CPU when using INT8 arithmetic.

One-sentence Summary: We provide a training technique that enables graph neural networks to use low precision integer arithmetic at inference time, yielding up to 4.7x latency improvements on CPU

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics

Data: [Reddit](https://paperswithcode.com/dataset/reddit)

Community Implementations: [ 1 code implementation](https://www.catalyzex.com/paper/arxiv:2008.05000/code)