MoP-JEPA: Hard-Assigned Predictor Mixtures for Stoc... | AI Research

MoP-JEPA: Hard-Assigned Predictor Mixtures for Stochastic JEPA World Models
JEPA (Joint-Embedding Predictive Architecture) world models are designed to help AI agents plan by predicting future states. However, these models typically use a single "predictor" to guess what happens next. This paper identifies a fundamental flaw in this approach: when an environment is stochastic—meaning one action could lead to several different, valid outcomes—a single predictor tries to calculate the "average" of all those possibilities. This average often results in a state that doesn't actually exist, making it impossible for the agent to plan effectively. The authors introduce MoP-JEPA, a method that replaces the single predictor with a mixture of "hard-assigned" heads, allowing the model to enumerate all possible future states rather than averaging them.

The Problem: Conditional-Mean Collapse

In standard JEPA models, the predictor is trained to minimize the distance between its output and the actual future state. When faced with multiple potential futures, the math forces the model to output the "conditional mean"—a point that sits in the middle of all possible outcomes. The authors prove that this "collapse" is a structural failure of both standard deterministic predictors and gated mixture-of-experts models. Because the model outputs a single, invalid point, any planning algorithm relying on this prediction will fail to navigate the environment correctly.

The Solution: Hard-Assigned Predictor Mixtures

To fix this, the authors propose MoP-JEPA. Instead of one predictor, the model uses a set of $K$ independent predictor heads and a "router." During training, the model uses a "hard-assignment" objective: each observed transition is assigned to the single closest predictor head. This forces each head to specialize in a specific type of outcome. The router then learns to predict which head should be active based on the current context. This creates an enumerable set of potential futures that a planner can actually use, as each head represents a distinct, reachable mode of the environment.

Verification and Results

Because simply having multiple predictions can sometimes lead to "coverage freeloading"—where a model covers many outcomes without being accurate—the authors implemented a rigorous verification protocol. They tested their method on official OGBench offline data, including complex maze tasks. While standard single-predictor models achieved very low success rates (0.02–0.09), MoP-JEPA reached up to 0.85. The authors also used a "realroute" criterion, which checks if the model’s proposed plan consists of transitions that actually exist in the environment. Under this strict audit, MoP-JEPA significantly outperformed other approaches, demonstrating that its predicted modes are not just guesses, but usable paths for planning.

Broader Applicability

The authors demonstrate that this mechanism is not limited to simple 2D mazes. By porting the MoP-JEPA head into existing frameworks like DINO-WM, they showed that the method works as a drop-in replacement for image-based tasks and more complex environments. The heads consistently specialize to match the branching factor of the environment, proving that the "collapse" is a general issue in stochastic world modeling and that hard-assigned mixtures provide a minimal, verifiable, and effective fix.