MoP-JEPA: Hard-Assigned Predictor Mixtures for Stochastic JEPA World Models
2026-07-06 • Artificial Intelligence
Artificial Intelligence
AI summaryⓘ
The authors show that a common way to predict future states in a model called JEPA doesn't work well in random (stochastic) environments. Instead of predicting a clear next state, it predicts an average that doesn't actually exist, causing problems. They prove that using multiple predictors, each responsible for a specific possible next state, works better and makes planning more accurate. Their method outperforms previous ones in tests and can even be verified to avoid false predictions. This work highlights that handling multiple possible futures is key for planning in uncertain situations.
JEPAworld modellatent statestochastic environmentmixture of predictorslatent regressionplanningmode collapsequantizerverification protocol
Authors
Zhi Song, Ximing Xing, Zhenchao Tang, hanbo Huang, Tianxu Lv, minghao Yang, Zhongzheng Niu, He Bing, Lusheng Wang, Jianhua Yao
Abstract
JEPA world models predict the next latent state with a single deterministic predictor trained by latent regression. We show that this fails structurally when the environment is stochastic: at a branching transition, the regression-optimal predictor outputs the conditional mean of the successor embeddings, a point between the true next states that corresponds to no state at all. We prove this collapse for deterministic and gated mixture-of-experts predictors, and prove that MoP-JEPA's hard-assigned predictors converge instead to a quantizer of the transition distribution: one head per successor mode, enumerable in a single forward pass, which is the interface a planner consumes. On official OGBench offline data with leak-free evaluation, planning over single-predictor rollouts performs poorly ($0.02$--$0.09$ success) while planning over our predicted modes reaches up to $0.85$, ahead of deterministic, gated-MoE, and variational predictors on every task. Because multi-prediction evaluation invites coverage freeloading, a verification protocol is part of the method: an input-agnostic codebook control, a shuffled-context test, router-gated readouts, transition-precision guards, and a verified-route criterion in which the model proposes its transition graph blind and ground truth is used only to check the result. Under this criterion our method outperforms the strongest soft alternative on all three mazes ($2$--$5\times$), and the protocol identifies the remaining gap in that baseline's raw scores as routes through predicted transitions that do not exist. The same model executes in the real environment, placing second of seven against the published OGBench baselines on the hardest maze. Multimodal dynamics decide whether a JEPA world model can plan at all; a mixture of predictors with hard assignment is a minimal and verifiable fix.