GraphSAINT and Cluster-GCN: Scalable Sampling Compared

Training GNNs on large graphs — millions of nodes and billions of edges — is not straightforward with mini-batch gradient descent. The standard approach of stacking message-passing layers creates an exponential receptive field: a 2-layer GNN with 10 neighbors per node requires aggregating over 100 neighbors per target node; a 3-layer GNN requires 1,000. Full-batch training becomes infeasible, and naive mini-batching produces biased gradient estimates because the boundary nodes have no neighbors in the mini-batch.

Two methods dominate scalable GNN training: GraphSAINT (Zeng et al., 2020) uses random subgraph sampling, and Cluster-GCN (Chiang et al., 2019) uses graph partitioning. Both are implemented in TGraphX (tgraphx.graphsaint and tgraphx.cluster_gcn). This article explains both methods, their bias-variance tradeoffs, and when each is the better choice.

The Core Problem: Neighborhood Explosion

When a GNN layer aggregates over a node's neighborhood, and the next layer does the same again, the set of nodes that contribute to a single target node's embedding grows exponentially with depth. Computing the full K-hop neighborhood for a mini-batch is prohibitively expensive for dense graphs.

Two complementary solutions exist:

1. Neighbor sampling: For each node in the mini-batch, sample a fixed number of neighbors per layer (NeighborSampler, GraphSAGE sampling). This limits the expansion but introduces bias at the boundary.

2. Subgraph sampling: Sample an entire subgraph and run the GNN on the full subgraph. The GNN sees complete local structure within the sample but may not see connections outside it.

Both GraphSAINT and Cluster-GCN fall into the second category, but they differ in how they sample subgraphs.

Prerequisites

Before reading this article, you should understand:

• Standard GNN mini-batch training concepts
• The neighborhood explosion problem (sketched above)
• Basic graph partitioning at a conceptual level

For a related comparison focused on neighbor-sampling specifically, see neighbor sampling with TGraphX for large graphs.

Cluster-GCN: Graph Partitioning for Batching

Cluster-GCN partitions the graph into densely connected clusters (using METIS or other partitioners) and trains each cluster as a mini-batch. Within a cluster, the GNN can aggregate over all intra-cluster neighbors, which are fully present in the batch.

The key insight is that densely connected clusters minimize the number of inter-cluster edges that are cut. During training, those cut edges are simply ignored, introducing a bias that is small when the clusters are internally dense.

Training step:
        1. Partition G into K clusters: C_1, ..., C_K
        2. For each mini-batch, sample m clusters: B = {C_i1, ..., C_im}
        3. Construct subgraph G_B = (nodes in B, all intra-cluster edges)
        4. Run full GNN on G_B, compute loss on labeled nodes in B
        5. Update gradients
        

The larger the clusters, the fewer inter-cluster edges are cut, but the larger the GPU memory requirement per batch. Smaller clusters reduce memory but increase bias.

Using Cluster-GCN in TGraphX

import torch
        import torch.nn as nn
        import torch.nn.functional as F
        from tgraphx.cluster_gcn import ClusterGCN, ClusterLoader
        from tgraphx.layers.sage import TensorGraphSAGELayer
        
        # Define a simple GNN model
        class NodeClassifier(nn.Module):
            def __init__(self, in_dim, hidden_dim, num_classes):
                super().__init__()
                self.conv1 = TensorGraphSAGELayer(in_dim, hidden_dim)
                self.conv2 = TensorGraphSAGELayer(hidden_dim, num_classes)
        
            def forward(self, x, edge_index):
                x = F.relu(self.conv1(x, edge_index))
                return self.conv2(x, edge_index)
        
        # Full graph
        num_nodes = 100000
        x = torch.randn(num_nodes, 128)
        edge_index = torch.randint(0, num_nodes, (2, 500000), dtype=torch.long)
        labels = torch.randint(0, 7, (num_nodes,))
        train_mask = torch.rand(num_nodes) < 0.6
        
        # Set up ClusterGCN with METIS partitioning
        cluster_gcn = ClusterGCN(
            x=x,
            edge_index=edge_index,
            labels=labels,
            num_parts=200,          # number of clusters
            batch_size=20,          # number of clusters per mini-batch
        )
        
        loader = cluster_gcn.loader(shuffle=True)
        model = NodeClassifier(128, 256, 7)
        optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
        
        for epoch in range(5):
            total_loss = 0.0
            for batch in loader:
                bx, bedge, blabels, bmask = batch
                optimizer.zero_grad()
                out = model(bx, bedge)
                loss = F.cross_entropy(out[bmask], blabels[bmask])
                loss.backward()
                optimizer.step()
                total_loss += loss.item()
            print(f"Epoch {epoch}: loss={total_loss:.4f}")
        

