PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Shyam A. Tailor, Javier Fernandez-Marques, Nicholas D. Lane

arXiv:2008.05000v3cs.LGstat.ML

TL;DR

GNN inference remains costly because computation is tied to graph size, while many deployment devices favor integer over floating-point arithmetic. This work analyzes quantization-specific degradation and introduces Degree-Quant, which supports strong low-precision performance, unseen-graph generalization, and faster CPU inference.

  • Problem

    GNN inference efficiency is limited by graph-size-dependent computation, motivating integer quantization for devices and serving environments where floating-point arithmetic is less efficient.

  • Method

    Degree-Quant is an architecture-agnostic quantization-aware training method that addresses degradation linked to STE implementation, node degree, and quantization-statistics tracking.

  • Results

    Across the evaluated settings, INT8 Degree-Quant models perform on par with or better than baselines, while INT4 models surpass strong baselines by margins regularly exceeding 20% and INT8 inference reaches up to 4.7× lower latency than FP32.

  • Takeaways & Limitations

    Degree-Quant provides a practical foundation for deploying quantized GNNs on off-the-shelf hardware and resource-constrained devices, including smartphones.

Abstract

from arXiv · show

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.

1 INTRODUCTION

GNNs model irregular data but remain expensive to run because inference compute scales with graph inputs, motivating integer quantization. The paper introduces Degree-Quant to address GNN-specific quantization problems and reports strong accuracy, generalization, and CPU efficiency.

  • Motivation: GNN inference compute remains coupled to input graph size despite relatively few model parameters.This limits deployment, including on devices that favor integer arithmetic over floating point.
  • Motivation: Integer quantization can reduce inference compute, memory, and energy without changing the model architecture.It is particularly relevant for smartphones, smaller devices, and data-center model serving.
  • Contributions: Degree-Quant explains how STE implementation, node degree, and quantization-statistics tracking affect low-precision GNN performance.The method is designed as an architecture-agnostic approach to quantization-aware training on graphs.
  • Contributions: INT8 Degree-Quant models often perform as well as FP32 counterparts, while INT4 models typically outperform quantized baselines by over 20%.The method also generalizes to unseen graphs, unlike previous attempts described by the authors.
  • Contributions: INT8 arithmetic delivers up to 4.7× CPU speedups relative to full-precision floating point, with 4-8× lower runtime memory usage.These results target more efficient inference without architectural modifications.

2 BACKGROUND

GNNs generalize message-passing ideas from CNNs to irregular graphs, where neighborhoods vary in size and ordering. Quantization-aware training simulates low-precision errors during training and uses straight-through estimators to optimize through rounding.

  • Message Passing Neural Networks: MPNN layers compute node representations through message passing, aggregation, and update phases.Messages are computed between nodes, aggregated with a permutation-invariant function, and used to update node features.
  • Message Passing Neural Networks: Unlike CNNs, GNN neighborhoods have variable topology, size, and ordering, although both architectures use weight sharing.This irregularity motivates treating GNNs as neural architectures for structured but non-grid data.
  • Message Passing Neural Networks: The paper considers GCN, GAT, and GIN architectures with different update rules involving degree normalization, attention coefficients, or learnable functions.These architectures represent distinct approaches within the MPNN paradigm.
  • Quantization: Integer quantization maps tensors into bounded discrete ranges using a scale and zero-point, then dequantizes them for approximate real-valued computation.The approximation becomes less accurate as bitwidth decreases.
  • Quantization: Quantization-aware training simulates quantization in the forward pass and uses a straight-through estimator to compute gradients through nondifferentiable operations.Gradient clipping, stochastic QAT, and layer re-ordering are additional techniques used to improve low-precision robustness.

3 QUANTIZATION FOR GNNS

