A Simplified Framework for Contrastive Learning for Node Representations

Contrastive learning has become one of the most popular self-supervised methods for learning rich and versatile data representations — a success largely owing to its impressive performance across domains ranging from image classification (Chen et al., 2020; Grill et al., 2020; He et al., 2020; Hjelm et al., 2018; Oord et al., 2018a; Tian et al., 2020) to sequence modeling (Kong et al., 2019; Oord et al., 2018b; Wu et al., 2020a) and adversarial learning (Ho and Nvasconcelos, 2020; Grill et al., 2020). By and large, contrastive learning’s general outline consists in generating two (or more) perturbed versions of the data and in subsequently learning to recognize representations from “positive pair” (i.e. embeddings corresponding to the same original entity) and to embed them close to one another in the embedding space. Conversely, “negative pairs” (i.e. embeddings corresponding to different original entity) are encouraged to lie farther apart. Remarkably, this simple “model-free” learning of differences (Gutmann et al., 2022) between data views has been shown to achieve state-of-the-art performance in many downstream supervised tasks when combined with a simple linear classifier (Chen et al., 2020). This result is quite surprising as the method performs well not only relative to other unsupervised embedding algorithms, but also to supervised end-to-end methods whose primary aim is to learn embeddings tailored to a particular task.

Contrastive Learning for Graphs.    Inspired by its success across various machine learning applications, contrastive learning has more recently been adopted by the graph representation learning community. In this case, contrastive learning is usually coupled with Graph Neural Networks (GNNs) (Scarselli et al., 2008; Kipf and Welling, 2016a; Zhou et al., 2020; Wu et al., 2020b) to learn a new class of node embeddings. Current contrastive methods for graphs can be classified into two categories, depending on which mutual information seeks to be maximized.

1. Local-to-global mutual information maximization, such as those used by Deep Graph InfoMax (DGI) (Velickovic et al., 2019) and Multi-View Graph Representation Learning (MVGRL) (Hassani and Khasahmadi, 2020). These methods seek to learn node embeddings by maximizing the mutual information between representations of the original data at different scales: DGI considers the mutual information between local (node) representations obtained by passing the original graph through a GNN encoder and a global graph embedding obtained by summing all local representations. Negative samples are generated by a corruption function separately. MVGRL, on the other hand, leverages diffusion processes to generate additional graphs and obtain corresponding additional pairs of node and graph representations. A pair of node and graph representations from different graph are considered as negative pairs. Both DGI and MVGRL use mutual information estimator parameterized by a neural network to obtain probability score for each local and global representation pair. More specifically, letting $G$ and $\widetilde{G}$ denote the original and the corrupted graphs respectively, the DGI loss writes as:

where $\mathcal{D}:{\mathbb{R}}^{F}\times{\mathbb{R}}^{F}\rightarrow{\mathbb{R}}$ is a mutual information estimator parameterized by MLPs, ${\bm{s}}\in{\mathbb{R}}^{F}$ is the summary vector of the original graph, and ${\bm{u}}_{i}$, $\widetilde{{\bm{u}}}_{j}\in{\mathbb{R}}^{F}$ are node representations from the original graph and its corrupted version respectively with $i=1,...,N$ and $j=1,...,M$.

2. Local-to-local mutual information maximization, such as Graph Contrastive Representation Learning (GRACE) (Zhu et al., 2020) and its extension, Graph Contrastive representation learning with Adaptive augmentation (GCA) (Zhu et al., 2021). Contrary to the previous approaches, these methods directly define similarity scores using the dot product between node pairs and foregoes the need for a parameterized mutual information estimator. These methods typically project their node representations onto the unit hypersphere, which is known to improve training stability and is more easily separable by downstream linear classifiers. GRACE and GCA generate two views of the original graph by randomly perturbing both edges and features (typically through a combination of random edge deletions and feature masking). Note that perturbations here are used to generate alternative versions of the same data rather than negative examples. Node representations are then learned through the maximization of a normalized temperature-scaled cross entropy loss (NT-Xent) (Chen et al., 2020) between node pairs, which amounts to solving a classification problem: for each node, in each version of the data, the algorithm seeks to recognize the embeddings from the same original node amongst all other embeddings in both versions of the graph. The objective function is provided in Equations (2) and (3). Despite the impressive performance showed by GRACE and GCA, they suffer from a few drawbacks. Both methods rely on row-wise postprocessing of a raw embedding matrix via a MLP with a nonlinear activation function. This choice was first inspired by the method’s vision counterpart (Chen et al., 2020) where it was empirically shown that utilizing a 2-layers MLPs with a nonlinear activation function after the ResNet (He et al., 2016) improved the expressive power and quality of representations compared to those with a linear activation function or without postprocessing step. However, the use of additional postprocessing via MLPs considerably complexifies the model architecture and its effect on the resulting representations is still ill-understood.

