Source-linked AI summary

Minimal Variance Sampling with Provable Guarantees for Fast Training of Graph Neural Networks

Weilin Cong, Rana Forsati, Mahmut Kandemir, Mehrdad Mahdavi

arXiv:2006.13866v2cs.LGstat.ML

TL;DR

Large-scale GNN sampling must contend with variance from approximating node embeddings and estimating gradients, while existing adaptive sampling is complicated by the multi-level composite objective. The paper proposes decoupled variance reduction using embedding history and gradient-based node sampling; it reports faster convergence and better generalization even with smaller mini-batches.

  • Problem

    Sampling-based GNNs face embedding-estimation bias in their multi-level composite objective, complicating adaptive sampling and motivating reduction of both embedding and gradient variance.

    Source: §1 · p. 2
  • Method

    MVS-GNN uses historical node embeddings to reduce embedding variance and approximate gradient norms to sample nodes with minimal variance.

    Source: §1 · p. 2
  • Results

    MVS-GNN achieves faster convergence and better generalization than existing sampling methods, even with smaller mini-batches.

    Source: §6 · p. 9
  • Takeaways & Limitations

    The method’s reported improvements extend across sampling methods and large-scale real graph datasets.

    Source: §1 · p. 2

    The optimal importance distribution requires N derivatives at each step, and TensorFlow and PyTorch do not support per-example gradients within a mini-batch.

    Source: §4.2 · p. 5, §4.3 · p. 6

Abstract

from arXiv · show

Sampling methods (e.g., node-wise, layer-wise, or subgraph) has become an indispensable strategy to speed up training large-scale Graph Neural Networks (GNNs). However, existing sampling methods are mostly based on the graph structural information and ignore the dynamicity of optimization, which leads to high variance in estimating the stochastic gradients. The high variance issue can be very pronounced in extremely large graphs, where it results in slow convergence and poor generalization. In this paper, we theoretically analyze the variance of sampling methods and show that, due to the composite structure of empirical risk, the variance of any sampling method can be decomposed into \textit{embedding approximation variance} in the forward stage and \textit{stochastic gradient variance} in the backward stage that necessities mitigating both types of variance to obtain faster convergence rate. We propose a decoupled variance reduction strategy that employs (approximate) gradient information to adaptively sample nodes with minimal variance, and explicitly reduces the variance introduced by embedding approximation. We show theoretically and empirically that the proposed method, even with smaller mini-batch sizes, enjoys a faster convergence rate and entails a better generalization compared to the existing methods.

1 INTRODUCTION

Neighbor dependence makes large-scale GNN training difficult, so mini-batch sampling limits computation by aggregating sampled neighbors. The paper analyzes the resulting variance and proposes a decoupled strategy that reduces both its embedding and gradient components.

1 INTRODUCTION

  • Sampling addresses neighbor explosion by aggregating only sampled neighbors, but GNN training remains difficult because node representations depend on one another across layers.This recursive dependence grows exponentially with the number of layers.Source: §1 · p. 1
  • Increasing per-layer sample sizes can reduce variance, but computation and memory requirements constrain this approach.Adaptive sampling using gradients or representations is another proposed route to variance reduction.Source: §1 · p. 1
  • Existing adaptive sampling methods are difficult to extend to GNNs because estimating embeddings across layers introduces bias into the multi-level composite training objective.This bias complicates estimation of the optimal sampling distribution.Source: §1 · pp. 1–2
  • Sampling-based GNNs introduce embedding approximation variance from sampled neighbors and stochastic gradient variance from mini-batch gradient estimation.The paper shows the gradient bias is proportional to embedding approximation variance and vanishes as that variance approaches zero.Source: §1 · p. 2
  • MVS-GNN reduces embedding variance using historical node embeddings and samples nodes by minimizing variance with approximate gradient norms.The method uses optimization’s dynamic information and is theoretically shown to outperform uniform and static importance sampling.Source: §1 · p. 2
  • MVS-GNN achieves faster convergence and better accuracy by reducing stochastic-gradient variance, even with significantly smaller mini-batches.The paper reports empirical comparisons across large-scale real graph datasets and sampling methods.Source: §1 · p. 2

2 ADDITIONAL RELATED WORK

Prior GNN sampling methods trade off variance, connectivity, and sampling or preprocessing costs across node-wise, layer-wise, and subgraph approaches.

