发表机构
City University of Hong Kong, China; Tencent, China（香港城市大学; 腾讯）
机构由 AI 辅助整理，请以论文原文为准。AI 中文总结
针对传统JEPA在随机环境下的状态坍缩问题，提出硬分配预测器混合的MoP-JEPA，可收敛到转移分布量化器，在OGBench测试中规划性能远超基线方案。
AI 中文摘要
JEPA世界模型通过单个确定性预测器、以隐回归训练来预测下一个隐状态。本文表明，当环境具有随机性时，该范式存在结构性缺陷：在分支转移处，回归最优预测器输出后继嵌入的条件均值，该点处于多个真实下一状态之间，完全不对应任何实际状态。本文针对确定性预测器与门控混合专家预测器证明了这种坍缩现象，并证明MoP-JEPA的硬分配预测器会收敛为转移分布的量化器：每个头对应一个后继模式，可通过单次前向传播枚举，为规划器提供可用接口。在无泄漏评估的官方OGBench离线数据集上，基于单预测器推演的规划性能极差（成功率0.02~0.09），而基于本文预测模式的规划成功率最高可达0.85，在所有任务上均优于确定性、门控MoE与变分预测器。由于多预测评估易出现覆盖率搭便车问题，本方法配套了验证协议：包含输入无关码本控制、打乱上下文测试、路由门控读出、转移精度防护，以及验证路由准则——模型盲生成转移图，仅用真值校验结果。在该准则下，本文方法在全部三个迷宫任务上性能是最优软分配基线的2~5倍，该协议还识别出基线原始得分的剩余虚高来自不存在的预测转移路径。同一模型可在真实环境中运行，在最难迷宫任务的7个公开OGBench基线中排名第二。多模态动力学决定了JEPA世界模型是否具备规划能力；采用硬分配的预测器混合是一种轻量且可验证的解决方案。
英文摘要
Joint-embedding predictive architectures (JEPAs) learn dynamics by predicting future observations in representation space. Yet most JEPA world models return one latent successor, even when hidden intent, partial observation, or stochastic dynamics make several futures plausible. We introduce Branch-JEPA, which replaces this point-valued transition with a context-weighted finite set of latent successors. Every branch is decoded independently, and the complete set is retained at inference. The architecture supports two complementary training regimes: specialization for recovering separated successors and full-set Energy-Score training for distributional fidelity. In a locked five-seed evaluation on the Argoverse~2 official validation split, full-set training improves trajectory Energy Score by $5.8$--$6.5\%$ and probability-weighted trajectory distance by $9.3$--$10.4\%$ over matched-$K{=}6$ assignment and transport objectives, while retaining $5.36$ endpoint-deduplicated effective branches. In a parameter-exact official-validation comparison, latent branching retains $10.3\%$ more effective modes and improves Energy Score, expected ADE, and Brier in all five paired seeds over branching only at the output decoder; every paired 95\% interval excludes zero. In an OGBench graph audit, Branch-JEPA increases teleport verified-route existence to $19.2\%$ versus $3.9\%$ for the MDN. Its raw-support advantage also persists with 29-D state and RGB observations. Together, latent branching preserves more distinct futures, while full-set scoring improves the quality of the resulting predictive distribution.
