PaperlayerView exact revision on arXiv

Source-linked AI summary

GNNAutoScale: Scalable and Expressive Graph Neural Networks via Historical Embeddings

Matthias Fey, Jan E. Lenssen, Frank Weichert, Jure Leskovec

arXiv:2106.05609v1cs.LG

TL;DR

Large graphs make full-batch GNN training difficult because neighbor explosion and GPU-memory limits force scalable methods to trade away edges or expressiveness. GNNAutoScale uses historical embeddings to prune computation while retaining the underlying message-passing model, and its theory and experiments show scalable, expressive training with substantially lower memory use.

  • Problem

    Scaling expressive message-passing GNNs is difficult because neighbor explosion and GPU-memory limits motivate edge sampling or other approaches that can restrict information or expressiveness.

  • Method

    GAS prunes computation sub-trees for mini-batches and fills out-of-batch dependencies with historical embeddings, independently of the underlying message-passing implementation.

  • Results

    GAS preserves the expressivity properties of the used GNN under stated conditions and achieves performance comparable to non-scalable full-batch models with much lower GPU-memory use.

  • Takeaways & Limitations

    GAS provides a general route to training deep and expressive GNNs on large graphs without sub-sampling edges.

  • Takeaways & Limitations

    Experiments focus on node-level tasks, while scaling edge-level and graph-level tasks remains empirically unverified.

Abstract

from arXiv · show

We present GNNAutoScale (GAS), a framework for scaling arbitrary message-passing GNNs to large graphs. GAS prunes entire sub-trees of the computation graph by utilizing historical embeddings from prior training iterations, leading to constant GPU memory consumption in respect to input node size without dropping any data. While existing solutions weaken the expressive power of message passing due to sub-sampling of edges or non-trainable propagations, our approach is provably able to maintain the expressive power of the original GNN. We achieve this by providing approximation error bounds of historical embeddings and show how to tighten them in practice. Empirically, we show that the practical realization of our framework, PyGAS, an easy-to-use extension for PyTorch Geometric, is both fast and memory-efficient, learns expressive node representations, closely resembles the performance of their non-scaling counterparts, and reaches state-of-the-art performance on large-scale graphs.

1. Introduction

GNNAutoScale addresses the difficulty of scaling expressive message-passing GNNs to large graphs. It uses historical embeddings to prune computation while retaining neighborhood information and the expressiveness of the underlying GNN.

  • Motivation: Large-graph GNN training is constrained by GPU memory and neighbor explosion, which make deeper architectures impractical.Mini-batch gradients remain expensive because node dependencies grow exponentially across layers and the full computation graph must be stored.
  • Limitations of prior approaches: Edge-sampling methods control neighbor explosion but may drop meaningful edges, remain shallow, restrict operators, or reduce expressiveness.Their applicability across the broad range of available GNN architectures remains an open question.
  • Limitations of prior approaches: Layer-isolated training accounts for all edges but cannot infer complex interactions across consecutive layers.
  • GNNAutoScale: GAS prunes each mini-batch computation graph to the batch and its direct 1-hop neighbors, using historical embeddings to fill out-of-batch dependencies.This makes memory consumption constant with respect to input node size while retaining available neighborhood information.
  • GNNAutoScale: GAS separates scalability from message passing, supporting arbitrary message-passing GNNs rather than requiring a specialized operator.The framework theoretically preserves the expressivity properties of the selected GNN under stated conditions.
  • Practical realization: PyGAS provides a PyTorch Geometric implementation that converts common and custom GNNs into scalable variants.Experiments report performance matching non-scalable full-batch equivalents while using orders of magnitude less GPU memory.

2. Scalable GNNs via Historical Embeddings

GAS replaces recursive computation for out-of-mini-batch neighbors with historical embeddings, preserving the general message-passing formulation while reducing memory demands. Its design uses local current-batch information, stored historical information, and selective history access to support scalable training and inference.

  • Message passing: A GNN updates each node by combining its current embedding with aggregated messages from neighboring nodes through parameterized message and aggregation functions.The formulation supports permutation-invariant aggregators such as sum, mean, and maximum across a wide range of GNN operators.
  • Historical embeddings: GAS partitions neighbors into current-batch nodes and out-of-batch nodes, approximating the latter with historical embeddings from previous training iterations.Newly computed embeddings are pushed to history for later iterations, while stored values are pulled for out-of-batch computation.
  • Memory scaling: Historical embeddings prune entire computation sub-trees, so required information remains local to each batch instead of being recomputed across the full graph.This avoids storing embeddings for the entire input graph as depth increases.
  • Memory scaling: GAS memory scales linearly with the number of layers for a batch and its immediate neighborhood, while most historical data can reside outside GPU memory.
  • Advantages: GAS processes all available graph edges rather than dropping edges, supporting lower-variance and more accurate estimations.For one epoch and layer, each edge is processed once, giving time complexity O(|E|) on par with full-batch training.
  • Advantages: Historical embeddings reduce inference complexity to a constant factor by directly using last-layer historical embeddings for test-node predictions.
  • Advantages: GAS provides theoretical guarantees concerning approximation errors and gradient bounds when model weights are fixed.The cited bound becomes less restrictive with depth because earlier errors propagate through later layers.