2 ADDITIONAL RELATED WORK

  • Node-wise methods restrict computation by sampling neighbors, but small samples can raise embedding variance and recursive sampling creates large sample complexity.VRGCN [4] uses historical activations to reduce variance, at the cost of additional graph convolutions.Source: §2 · p. 2
  • Layer-wise methods sample nodes independently or restrict candidates to upper-layer neighborhoods, but connectivity or per-layer distribution costs remain.FastGCN [3] needs large samples for cross-layer connectivity, while LADIES [33] tracks neighbors and recalculates importance distributions.Source: §2 · p. 2
  • Subgraph methods partition graphs or importance-sample nodes, with sensitivity to cluster size, costly partitioning, or disconnected sampled nodes.GraphSaint [28] uses normalization to eliminate bias and reduce variance, but independent node sampling cannot guarantee connectivity.Source: §2 · p. 2

3 PROBLEM STATEMENT

The section formulates sampled GNN training as a composite optimization problem and shows that sampling noise has both forward embedding-approximation and mini-batch gradient components.

3.1 Problem definition

  • A GNN layer propagates node features through the normalized Laplacian, and the problem is framed as semi-supervised node classification.The model learns per-layer weights by minimizing empirical loss over nodes.Source: §3.1 · p. 3
  • Sampling subsets of neighbors or layer nodes reduces the computation and memory costs of computing embeddings across GNN layers.Neighbor expansion otherwise grows exponentially with the number of layers.Source: §3.1 · p. 3

3.2 Variance analysis

  • Sampling creates forward embedding-approximation bias and mini-batch gradient variance, so reducing both is necessary to speed convergence.The stochastic gradient is biased because the exact intermediate composite embeddings are unavailable.Source: §3.2 · pp. 3–4
  • The paper motivates decoupled variance reduction because existing sampling strategies do not explicitly use dynamic training information to control embedding-approximation variance.The bias is analyzed in terms of per-layer embedding-approximation variance.Source: §3.2 · p. 4

4 ADAPTIVE MINIMAL VARIANCE SAMPLING

MVS-GNN uses optimization-time gradient information to sample nodes and embedding history to reduce embedding approximation variance. Its analysis and implementation describe minimal-variance sampling alongside checkpointed estimates and efficient per-sample-gradient calculation.

4.1 Decoupled variance reduction

  • MVS-GNN refreshes gradient and embedding estimates in large uniform checkpoints, then runs K small-batch updates using previous-layer embedding history to approximate current embeddings.Source: §4.1 · p. 4, §4 · p. 4
  • Against VRGCN, MVS-GNN samples once, keeps a constant node count per layer, and multiplies adjacency and embedding matrices once rather than twice per forward pass.Source: §4.1 · p. 5

4.2 Gradient-based minimal variance sampling

  • The ideal sampling distribution minimizes variance using per-node gradient norms, while the practical estimate uses last-layer stochastic gradients because exact per-step calculation needs N derivatives.Source: §4.2 · pp. 5–6
  • Importance sampling has lower stochastic-gradient variance than uniform sampling when per-sample gradient norms differ, with larger reductions as those differences grow.Source: §4.2 · p. 6

4.3 Implementation challenges

  • Because cited frameworks lack per-example gradients and per-node backward passes are inefficient, the implementation reuses one backward pass's intermediates, reducing gradient-calculation time despite extra storage.Source: §4.3 · p. 6

5 EXPERIMENTS

The experiments compare MVS-GNN with node-wise, layer-wise, and subgraph sampling baselines on node-classification datasets. They report higher accuracy with small batches, reduced gradient variance, and faster convergence, alongside evaluations of inner-loop intervals and training time.

Experiment setup

  • The experiments use semi-supervised node classification on Reddit, PPI, PPI-large, and Yelp, comparing MVS-GNN with GraphSage, VRGCN, LADIES, ClusterGCN, and GraphSaint.The setup uses two-layer GNNs by default, with K=20 and dataset-specific sampling and checkpoint settings.Source: §5 · pp. 6–7

Effectiveness of variance reduction

  • Figure 2 compares gradient variance and convergence for MVS-GNN and baselines on Reddit, PPI, PPI-large, and Yelp at batch size 512.The accompanying text says minimal variance sampling reduces mini-batch gradient variance and positively affects model performance.Source: Figure 2 · p. 7, §5 · p. 7

