RoSA: A Robust Self-Aligned Framework for Node-Node Graph Contrastive Learning

1 Introduction

Graph representation learning, which aims to learn low dimension representations of nodes and edges for downstream tasks, has become a popular method when dealing with graph-domain data recently. Among all these methods, unsupervised graph contrastive learning has received considerable research attention. It combines the new research trend of graph neural network (GNN) [Kipf and Welling, 2017] and contrastive self-supervised learning [Van den Oord et al., 2018; Chen et al., 2020; Grill et al., 2020] methods, and has achieved promising results on many graph-based tasks [Zhu et al., 2020b; Velickovic et al., 2019; You et al., 2020].

Contrastive learning aims to maximize the agreement between jointly sampled positive views and draw apart the distance between negative views, where in graph domain we refer augmented subgraph as a "view". Based on the scale of two contrasted views, graph contrasting learning can be classified as node-node, node-graph, and graph-graph level [Wu et al., 2021]. From another perspective, a pair of contrasted views is recognized as aligned or unaligned depending on the difference of their node sets. Two aligned views must have identical node indices, except the structure and some features may differ, and two unaligned views can have different node sets. Figure 1 gives an illustrative overview according to this taxonomy.

* Equal Contribution

† Corresponding Author

[Zhu et al., 2021a] indicates that for node-level tasks such as node classification, applying node-node contrasting can obtain the best performance gain. However, previous work for node-node graph contrastive learning all contrast nodes in the aligned scenario which may hinder the flexibility and variability of sampled views and restrict the expressive power of contrastive learning. Moreover, there exist certain circumstances where aligned views are unavailable, for instance the dynamic graphs where the nodes may appear/disappear as time goes by, and the random walk sampling where the views are naturally non-aligned. Compared with aligned node-node contrasting, the non-aligned scenario is able to sample different nodes and their relations more freely, which will assist the model in learning more representative and robust features.

However, applying non-aligned node-node contrasting faces three challenges. First, how to design sub-sampling methods that can generate unaligned views while maintaining semantic consistency? Second, how to contrast two nonaligned views even the number of nodes and correspondence between nodes are inconsistent? Third, how to boost the performance meanwhile enhance the robustness of model for unsupervised graph contrastive learning? None of them have been satisfactorily answered by previous work.

To tackle the challenges discussed above, we propose RoSA: a Robust Self-Aligned framework for node-node graph contrastive learning. Firstly, we utilize random walk sampling to obtain augmented views for non-aligned nodenode contrastive learning. Specifically, for a given graph, we sample a series of subgraphs based on a central node, and two different views of the same central node are treated as a positive pair, while views across different central nodes are selected as negative pairs. Note that even positive pairs are not necessarily aligned. Secondly, inspired by the message passing mechanism of graph neural networks, the node representation can be interpreted as the result of distribution transformation of its neighboring nodes. Intuitively, for a pair of views, we leverage the earth mover's distance (EMD) to model the minimum effort to transform the distribution of one view to the other as our objective, which can implicitly align different views and capture the changes in their distributions. Thirdly, we introduce unsupervised adversarial training that explicitly operates on node features to increase the diversity of samples and enhance the robustness of our model. To the best of our knowledge, this is the first work that fills the blank in non-aligned node-node graph contrastive learning.

Our main contributions are summarized as follows:

• We propose a robust self-aligned contrastive learning framework for node-node graph representation learning named RoSA. To the best of our knowledge, this is the first work dedicated to solving non-aligned node-node graph contrastive learning problems.
• To tackle the non-aligned problem, we introduce a novel graph-based optimal transport algorithm, g-EMD, which does not require explicit node-node correspondence and can fully utilize graph topological and attributive information for non-aligned node-node contrasting. Moreover, to compensate for the possible information loss caused by non-aligned sub-sampling, we propose a nontrivial unsupervised graph adversarial training to improve the diversity of sub-sampling and strengthen the robustness of the model.
• Extensive experimental results on various graph settings achieve promising results and outperform several baseline methods by a large margin, which validates the effectiveness and generality of our method.

2 RelatedWorks

2.1 Self-Supervised Graph Representation Learning

