# Branch-JEPA: Finite-Support Predictive Distributions for JEPA World Models

arXiv: https://arxiv.org/abs/2607.05238 · HTML: https://mlsanity.com/paper/2607.05238/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Latent predictive world models, one of 11 directions in World models.
- Builds on 3 of its 44 references that are in the map.
- First posted on arXiv 6 Jul 2026.

## Placed in

- World models / [Latent predictive world models](https://mlsanity.com/wm/latent-predictive-world-models/)

## Builds on (most cited first, 3 of 3)

- [V-JEPA 2: Self-Supervised Video Models Enable Understanding, Prediction and Planning](https://mlsanity.com/paper/2506.09985/) (2025)
- [UWM-JEPA: Predictive World Models That Imagine in Belief Space](https://mlsanity.com/paper/2605.25313/) (2026)
- [JEDI: Joint Embedding Diffusion World Model for Online Model-Based Reinforcement Learning](https://mlsanity.com/paper/2605.13013/) (2026)

## Abstract (arXiv)

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.