PaperlayerView exact revision on arXiv

Source-linked AI summary

Gaussian Joint Embeddings For Self-Supervised Representation Learning

Yongchao Huang

arXiv:2603.26799v1cs.LG

TL;DR

Deterministic predictive self-supervision can average multimodal targets and often relies on architectural asymmetries to avoid collapse. The paper introduces GJE and GMJE, which model joint latent densities for closed-form conditional inference, uncertainty, and covariance-aware geometry. Across synthetic and vision tasks, GMJE recovers complex conditional structure, learns strong discriminative representations, and supports unconditional latent sampling.

  • Problem

    Deterministic squared-loss prediction can average genuinely multimodal targets, while predictive methods lack explicit uncertainty and often rely on asymmetries to prevent collapse.

    Source: §1 · p. 5, §2 · p. 6

  • Method

    GJE and GMJE model joint Gaussian or Gaussian-mixture densities over context and target representations for probabilistic conditional inference and covariance-aware regularization.

    Source: §1 · pp. 5–6

  • Results

    GMJE recovers complex multimodal conditional structure, achieves strong discriminative performance, and defines a continuous latent density suitable for unconditional sampling.

    Source: §1 · p. 6, §7 · p. 48

  • Takeaways & Limitations

    Generative joint modeling provides a unified probabilistic approach to multimodal representation learning and latent-space sampling.

    Source: §6 · p. 46, §7 · p. 48

    Single-Gaussian GJE remains unimodal, and GMJE-MDN and prototypical EM rely on a fixed, predefined number of components.

    Source: §4.4.2 · p. 23, §5.1 · p. 37

Abstract

from arXiv · show

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.

1 Introduction

The paper replaces deterministic latent prediction with probabilistic joint modeling to address multimodal targets and representation collapse. It extends GJE to GMJE, develops optimization remedies, connects contrastive learning to the framework, and supports generative latent sampling.

Motivation

  • Deterministic squared-loss prediction can average across multimodal target modes, while predictive methods often require asymmetries or regularization to avoid collapse.The conditional mean may fall between typical modes, and predictive losses do not explicitly enforce global dispersion.Source: §1 · p. 5

GJE

  • GJE models the joint density of context and target representations, providing conditional inference, uncertainty estimates, and covariance-aware control of latent geometry.Its joint negative log-likelihood combines conditional matching with marginal density regularization.Source: §1 · pp. 5–6

Optimization

  • The paper identifies the Mahalanobis Trace Trap as a failure of naive empirical batch optimization that can weaken the intended cross-view learning signal.The proposed remedies are designed to preserve useful optimization under richer joint distributions.Source: §1 · p. 6

GMJE

  • GMJE extends GJE with Gaussian mixtures to represent multimodal structure and offers prototype, GMJE-MDN, GMJE-GNG, and SMC-GMJE variants.These variants span parametric, adaptive, and non-parametric settings while addressing the Mahalanobis Trace Trap.Source: §1 · p. 6

Contrastive connection

  • Standard contrastive learning is interpreted as a degenerate non-parametric limiting case of GMJE, motivating SMC memory banks with dynamically weighted samples.SMC replaces unweighted FIFO queues and prioritizes more informative samples.Source: §1 · p. 6

Generative densities

  • GMJE learns continuous latent densities that support unconditional image synthesis through direct sampling from the learned latent distribution.The paper reports this generative capability alongside representation-learning results.Source: §1 · p. 6

2 Background: From Deterministic Prediction to Generative Joint Modeling

The background contrasts deterministic JEPA-style prediction with probabilistic joint modeling. Classic prediction can average multimodal targets, lacks native uncertainty and global dispersion control, and commonly relies on architectural asymmetries, whereas GJE models the joint representation structure.

Classic JEPA

  • Classic JEPA predicts target representations from context through a deterministic latent-space objective.Its framework uses separate encoding and alignment through deterministic prediction.Source: §2 · p. 6, Figure 1 · p. 7

Deterministic prediction

  • Under squared-error risk, deterministic prediction returns the conditional mean, which can lie between typical modes in multimodal settings and provides no native predictive uncertainty.This limitation arises because no conditional distribution is modeled explicitly.Source: §2 · pp. 6–7

Representation geometry

  • Standard JEPA does not explicitly optimize a joint density or marginal term, so it lacks direct global dispersion control and makes collapse avoidance more delicate.Related predictive methods commonly add variance, covariance, or other auxiliary regularization.Source: §2 · p. 7

Training stability

  • Predictive self-supervised methods often use stop gradients, predictor asymmetry, or EMA target networks to stabilize training and prevent degenerate solutions.These mechanisms introduce architectural asymmetry rather than modeling the full joint representation distribution.Source: §2 · p. 7

Generative joint modeling

  • GJE replaces a black-box predictor with a joint distribution over concatenated context and target representations, deriving prediction from learned dependency structure.Its objective retains variance-covariance structure absent from purely deterministic prediction losses.Source: §2 · p. 7

3 Gaussian Joint Embeddings (GJE)

GJE offers dual sample-space and primal feature-space Gaussian formulations for probabilistic representation alignment. The dual view provides memory-based, context-dependent inference, while the primal view is compact and symmetric but exposes unimodality, homoscedasticity, and empirical optimization limitations that motivate GMJE.

3 Gaussian Joint Embeddings (GJE)

  • GJE encodes context and target views, concatenates their embeddings, and optimizes a joint likelihood that factorizes into conditional prediction and marginal regularization.The framework focuses on unconditional alignment here, while allowing optional conditioning on side information.Source: §3 · p. 8