First appeared in the field of computer vision [Van den Oord et al., 2018; He et al., 2020; Grill et al., 2020] and natural language processing [Gao et al., 2021], self-supervised learning showed promising performance in various tasks and applying it to graph domain quickly became a research hot-spot. GraphCL [You et al., 2020] uses different augmentations and applies a readout function to obtain graph-graph level representations, then optimizes the InfoNCE loss, which can be mathematically proved to be the lower bound of mutual information. Inspired by Deep InfoMax [Hjelm et al., 2019], DGI [Velickovic et al., 2019] maximizes the mutual information between patch and graph representations, which is nodegraph level contrasting. Recently, node-node level methods like GMI [Peng et al., 2020], GRACE [Zhu et al., 2020b], GCA [Zhu et al., 2021b] and BGRL [Thakoor et al., 2021] show superior performance on node classification task. Unlike DGI, GMI removes the readout function and maximizes the MI between inputs and outputs of the encoder at the nodenode level. With graph augmentation methods, GRACE focuses on contrasting aligned views using different nodes as the negative pairs, and the same nodes from different views are regarded as positive pairs, where each positive pair should be aligned first. GCA is similar to GRACE but with adaptive data augmentation. BGRL is a negative-sample-free method which borrows the idea from BGRL [Grill et al., 2020].

Previous works that involve graph level contrasting, usually have a readout function to obtain whole graph representation, which are naturally aligned, but when it comes to nodenode level contrasting, they always explicitly align nodes for positive pairs. The work of non-aligned node-node graph contrastive learning has not yet been explored.

2.2 Adversarial Training

Adversarial training (AT) has been found useful to improve the model's robustness. AT is a min-max training process, which aims to maintain the consistency of the model's output before and after adding adversarial perturbations. Previous works solve the adversarial perturbations from many different perspectives. [Goodfellow et al., 2015] gives a linear approximation of the perturbation under L2 norm (i.e.Fast Gradient method). Projected Gradient Descent method [Madry et al., 2018] tries to obtain a more precise perturbation in an iterative manner, but it takes more time, [Zhu et al., 2020a] provide more efficient methods. Lately, [Kong et al., 2020] adopts these methods into the graph domain in a supervised manner. However, unsupervised adversarial training for graphs is still unexplored. In this paper, we adopt AT into our contrastive method to improve the robustness of the model in an unsupervised manner.

3 Method

