A Simplified Framework for Contrastive Learning for Node Representations

Contrastive learning has recently established itself as a powerful self-supervised learning framework for extracting rich and versatile data representations. Broadly speaking, contrastive learning relies on a data augmentation scheme to generate two versions of the input data and learns low-dimensional representations by optimizing the contrastive loss to identify augmented samples corresponding to the same original entity. In this paper, we investigate the potential of deploying contrastive learning in combination with Graph Neural Networks (GNNs) for embedding nodes in a graph. Specifically, we show that the quality of node representations and model training time can be significantly improved by a simple column-wise post-processing of raw embeddings instead of the row-wise post-processing via multilayer perceptrons (MLPs) that is adopted by the majority of peer methods. This modification results in up to a 1.4% improvement in performance on downstream classification tasks, surpassing the state-of-the-art methods on 6 out of 8 different benchmarks. Additionally, it reduces model training time to at least one-third of the previous time. We justify our choice of post-processing by revisiting the alignment and uniformity paradigm [3], and show that column-wise post-processing improves both “alignment” and “uniformity” of the embeddings.