A Unified Lottery Ticket Hypothesis for Graph Neural Networks

Abstract

With graphs rapidly growing in size and deeper graph neural networks (GNNs) emerging, the training and inference of GNNs become increasingly expensive. Existing network weight pruning algorithms cannot address the main space and computational bottleneck in GNNs, caused by the size and connectivity of the graph. To this end, this paper first presents a unified GNN sparsification (UGS) framework that simultaneously prunes the graph adjacency matrix and the model weights, for effectively accelerating GNN inference on large-scale graphs. Leveraging this new tool, we further generalize the recently popular lottery ticket hypothesis to GNNs for the first time, by defining a graph lottery ticket (GLT) as a pair of core sub-dataset and sparse sub-network, which can be jointly identified from the original GNN and the full dense graph by iteratively applying UGS. Like its counterpart in convolutional neural networks, GLT can be trained in isolation to match the performance of training with the full model and graph, and can be drawn from both randomly initialized and self-supervised pre-trained GNNs. Our proposal has been experimentally verified across various GNN architectures and diverse tasks, on both small-scale graph datasets (Cora, Citeseer and PubMed), and large-scale datasets from the challenging Open Graph Benchmark (OGB). Specifically, for node classification, our found GLTs achieve the same accuracies with 20%~98% MACs saving on small graphs and 25%~85% MACs saving on large ones. For link prediction, GLTs lead to 48%~97% and 70% MACs saving on small and large graph datasets, respectively, without compromising predictive performance. Codes available at https://github.com/VITA-Group/Unified-LTH-GNN.

References in corpus (21)

• Semi-Supervised Classification with Graph Convolutional Networks
• Inductive Representation Learning on Large Graphs
• Neural Message Passing for Quantum Chemistry
• Spectral Networks and Locally Connected Networks on Graphs
• Graph Neural Networks: A Review of Methods and Applications
• Rethinking the Value of Network Pruning
• A Survey of Model Compression and Acceleration for Deep Neural Networks
• Interaction Networks for Learning about Objects, Relations and Physics
• Open Graph Benchmark: Datasets for Machine Learning on Graphs
• The State of Sparsity in Deep Neural Networks
• DeeperGCN: All You Need to Train Deeper GCNs
• Benchmarking Graph Neural Networks
• Comparing Rewinding and Fine-tuning in Neural Network Pruning
• Playing the lottery with rewards and multiple languages: lottery tickets in RL and NLP
• When Does Self-Supervision Help Graph Convolutional Networks?
• Winning the Lottery with Continuous Sparsification
• Graph InfoClust: Leveraging cluster-level node information for unsupervised graph representation learning
• The Lottery Tickets Hypothesis for Supervised and Self-supervised Pre-training in Computer Vision Models
• The Difficulty of Training Sparse Neural Networks
• EarlyBERT: Efficient BERT Training via Early-bird Lottery Tickets
• Improved large-scale graph learning through ridge spectral sparsification

Cited by in corpus (10)

• Chasing Sparsity in Vision Transformers: An End-to-End Exploration
• The Lottery Tickets Hypothesis for Supervised and Self-supervised Pre-training in Computer Vision Models
• LightGNN: Simple Graph Neural Network for Recommendation
• Sanity Checks for Lottery Tickets: Does Your Winning Ticket Really Win the Jackpot?
• Data-Efficient GAN Training Beyond (Just) Augmentations: A Lottery Ticket Perspective
• You are caught stealing my winning lottery ticket! Making a lottery ticket claim its ownership
• GANs Can Play Lottery Tickets Too
• Spending Your Winning Lottery Better After Drawing It
• Optimizing Connectivity through Network Gradients for Restricted Boltzmann Machines
• Membership Inference Attacks on Lottery Ticket Networks