A Unified Lottery Ticket Hypothesis for Graph Neural Networks

Tianlong Chen; Yongduo Sui; Xuxi Chen; Aston Zhang; Zhangyang Wang

Source-linked AI summary

A Unified Lottery Ticket Hypothesis for Graph Neural Networks

Tianlong Chen, Yongduo Sui, Xuxi Chen, Aston Zhang, Zhangyang Wang

arXiv:2102.06790v2Feb 12, 2021cs.LGcs.AIstat.ML

TL;DR

GNN scaling is hindered by the size and connectivity of graphs, which weight pruning alone does not address. The paper introduces UGS to prune graph adjacency and model weights jointly, then defines GLTs as jointly identified sparse graph-model pairs. Across node classification and link prediction, GLTs preserve predictive performance while substantially reducing inference MACs, though practical hardware acceleration remains open.

Problem
GNN training and inference become increasingly expensive as graphs grow, and existing weight-pruning methods do not address graph size and connectivity.
Method
UGS jointly prunes graph adjacency matrices and model weights, enabling GLTs composed of a core sub-dataset and sparse sub-network.
Results
GLTs preserve performance while reducing inference MACs across architectures, tasks, and dataset scales, including 20% ∼98% savings for node classification and 48% ∼97% plus 70% savings for link prediction.
Takeaways & Limitations
Joint graph-model sparsification extends the lottery ticket hypothesis to GNNs and identifies sparse graph-model pairs that retain task performance.
Takeaways & Limitations
How GLT sparsity translates into practical acceleration and energy savings remains open because many accelerators favor dense, regular computation.

Abstract

from arXiv · show

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.

1. Introduction

GNN inefficiency grows with graph size and connectivity, while prior pruning largely targets models rather than jointly simplifying graphs and weights. The paper introduces UGS and GLTs to co-prune both components while preserving performance across architectures, tasks, datasets, and initialization methods.

1. Introduction: GNN training and inference become inefficient at scale because neighborhood aggregation is costly on large, densely connected graphs.
1. Introduction: Prior work explored graph simplification and model compression separately, but did not jointly simplify input graphs and GNN models.
1. Introduction: UGS simultaneously prunes the graph adjacency matrix and model weights without assuming a particular GNN architecture or graph structure.
1. Introduction: A GLT jointly identifies a core sub-dataset and sparse sub-network that can be trained from initialization to match full-graph, full-model performance.
1. Introduction: 20% ∼98% MACs saving is achieved for node classification with little to no performance degradation.The reported graph sparsity reaches 5% ∼58.19% and GNN-model sparsity 20% ∼97.75%.
1. Introduction: 48% ∼97% and 70% MACs saving are reported for link prediction without performance loss, while the framework scales to GNNs up to 28 layers on large graphs.
1. Introduction: GLTs can also be found from self-supervised pre-trained GNN initializations, including with sparser graphs and models.

2. Related Work

Related work covers GNN architectures, graph sampling or sparsification, and lottery-ticket-based sparse subnetworks. The paper argues that GNNs require a distinct extension because they have smaller models and motivate joint data-model simplification.

2. Related Work: GNN research spans graph convolutions, anisotropic operations such as attention, and models addressing limitations of earlier architectures.
2. Related Work: Graph sampling and sparsification extract smaller subgraphs while aiming to preserve learning effectiveness and reduce cost as graph size increases.
2. Related Work: Earlier lottery-ticket work studies trainable sparse subnetworks that replace full models without sacrificing performance, including through rewinding techniques and applications across fields.
2. Related Work: The paper presents GNNs as a distinct LTH setting because their models are smaller and proposes extending LTH from model simplification to data-model co-simplification.

3. Methodology

The methodology formulates GNN computation over graph structure and introduces UGS to jointly sparsify graph edges and model weights, then iteratively identifies graph lottery tickets. The resulting GLTs combine a sparse graph, sparse model mask, and original initialization, with inference costs determined by the remaining edges and weights.

3.1. Notations and Formulations: A two-layer GNN uses normalized adjacency propagation, node features, learnable weights, activation, and softmax prediction.The graph is represented by adjacency matrix A and feature matrix X; the normalized adjacency is derived from A + I.
3.2. Unified GNN Sparsification: UGS introduces differentiable graph and weight masks whose shapes match the adjacency matrix and model weights, and co-optimizes them end to end.The framework is designed to reduce both graph edges and GNN parameters simultaneously.
3.2. Unified GNN Sparsification: UGS trains masked graph and model parameters with ℓ1 sparsity regularizers, then zeroes the lowest-magnitude mask elements at predefined pruning ratios.The two resulting masks are applied to the adjacency matrix and model weights to produce the sparse graph and model.
3.3. Graph Lottery Tickets: A GLT is a sparse graph and sparse GNN subnetwork that, initialized from Θ0, matches or surpasses the original model on the full graph.The GLT comprises the sparse graph Gs, model mask mθ, and original initialization Θ0.
3.3. Graph Lottery Tickets: GLTs are found by repeatedly applying UGS, updating graph and model masks, and rewinding GNN weights to the original initialization until target sparsities are reached.The procedure separately targets graph sparsity sg and GNN sparsity sθ.
3.3. Graph Lottery Tickets: GLT inference has time complexity O(L × ∥mg ⊙A∥0 × F + L × ∥mθ∥0 × |V| × F2) and memory complexity O(L×|V|×F+L×∥mθ∥0×F2).Pruned edges are removed from E and do not participate in subsequent computation.