In this chapter, we will introduce the framework of RoSA. Figure 2 gives an overview of RoSA.

)௜(ଵ܆෩ for ߜ perturbation

ࡴଶ(௝)

෩Gଶ(௝)

ࢆଶ(௝)

Figure 2: The overview of our proposed method: RoSA. The input is a series of subgraphs sampled from a full graph, where different random walk views from the same central node are recognized as positive pairs and views from different central nodes are treated as negative pairs. Then the subgraphs are fed into the encoder and projector to obtain node embeddings for contrasting. The self-aligned EMD-based contrastive loss will maximize the mutual information (MI) between positive pairs and minimize MI between negative pairs, guiding the model to learn rich representations. Besides, introducing adversarial training into this workflow enhances the robustness of the model.

3.1 Preliminaries

Given a graph G = (V, E), where V is the set of N nodes and E is the set of M edges. Also use G = (X,A) to represent graph features, where X = {x1, x2, . . . , xN} ∈ RN×d represents node feature matrix, each node's feature dimension is d and can be formulated as xi ∈ Rd , A ∈ RN×N represents the graph adjacency matrix, where Ai,j = 1 if an edge exists between node i and j, else Ai,j = 0. For subgraph sampling, each node i will be treated as central node to get subgraph G (i) . An augmented view of subgraph G (i) is represented as ˜ G (i) k where subscript k denotes the k-th augmented view.

3.2 Non-Aligned Node-Node Level Sub-Sampling

It has been proven that well-designed data augmentation plays a vital role in boosting the performance of contrastive learning [You et al., 2020]. However, different from the CV and NLP domain, where data is organized in a Euclidean fashion, graph data augmentation methods need to be redesigned and carefully selected.

It is worth noting that in this work, for a positive pair, we need to get different sets of nodes while preserving the consistency of their semantic meanings. Based on such a premise, we propose to utilize random walk with restart sampling as an augmentation method that selects nodes randomly and generates unaligned views. Specifically, random walk sampling starts from the central node v and generates a random path with a given step size s. Besides, at each step the walk returns to central node v with a restart probability α. The step size s should not be too large because we want to capture the local structure of the central node. Lastly, edge dropping and feature masking [Zhu et al., 2020b] are applied on subgraphs.

3.3 g-EMD: A Self-aligned Contrastive Objective

After obtaining two unaligned augmented views, we define a contrastive objective that measures the agreement of two different views. Prior arts mostly use cosine similarity as a metric to evaluate how far two feature vectors drift apart.

While under our setting, two views may have different and unaligned nodes, where a simple cosine similarity loses its availability. Hence we propose to leverage the earth mover's distance (EMD) as our similarity measure.

EMD [Rubner et al., 2000; Liu et al., 2020] is the measure of the distance between two discrete distributions, it can be interpreted as the minimum cost to move one pile of dirt to the other. Although prior work has introduced EMD to the CV domain, the adaptation in the graph domain has not been explored yet. Moreover, according to the characteristics of graph data, we also take topology distance into consideration while computing the cost matrix. Through a non-trivial solution, we combine the vanilla cost matrix and topology distance to obtain a rectified cost matrix which makes the cost related to the node similarity and the distance in the graph topology.

The calculation of g-EMD can be formulated as a linear optimization problem. In our case, the two augmented views have feature maps X ∈ RM×d and Y ∈ RN×d respectively, the goal is to measure the distance to transformX to Y. Suppose for each node xi ∈ Rd , it has ti units to transport, and node yj ∈ Rd has rj units to receive. For a given pair of nodes xi and yj , the cost of transportation per unit is Dij , and the amount of transportation is Γij . With above notations, we can define the linear optimization problem as follows:

where t ∈ RM and r ∈ RN are marginal weights for Γ respectively.

The set of all possible transportation matrices Γ can be de-

where 1 is all-one vector with corresponding size, and Π(t, r) is the set of all possible distributions whose marginal weights are t and r.

And the cost to transfer xi to yj is defined as

which indicates that nodes with similar representations prefer to generate fewer matching cost between each other. In addition to directly using node representations dissimilarity matrix as a distance matrix, we also take the topology distance Ψ ∈ RM×N (the smallest hop count between each pair of nodes) into consideration. Nodes are close in topology structure which indicates they may contain similar semantic information. How to combine the representation dissimilarity matrix and topology distance is not a trivial problem. In order not to adjust the original cost matrix D drastically, we adopt sigmoid function S with temperature on topology distance to get re-scale factors S ∈ [0.5, 1]M×N :

where τ ≥ 1 is the temperature factor to control the rate of curve change. We set τ as 2 empirically, and leave the choice of different of re-scale function and the tuning of different temperature factors in future work. With the re-scale factors S, we can update the cost matrix by

where ◦ is Hadamard product. In this way, we combine both topology distance and node representation dissimilarity matrix into distance matrix.

As D is fixed according to distributions X, Y and topology distance, to get g-EMD we need to find the optimal ˜Γ. To solve the optimal ˜Γ, we utilize Sinkhorn Algorithm [Cuturi, 2013] by introducing a regularization term:

where ⟨, ⟩F denotes Frobenius inner product, and λ is a hyperparameter that controls the strength of regularization. With this regularization, the optimal ˜Γ can be approximated as:

where P = e-λD, and v, u are two coefficient vectors whose values can be iteratively updated as

Then the question lies in how to get marginal weights t and r. The weight represents a node's contribution in comparison of two views, where a node should have larger weight if its semantic meaning is close to the other view. Based on this hypothesis, we define the node weight as dot product between its feature and the mean pooling feature from the other set:

where max is to make sure all weights are non-negative, and then both views will be normalized to ensure having the same amount of features to transport.

With optimal transportation amount ˜Γ, we obtain:

Now we can leverage EMD as the distance measure to contrastive loss objective. For any central node vi and its augmented graph views ( ˜ G (i) 1 , ˜ G (i) 2 ), an encoder fθ (e.g. GNN) is applied to get embeddings H(i) 1 andH(i) 2 respectively, then a linear projector gω is applied on top of that to get Z(i) 1 and Z(i) 2 to improve generality for downstream tasks as indicated in [Chen et al., 2020]. Formally, we define the EMD-based contrastive loss for node vi as

where s(x, y) is a function that calculates the similarity between x and y, here we use 1 - EMD(x, y) to replace s(x, y); 1 is an indicator function which returns 1 if i ̸= k otherwise returns 0; and τ is temperature parameter. Adding all nodes in N, the overall contrastive loss is given by:

We summarize our proposed algorithm for non-aligned node-node contrastive learning in Appendix A.

3.4 Unsupervised Adversarial Training

Adversarial training can be considered as an augmentation technique which aims to improve the model's robustness. [Kong et al., 2020] has empirically proven that graph adversarial augmentation on feature space can boost the performance of GNN under a supervised manner. Such a method can be modified for graph contrastive learning as

where θ, ω are the parameters of encoder and projector, D is data distribution, It = BX+δ0 (αt)∩BX(ϵ) where ϵ is the perturbation budget. For efficiency, the inner loop runs M times, the gradient of δ, θt-1 and ωt-1 will be accumulated in each time, and the accumulated gradients will be used for updating θt-1 and ωt-1 during outer update. Equipped with such adversarial augmentation, we complete a more robust selfaligned task. The energy is hopefully transferred between nodes belonging to different categories during max-process, and min-process will remedy such a bad situation to make the alignment more robust. In this way, the adversarial augmentation increases the diversity of samples and improves the robustness of the model.

4 Experiments

We conduct extensive experiments on ten public benchmark datasets to evaluate the effectiveness of RoSA. We use RoSA to learn node representations in an unsupervised manner and assess their quality by a linear classifier trained on top of that. Some more detailed information about datasets and experimental setup can be found in Appendix B, C.

4.1 Datasets

We conduct experiments on ten public benchmark datasets that include four homophilous datasets (Cora, Citeseer, Pubmed and DBLP), three heterophilous datasets (Cornell, Wisconsin and Texas), two large-scale inductive datasets (Flickr and Reddit) and one dynamic graph dataset (CIAW) to evaluate the effectiveness of RoSA. Details of datasets can be found in Appendix B.

4.2 Experimental Setup

Models. For small-scale datasets, we apply a two-layer GCN as our encoder fθ and for the large-scale datasets (Flickr and Reddit), we adopt a three-layer GraphSAGEGCN [Hamilton et al., 2017] with residual connections as the encoder following DGI [Velickovic et al., 2019] and GRACE [Zhu et al., 2020b]. The formulas of encoders can be found in Appendix C. Specifically, similar to [Chen et al., 2020], a projection head which comprises a two-layer nonlinear MLP with BN is added on top of the encoder. Detailed hyperparameter settings are in Appendix C.

Baselines. We compare RoSA with two node-graph constrasting methods DGI [Velickovic et al., 2019], SUBGCON [Jiao et al., 2020]), and four node-node methods GMI [Peng et al., 2020], GRACE [Zhu et al., 2020b], GCA [Zhu et al., 2021b] and BGRL [Thakoor et al., 2021].