Quantizing GNNs introduces aggregation and gradient-estimation errors that depend on node degree and quantization-range tracking. Degree-Quant addresses these issues with stochastic degree-based protection and percentile-tracked ranges while remaining architecture-agnostic.

  • Sources of Error: High in-degree nodes create larger aggregation errors that can distort weight-gradient updates through the GNN layer.The paper identifies aggregation as a substantial numerical-error source, especially for high in-degree nodes.
  • Sources of Error: Degree-dependent outliers distort min-max and momentum quantization ranges, increasing rounding error for nodes with smaller in-degrees.Choosing qmin and qmax therefore trades off truncation error against rounding error.
  • Sources of Error: Aggregation outputs vary with node in-degree, with mean and variance growing fastest for GIN and remaining roughly constant for GAT.Empirical results on Cora matched the predicted trends for GIN, GCN, and GAT.
  • Degree-Quant: Degree-Quant stochastically protects high in-degree nodes by performing their message passing, aggregation, and update computations at full precision during training.Protected nodes are sampled using degree-dependent Bernoulli probabilities, while quantized weights remain shared across nodes.
  • Degree-Quant: Percentile tracking clips the top and bottom 0.1% of tensor values to make quantization ranges representative of most values and reduce rounding error.The method is proposed as a core QAT component for graphs because aggregation variance can fluctuate substantially with degree.
  • Degree-Quant: Degree-Quant is architecture-agnostic and complementary to sampling-based memory reduction and weight-pruning approaches.The method targets low-precision arithmetic at inference time without changing the model architecture.
  • Quantization Choices: Quantization implementation choices can cause dramatic performance changes across INT8 and INT4 experiments.The study evaluates four typical quantization implementations, with independently tuned hyperparameters for each dataset-model pair.

4 EXPERIMENTS

Experiments compare Degree-Quant with FP32, QAT, and stochastic QAT across six datasets and multiple GNN architectures. Degree-Quant generally recovers or improves quantized performance, with especially large gains at INT4.

  • Experimental setup: Six datasets evaluate FP32, INT8-QAT, INT4-QAT, stochastic QAT, and Degree-Quant across node-level and graph-level tasks.The datasets are Cora, CiteSeer, ZINC, MNIST, CIFAR10 superpixels, and REDDIT-BINARY.
  • Degree-Quant results: INT8 Degree-Quant models recover most accuracy lost to quantization and sometimes outperform extensively tuned FP32 baselines.This pattern is reported across the evaluated datasets, with some architecture- and dataset-specific exceptions discussed for GIN.
  • Quantization implementation: The choice of quantization implementation produces substantially more performance variance across GNN architectures and problems than is typically observed for CNNs.The study evaluates four off-the-shelf quantization procedures on GCN, GAT, and GIN.
  • Training behavior: INT4 training generally benefits from momentum, while stochastic masking recovers some citation-network accuracy but has little impact elsewhere and can harm performance.For ZINC, GIN has a 16% regression-loss increase versus FP32, while W4A4 accuracy drops exceed 35% on Cora and 47% on Citeseer.
  • Degree-Quant results: +26.2% is the reported Degree-Quant gain over INT4 baselines for GIN on Citeseer.Corresponding gains are +24.9% on Cora and +23.0% on REDDIT-BINARY.

5 DISCUSSION

The discussion explains how quantization-range choices and element sensitivity shape GNN degradation. Percentile ranges and Degree-Quant’s stochastic masking improve stability, while precision-sensitive aggregation and message elements require particular protection.

  • Latency and memory: CPU latency gains depend on graph topology and feature dimension, with reported INT8 speedups ranging from 1.3× on Cora and Citeseer to 4.3× on Reddit.The discussion attributes CPU gains to improved cache performance for sparse operations; GPUs benefit less.
  • Percentile ranges: Absolute min/max ranges can exceed the range needed for 99.9% of values by more than 2×, halving effective quantization resolution.Percentile-based ranges avoid allowing infrequent outliers to dominate the range.
  • Percentile ranges: Degree-Quant is more stable than the baselines and achieves strong results with an order of magnitude less tuning.The percentile-range component primarily improves stability, while sometimes also improving performance.
  • INT4 degradation: In GAT, aggregation and message elements are especially sensitive to reduced precision, whereas weight elements are consistently unaffected.The INT4 analysis converts individual elements of an INT8 GAT without Degree-Quant to isolate their effects.
  • Degree-Quant components: Stochastic protective masking alone often captures most of the gain over QAT, while percentile-based ranges contribute stability and the full method performs consistently across architectures and datasets.The full Degree-Quant method does not require the extensive preliminary analysis used to identify the individual effects.

6 CONCLUSION