Non-contrastive Learning for Graphs. The previous set of methods have some computational difficulties when processing large graph as the evaluation of the normalizing factor for each node requires the computation of pairwise cosine similarities between all negative pairs. Consequently, newer methods such as Bootstrapped Graph Latents (BGRL) (Thakoor et al., 2021) and Canonical Correlation Analysis inspired Self-Supervised Learning on Graphs (CCA-SSG) (Zhang et al., 2021) propose to leverage non-contrastive objectives which do not incorporate the negative samples. In CCA-SSG, for instance, the goal is to directly minimize the distance between positive pairs, while additionally decorrelating features in different dimensions. The objective function writes as:

where $U,V\in{\mathbb{R}}^{N\times F}$ are node representation matrix of two perturbed versions of the original graph, $\lambda>0$ be the hyperparameter, and $\|\cdot\|_{F}$ denotes the Frobenius norm. In consequences, both BGRL and CCA-SSG enjoy low computational complexities by disregarding any computations with negative samples. However, BGRL requires additional components such as exponential moving average (EMA) and Stop-Gradient — making the method both more expensive and harder to interpret. Although CCA-SSG relies on a simple and symmetric architecture without additional components, this feature-wise contrastive method typically requires larger embedding dimension to improve the quality of its representations.

Contributions. The contributions of our paper are threefold. (1) We propose a simplified framework for contrastive learning for node representations and refer to our method as Contrastive Learning for Node Representations (CLNR). CLNR follows the core structure of GRACE (Zhu et al., 2020) and GCA (Zhu et al., 2021) which generate two views of an input graph via random data augmentation, use a shared GNN encoder, and find node representations by maximizing the NT-Xent loss (Chen et al., 2020). In contrast to GRACE and GCA, CLNR simplifies the model architecture by replacing row-wise postprocessing of a node representation matrix via a MLPs with a simple column-wise standardization. We show this substitution considerably improves the quality of representations and model training time. (2) We provide a finer understanding on the effect of various postprocessing steps and justify our choice of postprocessing (column-wise standardization) by revisiting the “alignment and uniformity paradigm” of Wang and Isola (2020) which was proposed to analyze the quality of representations on the unit hypersphere. (3) We provide experimental evidence of the quality of the induced node embeddings by using them as input for downstream classification. We show that CLNR provides superior results compared to state-of-the-art approaches, including CCA-SSG and GRACE.

Notations. For a given graph $G=(X,A)$, $X\in{\mathbb{R}}^{N\times D}$ is the feature matrix and $A\in\{0,1\}^{N\times N}$ is the adjacency matrix. Here, $N$ denotes the number of nodes in the graph $G$, while $D$ denotes the feature dimension.

Structure of the paper. We introduce the general contrastive learning framework for node representations and propose our simplified framework in Section 2. Experiments are given in Section 3 and discussion on various postprocessing step is provided in Section 4. Additional experiments are present in Section 5 and conclusions are presented in Section 6.

We begin by presenting the general framework for contrastive learning for node embddings, and focus on studying the choice of the embedding postprocessing.

In this section, we showcase the potential of our simplified framework for node constrative learning by deploying it on a number real-life datasets and comparing it against existing peer methods.

Real-life Datasets. We evaluate the quality of our node representations using a node classification task on real-life data. We consider eight node classification benchmarks: Cora, CiteSeer, Pubmed, Amazon Computers, Amazon Photo, Coauthor CS, Coauthor Physics, ogbn-arxiv. The first seven benchmarks are used in Yang et al. (2016) and Shchur et al. (2018), and ogbn-arxiv is from the Open Graph Benchmark datasets (Hu et al., 2020). We adopt the public splits for the citation network datasets (Cora, Citeseer, Pubmed), and 1:1:8 train/validation/test splits for Amazon and Coauthor network datasets (Amazon Computer, Amazon Photo, Coauthor CS, Coauthor Physics). Details of the datasets are given in Appendix A.1.

Experimental Setup. We follow the evaluation protocol suggested by Velickovic et al. (2019), which consists of two stages: 1) Pretraining and 2) Linear evaluation. First, we begin by training a GNN encoder using the NT-Xent loss described in Equation (2), thereby yielding node embeddings for the entire graph. We then fit a simple linear classifier on a subset of nodes (training set). We then evaluate the classification accuracy on the remaining nodes (test set). All experiments are done with Pytorch and Pytorch Geometric module on GPU with 16 GB memory.

In this section, we propose to delve deeper in the comparison between CLNR and GRACE, and try to explain the superiority of the embeddings obtainted by CLNR. To this end, we revisit the “alignment and uniformity paradigm” proposed by Wang and Isola (2020).

