Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Efficient Mixed Precision Quantization in Graph Neural Networks

Samir Moustafa, Nils M. Kriege, Wilfried N. Gansterer

TL;DR

This work addresses the high computational cost of inference in Graph Neural Networks by proposing MixQ-GNN, a mixed-precision quantization framework guided by a theorem that enables efficient, integer-based message passing with FP32-equivalent results. It introduces a relaxed, differentiable architecture search over per-component bit-widths and a memory-oriented penalty to balance accuracy and efficiency, enabling scalable selection of bit-widths across all GNN components. Empirical results across node- and graph-level tasks show MixQ-GNN achieving substantial BitOPs reductions (averaging around 5.5x for node tasks and 5.1x for graph tasks) while maintaining or improving predictive performance on several datasets, and demonstrating compatibility with existing quantization approaches like DQ. The framework proves robust across small to large graphs, including synthetic CSL data and large-scale real-world benchmarks, highlighting its practical impact for efficient GNN deployment on diverse hardware.

Abstract

Graph Neural Networks (GNNs) have become essential for handling large-scale graph applications. However, the computational demands of GNNs necessitate the development of efficient methods to accelerate inference. Mixed precision quantization emerges as a promising solution to enhance the efficiency of GNN architectures without compromising prediction performance. Compared to conventional deep learning architectures, GNN layers contain a wider set of components that can be quantized, including message passing functions, aggregation functions, update functions, the inputs, learnable parameters, and outputs of these functions. In this paper, we introduce a theorem for efficient quantized message passing to aggregate integer messages. It guarantees numerical equality of the aggregated messages using integer values with respect to those obtained with full (FP32) precision. Based on this theorem, we introduce the Mixed Precision Quantization for GNN (MixQ-GNN) framework, which flexibly selects effective integer bit-widths for all components within GNN layers. Our approach systematically navigates the wide set of possible bit-width combinations, addressing the challenge of optimizing efficiency while aiming at maintaining comparable prediction performance. MixQ-GNN integrates with existing GNN quantization methods, utilizing their graph structure advantages to achieve higher prediction performance. On average, MixQ-GNN achieved reductions in bit operations of 5.5x for node classification and 5.1x for graph classification compared to architectures represented in FP32 precision.

Efficient Mixed Precision Quantization in Graph Neural Networks

TL;DR

This work addresses the high computational cost of inference in Graph Neural Networks by proposing MixQ-GNN, a mixed-precision quantization framework guided by a theorem that enables efficient, integer-based message passing with FP32-equivalent results. It introduces a relaxed, differentiable architecture search over per-component bit-widths and a memory-oriented penalty to balance accuracy and efficiency, enabling scalable selection of bit-widths across all GNN components. Empirical results across node- and graph-level tasks show MixQ-GNN achieving substantial BitOPs reductions (averaging around 5.5x for node tasks and 5.1x for graph tasks) while maintaining or improving predictive performance on several datasets, and demonstrating compatibility with existing quantization approaches like DQ. The framework proves robust across small to large graphs, including synthetic CSL data and large-scale real-world benchmarks, highlighting its practical impact for efficient GNN deployment on diverse hardware.

Abstract

Graph Neural Networks (GNNs) have become essential for handling large-scale graph applications. However, the computational demands of GNNs necessitate the development of efficient methods to accelerate inference. Mixed precision quantization emerges as a promising solution to enhance the efficiency of GNN architectures without compromising prediction performance. Compared to conventional deep learning architectures, GNN layers contain a wider set of components that can be quantized, including message passing functions, aggregation functions, update functions, the inputs, learnable parameters, and outputs of these functions. In this paper, we introduce a theorem for efficient quantized message passing to aggregate integer messages. It guarantees numerical equality of the aggregated messages using integer values with respect to those obtained with full (FP32) precision. Based on this theorem, we introduce the Mixed Precision Quantization for GNN (MixQ-GNN) framework, which flexibly selects effective integer bit-widths for all components within GNN layers. Our approach systematically navigates the wide set of possible bit-width combinations, addressing the challenge of optimizing efficiency while aiming at maintaining comparable prediction performance. MixQ-GNN integrates with existing GNN quantization methods, utilizing their graph structure advantages to achieve higher prediction performance. On average, MixQ-GNN achieved reductions in bit operations of 5.5x for node classification and 5.1x for graph classification compared to architectures represented in FP32 precision.
Paper Structure (19 sections, 1 theorem, 10 equations, 9 figures, 11 tables, 1 algorithm)

This paper contains 19 sections, 1 theorem, 10 equations, 9 figures, 11 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Related Work
  4. Methodology
  5. Relaxed Message Passing for Mixed Precision Quantization
  6. MixQ-GNN Bit-Width Selection
  7. Experimental Evaluation
  8. Efficiency Metric
  9. Hardware and System Configuration
  10. Node level tasks
  11. MixQ-GNN versus A$^2$Q
  12. GraphSAGE Case Study
  13. Large-Scale Experiments
  14. Graph-level Tasks
  15. Synthetic Dataset
  16. ...and 4 more sections

Key Result

Theorem 1

Given the adjacency matrix $A \in \mathbb{R}^{n \times n}$ and the features matrix $X \in \mathbb{R}^{n \times f}$, with their quantized forms $Q_a(A)$ and $Q_x(X)$ respectively, and quantization parameter vectors $\{S_{a}, Z_{a}\}$, $\{S_{x}, Z_{x}\}$, $\{S_{y}, Z_{y}\}$, the quantized product $Q_y where $C_{1}$, $C_{2}$, and $C_{3}$ are either constants or involve operations with vectors, both o

Figures (9)

  • Figure 1: Accuracies of six types of GNNs on the Cora dataset over the number of operations required for a single forward pass (inference). Each architecture uses 1 to 5 layers of the same type, and each was tested five times with the best hyper-parameter settings. Colors distinguish the layer types, and the radius of the circles indicates the model size.
  • Figure 2: Accuracy vs. average bit-width for a two layer GCN architecture on the Cora dataset, using mixed precision quantization across 9 components within the architecture. Bit-width options: $\{2, 4, 8\}$, totaling $3^9$ possible combinations, each point $\odot$ is an instance of this combination.
  • Figure 3: Histograms of the bit-widths for each component in the two layer GCN architecture on the Pareto front of Figure \ref{['figure:all_options_for_gcn_over_cora']}. The 16 different choices along the Pareto front line are considered the optimal bit-width combinations.
  • Figure 4: An example of unfolding tree quantization using DQ, A$^2$Q, and MixQ-GNN methods. Each color represents a different bit-width (Bit$_i$) for node or edge weights.
  • Figure 5: An overview of a quantization example of one-layer message passing: (a) Complete graph $\mathcal{G}$ with three nodes and three edges. (b) Unfolding tree with depth one. (c) Simulated quantization for the unfolding tree with three quantizers: $Q_{x}$, $Q_{a}$, and $Q_{y}$. (d) Quantized inference for the unfolding tree with dequantizer as output.
  • ...and 4 more figures

Theorems & Definitions (1)

  • Theorem 1: Quantized Message Passing Schema