论文
Branch-JEPA:JEPA 世界模型的有限支持预测分布
Branch-JEPA: Finite-Support Predictive Distributions for JEPA World Models
摘要
联合嵌入预测架构(JEPA)通过预测表示空间中的未来观察来学习动态。然而,大多数 JEPA 世界模型都会返回一个潜在的后继者,即使隐藏的意图、部分观察或随机动态使多个未来变得合理。我们引入了 Branch-JEPA,它将点值转换替换为上下文加权的有限潜在后继集。每个分支都是独立解码的,并且在推理时保留完整的集合。该架构支持两种互补的训练机制:用于恢复分离后继者的专业化和用于分布保真度的全套能量评分训练。在 Argoverse~2 官方验证分割的锁定五种子评估中,全套训练将轨迹能量分数提高了 $5.8$--$6.5\%$,概率加权轨迹距离比匹配的 $K{=}6$ 分配和传输目标提高了 $9.3$--$10.4\%$,同时保留了 $5.36$ 端点重复数据删除的有效分支。在参数精确的官方验证比较中,与仅在输出解码器处的分支相比,潜在分支保留了 10.3%\%$ 的更有效模式,并提高了所有五对种子中的能量得分、预期 ADE 和 Brier;每个配对的 95\% 区间不包括零。在 OGBench 图形审计中,Branch-JEPA 将传送验证路由的存在增加到 $19.2\%$,而 MDN 的 $3.9\%$。其原始支持优势在 29 维状态和 RGB 观察中也持续存在。总之,潜在分支保留了更独特的未来,而全套评分提高了所得预测分布的质量。