Source-linked AI summary

Large-Scale Representation Learning on Graphs via Bootstrapping

Shantanu Thakoor, Corentin Tallec, Mohammad Gheshlaghi Azar, Mehdi Azabou, Eva L. Dyer, Rémi Munos, Petar Veličković, Michal Valko

arXiv:2102.06514v3cs.LGcs.SIstat.ML

TL;DR

Graph self-supervised learning seeks label-free representations, but strong contrastive methods can require many negatives and become expensive on large graphs. BGRL predicts representations across simply augmented views without negative examples, matching or exceeding prior performance with lower memory and scaling to very large graphs. The paper reports state-of-the-art benchmark results, including 73.89% single-model performance and a second-place KDD Cup result using BGRL-based representations.

  • Problem

    Graph self-supervised learning aims to reduce reliance on costly labels, but contrastive methods often require many negative examples and can be prohibitively expensive for large graphs.

  • Method

    BGRL predicts target representations from online representations across two simply augmented graph views, updating the target encoder by exponential moving average without negative examples.

  • Results

    BGRL matches or outperforms prior methods on established benchmarks with 2-10x lower memory costs, while scaling to extremely large graphs and achieving state-of-the-art performance.

  • Takeaways & Limitations

    BGRL provides a scalable non-contrastive approach for graph representation learning and supports semi-supervised use of unlabeled data on graphs with hundreds of millions of nodes.

  • Takeaways & Limitations

    BGRL’s bootstrapping dynamics are not fully understood, which may make black-box use and harmful downstream effects difficult to diagnose and resolve.

Abstract

from arXiv · show

Self-supervised learning provides a promising path towards eliminating the need for costly label information in representation learning on graphs. However, to achieve state-of-the-art performance, methods often need large numbers of negative examples and rely on complex augmentations. This can be prohibitively expensive, especially for large graphs. To address these challenges, we introduce Bootstrapped Graph Latents (BGRL) - a graph representation learning method that learns by predicting alternative augmentations of the input. BGRL uses only simple augmentations and alleviates the need for contrasting with negative examples, and is thus scalable by design. BGRL outperforms or matches prior methods on several established benchmarks, while achieving a 2-10x reduction in memory costs. Furthermore, we show that BGRL can be scaled up to extremely large graphs with hundreds of millions of nodes in the semi-supervised regime - achieving state-of-the-art performance and improving over supervised baselines where representations are shaped only through label information. In particular, our solution centered on BGRL constituted one of the winning entries to the Open Graph Benchmark - Large Scale Challenge at KDD Cup 2021, on a graph orders of magnitudes larger than all previously available benchmarks, thus demonstrating the scalability and effectiveness of our approach.

1 INTRODUCTION

Graph representation learning needs label-free methods because supervised GNNs can overfit, while contrastive approaches require many negative comparisons. BGRL addresses these challenges with simple augmentations and prediction between online and target representations.

  • Supervised GNNs can overfit and fail to generalize, motivating simplified graph representations that do not require labels.
  • Contrastive graph methods pull together related representations and push apart unrelated ones, requiring comparisons with many negative examples.
  • BGRL learns representations by predicting an alternative augmented view with online and target encoders, while updating the target through an exponential moving average.
  • BGRL uses simple graph augmentations and avoids negative examples, enabling scalable self-supervised representation learning on very large graphs.

2 BOOTSTRAPPED GRAPH LATENTS

BGRL creates two augmented graph views, predicts target embeddings from online embeddings, and updates only the online network by gradient descent while tracking it with a moving-average target encoder. This non-contrastive design scales linearly and avoids negative-pair computation.

  • BGRL components: BGRL applies stochastic augmentations to create two graph views, then encodes them with distinct online and target graph encoders.
  • BGRL components: The predictor transforms the online representation into a prediction of the target representation for each node.
  • BGRL components: BGRL omits the projector network used by BYOL because graph embeddings are already small, with no observed performance loss.
  • BGRL update step: Only online parameters are optimized toward target representations, while target parameters follow an exponential moving average of the online parameters.
  • BGRL update step: The method does not optimize the full cosine objective with respect to target parameters, and empirical analysis reports no collapse to trivial solutions.
  • Scalable objective: BGRL avoids quadratic all-pairs contrasting and therefore scales linearly with graph size.

