cs.LGMar 26, 2026

Gaussian Joint Embeddings For Self-Supervised Representation Learning

Authors: Yongchao Huang

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

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.

Figures & tables

Figure 1: A comparison of (a) the classic JEPA framework, based on separate encoding and alignment via deterministic latent prediction, and (b) our proposed GJE framework, which models the joint distribution of context and target representations and derives predictions probabilistically. The component with a red outline in (a) represents the deterministic distance loss specific to classic JEPA. ‘sg’ denotes stop gradient , and ε\varepsilon represents injected side information (e.g. physics, action conditioning or noise).
Figure 2: The general Gaussian Mixture Joint Embeddings (GMJE) framework. Dual encoders map the context and target views to a joint embedding ZZ . This embedding is evaluated against a set of KK learnable global mixture parameters ( μk,Σk,πk\mu_{k},\Sigma_{k},\pi_{k} ) to natively model the joint probability density p(zc,zt)p(z_{c},z_{t}) within a symmetric architectural framework.
Figure 3: The Information Bottleneck in the GMJE-MDN architecture. To prevent identity collapse, the parameter network must take only the context embedding zcz_{c} as input to predict the conditional mixture parameters. These generated parameters are then evaluated against the target embedding ztz_{t} inside the conditional mixture model. To prevent representation collapse, this conditional objective is coupled with a marginal loss on zcz_{c} .
Figure 4: Predictive distributions on Dataset A (Separated Branches). (a) Classic JEPA collapses toward the conditional average and fails to recover the three valid branches. Ground-truth samples are shown in gray. (b) The dual-space RBF kernel baseline learns a flexible mean but retains a unimodal Gaussian predictive distribution, producing a broad uncertainty band rather than resolving separate branches. (c) GMJE-EM ( K=1K=1 ) fits a single global ellipse and heavily over-smooths the manifold. (d) GMJE-EM ( K=3K=3 ) separates the density across multiple components, but the fixed covariance ellipses remain geometrically rigid. (e) GMJE-GNG adaptively places prototypes along the non-linear ridges and captures the manifold topology more faithfully. (f) GMJE-MDN produces the closest qualitative match to the ground-truth conditional density by predicting instance-dependent mixture parameters.
Figure 5: Internal parameter routing of GMJE-MDN on Dataset A. (a) The conditional mixing weights remain close to the ground-truth uniform value of 1/31/3 . (b) The learned conditional means align closely with the three underlying branch functions. (c) The predicted conditional standard deviations remain near the true Gaussian observation noise level ϵ=0.05\epsilon=0.05 .
Figure 6: Predictive distributions on Dataset B (Intersecting Branches). The overall trends are consistent with Dataset A but the intersection structure makes the task more challenging. (a) Classic JEPA again collapses toward the conditional average near zero. (b) The dual-space RBF kernel baseline remains unimodal and cannot isolate the intersecting branches. (c–d) The fixed-covariance parametric GMJE baselines capture part of the multi-modal structure but remain limited by global geometric rigidity. (e) GMJE-GNG provides an adaptive topological approximation of the manifold. (f) GMJE-MDN most accurately captures the intersecting conditional density among the compared methods.
Figure 7: Internal parameter routing of GMJE-MDN on Dataset B. The learned means track the three intersecting target branches, the standard deviations remain close to the true observation noise, and the mixing weights remain near the ground-truth value of 1/31/3 across the context domain.
Figure 8: Empirical training dynamics comparing a standard FIFO queue (MoCo) against our proposed SMC particle bank (SMC-GMJE) under a constrained memory capacity ( M=256M=256 ). (Left) Pre-training InfoNCE loss. Note the severe volatility of the FIFO queue starting around Epoch 500. (Center) Linear probing test accuracy, showing SMC maintaining a distinct advantage. (Right) The Effective Sample Size (ESS) of the SMC filter. The sharp ESS drop at ∼100,000\sim 100,000 batches (Epoch 500, after the warmup stage) indicates increasing concentration on a smaller subset of informative particles, coincides with the stabilization of the latent space and the onset of targeted hard-negative mining.
MethodArchitecture TypeCIFAR-10 Accuracy
SimCLR [ 19 ]Symmetric + Large Batch Contrastive81.50%
BYOL [ 30 ]Asymmetric EMA + Predictor79.04%
MoCo v2 [ 35 ]Asymmetric EMA + FIFO Queue79.02%
SMC-GMJE (Ours)Asymmetric EMA + Dynamic Particle Bank74.52%
Table 1: Linear Probing Top-1 accuracy on CIFAR-10 using a ResNet-18 backbone.
Figure 9: t-SNE visualization of real image embeddings (colored by class) and synthetic latent samples (black crosses). (a) A post-hoc GMM fitted to SimCLR places substantial mass between semantic clusters (e.g. empty regions). (b) Unimodal GJE spans the latent space broadly but lacks clear multi-cluster organization. (c) Parametric GMJE generates samples that align more closely with the clustered semantic structure of the learned manifold.
Figure 10: Unconditional image synthesis from latent samples. (a) SimCLR produces blurry, out-of-distribution chimeras. (b) Unimodal GJE suffers from mean collapse, producing nearly identical, heavily smoothed averages with reduced diversity. (c) Parametric GMJE yields the sharpest and most diverse decoded samples among the three methods.
Model FrameworkAvg. log⁡∣Σk∣\log|\Sigma_{k}| (Volume)Avg. Tr(Σk)Tr(\Sigma_{k}) (Variance)Prototype Spread ( L2L^{2} )
SimCLR (Post-Hoc GMM)-878.990.850.55
Unimodal GJE ( K=1K=1 )-636.150.99N/A
Parametric GMJE (EM)-621.581.001.43
Table 2: Geometric properties of the learned latent distributions ( d=128d=128 ). These statistics provide a compact summary of component volume, total variance, and cross-component separation.
More in this topic

