GRAPES: Learning to Sample Graphs for Scalable Graph Neural Networks

TL;DR

GRAPES introduces an adaptive, learnable sampling method for GNNs that improves scalability and accuracy across diverse graph types by dynamically identifying crucial nodes during training.

Contribution

It proposes GRAPES, a novel adaptive sampling technique that trains a secondary GNN to predict node importance, enhancing GNN scalability and performance on various graph structures.

Findings

GRAPES outperforms static sampling methods in accuracy and scalability.

It maintains high accuracy with smaller sample sizes.

Effective on both homophilous and heterophilous graphs.

Abstract

Graph neural networks (GNNs) learn to represent nodes by aggregating information from their neighbors. As GNNs increase in depth, their receptive field grows exponentially, leading to high memory costs. Several existing methods address this by sampling a small subset of nodes, scaling GNNs to much larger graphs. These methods are primarily evaluated on homophilous graphs, where neighboring nodes often share the same label. However, most of these methods rely on static heuristics that may not generalize across different graphs or tasks. We argue that the sampling method should be adaptive, adjusting to the complex structural properties of each graph. To this end, we introduce GRAPES, an adaptive sampling method that learns to identify the set of nodes crucial for training a GNN. GRAPES trains a second GNN to predict node sampling probabilities by optimizing the downstream task objective. We evaluate GRAPES on various node classification benchmarks, involving homophilous as well as heterophilous graphs. We demonstrate GRAPES' effectiveness in accuracy and scalability, particularly in multi-label heterophilous graphs. Unlike other sampling methods, GRAPES maintains high accuracy even with smaller sample sizes and, therefore, can scale to massive graphs. Our code is publicly available at https://github.com/dfdazac/grapes.