3 COMPUTATIONAL COMPLEXITY ANALYSIS

BGRL has linear computational scaling in graph size, unlike GRACE’s quadratic all-pairs contrastive computation. Its update cost combines encoder, prediction, and linear BGRL terms, while GRACE includes a quadratic contrastive term.

  • GRACE performs quadratic all-pairs contrastive computation at each update step, whereas BGRL avoids this contrastive operation.
  • For graphs with N nodes and M edges, BGRL’s encoder computations use terms proportional to M + N, with prediction and BGRL costs proportional to N.
  • BGRL’s per-step complexity is 6Cencoder(M + N) + 4CpredictionN + CBGRLN, compared with GRACE’s 4Cencoder(M + N) + 4CprojectionN + CGRACEN^2.
  • Table 1 compares BGRL and GRACE computational requirements across standard benchmark graphs and marks runs exceeding a 16GB V100 GPU as OOM.

4 EXPERIMENTAL ANALYSIS

The experiments evaluate BGRL across dataset scales, graph architectures, training regimes, and augmentation settings. BGRL matches or outperforms prior methods while reducing memory costs, remains competitive when contrastive objectives are approximated, and scales to MAG240M with strong semi-supervised results.

  • 4.1 PERFORMANCE AND EFFICIENCY GAINS WHEN SCALABILITY IS NOT A BOTTLENECK: BGRL achieves state-of-the-art performance on 4 of 5 medium-scale benchmark datasets while using 2-10x less memory.This memory-efficiency improvement incurs no reported performance cost relative to prior methods such as GRACE.
  • 4.1 PERFORMANCE AND EFFICIENCY GAINS WHEN SCALABILITY IS NOT A BOTTLENECK: BGRL safely uses simple augmentations on larger graphs without sacrificing performance, unlike more complex augmentation heuristics that can be expensive to scale.The comparison focuses on simple, easily scalable augmentations used by both BGRL and GRACE.
  • 4.2 PERFORMANCE WHEN SCALABILITY IS A BOTTLENECK: GRACE-SUBSAMPLING is highly sensitive to the number of negatives, while BGRL exceeds its performance at k = 2 and larger k values cause out-of-memory errors on a 16GB GPU.These results suggest that scaling contrastive objectives can reduce performance through necessary approximations.
  • 4.2 PERFORMANCE WHEN SCALABILITY IS A BOTTLENECK: On PPI, BGRL is competitive with GRACE using MeanPooling networks and reaches new state-of-the-art performance with a GAT model.GRACE does not improve the GAT model over the non-attentional MeanPooling encoder, while BGRL does.
  • 4.3 SCALING TO EXTREMELY LARGE GRAPHS: 73.89% is the highest reported single-model performance on MAG240M, compared with 70.02% for OGB baselines and 73.71% for the KDD Cup 2021 first-place solution before ensembling.The BGRL-based result uses no ensembling.
  • 4.3 SCALING TO EXTREMELY LARGE GRAPHS: The BGRL-based MAG240M solution achieved second place in the KDD Cup 2021 track with 75.19% test accuracy.The first- and third-place solutions achieved 75.49% and 74.60%, respectively, though direct comparison is confounded by other design factors.

5 RELATED WORK

Graph self-supervised learning progressed from random-walk objectives to contrastive GNN methods, while BGRL extends the area toward scalable semi-supervised learning on very large graphs.

  • Early graph representation methods used random-walk objectives, but combining random walks with GNNs could degrade performance.
  • Contrastive graph methods learn by bringing related representations together and separating unrelated ones, often through augmented graph views.
  • BGRL extends graph self-supervised learning to sampled subgraphs, semi-supervised training, and very high-data regimes such as MAG240M.
  • BGRL's contribution is evaluated against the computation and memory trade-offs of contrastive methods, particularly their reliance on negative examples.

