[Paper Review] Representation Uncertainty in Self-Supervised Learning as Variational Inference

Hiroki Nakamura, Masashi Okada+1 more|arXiv (Cornell University)|Mar 22, 2022

Domain Adaptation and Few-Shot Learning4 citations

TL;DR

This paper proposes VI-SimSiam, a novel self-supervised learning method that frames SimSiam as variational inference over a power spherical latent space, enabling uncertainty-aware representation learning without labels. By modeling representations probabilistically, VI-SimSiam predicts both features and uncertainty (via concentration parameter κ), achieving competitive accuracy while qualitatively linking lower κ to ambiguous or low-salience images.

ABSTRACT

In this study, a novel self-supervised learning (SSL) method is proposed, which considers SSL in terms of variational inference to learn not only representation but also representation uncertainties. SSL is a method of learning representations without labels by maximizing the similarity between image representations of different augmented views of an image. Meanwhile, variational autoencoder (VAE) is an unsupervised representation learning method that trains a probabilistic generative model with variational inference. Both VAE and SSL can learn representations without labels, but their relationship has not been investigated in the past. Herein, the theoretical relationship between SSL and variational inference has been clarified. Furthermore, a novel method, namely variational inference SimSiam (VI-SimSiam), has been proposed. VI-SimSiam can predict the representation uncertainty by interpreting SimSiam with variational inference and defining the latent space distribution. The present experiments qualitatively show that VI- SimSiam could learn uncertainty by comparing input images and predicted uncertainties. Additionally, we described a relationship between estimated uncertainty and classification accuracy.

Motivation & Objective

• To establish a theoretical connection between self-supervised learning (SSL) and variational inference, particularly for non-contrastive SSL methods like SimSiam.
• To develop a method that learns not only representations but also their uncertainty in a self-supervised manner, addressing the lack of uncertainty estimation in existing SSL frameworks.
• To demonstrate that uncertainty estimates (κ) correlate with downstream classification accuracy and image saliency, enhancing model interpretability.
• To unify SimSiam, SimCLR, and DINO under a common variational inference framework, revealing their underlying probabilistic structure.

Proposed method

• The method interprets SimSiam as variational inference by modeling the latent representation as a power spherical distribution, enabling uncertainty estimation through the concentration parameter κ.
• It formulates the SSL objective as an evidence lower bound (ELBO) with three components: alignment (J_align), uniformity (J_uniform), and KL divergence (J_KL) to regularize the posterior.
• The posterior p(z|x,θ) is modeled as a Product-of-Experts (PoE) of single-modal posteriors, while the approximate posterior q(z|x,φ) is a Mixture-of-Experts (MoE), enabling multi-augmentation consistency.
• The representation uncertainty is encoded via the concentration parameter κ of the power spherical distribution, which is jointly optimized during training.
• The method uses a differentiable approximation of the log-normalizing constant of the power spherical distribution to enable end-to-end training.
• The framework generalizes existing SSL methods: SimSiam as deterministic output with uniform prior on the hypersphere, SimCLR as categorical output, and DINO as a deterministic variant with a momentum encoder.

Experimental results

Research questions

• RQ1 Can self-supervised learning methods like SimSiam be interpreted through the lens of variational inference?
• RQ2 How can representation uncertainty be explicitly modeled and predicted in a self-supervised setting without labeled data?
• RQ3 What is the relationship between the estimated uncertainty (κ) and the model's downstream classification performance?
• RQ4 How do the latent space distributions (e.g., power spherical, von Mises-Fisher) unify different SSL objectives under a common probabilistic framework?

Key findings

• VI-SimSiam successfully learns representation uncertainty without labels, with lower uncertainty (higher κ) corresponding to images that are more salient or well-structured.
• The uncertainty parameter κ is qualitatively linked to image content: images with low saliency or ambiguous features exhibit lower κ values, indicating higher uncertainty.
• The method achieves competitive linear evaluation accuracy on ImageNet-1K, comparable to standard SimSiam, demonstrating that uncertainty learning does not compromise performance.
• A positive correlation is observed between the estimated uncertainty (κ) and classification accuracy, suggesting that higher uncertainty predictions are associated with more confident and accurate predictions.
• The theoretical framework unifies SimSiam, SimCLR, and DINO as special cases of variational inference over different latent distributions, with VI-SimSiam extending SimSiam to include uncertainty estimation.