MoP-JEPA: Hard-Assigned Predictor Mixtures for Stochastic JEPA World Models
作者: Zhi Song, Ximing Xing, Zhenchao Tang, hanbo Huang, Tianxu Lv, minghao Yang, Zhongzheng Niu, He Bing, Lusheng Wang, Jianhua Yao
分类: cs.AI
发布日期: 2026-07-06
💡 一句话要点
提出MoP-JEPA以解决随机环境下预测失效问题
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)
关键词: JEPA世界模型 随机环境 硬分配预测器 混合预测 规划算法 机器学习 动态决策
📋 核心要点
- 现有的JEPA世界模型在随机环境中使用单一确定性预测器时,无法有效预测下一个潜在状态,导致规划失败。
- 论文提出MoP-JEPA,通过硬分配的混合预测器来解决这一问题,使每个后继模式都有独立的预测头,从而更准确地捕捉转移分布。
- 在OGBench数据集上,基于单一预测器的规划成功率仅为0.02至0.09,而使用MoP-JEPA的规划成功率高达0.85,显示出显著的性能提升。
📝 摘要(中文)
JEPA世界模型通过单一确定性预测器进行潜在状态的预测,但在随机环境中,这种方法在结构上存在缺陷。论文证明了在分支转移时,回归最优预测器输出的条件均值并不对应于真实的下一个状态。为此,MoP-JEPA采用硬分配的预测器,能够收敛到转移分布的量化器,使得每个后继模式对应一个预测头。实验证明,基于我们预测模式的规划在OGBench离线数据集上表现优异,成功率高达0.85,显著优于现有的确定性和变分预测器。
🔬 方法详解
问题定义:论文要解决的问题是现有JEPA世界模型在随机环境中使用单一确定性预测器时的预测失效,特别是在分支转移情况下,回归最优预测器输出的条件均值无法对应真实状态。
核心思路:论文的核心思路是采用硬分配的混合预测器,使每个后继模式对应一个独立的预测头,从而有效捕捉转移分布,避免了单一预测器的局限性。
技术框架:整体架构包括多个硬分配的预测器,每个预测器负责一个后继模式。模型通过输入无关的代码本控制、洗牌上下文测试、路由门控读取等模块来确保预测的准确性和可靠性。
关键创新:最重要的技术创新点在于引入了硬分配的混合预测器,这一设计使得模型能够在随机环境中有效地进行规划,克服了传统方法的局限性。
关键设计:关键设计包括输入无关的代码本控制、路由门控读取、转移精度保护等技术细节,确保模型在预测时的准确性和有效性。
🖼️ 关键图片
📊 实验亮点
实验结果显示,基于MoP-JEPA的规划在OGBench离线数据集上成功率高达0.85,显著优于传统的确定性预测器和变分预测器,后者的成功率仅为0.02至0.09,提升幅度达到数倍。
🎯 应用场景
该研究的潜在应用领域包括机器人导航、自动驾驶和智能决策系统等,能够在复杂和动态的环境中进行有效的规划与决策。未来,该方法有望推动更高效的智能体设计,提升其在不确定环境中的表现。
📄 摘要(原文)
JEPA world models predict the next latent state with a single deterministic predictor trained by latent regression. We show that this fails structurally when the environment is stochastic: at a branching transition, the regression-optimal predictor outputs the conditional mean of the successor embeddings, a point between the true next states that corresponds to no state at all. We prove this collapse for deterministic and gated mixture-of-experts predictors, and prove that MoP-JEPA's hard-assigned predictors converge instead to a quantizer of the transition distribution: one head per successor mode, enumerable in a single forward pass, which is the interface a planner consumes. On official OGBench offline data with leak-free evaluation, planning over single-predictor rollouts performs poorly ($0.02$--$0.09$ success) while planning over our predicted modes reaches up to $0.85$, ahead of deterministic, gated-MoE, and variational predictors on every task. Because multi-prediction evaluation invites coverage freeloading, a verification protocol is part of the method: an input-agnostic codebook control, a shuffled-context test, router-gated readouts, transition-precision guards, and a verified-route criterion in which the model proposes its transition graph blind and ground truth is used only to check the result. Under this criterion our method outperforms the strongest soft alternative on all three mazes ($2$--$5\times$), and the protocol identifies the remaining gap in that baseline's raw scores as routes through predicted transitions that do not exist. The same model executes in the real environment, placing second of seven against the published OGBench baselines on the hardest maze. Multimodal dynamics decide whether a JEPA world model can plan at all; a mixture of predictors with hard assignment is a minimal and verifiable fix.