Graph Sampling Strategies for Graphs Too Large to Fit in Memory

Full-batch training on a graph computes every node’s representation at every layer, which needs the whole graph resident. The two ways out sample different things, and the choice changes both your memory curve and what your gradients are estimating.

The problem sampling solves

A graph convolution needs a node’s neighbours to compute its layer-1 output, its neighbours’ neighbours for layer 2, and so on. The receptive field of an L-layer model is the L-hop neighbourhood, and on a graph with average degree d that is roughly d^L nodes. With d = 50 and L = 3 that is 125,000 nodes for one prediction — and on a graph with a heavy-tailed degree distribution, a single hub in the neighbourhood pulls in millions. Sampling exists to bound that number.

Neighbour sampling

Pick a batch of seed nodes. For each, sample a fixed number of its neighbours; for each of those, sample a fixed number again; repeat once per layer. This is GraphSAGE, from Hamilton, Ying and Leskovec (NeurIPS 2017), and it is the default in every major graph library. The per-hop sample sizes are the fan-out.

The bound is per seed and it is a product across hops. The paper reports that in practice high performance was achievable with “K=2 and S1·S2 ≤ 500” — two layers, with the product of the two per-hop sample sizes kept under 500 — which is the authors telling you directly that the product, not the individual numbers, is the budget.

The bias is subtle and worth stating precisely. Sampling k neighbours uniformly gives an unbiased estimate of the mean over neighbours, so a mean-aggregating layer is fine in expectation. It is not unbiased for a sum aggregator, for max, or for anything else non-linear, and it is not unbiased through the composition of layers — the estimator’s variance compounds with depth, which is part of why deep sampled GNNs are noisier to train than the depth alone suggests.

There is also a per-batch cost that has nothing to do with the model: the sampled neighbourhood has to be relabelled into a dense block. The sampler collects a set of global node ids, deduplicates them, assigns each a local index, and rewrites the sampled edges against those indices so the layer can be a dense matrix multiply. On a large graph that relabelling and the feature gather that follows it usually cost more wall-clock time than the forward and backward passes combined, which is why fan-out is a throughput parameter as much as a modelling one.

Subgraph sampling

The alternative inverts the order: sample a subgraph first, then run a normal full-batch GNN on it. No per-layer expansion happens, because the subgraph is the whole world for that step.

Wei-Lin Chiang and colleagues’ Cluster-GCN (KDD 2019) partitions the graph with a graph-partitioning algorithm such as METIS into many small clusters, then forms each batch from a random handful of clusters and the edges among them. Because partitioning minimises cut edges, most edges survive inside a batch. Sampling several clusters per batch rather than one is what restores some of the between-cluster edges, and the paper is explicit that this stochastic multi-cluster step matters.

Hanqing Zeng and colleagues’ GraphSAINT (ICLR 2020) samples the subgraph directly — by random nodes, random edges, or random walks — and then corrects the bias analytically, applying normalisation coefficients derived from each node’s and edge’s sampling probability so that the aggregation is unbiased despite the subgraph being non-uniform. That correction is the contribution; a subgraph sampler without it is systematically wrong in favour of densely connected regions.

The two memory curves

Same inputs for both: 128-dimensional float32 features, hidden width 256, three layers, batch of 512 seed nodes, average degree 50.

NEIGHBOUR SAMPLING, fan-out [25, 10, 5]

nodes per seed = 1 + 25 + 250 + 1,250 = 1,526
batch slots    = 512 x 1,526          = 781,312
input features = 781,312 x 128 x 4    = 400 MB
layer-1 output = (512 x 276) x 256 x 4 =  145 MB
                  ^ only hops 0..2 need a layer-1 output

drop to fan-out [25, 10]:
nodes per seed = 1 + 25 + 250          = 276
batch slots    = 512 x 276             = 141,312
input features = 141,312 x 128 x 4     =  72 MB

SUBGRAPH SAMPLING, 5,000-node subgraph

input features = 5,000 x 128 x 4       = 2.6 MB
layer outputs  = 3 x 5,000 x 256 x 4   = 15.4 MB
edges kept     = 5,000 x 50 x (fraction inside subgraph)

The shapes of the two curves differ, and that is the whole decision. Neighbour sampling grows multiplicatively with depth: one more layer at fan-out 5 multiplies the working set by about five, and the jump from two layers to three above is 72 MB to 400 MB. Subgraph sampling is flat in depth — a fourth layer adds one more 5 MB activation tensor and nothing else, because the node set never grows. If you need depth, subgraph sampling is the only one of the two whose cost you can predict.

Choosing, and the biases you accept

• Two layers, labels scattered thinly across a huge graph: neighbour sampling. You want a batch of specific labelled seeds, the depth is shallow enough that the product stays small, and you never materialise a subgraph you did not need.
• Three or more layers, or dense labels: subgraph sampling. The multiplicative term is what kills you, and removing it is worth accepting a partition-shaped bias.
• Heavy-tailed degrees: cap or weight before you sample. A uniform neighbour sample from a node with two million neighbours is a fair sample of a neighbourhood that means almost nothing. Degree-capping high-degree nodes, or importance-sampling by inverse degree, changes the estimate you are computing — deliberately, and usually for the better.
• Cluster-GCN’s specific bias. Partitioning puts structurally similar nodes together, so a batch is not a random sample of the graph and its gradient is correlated within the batch. Combining several clusters per batch is the mitigation, and it is not optional.
• Inference is a separate decision. At prediction time there is no need to sample at all — layer-by-layer full-graph inference, computing every node’s layer-1 output before starting layer 2, costs one pass per layer over the whole graph and gives the exact answer. Sampling at inference introduces variance into a prediction that did not need it.

One failure that catches people: training with a sampled neighbourhood and serving with the full one is a train/serve mismatch. The model saw means over 25 neighbours and now receives means over 2,000. Mean aggregation is fairly robust to this; sum aggregation is not, and will shift its output scale by orders of magnitude. If you sample in training, either normalise so the aggregate is scale-free or sample the same way at inference. The capacity side of all this is worked in scaling a GNN to a billion nodes.