VJEPA: Variational Joint Embedding Predictive Architectures as Probabilistic World Models

📄 arXiv: 2601.14354v1 📥 PDF

作者: Yongchao Huang

分类: cs.LG

发布日期: 2026-01-20

备注: 77 pages

💡 一句话要点

VJEPA：概率世界模型，通过变分联合嵌入预测架构实现稳健的不确定性感知规划。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 联合嵌入预测架构 变分推断 概率世界模型 不确定性感知 机器人控制 自监督学习 贝叶斯滤波

📋 核心要点

  1. 现有JEPA方法依赖确定性回归目标，忽略了概率语义，限制了其在随机控制中的应用。
  2. VJEPA通过变分目标学习未来潜在状态的预测分布，统一了表征学习、预测状态表示和贝叶斯滤波。
  3. VJEPA和BJEPA在噪声环境中成功过滤高方差干扰因素，避免了生成模型的表征崩溃问题。

📝 摘要（中文）

联合嵌入预测架构(JEPA)通过预测潜在表示而非重建高熵观测，为自监督学习提供了一种可扩展的范式。然而，现有的公式依赖于确定性回归目标，这掩盖了概率语义并限制了其在随机控制中的适用性。本文提出了变分JEPA (VJEPA)，这是一种概率泛化，它通过变分目标学习未来潜在状态的预测分布。我们证明了VJEPA统一了表征学习与预测状态表示(PSR)和贝叶斯滤波，确立了序列建模不需要自回归观测似然。理论上，我们证明了VJEPA表示可以作为最优控制的充分信息状态，而无需像素重建，同时为避免崩溃提供正式保证。我们进一步提出了贝叶斯JEPA (BJEPA)，它将预测信念分解为学习的动力学专家和模块化先验专家，从而通过专家乘积实现零样本任务迁移和约束（例如，目标、物理）满足。通过一个噪声环境实验，我们证明了VJEPA和BJEPA成功地过滤掉了导致生成基线中表征崩溃的高方差干扰因素。通过实现有原则的不确定性估计（例如，通过采样构建可信区间）同时保持关于观测的无似然性，VJEPA为高维噪声环境中可扩展、稳健、不确定性感知规划提供了一个基础框架。

🔬 方法详解

问题定义：现有的联合嵌入预测架构(JEPA)主要采用确定性的回归目标进行自监督学习，这种方式忽略了环境中的不确定性，限制了其在需要处理随机性的任务，如机器人控制中的应用。此外，确定性目标也可能导致模型对噪声敏感，产生不稳定的表征。

核心思路：VJEPA的核心思路是将JEPA从确定性预测扩展到概率预测。通过学习未来潜在状态的预测分布，VJEPA能够更好地捕捉环境的不确定性，并提供更鲁棒的表征。这种概率建模方法允许模型进行不确定性估计，例如通过采样构建可信区间，从而支持更稳健的决策。

技术框架：VJEPA的整体框架包括一个编码器，用于将观测映射到潜在状态；一个预测模型，用于预测未来潜在状态的分布；以及一个变分目标函数，用于训练模型。BJEPA在此基础上，将预测信念分解为动力学专家和先验专家，通过专家乘积进行推理。这种模块化设计允许模型利用先验知识，并实现零样本任务迁移。

关键创新：VJEPA的关键创新在于其概率建模方法，它将JEPA从确定性回归扩展到变分推断。这种方法允许模型学习未来潜在状态的预测分布，从而更好地捕捉环境的不确定性。此外，BJEPA的专家乘积设计允许模型利用先验知识，并实现零样本任务迁移。

关键设计：VJEPA的关键设计包括变分目标函数，它鼓励模型学习准确的预测分布；以及BJEPA的专家乘积设计，它允许模型利用先验知识。具体的损失函数可能包括KL散度，用于衡量预测分布与真实分布之间的差异。网络结构可能采用Transformer或RNN等序列模型，用于捕捉时间依赖性。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

实验结果表明，VJEPA和BJEPA在噪声环境中能够有效过滤高方差干扰因素，避免生成模型的表征崩溃问题。通过不确定性估计，VJEPA能够提供更稳健的规划能力。具体性能数据未知，但论文强调了VJEPA在处理噪声方面的优势。

🎯 应用场景

VJEPA适用于高维、噪声环境下的机器人控制、自动驾驶等领域。通过提供稳健的不确定性估计，VJEPA可以帮助智能体更好地应对环境中的不确定性，并做出更可靠的决策。此外，BJEPA的专家乘积设计允许模型利用先验知识，从而加速学习过程，并实现零样本任务迁移。未来，VJEPA有望应用于更复杂的现实世界场景。

📄 摘要（原文）

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.