3.1 The Dual Formulation: Sample-Space Gaussian Process Joint Embeddings (GPJE)

  • GPJE models sample similarities with a Gaussian Process, using context Gram matrices and target embeddings to produce non-parametric conditional predictions.Inference combines cross-similarities, the inverse training Gram matrix, and the stored target memory bank.Source: §3.1 · p. 9
  • The dual GP objective is asymmetric and can collapse because it regularizes context-space volume without an equivalent target-space expansionary term.The described failure maps target embeddings toward zero and subsequently collapses the context space.Source: §3.1 · pp. 9–10
  • GPJE inference is a weighted combination of the training target memory bank, but exact Gram-matrix inversion costs O(N3), motivating Random Fourier Feature approximations.The predictive output has dimensions [1 × dt].Source: §3.1 · p. 10

3.2 The Primal Formulation: Feature-Space Covariance

  • Primal GJE models feature covariance in a d×d joint space, yielding a closed-form linear conditional predictor and constant predictive covariance without memory at inference.Its covariance is partitioned into context, target, and cross-covariance blocks.Source: §3.2 · pp. 10–12, §3.3 · pp. 12–13
  • Naive empirical primal optimization creates the Mahalanobis Trace Trap: covariance cancellation makes the data-fit term constant, leaving volume minimization to drive dimensional collapse.The proposed remedy is to invert the regularizer by maximizing differential entropy or optimizing cross-covariance directly.Source: §3.2 · pp. 11–12

3.3 Sample Space vs. Feature Space

  • Dual and primal GJE are geometrically distinct but become algebraically equivalent under a standard linear kernel, with the dual O(N3) formulation reducing to the primal O(d3) covariance framework.The equivalence links sample-space similarity modeling to feature-space covariance modeling.Source: §3.3 · pp. 12–13
  • Both single-Gaussian formulations are unimodal, so they average distinct target branches; primal GJE also has context-independent covariance that can reduce certainty for unambiguous inputs.These limitations motivate the transition to Gaussian Mixture Joint Embeddings.Source: §3.3 · p. 13

4 (Primal Space) Gaussian Mixture Joint Embeddings (GMJE)

This section develops GMJE as a symmetric primal-feature-space Gaussian mixture framework for modeling multimodal conditional structure with probabilistic marginal regularization and a normalized generative latent density. It covers prototype-based, conditional MDN, topology-adaptive, and non-parametric memory-bank formulations, including connections to contrastive learning, SMC-based memory management, and unconditional embedding generation.

4 (Primal Space) Gaussian Mixture Joint Embeddings (GMJE)

  • GMJE generalizes GJE with a symmetric primal-space joint Gaussian mixture that models complex dependencies and multimodal alignments without architectural anti-collapse heuristics.The joint objective preserves marginal distributions and allows gradients to flow through both encoders while the context embedding routes predictions across mixture components.Source: §4 · pp. 13–14, §4.2 · p. 16

4.1 From Unimodal to Multi-modal Joint Modeling

  • Gaussian mixtures provide a universal finite-component approximation for sufficiently smooth continuous densities, with required component count determined by geometric complexity.This theoretical capacity motivates using mixtures to represent irregular, multimodal self-supervised representation distributions.Source: §4.1 · pp. 14–15

4.2 The GMM Formulation for Joint Embeddings

  • The joint GMM preserves explicit context and target marginals while yielding a closed-form conditional GMM whose context-dependent responsibilities select appropriate target modes.This avoids replacing symmetric joint modeling with a conditional-only predictor that would lose marginal geometric regularization.Source: §4.2 · pp. 15–16

4.3 Scalable Optimization via Learnable Prototypes

  • Prototype GMJE optimizes global mixture means and covariance through differentiable relaxed EM, balancing soft-clustering data fit against covariance-volume regularization that spreads prototypes across the latent manifold.Shared covariance improves tractability and stability; component-specific covariances increase expressiveness but require K matrix inversions and singularity safeguards.Source: §4.3 · pp. 17, 19

4.4 GMJE-MDN: Mixture Density Network (MDN) for GMJE Learning

  • Because MSE and cross-entropy converge toward conditional averages, GMJE-MDN introduces instance-conditioned mixture parameters to represent complete multimodal conditional distributions for multi-valued mappings.The MDN extension addresses the limitation that averaging valid modes can produce an invalid target, while retaining a fixed number of mixture components.Source: §4.3 · pp. 19–20, §4.4 · p. 20

4.4.1 MDN

  • MDNs can represent arbitrary conditional densities using Gaussian kernels, and isotropic components remain globally expressive because their centers capture cross-dimensional dependencies through spatial arrangement.Full covariance matrices are therefore optional rather than necessary for universal approximation, although each kernel assumes local component independence.Source: §4.4.1 · p. 20
  • Standard MDNs dynamically predict conditional mixture means, variances, and weights from x, with exponential variances enabling input-dependent uncertainty instead of constant variance.Training minimizes the negative log-likelihood of the conditional mixture; softmax weights remain normalized and exponential variances remain positive.Source: §4.4.1 · pp. 20–21

4.4.2 GMJE-MDN

  • GMJE-MDN factorizes joint likelihood into a context marginal and context-conditioned target mixture, preventing identity leakage while preserving symmetric anti-collapse regularization.The parameter network observes only zc and predicts µk(zc), Σk(zc), and πk(zc); the marginal covariance term maintains a well-conditioned context space.Source: §4.4.2 · pp. 21–22, Figure 3 · p. 23

4.5 GMJE-GNG: Dynamic Prototype Discovery via Growing Neural Gas

  • GMJE-GNG removes the fixed-K assumption by growing topology-preserving prototypes in high-error regions, representing both mixture modes and connecting manifold ridges.GNG also reduces initialization sensitivity associated with gradient-based optimization or k-means on overlapping or elongated joint distributions.Source: §4.5 · pp. 23–24

4.5.1 Mathematical Formulation

  • GNG updates winners and their graph neighbors toward sampled joint embeddings, then inserts prototypes between high-error neighboring nodes to reduce local density mismatch.After growth, the graph’s vertices directly provide the learned global prototypes.Source: §4.5.1 · p. 24, §4.5.2 · p. 24

