Timezone: »

A Unified Lottery Ticket Hypothesis for Graph Neural Networks
Tianlong Chen · Yongduo Sui · Xuxi Chen · Aston Zhang · Zhangyang “Atlas” Wang

Tue Jul 20 05:25 AM -- 05:30 AM (PDT) @

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 are at https://github.com/VITA-Group/Unified-LTH-GNN.

Author Information

Tianlong Chen (University of Texas at Austin)
Yongduo Sui (University of Science and Technology of China)
Xuxi Chen (University of Texas at Austin)
Aston Zhang (AWS AI)
Aston Zhang

Aston Zhang is a senior scientist at Amazon Web Services AI. His research interests are in deep learning. He received a Ph.D. in computer science from University of Illinois at Urbana-Champaign. He has served as an editorial board member for Frontiers in Big Data and a program committee member (reviewer) for ICML, NeurIPS, ICLR, WWW, KDD, SIGIR, and WSDM. His book Dive into Deep Learning (www.d2l.ai) has been used as a textbook in worldwide universities.

Zhangyang “Atlas” Wang (University of Texas at Austin)

Related Events (a corresponding poster, oral, or spotlight)

More from the Same Authors