Uncertainty in Contrastive Learning: On the Predictability of Downstream Performance

TL;DR

This paper investigates whether the uncertainty of contrastive learning representations can be quantified to predict downstream performance, proposing a method based on local data distribution and consistency in embedding space.

Contribution

It introduces a novel approach to estimate uncertainty in contrastive learning embeddings, enabling prediction of downstream task performance from pre-trained representations.

Findings

Uncertainty estimates correlate strongly with downstream accuracy.

Local data distribution and consistency are effective for uncertainty quantification.

Method enhances reliability assessment of contrastive models in critical applications.

Abstract

The superior performance of some of today's state-of-the-art deep learning models is to some extent owed to extensive (self-)supervised contrastive pretraining on large-scale datasets. In contrastive learning, the network is presented with pairs of positive (similar) and negative (dissimilar) datapoints and is trained to find an embedding vector for each datapoint, i.e., a representation, which can be further fine-tuned for various downstream tasks. In order to safely deploy these models in critical decision-making systems, it is crucial to equip them with a measure of their uncertainty or reliability. However, due to the pairwise nature of training a contrastive model, and the lack of absolute labels on the output (an abstract embedding vector), adapting conventional uncertainty estimation techniques to such models is non-trivial. In this work, we study whether the uncertainty of such a representation can be quantified for a single datapoint in a meaningful way. In other words, we explore if the downstream performance on a given datapoint is predictable, directly from its pre-trained embedding. We show that this goal can be achieved by directly estimating the distribution of the training data in the embedding space and accounting for the local consistency of the representations. Our experiments show that this notion of uncertainty for an embedding vector often strongly correlates with its downstream accuracy.