Alignment vs. Uniformity. Noting the success of mapping features to the unit hypersphere and restricting their analysis to those, Wang and Isola (2020) propose an analysis of contrastive losses in terms of two quantities: (a) Alignment, which quantifies how close representations from positive pairs (i.e. embeddings corresponding to the same entity) are mapped to one another:

In the previous expression, $f(x)$ denotes the embedding for the nodes, $\alpha>0$ is any positive exponent (usually chosen to be equal to 2); And (b) uniformity, defined as the logarithm of the average pairwise Gaussian potential:

Uniformity essentially measures uniformly distributed the features are on the unit sphere. Wang and Isola (2020) indeed posit that data distributions that are more uniformly spread on the hypersphere preserve more information. They then proceed to show that the NT-Xent loss can asymptotically be written as the sum of a uniformity and an alignment term, so that it optimizes for both. Building off of these observations, the authors suggest that explicitly optimizing for both uniformity and alignment might yield better results than using the classical negative cross-entropy loss. However, the weighted loss that they propose $\ell=\ell_{\text{uniform}}+\lambda\ell_{\text{align}}$ has the disadvantage of introducing yet another parameter $\lambda$ that has to be cross-validated for.

We begin by plotting the learned representations on Cora dataset using GRACE and our method, CLNR in Figure 1. As expected, our batch-normalization-based method significantly improved the uniformity of the learned representations.

Figure 2 shows the alignment (Equation (4)) vs uniformity (Equation (5)) metrics achieved for each of the methods on 6 different datasets. We observe that sole CLNR is located in left-bottom side of nCLNR (no postprocessing) for 5 out of 6 benchmark datasets. This implies that column-wise postprocessing improves both alignment and uniformity compared to the one without postprocessing. On the other hand, GRACE is shown consistently at the right-bottom corner of the figures. This implies that row-wise postprocessing focuses on improving alignment of representations at the expense of losing significant amount of uniformity. The only exception is Photo where GRACE is also located in left-bottom side of nCLNR and exhibits better node classification accuracy compared to CLNR. We might expect that GCLNR which uses both row-wise and column-wise postprocessing improve the quality of representations compared with the ones using just one of them. Interestingly, we observe that GCLNR is consistently located at the middle of CLNR and GRACE which implies that we can not fully recover the uniformity by column-wise postprocessing once the row-wise postprocessing is done. Surprisingly, dCLNR which additionally decorrelates different features has a similar value of alignment and uniformity metrics with nCLNR.

We conclude this paper with a discussion on the additional benefits of CLNR compared to existing approaches. We focus this discussion on two points: (a) efficiency of the embedding method (as measured by the accuracy of downstream node classifiers as a function of the embedding size) and (b) robustness to perturbations. Throughout this discussion, we include in the comparison a non-contrastive competitor, CCA-SSG (Zhang et al., 2021) due to its impressive performance and scalability.

Efficiency of Embedding. We now study the effect of embedding dimension on the quality of node representations. From Figure 3, we observe that CLNR shows consistent competitive performance across all embedding dimensions on all benchmark datasets. GRACE also shows good performance across all embedding dimensions on all benchmark datasets except CiteSeer. Conversely, we observe that nCLNR which does not postprocess the raw embeddings and CCA-SSG which uses a non-contrastive loss are considerably more sensitive to the dimension.

Robustness of Methods. Finally, we compare nCLNR, CLNR, GRACE, and CCA-SSG in terms of robustness on edge perturbation. Let $|A|$ be the number of edges in an input graph $G$. For given $p\in(0,1)$ we generate adjacency matrix noise $\Delta_{A}\in{\mathbb{R}}^{N\times N}$ containing $|A|\times p$ number of random undirected edges. We then generate a perturbed adjacency matrix with $\widetilde{A}=A+\Delta_{A}$ and we train GCN encoder with the perturbed graph $\widetilde{G}=(X,\widetilde{A})$ instead of $G$. The results are summarized in Figure 4. Figure 4 shows node classification accuracy of the four models on 5 benchmark datasets along with an increasing value of $p$. We observe that CLNR consistently outperforms all other methods on all benchmark datasets except CiteSeer.

We have introduced self-supervised contrastive learning for node representations (CLNR), a simplified but powerful method for learning effective node representations. By replacing the expensive MLP-based row-wise postprocessing of GRACE (Zhu et al., 2020) with a simple batch normalization, we simplify the architecture of existing contrastive learning methods based on NT-Xent loss and considerably its performance, scalability and robustness. To explain this success, we revisit the alignment and uniformity paradigm. We provided empirical evidence that the column-wise postprocessing was sufficient to improve both alignment and uniformity of the embeddings, while row-wise postprocessing improves the alignment of node representations at the expense of losing significant amounts of uniformity.