A Probabilistic Model Behind Self-Supervised Learning

TL;DR

This paper introduces a probabilistic generative model for self-supervised learning, unifying various methods and demonstrating improved representation quality, especially in style-dependent tasks.

Contribution

It proposes a generative latent variable model that unifies different SSL approaches and introduces SimVAE, a generative method that enhances representation learning.

Findings

SimVAE outperforms existing SSL methods on simple benchmarks.

The model provides a theoretical framework linking SSL to mutual information.

SimVAE narrows the gap between generative and discriminative methods.

Abstract

In self-supervised learning (SSL), representations are learned via an auxiliary task without annotated labels. A common task is to classify augmentations or different modalities of the data, which share semantic content (e.g. an object in an image) but differ in style (e.g. the object's location). Many approaches to self-supervised learning have been proposed, e.g. SimCLR, CLIP, and DINO, which have recently gained much attention for their representations achieving downstream performance comparable to supervised learning. However, a theoretical understanding of self-supervised methods eludes. Addressing this, we present a generative latent variable model for self-supervised learning and show that several families of discriminative SSL, including contrastive methods, induce a comparable distribution over representations, providing a unifying theoretical framework for these methods. The proposed model also justifies connections drawn to mutual information and the use of a ''projection head''. Learning representations by fitting the model generatively (termed SimVAE) improves performance over discriminative and other VAE-based methods on simple image benchmarks and significantly narrows the gap between generative and discriminative representation learning in more complex settings. Importantly, as our analysis predicts, SimVAE outperforms self-supervised learning where style information is required, taking an important step toward understanding self-supervised methods and achieving task-agnostic representations.