🤖 AI

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

MoP-JEPA addresses the limitations of single-output predictors in stochastic JEPA world models by employing a context-only router with $K$ hard-assigned heads to generate a finite set of candidate successors, significantly outperforming standard regression and MDN approaches in both raw coverage and verified-route success on graph search tasks.

Original authors: Zhi Song, Ximing Xing, Zhenchao Tang, hanbo Huang, Weilong Yan, Tianxu Lv, Minghao Yang, Zhongzheng Niu, Bing He, Lusheng Wang, Jianhua Yao

Published 2026-07-14

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Zhi Song, Ximing Xing, Zhenchao Tang, hanbo Huang, Weilong Yan, Tianxu Lv, Minghao Yang, Zhongzheng Niu, Bing He, Lusheng Wang, Jianhua Yao

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to navigate a giant, twisting maze. The robot has a "crystal ball" (a world model) that tries to guess what happens next after it takes a step.

The Problem: The "Average" Trap
Most current robots use a crystal ball that works like a very polite, indecisive librarian. If you ask, "If I go left here, where will I end up?" and there are two possible answers (maybe a teleporter sends you to the kitchen, or maybe a trap sends you to the garden), the librarian doesn't give you two options. Instead, it gives you one answer: the exact middle point between the kitchen and the garden.

In the real world, there is no "middle" between a kitchen and a garden. If the robot tries to walk to that middle point, it ends up crashing into a wall that doesn't exist in either reality. The paper calls this the "single-output restriction." It argues that when the future is messy and has multiple possibilities (stochastic), trying to predict just one "average" future is a recipe for failure.

The Solution: The "Hard-Assigned" Team
The authors, Zhi Song and his team at Tencent and City University of Hong Kong, built a new kind of crystal ball called MoP-JEPA. Instead of one librarian, they hired a team of K different experts (heads).

Here is how it works:

The Router: A smart manager looks at the current situation and decides which experts are likely active. Crucially, this manager only sees the current context; it cannot peek at the future to know which outcome actually happened.
The Experts: Each expert is trained to be very good at predicting one specific type of future.
The Candidate Set: When a new situation happens, the system doesn't blend the experts' opinions into one answer. Instead, it outputs a list of distinct possibilities from the active experts.

This creates a list of distinct, valid possibilities (a "candidate set") rather than a blurry, fake average. It's like asking a travel agent for a list of possible destinations rather than a map to a place that is halfway between Paris and Tokyo.

The Evidence: Does it Actually Work?
The team tested this on some tricky mazes (from a dataset called OGBench) where the robot could teleport or get stitched into different paths.

The Old Way: When the standard robot tried to plan a path using its "average" prediction, it succeeded in only 0.02 to 0.09 (2% to 9%) of the queries. It was mostly guessing wrong.
The New Way: The MoP-JEPA robot, using its list of distinct possibilities, succeeded in 0.85 (85%) of the queries.

But the authors were careful. They knew that just having a long list of guesses isn't enough; you could just guess everything and get lucky. So, they added a strict test called "realroute." This checks if the robot's list of guesses actually contains a path made of real transitions that can be walked through.

The Result: MoP-JEPA passed this strict test on all three mazes.
The Comparison: Another method called a "Mixture Density Network" (MDN) could guess a lot of things (high coverage), but many of those guesses were fake edges that didn't exist in the real maze. MoP-JEPA's guesses were useful and real.

