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

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. We identify the sources of error that uniquely arise when attempting to quantize GNNs, and propose an architecturally-agnostic 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 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.

arXiv commentPublished as a conference paper at ICLR 2021

References in corpus (17)

• Inductive Representation Learning on Large Graphs
• Neural Message Passing for Quantum Chemistry
• Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding
• Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation
• BinaryConnect: Training Deep Neural Networks with binary weights during propagations
• Fast Graph Representation Learning with PyTorch Geometric
• Graph Convolutional Matrix Completion
• Mixed Precision Training
• Quantizing deep convolutional networks for efficient inference: A whitepaper
• DropEdge: Towards Deep Graph Convolutional Networks on Node Classification
• Bayesian Compression for Deep Learning
• Benchmarking Graph Neural Networks
• Integer Quantization for Deep Learning Inference: Principles and Empirical Evaluation
• A Study of BFLOAT16 for Deep Learning Training
• Robust Quantization: One Model to Rule Them All
• Learning Accurate Low-Bit Deep Neural Networks with Stochastic Quantization
• Gradient ${\ell }_1$ Regularization for Quantization Robustness

Cited by in corpus (2)

• The Evolution of Distributed Systems for Graph Neural Networks and their Origin in Graph Processing and Deep Learning: A Survey
• VQ-GNN: A Universal Framework to Scale up Graph Neural Networks using Vector Quantization