A Simplified Framework for Contrastive Learning for Node Representations

TL;DR

This paper introduces a simplified contrastive learning framework for node embeddings in graphs, using column-wise postprocessing to improve embedding quality and training efficiency, outperforming many existing methods.

Contribution

It proposes a novel column-wise postprocessing technique for graph contrastive learning, replacing complex MLP-based methods and enhancing performance and training speed.

Findings

Up to 1.5% improvement in downstream classification accuracy

Outperforms state-of-the-art on 6 out of 8 benchmarks

Column-wise postprocessing enhances both alignment and uniformity

Abstract

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.