4.5.2 The GNG Algorithm for GMJE Component Learning

  • Edge aging removes obsolete connections, allowing GMJE-GNG to separate disjoint semantic clusters while retaining O(Kmax · D) per-sample time and O(Kmax · D + Kmax^2) space complexity.The method stores prototype vectors and an adjacency matrix, making its memory overhead smaller than large non-parametric memory banks when Kmax is bounded.Source: §4.5.2 · p. 25

4.6 Data as Mode GMJE (DaM-GMJE): Contrastive Learning as Non-Parametric GMJE

  • DaM-GMJE treats every memory-bank or batch pair as an equiprobable Gaussian component, interpreting contrastive learning as a non-parametric KDE-like GMJE limit.Standard contrastive assumptions impose isotropic block-diagonal covariance, zero cross-covariance, L2 normalization, and low temperature, yielding the familiar instance-discrimination structure.Source: §4.6 · pp. 26–27
  • Standard contrastive learning is mathematically equivalent to a non-parametric, isotropic Gaussian mixture model within GMJE.This equivalence requires isotropic block-diagonal covariance with zero cross-covariance and instance-level modes, limiting relational modeling and causing class collisions.Source: §4.6 · p. 28

4.7 Sequential Monte Carlo (SMC) for Dynamic Memory Bank Optimization

  • SMC replaces FIFO’s uniform prior with a probability-weighted particle system that retains informative hard negatives instead of discarding them solely due to age.FIFO can discard rare classes or hard negatives while retaining redundant easy negatives; SMC dynamically weights and resamples particles according to current-batch likelihoods.Source: §4.7 · p. 29

4.7.1 The Memory Bank as a Particle System

  • SMC treats stored embeddings as weighted particles approximating the latent marginal and upgrades the conditional contrastive objective into an SMC-weighted InfoNCE loss.The particle weights are normalized to sum to one, yielding a weighted kernel density estimate of the marginal distribution.Source: §4.7.1 · p. 30

4.7.2 Sequential Importance Weighting and Resampling

  • At each iteration, SMC updates particle weights from query likelihoods, resamples M particles proportionally, and concentrates capacity on dense hard-negative regions while preventing weight degeneracy.Highly weighted particles may be retained or cloned, low-weight particles are discarded, and Effective Sample Size monitors concentration before resampling.Source: §4.7.2 · p. 30

4.7.3 Algorithm: SMC-GMJE for Dynamic Memory Bank Optimization

  • SMC-GMJE is a probabilistically rigorous, drop-in contrastive-learning upgrade with O(B · M · d) time complexity, while SMC can also optimize full-covariance parametric GMJE.Resampling transfers probability mass into particle multiplicity, allowing uniform weights after resampling; generalized SMC reweights particles using the full joint negative log-likelihood.Source: §4.7.3 · pp. 32–33

4.8 Generative GMJE: Sampling from the Learned Latent Manifold

  • GMJE explicitly learns a continuous, closed-form latent density whose marginalized target mixture supports efficient unconditional sampling of novel embeddings without real input data.Sampling selects a mixture component from the learned categorical weights before drawing from its Gaussian distribution.Source: §4.8 · p. 33
  • Given a selected mixture mode, GMJE samples the continuous target embedding z_t from its marginalized Gaussian component N(µ_t,k, Σ_tt,k).The target covariance is computed through a Cholesky decomposition, Σ_tt,k = A_kA_k^T.Source: §4.8 · p. 33
  • GMJE’s normalized latent density forms a strict generative manifold, unlike discriminative self-supervised spaces where random latent coordinates yield OOD vectors.This supports generative replay, latent data augmentation, and unconditional image synthesis, evaluated later in Section 5.Source: §4.8 · p. 34
  • The parametric GMJE branch learns K compressed full-covariance prototypes without a memory bank for discriminative or generative downstream tasks.Its variants are globally static, conditionally dynamic, or topologically grown.Source: §4.8 · p. 34
  • The non-parametric GMJE branch treats M encoded instances as fixed isotropic modes and therefore requires a dynamic memory bank to estimate the marginal distribution.This marginal distribution supplies the contrastive denominator.Source: §4.8 · p. 34

The Primal-Dual Paradigm Shift

The paper contrasts dual sample-space GJE with primal feature-space GMJE, arguing that feature covariance provides stable coordinates for multi-modal semantic structure while improving computational scaling.

The Primal-Dual Paradigm Shift

  • Dual sample-space GJE evaluates the N × N Gram matrix and uses a diversity regularizer to guard against representation collapse.Its sample-similarity perspective is theoretically elegant for continuous mappings but is poorly suited to persistent semantic mixtures.Source: The Primal-Dual Paradigm Shift · pp. 34–35
  • Primal GMJE prevents dimensional collapse by maximizing latent-space volume through a feature-covariance regularizer.This targets perfectly correlated features that would otherwise flatten representations into a lower-dimensional subspace.Source: The Primal-Dual Paradigm Shift · p. 34
  • Compared with dual-space GJE, primal GMJE scales with O(d3) rather than O(N3) and directly supports K discrete prototypes for semantic routing.Dual GJE also struggles to isolate discrete multi-modal clusters and requires approximations such as random Fourier features.Source: The Primal-Dual Paradigm Shift · p. 34
  • Changing mini-batch reference samples make dual-space prototypes unstable, while Gram-matrix mixtures model similarity patterns rather than distinct semantic categories.These structural issues motivate the transition from dual GJE to primal GMJE.Source: The Primal-Dual Paradigm Shift · p. 35
  • GMJE operates in d×d feature space, where K prototypes can represent stable spatial semantic regions rather than batch-dependent pairwise similarities.The primal formulation supplies a fixed global coordinate system for multi-modal semantic modeling.Source: The Primal-Dual Paradigm Shift · p. 35

5 Experiments42

