Source-linked AI summary
SGQuant: Squeezing the Last Bit on Graph Neural Networks with Specialized Quantization
Boyuan Feng, Yuke Wang, Xu Li, Shu Yang, Xueqiao Peng, Yufei Ding
TL;DR
GNNs’ high memory footprint limits their application on resource-constrained devices, while existing quantization approaches may trade accuracy for savings or leave savings suboptimal. SGQuant combines GNN-tailored quantization, multiple granularities, and automatic bit selection. It reports substantial memory reductions with small average accuracy losses across GCN, AGNN, and GAT.
Problem
GNN quantization lacks an efficient approach that balances memory savings with accuracy while accounting for GNN-specific properties.
Source: §I · p. 1Method
SGQuant combines GNN-tailored quantization and fine-tuning with component-wise, topology-aware, and layer-wise quantization and automatic bit selection.
Source: §I · p. 1Results
Memory consumption falls by up to 31.9×, 29.59×, and 13.2× on GCN, AGNN, and GAT, with average accuracy losses of 0.34%, 0.31%, and 0.47%, respectively.
Source: B. Overall Performance · p. 7Takeaways & Limitations
SGQuant provides a multi-granularity approach to compressing GNN features while minimizing accuracy drop.
Source: §VII · p. 8
Abstract
from arXiv · showhide
With the increasing popularity of graph-based learning, Graph Neural Networks (GNNs) win lots of attention from the research and industry field because of their high accuracy. However, existing GNNs suffer from high memory footprints (e.g., node embedding features). This high memory footprint hurdles the potential applications towards memory-constrained devices, such as the widely-deployed IoT devices. To this end, we propose a specialized GNN quantization scheme, SGQuant, to systematically reduce the GNN memory consumption. Specifically, we first propose a GNN-tailored quantization algorithm design and a GNN quantization fine-tuning scheme to reduce memory consumption while maintaining accuracy. Then, we investigate the multi-granularity quantization strategy that operates at different levels (components, graph topology, and layers) of GNN computation. Moreover, we offer an automatic bit-selecting (ABS) to pinpoint the most appropriate quantization bits for the above multi-granularity quantizations. Intensive experiments show that SGQuant can effectively reduce the memory footprint from 4.25x to 31.9x compared with the original full-precision GNNs while limiting the accuracy drop to 0.4% on average.
I. INTRODUCTION
GNNs offer high accuracy and generality, but their graph structures and node features strain the memory and power budgets of resource-constrained devices. SGQuant combines GNN-tailored quantization, multiple quantization granularities, and automatic bit selection to reduce memory while limiting accuracy loss.
I. INTRODUCTION
- GNNs are presented as accurate and general graph-learning tools applicable to different graph structures and dynamic graphs without much recomputing overhead.The cited examples include node classification and link prediction.Source: §I · p. 1
- GNN graph structures and node features can strain the memory budgets of embedded systems and IoT devices.The passage also identifies data operations and movements as consumers of these devices’ limited energy.Source: §I · p. 1
- Existing quantization approaches risk either high accuracy loss from aggressive uniform quantization or suboptimal savings from conservative quantization.The introduction also cautions that CNN quantization techniques do not account for GNN-specific properties.Source: §I · p. 1
- SGQuant combines a GNN-tailored quantization algorithm and fine-tuning with component-wise, topology-aware, and layer-wise quantization.The design aims to reduce memory consumption while maintaining accuracy and accommodate diverse precision demands.Source: §I · p. 1
- Automatic bit selection chooses appropriate bits for the different quantization granularities.The contribution is framed as end-to-end selection.Source: §I · p. 1
- SGQuant reduces memory consumption by up to 31.9× while limiting average accuracy loss to 0.4%.The reported comparison is against the original full-precision model.Source: §I · p. 2
A. Graph Neural Network
GNN layers use attention to relate neighboring nodes and combine their embeddings, while storing node embeddings creates memory overhead that grows with graph size. GNN structure and topology also offer quantization opportunities distinct from CNN quantization.
A. Graph Neural Network
- GNNs use attention to measure relationships between neighboring nodes, then combine attention-weighted neighboring embeddings to compute each node’s next-layer embedding.Source: A. Graph Neural Network · p. 2
- Storing embeddings for N nodes requires an N × N_k matrix at layer k, creating heavy memory overhead on large graphs.Source: A. Graph Neural Network · p. 2
B. Quantization
- GNN topology and architectural hierarchy permit quantization tailored to nodes, layers, and data categories, unlike CNNs’ fixed feature-map shapes.Source: B. Quantization · p. 2
III. GNN-TAILORED QUANTIZATION
SGQuant quantizes both GNN attention and embedding matrices, unlike CNN approaches described here that generally quantize features alone. It addresses mismatched bit widths through rematching before combination computation.
A. Quantization Algorithm Design
- SGQuant quantizes both the N × N attention matrix and the N × D_k embedding matrix, whereas CNN quantization generally considers only features.Source: A. Quantization Algorithm Design · p. 2
- For q-bit attention and p-bit embeddings, quantization reduces each matrix’s storage to q/32 and p/32 of its full-precision size, respectively.Source: A. Quantization Algorithm Design · p. 3
- When attention and embeddings use different bit widths, rematching recovers their quantized values to 32-bit before the combination component uses them.Source: A. Quantization Algorithm Design · p. 3
- The rematching overhead is negligible because computation recovers only nodes connected to the current node.Source: A. Quantization Algorithm Design · p. 3
B. GNN Quantization Finetuning
Directly quantizing GNNs at inference can cause substantial accuracy loss; fine-tuning the quantized model recovers most of it. SGQuant enables end-to-end fine-tuning with a tailored quantization layer.
B. GNN Quantization Finetuning
- An accuracy loss up to 10% can be largely recovered to less than 0.5% when the quantized GNN is fine-tuned, and this procedure is needed only once per quantized GNN model.Source: B. GNN Quantization Finetuning · p. 3
- SGQuant uses the original GNN loss and a straight-through estimator to enable backpropagation through the floor-based quantization operation.Source: B. GNN Quantization Finetuning · p. 3
- A tailored PyTorch-Geometric quantization layer supports both quantized inference and backpropagation for end-to-end fine-tuning.Source: B. GNN Quantization Finetuning · p. 3
IV. MULTI-GRANULARITY QUANTIZATION
SGQuant uses finer quantization granularity because layers, nodes, and components differ in sensitivity to quantization error. Component-wise quantization assigns separate bit widths to attention and combination operations.
IV. MULTI-GRANULARITY QUANTIZATION
- Uniform quantization applies the same bit width across GNN layers and components, but ignores differing sensitivities and can degrade accuracy.Source: §IV · p. 3
A. Component-wise Quantization
- Component-wise quantization gives attention and combination components different bit widths because attention is more robust to numerical error than combination.Source: A. Component-wise Quantization · pp. 3–4
IV. MULTI-GRANULARITY QUANTIZATION
- Figure 4 contrasts component-wise, topology-aware, layer-wise, and uniform quantization; shared colors denote the same bit width.Source: Figure 4 · p. 4
A. Component-wise Quantization
- The component-wise configuration assigns separate bits to attention and combination, then rematches their values to 32-bit before combination computation.Source: A. Component-wise Quantization · p. 4
B. Topology-aware Quantization
Topology-aware Quantization (TAQ) assigns node embeddings different bit widths according to node degree. It uses lower precision for high-degree nodes because their larger neighborhoods make them more robust to quantization error.
B. Topology-aware Quantization
- TAQ assigns quantization bits to node features according to degree, using higher precision for low-degree nodes and lower precision for high-degree nodes.Degrees are divided into intervals, and Fbit maps each interval to a bit width.Source: B. Topology-aware Quantization · p. 4
- Large-degree nodes are more robust to quantization error because aggregation averages errors across many neighbors toward zero.The passage attributes this convergence to the law of large numbers.Source: B. Topology-aware Quantization · p. 4
- TAQ leaves the attention matrix unquantized because it considers first-order topology information and omits second-order information.The described topology signal is node degree.Source: B. Topology-aware Quantization · pp. 4–5
C. Layer-wise Quantization
Layer-wise Quantization (LWQ) assigns bits by layer, motivated by differing layer sensitivities. SGQuant can combine granularities, including LWQ with component-wise quantization (CWQ) and topology-aware quantization (TAQ).
C. Layer-wise Quantization
- LWQ gives higher bits to leading layers and lower bits to succeeding layers, reflecting their low-level-detail and high-level-feature roles.At equal memory consumption, the evaluation reports higher accuracy when higher bits go to leading rather than succeeding layers.Source: C. Layer-wise Quantization · p. 5
- LWQ uses the same layer-specific bit width for the attention and embedding matrices at that layer.The layer bit is denoted q_k, with n denoting the number of GNN layers.Source: C. Layer-wise Quantization · p. 5
D. Combine Multiple Granularities
- LWQ and CWQ can be combined to provide finer-grained configurations; LWQ+CWQ achieves lower bits at the same accuracy than either alone.The passage says the combination could potentially yield higher accuracy under the same memory budget.Source: D. Combine Multiple Granularities · p. 5
- SGQuant also describes combining LWQ, TAQ, and CWQ, with attention-matrix bits independent of topology and embedding bits conditioned on node degree.The text also notes LWQ+TAQ and CWQ+TAQ combinations but omits their details.Source: D. Combine Multiple Granularities · p. 5
V. AUTO-BIT SELECTION
SGQuant frames bit assignment across quantization granularities as a combinatorial optimization problem: select discrete bit widths to minimize end-to-end loss. Its automatic bit selection uses a learned cost model and configuration exploration to address the design-space and prediction challenges.
V. AUTO-BIT SELECTION
- Bit selection across granularities is formulated as choosing discrete quantization widths to minimize end-to-end loss.The loss can be cross-entropy for classification or L2 norm for regression, and the formulation includes weight matrices because quantized GNNs are fine-tuned end-to-end.Source: §V · p. 5
- The search space grows exponentially when layer-, component-, and topology-aware quantization are used together, while model and graph diversity complicate prediction of quantization effects.Graph size and degree distributions vary, and topology can substantially affect quantization error.Source: §V · p. 5
- ABS combines a machine-learning cost model that predicts quantized-GNN accuracy with an exploration scheme that selects promising configurations.This design is proposed to address the challenges in configuration selection.Source: §V · p. 5
A. Machine Learning Cost Model
The machine-learning cost model learns interactions among quantization bits, GNN models, and graph topology from measured configurations. An iterative exploration scheme uses its predictions to select promising configurations while balancing prediction accuracy and selection overhead.
A. Machine Learning Cost Model
- Random search may need many samples, while a predefined cost model may not generalize across GNN models and graph inputs.The predefined model is described as specific to a particular GNN and graph topology.Source: A. Machine Learning Cost Model · pp. 5–6
- The cost model learns from sampled quantization configurations, using their bit settings as features and measured accuracy as labels to predict remaining configurations.It treats prediction as regression and uses a regression tree for faster inference and lower training-data demands than neural networks.Source: A. Machine Learning Cost Model · p. 6
B. Exploration Scheme
- The exploration scheme iteratively trains the cost model, predicts accuracy for sampled configurations, evaluates top-ranked candidates, and repeats the process.It keeps configurations with accuracy drop below 0.5% and selects the remaining configuration with the lowest memory consumption.Source: B. Exploration Scheme · p. 6
- The iterative search balances model-training overhead against prediction precision, with N_measure=40 and N_iter=5 reported as a balance in the experiments.The default N_sample is 2000, with inference latency below 0.1 seconds per iteration.Source: B. Exploration Scheme · p. 6
VI. EVALUATION
The evaluation tests SGQuant across GCN, AGNN, and GAT on small feature-rich and large graph datasets, measuring memory savings alongside accuracy loss.
1) GNN Architectures:
- The evaluated architectures are GCN, AGNN, and GAT, with GCN also serving as a backbone for GraphSage and Diffpool.The paper describes AGNN as using specialized propagation layers and GAT as an attention-based reference architecture.Source: 1) GNN Architectures · p. 6
2) Datasets:
- The evaluation covers small, feature-rich Citeseer, Cora, and Pubmed graphs and larger Amazon-computer and Reddit graphs.The passage characterizes the first group as small in nodes and edges, but high-dimensional in feature embeddings.Source: 2) Datasets · p. 6
B. Overall Performance
- GCN, AGNN, and GAT reduce memory by 31.9×, 29.59×, and 13.2×, with average accuracy losses of 0.34%, 0.31%, and 0.47%, respectively.These comparisons are against the original full-precision models.Source: Table III · p. 7, B. Overall Performance · p. 7
- Smaller datasets allow lower average bits, while more complex models select higher bits; for example, GCN uses 1.22 bits on Cora, and GAT uses 7.53 bits on Amazon-computer.The paper attributes this pattern to smaller graphs having less significant quantization precision loss, and complex models having more intricate computations that can enlarge it.Source: B. Overall Performance · p. 7
- On Reddit, SGQuant reduces memory occupation by up to 530MB.The paper identifies this as an example of the absolute savings on large datasets.Source: B. Overall Performance · p. 7
C. Breakdown Analysis of Multi-granularity Quantization
The breakdown compares quantization granularity and automatic bit selection, then concludes that SGQuant combines layer-, component-, and topology-aware choices to reduce memory with negligible accuracy loss.
C. Breakdown Analysis of Multi-granularity Quantization
- Compared with uniform quantization, layer-wise quantization lowers error, layer-plus-component quantization further mitigates accuracy degradation, and adding topology awareness lowers error at each memory size.The study attributes the topology-aware result to incorporating node degree, since high-degree nodes gather more information from neighbors and can tolerate more aggressive quantization with minor information loss.Source: C. Breakdown Analysis of… · pp. 7–8, Figure 7 · p. 8
- At around 2MB, LWQ+CWQ+TAQ reduces GAT-on-Cora error rate by 2.2% versus uniform quantization at the same memory size.Table IV reports allocated bit-width and error rate for the different granularities.Source: Table IV · p. 8, C. Breakdown Analysis of… · p. 8
D. Effectiveness of Auto-Bit Selection
- On AGNN and Cora, the ML cost model converges within 200 quantization-configuration trials, locates appropriate quantization bits more swiftly than naive random search, and achieves 25× memory saving versus random search’s 20×.Source: D. Effectiveness of Auto-Bit Selection · p. 8
VII. CONCLUSION
- SGQuant combines layer-wise, component-wise, and topology-aware quantization with ML-based automatic bit selection; experiments report up to 31.9× memory reduction under negligible accuracy drop.The conclusion presents ABS as a way to reduce users’ design-exploration effort.Source: §VII · p. 8