Graph Auto-Encoder via Neighborhood Wasserstein Reconstruction

1 INTRODUCTION . Network/Graph representation learning ( a.k.a . embedding ) aims to preserve the high-dimensional complex graph information involving node features and link structures in a low-dimensional embedding space , which requires effective feature selection and dimension reduction ( Hamilton et al. , 2017b ) . Graph neural networks ( GNNs ) have done great jobs to this end , but most of them rely on node labels from specific downstream tasks to be trained in a semi-supervised fashion ( Kipf & Welling , 2017 ; Hamilton et al. , 2017a ; Wu et al. , 2019 ; Veličković et al. , 2018a ; Klicpera et al. , 2019 ; Chien et al. , 2021 ) . However , similar to other domains , unsupervised representation learning is preferred in many cases , not only because labeled data is not always available ( Hu et al. , 2020 ; Xie et al. , 2021 ) , but also task-agnostic representations can better transfer and generalize among different scenarios ( Erhan et al. , 2010 ; Bengio , 2012 ; Radford et al. , 2016 ) . To train GNNs in an unsupervised fashion , the classic auto-encoder framework ( Baldi , 2012 ; Goodfellow et al. , 2016 ) provides a natural solution and has been widely explored such as the prominent work ( V ) GAE ( Kipf & Welling , 2016 ) . Specifically , classic auto-encoders aim to decode from the low-dimensional representations information in the entire receptive field of the neural networks . For GNNs , the receptive field of a node representation is its entire neighborhood . However , existing graph auto-encoders appear away from such a motivation and are designed to merely decode the direct links between the node pairs by minimizing a link reconstruction loss . The fundamental difficulty to reconstruct the entire receptive fields of GNNs is due to the non-trivial design of a reconstruction loss on the irregular graph structures . Unfortunately , the over-simplification into link reconstruction makes the learned node representations drop much information and thus provides undesired performance in many downstream tasks . ∗Equal contribution . †Corresponding author . 1Code available at https : //github.com/mtang724/NWR-GAE . Take Figure 1 as an example , where different types of information are mixed in a graph ( e.g. , proximity and structure information as illustrated in Figure 5 in Appendix B ( Cui et al. , 2021 ) ) . The node representations learned by existing graph auto-encoders such as GAE ( Kipf & Welling , 2016 ) are driven too much to be similar on linked nodes due to their simple link reconstruction objective , and thus fail to distinguish node pairs like ( 2 , 4 ) and ( 3 , 5 ) in the cliques , though they clearly have different structural roles and node features . On the other hand , structure-oriented embedding models like GraphWave ( Donnat et al. , 2018 ) can not consider node features and spatial proximity , and thus fail to distinguish node pairs like ( 0 , 1 ) , ( 2 , 4 ) and ( 3 , 5 ) though they have different features , as well as ( 2 , 5 ) and ( 3 , 4 ) though they are further apart . An ideal unsupervised node representation learning model as we advocate in this work is expected to be task-agnostic and encode as much information as possible of all types in a low-dimensional embedding space . In this work , we aim to fundamentally address the above limitations of existing unsupervised node representation learning models by proposing a novel graph auto-encoder framework for unsupervised GNN training . The new framework is equipped with a powerful decoder that fully reconstructs the information from the entire receptive field of a node representation . Our key technical contribution lies in designing a principled and easy-to-compute loss to reconstruct the entire irregular structures of the node neighborhood . Specifically , we characterize the decoding procedure as iteratively sampling from a series of probability distributions defined over multi-hop neighbors ’ representations obtained through the GNN encoder . Then , the reconstruction loss can be decomposed into three parts , for sampling numbers ( node degrees ) , neighbor-representation distributions and node features . All of these terms are easy to compute but may represent the entire receptive field of a node instead of just the linkage information to its direct neighbors . For the most novel and important term , neighborrepresentation distribution reconstruction , we adopt an optimal-transport loss based on Wasserstein distance ( Frogner et al. , 2015 ) and thus name this new framework as Neighborhood Wasserstein Reconstruction Graph Auto-Encoder ( NWR-GAE ) . As also illustrated in Figure 1 , NWR-GAE can effectively distinguish all pairs of nodes dissimilar in different perspectives , and concisely reflect their similarities in the low-dimensional embedding space . We have conducted extensive experiments on four synthetic datasets and nine real-world datasets . Among the real-world datasets , three have proximity-oriented tasks , three have structure-oriented tasks , and three have proximity-structure-mixed tasks . We can observe significant improvements brought by NWR-GAE over the best method among the state-of-the-art baselines on all structureoriented tasks ( 8.74 % to 18.48 % ) and proximity-structure-mixed tasks ( -2.98 % to 8.62 % ) , and competitive performance on proximity-oriented tasks ( -3.21 % to -0.32 % ) . In-depth ablation and hyper-parameter studies further consolidate the claimed advantages of NWR-GAE . 2 PRELIMINARIES , MOTIVATIONS & OTHER RELATED WORKS . In this work , we focus on the auto-encoder framework for unsupervised task-agnostic graph representation learning . The original motivation of auto-encoders is to perform neural-network-based dimension reduction of the data that originally lies in a high-dimensional space ( Hinton & Salakhutdinov , 2006 ) . Specifically , an auto-encoder consists of two components , an encoder and a decoder . The encoder works to compress each data point into a low-dimensional vector representation , while the decoder works to reconstruct the original information from this vector . By minimizing the reconstruc- tion error , the encoder automatically converges to a good compressor that allows the low-dimensional representations to capture as much information as possible from the original data . Although the above high-level idea of auto-encoders is clear , when it is applied to graph structured data , the problem becomes challenging . This is because in graph-structured data , information of data points ( nodes to be specific as most widely studied ) is correlated due to the ambient graph structure . Without a specific task needed in a priori , the learned low-dimensional representation of a node should carry as much information as possible from not only its own features but also the features of the nodes it connects to ( both directly and indirectly ) . This implies that when building auto-encoders for graph-structure data , we expect the node representations to be able to reconstruct all correlated node features . However , existing graph auto-encoders seem to be away from this motivation . Previous prominent works such as unsupervised GraphSAGE ( Hamilton et al. , 2017a ) , GAE ( Kipf & Welling , 2016 ) , their generative variants such as VGAE ( Kipf & Welling , 2016 ) , CondGen ( Yang et al. , 2019 ) ( ) , and many others ( Grover et al. , 2019 ; Pan et al. , 2018 ; Shi et al. , 2020 ; Yang et al. , 2021 ) , use GNNs to encode graph structured data into node representations . Without exception , they follow the rationale of traditional network embedding techniques ( Perozzi et al. , 2014 ; Qiu et al. , 2018 ; Grover & Leskovec , 2016 ) and adopt link reconstruction in the decoder as the main drive to optimize their GNN encoders . The obtained node representations best record the network linkage information but lose much of other important information , such as local structures , neighbors ’ features , etc . Hence , these auto-encoders will most likely fail in other tasks such as node classifications ( especially structure-oriented ones as manifested in Figure 1 ) . To better understand this point , we carefully analyze the source of information encoded in each node representation via a GNN . Suppose a standard message-passing GNN ( Gilmer et al. , 2017 ) is adopted as the encoder , which is a general framework that includes GCN ( Kipf & Welling , 2017 ) , GraphSAGE ( Hamilton et al. , 2017a ) , GAT ( Veličković et al. , 2018a ) , GIN ( Xu et al. , 2019c ) and so on . After k-hop message passing , the source of information encoded in the representation of a node v essentially comes from the k-hop neighborhood of v ( Fig . 2 ) . Therefore , a good representation of node v should capture the information of features from all nodes in its k-hop neighborhood , which is agnostic to downstream tasks . Note that this may not be ideal as nodes out of k-hop neighborhood may also provide useful information , but this is what GNN-based graph auto-encoders can be expected to do due to the architectures of GNN encoders . This observation motivates our study on a novel graph decoder that can better facilitate the goal of GNN-based graph auto-encoders , based on the neighborhood reconstruction principle . We will formalize this principle in Sec . 3 . Relation to the InfoMax principle . Recently , DGI ( Veličković et al. , 2018b ) , EGI ( Zhu et al. , 2021 ) and others ( Sun et al. , 2020 ; Hu et al. , 2020 ; You et al. , 2020 ; Hassani & Khasahmadi , 2020 ; Suresh et al. , 2021 ) have used constrasive learning for unsupervised GNN training methods and may capture information beyond the directed links . They adopt the rule of mutual information maximization ( InfoMax ) , which essentially works to maximize certain correspondence between the learned representations and the original data . For example , DGI ( Veličković et al. , 2018b ) maximizes the correspondence between a node representation and which graph the node belongs to , but this has no guarantee to reconstruct the structural information of node neighborhoods . Recent works even demonstrate that maximizing such correspondence risks capturing only the noisy information that is irrelevant to the downsteam tasks because noisy information itself is sufficient for the models to achieve InfoMax ( Tschannen et al. , 2020 ; Suresh et al. , 2021 ) , which gets demonstrated again by our experiments . Our goal instead is to let node representations not just capture the information to distinguish nodes but capture as much information as possible to reconstruct the features and structure of the neighborhood . Optimal-transport ( OT ) losses . Many machine learning problems depend on the characterization of the distance between two probability measures . The family f -divergence has the non-continuous issue when the two measures of interest have non-overlapped support ( Ali & Silvey , 1966 ) . Therefore , OT-losses are often adopted and have shown great success in generative models ( Gulrajani et al. , 2017 ) and domain adaptation ( Courty et al. , 2016 ) . OT-losses have been used to build variational auto-encoders for non-graph data ( Tolstikhin et al. , 2018 ; Kolouri et al. , 2018 ; Patrini et al. , 2020 ) . But one should note the our work is not a graph-data-oriented generalization of these works : They use OT-losses to characterize the distance between the variational distribution and the data empirical distribution while our model even does not use a variational distribution . Our model may be further improved by being reformulated as a variational autoencoder but we leave it as a future direction . Here , we give a frequently-used OT loss based on 2-Wasserstein distance that will be used later . Definition 2.1 . Let P , Q denote two probability distributions with finite second moment defined on Z ⊆ Rm . The 2-Wasserstein distance between P and Q defined on Z , Z ′ ⊆ Rm is the solution to the optimal mass transportation problem with ` 2 transport cost ( Villani , 2008 ) : W2 ( P , Q ) = ( inf γ∈Γ ( P , Q ) ∫ Z×Z′ ‖Z − Z ′‖22dγ ( Z , Z ′ ) ) 1/2 ( 1 ) where Γ ( P , Q ) contains all joint distributions of ( Z , Z ′ ) with marginals P and Q respectively .

Summary: The paper proposes a novel approach to graph representation learning. In particular, a graph auto-encoder is proposed that aims to better capture the topological structure by utilising a neighbourhood reconstruction and a degree reconstruction objective. An optimal-transport based objective is proposed for the neighbourhood reconstruction that optimises the 2-Wasserstein distance between the decoded distribution and an empirical estimate of the neighbourhood distribution. An extensive experimental analysis is performed, highlighting the benefits of the proposed approach on a range of synthetic datasets to capture structure information. The experimental results also highlight its robustness across 9 different real-world graph datasets (ranging from proximity-oriented to structure-oriented datasets).