Experiments evaluate GMJE variants across synthetic and vision settings, focusing on multimodal conditional structure, memory-bank stability, representation accuracy, and unconditional latent sampling. The results show improvements in modeling multimodal structure and latent generations, while also identifying limitations under severe capacity constraints and at short prediction horizons.

5.1 Ambiguous Alignment of Synthetic Embedding Representations

  • GMJE-MDN most faithfully recovers the separated and intersecting multimodal conditional densities, while GMJE-GNG tracks topology and fixed-K mixtures remain geometrically rigid.Classic JEPA collapses toward conditional means; unimodal RBF and K = 1 models cannot resolve branches, whereas GMJE-MDN adapts local means, variances, and mixture weights.Source: §5.1 · pp. 37, 39

5.2.1 The Efficiency of SMC Memory Banks vs. FIFO

  • SMC-GMJE reaches 81.99% final linear-probing accuracy versus FIFO’s 78.08% under M = 256, a +3.91% gain with more stable late training.SMC reweights and resamples informative hard-negative particles as ESS decreases around epoch 500, whereas FIFO loses rare or informative negatives and becomes volatile.Source: §5.2.1 · pp. 40–41, Figure 8 · p. 40

5.2.2 SMC-GMJE vs. Standard Baselines

  • After 200 pre-training epochs, SMC-GMJE reaches 74.52% CIFAR-10 linear-probe accuracy, below SimCLR, BYOL, and MoCo v2 at approximately 79–82%.These are short-horizon comparisons under limited compute; SMC-GMJE trains stably but remains under-optimized relative to mature heuristic baselines.Source: §5.2.2 · pp. 41–42, Table 1 · p. 41

5.2 Representation Learning on Vision Benchmarks

  • The vision experiments therefore support dynamic particle reweighting under severe memory constraints but not superiority over established SSL methods under short CIFAR-10 pre-training.The broader benchmark indicates viability and stable improvement, while the remaining performance gap suggests a need for stronger optimization and larger-scale tuning.Source: §5.2 · pp. 39–40, §5.2.2 · p. 42

5.3 Generative GMJE: Unconditional Image Synthesis via Latent Sampling

  • The MNIST unconditional-generation experiment introduces a comparison of Parametric GMJE, unimodal GJE, and SimCLR latent spaces using frozen-encoder decoders, but this part provides no synthesis results.Each encoder is trained for 50 epochs, then its frozen embeddings are decoded by a lightweight convolutional decoder trained for 50 epochs with MSE reconstruction loss.Source: §5.3 · p. 42
  • The experiment samples latent points from learned densities, projects them onto the hypersphere, and decodes them after discarding the original image dataset.Sample quality depends on both latent geometry and whether sampled points remain within regions where the decoder learned to invert the manifold.Source: §5.3 · p. 43
  • Compared methods use a post-hoc 50-component full-covariance GMM for SimCLR, a single global Gaussian for unimodal GJE, or GMJE’s natively learned mixture distribution.These strategies distinguish externally fitted, unimodal, and natively multimodal latent density modeling.Source: §5.3 · p. 43
  • GMJE generates the sharpest, most diverse decoded samples and better aligns synthetic embeddings with semantic clusters than SimCLR and unimodal GJE.SimCLR produces blurry chimeras, while unimodal GJE exhibits mean collapse; GMJE yields the strongest qualitative unconditional generations.Source: Figure 9 · p. 43, Figure 10 · p. 44, §5.3 · p. 45
  • The results show that discriminative representations do not automatically provide good unconditional sampling distributions, whereas GMJE learns multimodal latent geometry during representation learning.Post-hoc density fitting only partially recovers structure, and a single Gaussian captures coarse global geometry; explicit multimodal modeling bridges representation learning and generation.Source: §5.3 · pp. 45–46

6 Discussion

The discussion presents GMJE as a probabilistic framework unifying contrastive and predictive SSL while addressing multimodality, collapse, and adaptive latent-density modeling.

Unifying SSL Paradigms

  • GMJE interprets standard contrastive learning as a degenerate non-parametric case and replaces age-based FIFO memory updates with SMC reweighting that emphasizes informative samples.The SMC Particle Bank downweights stale or redundant samples while emphasizing informative hard negatives.Source: §6 · p. 46
  • GMJE-MDN replaces deterministic predictors with dynamic conditional mixtures, while GMJE-GNG adapts prototype capacity using local quantization error.These variants target multimodal inverse-style alignment and avoid fixing the mixture cardinality entirely in advance.Source: §6 · p. 46

The Elimination of Asymmetric Heuristics

  • Unlike asymmetric predictive methods, GMJE uses covariance-aware geometry to promote non-degenerate full-rank representations and permit symmetric encoder gradients.The balance between Mahalanobis fitting and the log-determinant covariance penalty supplies the relevant geometric regularization.Source: §6 · p. 46

Generative vs. Discriminative Manifolds

  • GMJE equips latent representations with an explicit Gaussian-mixture density, avoiding low-density sampling gaps in contrastive spaces and over-smoothed samples from unimodal models.This connects discriminative representation learning with explicit generative modeling in latent space.Source: §6 · pp. 46–47

Dynamic Component Selection and GMJE-GNG

  • Choosing the mixture count K is a central parametric limitation: overestimation fragments semantic classes, whereas underestimation blurs distinct modes.SMC-GMJE avoids declaring K, while GMJE-GNG recruits or prunes prototypes according to evolving topology and localized error.Source: §6 · p. 47

The Gaussian Beauty

  • Gaussian and mixture-density structures provide a common probabilistic lens linking uncertainty-aware inference, MDNs, and generative representation learning.The discussion positions density estimation as a unifying perspective across these Gaussian-based methods.Source: §6 · p. 47

7 Conclusion

The conclusion presents GJE and GMJE as probabilistic alternatives to deterministic SSL prediction, with remedies for optimization failures and multimodal structure. Across synthetic and vision tasks, GMJE recovers complex conditional structure, supports competitive representations, and defines densities suitable for unconditional sampling.