Degree-Quant is presented as an architecture-agnostic, stable approach for deploying quantized GNNs with off-the-shelf hardware. It combines strong low-bit performance with compression and latency benefits relevant to resource-constrained devices.

  • Conclusion: 8× compression is achieved with 4-bit weights and activations while surpassing strong baselines by margins regularly exceeding 20%.At 8 bits, Degree-Quant models perform on par with or better than baselines.
  • Conclusion: INT8 Degree-Quant models achieve up to 4.7× lower latency than FP32 models.The conclusion frames this result as supporting acceleration on off-the-shelf hardware.
  • Conclusion: The paper positions Degree-Quant as a first step toward wider GNN deployment, including on resource-constrained devices such as smartphones.Implementation advice and additional quantization best practices are provided in appendix A.5.

A.1 EXPERIMENTAL SETUP

The experiments reproduce established FP32 baselines across several graph architectures and datasets, then evaluate quantization-aware training under specified implementation and hardware settings.

  • A.1 EXPERIMENTAL SETUP: Baselines use architectures and results from prior work on citation networks, MNIST, CIFAR-10, ZINC, and REDDIT-BINARY.The authors re-implemented the referenced architectures and datasets and replicated their FP32 results.
  • A.1 EXPERIMENTAL SETUP: The evaluation covers GCN, GAT, and GIN architectures, with GIN models using learnable parameter ϵ.The architecture descriptions and parameter counts are summarized in Tables 5 and 6.
  • A.1 EXPERIMENTAL SETUP: Hyperparameters are selected with random search and asynchronous Hyperband pruning over learning rate, weight decay, dropout, and drop-edge probabilities.Search ranges are initialized around values from the reference baseline implementations.
  • A.1 EXPERIMENTAL SETUP: Quantization-aware training quantizes layer inputs, weights, messages, aggregation inputs and outputs, and update-stage outputs.The attention mechanism after GAT softmax is excluded because of its numerical precision requirements.
  • A.1 EXPERIMENTAL SETUP: Batch-normalization layers remain unfolded, while their inputs and outputs are quantized before multiplication with layer weights.The authors identify BN folding for GNN deployments as future work.
  • A.1 EXPERIMENTAL SETUP: REDDIT-BINARY GIN models leave the first MLP input layer unquantized to avoid severe degradation for scalar node features.All other layers use quantization-aware training.

A.2 DATASETS

The study evaluates node- and graph-level tasks across citation, image-derived, molecular, and social-network datasets using standard or task-specific data splits.

  • A.2 DATASETS: Cora and Citeseer represent documents as nodes, citations as edges, and bag-of-words vectors as node features for node classification.The dataset statistics are reported in Table 7.
  • A.2 DATASETS: MNIST and CIFAR-10 are converted into graphs with SLIC, while ZINC supports graph regression and REDDIT-BINARY supports graph classification.The image datasets are transformed so that each node represents a region in the image.
  • A.2 DATASETS: MNIST, CIFAR-10, and ZINC use standard splits, citation datasets use the Kipf–Welling splits, and REDDIT-BINARY uses 10-fold cross-validation.These split choices are used throughout the evaluation.

A.3 QUANTIZATION IMPLEMENTATIONS

The quantization study compares gradient estimators, range-statistics observers, stochastic quantization, and percentile handling to characterize their effects on QAT for GNNs.

  • A.3 QUANTIZATION IMPLEMENTATIONS: The study compares vanilla STE, which passes gradients unchanged, with gradient clipping, which limits gradients outside the representable [qmin, qmax] range.These implementations are evaluated for their impact on quantization-aware training results.
  • A.3 QUANTIZATION IMPLEMENTATIONS: The update rules compare observed min/max statistics with momentum-based statistics for quantizer calibration.The supplied formulations include separate update rules for STE min/max and STE momentum.
  • A.3 QUANTIZATION IMPLEMENTATIONS: Quantization modules track tensor extrema to compute qmin, qmax, zero-point, and scale parameters.The experiments compare direct min/max tracking with momentum-based moving averages during training.
  • A.3 QUANTIZATION IMPLEMENTATIONS: Stochastic QAT samples a Bernoulli mask each training step to decide which weight elements are quantized or retained at full precision.Reported stochastic-QAT results use block size one because larger blocks often caused a severe performance drop.
  • A.3 QUANTIZATION IMPLEMENTATIONS: Percentile-based quantization remains important for graph-level tasks because variable graph sizes can produce more-tailed summarized representations.The discussion connects this effect to graph regression and graph classification settings.

