Mini-Batch Sampling Math: NeighborLoader, GraphSAINT, and Cluster-GCN in TGraphX

Full-batch GNN training has a scaling problem that is structural, not just engineering. To compute one node's representation after L layers you need its L-hop neighbourhood, and in a well-connected graph that neighbourhood grows exponentially — the "neighbour explosion." For tensor-valued nodes the problem is worse, because each node in that exploding set is a [C, H, W] map rather than a small vector. TGraphX ships three classic answers, and they are worth understanding mathematically because they make different bias/variance trade-offs. This note works through the math of each, grounded in the source.

The introductory treatment is in Neighbor Sampling for Large Graphs; this article is the deeper, comparative one.

Strategy 1: NeighborLoader — sample the receptive field

NeighborLoader (in tgraphx/loaders.py) picks a small set of seed nodes and samples a bounded number of neighbours per hop around them, producing a GraphMiniBatch. The key idea is that loss is computed only on the seed nodes: the batch carries seed_node_ids, seed_local_indices, and a seed_logits(logits) helper that extracts just the seed rows. So the sampled neighbours exist only to compute the seeds' representations; they are not themselves training targets in that batch.

for batch in loader:                       # GraphMiniBatch per step
            logits = model(batch.node_features, batch.edge_index)
            loss = F.cross_entropy(batch.seed_logits(logits), batch.seed_y)
        

The math: bounding fan-out to s neighbours per hop caps the receptive field at O(s^L) instead of the full exponential growth, which is the explicit purpose of the method. The cost is that you compute an estimate of each seed's full-neighbourhood representation, so there is sampling variance. The loaders are documented to be deterministic when a seed is given, which makes that variance reproducible.

Strategy 2: GraphSAINT — sample subgraphs, then debias

GraphSAINT (tgraphx/graphsaint.py) takes a different route: instead of expanding around seeds, it samples an entire subgraph per step and trains on it as if it were the whole graph. TGraphX provides three samplers — GraphSAINTNodeSampler, GraphSAINTEdgeSampler, and GraphSAINTRandomWalkSampler — each of which .sample()s a Graph.

The subtlety is statistical. Sampling subgraphs biases the estimator: high-degree nodes and frequently-sampled edges are over-represented. GraphSAINT corrects this with normalisation coefficients, and TGraphX exposes estimate_norm_coefficients, described in the source as a Monte-Carlo estimate of how often each node and edge is sampled. Weighting each node's and edge's contribution by the inverse of its sampling probability yields an approximately unbiased estimate of the full-graph loss and aggregation — the central trick from the GraphSAINT paper. GraphSAINTLoader wires the sampler and coefficients together so you do not assemble them by hand.

Strategy 3: Cluster-GCN — partition, then batch clusters

Cluster-GCN (tgraphx/cluster_gcn.py) partitions the node set into clusters that keep most edges inside a cluster, then forms each mini-batch from one or more whole clusters. Because intra-cluster edges are preserved, message passing inside a batch is close to exact; only the relatively few cut edges between clusters are dropped per step. TGraphX offers four partitioners with different trade-offs:

The honest boundary is built into the code: the spectral partitioner is O(N³) and the framework restricts it to small graphs (the README notes a ≤ 4096-node cap and a warning). That is the right behaviour — a method that does not scale should refuse to pretend it does.

Choosing among them

There is no universally best choice; each trades a different error for memory savings. The right move is to pick by graph structure and then profile — none of these is faster in the abstract, and the package ships the mechanisms rather than a benchmark that would justify a blanket speed claim.

A decision flow

If you are unsure which sampler to start with, a short flow helps. Is your graph small enough to train full-batch? Then do that — sampling only adds variance you do not need. Is it a node-level task on a large graph with bounded degree? Start with NeighborLoader, since seed-node loss and bounded fan-out map cleanly onto node prediction. Is the graph dense, so neighbour fan-out explodes even with capped sampling? GraphSAINT's subgraph sampling with normalisation coefficients controls that variance directly. Does the graph have strong community structure? Cluster-GCN's partitions keep most edges intra-batch and waste little computation on cut edges.

These are starting points, not verdicts. The interaction between sampler, model depth, and graph structure is empirical, so measure validation accuracy and step time on your own data — and read the source for each sampler's exact bias — rather than assuming one approach dominates. The three loaders are interchangeable enough that trying two is usually cheap, and the comparison is often more informative than any rule.

Related guides

• Neighbor Sampling for Large Graphs in TGraphX
• Scaling Tensor Node Features: The Feature Store
• The TGraphX Graph and GraphBatch

Conclusion

TGraphX gives you three principled ways to train GNNs on graphs that do not fit in memory: bound the receptive field (NeighborLoader), sample-and-debias subgraphs (GraphSAINT with estimate_norm_coefficients), or partition into mostly-self-contained clusters (Cluster-GCN). Each has a clean mathematical story and an honest cost, and the source — including the spectral partitioner's size cap — is candid about where the methods stop being practical.