3. Approximation Error and Expressiveness

GAS bounds approximation error from historical embeddings and shows that, under suitable separation and staleness conditions, scalable GNNs can retain WL-level expressiveness while accounting for approximation effects.

  • Approximation Error: Historical-embedding error is decomposed into input-estimation error and history staleness, with bounds depending on network smoothness and neighborhood size.The analysis assumes Lipschitz-continuous message and update functions and tracks errors across layers.
  • Approximation Error: Deeper GNNs can accumulate larger inaccuracies because errors propagate through successive layers and depend exponentially on Lipschitz constants and neighborhood size.The paper proposes tightening the bounds to support deep and nonlinear GNNs.
  • Scope: GAS does not guarantee convergence to the same optimum because it targets arbitrary GNNs solving non-convex problems.The framework is therefore theoretically expressive without implying identical optimization outcomes.
  • Expressiveness: Unlike edge-sampling methods, GAS accounts for approximation error while retaining theoretical access to message-passing functions as expressive as the WL test.The framework separates scalability and expressiveness from edge-dropping regularization.
  • Expressiveness: If exact inputs are sufficiently separated and historical embeddings remain sufficiently close, historical-based operators can distinguish equal from non-equal inputs.The argument relies on non-overlapping error neighborhoods around exact inputs.
  • Expressiveness: Under Lemma 4 conditions, multi-layer historical-embedding GNNs can distinguish structures distinguishable by the WL test during training.This connects scalable message passing to structural reasoning without requiring edge dropping.
  • Tightening the Bounds: GAS reduces history access by minimizing inter-connectivity between mini-batches, using graph clustering to favor intra-cluster neighbors.METIS-style clustering can be computed in O(|E|) time and applied once during preprocessing.
  • Tightening the Bounds: An auxiliary loss can encourage locally Lipschitz intermediate functions, making their outputs less sensitive to small input perturbations.This is presented as a practical way to control approximation error.

4. Related Work

Related work addresses GNN scalability through historical embeddings, concurrent execution, and subgraph sampling, but these approaches differ in how they handle neighborhoods, edges, and transfer overhead.

  • Historical Embeddings: Historical embeddings reduce neighbor-sampling variance and avoid recursively sampling large neighborhoods at every layer.The approach originates in VR-GCN and was later simplified to one-shot sampling.
  • Execution Strategies: Concurrent mini-batch execution overlaps memory transfers with computation, avoiding the I/O bottleneck of serial mini-batches and matching full-batch performance.An additional worker thread enables the overlap and provides a two-times improvement over serial execution.
  • Subgraph Sampling: CLUSTER-GCN uses clustering to form meaningful subgraphs but restricts message passing to intra-connected nodes, potentially ignoring information outside the current mini-batch.GAS instead uses clustering to reduce history accesses while retaining broader neighborhood information.

5. PyGAS: Auto-Scaling GNNs in PyG

PyGAS packages GAS as an accessible PyTorch Geometric extension, while optimizing historical-embedding transfers so scalable models can execute efficiently on large graphs.

  • Interface: PyGAS converts common and custom PyTorch Geometric GNNs into scalable variants through an easy-to-use interface.It also provides a deterministic test bed for evaluating models on large-scale graphs.
  • Transfer Optimization: PyGAS asynchronously transfers historical embeddings at each optimization step so GPUs do not remain idle during memory movement.The implementation uses non-blocking device transfers for pulling and pushing histories.
  • Transfer Optimization: The non-blocking transfer scheme can be twice as fast as serial non-overlapping execution.Custom C++/CUDA code avoids Python’s global interpreter lock.

6. Experiments

Experiments show that GAS preserves full-batch predictive performance while reducing memory demands and supporting efficient training of deep, expressive GNNs on large graphs.

  • 6.1. GAS resembles full-batch performance: GAS reaches full-batch quality for deep GCNII and expressive GIN models, whereas the naive history baseline falls short.The evaluation uses 64-layer GCNII and 4-layer GIN models on CORA and CLUSTER.
  • 6.1. GAS resembles full-batch performance: GAS performance closely matches full-batch training across evaluated GNN models and small transductive datasets.The comparison covers four models and multiple datasets across 20 initializations.
  • 6.1. GAS resembles full-batch performance: Both minimizing inter-connectivity and applying regularization improve GCNII performance relative to full-batch training, with their combination strongest.Table 2 reports these improvements in percentage points.
  • 6.2. GAS is fast and memory-efficient: GAS uses low GPU memory while incorporating all available neighborhood information in one optimization step; its memory grows linearly with depth.CLUSTER-GCN uses lower memory but averages only ≈23% of available receptive-field information.
  • 6.2. GAS is fast and memory-efficient: Concurrent history transfer nearly eliminates I/O overhead, while serial transfer can increase runtime by up to 350%.For typical inter-/intra-connectivity ratios, added aggregation may increase runtime by up to 25%.
  • 6.2. GAS is fast and memory-efficient: GAS is faster and uses less memory than GTTF because GTTF’s recursive neighborhood construction scales exponentially with GNN depth.The comparison uses a four-layer GCN with equal mini-batch and receptive-field sizes.
  • 6.3. GAS scales to large graphs: Deep and expressive GAS models advance state-of-the-art performance on REDDIT and FLICKR and perform equally well on PPI.The evaluated models include GCNII and PNA, which are difficult to scale with full-batch training.

