让世界模型能预测多个未来轨迹,提升自动驾驶路径预测精度。
Branch-JEPA: Finite-Support Predictive Distributions for JEPA World Models

- 用有限分支替代单一预测,每个分支独立解码并保留完整集。
- 在Argoverse2上轨迹能量得分提升5.8%~6.5%,距离误差降9.3%~10.4%。
- 适合需要多可能性预测的场景,如自动驾驶、机器人规划。
联合嵌入预测架构(JEPAs)通过在表征空间中预测未来观测来学习动态。然而,多数JEPA世界模型仅输出一个潜在后继,即使隐藏意图、部分观测或随机动态使多个未来可能。本文提出Branch-JEPA,将点值转移替换为上下文加权的有限个潜在后继集合。每个分支独立解码,推理时保留完整集合。该架构支持两种互补训练方式:分离后继恢复的专化训练,以及保证分布保真度的全集能量得分训练。在Argoverse²官方验证集的五种子锁测试中,全集训练使轨迹能量得分提升5.8%–6.5%,概率加权轨迹距离降低9.3%–10.4%,相比匹配K=6分配与传输目标,同时保持5.36个去重有效分支。参数完全一致的官方验证对比显示,潜空间分叉比仅在输出解码器分叉多出10.3%有效模式,且所有五组配对种子的95%置信区间均排除零值,显著提升能量得分、期望平均位移误差(ADE)和布里尔分数。在OGBench图审计中,Branch-JEPA将可验证路线存在率从MDN的3.9%提升至19.2%。其原始支持优势在29维状态与RGB观测下依然成立。结果表明,潜空间分叉不仅能保留更多不同未来,全集评分还能提升预测分布质量。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。