7 Conclusion

  • GJE and GMJE replace deterministic latent prediction with probabilistic joint modeling designed for multimodal inverse problems and representation-collapse challenges.The framework models context-target joint densities and extends the approach from Gaussian to Gaussian-mixture structure.Source: §7 · p. 47
  • The framework combines closed-form conditional inference and uncertainty estimation with remedies for the Mahalanobis Trace Trap across parametric, adaptive, and non-parametric settings.The remedies include global prototypes, dynamic MDNs, GMJE-GNG, and SMC particle filtering.Source: §7 · pp. 47–48
  • Across synthetic tasks and high-dimensional vision benchmarks, GMJE recovered complex multimodal conditional structure and demonstrated strong discriminative performance.The reported results also describe a continuous latent density suitable for unconditional sampling.Source: §7 · p. 48
  • GMJE provides a continuous latent density for unconditional sampling while retaining discriminative utility, linking representation learning with generative modeling.The conclusion frames generative joint modeling as a promising route for accommodating multimodal real-world data.Source: §7 · p. 48

Future Work

  • Future work will extend GMJE to sequential and autoregressive settings and combine its continuous densities with discrete neuro-symbolic constraints.The named targets include video prediction, LLMs, and RiJEPA-style structured generative manifolds.Source: §7 · p. 48

8 Related Works

This section situates GJE/GMJE among related representation-learning methods and establishes Gaussian modeling and inference as its mathematical foundation. It also derives the conditional Gaussian distribution underlying the framework, highlighting how affine structure and residual independence determine conditional mean and covariance.

8 Related Works

  • GMJE provides a probabilistic alternative to contrastive and non-contrastive methods by modeling full joint covariance rather than relying on negative sampling or asymmetric collapse-prevention mechanisms.The framework is contrasted with CPC, MoCo, SimCLR [19], BYOL [30], SwAV, SimSiam [20], DINO, and VICReg.Source: §8 · pp. 48–49
  • Deterministic JEPA prediction is limited in genuinely multi-modal settings because squared-loss regression approaches conditional means that may lie in low-density regions, motivating probabilistic joint-density extensions.VJEPA, BiJEPA, and RiJEPA address probabilistic, symmetric, and neuro-symbolic directions, but a fully probabilistic joint-density view remains underdeveloped.Source: §8 · pp. 49–50
  • GMJE builds on Gaussian mixture and inference mechanisms: MDNs model conditional multi-valued densities, GNG grows topology-preserving prototypes, RFF reduces kernel costs, and SMC supports adaptive memory weighting.GMJE-MDN addresses information leakage, GMJE-GNG avoids fixed component counts and EM local minima, and SMC-GMJE prioritizes informative negatives while down-weighting stale ones.Source: §8 · p. 50

A Properties of Gaussian Distributions

  • A multivariate Gaussian is defined by an affine transformation of standard normal variables, remains closed under affine transformations, and has covariance Σ = AA⊤.These properties provide the basis for Gaussian generative modeling and closed-form inference.Source: §A.1 · p. 56

A.2 Density Function and Geometric Interpretation

  • Gaussian density geometry is governed by Mahalanobis distance and covariance eigenstructure, while joint normality requires explicit joint modeling rather than merely normal marginals.For jointly normal variables, zero cross-covariance implies strict independence.Source: §A.2 · p. 56, §A.3 · p. 56

A.4 Marginal and Conditional Distributions

  • A single global Gaussian has conditional means that vary with observations but conditional covariance that remains constant, making it unable to represent manifolds with dynamically changing variance.This homoscedasticity trap motivates mixture-based modeling for complex multi-modal structure.Source: §A.4 · p. 57

A.5 Information-Theoretic Properties

  • The Gaussian information-theoretic formulation links the log-determinant term in GJE’s objective to entropy and, when marginal volumes do not collapse, minimizing joint complexity to maximizing mutual information.The mutual-information expression depends on the determinants of the joint covariance and its marginal covariance blocks.Source: §A.5 · pp. 57–58

A.6 Log-Likelihood and Sampling

  • Gaussian likelihoods support generative sampling through Cholesky factorization, and block-matrix Schur complements provide the covariance inverses and conditional-covariance terms needed for joint Gaussian inference.The affine-transform derivation constructs a variable independent of the conditioning block because jointly Gaussian zero cross-covariance implies independence.Source: §A.6 · p. 58, §B.1 · p. 59, §C.1 · p. 60

C.1 Method 1: Derivation via Affine Transformation

  • The resulting conditional covariance is Σ_11 − Σ_12Σ_22^-1Σ_21.The covariance calculation follows by expanding the affine transformation and using the residual’s independence from x_2.Source: §C.1 · p. 60
  • The derivation rewrites x_1 as an independent residual plus the affine term Σ_12Σ_22^-1x_2.Conditioning on x_2 = a fixes the affine term while leaving the residual random.Source: §C.1 · p. 60
  • Because the residual is independent of x_2, observing x_2 = a leaves its conditional distribution equal to its marginal Gaussian distribution.Adding the conditioned affine term shifts the residual’s mean without changing its Gaussian form.Source: §C.1 · p. 60

C.2 Method 2: Derivation via Joint Distribution and Completing the Square

The paper derives Gaussian conditional inference by expanding the joint density and matching terms, then applies the result to GJE's zero-mean context-target embeddings. This yields joint and conditional objectives whose covariance-aware structure supports closed-form prediction and collapse resistance.

C.2 Method 2: Derivation via Joint Distribution and Completing the Square

  • Joint-density expansion and coefficient matching recover the conditional Gaussian parameters, agreeing exactly with the alternative derivation.The two methods produce identical conditional means and covariances.Source: §C.2 · p. 61

C.3 Application to Gaussian Joint Embeddings (GJE)

  • GJE maps the target to zt and the condition to zc, using zero-mean priors and covariance blocks to obtain closed-form conditional mean and epistemic uncertainty.Source: §C.3 · pp. 61–62