4.3 Results and Analysis

Results for homophilous datasets. Table 1 shows the node classification results on four homophilous datasets, some of the reported statistics are borrowed from [Zhu et al., 2020b]. Experiment results show that N-N methods surpass N-G on node classification tasks. And RoSA is superior to all baselines and achieves state-of-the-art performance, and even surpasses the supervised method (GCN), which proves the effectiveness of leveraging EMD-based contrastive loss and adversarial training in non-aligned node-node scenarios. Different from other node-node methods that train on full graphs, our method is trained on various non-aligned subgraphs, which brings more flexibility but also non-alignment challenge. RoSA learns more information from the challenging pretext task. The visualization of cost matrix and transportation matrix in EMD during training is in Appendix E.

Results for heterophilous dataset. Previous works have shown that GCN performs poorly on heterophilous graphs [Pei et al., 2020], because there are a lot of highfrequency signals on such graphs, and GCN is essentially a low-pass filter, where a lot of useful information will be filtered out. Since the design of the encoder is not the focus of our work, we use both GCN and MLP as our encoders in this part.

We compare the performance of our model with DGI, SUBG-CON, GMI, GRACE using either GCN or MLP as encoder, see Table 2. From the statistics, we can summarize three major conclusions: Firstly, the overall performance of SUBG-CON and DGI lags behind the others. This is because SUBG-CON and DGI are the node-graph level contrasting methods that maximize the mutual information between central node representation and its contextual subgraph representation, and under the heterophilous circumstance, the contextual graph representation gathers highly variant features from different kinds of nodes, which renders wrong and meaningless signals.

