Sampling and Scaling to Large Graphs
Why full-batch training does not scale past a certain graph size, the neighbourhood explosion that makes naive mini-batching worse, and the three sampling strategies that resolve it.
Training a graph network on a small graph is easy: hold the whole thing in memory, compute all node representations at every layer, backpropagate. On a graph with a hundred million nodes this is impossible, and the obvious fix of mini-batching runs into a problem specific to graphs.
Neighbourhood explosion
To compute one node's representation at layer \(L\), you need its neighbours' representations at layer \(L-1\), which need their neighbours' at \(L-2\), and so on. The receptive field is the \(L\)-hop neighbourhood, and in a graph with average degree \(d\) it contains roughly \(d^L\) nodes.
With average degree 20 and three layers that is 8,000 nodes per target node. In a social network with hubs of degree tens of thousands, a single batch's receptive field can be a substantial fraction of the entire graph. Mini-batching by target node therefore does not reduce the computation the way it does for independent examples, which is the core difficulty.
The three strategies
Node-wise sampling, as in GraphSAGE, samples a fixed number of neighbours per node per layer. With a sample size of 10 and three layers, the receptive field is bounded at 1,000 regardless of degree, which makes the cost predictable and independent of the graph's degree distribution. It introduces variance in the estimated aggregation, and neighbours sampled independently per layer means the total node count still grows with depth.
Layer-wise sampling samples a fixed set of nodes for each layer jointly, so the count per layer is constant and does not compound with depth. Importance sampling by node degree reduces the variance. It requires the sampled layer to cover the previous layer's needs, which makes the implementation more involved.
Subgraph sampling, as in Cluster-GCN and GraphSAINT, partitions or samples a subgraph and runs full-batch training on it. Memory is bounded by the subgraph, computation is dense and efficient, and edges crossing the subgraph boundary are lost, which biases the aggregation toward within-cluster structure. Randomising the partition across epochs mitigates it.
When it breaks
Sampling variance interacts with depth. Each layer's sampling adds noise, and three layers of sampled aggregation is a noisier estimate than one. Larger samples reduce it and cost the memory the sampling was meant to save.
Hubs are both important and expensive. High-degree nodes appear in many receptive fields, so they dominate computation, and downsampling them saves a great deal and loses the information that made them central. Degree-aware sampling is the compromise.
Inference is a different problem from training. Sampled training with full-graph inference produces a distribution mismatch: the model was trained on partial neighbourhoods and applied to complete ones. Sampling identically at inference costs accuracy for consistency, and either choice needs to be deliberate.
Distributed training partitions the graph, and edges cross partitions. Cross-partition edges require communication per layer per batch, and the partition quality determines how much. Graph partitioning to minimise edge cuts is a hard problem that becomes part of the training infrastructure.
12 flashcards for this concept
Click a card to reveal the answer.