VJEPA: Variational Joint Embedding Predictive Architectures as Probabilistic World Models

TL;DR

This work addresses the limitations of deterministic JEPA methods by introducing Variational JEPA (VJEPA), a probabilistic framework that learns a predictive distribution over future latent states without reconstructing observations. It unifies JEPA with Predictive State Representations and Bayesian filtering, enabling belief propagation and distributional planning in latent space, while preserving a reconstruction-free paradigm. The authors extend VJEPA with Bayesian JEPA (BJEPA), which factors predictive belief via a Product of Experts to incorporate structural priors and enable zero-shot task transfer and constraint satisfaction. Theoretical results establish predictive sufficiency for control and formal collapse avoidance, and experiments on a Noisy TV-like linear system demonstrate robust nuisance filtering and uncertainty-aware planning, suggesting principled uncertainty estimation in high-dimensional, noisy environments without autoregressive observation likelihoods.

Abstract

Joint Embedding Predictive Architectures (JEPA) offer a scalable paradigm for self-supervised learning by predicting latent representations rather than reconstructing high-entropy observations. However, existing formulations rely on \textit{deterministic} regression objectives, which mask probabilistic semantics and limit its applicability in stochastic control. In this work, we introduce \emph{Variational JEPA (VJEPA)}, a \textit{probabilistic} generalization that learns a predictive distribution over future latent states via a variational objective. We show that VJEPA unifies representation learning with Predictive State Representations (PSRs) and Bayesian filtering, establishing that sequential modeling does not require autoregressive observation likelihoods. Theoretically, we prove that VJEPA representations can serve as sufficient information states for optimal control without pixel reconstruction, while providing formal guarantees for collapse avoidance. We further propose \emph{Bayesian JEPA (BJEPA)}, an extension that factorizes the predictive belief into a learned dynamics expert and a modular prior expert, enabling zero-shot task transfer and constraint (e.g. goal, physics) satisfaction via a Product of Experts. Empirically, through a noisy environment experiment, we demonstrate that VJEPA and BJEPA successfully filter out high-variance nuisance distractors that cause representation collapse in generative baselines. By enabling principled uncertainty estimation (e.g. constructing credible intervals via sampling) while remaining likelihood-free regarding observations, VJEPA provides a foundational framework for scalable, robust, uncertainty-aware planning in high-dimensional, noisy environments.

Figures (5)

• Figure 1: Variational JEPA (VJEPA) architecture. The context $x_C$ is encoded into $Z_C$, conditioning a probabilistic predictor $p_\phi(Z_T \mid Z_C,\xi_T)$. The target $x_T$ is encoded into a distribution $q_{\theta'}$ (typically via EMA), from which a sample $Z_T$ is drawn. Training minimizes the sum of a negative log-likelihood term (red box) and a KL regularization term (blue box). $f_{\theta'}$ can output sufficient statistics (e.g. mean and covariance for Gaussians) for parameterizing $q_{\theta'}$.
• Figure 2: Bayesian JEPA (BJEPA) architecture. A predictive likelihood models dynamics while a prior encoder maps auxiliary input $\eta$ to a latent constraint. These distributions are fused via a PoE operator to yield a posterior predictive distribution. During training, the prediction is matched against the target encoder output using a JEPA loss.
• Figure 3: Performance metrics across noise scales.Top Row: Training set $R^2$. Bottom Row: Test set $R^2$ (Generalization). The generative models (VAE, AR) degrade linearly as noise increases, tracking the distractor (Bottom Right). The JEPA-based models (Blue/Cyan/Purple) maintain high signal recovery (Bottom Left) even at high noise scales, demonstrating invariance to nuisance variability.
• Figure 4: Latent Reconstructions at varying noise scales. At $\sigma=8.0$ (Right), the VAE and AR reconstructions (dashed lines) track the high-frequency noise. In contrast, BJEPA and VJEPA (solid lines) successfully filter the noise and track the underlying true signal (black line).
• Figure 5: Schematic overview of a Particle Filter (Sequential Monte Carlo) algorithm applied within the VJEPA framework. The process begins with a set of weighted particles representing the current belief. The Prediction step uses the learned probabilistic dynamics model $p_\phi$ as a proposal distribution to propagate particles forward. The Update step re-weights these particles based on new observation information, which is incorporated via a surrogate likelihood derived from the target encoder $q_{\theta'}$ (or an optional decoder $p_\psi$). Finally, Resampling is performed to avoid particle degeneracy, resulting in an updated particle set representing the posterior belief for the next time step.

Theorems & Definitions (19)

• Definition 1: Collapse
• Theorem 1: No Collapsed Global Optimum under Target Diversity
• Remark 1
• Definition 2: Control-Relevant Predictive Sufficiency
• Lemma 1: Latent sufficiency implies cost sufficiency
• Theorem 2: Sufficiency for Optimal Control from a Predictive Information State
• Theorem 3: Optimality of MAP (Mean) Control under Quadratic Costs
• proof
• proof
• Theorem 4: Variational Mutual Information Lower Bound