Understanding InfoNCE: Transition Probability Matrix Induced Feature Clustering

TL;DR

This paper provides a theoretical understanding of InfoNCE by modeling data augmentation dynamics with a transition probability matrix, and introduces SC-InfoNCE, a new loss that improves feature clustering and performance across multiple domains.

Contribution

It introduces a feature space model with a transition matrix to explain InfoNCE's behavior and proposes SC-InfoNCE, a flexible loss function for better feature alignment.

Findings

SC-InfoNCE improves feature clustering in various domains.

Theoretical model explains InfoNCE's optimization behavior.

Experiments show consistent performance gains across datasets.

Abstract

Contrastive learning has emerged as a cornerstone of unsupervised representation learning across vision, language, and graph domains, with InfoNCE as its dominant objective. Despite its empirical success, the theoretical underpinnings of InfoNCE remain limited. In this work, we introduce an explicit feature space to model augmented views of samples and a transition probability matrix to capture data augmentation dynamics. We demonstrate that InfoNCE optimizes the probability of two views sharing the same source toward a constant target defined by this matrix, naturally inducing feature clustering in the representation space. Leveraging this insight, we propose Scaled Convergence InfoNCE (SC-InfoNCE), a novel loss function that introduces a tunable convergence target to flexibly control feature similarity alignment. By scaling the target matrix, SC-InfoNCE enables flexible control over feature similarity alignment, allowing the training objective to better match the statistical properties of downstream data. Experiments on benchmark datasets, including image, graph, and text tasks, show that SC-InfoNCE consistently achieves strong and reliable performance across diverse domains.