Mini-Batching and Scaling GNNs to Large Graphs

Training a neural network on ordinary tabular data is easy to mini-batch: grab any random 256 rows, they're independent of each other. A graph neural network breaks this, because a node's prediction depends on its neighbors, whose predictions depend on their neighbors, and so on — a randomly chosen batch of nodes drags along an ever-expanding web of dependencies, and on a graph with millions of nodes this "neighborhood explosion" can mean a single training batch effectively touches the whole graph.

The standard fix is neighborhood sampling: instead of using every neighbor at every layer, randomly sample a fixed, small number of neighbors per node per layer (GraphSAGE popularized sampling, say, 10 neighbors at the first layer and 5 at the second), which bounds the computational blow-up to a manageable, fixed size regardless of how large or densely connected the full graph is. A second common approach is graph clustering-based batching (as in Cluster-GCN), which partitions the graph into densely-connected sub-clusters ahead of time and trains on one cluster at a time, trading a little accuracy at cluster boundaries for training batches that fit in memory.

Either approach lets a GNN train on graphs — a payments network, a corporate ownership graph — with far more nodes and edges than would fit if every layer required the full neighborhood, at the cost of some noise from only ever seeing a sampled slice of each node's true neighborhood.