VQ-GNN: A Universal Framework to Scale up Graph Neural Networks using Vector Quantization

Abstract

Most state-of-the-art Graph Neural Networks (GNNs) can be defined as a form of graph convolution which can be realized by message passing between direct neighbors or beyond. To scale such GNNs to large graphs, various neighbor-, layer-, or subgraph-sampling techniques are proposed to alleviate the "neighbor explosion" problem by considering only a small subset of messages passed to the nodes in a mini-batch. However, sampling-based methods are difficult to apply to GNNs that utilize many-hops-away or global context each layer, show unstable performance for different tasks and datasets, and do not speed up model inference. We propose a principled and fundamentally different approach, VQ-GNN, a universal framework to scale up any convolution-based GNNs using Vector Quantization (VQ) without compromising the performance. In contrast to sampling-based techniques, our approach can effectively preserve all the messages passed to a mini-batch of nodes by learning and updating a small number of quantized reference vectors of global node representations, using VQ within each GNN layer. Our framework avoids the "neighbor explosion" problem of GNNs using quantized representations combined with a low-rank version of the graph convolution matrix. We show that such a compact low-rank version of the gigantic convolution matrix is sufficient both theoretically and experimentally. In company with VQ, we design a novel approximated message passing algorithm and a nontrivial back-propagation rule for our framework. Experiments on various types of GNN backbones demonstrate the scalability and competitive performance of our framework on large-graph node classification and link prediction benchmarks.

NeurIPS 2021

References in corpus (21)

• Neural Discrete Representation Learning
• Unsupervised Learning of Visual Features by Contrasting Cluster Assignments
• Cluster-GCN: An Efficient Algorithm for Training Deep and Large Graph Convolutional Networks
• Fast Graph Representation Learning with PyTorch Geometric
• Simplifying Graph Convolutional Networks
• Linformer: Self-Attention with Linear Complexity
• Open Graph Benchmark: Datasets for Machine Learning on Graphs
• Graph Neural Networks with convolutional ARMA filters
• Scaling Graph Neural Networks with Approximate PageRank
• vq-wav2vec: Self-Supervised Learning of Discrete Speech Representations
• Adaptive Sampling Towards Fast Graph Representation Learning
• SIGN: Scalable Inception Graph Neural Networks
• Graph-Bert: Only Attention is Needed for Learning Graph Representations
• BIVA: A Very Deep Hierarchy of Latent Variables for Generative Modeling
• Generating Diverse High-Fidelity Images with VQ-VAE-2
• What Can Neural Networks Reason About?
• Is Attention Better Than Matrix Decomposition?
• Degree-Quant: Quantization-Aware Training for Graph Neural Networks
• Global Attention Improves Graph Networks Generalization
• Lipschitz Normalization for Self-Attention Layers with Application to Graph Neural Networks
• Learning Product Codebooks using Vector Quantized Autoencoders for Image Retrieval

Cited by in corpus (3)

• The Evolution of Distributed Systems for Graph Neural Networks and their Origin in Graph Processing and Deep Learning: A Survey
• Low-bit Quantization for Deep Graph Neural Networks with Smoothness-aware Message Propagation
• Hierarchical Vector Quantized Graph Autoencoder with Annealing-Based Code Selection