MoP-JEPA: Hard-Assigned Predictor Mixtures for Stochastic JEPA World Models

2026-07-06

AI Summary

This work addresses the limitation of conventional JEPA world models, which employ a single deterministic predictor and consequently produce meaningless mean embeddings at stochastic state branches, hindering effective planning. To overcome this, the authors propose MoP-JEPA, a novel approach that introduces a hard-assignment mixture-of-predictors architecture to explicitly quantize stochastic transition distributions into multiple modes, with each prediction head corresponding to a plausible successor state. The study provides the first theoretical proof that this hard-assignment mechanism converges to a quantized representation of the true transition distribution. Furthermore, it devises validation protocols—including input-agnostic codebook control and shuffled-context testing—to verify the authenticity of multimodal predictions. Evaluated on OGBench offline data, MoP-JEPA achieves a planning success rate of 0.85, substantially outperforming existing methods, and ranks second in real-world trials on the most challenging maze, with validation protocols exposing spurious path generation in competing approaches.

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.

Problem

• stochastic environments
• JEPA world models
• multimodal dynamics
• latent state prediction
• branching transitions

Innovation

• MoP-JEPA
• hard-assigned mixture
• stochastic world models
• multimodal prediction
• verifiable planning