Effect of mini-batch size

  • MVS-GNN achieves significantly higher accuracy than the comparison methods even with small batch sizes.Table 3 reports test-set F1-micro across batch sizes; the supplied prose states the small-batch accuracy advantage without giving values.Source: §5 · p. 7, Table 3 · p. 8

Evaluation of total time

  • On PPI, the timing comparison covers sampling, training, and distribution-calculation time for a five-layer GNN with batch size 512.The paper states that MVS-GNN requires less time because it does not recursively sample neighbors for each layer and mini-batch node.Source: §5 · pp. 7–8, Table 4 · p. 8

Evaluation on inner-loop interval

  • On Reddit, Figure 3 compares gradient variance, training loss, and testing loss across inner-loop intervals K=10, 20, 30, and 40.The text notes that larger K can speed training but may leave the importance-sampling distribution stale.Source: Figure 3 · p. 8, §5 · p. 8

Evaluation on small mini-batch size

  • With all neighbors used to infer embeddings, the Cora, Citeseer, and Pubmed small-batch experiment isolates node-selection randomness and reports reduced gradient variance and faster convergence for MVS.The importance-sampling ratios are 50% for Pubmed and 100% for Cora and Citeseer, with distributions updated every 10 iterations.Source: §5 · p. 8

6 CONCLUSION

The conclusion decomposes sampling variance into embedding-approximation and stochastic-gradient components, then proposes a decoupled strategy to mitigate both. The authors report faster convergence and better generalization, including with smaller mini-batches, and identify evaluation on other GNN variants as future work.

Conclusion

  • Sampling variance is decomposed into embedding-approximation variance and stochastic-gradient variance because of the composite structure of empirical risk.Figure 4 compares gradient variance, training loss, and testing loss with small mini-batches on Cora, Citeseer, and Pubmed.Source: §6 · p. 8, Figure 4 · p. 9
  • The proposed decoupled strategy uses gradient information to sample nodes with minimal variance and explicitly reduces embedding-approximation variance.The strategy targets both components identified in the variance analysis.Source: §6 · p. 9
  • The authors report faster convergence and better generalization than existing sampling methods, even with smaller mini-batches.Source: §6 · p. 9
  • Empirical efficiency on graph-classification and attention-based GNN variants is left for future study.Source: §6 · p. 9

A PROOF OF LEMMA 3.1

The proof bounds stochastic-gradient bias by relating it to embedding-approximation variance across layers. It constructs the comparison through intermediate products that differ in one factor at a time.

A PROOF OF LEMMA 3.1

  • To bound the difference between sampled and full gradients, the proof adds and subtracts intermediate terms so adjacent products differ in at most one factor.Source: A PROOF OF LEMMA 3.1 · p. 11
  • The proof bounds the squared stochastic-gradient bias by a linear combination of per-layer embedding-approximation variances.Source: A PROOF OF LEMMA 3.1 · p. 11
  • The proof identifies each layer’s embedding-approximation variance as a term in that bound.Source: A PROOF OF LEMMA 3.1 · p. 11

B EMBEDDING APPROXIMATION VARIANCE ANALYSIS

The appendix analyzes embedding-approximation variance for sampled nodes at a layer. Its stated bound depends on average node degree and upper bounds on layer-specific quantities; it also relates MVS-GNN’s approximation variance to VRGCN with one-neighbor sampling.

B EMBEDDING APPROXIMATION VARIANCE ANALYSIS

  • The appendix section analyzes approximation variance for sampled nodes at the lth layer.Source: B EMBEDDING APPROXIMATION VARIANCE… · p. 11
  • For MVS-GNN sampling N_l nodes at layer l, Lemma B.1 gives an upper bound on V_l in terms of average node degree and layer-specific bounds.The passage identifies Δγ_l as an upper bound on a norm and β_l as an upper bound on the norm of a normalized-Laplacian row.Source: B EMBEDDING APPROXIMATION VARIANCE… · p. 11
  • Because only mini-batch nodes are guaranteed to be sampled in inner layers, MVS-GNN’s embedding-approximation variance is described as similar to VRGCN with neighbor sampling size s=1.Source: B EMBEDDING APPROXIMATION VARIANCE… · p. 12

C PROOF OF LEMMA 4.3

During Cora training, per-sampler gradients vary in their impact on the full-batch gradient. This motivates more frequent sampling of nodes with approximately large gradients to reduce mini-batch variance.

Loading 2006.13866v2…