Gaussian Joint Embeddings For Self-Supervised Representation Learning

TL;DR

This paper introduces Gaussian Joint Embeddings (GJE) and Gaussian Mixture Joint Embeddings (GMJE), probabilistic models for self-supervised learning that improve multi-modal representation and uncertainty estimation.

Contribution

It proposes a probabilistic framework for self-supervised learning that models joint densities, addressing collapse issues and enabling principled uncertainty and latent geometry control.

Findings

GMJE recovers complex conditional structures in multi-modal tasks.

GMJE learns competitive discriminative representations.

Latent densities from GMJE are better suited for unconditional sampling.

Abstract

Self-supervised representation learning often relies on deterministic predictive architectures to align context and target views in latent space. While effective in many settings, such methods are limited in genuinely multi-modal inverse problems, where squared-loss prediction collapses towards conditional averages, and they frequently depend on architectural asymmetries to prevent representation collapse. In this work, we propose a probabilistic alternative based on generative joint modeling. We introduce Gaussian Joint Embeddings (GJE) and its multi-modal extension, Gaussian Mixture Joint Embeddings (GMJE), which model the joint density of context and target representations and replace black-box prediction with closed-form conditional inference under an explicit probabilistic model. This yields principled uncertainty estimates and a covariance-aware objective for controlling latent geometry. We further identify a failure mode of naive empirical batch optimization, which we term the Mahalanobis Trace Trap, and develop several remedies spanning parametric, adaptive, and non-parametric settings, including prototype-based GMJE, conditional Mixture Density Networks (GMJE-MDN), topology-adaptive Growing Neural Gas (GMJE-GNG), and a Sequential Monte Carlo (SMC) memory bank. In addition, we show that standard contrastive learning can be interpreted as a degenerate non-parametric limiting case of the GMJE framework. Experiments on synthetic multi-modal alignment tasks and vision benchmarks show that GMJE recovers complex conditional structure, learns competitive discriminative representations, and defines latent densities that are better suited to unconditional sampling than deterministic or unimodal baselines.