Secondly, with the same method, the MLP version performs significantly better than its GCN counterpart, which confirms the statement that MLP is more suitable for heterophilous graphs. Furthermore, we can observe that the performance gap between node-global and node-node methods widens when using MLP as the encoder. We suspect such a phenomenon is caused because the GCN encoder loses a large amount of information under a heterophilous setting and makes the effort of other modules in vain.

Thirdly and most importantly, RoSA outperforms other benchmarks on all three datasets, no matter the choice of the encoder, which validates the effectiveness of RoSA for heterophilous graphs. We speculate that RoSA will tighten the distance of nodes of the same class.

Result for inductive learning on large-scale datasets. The experiments conducted above are all under the transductive setting. In this part, the experiments are under the inductive setting where tests are conducted on unseen or untrained nodes. The micro-averaged F1 score is used for both of these two datasets. The results are shown in Table 3, we can see that RoSA works well on large-scale graphs under inductive setting and reaches state-of-the-art performance. An explanation is DGI, GMI and GRACE can not directly work on full graphs, they use the sampling strategy proposed by [Hamilton et al., 2017] in their original work. However, we adopt subsampling (random walk) as our augmentation technique which means our method can seamlessly work on these large graphs. Furthermore, our pretext task is designed for such subsampling which is more suitable for large graphs.

Results for dynamic graphs dataset. In addition, we also test our method on dynamic graphs. For the contrastive task, we consider the adjacent snapshots as positive views because the evolution process is generally "smooth", and the snapshots far away from the anchor are considered as negative views. In CIAW, each snapshot maintains all nodes appeared in the timeline, however, in real-world scenarios, the addition or deletion of nodes happens as time goes by. So in CIAW*, we remove isolated nodes in each snapshot to emulate such a situation. Note that GRACE can not work on CIAW* because CIAW* creates a non-aligned situation, while GRACE is inherently an aligned method. From the statistics in Table 4, RoSA surpasses other competitors and can work well in both situations. Currently, we simply use static GNN encoder with discrete-time paradigms which can be replaced with temporal GNN encoders, and we will leave it for future work.

4.4 Ablation Study

To prove the effectiveness of the design of RoSA, we conduct ablation experiments masking different components under the same hyperparameters. First we replace the EMD-based InfoNCE loss with a regular cosine similarity metric, represented as RoSA w/o EMD (In order to make it computable under such situation, we restrict the same amount of nodes for contrasted views). Second we use the vanilla cost matrix for EMD, named as RoSA w/o TD. Then we remove the adversarial training process, denoted as RoSA w/o AT. Finally, we adopt aligned views contrasting instead of the original nonaligned random walking, named as RoSA Aligned. For a fair comparison, we keep other hyperparameters and the training scheme same. The results is summarized in Figure 3. As we can see, the performance degrades without either EMD, adversarial training or rectified cost matrix, which indicates the effectiveness of the corresponding components. Furthermore, compared to aligned views, the model achieves comparable or even better results under the non-aligned condition, which demonstrates that our model, to a certain degree, solves the non-aligned graph contrasting problem. The experiments of sensitivity analysis are in Appendix D.

5 Conclusion

In this paper, we propose a robust self-aligned framework for node-node graph contrastive learning, where we design and utilize the graph-based earth mover's distance (g-EMD) as a similarity measure in the contrastive loss to avoid explicit alignment between contrasted views. Then we introduce unsupervised adversarial training into graph domain to further improve the robustness of the model. Extensive experiment results on homophilous, non-homophilous and dynamic graphs datasets demonstrate that our model can effectively be applied to non-aligned situations and outperform other competitors. Moreover, in this work we adopt simple random walk with restart as the subsampling technique, and RoSA may achieve better performance if equipped with more powerful sampling methods in future work.

Acknowledgments

This work has been supported in part by National Key Research and Development Program of China (2018AAA0101900), Zhejiang NSF (LR21F020004), Key Research and Development Program of Zhejiang Province, China (No. 2021C01013), Chinese Knowledge Center of Engineering Science and Technology (CKCEST), Zhejiang University iFLYTEK Joint Research Center.