D.1 The Generative Objective (Joint NLL)

D.2 The Predictive Objective (Conditional NLL)

  • The conditional objective replaces black-box prediction with a closed-form conditional mean and a dynamic covariance that quantifies epistemic uncertainty.Source: §D.2 · p. 63
  • With fixed identity conditional covariance, the conditional NLL reduces to standard JEPA's MSE, making Dual-GJE a probabilistically grounded generalization of JEPA.Source: §D.2 · p. 63

D.3 Joint NLL: Closed-Form Inference via Block Inversion

  • Block inversion and the Schur complement provide exact Gaussian inference for new contexts, coupling the optimal linear projection to dataset-driven epistemic uncertainty.Source: §D.3 · p. 63

D.4 Combating Collapse: Why Joint Optimization is Required

  • Joint optimization combines marginal and conditional likelihoods: the marginal term guards context collapse, while covariance-dependent terms preserve dispersion and mutual-information structure.The full joint objective retains probabilistic guards that standard JEPA discards.Source: §D.4 · p. 64

E Representation Collapse in Symmetric Dual GJE

Symmetric optimization of Dual-GJE collapses because its conditional likelihood lacks output-space marginal regularization. EMA target updates and stop-gradients prevent this failure, whereas primal feature-space objectives provide symmetric regularization directly.

E Representation Collapse in Symmetric Dual GJE

E.1 The Root Cause: The Asymmetry of the Conditional Likelihood

  • The root cause is conditional asymmetry: learned targets lack the marginal regularizer that would exert an expansive force on the target space.Source: §E.1 · p. 65

E.2 The Structural Resolution: EMA and Stop-Gradients

  • EMA target encoders and stop-gradients treat target representations as fixed, diverse anchors, preventing the trace penalty from being bypassed through target collapse.Source: §E.2 · pp. 65–66
  • Primal models such as GMJE and HSIC avoid this asymmetry by evaluating feature covariance directly and regularizing both representation spaces symmetrically.Source: §E.2 · p. 66

F Training Dual-GJE: Kernel Optimization and EMA

Dual-GJE optimizes a kernelized conditional likelihood whose exact computation is cubic in batch size. Random Fourier Features replace the large Gram matrix with a low-dimensional feature representation, reducing the bottleneck to linear batch scaling up to feature dimension.

F Training Dual-GJE: Kernel Optimization and EMA

  • Dual-GJE trains context encoder and kernel parameters through conditional NLL while updating the target encoder exclusively by EMA after detaching it from gradients.Source: F Training Dual-GJE: Kernel… · p. 66
  • Exact Dual-GJE requires Gram-matrix inversion and log-determinant computation with complexity O(N^2d_c + N^2d_t + N^3), dominated by the cubic term for large batches.Source: F Training Dual-GJE: Kernel… · p. 67

F.1 Scalable Dual-GJE via Random Fourier Features (RFF)

F.1.2 Applying RFF to the Dual GJE Matrix

  • Woodbury and Weinstein–Aronszajn identities move inverse and log-determinant calculations from the N × N Gram matrix to the D × D covariance matrix.Source: §F.1.2 · pp. 67–68
  • The RFF reformulation reduces complexity from O(N^3) to O(ND^2), enabling linear scaling with batch size N.Source: §F.1.2 · p. 68

G Training Primal-GJE: Symmetric Optimization and the Entropy Fix

Primal-GJE uses symmetric feature-space covariance optimization, eliminating the need for EMA target networks or stop-gradients. Its empirical training requires an entropy-based fix for the Mahalanobis Trace Trap and scales linearly with batch size when embedding dimension is smaller.

G Training Primal-GJE: Symmetric Optimization and the Entropy Fix

H Derivations of Scalable RFF Matrix Identities

H.1 Derivation of the RFF Woodbury Identity

This section derives low-dimensional alternatives for Gaussian joint objectives and presents HSIC as a differentiable dependence-maximization objective for Primal-GJE.

H.1 Derivation of the RFF Woodbury Identity

  • O(ND^2) inversion replaces direct O(N^3) inversion by requiring only the D × D feature covariance matrix C.The RFF approximation uses Woodbury to avoid directly inverting the large N × N covariance matrix.Source: §H.1 · p. 70

H.2 Derivation of the Log-Determinant via Weinstein-Aronszajn identity

  • The N-dimensional log-determinant is evaluated exactly through the D-dimensional log-determinant of C plus a scalar correction involving dimensionality and jitter.The Weinstein-Aronszajn identity changes the determinant computation from sample space to feature space.Source: §H.2 · pp. 70–71

I Primal-GJE: Alternative Optimization via the Hilbert-Schmidt Independence Criterion (HSIC) Objective

  • HSIC provides a gradient-based alternative to joint NLL by maximizing statistical dependence between context and target representations.The objective uses HSIC as a proxy for mutual information rather than directly optimizing the joint likelihood.Source: I Primal-GJE: Alternative Optimization… · p. 71, §I.4 · p. 72

I.1 The Cross-Covariance Operator

  • HSIC measures dependence through the squared Hilbert-Schmidt norm of an RKHS cross-covariance operator and is zero exactly under independence with universal kernels on compact domains.The cross-covariance operator generalizes covariance matrices to function spaces.Source: §I.1 · p. 71, §I.2 · p. 72

I.2 Population and Empirical HSIC

  • The empirical HSIC estimator uses context and target Gram matrices with a centering matrix, converging at O(N^-1/2) with finite-sample bias bounded by O(N^-1).Its computation requires kernel evaluations and matrix multiplications without user-defined matrix-inversion regularization.Source: §I.3 · p. 72, §I.2 · p. 72

I.4 HSIC as a Differentiable Self-Supervised Objective

  • Maximizing HSIC penalizes constant representations because collapsed Gram matrices become centered to zero, while marginal log-determinant regularization can preserve spherical marginals.This provides a collapse-avoidance mechanism without necessarily using EMA target networks or stop-gradient heuristics.Source: §I.4 · pp. 72–73

