Large-Scale Representation Learning on Graphs via Bootstrapping

1 INTRODUCTION . Graphs provide a powerful abstraction for complex datasets that arise in a variety of applications such as social networks , transportation networks , and biological sciences ( Hamilton et al. , 2017 ; DerrowPinion et al. , 2021 ; Zitnik & Leskovec , 2017 ; Chanussot et al. , 2021 ) . Despite recent advances in graph neural networks ( GNNs ) , when trained with supervised data alone , these networks can easily overfit and may fail to generalize ( Rong et al. , 2019 ) . Thus , finding ways to form simplified representations of graph-structured data without labels is an important yet unsolved challenge . Current state-of-the-art methods for unsupervised representation learning on graphs ( Veličković et al. , 2019 ; Peng et al. , 2020 ; Hassani & Khasahmadi , 2020 ; Zhu et al. , 2020b ; a ; You et al. , 2020 ) are contrastive . These methods work by pulling together representations of related objects and pushing apart representations of unrelated ones . For example , current best methods Zhu et al . ( 2020b ) and Zhu et al . ( 2020a ) learn node representations by creating two augmented versions of a graph , pulling together the representation of the same node in the two graphs , while pushing apart every other node pair . As such , they inherently rely on the ability to compare each object to a large number of negatives . In the absence of a principled way of choosing these negatives , this can require computation and memory quadratic in the number of nodes . In many cases , the generation of a large number of negatives poses a prohibitive cost , especially for large graphs . In this paper , we introduce a scalable approach for self-supervised representation learning on graphs called Bootstrapped Graph Latents ( BGRL ) . Inspired by recent advances in self-supervised learning in vision ( Grill et al. , 2020 ) , BGRL learns node representations by encoding two augmented versions of a graph using two distinct graph encoders : an online encoder , and a target encoder . The online encoder is trained through predicting the representation of the target encoder , while the target encoder is updated as an exponential moving average of the online network . Critically , BGRL does not require contrasting negative examples , and thus can scale easily to very large graphs . Our main contributions are : • We introduce Bootstrapped Graph Latents ( BGRL ) , a graph self-supervised learning method that effectively scales to extremely large graphs and outperforms existing methods , while using only simple graph augmentations and not requiring negative examples ( Section 2 ) . • We show that contrastive methods face a trade-off between peak performance and memory constraints , due to their reliance on negative examples ( Section 4.2 ) . Due to having time and space complexity scaling only linearly in the size of the input , BGRL avoids the performancememory trade-off inherent to contrastive methods altogether . BGRL provides performance competitive with the best contrastive methods , while using 2-10x less memory on standard benchmarks ( Section 3 ) . • We show that leveraging the scalability of BGRL allows making full use of the vast amounts of unlabeled data present in large graphs via semi-supervised learning . In particular , we find that efficient use of unlabeled data for representation learning prevents representations from overfitting to the classification task , and achieves significantly higher , state-of-the-art performance . This was critical to the success of our solution at KDD Cup 2021 in which our BGRL-based solution was awarded one of the winners , on the largest publicly available graph dataset , of size 360GB consisting of 240 million nodes and 1.7 billion edges ( Section 4.3 ) . 2 BOOTSTRAPPED GRAPH LATENTS 2.1 BGRL COMPONENTS BGRL builds representations through the use of two graph encoders , an online encoder Eθ and a target encoder Eφ , where θ and φ denote two distinct sets of parameters . We consider a graph G = ( X , A ) , with node features X ∈ RN×F and adjacency matrix A ∈ RN×N . BGRL first produces two alternate views of G : G1 = ( X̃1 , Ã1 ) and G2 = ( X̃2 , Ã2 ) , by applying stochastic graph augmentation functions T1 and T2 respectively . The online encoder produces an online representation from the first augmented graph , H̃1 : = Eθ ( X̃1 , Ã1 ) ; similarly the target encoder produces a target representation of the second augmented graph , H̃2 : = Eφ ( X̃2 , Ã2 ) . The online representation is fed into a node-level predictor pθ that outputs a prediction of the target representation , Z̃1 : = pθ ( H̃1 ) . BGRL differs from prior bootstrapping approaches such as BYOL ( Grill et al. , 2020 ) in that it does not use a projector network . Unlike vision tasks , in which a projection step is used by BYOL for dimensionality reduction , common embedding sizes are quite small for graph tasks and so this is not a concern in our case . In fact , we observe that this step can be eliminated altogether without loss in performance ( Appendix B ) . The augmentation functions T1 and T2 used are simple , standard graph perturbations previously explored ( You et al. , 2020 ; Zhu et al. , 2020b ) . We use a combination of random node feature masking and edge masking with fixed masking probabilities pf and pe respectively . More details and background on graph augmentations is provided in Appendix D. 2.2 BGRL UPDATE STEP Updating the online encoder Eθ : The online parameters θ ( and not φ ) , are updated to make the predicted target representations Z̃1 closer to the true target representations H̃2 for each node , by following the gradient of the cosine similarity w.r.t . θ , i.e. , ` ( θ , φ ) = − 2 N N−1∑ i=0 Z̃ ( 1 , i ) H̃ > ( 2 , i ) ‖Z̃ ( 1 , i ) ‖‖H̃ ( 2 , i ) ‖ ( 1 ) θ ← optimize ( θ , η , ∂θ ` ( θ , φ ) ) , ( 2 ) where η is the learning rate and the final updates are computed from the gradients of the objective with respect to θ only , using an optimization method such as SGD or Adam ( Kingma & Ba , 2015 ) . In practice , we symmetrize this loss , by also predicting the target representation of the first view using the online representation of the second . Updating the target encoder Eφ : The target parameters φ are updated as an exponential moving average of the online parameters θ , using a decay rate τ , i.e. , φ← τφ+ ( 1− τ ) θ , ( 3 ) Figure 1 visually summarizes BGRL ’ s architecture . Note that although the objective ` ( θ , φ ) has undesirable or trivial solutions , BGRL does not actually optimize this loss . Only the online parameters θ are updated to reduce this loss , while the target parameters φ follow a different objective . This non-collapsing behavior even without relying on negatives has been studied further ( Tian et al. , 2021 ) . We provide an empirical analysis of this behavior in Appendix A , showing that in practice BGRL does not collapse to trivial solutions and ` ( θ , φ ) does not converge to 0 . Scalable non-contrastive objective : Here we note that a contrastive approach would instead encourage Z̃ ( 1 , i ) and H̃ ( 2 , j ) to be far apart for node pairs ( i , j ) that are dissimilar . In the absence of a principled way of choosing such dissimilar pairs , the naïve approach of simply contrasting all pairs { ( i , j ) | i 6= j } , scales quadratically in the size of the input . As BGRL does not rely on this contrastive step , BGRL scales linearly in the size of the graph , and thus is scalable by design . 3 COMPUTATIONAL COMPLEXITY ANALYSIS . We provide a brief description of the time and space complexities of the BGRL update step , and illustrate its advantages compared to previous strong contrastive methods such as GRACE ( Zhu et al. , 2020b ) , which perform a quadratic all-pairs contrastive computation at each update step . The same analysis applies to variations of the GRACE method such as GCA ( Zhu et al. , 2020a ) . Consider a graph with N nodes and M edges , and simple encoders E that compute embeddings in time and space O ( N +M ) . This property is satisfied by most popular GNN architectures such as convolutional ( Kipf & Welling , 2017 ) , attentional ( Veličković et al. , 2018 ) , or message-passing ( Gilmer et al. , 2017 ) networks . BGRL performs four encoder computations per update step ( twice for the target and online encoders , and twice for each augmentation ) plus a node-level prediction step ; GRACE performs two encoder computations ( once for each augmentation ) , plus a node-level projection step . Both methods backpropagate the learning signal twice ( once for each augmentation ) , and we assume the backward pass to be approximately as costly as a forward pass . We ignore the cost for computation of the augmentations in this analysis . Thus the total time and space complexity per update step for BGRL is 6Cencoder ( M +N ) + 4CpredictionN +CBGRLN , compared to 4Cencoder ( M +N ) + 4CprojectionN +CGRACEN2 for GRACE , where C· are constants depending on architecture of the different components . Table 1 shows an empirical comparison of BGRL and GRACE ’ s computational requirements on a set of benchmark tasks . 4 EXPERIMENTAL ANALYSIS . We present an extensive empirical study of performance and scalability , showing that BGRL is effective across a wide range of settings from frozen linear evaluation to semi-supervised learning , and both when performing full-graph training and training on subsampled node neighborhoods . We give results across a range of dataset scales and encoder architectures including convolutional , attentional , and message-passing neural networks . We analyze the performance of BGRL on a set of 7 standard transductive and inductive benchmark tasks , as well as in the very high-data regime by evaluating on the MAG240M dataset ( Hu et al. , 2021 ) . We present results on medium-sized datasets where contrastive objectives can be computed on the entire graph ( Section 4.1 ) , on larger datasets where this objective must be approximated ( Section 4.2 ) , and finally on the much larger MAG240M dataset designed to test scalability limits ( Section 4.3 ) , showing that BGRL improves performance across all scales of datasets . In Appendix C , we show that BGRL achieves state-of-the-art performance even in the low-data regime on a set of 4 small-scale datasets . Dataset sizes are summarized in Table 2 and described further in Appendix E. Evaluation protocol : In most tasks , we follow the standard linear-evaluation protocol on graphs ( Veličković et al. , 2019 ) . This involves first training each graph encoder in a fully unsupervised manner and computing embeddings for each node ; a simple linear model is then trained on top of these frozen embeddings through a logistic regression loss with ` 2 regularization , without flowing any gradients back to the graph encoder network . In the more challenging MAG240M task , we extend BGRL to the semi-supervised setting by combining our self-supervised representation learning loss together with a supervised loss . We show that BGRL ’ s bootstrapping objective obtains state-of-the-art performance in this hybrid setting , and even improves further with the added use of unlabeled data for representation learning - properties which have not been previously demonstrated by prior works on self-supervised representation learning on graphs . Implementation details including model architectures and hyperparameters are provided in Appendix F. Algorithm implementation and experiment code for most tasks have been provided in supplementary material , while code for our solution on MAG240M has been open-sourced1 as part of the KDD Cup 2021 .

Summary: This paper proposes BRGL, a method for graph representation learning based on ‘bootstrapping’ (in the same sense BYOL is). Different from the prior art, which mainly relies on contrastive learning, BRGL naturally scales linearly with graph size.  An extensive suite of experiments shows that BRGL scales better than previous works while yielding state-of-the-art (SOTA) performance in node classification.