ICLR ETHICS STATEMENT

The paper notes that label-free graph representations may help where labels are difficult or expensive to obtain, but warns that downstream misuse can create fairness and diagnostic risks.

  • BGRL may help applications where obtaining labels is challenging or expensive.
  • Using unsupervised embeddings without considering their training context may contribute to stereotyping or unfair biases in downstream tasks.
  • Because BGRL's bootstrapping dynamics are not fully understood, harmful downstream effects may be difficult to diagnose and resolve.

ICLR REPRODUCIBILITY STATEMENT

The paper reports reproducible experiments supported by released code, configurations, checkpoints, and independently verified large-scale implementation, alongside analyses of BGRL training behavior.

  • The authors report that their results are reproducible and provide implementation details and experimental materials.
  • Most experiments finish within 30 minutes on a single V100 GPU, and trained checkpoints allow performance verification without retraining.
  • The large-scale MAG240M solution was open-sourced through the KDD Cup 2021 and independently verified by contest organizers.
  • BGRL's loss does not converge to zero, while embedding spread remains comparable to embedding norms across datasets.
  • Embedding norms neither collapse to zero nor diverge as training progresses.

B ABLATIONS ON PROJECTOR NETWORK

The projector network used by some related methods is unnecessary for BGRL: the predictor can serve its role, and removing the projector does not reduce performance.

  • BGRL omits a projector network, unlike BYOL and GRACE, because its predictor can prevent augmentation invariance.
  • The projector's dimensionality-reduction role is less necessary for graph tasks because common graph embedding sizes are relatively small.

C COMPARISON ON SMALL DATASETS

On small graph datasets, BGRL remains effective in the low-data regime, while the study uses repeated random splits and examines simple and adaptive augmentation choices.

  • C COMPARISON ON SMALL DATASETS: BGRL performs well in the low-data regime and achieves new state-of-the-art performance on two small datasets.The evaluation covers Cora, CiteSeer, PubMed, and DBLP.
  • C COMPARISON ON SMALL DATASETS: Table 7 reports results averaged over 20 random dataset splits and model initializations.This evaluation follows the split-averaging procedure used in prior work rather than standard fixed splits.
  • C COMPARISON ON SMALL DATASETS: The augmentation pipeline primarily uses standard graph augmentations, while additional experiments consider more complex adaptive augmentations.The standard pipeline is motivated by the limited exploration of meaningful graph augmentations and their computational cost.
  • C COMPARISON ON SMALL DATASETS: The standard augmentations include node feature masking and edge masking, with edge masking leveraging graph topology.These graph-wise transformations do not operate independently on each node.

E DATASET DETAILS

The experiments cover citation, co-purchase, coauthorship, biological interaction, and Wikipedia-derived graphs with task-specific node meanings, features, labels, and data splits.

  • E DATASET DETAILS: WikiCS represents Wikipedia Computer Science articles as nodes and references as edges, with 10 subfield classes and 20 canonical splits.Node features are averages of GloVE embeddings of article words.
  • E DATASET DETAILS: Amazon Computers and Amazon Photos are co-purchase graphs whose products have 10 and 8 category classes, respectively.Product features are bag-of-words representations of reviews, and both datasets use 10/10/80% random train/validation/test splits.
  • E DATASET DETAILS: Coauthor CS and Coauthor Physics connect authors who co-authored papers, with 15 and 5 research-field classes, respectively.Their node features are bag-of-words representations of paper keywords, using 10/10/80% random splits.
  • E DATASET DETAILS: ogbn-arXiv is a symmetrized citation network of computer-science papers with 40 subject-area classes and skip-gram-based average word-embedding features.An edge connects papers when either paper cites the other.
  • E DATASET DETAILS: PPI contains 24 tissue-specific protein-interaction graphs, using 20 graphs for training, 2 for validation, and 2 for testing.Each node has 50 biological-property features, and the multilabel task permits up to 121 labels per node.