J Exact Dual GJE with a Linear Kernel: The Primal-Dual Equivalence

This section proves that linear-kernel dual GJE is exactly equivalent to unimodal primal GJE, then derives marginal distributions and conditional-mixture structure for GMJE.

J Exact Dual GJE with a Linear Kernel: The Primal-Dual Equivalence

  • Linear-kernel dual GJE is algebraically identical to a single-component Gaussian mixture model operating on the primal feature covariance.The equivalence validates the K = 1 visualization baseline used in synthetic experiments.Source: J Exact Dual GJE with a Linear Kernel:… · p. 73

J.1 Setup and the Empirical Loss

  • The dual N × N data-fit and log-determinant terms can both be projected exactly into the d × d primal covariance C using push-through and determinant identities.The derivation rewrites the sample-space objective in terms of feature-space matrix operations.Source: §J.1 · p. 73, §J.2 · pp. 73–74, §J.3 · p. 74

J.4 Conclusion: The Unimodal Constraint

  • The resulting linear-kernel objective has no nonlinear flexibility and restricts learned representations to the geometry of a single global multivariate Gaussian.Its gradients are locked to the same global d × d covariance structure.Source: §J.4 · pp. 74–75

K Properties of GMM

  • GMJE models concatenated context and target embeddings with a K-component joint Gaussian mixture, yielding component-based marginal and conditional distributions for multimodal uncertainty.The mixture formulation is designed to represent multimodal structure through multiple Gaussian components.Source: K Properties of GMM · p. 75

K.1 Derivation of the Marginal Distribution of a Joint GMM

  • The target marginal of a joint GMM is another GMM whose component means and covariances are the target blocks, weighted by the original mixture coefficients.This marginal supports direct sampling in a single-view target representation space.Source: §K.1 · pp. 75–76

K.2 Derivation of the Conditional Distribution of a Joint GMM

This section derives conditional inference for joint Gaussian mixtures and examines how deterministic prediction and dynamic covariance estimation can collapse or eliminate representation-learning signals. It also presents conditional bottleneck and SMC-based extensions, highlighting collapse mechanisms arising from information leakage and imbalanced attraction and repulsion.

K.2 Derivation of the Conditional Distribution of a Joint GMM

  • A joint GMM conditions exactly to a dynamically weighted GMM, with posterior component weights γk(zc) providing probabilistic routing among closed-form Gaussian conditionals.The weights equal p(k | zc) and sum to one; each component conditional follows standard Gaussian conditioning with learnable means and covariances.Source: §K.2 · pp. 76–77

K.3 Uncertainty Quantification: The Law of Total Variance

  • GMJE represents both local noise and multimodal ambiguity, whereas classic JEPA has zero uncertainty and unimodal GJE absorbs mean-spread ambiguity into an over-smoothed covariance.Component covariances encode aleatoric noise, while separated means and mixing weights encode epistemic or multimodal uncertainty; GMJE-MDN preserves these components separately.Source: §K.3 · pp. 77–78

L MSE Minimization Yields the Conditional Mean

  • Minimizing expected squared error yields the conditional mean because irreducible conditional variance is fixed and the remaining squared distance is minimized only when g(X)=E[Y|X].The cross-term vanishes under conditional expectation, leaving an irreducible term plus a nonnegative reducible term.Source: L MSE Minimization Yields the… · p. 78, §L.1 · p. 79

L.2 Application to Classic JEPA

  • In classic JEPA, the optimal deterministic predictor is g(zc)=E[zt|zc], which can fall into empty space between valid modes when p(zt|zc) is multimodal.This follows from the same risk decomposition applied to context and target embeddings, and is observed in the first experiment’s shape task.Source: §L.2 · pp. 79–80

M.2 The Pathological Optimum (Identity Collapse)

  • Feeding the full joint embedding into an MDN parameter network enables identity collapse: the model can set every mean to the input, shrink covariances, and drive NLL toward −∞.Because the exact target is available as input, the Mahalanobis term becomes zero, the Gaussian exponent reaches exp(0)=1, and density diverges as |Σk(Z)|→0.Source: M The Identity Collapse Trap in… · p. 80, §M.2 · p. 80

M.3 Empirical Symptoms

  • Identity collapse memorizes input noise and makes dynamic mixture weights uniform, eliminating semantic branch routing; restricting the network to context zc forces coverage of valid target possibilities.The conditional information bottleneck prevents access to zt, requiring components to represent distinct branches and route probability mass according to ambiguity.Source: §M.3 · pp. 80–81, §M.4 · p. 81

N Generalizing SMC to the Full GMJE Objective

  • Full-covariance SMC-GMJE replaces contrastive similarity with the exact GMJE NLL, storing joint embeddings as particle means and learning a shared covariance for nonlinear manifold dependencies.Normalized particle weights depend only on the Mahalanobis exponent because shared Gaussian normalization constants and log-determinant terms cancel.Source: N Generalizing SMC to the Full GMJE… · p. 81, §N.1 · pp. 81–82

N.2 Algorithm: General SMC-GMJE

  • The general SMC-GMJE algorithm initializes a weighted memory bank of joint-embedding means and optimizes a log-sum-exp objective using particle weights as mixing priors and a learnable covariance.The forward pass processes incoming query batches against the memory-bank means under the current particle weights.Source: §N.2 · pp. 82–83
  • General SMC-GMJE updates encoder and covariance parameters, assigns prior weights, performs importance resampling, replaces the memory bank, and resets weights uniformly.Source: §N.2 · p. 83

O The “Mahalanobis Trace” Trap in Empirical Batch Optimization

  • Recomputing empirical covariance from the same mini-batch collapses Mahalanobis calculations into constants, catastrophically eliminating data-fit gradients.This differs from fixed population covariance, which changes as representations move and therefore provides valid encoder gradients.Source: O The “Mahalanobis Trace” Trap in… · p. 83