What This Means (and What It Doesn't)
The paper proves that for robots navigating uncertain worlds, you need a system that can say, "It could be A, or it could be B," rather than "It will be somewhere between A and B."

What it rules out: The paper explicitly argues against using "gated" mixtures (like M3-JEPA) that blend experts into a single output. Even if you have many experts inside, if they mash their answers together into one vector, you still end up with the "average trap."
What it suggests: The authors suggest that this "hard-assigned" approach is a better way to handle the messy, multi-path nature of real-world planning.
The Confidence: The results are measured on specific, held-out datasets (OGBench). The paper shows that in these simulations, the new method is vastly superior to the old "average" predictors. It doesn't claim to have solved all robot planning forever, but it shows a clear, measured win in these specific maze scenarios.

In short: If you want a robot to plan for a future that has many possibilities, stop asking it for the average. Give it a team of specialists, let them generate a list of real options, and let the robot choose the best path from that list.

Technical Summary: MoP-JEPA

Problem Statement
JEPA (Joint Embedding Predictive Architecture) world models typically predict the next latent state using a single regressor. While effective in deterministic environments, this approach fails under stochastic transitions where a single context (state-action pair) can lead to multiple valid successors. Standard regression objectives (squared loss or cosine loss) force the model to predict the conditional mean or the normalized direction of that mean. In multimodal distributions, this "compromise" output often corresponds to a nonexistent state (e.g., an invalid midpoint between two distinct destinations), rendering the predictor useless for planning tasks that require enumerating valid future paths. Existing attempts to introduce stochasticity, such as gated Mixture-of-Experts (MoE) fused by a weighted sum or single-Gaussian variational distributions, still collapse to a single output vector per context, retaining the same fundamental limitation.

Methodology
The paper proposes MoP-JEPA (Mixture-of-Predictors JEPA), which replaces the single-output predictor with a structure capable of producing a finite set of candidate successors in a single forward pass.

Architecture: The model retains the standard JEPA backbone (encoder $f_\theta$ and EMA target encoder $f_\xi$ ) but replaces the single predictor $g_\phi$ with $K$ parallel predictor heads $\{g_{\phi_k}\}_{k=1}^K$ .
Router: A context-only router $r_\psi$ takes the context $c = [z, a]$ (latent state and action) and outputs a probability distribution $\pi$ over the $K$ heads. Crucially, the router does not see the target future state $z^+$ , ensuring it learns to predict valid successors based solely on the current context.
Hard Assignment Objective: Training utilizes a "Winner-Take-All" (WTA) mechanism inspired by conditional vector quantization. For each target $z^+_i$ $z_{i}^{+}$ , the nearest head $k^*_i$ $k_{i}^{*}$ is identified. The loss function consists of three terms:
1. Prediction Loss: The distance between the winning head's output and the target.
2. Router Loss: Cross-entropy to train the router to predict the index of the winning head.
3. Load Balancing: A penalty to prevent the router from collapsing to a single head, encouraging the utilization of all $K$ heads.
Inference: At test time, the model outputs a set of $K$ candidate latents $\{u_k\}$ with associated router probabilities $\{\pi_k\}$ . This finite set serves as the transition graph for downstream planners (e.g., BFS or MPC) without requiring a decoder or multiple rollouts.

Key Contributions

Finite Successor Interface for JEPA: The paper introduces a mechanism to convert the single-output restriction of JEPA into a finite candidate set, allowing the model to represent multimodal transitions explicitly.
Theoretical Characterization of Collapse: The authors formally demonstrate (via Propositions 1–3) that deterministic predictors and fused MoEs inherently optimize for the conditional mean, which is often an invalid state in stochastic environments. They identify hard-assigned heads as a form of conditional vector quantization that can cover distinct modes.
Verified-Route Evaluation Protocol: Recognizing that high coverage can be achieved by indiscriminate prediction (e.g., predicting all possible edges regardless of context), the paper introduces a "realroute" metric. This metric evaluates whether the proposed candidate set contains a valid path of real transitions from start to goal, distinguishing context-dependent learning from mere coverage.

Results
The authors evaluate MoP-JEPA on OGBench datasets (including PointMaze and AntMaze variants with teleportation and stitching, which induce multimodal transitions).

Planning Performance: On held-out OGBench transitions, graph search using single-output predictors succeeds on only 0.02–0.09 of queries. In contrast, MoP-JEPA achieves a success rate of 0.85.
Mode Recovery: Visualizations in latent space (Figures 3 and 4) show that dense predictors collapse to a single invalid midpoint between branches, whereas MoP-JEPA successfully recovers distinct peaks corresponding to each valid branch.
Verified-Route Success: MoP-JEPA leads on the "realroute" metric across all three maze environments. While a Mixture Density Network (MDN) achieves high raw coverage, it predicts many nonexistent edges, failing the verified-route test. MoP-JEPA maintains high coverage while ensuring the predicted edges correspond to actual transitions.
Ablation: The router gating mechanism is shown to be essential; without it, the model may retain heads that do not correspond to valid context-dependent transitions.

Significance and Claims
The paper claims that the primary limitation of current JEPA world models in stochastic settings is the single-output restriction, which forces a compromise that planning algorithms cannot utilize. By introducing hard-assigned heads, MoP-JEPA allows the model to "enumerate" discrete successor modes rather than averaging them. The authors emphasize that their contribution is not a new encoder or anti-collapse machinery, but a specific predictor-level mechanism that enables planning-level evidence of success. The work provides an empirical account of how single-output prediction degrades planning under stochasticity and offers a protocol to distinguish between useful, context-dependent successors and indiscriminate coverage.

Technical Summary: MoP-JEPA

More like this