A Simplified Framework for Contrastive Learning for Node Representations
This work addresses the challenge of efficient and effective node embedding for graph-based tasks, offering an incremental improvement over existing methods.
The paper tackles the problem of improving node representation learning in graphs by proposing a simple column-wise postprocessing of embeddings in contrastive learning with Graph Neural Networks, resulting in up to 1.5% improvement in downstream classification tasks and outperforming state-of-the-art methods on 6 out of 8 benchmarks.
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 maximizing a normalized temperature-scaled cross entropy loss (NT-Xent) 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 for embedding nodes in a graph. Specifically, we show that the quality of the resulting embeddings and training time can be significantly improved by a simple column-wise postprocessing of the embedding matrix, instead of the row-wise postprocessing via multilayer perceptrons (MLPs) that is adopted by the majority of peer methods. This modification yields improvements in downstream classification tasks of up to 1.5% and even beats existing state-of-the-art approaches on 6 out of 8 different benchmarks. We justify our choices of postprocessing by revisiting the "alignment vs. uniformity paradigm", and show that column-wise post-processing improves both "alignment" and "uniformity" of the embeddings.