4. Experiments

Experiments evaluate UGS and GLTs across GNN architectures, graph scales, tasks, initialization schemes, and pruning settings. GLTs generally preserve predictive performance while substantially reducing inference MACs, with graph sparsification emerging as the main bottleneck and denser graphs showing greater resilience.

4.1. The Existence of Graph Lottery Ticket: GLTs exist across GCN, GIN, and GAT on Cora, Citeseer, and PubMed for node classification and link prediction, reducing inference MACs by 20%–98%.They are identified over graph sparsities from 5% to 58.19% without performance deterioration.
4.1. The Existence of Graph Lottery Ticket: UGS consistently outperforms random pruning, while ADMM is competitive at moderate graph sparsity but performs 3%–4% worse under heavy pruning.UGS is also described as flexible beyond ADMM’s adjacency-matrix and aggregation assumptions.
4.1. The Existence of Graph Lottery Ticket: GAT tolerates higher sparsity than GCN and GIN, while Cora is more pruning-sensitive and PubMed is more robust.The paper suggests attention-based aggregation may re-identify important connections in pruned graphs.
4.2. Scale Up Graph Lottery Tickets: 85%, 25%, and 70% MACs savings are achieved on Ogbn-ArXiv, Ogbn-Proteins, and Ogbl-Collab, respectively, while matching performance with 28-layer ResGCNs.These experiments test scalability on large graphs with millions of connections.
4.2. Scale Up Graph Lottery Tickets: Denser graphs are more resilient to sparsification: Ogbn-Proteins, with average degree 597.00, shows a negligible UGS–random-pruning performance gap even when heavily pruned.The paper relates this robustness to high-degree nodes being less likely to become isolated.
4.3. Graph Lottery Ticket from Pre-training: GraphCL initialization consistently improves UGS performance at moderate sparsity levels, benefiting GLTs on node classification and link prediction.At extreme sparsity, initialization becomes less influential; on Citeseer, suitable sparsity can enlarge the pre-training gain.
4.4. Ablation Study: With fixed pθ = 20%, pg = 5% identifies the GLT, whereas pg = 20% substantially degrades accuracy at large sparsities.With fixed pg = 20%, pθ values of 10%, 20%, and 40% perform similarly, supporting graph sparsification as the key bottleneck.
4.4. Ablation Study: GLTs outperform randomly re-initialized Random GLTs, and UGS-sparsified graphs appear to preserve critical vertices that previously had denser connections.These findings provide evidence that the identified structure and initialization both matter.

5. Conclusion and Discussion

The paper presents GLTs as a unified approach that trims graph and GNN complexity while preserving task performance, but practical acceleration and energy savings remain open.

UGS simultaneously sparsifies graph adjacency matrices and model weights, providing a generalized pruning framework for GNNs.
GLTs broadly reduce inference MACs without sacrificing task performance across the reported settings.
Practical acceleration and energy-saving benefits of GLT sparsity remain unresolved because sparse, irregular graph operations are difficult to implement efficiently.

A1. More Implementation Details

The experiments use standard small and large graph benchmarks, specified train-validation-test splits, selected GNN architectures, and validation-based reporting.

Datasets: Experiments cover Cora, Citeseer, PubMed, Ogbn-ArXiv, Ogbn-Proteins, and Ogbl-Collab graph datasets.
Dataset Splits: Small- and medium-scale node classification uses 140, 120, and 60 training labels for Cora, Citeseer, and PubMed, respectively.
GNN Configurations: Small- and medium-scale experiments use two-layer GCN, GIN, and GAT networks with 512 hidden units, while large-scale experiments use 28-layer ResGCNs.
Training: For Ogbn-Proteins, training samples ten subgraphs from the whole graph and feeds one subgraph to the GCN per iteration.
Evaluation: Test accuracy, ROC-AUC, and Hits@50 are reported at the best validation result, with all training and evaluation conducted for one run.

A2. More Experiment Results

Additional experiments find lottery tickets across architectures, tasks, graph scales, sparsity levels, and initialization settings, while graph measurements characterize the resulting sparse structures.

Node Classification: UGS finds graph lottery tickets across 20%–90% GNN sparsity on small datasets without performance deterioration.The result is reported for GCN, GIN, and GAT on Cora, Citeseer, and PubMed node classification.
Link Prediction: Link-prediction GLTs occur across 5%–50% graph sparsity and 20%–90% GNN sparsity without performance deterioration.The experiments use GCN, GIN, and GAT on Cora, Citeseer, and PubMed.
Large-Scale and Deep GNNs: Deep ResGCN experiments on larger graphs find GLTs with 30%–50% graph sparsity and 20%–80% weight sparsity without performance deterioration.
Initialization: Self-supervised GraphCL pre-training enables UGS to identify higher-quality graph lottery tickets.The comparison covers node classification with low label rates and link prediction.
Graph Structure: UGS-produced sparse graphs appear to retain more critical vertices with higher connectivity than comparison graphs.The measurements include clustering coefficient and node and edge betweenness centrality, compared with original, random-pruned, and ADMM-sparsified graphs.