[論文レビュー] VJEPA: Variational Joint Embedding Predictive Architectures as Probabilistic World Models

Yongchao Huang|arXiv (Cornell University)|Jan 20, 2026

Adversarial Robustness in Machine Learning被引用数 0

ひとこと要約

VJEPAはJEPAを拡張し、将来の潜在状態に対する確率的予測モデルを提供することで不確実性を考慮した計画を可能にし、JEPAを再構成 Observation なしでベイズフィルタリングと予測状態表現にリンクします。

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 extit{deterministic} regression objectives, which mask probabilistic semantics and limit its applicability in stochastic control. In this work, we introduce \emph{Variational JEPA (VJEPA)}, a extit{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.

研究の動機と目的

• JEPAを確率的予測状態空間モデルとして動機づけ formalize する。
• JEPA 表現がピクセル再構成なしで最適制御の十分情報状態になり得ることを示す。
• JEPAを Predictive State Representations および Bayesian filtering と統合する。
• Bayesian JEPA (BJEPA) を導入しモジュラー-priors とゼロショットタスク転送を可能にする。
• ノイズ環境における崩壊回避と不確実性を考慮した予測を示す。

提案手法

• p_phi(Z_T | Z_C, xi_T) を将来潜在状態の learned predictive distribution として導入する。
• EMA 更新を用いたターゲットエンコーダからのアマートライド推論分布 q_theta'(Z_T | x_T) を使用。
• variational objective L_VJEPA = E[-log p_phi(Z_T|Z_C,xi_T)] + beta E[KL(q_theta'(Z_T|x_T) || p(Z_T))] で訓練。
• 文脈 Z_C = f_theta(x_C) およびターゲット構造 xi_T によって JEPA 構造を維持。
• 観測モデル p_psi(x_T|Z_T) を任意で含めることができるが最適化は行わず；学習は表現空間予測に依存。
• 潜在空間での予測と不確実性伝搬を提供して計画を行う。

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

実験結果

リサーチクエスチョン

• RQ1 決定論的 JEPA が暗黙的に最適化する確率的目的は何か、そして不確実性を扱うようにどのように一般化できるか。
• RQ2 JEPA を観測再構成なしで最適制御の十分情報状態となる潜在動力学系として形式化できるか。
• RQ3 JEPA は Bayesian filtering および Predictive State Representations とどう関係し、構造的 priors を Bayesian factor で注入できるか。
• RQ4 時間的構造を導入すると自己回帰的観測尤度を課すのか、それとも JEPA は尤度を持たずに妨害要因を避けつつ機能するか。

主な発見

• VJEPA は将来の潜在状態に対する確率的予測モデルを提供し、不確実性推定や多峰的未来を可能にする。
• 観測再構成を必要とせず、JEPA を Predictive State Representations および Bayesian filtering と統合する枠組みを提供する。
• BJEPA は VJEPA を拡張し、予測信念をダイナミクスエキスパートとモジュラー priors エキスパートに分解することで制約とゼロショットタスク転送を Product of Experts を介して実現する。
• 変分目的での訓練はターゲットの多様性と非自明な条件付けの下で崩壊回避の保証を生み出す。
• 実証的なおもちゃの結果は、VJEPA と BJEPA が高分散のノイズ要素を除外し、ノイズ環境での不確実性を考慮した計画をサポートすることを示す。

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).