7. Conclusion and Future Work

GAS scales arbitrary message-passing GNNs to large graphs without sub-sampling edges and supports scalable training of deep and expressive models. The approach is also technically applicable to edge-level and graph-level tasks, although those settings still require empirical verification.

  • GAS scales arbitrary message-passing GNNs to large graphs without sub-sampling edges.
  • The framework can train deep and expressive GNNs in a scalable fashion.
  • GAS is technically able to scale edge-level and graph-level GNN training, but this remains empirically unverified.
  • Future work includes combining GAS with distributed training and accessing histories from disk rather than CPU memory.

8. Proofs

The proofs bound historical-embedding errors under Lipschitz continuity and establish conditions under which GAS preserves the expressiveness of the original GNN.

  • Additional layers can amplify approximation errors because earlier errors propagate through later layers.
  • The error bound depends on Lipschitz constants, neighborhood size, and input and historical-embedding errors.
  • Sum aggregation introduces a neighborhood-size factor, whereas mean and max aggregation yield tighter bounds.
  • Theorem 2 bounds approximation error when historical embeddings remain sufficiently fresh and the GNN functions are Lipschitz continuous.
  • Under suitable conditions, a GNN as expressive as the WL test can retain distinguishable node representations with historical embeddings.

9. Algorithm

The GAS mini-batch algorithm processes current-batch nodes and their neighbors while exchanging embeddings through push and pull operations across layers.

  • Algorithm 1 takes a graph, input node features, number of batches, and number of layers as inputs.
  • For each intermediate layer, GAS computes updated embeddings for current-batch nodes using their direct neighbors.
  • The algorithm pushes updated embeddings and pulls embeddings for nodes outside the current batch.

10. GNN Operators

The paper evaluates several message-passing operators, including GCN, GAT, GIN, PNA, APPNP, and GCNII, whose aggregation and propagation mechanisms differ.

  • GCN: GCN uses symmetrically normalized mean aggregation followed by a linear transformation.
  • GAT: GAT performs anisotropic aggregation using learnable attention coefficients for normalization.
  • APPNP: APPNP first predicts node labels from features and then smooths those predictions through propagation.
  • GCNII: GCNII extends APPNP with trainable propagation and initial residual connections.
  • GIN: GIN combines sum aggregation with MLPs to obtain a maximally powerful GNN operator.
  • PNA: PNA combines multiple aggregators with degree-scalers to capture graph structural properties.

11. PyGAS Programming Interface

PyGAS auto-scales a PyTorch Geometric GCN with minimal code changes by adding scalable history management and push-and-pull updates. The framework also uses graph clustering and Lipschitz continuity to reduce history staleness and preserve expressive-model performance.

  • Programming interface: ScalableGNN creates history embeddings and provides concurrent history access through push_and_pull().The forward pass applies push-and-pull history updates after each intermediate convolution.
  • Ablation studies: METIS-based graph clustering minimizes inter-connectivity between mini-batches, reducing history accesses and tightening approximation errors.The ablation compares inter-/intra-connectivity ratios across random and METIS-based mini-batch sampling.
  • Programming interface: PyGAS converts a full-batch PyTorch Geometric GCN into a scalable model with minimal changes.The scalable version inherits from ScalableGNN, while Listing 2 identifies changed and newly added lines.
  • Ablation studies: Combining reduced inter-connectivity with Lipschitz continuity helps a 4-layer GIN resemble full-batch performance on CLUSTER.The reported ablation evaluates the combined techniques for expressive models with highly nonlinear message passing.
  • Datasets and tasks: The evaluation covers node-classification, community-detection, and graph-based prediction tasks across citation, Wikipedia, coauthor, product, social, protein, image, and large-scale datasets.Listed datasets include CORA, CITESEER, PUBMED, WIKI-CS, COAUTHOR-CS, COAUTHOR-PHYSICS, AMAZON variants, CLUSTER, PATTERN, REDDIT, PPI, FLICKR, YELP, ogbn-arxiv, and ogbn-products.
Loading 2106.05609v1…