Joint Embedding Variational Bayes

TL;DR

VJE introduces a probabilistic, reconstruction-free framework for self-supervised learning that models embeddings with a Student-t distribution, providing uncertainty estimates and strong out-of-distribution detection.

Contribution

It presents a novel variational approach with a heavy-tailed distribution and shared uncertainty modeling, advancing non-contrastive self-supervised learning.

Findings

VJE achieves competitive accuracy on ImageNet-1K, CIFAR-10/100, STL-10.

Representation likelihoods enable effective out-of-distribution detection.

The framework models structured feature uncertainty directly in embedding space.

Abstract

We introduce Variational Joint Embedding (VJE), a reconstruction-free latent-variable framework for non-contrastive self-supervised learning in representation space. VJE maximizes a symmetric conditional evidence lower bound (ELBO) on paired encoder embeddings by defining a conditional likelihood directly on target representations, rather than optimizing a pointwise compatibility objective. The likelihood is instantiated as a heavy-tailed Student--\(t\) distribution on a polar representation of the target embedding, where a directional--radial decomposition separates angular agreement from magnitude consistency and mitigates norm-induced pathologies. The directional factor operates on the unit sphere, yielding a valid variational bound for the associated spherical subdensity model. An amortized inference network parameterizes a diagonal Gaussian posterior whose feature-wise variances are shared with the directional likelihood, yielding anisotropic uncertainty without auxiliary projection heads. Across ImageNet-1K, CIFAR-10/100, and STL-10, VJE is competitive with standard non-contrastive baselines under linear and \(k\)-NN evaluation, while providing probabilistic semantics directly in representation space for downstream uncertainty-aware applications. We validate these semantics through out-of-distribution detection, where representation-space likelihoods yield strong empirical performance. These results position the framework as a principled variational formulation of non-contrastive learning, in which structured feature-wise uncertainty is represented directly in the learned embedding space.