O.1 Case 1: The Trap in Unimodal Gaussian (Primal-GJE)

  • In Primal-GJE, substituting the batch covariance into the average Mahalanobis term yields the identity matrix and trace d, a constant intrinsic dimensionality.Source: §O.1 · p. 84

O.2 Case 2: The Trap in GMM (GMJE)

  • In GMJE, responsibility-weighted covariance substitution makes each component’s data-fit term constant, producing zero gradient toward its cluster center µ_k.Source: §O.2 · p. 84
  • The joint objective balances conditional attraction, which pulls augmented views together, against marginal repulsion, which expands geometric volume.Removing attraction or misaligning repulsion causes complete collapse or dimensional collapse, respectively; the latter yields ill-conditioned covariance and entangled features.Source: §O.2 · p. 85

O.3 The Dual Consequences of the Trace Trap

The Mahalanobis Trace Trap makes the empirical data-fit term constant, eliminating gradients and causing both dimensional collapse and loss of cross-view attraction. The section presents architectural and entropy-based remedies, followed by reproducibility details for the experiments.

O.3 The Dual Consequences of the Trace Trap

  • The empirical Mahalanobis data-fit term evaluates to the constant d, so its gradients vanish and optimization loses the intended signal.This cancellation occurs because the covariance is derived directly from the same embeddings being optimized.Source: §O.3 · p. 85
  • Dimensional collapse follows when the unopposed log-determinant term shrinks latent volume toward zero.The dead data-fit term removes the expansive barrier that would otherwise counteract this shrinkage.Source: §O.3 · p. 85
  • Entropy maximization reverses the log-determinant regularizer, producing an outward repulsive force that expands embeddings across the available hypersphere.This remedy restores dispersion but does not by itself restore matched-view attraction.Source: §O.3 · pp. 85–86
  • Parametric prototypes, context-conditioned MDNs, and SMC memory banks avoid cancellation by decoupling covariance or density components from current-batch embeddings.The three remedies use learned parameters, context-generated covariances, or historical representations respectively.Source: §O.3 · p. 86

O.4 Beyond the Trace Trap: The Log-Determinant Cheat

  • Exact NLL optimization creates a second collapse route because the network can jointly shrink representations and covariance toward zero, driving log |Σ| toward −∞.This shortcut motivates replacing rigid variance floors with continuous entropy maximization.Source: §O.4 · p. 86

P Experimental Details

  • The appendix documents shared hardware, software, seeds, data-generation settings, and model hyperparameters for synthetic, vision, and sampling experiments.Synthetic datasets use 3000 training pairs, three equally probable branches, Gaussian target noise, and a 300-point evaluation grid.Source: P Experimental Details · p. 87, §P.1 · p. 87, §P.2.1 · p. 87

P.2.3 Models and Hyperparameters

  • The synthetic experiments compare MSE JEPA, an RBF GJE baseline, Gaussian-mixture models, GMJE-GNG, and a conditional MDN under specified training protocols.The listed implementations include K = 1 and K = 3 Gaussian mixtures and a three-component conditional MDN.Source: §P.2.3 · p. 88

P.3.1 Random Seeds

  • CIFAR-10 experiments used a modified ResNet-18 backbone adapted to 32 × 32 images, standard SSL augmentations including horizontal flips and CIFAR-10 normalization, and seeds set to 42.Source: §P.3.1 · p. 89, §P.3.2 · p. 89, §P.3.3 · p. 89

P.3.4 Experiment 2a: SMC Memory Banks vs. FIFO

The experiments compare SMC-GMJE with FIFO and other SSL baselines under controlled memory-bank, training, and probing protocols. They also specify the vision architecture, latent-density extraction, and unconditional sampling procedure.

P.3.4 Experiment 2a: SMC Memory Banks vs. FIFO

  • The long-horizon comparison trains for 1000 epochs with batch size 256, memory-bank size M = 256, and EMA momentum m = 0.999.These settings define the constrained-bank experiment.Source: §P.3.4 · p. 89
  • SMC-GMJE isolates the bank-update rule by replacing FIFO replacement with likelihood-based reweighting and multinomial resampling.MoCo v2 and SMC-GMJE share the backbone and augmentations in the long-horizon comparison.Source: §P.3.4 · p. 89

P.3.5 Experiment 2b: 200-Epoch Benchmark Comparison

  • The benchmark comparison uses 200 pre-training epochs followed by 100 linear-probing epochs across SimCLR, MoCo v2, VICReg, Parametric GMJE, GMJE-MDN, and SMC-GMJE.After pre-training, encoders are frozen and a linear classifier is trained.Source: §P.3.5 · p. 90
  • The compared methods use distinct baseline-specific hyperparameters, including K = 5 dynamic GMJE-MDN components, M = 4096 for SMC-GMJE, and K = 1000 global prototypes for Parametric GMJE.SMC-GMJE also uses EMA momentum 0.999 and temperature τ = 0.1.Source: §P.3.5 · p. 90

P.4.2 Data Pipeline

  • For MNIST, images were padded from 28 × 28 to 32 × 32, with two pipelines: SSL manifold-learning and plain reconstruction, using distinct preprocessing procedures.The SSL pipeline used RandomResizedCrop with scale range [0.8, 1.0], RandomRotation up to 15◦, tensor conversion, and normalization with mean 0.5 and standard deviation 0.5. The reconstruction pipeline used deterministic padding and tensor conversion.Source: §P.4.2 · p. 91

P.4.4 Training and Sampling Protocol

  • Latent-density extraction fits a post-hoc K = 50 GMM for SimCLR, empirical moments for unimodal GJE, and learned prototypes with EMA covariance for Parametric GMJE.The extracted densities are then sampled, L2-normalized, and decoded into image space.Source: §P.4.4 · p. 92

Related papers

Loading 2603.26799v1…