F IMPLEMENTATION DETAILS

Implementation uses standard GNN encoders, optimization schedules, normalization, and frozen linear evaluation protocols tailored to the dataset and encoder setting.

  • F IMPLEMENTATION DETAILS: BGRL uses an MLP predictor, AdamW with weight decay 10^-5, and a target-parameter decay rate τ scheduled from 0.99 toward 1.0.Models use Glorot initialization, and the predictor has one hidden layer.
  • F IMPLEMENTATION DETAILS: GCN propagation uses self-looped adjacency, degree normalization, a nonlinearity, and a learned weight matrix at each layer.The stated nonlinearity may be ReLU.
  • F IMPLEMENTATION DETAILS: The inductive encoder is a 3-layer mean-pooling network with skip connections, 512-dimensional layers, and PReLU activation.This architecture follows the cited prior proposal.
  • F IMPLEMENTATION DETAILS: GAT encoders transform node features and compute normalized pairwise attention coefficients using learned attention functions.The reported architecture uses three GAT layers, with four heads of size 256 in the first two layers.
  • F IMPLEMENTATION DETAILS: Training uses 10,000 total steps and 1,000 warmup steps, except GAT experiments on PPI, which use 20,000 and 2,000 steps.Learning-rate and target-decay schedules follow cosine annealing procedures.
  • F IMPLEMENTATION DETAILS: Frozen linear evaluation fits a regularized logistic-regression classifier on embeddings without backpropagating into the encoder.For larger datasets, AdamW replaces the liblinear solver; embeddings are ℓ2-normalized before classification.

G MAG240M EXPERIMENT DETAILS

MAG240M experiments test BGRL on a heterogeneous graph with 121 million papers, 122 million authors, and 26 thousand institutions, using message-passing and GCN encoders under multiple evaluation settings.

  • G MAG240M EXPERIMENT DETAILS: MAG240M contains 121 million papers, 122 million authors, and 26 thousand institutions, with 768-dimensional BERT paper features and 153 paper categories.The task classifies arXiv papers, which comprise 1% of paper nodes.
  • G MAG240M EXPERIMENT DETAILS: The main MAG240M experiments use bidirectional MPNNs with four message-passing steps, 256 hidden dimensions, and MLP node and edge updates.The update MLPs have two hidden layers of size 512.
  • G MAG240M EXPERIMENT DETAILS: Because full-graph training is infeasible, each batch samples 1,024 central nodes across eight devices with fixed-depth neighborhoods.The depth-2 sampling fanouts depend on whether neighbors are papers, authors, or institutions.
  • G MAG240M EXPERIMENT DETAILS: MAG240M settings use edge masking probability 0.2, feature masking probability 0.4, target decay starting at 0.999, and a learning rate annealed from 0.01 to zero.AdamW uses weight decay 10^-5, with warmup equal to 10% of the learning period.
  • G MAG240M EXPERIMENT DETAILS: With a 2-layer, 128-dimensional GCN, both BGRL and GRACE outperform fully supervised training, while BGRL learns faster and more stably.This tests whether BGRL remains effective with a weaker encoder.
  • G MAG240M EXPERIMENT DETAILS: Figure 11 compares BGRL and GRACE memory usage across five standard datasets and shows empirical scaling behavior consistent with theoretical predictions.GRACE is omitted for the largest dataset because it runs out of memory.
  • G MAG240M EXPERIMENT DETAILS: Under frozen evaluation on MAG240M, GRACE outperforms BGRL, but both methods perform poorly and underperform LabelProp.This motivates evaluating the methods in the semi-supervised setting.
Loading 2102.06514v3…