Branch-JEPA: Finite-Support Predictive Distributions for JEPA World Models
TL;DR
Branch-JEPA is introduced, which replaces this point-valued transition with a context-weighted finite set of latent successors, and preserves more distinct futures, while full-set scoring improves the quality of the resulting predictive distribution.
Abstract
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.