Representation Learning

Representation learning develops data transformations that expose useful structure for prediction, transfer, generation, and control.

See all 799 papers

Explore similar work

CardsList
Mar 20, 2026

Var-JEPA: A Variational Formulation of the Joint-Embedding Predictive Architecture - Bridging Predictive and Generative Self-Supervised Learning

The Joint-Embedding Predictive Architecture (JEPA) is often seen as a non-generative alternative to likelihood-based self-supervised learning, emphasizing prediction in representation space rather than reconstruction in observation space. We argue that the resulting separation from probabilistic generative modeling is largely rhetorical rather than structural: the canonical JEPA design (coupled encoders with a context-to-target predictor) mirrors the variational posteriors and learned conditional priors obtained when variational inference is applied to a particular class of coupled latent-variable models, and standard JEPA can be viewed as a deterministic specialization in which regularization is imposed via architectural and training heuristics rather than an explicit likelihood. Building on this view, we derive the Variational JEPA (Var-JEPA), which makes the latent generative structure explicit by optimizing a single Evidence Lower Bound (ELBO). This yields meaningful representations without ad-hoc anti-collapse regularizers and allows principled uncertainty quantification in the latent space. We instantiate the framework for tabular data (Var-T-JEPA) and achieve strong representation learning and downstream performance, improving over T-JEPA across real-world tabular benchmarks while remaining competitive with strong raw-feature baselines.

Moritz Gögl, Christopher Yau

University of Oxford, UK · Health Data Research UK

Joint-Embedding Predictive ArchitectureSelf-Supervised LearningLatent Variable ModelsVariational Inference
AskPDF
Oct 17, 2026

MM-LeJEPA: Correlated Gaussian Regularization for Multimodal Joint-Embedding Predictive Architectures

LeJEPA derives an isotropic Gaussian embedding target by minimizing, under its task assumptions, a Fisher information criterion that controls the task-averaged squared prediction bias. This enables unimodal self-supervised learning without reconstruction decoders, exponential-moving-average teachers, or contrastive negatives. However, its isotropic task prior does not account for the cross-modal structure of multimodal tasks that read shared content. We introduce MM-LeJEPA by extending LeJEPA's task assumptions to multimodal learning. Specifically, we model multimodal tasks as strength-weighted combinations of components that read independent information from the two modalities and components that read the same content from both. Under explicit gradient-moment, correspondence, and symmetry assumptions, we derive a correlated Gaussian target that minimizes a weighted Fisher information functional under a fixed total variance constraint. An analytic transform adapts Sketched Isotropic Gaussian Regularization (SIGReg) to this target, yielding a two-term objective that combines distributional regularization with local-to-global prediction. The extension adds a single scalar hyperparameter, the target cross-modal correlation, and retains the simplicity of LeJEPA. On public audio-visual classification benchmarks, our results are competitive with representative open-source methods. Code will be released.

Haowen Bai, Rui Gong$, Meichen Liu +3

Nanyang Technological University · Hyundai Motor Group Innovation Center Singapore (HMGICS)

Multimodal SSLJoint-Embedding Predictive ArchitectureCross-Modal Representation Learning
AskPDF
Aug 16, 2026

Gaussian-JEPA: Joint-Embedding Predictive Learning for 3D Gaussian Splats

3D Gaussian Splatting (3DGS) represents 3D content with anisotropic primitives that jointly encode geometry and appearance. Fixed-budget encoders consume sampled observations of Gaussian assets, so the same object may be observed through different primitive realizations. Existing self-supervised methods mainly reconstruct masked Gaussian attributes, tying supervision to one sampled realization and requiring an input-space decoder. Latent prediction offers an alternative, but its application to Gaussian tokens requires targets that accommodate coupled attributes and heterogeneous spatial support. We introduce Gaussian-JEPA, which predicts representations of held-out Gaussian token blocks from visible context. An online encoder processes the context, while a shared exponential-moving-average encoder supplies stop-gradient features for multi-scale targets. Complementary target projections and feature-space grounding provide latent supervision without reconstructing Gaussian attributes. We evaluate the features under Gaussian resampling, partial observations, and renderable shape completion, together with transfer to part segmentation and object classification. Compared with matched reconstruction pretraining, Gaussian-JEPA is more consistent across resampled inputs, retains more instance information under partial observations, and provides stronger frozen features for Gaussian completion. These results support latent prediction as an effective objective for reusable 3D Gaussian representations. Code is on the project page (https://amazingren.github.io/Gaussian-JEPA/).

Bin Ren, Qi Ma, Yue Li +7

Mohamed bin Zayed University of Artificial Intelligence (MBZUAI), AE · ETH Zürich, CH · University of Amsterdam, NL +2

Self-Supervised Pre-Training3D Gaussian SplattingJoint-Embedding Predictive Architecture
AskPDF