A.5 IMPLEMENTATION ADVICE

The implementation advice recommends less aggressive quantization choices and identifies masking and percentile observers as components of Degree-Quant, with masking providing additional performance gains.

  • A.5 IMPLEMENTATION ADVICE: The released implementation is intended for downstream use and includes practical details for improving results and accelerating training.The authors provide a clean implementation on GitHub.
  • A.5 IMPLEMENTATION ADVICE: The experiments quantize nearly everything possible, a more aggressive setup than is necessary for practical deployment.This choice is made to study quantization pitfalls in GNNs.
  • A.5 IMPLEMENTATION ADVICE: Not quantizing the final layer, and potentially the first layer, improves accuracy, especially at INT4.The recommendation follows common practice for CNNs and Transformers.
  • A.5 IMPLEMENTATION ADVICE: Higher precision in summarization stages and mixed precision focused on message passing can improve accuracy with limited runtime cost.The authors identify message passing as the main source of quantization benefits.
  • A.5 IMPLEMENTATION ADVICE: The authors advise using a less aggressive convention than the one used in the study, particularly for the first two recommendations.This recommendation bounds how directly the reported setup should be transferred to deployment.
  • A.5 IMPLEMENTATION ADVICE: Degree-Quant combines topology-aware stochastic masking with percentile-based range observers for quantizers.Percentiles primarily stabilize optimization and reduce hyperparameter sensitivity, while masking adds further performance improvement.

A.5.3 PERCENTILES

Percentile-based range observers improve quantization robustness but add substantial training cost. The paper suggests sampling as a possible speedup, while learned step sizes and robust quantization remain unevaluated alternatives.

  • Percentile-based range observers can significantly slow training because computing the operation takes substantial time.
  • Sampling 10%—and potentially 1%—of the data may speed percentile computation, but this setup was not evaluated in the paper.
  • Learned step sizes are expected to improve GNN quantization performance beyond the percentile-based ranges used in this paper.
  • Robust quantization is another proposed direction for reducing sensitivity to changing quantization ranges.
  • Aggressively quantizing message phases is preferable because sparse operations typically dominate GNN inference, particularly on cache-dependent CPUs.
  • Reported CPU and GPU latency results are non-optimal because the implementations did not substantially use cache blocking or kernel fusion.

A.5.6 PITFALLS

The paper identifies instability and architecture-dependent sensitivity as important pitfalls in quantized GNNs. Its figures, tables, and implementation details document how precision, stochastic quantization, topology-aware masking, and graph-level output stages are handled.

  • A.5.6 PITFALLS: Standard deviations reached 18% in some experiments, and citation-network models failed to converge on GPUs, requiring CPU training.
  • A.5.6 PITFALLS: GCN is more tolerant of INT4 quantization, whereas GIN and GAT suffer when accurate post-update representations are further quantized.
  • A.5.6 PITFALLS: Figures 7 and 8 show INT8 GCN and GIN degradation on Cora as individual elements are converted to INT4 without Degree-Quant.
  • A.5.6 PITFALLS: Degree-Quant validation curves for REDDIT-BINARY average results across 10-fold cross-validation and compare four DQ-INT8 settings with an FP32 baseline.
  • A.5.6 PITFALLS: The DQ ablation separates stochastic masking from percentile range tracking and evaluates their contributions at INT8 and INT4.
  • A.5.6 PITFALLS: The latency table reports INT8 results on a 22-core Intel Xeon Gold 6152 and GTX 1080Ti, using 128 input/output features and 1K CIFAR-10 graphs per batch.
  • A.5.6 PITFALLS: DQ generates degree-based masks at each training stage and reuses them in subsequent quantization layers, retaining higher-degree nodes at FP32 more often.
  • A.5.6 PITFALLS: nQAT applies stochastic quantization only to weights, while its remaining quantization modules follow standard QAT.
Loading 2008.05000v3…