Self-Supervised Representation Learning via Latent Graph Prediction

1 INTRODUCTION . Self-supervised learning ( SSL ) methods seek to use supervisions provided by data itself and design effective pretext learning tasks . These methods allow deep models to learn from a massive amount of unlabeled data and have achieved promising successes in natural language processing ( Devlin et al. , 2019 ; Wu et al. , 2019 ; Wang et al. , 2019 ) and image tasks ( Batson and Royer , 2019 ; Xie et al. , 2020 ; He et al. , 2020 ; Chen et al. , 2020 ) . To use unlabeled graph data , earlier studies ( Perozzi et al. , 2014a ; Grover and Leskovec , 2016 ) adapt sequence-based SSL methods ( Mikolov et al. , 2013b ; a ) to learn node representations . Inspired by the recent success of SSL in the image domain , a variety of SSL methods based on graph neural networks ( GNNs ) have been proposed in different learning paradigms . In particular , ( Veličković et al. , 2019 ; Zhu et al. , 2020 ; Thakoor et al. , 2021 ; Hassani and Khasahmadi , 2020 ; You et al. , 2020 ) construct SSL tasks as unsupervised approaches to learn representations from graph data at either node-level or graph-level ; ( Hu et al. , 2020 ) proposes SSL strategies to pre-train GNNs for downstream tasks ; and ( Jin et al. , 2020 ; Kim and Oh , 2021 ) employ SSL as auxiliary tasks to boost the performance of main learning tasks . Common taxonomies in recent survey works ( Xie et al. , 2021 ; Liu et al. , 2020 ) consider two categories of SSL methods to train GNNs ; namely , contrastive methods and predictive methods . Contrastive methods employ pair-wise discrimination as their pretext learning tasks . It performs transformations or augmentations to obtain multiple views from a graph and trains GNNs to discriminate between jointly sampled view pairs and independently sampled view pairs . In contrast , predictive methods ( Hamilton et al. , 2017 ; Hwang et al. , 2020 ; Rong et al. , 2020 ) train GNNs to predict certain labels obtained from the input graph , such as node reconstruction , connectivity reconstruction , graph statistical properties , and domain knowledge-based targets . Adapted from the image domain , current state-of-the-art SSL methods for graphs are mostly contrastive . As a drawback , they usually depend on large training sample size to include a sufficient number of negative samples . With limited computing resources , contrastive methods may not be applicable to large-scale graphs without suffering from performance loss . To address the drawback , BGRL ( Thakoor et al. , 2021 ) adapts BYOL ( Grill et al. , 2020 ) to the graph domain . BGRL still obtains different views from each given graph , but it eliminates the requirement of negative samples by replacing contrastive objectives with the prediction of offline embedding . BGRL has achieved competitive performance to the contrastive methods . However , unlike contrastive methods grounded by mutual information estimation and maximization , BYOL and BGRL lack theoretical guidance and require implementation measures to prevent collapsing to trivial representations , such as stop gradient , EMA , and normalization layers . In this work , we propose LaGraph , a predictive SSL framework for representation learning of graph data , based on self-supervised latent graph prediction . In particular , we describe the notion of the latent graph and introduce the latent graph prediction as a pretext learning task . We adapt the supervised objective of latent graph prediction into a self-supervised setting by deriving its self-supervised upper bounds , according to which we present the learning framework of LaGraph . We provide further justifications of LaGraph by comparing it with theoretically sound methods in different domains . Our experimental results demonstrate the effectiveness of LaGraph on both graph-level and node-level representation learning , where a remarkable performance boost is achieved on a majority of datasets with higher stability to smaller batch size or training on subsets of nodes . Relations with Prior Work : Both LaGraph and some existing contrastive methods ( You et al. , 2020 ; Zhu et al. , 2020 ; Hu et al. , 2020 ) apply node masking . A major difference is that LaGraph is a predictive method instead of contrastive . While those contrastive methods use node masking as an augmentation to obtain different views for contrast , LaGraph employ it for the computation of the invariance regularization . In addition , the objective of BGRL can be considered similar to the invariance regularization term in our objective . While the objective of BGRL is designed without theoretical guidance , our derived theorems associated with LaGraph objectives can explain the success of BGRL to some extent and provide guidance on how to better adopt objectives related to the invariance regularization on graphs . 2 METHODS . 2.1 NOTATIONS AND PROBLEM FORMULATION . We consider an undirected graph G = ( V , E ) with a set of attributed nodes V and a set of edges E. We formulate the graph data as a tuple of matrices ( A , X ) , where A ∈ R|V |×|V | denotes the adjacency matrix and X ∈ R|V |×d denotes the node features of dimension d. We employ a graph encoder E based on graph neural networks ( GNNs ) to encode each node or graph into a corresponding representation . Namely , we compute the node-level representations or node embedding by H = E ( A , X ) ∈ R|V |×q and the graph-level representation or graph embedding by z = R ( H ) ∈ R1×q , where q denotes the embedding dimension andR : R|V |×q → R1×q is a readout function . Self-supervised representation learning is employed to train the graph encoder E on a set of K graphs { Gi } Ki=1 without labels from downstream tasks . In particular , we seek to design effective pre-text learning tasks , whose labels are obtained by task designation or from given data , to train the graph encoder E and produces informative representations for downstream tasks . Depending on the pre-text learning tasks , the encoder E is usually trained together with some prediction head D for predictive SSL or a discriminator for contrastive SSL . 2.2 LATENT GRAPH PREDICTION . Our method considers latent graph prediction as a pretext task to train graph neural networks . In this subsection , we introduce the general notion of latent data , followed by its specific definition for graph data , and the construction of the learning task . For any observed data instance x , we assume that there exists a corresponding latent data xI , determining the semantic of x , such that the latent data xI is generated from a prior p ( xI ) and the observed data instance is further generated from a certain distribution conditioned on the latent data , i.e. , p ( x|xI ) . The most common case for the pair of observed data and latent data is the noisy data and its clean version . When it comes to graph data , we consider the case that an observed graph data G = ( A , X ) is ( noisily ) generated from its latent graph G ` = ( A , F ) with the same node set and edge set , where node feature matrices X and F for the two graphs have the same dimensionality . We make two assumptions to the graphs without loss of generality . First , we assume that the observed feature vector xv of each node v in an observed graph is independently generated from a certain distribution conditioned on the corresponding latent graph . In other words , how xv is generated from the latent feature fv is not affected by the generation of other observed feature vectors . Second , we assume that the conditional distribution of the observed graph is centered at the latent graph , i.e. , E [ X|G ` ] = F . The above assumptions are natural when we have little knowledge about the generation process and are commonly used in other types of data such as the non-structural and zero-mean noise in images . In cases that the generation processes of different nodes are related or the distribution is not centered at F , we can still consider the related or biased components into the latent feature and therefore have the assumptions satisfied . As the latent data usually determine the semantic meaning of observed data , we believe the prediction of latent graph can provide informative supervision for the learning of both graph-level and nodelevel representations . We are hence interested in constructing the learning task of latent graph prediction . To perform latent graph prediction , it is straightforward to employ a graph neural network f : { 0 , 1 } |V |×|V | × R|V |×d → R|V |×d that takes an observed graph G = ( A , X ) as inputs and predicts the feature matrix of its latent graph GI = ( A , F ) . When the ground truth of the latent feature matrix F is known , the learning objective can be designed as f∗ = argmin f E ‖f ( A , X ) − F ‖2 . ( 1 ) Intuitively , the latent graph prediction can be considered as a generalized task from noisy data reconstruction that predicts the signal from the noisy data with the objective argminf E ‖f ( x ) − s‖2 , where the mapping from the signal to the noisy data p ( x|s ) can usually be explicitly modeled and samples of signal ( ground truth ) can usually be captured . In the data reconstruction case , pairs of ( x , s ) can be therefore directly captured or synthetically generated given a certain noise model p ( x|s ) . However , when the task is generalized to latent graph prediction , there is a key challenge preventing us from directly applying the prediction task . That is , whereas there are natural supervisions for noisy data reconstruction , the latent graph is not observed and we are unable to explicitly model the mapping from latent graphs to observed graphs , i.e. , the conditional distribution p ( G|GI ) . 2.3 SELF-SUPERVISED UPPER BOUNDS FOR LATENT GRAPH PREDICTION . As discussed in the previous subsection , unlike typical noisy data reconstruction tasks , the latent graph is not observed and p ( G|GI ) can not be modeled explicitly . This makes it difficult to construct a direct learning task for latent graph prediction using the objective in Equation ( 1 ) . We therefore seek to optimize an alternative objective that approximately optimizes the objective in Equation ( 1 ) without requiring the distribution p ( G|GI ) , nor features F of the latent graph . We now introduce the proposed self-supervised objective for latent graph prediction . We derive our self-supervised objective without involving F by constructing an upper bound of the objective in Equation ( 1 ) . Specifically , we let J ⊂ { 0 , · · · , |V | − 1 } be an arbitrary subset of node indices , Jc denote the complement of set J , and XJc : = 1Jc X + 1J M be the feature matrix with features of nodes in VJ masked , where denotes element-wise multiplication , M ∈ R|V |×d denotes a matrix consisting of independent random noise or zeros as masking values , and 1J ∈ R|V |×d denotes an indicator matrix such that 1J [ i , : ] = 1 , ∀i ∈ J and 1J [ i , : ] = 0 , ∀i /∈ J . We describe the self-supervised upper bound in Theorem 1 , whose proof is provided in Appendix A. Theorem 1 . Consider a graph G = ( A , X ) and its latent graph GI = ( A , F ) . We let the variance of any elements in X be bounded by σ2 and J be a subset of nodes V in the graph G. For any graph neural network f : { 0 , 1 } |V |×|V | ×R|V |×d → R|V |×d , we have the following inequality EA , X , F [ ‖f ( A , X ) − F ‖2 + ‖X − F ‖2 ] ≤ EA , X ‖f ( A , X ) −X‖2 + 2σ|V |EJ [ EA , X ‖fJ ( A , X ) − fJ ( A , XJc ) ‖2 |J | ] 1/2 . ( 2 ) Intuitively , the first component in the upper bound derived in Theorem 1 measures the reconstruction error on the feature matrix X of the given observed graph G , enforcing the intermediate representations to be informative . The second component controls how much information is accessible from the input feature of a node vi when reconstructing the feature of vi , by encouraging the output of a node to be invariant to the missing of its features in the input graph . We then call the first component as a reconstruction term and the second component as an invariance regularization term . Note that the invariance regularization is only computed on masked nodes in contrast to the BGRL objective , based on different theoretical grounding and leading to a different effect ( see Appendix D ) . In tasks of self-supervised representation learning , we are more interested in graph-level or node-level representations than predicted latent graphs . In these cases , we expect the representations also holds the invariance property held by the final outputs . We therefore seek to apply the invariance regularization to the representations , since a regularization applied to the output does not necessarily control the information accessibility of representations produced intermediately in the graph neural network . To do so , we separately consider the encoder E and decoder D in the graph neural network f . We introduce certain assumptions to the decoder networkD and the readout functionR , and derive two additional upper bounds for node-level and graph-level representation learning , respectively in following corollaries . Proofs of the corollaries are provided in Appendix B. Corollary 1 . Let G = ( A , X ) be a given graph , GI = ( A , F ) be its latent graph , E and D be a graph encoder and a prediction head ( decoder ) consisting of fully-connected layers . If the prediction head D is ` -Lipschitz continuous with respect to l2-norm , we further have the following inequality , E [ ‖D ( H ) − F ‖2 + ‖X − F ‖2 ] ≤ E ‖D ( H ) −X‖2 + 2σ|V | ` EJ [ E ‖HJ −H ′J‖ 2 |J | ] 1/2 , ( 3 ) where H = E ( A , X ) and H ′ = E ( A , XJc ) denote the node embedding of the given graph and the masked graph , respectively , and HJ : = H [ J , : ] selects rows with indices in J . Corollary 2 . Let G = ( A , X ) be a given graph , GI = ( A , F ) be its hidden latent graph , E be a graph encoder , R be a readout function satisfying k-Bilipschitz continuity with respect to l2-norm , and D be a prediction head ( decoder ) . If the prediction head D is ` -Lipschitz continuous with respect to l2-norm , we have the following inequality , E [ ‖D ( H ) − F ‖2 + ‖X − F ‖2 ] ≤ E ‖D ( H ) −X‖2 + 2σ|V |k ` EJ [ E ‖z − z′‖2 |J | ] 1/2 , ( 4 ) where z = R ( H ) and z′ = R ( H ′ ) denote the graph-level representations of the given graph and the masked graph , respectively . We note that the assumptions and restrictions are natural or practically satisfiable . The assumption that variance of each element in X is bounded by σ holds when node features are from { 0 , 1 } d or when feature normalization is applied . The ` -Lipschitz continuous property is common for neural networks . And the k-Bilipschitz continuity can be satisfied by applying an injective readout function such as global sum pooling , which is commonly used in graph-level tasks . 2.4 THE LaGraph FRAMEWORK We design our self-supervised learning framework according to upper bounds derived in Corollary 1 and Corollary 2 . To train encoder E together with decoder D under self-supervision , we input to the encoder both the given graph ( A , X ) and its variation ( A , XJc ) with a random subset J of node indices for nodes to be masked and obtain node-level representations H = E ( A , X ) and H ′ = E ( A , XJc ) for the two graphs respectively . The self-supervised losses are computed on input node features , reconstructed node features , and representations , as demonstrated in Figure 1 . In particular , we consider a mini-batch of N graphs { ( Ai , Xi ) } Ni=1 and their corresponding masked variation { ( Ai , X ( i , Jic ) ) } Ni=1 where Ji denotes the node indices subset for the i-th graph . The self-supervised loss for node-level representation learning follows Corollary 1 and is computed as Lnode ( E , D ) = 1 N N∑ i=1 ‖D ( Ai , Hi ) −Xi‖2 /|Vi|+ α [ ∑ i ‖1Ji Hi − 1Ji H ′i‖ 2∑ i |Ji| ] 1/2 , ( 5 ) where α is a hyper-parameter corresponding to the multiplier 2σ ` in Corollary 1 . To fulfill the conditions in Corollary 1 , we employ fully-connected layers instead of graph convolutional layers in the decoder D. The pseudo-code are provided in Appendix C. Similarly , using the same notations above , the self-supervised loss for graph-level representation learning follows Corollary 2 and is computed as Lgraph ( E , D ) = 1 N N∑ i=1 ‖D ( Ai , Hi ) −Xi‖2 /|Vi|+ α′ [ ∑ i ‖zi − z′i‖ 2 / ∑ i |Ji| ] 1/2 , ( 6 ) where zi = R ( Hi ) and z′i = R ( H ′i ) denote the graph-level representations obtained by applying readout function R to the node-level representations , respectively , and α′ is a hyper-parameter corresponding to the multiplier 2σk ` in Corollary 2 . To fulfill the conditions in Corollary 2 , we employ global sum pooling as the readout function R , where as the decoder D here can consists of either fully-connected layers or graph convolutional layers . The pseudo-code are provided in Appendix C .

Summary: The authors propose a self-supervised learning (SSL) framework for graph neural networks by predicting an unobserved latent graph. The proposed self-supervised learning loss is derived from upper bounds to the original objective functions for predicting latent graphs. The proposed method is robust to small sample size and small batch size.