Understanding Contrastive Representation Learning through Alignment and Uniformity on the Hypersphere

TL;DR

This paper analyzes contrastive representation learning by focusing on the properties of alignment and uniformity on the hypersphere, providing theoretical insights and practical metrics that improve understanding and performance.

Contribution

It introduces a theoretical framework linking contrastive loss to alignment and uniformity, along with metrics to quantify these properties and demonstrate their importance.

Findings

Alignment and uniformity are key to contrastive learning success.

Optimizing these properties improves downstream task performance.

Metrics for alignment and uniformity correlate with better representations.

Abstract

Contrastive representation learning has been outstandingly successful in practice. In this work, we identify two key properties related to the contrastive loss: (1) alignment (closeness) of features from positive pairs, and (2) uniformity of the induced distribution of the (normalized) features on the hypersphere. We prove that, asymptotically, the contrastive loss optimizes these properties, and analyze their positive effects on downstream tasks. Empirically, we introduce an optimizable metric to quantify each property. Extensive experiments on standard vision and language datasets confirm the strong agreement between both metrics and downstream task performance. Remarkably, directly optimizing for these two metrics leads to representations with comparable or better performance at downstream tasks than contrastive learning. Project Page: https://tongzhouwang.info/hypersphere Code: https://github.com/SsnL/align_uniform , https://github.com/SsnL/moco_align_uniform