GraphSAINT: Subgraph Sampling with Normalization

GraphSAINT takes a different approach: it samples subgraphs randomly and applies a normalization correction to make the gradient estimator unbiased. Three sampling strategies are provided:

• Node sampling: Sample nodes uniformly at random, then include all edges between sampled nodes.
• Edge sampling: Sample edges with probability proportional to their product of node degrees, then include all nodes of sampled edges.
• Random walk sampling: Perform multiple random walks, collect all visited nodes, include all edges between them.

The normalization corrects for the fact that different nodes and edges have different probabilities of appearing in a subgraph. Without normalization, high-degree nodes would be over-represented in the gradient.

Normalization factor for node v: α_v = 1 / p(v ∈ subgraph)
        Normalized loss: L = Σ_v (1/α_v) · l(ŷ_v, y_v)
        

Using GraphSAINT in TGraphX

import torch
        import torch.nn.functional as F
        from tgraphx.graphsaint import GraphSAINT, SAINTSampler
        from tgraphx.layers.sage import TensorGraphSAGELayer
        import torch.nn as nn
        
        class NodeClassifier(nn.Module):
            def __init__(self, in_dim, hidden_dim, num_classes):
                super().__init__()
                self.conv1 = TensorGraphSAGELayer(in_dim, hidden_dim)
                self.conv2 = TensorGraphSAGELayer(hidden_dim, num_classes)
        
            def forward(self, x, edge_index):
                x = F.relu(self.conv1(x, edge_index))
                return self.conv2(x, edge_index)
        
        num_nodes = 100000
        x = torch.randn(num_nodes, 128)
        edge_index = torch.randint(0, num_nodes, (2, 500000), dtype=torch.long)
        labels = torch.randint(0, 7, (num_nodes,))
        
        # GraphSAINT with random walk sampling
        saint = GraphSAINT(
            x=x,
            edge_index=edge_index,
            labels=labels,
            sampler="walk",          # "node", "edge", or "walk"
            walk_length=3,           # for walk sampler
            num_walks=200,           # number of walks per batch
            sample_coverage=100,     # used to estimate normalization factors
        )
        
        loader = saint.loader(batch_size=5000, shuffle=True)
        model = NodeClassifier(128, 256, 7)
        optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
        
        for epoch in range(5):
            total_loss = 0.0
            for batch in loader:
                bx, bedge, blabels, norm_weights = batch
                optimizer.zero_grad()
                out = model(bx, bedge)
                # Apply normalization weights to loss
                loss = (F.cross_entropy(out, blabels, reduction='none') * norm_weights).mean()
                loss.backward()
                optimizer.step()
                total_loss += loss.item()
            print(f"Epoch {epoch}: loss={total_loss:.4f}")
        

The norm_weights tensor contains the correction factors for each node in the subgraph. Using them makes the gradient estimator unbiased in expectation.

Head-to-Head Comparison

The right choice often depends on graph structure. If your graph has natural community structure (citation networks, social networks), Cluster-GCN's partitioning aligns with that structure and inter-cluster edges are genuinely less important. If your graph is more random or expander-like (few natural clusters), GraphSAINT's unbiased estimator tends to give better gradient quality.

Bias-Variance Analysis

Both methods trade off bias and variance:

Cluster-GCN introduces bias by ignoring inter-cluster edges. This bias is reduced by using more clusters or by including edges between neighboring clusters (multi-cluster batching). However, multi-cluster batching increases memory. The bias is not corrected by training longer — it is a systematic error due to the sampling scheme.

GraphSAINT corrects bias through normalization but variance in the gradient estimator remains. Variance decreases with larger subgraphs. Edge sampling tends to produce denser subgraphs than node sampling, which reduces variance at similar computational cost. Walk sampling produces connected subgraphs that are the best approximation of the local neighborhood structure.

In practice, the distinction matters most in the early training epochs. Late in training, when gradients are small, both methods converge to similar parameter values if the learning rate and batch size are tuned appropriately.

Limitations and Honest Notes

GraphSAINT normalization computation can be expensive. The sample_coverage parameter controls how many subgraphs are sampled to estimate normalization factors. Too few samples produces inaccurate factors; too many is computationally expensive. The default of 50–100 is a reasonable starting point.

Cluster-GCN's partitioning is a preprocessing cost. METIS partitioning can take minutes to hours on very large graphs. This cost is paid once before training begins.

Both methods assume the GNN fits in memory on a per-batch basis. If individual clusters or subgraphs are too large for GPU memory, further sub-sampling within the batch is required.

