Two-level local/coarse GNN training

(1) Integration point: modify the GNN training loop before message passing, not the GNN layers themselves. Store a node partition and construct a one-hop halo for each partition so local updates retain boundary context. Use local subgraphs for 3–7 consecutive optimizer steps, then run one coarse step on 5–20% of nodes sampled within each partition. Aggregate gradients across workers before applying the coarse update.

(2) Pseudocode:

partition V into V_j; build G_j with one-hop halos
sample C_j from V_j with probabilities p_v; build G_c
for k = 1,...,K:
    for r = 1,...,R:
        for workers j in parallel:
            g_j = grad_theta loss(G_j, theta)
            theta_j = local_update(theta, g_j)
        theta = average_j(theta_j)
    g_c = sum_{v in C} (1/p_v) grad loss_v(G_c, theta) / |C|
    theta = coarse_update(theta, g_c)

Here \(U\) can be SGD, Adam, AdaGrad, or the paper's AG2m update. The inverse-probability factor is exactly the displayed estimator and prevents bias from nonuniform coarse sampling. (3) Compute partitions, halos, and coarse edges exactly; estimate only the stochastic coarse gradient empirically. Tune the local-to-coarse ratio \(R\), halo width, and sampling rate. (4) First cheap experiment: train a 3-layer GraphSAGE or GCN on Reddit, ogbn-products, or a synthetic power-law graph, comparing full-batch training and standard neighbor sampling at equal GPU-hours. Measure wall-clock time, peak memory, communication volume, validation accuracy, and loss versus processed edges. The expected signal is 2–4x lower message-passing cost or memory at matched accuracy, without persistent degradation caused by cross-partition edges.