VJEPA: Variational Joint Embedding Predictive Architectures as Probabilistic World Models

2026-01-20

AI Summary

This work addresses the limitation of existing Joint Embedding Predictive Architectures (JEPAs), which rely on deterministic targets and thus struggle to model environmental stochasticity, hindering their applicability in uncertainty-aware control. To overcome this, we propose the Variational Joint Embedding Predictive Architecture (VJEPA), the first probabilistic extension of JEPA, which leverages variational inference to learn predictive distributions over future latent states. This framework unifies representation learning, predictive state modeling, and Bayesian filtering—enabling sequential inference and optimal control without pixel-level reconstruction. We further introduce Bayesian JEPA (BJEPA), incorporating a product-of-experts mechanism to yield informative, non-collapsed state representations that support zero-shot task transfer and constraint satisfaction. Experiments demonstrate that VJEPA and BJEPA effectively filter noise, construct reliable confidence intervals, and enable robust, uncertainty-aware planning in high-variance disturbance environments.

Technology Category

Reasoning under Uncertainty: Relational Probabilistic ModelsHumans and AI: Human-Aware Planning and Behavior PredictionIntelligent Robots: State Estimation

Application Category

Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphsSystems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applicationsResponsible Web: Machine-in-the-loop, human agency and autonomy

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.

Problem

• Joint Embedding Predictive Architectures
• probabilistic world models
• representation collapse
• stochastic control
• uncertainty estimation

Innovation

• Variational JEPA
• Probabilistic World Models
• Predictive State Representations
• Bayesian Filtering
• Representation Collapse Avoidance