Scalability to large networks

"My forward pass wanted the whole city in memory at once. The city, it turned out, did not fit. So I learned to look at neighborhoods."

A Right-Sized AI Agent

The scaling wall: neighborhood explosion, not node count

What makes large-graph training hard is not simply that there are many nodes. It is that message passing couples them. In a full-graph forward pass, layer \(\ell\) computes, for every node \(i\),

so to get one node's output at the final layer you need its neighbors' representations at the previous layer, which need their neighbors, and so on. Why this is fatal: with a \(k\)-layer network and average degree \(\bar{d}\), computing the loss for a single node touches on the order of \(\bar{d}^{\,k}\) nodes. On a traffic graph with \(\bar{d} \approx 6\) and \(k = 4\) that is already more than a thousand nodes per target, and the receptive fields of nearby targets overlap so heavily that the honest ceiling is the full graph. The activations you must hold for backpropagation therefore scale with \(N \cdot F\) (nodes times feature width) per layer, and for \(N\) in the hundreds of thousands that alone overruns a commodity GPU. When the wall bites: any network past a few thousand nodes, or any spatiotemporal graph where each node also carries a long time series, since the temporal dimension multiplies the per-node cost (see Section 54.3).

Neighbor sampling: bound the branching (GraphSAGE)

What GraphSAGE does is refuse to look at all of a node's neighbors. At each layer it samples a fixed number \(S_\ell\) of neighbors per node, so a \(k\)-layer forward pass over a mini-batch of target nodes expands into a bounded tree of size \(\prod_{\ell} S_\ell\) per target rather than \(\bar{d}^{\,k}\). Why this is sound: the neighborhood aggregation \(\bigoplus\) is an expectation over neighbors (a mean, or a normalized sum), and a random sample is an unbiased-in-expectation estimate of that aggregate, so the model still learns the right function while paying a controlled, constant per-node cost. Choose \(S_1 = 25, S_2 = 10\) and each target's compute tree has at most 250 leaves regardless of whether the graph has ten thousand nodes or ten million. How you deploy it: build mini-batches of target nodes, sample their sampled subtrees on the fly, and train with ordinary stochastic gradient descent, exactly as you would batch time-series windows in Chapter 14. When to prefer it: graphs where degree is high and skewed (a few hub sensors wired to hundreds of others), because a fixed sample size tames the hubs that would otherwise dominate both memory and the aggregate.

Subgraph training: bound the extent (Cluster-GCN and GraphSAINT)

What these methods do is cut the graph into pieces and train on one piece at a time. Cluster-GCN runs a graph-partitioning step (typically METIS) to group densely-connected nodes into clusters, then treats each cluster (or a small union of clusters) as a self-contained mini-batch: message passing happens only within the subgraph, so there is no neighbor explosion because there are simply no edges leaving the batch during that step. GraphSAINT instead samples a subgraph by a random walk or node/edge sampler and corrects the resulting bias with normalization coefficients on nodes and edges, so the expected gradient matches the full-graph gradient. Why subgraph methods win over neighbor sampling on some workloads: neighbor sampling still recomputes overlapping subtrees for nearby targets and can suffer high variance when samples are small, whereas a subgraph batch amortizes one clean forward pass across all its nodes at once. How the tradeoff reads: partitioning removes the edges that cross cluster boundaries, so you lose some long-range coupling per step; you recover it by re-partitioning across epochs or by unioning several clusters per batch so boundaries land in different places each time. When to reach for it: very large, relatively homogeneous sensor graphs (a national metering grid) where clusters correspond to real regions (a substation's downstream feeders, a district metering area), so the partition respects physics instead of fighting it.

The snippet below trains a 2-layer GraphSAGE model on a large sensor graph using PyTorch Geometric's neighbor-sampling loader, showing how few lines separate a toy full-graph script from one that scales.

import torch
from torch_geometric.nn import SAGEConv
from torch_geometric.loader import NeighborLoader

class SensorSAGE(torch.nn.Module):
    def __init__(self, in_dim, hid, out_dim):
        super().__init__()
        self.conv1 = SAGEConv(in_dim, hid)
        self.conv2 = SAGEConv(hid, out_dim)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index).relu()
        return self.conv2(x, edge_index)

# `data` is a large PyG graph: data.x = [N, F] sensor features, data.y targets.
# fan-out (25, 10): sample 25 neighbors at layer 1, 10 at layer 2 per seed node.
loader = NeighborLoader(data, num_neighbors=[25, 10],
                        batch_size=1024, shuffle=True,
                        input_nodes=data.train_mask)

model = SensorSAGE(data.num_features, 128, data.num_classes)
opt = torch.optim.Adam(model.parameters(), lr=1e-3)

for batch in loader:                       # each batch is a bounded subgraph
    opt.zero_grad()
    out = model(batch.x, batch.edge_index)
    # loss only on the seed nodes of this batch (first batch.batch_size rows)
    seed = slice(0, batch.batch_size)
    loss = torch.nn.functional.cross_entropy(out[seed], batch.y[seed])
    loss.backward()
    opt.step()

Historical embeddings and linear-time attention: bound the freshness cost

What historical-embedding methods (GNNAutoScale and its relatives) exploit is that a node's representation changes slowly between gradient steps. Instead of recomputing every out-of-batch neighbor's embedding, they store the last computed embedding for each node in a cache (often in CPU memory or on disk) and reuse it as a stale-but-serviceable stand-in during message passing. Why this recovers full-graph accuracy that plain sampling can lose: the target node still aggregates over its entire real neighborhood, not a random subset, so there is no sampling variance; the only error is the small staleness of cached neighbors, which the theory bounds and which shrinks as training converges. How the bookkeeping works: after each forward pass you write the freshly computed embeddings back into the cache, so the pool of stale values is continuously refreshed. When it is the right tool: when accuracy matters more than the last increment of speed, and you have cheap host memory to hold one embedding vector per node.

A complementary lever attacks the other axis, the temporal one. The BigST model you met in Section 54.3 replaces the quadratic spatial attention of a vanilla spatiotemporal transformer with a linearized attention whose cost grows linearly in the number of nodes, so a single model can forecast over more than ten thousand sensors without the \(O(N^2)\) attention matrix that makes full attention (the mechanism of Chapter 15) impossible at that scale. The general principle: wherever an operation costs \(O(N^2)\) in the sensor count (dense attention, full pairwise distances, a dense adjacency matmul), find the linear-time surrogate before you find a bigger GPU.

Inference and deployment at scale

Training is only half the budget; a deployed sensor-network model must also serve. Two facts help. First, inference has no backward pass, so activation memory drops sharply and you can afford larger neighborhoods or even a full-graph pass in chunks. Second, at inference you frequently need to update only the nodes whose inputs just changed: a streaming sensor network delivers new readings continuously (the streaming setting of Chapter 60), and message passing lets you recompute the affected \(k\)-hop region rather than the whole graph. When even that is too much for an edge gateway, the compression toolkit of Chapter 59 applies unchanged: quantize the weights, prune the least-used edges, and distill the large GNN into a smaller student. And throughout, keep the leakage-safe discipline of Chapter 5: when you partition a graph for training, make sure your train and test node splits do not leak across cluster boundaries, or an over-optimistic score will follow.

What's Next

In Section 54.6, we take these scalable models to the field, walking through how graph neural networks are actually deployed on traffic, energy, water, and industrial grids: what the graph looks like in each domain, which scaling method each network's size and topology demands, and the operational realities (missing sensors, reconfigured topology, regulatory constraints) that separate a benchmark score from a system a utility will trust.