[Paper Review] A Unified Lottery Ticket Hypothesis for Graph Neural Networks

Tianlong Chen, Yongduo Sui+3 more|arXiv (Cornell University)|Feb 12, 2021

Advanced Graph Neural NetworksComputer Science71 references28 citations

TL;DR

This paper introduces Unified GNN Sparsification (UGS) to co-prune graphs and GNN weights and generalizes the Lottery Ticket Hypothesis to Graph Neural Networks, yielding Graph Lottery Tickets (GLTs) with substantial inference savings and preserved performance.

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.

Motivation & Objective

• Motivate and reduce the computational cost of GNN inference on large-scale graphs by jointly pruning graph structure and model parameters.
• Generalize the Lottery Ticket Hypothesis to GNNs by defining Graph Lottery Tickets (GLTs) that combine a sparse sub-graph with a sparse sub-network.
• Demonstrate that GLTs can be found from random initialization and self-supervised pre-training across diverse GNN architectures and tasks.
• Show scalability of UGS/GLTs to deep GNNs on large datasets and provide insights into sparsity levels and hardware implications.

Proposed method

• Propose Unified GNN Sparsification (UGS) with differentiable masks for the graph (m_g) and the model weights (m_theta) to jointly prune A and Theta via a loss L_UGS = L({m_g ⊙ A, X}, m_theta ⊙ Θ) + γ1||m_g||1 + γ2||m_theta||1.
• Use iterative pruning (Algorithm 2) to locate Graph Lottery Tickets by repeatedly applying UGS and rewinding weights to the original initialization, yielding GLTs with specified graph sparsity s_g and model sparsity s_theta.
• Define GLT as the triplet (sparse graph G_s, sparse mask m_theta, initialization Θ0) that achieves performance comparable to training on the full graph and dense model.
• Evaluate GLTs across GCN/GIN/GAT architectures on small graphs (Cora, Citeseer, PubMed) and large graphs (OGB), and on node classification and link prediction tasks.

Experimental results

Research questions

• RQ1 Can a unified sparsification framework prune both graph structure and GNN weights without sacrificing performance?
• RQ2 Do Graph Lottery Tickets exist across different GNN architectures and tasks, and can they be found from random initialization or self-supervised pre-training?
• RQ3 How do sparsity levels in graphs and models affect node classification and link prediction on small vs. large graphs?
• RQ4 Is UGS scalable to deep GNNs on large-scale datasets, and how does graph density influence robustness to pruning?

Key findings

• GLTs exist with substantial MACs savings and little to no performance loss across GCN, GIN, and GAT on Cora, Citeseer, and PubMed.
• GLTs achieve 20%–98% MACs savings for nodes and 25%–85% MACs savings for large graphs, with up to 58.19% graph sparsity and 97.75% model sparsity.
• For link prediction, GLTs yield 48%–97% MACs savings on small graphs and 70% savings on large graphs without accuracy loss.
• UGS outperforms random pruning and is competitive with ADMM, with advantages in scalability to large graphs and modern GNN variants.
• GLTs can be derived from self-supervised pre-training (GraphCL) and can offer robustness and performance gains over random initializations in many settings.
• UGS scales to deep 28-layer ResGCNs on large graphs, achieving substantial MACs savings (e.g., up to 85% on some datasets) while preserving performance.