Probabilistic Contrastive Learning with Explicit Concentration on the Hypersphere

TL;DR

The paper tackles the lack of explicit uncertainty estimation in self-supervised contrastive learning by placing representations on the unit hypersphere and modeling them with a von Mises-Fisher distribution. It introduces an unnormalized form $\psi(\boldsymbol{x}; \boldsymbol{\mu}, \kappa) = \exp(\kappa \boldsymbol{\mu}^T \boldsymbol{x})$ and a learnable concentration parameter $\kappa$ as a direct measure of uncertainty, regulated by an $\ell_2$ penalty to avoid degenerate solutions. A probabilistic embedding alignment loss is proposed, $L_{align} = - \lambda_{align} (\kappa_1 + \kappa_2) \boldsymbol{\mu}_1^T \boldsymbol{\mu}_2$, which is combined with the standard SimCLR objective to yield a total loss that preserves discriminativeness while encoding uncertainty. Empirical results on CIFAR-10-C demonstrate that $\kappa$ tracks corruption severity and enables failure analysis, and concatenating $\kappa$ with features improves OOD detection across several benchmarks. The approach is compatible with multiple contrastive frameworks, offering a scalable path to uncertainty-aware SSL applicable to high-stakes domains like autonomous driving and medical imaging.

Abstract

Self-supervised contrastive learning has predominantly adopted deterministic methods, which are not suited for environments characterized by uncertainty and noise. This paper introduces a new perspective on incorporating uncertainty into contrastive learning by embedding representations within a spherical space, inspired by the von Mises-Fisher distribution (vMF). We introduce an unnormalized form of vMF and leverage the concentration parameter, kappa, as a direct, interpretable measure to quantify uncertainty explicitly. This approach not only provides a probabilistic interpretation of the embedding space but also offers a method to calibrate model confidence against varying levels of data corruption and characteristics. Our empirical results demonstrate that the estimated concentration parameter correlates strongly with the degree of unforeseen data corruption encountered at test time, enables failure analysis, and enhances existing out-of-distribution detection methods.

Figures (4)

• Figure 1: The von Mises-Fisher ($vMF$) distribution with a fixed mean vector and varied concentration values $\kappa$ on a sphere. Our framework equips contrastive learning representations to have the $vMF$ distribution and utilizes the estimated $\kappa$'s as the indications of uncertainty.
• Figure 2: A. Decreasing $\kappa$ implies less concentration and therefore more uncertainty in the representation (left). The associated image corruption is from mild to severe (right). B. The two groups of kappa values (i.e., correctly classified and misclassified) from the test set are significantly different.
• Figure 3: Additional augmentation degrades the quality of uncertainty estimation for specific types of corruptions. For instance, introducing Gaussian noise during training causes the correlation with both Gaussian and Speckle noise to shift from negative to positive.
• Figure 4: Left: Comparison of top-1 classification accuracy on the downstream task over the 1000 training epochs. The deterministic approach represents the original SimCLR approach that learns a one-to-one mapping from an image to a representation. Right: Comparison of correlation between $\kappa$ and levels of data corruption (i.e., uncertainty estimation quality) over the 1000 training epochs.