Neither method is strictly superior across all benchmarks. The original papers report competitive results on different datasets. For reliable comparison, run both on your specific dataset and hardware configuration.

For graphs with tensor-valued node features, memory cost per batch scales with feature spatial size. Use smaller spatial features or reduce batch size accordingly.

Frequently Asked Questions

Can I use GraphSAINT and Cluster-GCN with any GNN architecture?
Yes — both methods produce (x_subgraph, edge_index_subgraph) pairs that can be fed to any GNN layer. They are agnostic to the GNN architecture.

How does GraphSAINT compare to NeighborSampler?
NeighborSampler (used in GraphSAGE mini-batch training) samples neighbors layer by layer and introduces bias at the boundary of each sampled subtree. GraphSAINT samples a complete subgraph and corrects for sampling bias via normalization. See the neighbor sampling guide for details on NeighborSampler.

Is graph partitioning recomputed every epoch?
No. Cluster-GCN computes the partition once and reuses it for all training epochs. Only the order in which clusters are batched changes per epoch.

What if my graph does not have clear community structure?
Use GraphSAINT with walk or edge sampling. Cluster-GCN's bias is small on community-structured graphs but can be significant on random graphs.

Where is the source?
GitHub, package: PyPI, preprint: arXiv:2504.03953.

Practical Tuning Guide

Both methods have hyperparameters that significantly affect convergence quality. Getting them right avoids weeks of debugging slow training runs.

For Cluster-GCN:

• Number of partitions (num_parts): A good starting point is num_parts = N / 200 — i.e., roughly 200 nodes per cluster. Too few large clusters → memory pressure per batch. Too many small clusters → high inter-cluster edge bias.
• Batch size (clusters per batch): Start with 10–20 clusters per batch. Increasing the batch size reduces gradient variance at the cost of higher memory.
• Multiple epochs per partition: Cluster-GCN typically uses a different random ordering of clusters per epoch. Do not shuffle within a cluster.

For GraphSAINT:

• Sampler type: Walk sampling generally produces the best subgraph quality (high connectivity, representative of the full graph). Edge sampling is faster but produces sparser subgraphs.
• Walk length: For 2-layer GNNs, walk length 2–3 is usually sufficient. Longer walks produce larger and denser subgraphs.
• sample_coverage: A value of 50–100 provides good normalization estimates without excessive pre-computation.
• Batch size (nodes per subgraph): Aim for 2–5% of the total graph per batch. For a 100k-node graph, batches of 2,000–5,000 nodes are typical.

Using GraphSAINT and Cluster-GCN for Graph Classification

Both methods are primarily designed for node-level tasks on a single large graph. For graph classification (many small graphs), standard batching with GraphBatch is more appropriate. However, if individual graphs in your collection are very large (thousands of nodes per graph), you can apply these sampling methods per-graph before batching:

from tgraphx import GraphBatch
        
        # For very large individual graphs in a graph classification dataset:
        # Sample a representative subgraph from each graph before batching
        def sample_subgraph(g, budget=500):
            """Sample ~budget nodes from graph g via random walk."""
            from tgraphx.graphsaint import SAINTSampler
            sampler = SAINTSampler(g.edge_index, g.num_nodes, method="walk",
                                   walk_length=3, num_walks=budget // 10)
            sub_edge_index, sub_nodes = sampler.sample()
            sub_x = g.node_features[sub_nodes]
            return sub_x, sub_edge_index
        
        # Process each graph in a batch
        graphs = [...]  # list of Graph objects, each with ~10k nodes
        batched = GraphBatch.from_graphs([
            (lambda g, sg=sample_subgraph(g): sg)(g) for g in graphs
        ])
        

This pattern is not recommended when graphs are small enough to fit in memory directly. Use GraphBatch.from_graphs() directly for standard graph classification.

Reproducibility in Scalable Training

Results from GraphSAINT and Cluster-GCN can vary significantly across runs due to:

• Random subgraph sampling (GraphSAINT)
• Random cluster ordering per epoch (Cluster-GCN)
• Stochastic gradient descent itself

For reproducible experiments, fix seeds and document the sampling configuration:

from tgraphx.reproducibility import set_reproducible
        
        with set_reproducible(seed=42):
            saint = GraphSAINT(
                x=x,
                edge_index=edge_index,
                labels=labels,
                sampler="walk",
                walk_length=3,
                num_walks=200,
                sample_coverage=100,
            )
            # ... training ...
        

For multi-worker data loading, also set num_workers=0 during debugging to ensure deterministic ordering. For published results, report mean ± std over at least three seeds. See the GNN research reproducibility guide for best practices.