SGS-GNN: A Supervised Graph Sparsifier for Graph Neural Networks
TL;DR
SGS-GNN improves F1-scores by 4% relative to full training and up to 30% on heterophilic graphs and outperforms state-of-the-art methods by 4–7% at similar sparsity levels while reducing peak memory usage by up to 3.9×.
Abstract
We propose SGS-GNN, a supervised graph sparsifier for Graph Neural Networks (GNNs) to improve predictive performance and reduce the cost of message passing by removing task-irrelevant edges. Existing unsupervised sparsifiers are not task-aware, while existing supervised sparsifiers suffer from significant memory overhead, poor sparsity control, and a lack of homophily/heterophily awareness. SGS-GNN addresses these limitations by adopting a feature- and structure-aware edge-probability encoder, a sparse subgraph sampler that strictly adheres to a global sparsity constraint, and a homophily-aware regularizer to improve prediction accuracy across homophilic and heterophilic graphs. A key scalability-enhancing feature of SGS-GNN is that it ensures encoder updates are computed by backpropagating through the sampled subgraph, and avoids retaining edge-level computation graphs for all edges via gradient checkpointing. A key efficiency-enhancing feature of SGS-GNN is that the edge-probability encoder is updated only when it outperforms a degree-based edge sampler, ensuring performance no worse than a strong unsupervised baseline. Experiments on 33 homophilic and heterophilic graphs show that SGS-GNN improves F1-scores by 4% relative to full training and up to 30% on heterophilic graphs. Furthermore, SGS-GNN outperforms state-of-the-art methods by 4–7% at similar sparsity levels while reducing peak memory usage by up to 3.9×.