MoP-JEPA: Hard-Assigned Predictor Mixtures for Stochastic JEPA World Models Zhi Song1,2 , Ximing Xing2 , Zhenchao Tang2 , Hanbo Huang2 , Tianxu Lv2 , Minghao Yang2 , Zhongzheng Niu2 , He Bing* 2 , Lusheng Wang* 1 , Jianhua Yao* 2 1
City University of Hong Kong, China
arXiv:2607.05238v1 [cs.AI] 6 Jul 2026
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 leakfree 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×), 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.
Introduction In the autonomous-intelligence roadmap of LeCun (2022), the world model is the centerpiece: an agent plans by rolling a learned latent predictor forward. Modern JEPA world models (Zhou et al. 2024; Assran et al. 2025) instantiate this idea with latent regression: encode the present, predict the next latent, and plan in that learned space. The design is powerful when the future is effectively single-valued. It is structurally mismatched to stochastic dynamics, where the same context can have several valid successors and a planner needs to enumerate them rather than average them. * Corresponding author.
2
Tencent, China
failure Branch one context many futures
fix Collapse single head predicts mean
predictor-level mechanism
Enumerate hard heads cover modes
Audit router-gated realroute plans
planning-level evidence
Figure 1: Argument outline. Stochastic branchings create a one-to-many latent prediction problem; a single JEPA predictor collapses to the conditional mean, while hardassigned predictor heads enumerate the successor modes. The paper then audits whether those modes are actually usable for planning, rather than merely improving multisample coverage. This paper isolates that mismatch at the predictor interface. A deterministic JEPA head trained by squared or cosine regression returns the conditional mean of the successor embeddings, which can lie between all valid futures. A gated weighted-sum MoE still emits one vector and collapses for the same reason. The minimal repair is not a larger encoder or a softer density head, but a hard-assigned set of predictor heads: each observed transition trains only its nearest head, and a router estimates which heads are active. The result is an enumerable set of successor latents, exactly the object a graph search or MPC rollout can consume. Figure 1 sketches the full causal chain. Evidence chain. We prove the collapse for deterministic and gated-MoE predictors (Props. 1 and 2), and prove that MoP-JEPA’s hard assignment is a per-context quantizer of the successor distribution (Prop. 3). We then test the mechanism before emphasizing benchmark scores: two-step branchings produce three or four valid successor modes, the dense predictor places one peak, and MoP-JEPA re-opens the fan-out (Fig. 2); antmaze, pixel observations, ETH/UCY, SVHN, and a DINO-WM port check that the phenomenon is not a 2-D maze artifact (Table 1). Finally, we ask whether the recovered modes support planning. Single-output planners reach only 0.02–0.09 success on OGBench stitch and teleport tasks, whereas planning over our enumerated modes recovers (Table 2) and the recovered plans use real transitions (Table 3).
Why the audit matters. Multi-prediction evaluation has a loophole: more samples can cover more outcomes without being context-conditional enough to plan. Our protocol therefore audits everyone, ourselves first, with an inputagnostic codebook control, a shuffled-context test, a routergated readout, transition-precision guards, and a verifiedroute criterion in which the model proposes a transition graph blind and ground truth is used only to check whether the route exists. This demotes raw coverage before it adjudicates any baseline. Contributions. (1) Diagnosis: conditional-mean collapse of deterministic and gated-MoE JEPA predictors under stochastic dynamics. (2) Fix: MoP-JEPA, a hard-assignment multi-predictor head whose loss is the classical multiplechoice / winner-take-all objective (Guzman-Rivera, Batra, and Kohli 2012; Lee et al. 2016) but whose role here is world-model mode enumeration. (3) Causal evaluation: a collapse-to-planning measurement chain, with controls, that separates usable context-conditional modes from coverage freeloading.
Related Work JEPA world models. I-JEPA (Assran et al. 2023) and VJEPA (Bardes et al. 2024) established latent-regression prediction; DINO-WM (Zhou et al. 2024) and V-JEPA 2 (Assran et al. 2025) use the recipe as an action-conditioned world model for planning. Notably, the original JEPA blueprint (LeCun 2022) includes a latent variable whose stated role is to carry exactly the information a multimodal future leaves undetermined — yet every released system above implements the predictor without it, deterministically. We verified this directly on the flagship: loading the released V-JEPA 2AC checkpoint (a 305M-parameter action-conditioned predictor) with the authors’ own code, the predictor emits a single next-latent per (state, action) — the single-head design our analysis shows collapses under stochastic dynamics. We study what that omission costs and supply the missing component in an enumerable form; we are orthogonal to the encoder and its anti-collapse machinery (Bardes, Ponce, and LeCun 2022). Probabilistic JEPA. Huang et al. (2026) argue JEPA should be probabilistic and add a variational single-Gaussian predictive distribution; the problem framing is theirs. We benchmark their mechanism (with fixes and a steelmanned implementation; Var-JEPA hereafter, to avoid confusion with Meta’s V-JEPA) and find that a single Gaussian spreads but cannot separate discrete successor modes. Enumerable modes, not samples, are what discrete search consumes. MoE predictors in JEPA. M3-JEPA (Lei et al. 2024) implements the predictor as a gated MoE fused by a weighted sum into one output. We prove the fusion collapses identically to the single head (Prop. 2) and confirm this empirically; the fix is the loss (keep K outputs, assign hard), not the expert count. Multiple-choice and mixture regression. MCL/WTA (Guzman-Rivera, Batra, and Kohli 2012; Lee et al. 2016)
and mixture density networks (Bishop 1994) are classical answers to one-to-many regression. Our contribution is the JEPA world-model instantiation, the collapse diagnosis, and the verification protocol; the multimodal-regression literature typically reports oracle coverage without controls, which is the loophole the protocol closes. Offline GCRL (external anchor). OGBench’s published baselines (Park et al. 2025, 2023; Wang et al. 2023; Eysenbach et al. 2022; Ghosh et al. 2021) solve these tasks with reward-free policy learning. We do not claim to beat this line; we report their published numbers as an external anchor at execution time and note the protocol differences.
Setup and Notation Offline transitions D = {(s, a, s′ )} come from the official OGBench datasets (Park et al. 2025). An encoder fθ (s) → z ∈ Rd is trained with an EMA target copy fξ ; a predictor g maps (z, a) to a prediction of z ′ = fξ (s′ ). Stochastic or stitched data makes the conditional law p(z ′ | z, a) multimodal: teleport cells map one (s, a) to one of three fixed destinations, and stitched datasets give one state several observed continuations (measured average 2.1–2.2 successor modes per cell). We write c = (z, a) for the context; µ1 , . . . , µM for the M successor-mode centers with weights w1 , . . . , wM and within-mode standard deviation σ; and ẑ ′ = z ′ /∥z ′ ∥ for normalized targets. The downstream task is goal reaching: plan to the official evaluation goals using only the model (graph search over predicted successors, or model-predictive control in the environment), with success judged without giving the model access to ground truth.
Conditional-Mean Collapse Proposition 1 (deterministic predictor). Under squared loss thePoptimal single predictor is g ∗ (c) = E[z ′ | c] = m wm µm ; its error is lower-bounded by the between-mode variance, and for well-separated modes (minm̸=m′ ∥µm − µm′ ∥ ≫ σ) the optimum lies far from every mode. Under cosine loss with normalized targets and a unit-norm predictor, the exact minimizer is u∗ = E[ẑ ′ ]/∥E[ẑ ′ ]∥, whose similarity to every mode direction is bounded away from 1 whenever the mode directions are mutually separated (Lemma 1, technical appendix). Proposition 2 (gated MoE). A weighted sum ĝ(c) = P π (c) g (c) outputs a single vector; its objective equals k k k Prop. 1’s with v = ĝ(c). Additional experts and router capacity change which single vector is reachable, not how many are output. Proposition 3 (best-of-K). Assume M well-separated modes with bounded within-mode variance, mode weights bounded away from zero, K ≥ M , andpredictor heads with sufficient capacity. Then L({gk }) = E mink ∥gk (c) − z ′ ∥2 is, per context, the K-means distortion of p(· | c) (Lloyd 1982), and every optimum assigns at least one head to each mode, with distortion approaching the within-mode variance. A context-only router trained on the winning index estimates the mode weights. The load-balance and router
dense predictor (standard JEPA)
best-of-K (ours, router )
1.00 0.95
Where a branching maze forks three-to-four ways, a dense predictor collapses onto a single mode; best-of-K re-opens every branch.
GROUND
TRUTH
DENSE
true 2-step successor law
19.8%
cell (0, 3)
Teleport maze
18.8%
100.0%
22.2%
1.4
1.3
cell (5, 1)
Teleport maze
100.0%
0
-1.20 -1.2
0
0
0 1.2
77.4%
100.0%
one peak
11.2%
11.6%
cell (5, 3)
Medium-stitch maze
1.2
1.0
77.2%
1.3
0.5
latent y
0
0
latent y
true dest.
true dest.
0.75 0.70 0.65
Figure 3: The collapse and the fix on official data (OGBench pointmaze-teleport; trained models). Heat: cosine similarity of the predicted next-latent to each cell’s latent. Left: at the teleport cell (star) with two destinations (green circles), the dense predictor decodes to the invalid midpoint cell (red cross). Right: our two router-activated heads (blue crosses) land on both true destinations with weights π = 0.57/0.43, close to the empirical branch frequencies 0.42/0.40. The same geometry viewed in latent space is in the technical appendix.
0
latent x
latent x 1.3
0
-1.30 -1.3 0
latent x
1.3
0.5
0
-1.30 -1.3
-1.30 -1.3
11.4%
1.0
1.0
1.3
latent y
1.2
0.5
0
1.0
7.3%
23.1%29.1%
1.0
0.5
=0.57
0.55
1.0
-1.00 -1.0
0.80
0.60
one peak 1.0
11.3%
true dest.
1.4
7.5%
1.2
=0.43
0
36.4%
23.1%30.3%
0.5
true dest.
1.3
0
density
prediction = invalid midpoint
0 1.4
0.5
teleport cell
-1.40 -1.4
0
-1.20 -1.2
teleport cell
20.3%
0
-1.30 -1.3
1.0
0.85
44.4%
1.4
0.5
0
39.1%
0.90
OURS
1.0
0
density
17.5%
one peak
0.5
-1.40 -1.4
·
belief re-opens the full fan-out
17.8%
1.0
0.5
BEST-OF-K
39.2%
1.0
density
PREDICTOR
belief collapses to one mode
cosine similarity of predicted next-latent to each cell ( 0.55)
Two-step successor beliefs
1.3
1.3
Green pins = true successors. Dense puts 100% on a single mode; best-of-K restores a peak under every branch.
Figure 2: Two-step successor beliefs at high-branching cells. Rows: three cells whose two-step successor law has three or four valid modes. Columns: empirical ground truth (green), the dense single predictor (red), and MoP-JEPA (blue). Each panel is the local 2-D latent plane spanned by the true successor embeddings; floor axes give the latent coordinates, and the vertical coordinate is unit-normalized predictive density. Green pins mark true successor latents. Percentages in the left column are empirical two-step branch frequencies; percentages in the right column are router weights πk for active heads. The dense predictor collapses to one peak, while MoP-JEPA re-opens the multimodal fan-out with one predicted peak per branch.
cross-entropy terms used in training are regularizers outside this statement. Proofs and the precise separation conditions are in the technical appendix. Figure 2 measures Props. 1 and 3 in the latent space of trained models; Figure 3 shows the same models’ decoded predictions on the maze (a controlled bimodal illustration is in the technical appendix); the K-sweep of Figure 9 measures the rise-until-M -then-plateau behavior that Prop. 3 predicts.
Method Variables and standard JEPA objective. Let D = N {(si , ai , s+ i )}i=1 be an offline transition dataset, where si ∈ S is the current observation or state, ai ∈ A ⊆ Rp is the action, and s+ i is the observed successor. In domains without actions, ai is omitted and the context below is just the current latent. The online encoder is fθ : S → Rd and the target encoder is an exponential-moving-average copy fξ : S → Rd .
For a minibatch B, define fθ (si ) , ∥fθ (si )∥2 fξ (s+ i ) zi+ = sg , ∥fξ (s+ i )∥2 ci = [zi , ai ], zi =
(1)
where zi ∈ Rd is the online context latent, zi+ ∈ Rd is the stop-gradient target latent, ci ∈ Rd+p is the predictor input, and sg(·) stops gradients into the target branch. A standard deterministic JEPA predictor is a function gϕ : Rd+p → Rd trained by latent regression, 1 X ℓ ḡϕ (ci ), zi+ , Ldense (θ, ϕ; ξ) = |B| i∈B (2) gϕ (c) ḡϕ (c) = , ∥gϕ (c)∥2 with cosine distance ℓ(u, v) = 1 − u⊤ v in our implementation. After each optimizer step on (θ, ϕ), the target encoder is updated by ξ ← τ ξ + (1 − τ )θ, (3) with EMA momentum τ (we use τ = 0.996). Thus the encoder side is the usual JEPA recipe; the only architectural change below is the predictor. Mixture-of-Predictors parameterization. MoP-JEPA replaces the single predictor gϕ with K predictor heads {gϕk }K k=1 , each mapping the same context ci to one candidate successor latent. It also learns a context-only router rψ : Rd+p → RK . For each example i and head k, define gϕk (ci ) , ∥gϕk (ci )∥2 exp(rψ (ci )k ) πik = PK , j=1 exp(rψ (ci )j )
uik =
dik = ℓ(uik , zi+ ).
(4)
Here uik is the kth predicted successor latent, πik is the router’s deployment-time weight for that head, and dik is the distance from that head to the observed target. The router is deliberately restricted to ci ; it never sees zi+ and therefore cannot choose a head by peeking at the future. Hard-assignment objective. For fixed predictions, the latent target assigns itself to the closest head, ki∗ = arg
min dik ,
1≤k≤K
γik = 1[k = ki∗ ],
(5)
where ki∗ is the winning head and γik is its one-hot assignment. The MoP-JEPA loss is K
LMoP =
1 XX 1 X γik dik + λroute − log πiki∗ |B| |B| i∈B k=1
+ λbal
K X
i∈B
γ̄k log(K γ̄k ),
k=1
(6) P where γ̄k = |B|−1 i∈B γik is the minibatch usage of head k, λroute weights the router cross-entropy, and λbal weights the load-balancing penalty KL(γ̄∥Unif(K)). The first term is the best-of-K JEPA regression loss. It gives gradient only to the winning head for each target, so with assignments fixed the update for head k is exactly the ordinary JEPA predictor update restricted to the subset {i : ki∗ = k}. The second term trains the router to predict the hard assignment from the context alone; the third prevents unused heads during minibatch optimization. When K = 1, Eq. (6) reduces to Ldense up to constants. This is a hard-EM procedure in latent space. The assignment step chooses the nearest current head (ki∗ ), and the update step moves that head toward the assigned target while fitting the router to the same assignments. Under the multimodal conditional law in Prop. 3, this objective is the empirical K-means distortion of p(z + | c); hence the heads specialize to distinct successor modes instead of averaging them into one vector. Baselines and protocol. We compare predictor mechanisms, not codebases. Because none of these predictors has been applied to JEPA world-model planning before, no prior numbers exist for this setting; we therefore reimplement each under one common protocol — identical encoder, offline data, EMA target, and planner — so that the predictor head is the only variable. The arms are: dense (K=1, standard JEPA); M3-JEPA (Lei et al. 2024), its gated mixtureof-experts fused by a weighted sum; Var-JEPA (Huang et al. 2026), a variational single-Gaussian predictor (steel√ manned: cosine-trained mean, residual-calibrated σ, 1/ dscaled sampling); MDN (Bishop 1994), a soft mixture with load balancing (best of a K × λ grid); and MoP-JEPA (ours). These are faithful same-protocol reimplementations, not runs of the authors’ released code; the published OGBench GCRL numbers we quote later are used only as an external anchor, never as a same-protocol comparison. Prediction interface. At test time the model receives only (s, a). It computes z = fθ (s)/∥fθ (s)∥2 , c = [z, a], and
emits the set P(s, a) = {(uk , πk )}K k=1 , gϕk (c) , uk = ∥gϕk (c)∥2 πk = softmax(rψ (c))k .
(7)
The output is therefore an enumerable successor set in one forward pass, not samples from a post-hoc decoder and not a weighted average. A downstream planner or evaluator can consume all heads, or only router- active heads above a fixed threshold, without access to the realized successor.
Verification Protocol Multi-prediction can look good by covering outcomes for free, so every positive result is audited. We use planAll for model-internal graph searchability and official-goal success for held-out goals, but the load-bearing metric is realroute: the model first proposes a transition graph blind, and ground truth is used only afterward to check whether a route made entirely of real transitions exists inside that proposal. Transition precision guards against graph inflation by non-existent edges, and execution success runs the planner closed-loop in the environment. Three controls separate usable mode enumeration from freeloading: an input-agnostic codebook, an off-context shuffle test, and router gating, where only heads with context-predicted mass πk > 0.5/K count.
Experiments Setup. Official OGBench offline datasets: pointmazemedium-stitch, pointmaze-large-stitch, pointmaze-teleportnavigate, and antmaze-teleport (29-D). Leak-free evaluation: coordinate features (no cell-identity memorization) and 20% of unique transitions held out of training. Five seeds for planning, three for coverage and execution; tables give mean±std over seeds and figures show per-seed points. No metric gives the model access to ground truth at decision time.
Evidence Chain and Generality The central question is not whether one predictor variant wins one maze benchmark, but whether the same mechanism repeats across stochastic predictors, branching degrees, and domains. Table 1 summarizes the evidence chain we test before reporting planner scores; Fig. 4 plots the three main diagnostics. The same collapse and fix appear beyond lowdimensional mazes. On a genuine JEPA latent forecaster (ETH/UCY pedestrians (Pellegrini et al. 2009), leave-onescene-out) the mixture beats the codebook control on every scene and seed (0.053 vs. 0.070 vs. dense 0.118, shuffle 2.0×), and on raw multimodal regression it improves pedestrian forecasting (+37.7%) and image inpainting (+24.1%), both fully context-conditional under the controls. Conversely, standard masked-SSL pretext is near-deterministic given context: there is no collapse to fix and multiple predictors tie the single one. This scopes the contribution precisely: JEPA as a forecaster or world model of genuinely multimodal futures.
Beyond 2-D mazes
Branching complexity
101 2.5x
2.5x
1.6x
100 pixel SVHN ETH/UCY DINO-WM COV-R modes ADE head gap
1.0
successor recall
102
Head specialization
0.9
mode coverage
relative gain (log scale)
360.7x
0.8 0.7 0.6
dense MDN MoP
0.5 2
3
4
true modes
5
0.89
0.8
0.82
0.6 0.4 0.2 DINO-WM winner 0.0
361x below mean head
3-way
dense MoP
4-way
Figure 4: Evidence chain before benchmark scores. Left: beyond 2-D mazes, MoP improves the appropriate domain metric (coverage for pixel/SVHN, inverse ADE for ETH/UCY); in DINO-WM, the winning head is hundreds of times closer than the mean head, showing specialization in the official world-model code. Middle: as the true number of modes grows in a controlled conformer branch sweep, dense coverage degrades fast, MDN partially covers, and hard-assigned MoP remains highest. Right: on real two-step maze branchings, dense recovers no successor modes while MoP recovers most 3/4-way successors; this is the empirical specialization predicted by Prop. 3.
Link in the chain
Evidence in the paper
planAll ↑
med-stitch
teleport
large-stitch
Collapse
Dense/gated-MoE emit one latent mean; the decoded maze state is an invalid midpoint (Figs. 2, 3). MoP-JEPA re-opens three/four-way fan-outs; independently trained heads collapse to the same mean. Coverage rises until K matches the mode count; twostep cells expose the rare-mode ceiling. ETH/UCY, SVHN, antmaze, pixel observations, and DINO-WM show the same mechanism beyond 2-D point mazes. Realroute verifies that enumerated modes form a searchable graph of real transitions.
dense (JEPA) M3-JEPA Var-JEPA MDN MoP-JEPA (ours)
.084±.03 .075±.02 .085±.01 .566±.20 .851±.08
.037±.01 .036±.01 .035±.01 .512±.12 .748±.09
.016±.01 .031±.01 .033±.00 .717† ±.12 .386±.45
codebook (control)
.267±.02
.167±.02
.149±.04
Specialization Branching complexity Generality
Planning consequence
Table 1: The paper’s evidence chain: stochastic targets create conditional-mean collapse; hard-assigned heads specialize to modes; mode enumeration restores a transition graph that can be audited by planning.
The mechanism also transfers to a second codebase and to image observations. Porting the head into DINO-WM’s official world model (Zhou et al. 2024) is a one-line predictor override (predictor=mixture) with no other pipeline change; it trains as a drop-in, and its heads specialize rather than duplicate — on held-out pixel point-maze the winning head’s next-latent error is two orders of magnitude below the head average, with ≈ 2 heads active per state, matching the maze’s local branching factor. On masked-digit SVHN (real images, bottom-half inpainting), the deterministic head realizes only 0.31 of the conditional digit-modes while the mixture heads realize 0.78 (2.5×, 3 seeds); density and ensemble baselines reach 0.85–0.89 but expose no enumerable, router-gated successor set to plan through. Finally, the collapse and its fix persist as the number of futures grows. Aggregated over all three- and four-way twostep maze cells (3 seeds, technical appendix), MoP-JEPA
Table 2: planAll (searchability of the predicted graph; 5 seeds). Official-goal success shows the same ordering (Fig. 5); the codebook control reaches near-zero official-goal success (0.00–0.15) despite its nonzero planAll. † MDN’s large-stitch score depends on low-precision predicted edges and is adjudicated in Table 3.
recovers 82–88% of the successors while the dense predictor recovers none: its single mean prediction lands between the modes, on no valid state. The hard-assigned heads retain even low-probability (< 10%) successors; the residual gap is a rare-mode ceiling at four-way branching (a single tail mode missed in some cells), consistent with the winnertake-all limit we report in the limitations.
Planning on the Predicted Graph Table 2: deterministic, gated-MoE, and Var-JEPA predictors reach 0.02–0.09, consistent with Props. 1–2; ours reaches up to 0.85.
Adjudication: Verified Routes MDN’s raw numbers depend on predicted edges that do not exist: its transition precision is 0.14–0.21 (ours 0.42–0.56), and graph search will happily route through hallucinated walls. Under realroute, all three mazes favor ours (Table 3,
N rs ok se PA PA den M3-JE Var-JE MD ou odebo c
teleport-navigate
large-stitch
0.75
0.72
0.51 0.04 0.04 0.03
realroute (real route inside predicted graph)
planning success (planAll)
medium-stitch 1.0 0.85 0.8 0.57 0.6 0.4 0.27 0.2 0.08 0.07 0.09 0.0
0.39 0.17
N rs ok se PA PA den M3-JE Var-JE MD ou odebo c
0.02 0.03 0.03
0.15
N rs ok se PA PA den M3-JE Var-JE MD ou odebo c
Figure 5: planAll by arm (bars: mean; dots: seeds). The codebook control (hatched) shows that coverage alone does not plan: it maxes raw coverage (Table 5) yet reaches nearzero official-goal success.
dense
MDN
MoP-JEPA
medium-stitch teleport-navigate large-stitch
.083±.03 .035±.01 .016±.01
.134±.03 .039±.01 .057±.03
.247±.08 .192±.09 .139±.15
Table 3: Verified planning: a real route must exist inside the model’s blind proposal (5 seeds). The failure modes separate: dense is precise but mode-collapsed (a graph too sparse to route); MDN covers but its routes rely on nonexistent edges; MoP-JEPA is the only predictor whose proposed graph contains real routes at a useful rate.
Fig. 6). The paired advantage over MDN is statistically significant on medium-stitch (bootstrap 95% CI of the per-seed difference [0.06, 0.18]) and teleport ([0.07, 0.24]); on largestitch the mean favors ours but the CI includes zero at five seeds, reflecting the seed-bimodality we report below. Under fully matched capacity and training (K=16, 60k steps, 5 seeds, both arms’ best load-balance settings) the largestitch verdict is unambiguous: ours reaches realroute 0.208 vs. MDN’s 0.095 (2.2×), with MDN’s found rate high (0.88) but precision low (0.21) — it routes through non-existent edges — against our 0.42 precision. More capacity buys the MDN more hallucinated edges, not more real routes. Long horizons. Fig. 7: at distance 6 on medium-stitch ours retains 0.18 verified success versus MDN’s 0.01; on teleport, where long routes cross the stochastic branching repeatedly, ours holds 0.05–0.15 out to distance 14 while both baselines reach zero by distance 6–8. Mode errors compound multiplicatively with rollout depth, which is why a mode-faithful one-step predictor is the load-bearing requirement for long-horizon planning. Is the MDN under-tuned? A fairness grid over components K ∈ {8, 16} and load-balance weight λ ∈ {0, 0.01, 0.1} (learned σ throughout, 3 seeds each) leaves the ordering unchanged on every task: the strongest variant reaches realroute 0.125 on medium (ours 0.247) and 0.114 on teleport (ours 0.192). On large-stitch, strong loadbalancing genuinely helps MDN (best variant 0.191); under fully matched capacity and training it remains behind ours.
dense (JEPA) MDN best-of-K (ours)
0.3 0.2 0.1 0.0
medium-stitch teleport-navigate large-stitch
Figure 6: Verified routes by task (bars: mean; dots: seeds). medium-stitch
0.6 verified planning (realroute)
realroute ↑
0.4
teleport-navigate
large-stitch
dense (JEPA) MDN best-of-K (ours)
0.4 0.2 0.0
2
4 6 8 10 12 14 true shortest-path distance
2
4 6 8 10 12 14 true shortest-path distance
2
4 6 8 10 12 14 true shortest-path distance
Figure 7: Verified (realroute) success vs. the true shortestpath distance of the goal (3 seeds). The multimodal advantage compounds with horizon: on teleport, ours still plans at distances 10–14 where both baselines reach zero by d=6; on large-stitch all arms decay quickly.
Verified Planning Across the JEPA Predictor Family The family comparison is same-protocol by construction: every arm shares the encoder, the offline data, and the planner, and differs only in the predictor head, so Table 4 isolates the effect of the predictor’s output distribution. Our executor is verified sound (15/15 on true paths; on a truth graph corrupted with 10% false edges, replanning lifts success 0.40 → 1.00), and a uniform-edge-cost ablation leaves all execution rankings unchanged. The published-baseline numbers are quoted only as external context, not a sameprotocol comparison; the load-bearing result is that a oneline predictor change turns a world model that cannot plan into one that plans verifiably.
Coverage of Stochastic Dynamics Destinations are fixed global states, so the input-agnostic codebook maxes raw COV-R at 1.0; our own protocol flags naive coverage as freeloading. The verdict is thresholdinsensitive: sweeping the gate over t ∈ [1/64, 1/4], our gated coverage degrades gracefully (0.74 → 0.40) while precision (≈ 0.92) and shuffle leakage (≈ 0.26) stay flat, and every baseline is invariant or pinned below. On antmaze-teleport (29-D ant, destinations as manifolds of 32 real samples) the ranking persists with honest attenuation: ours 0.576 raw, about 2× dense (0.294) and 1.5× MDN (0.386); gated 0.342 at 0.86 precision; Var-JEPA 0. On pixel observations (visual-antmaze-teleport, 64×64×3;
verified planning (realroute) ↑ JEPA predictor (same encoder, data, planner) dense (standard JEPA) M3-JEPA Var-JEPA (variational) MDN (best of a K × λ grid) MoP-JEPA (ours)
execution ↑
med-stitch
teleport
large-stitch
med
large
.083 .075 .034 .134 .247
.035 .036 .001 .039 .192
.016 .031 .033 .095 .208
.27 — — .87 .67
.13 — — .27 .33
Table 4: Within the JEPA world-model family — same encoder, data, and planner, only the predictor changes — ours wins verified planning on all three mazes (2–5×; 5 seeds.) Left block (realroute): a plan counts only if the model’s blindly-proposed transition graph contains a route of real transitions, so coverage cannot be freeloaded. Dense/gated-MoE/VarJEPA are mode-collapsed (Props. 1–2); MDN covers but routes through non-existent edges (precision 0.14–0.21 vs. ours 0.42–0.56; large-stitch adjudicated at matched K=16/60k steps). Right block: real-environment execution success (official info[’success’], our replanning executor); MDN’s higher medium-stitch number is the executor’s online repair compensating its low-precision graph, which is precisely why model fidelity is judged by realroute. External context: on large-stitch our execution (33) sits second among the seven published OGBench offline-GCRL baselines (Park et al. 2025) (QRL 84; then ours, GCIQL 31, HIQL 13, GCIVL 12, GCBC 7, CRL 0) — notable for a generic world model plus search rather than a policy trained end-to-end for the benchmark.
raw
π-gated
COV-P
shuffle
dense M3-JEPA Var-JEPA MDN MoP-JEPA
.244±.03 .242±.04 .001±.00 .596±.02 .997±.00
.244±.03 .242±.04 .001±.00 .292±.01 .632±.22
.733 .727 .000 .864 .924
2.4× 1.4× — 1.8× 2.4×
codebook
1.00±.00
1.00‡
.417
1.0×
pointmaze-teleport
Table 5: Coverage of the three teleport destinations (COVR, 3 seeds). ‡ The codebook has no context and cannot be gated. Raw coverage can be freeloaded (the codebook maxes it), so raw COV-R is a diagnostic, never a headline; under the router-gated readout ours is the only arm simultaneously high on coverage, precision, and context-conditionality.
dense M3-JEPA Var-JEPA MDN MoP-JEPA codebook
modes enum.?
COV-R (gated)
prec.
realroute
exec (hard)
no no samples soft yes static
.24 .24 .00 .29 .63 n/a
high high — .14–.21 .42–.56 —
0/3 0/3 0/3 0/3 3/3 0/3
.13 — — .27 .33 —
Table 6: Trade-off summary. Each baseline misses a different requirement of a plannable stochastic world model; MoP-JEPA is the only arm that enumerates contextconditional modes with usable precision and leads verified planning on all three mazes.
Summary and Ablations teleport-destination COV-R
pointmaze-teleport (2-D)
antmaze-teleport (29-D ant) 1.00
raw K modes router- gated
1.0 0.8 0.6 0.4 0.2 0.0
1.00
0.63 0.24
0.29
0.29
0.24 0.00
se
den
M3-
JEPA
EPA
-J Var
0.34
0.24
0.24 0.00
N
MD
our
s
k
boo
e cod
se
den
M3-
JEPA
-J Var
EPA
N
MD
s
our
e cod
k
boo
Figure 8: Raw (light) vs. router-gated (solid) coverage on pointmaze and 29-D antmaze.
CNN JEPA trained from scratch, teleport transitions oversampled identically for every arm), the ranking again persists: ours 0.613 raw COV-R vs. 0.24–0.27 for every baseline (2.4×), with partial mode recovery (gated ≈ 1 of 3 modes at 0.97 precision). High-dimensional and pixel-level JEPA world models capture the branching direction but not yet full mode-fidelity; closing that gap is open.
K-sweep (real data). Teleport COV-R rises from 0.231 at K=1 (matching dense 0.258) to 0.712 at K=3 (the number of modes) and plateaus at 0.945 by K=8 (MDN at K=8: 0.531). Design ablations. An explicit head-separation loss is redundant (WTA already separates, and it can hurt recall). Slow soft-to-hard annealing over 80% of training drags every head toward the mean; a fast warmup accelerates smallmodel convergence dramatically while plain hard EM is already sufficient at full scale. An ensemble of K independently trained heads (no assignment) collapses entirely: realroute 0.03–0.10, at dense level, with zero spread across heads — each head converges to the same conditional mean, Prop. 1 applied K times. Independence is not diversity; hard assignment is the load-bearing design choice.
The Protocol Fires Where It Should It audits the strongest baseline: MDN’s raw planning scores are traced to non-existent predicted edges (§Adjudication). It audits us: raw coverage is demoted as freeloadable
teleport COV-R
1.0 0.8 0.6 0.4 0.2 0.0
K = #modes
1
best-of-K (ours) dense (K=1) MDN (K=8)
2 3 4 number of predictor heads K
8
Figure 9: K-sweep on real data (teleport COV-R): coverage rises until K equals the number of modes, then plateaus, with no penalty for over-provisioning, as Prop. 3 predicts.
(Table 5). It detects metric artifacts: on antmaze, fullobservation destination centroids score every arm and every control at exactly zero, the fingerprint of an off-manifold reference point (ant joints are arbitrary at a fixed-xy destination), which we caught and replaced with real-sample manifolds. It rejects illusory multimodality: on road-constrained taxi forecasting an apparent +36% gain is 84% reproducible by the codebook (shuffle 1.23×) and is rejected; on persistence-dominated next-frame video no gain is claimed. The verdict tracks whether futures genuinely branch, which is what makes the positive results trustworthy.
Discussion and Limitations Why this matters for scaling JEPA world models. Video- and robot-scale world models (Assran et al. 2025; Zhou et al. 2024) are trained on data where stochasticity is the rule. Our results indicate that the deterministic predictor’s failure there is not graceful: the error concentrates on exactly the branchings a planner must resolve, and it silently produces states that do not exist. The broader lesson is that a stochastic world model should expose a set-valued transition interface, not only a latent expectation. MoP-JEPA is one minimal implementation of that interface: a head-level change, orthogonal to encoder architecture and scale, that yields enumerable branch outcomes with calibrated router weights. The verification protocol is architecture-agnostic: any world model that emits multiple futures (mixture, variational, diffusion, or generative) can be audited with the same suite, and we would encourage its adoption wherever multisample prediction numbers are reported. Limitations. The claim concerns the transition distribution and what it enables; single-shot selection of an aleatoric branch is impossible for any model and is not claimed. WTA optimization on the hardest maze is seed-bimodal, and at four-way two-step branching a rare (< 10%) successor can fall below head activation and be missed (full four-mode recovery in 37% of such cells); annealing does not fix it at scale, an explicit head-repulsion regularizer hurts, and headwise learning rates or replay biased toward rare modes remain untested. Mazes are low-dimensional state spaces; the
29-D ant extends the mechanism with attenuation, and the DINO-WM port confirms the head-level fix on pixel observations, though scaling the verification protocol to videoand robot-scale world models remains future work. Comparisons within the predictor family are same-protocol reimplementations (no prior numbers exist for these predictors in this setting); the published-baseline anchor is quoted, not rerun, and is not a same-protocol comparison.
Conclusion A deterministic JEPA world model cannot represent a stochastic branching: it predicts a latent expectation where planning needs enumerable futures. We prove this conditional-mean collapse, show it across branch counts and beyond 2-D mazes, and measure the planning failure it causes. Hard-assigned predictors recover one successor mode each, turning the predictor output from a single invalid average into a searchable transition set; verified planning succeeds when that set contains real routes. The method is deliberately minimal, but the point is broader: stochastic JEPA world models need mode enumeration, and claims about multimodal prediction should be audited for contextconditional, real-transition structure rather than raw oracle coverage.
References Assran, M.; Bardes, A.; Fan, D.; Garrido, Q.; Howes, R.; et al. 2025. V-JEPA 2: Self-supervised video models enable understanding, prediction and planning. arXiv preprint arXiv:2506.09985. Assran, M.; Duval, Q.; Misra, I.; Bojanowski, P.; Vincent, P.; Rabbat, M.; LeCun, Y.; and Ballas, N. 2023. Self-supervised learning from images with a joint-embedding predictive architecture. In CVPR. Bardes, A.; Garrido, Q.; Ponce, J.; Chen, X.; Rabbat, M.; LeCun, Y.; Assran, M.; and Ballas, N. 2024. Revisiting feature prediction for learning visual representations from video. arXiv preprint arXiv:2404.08471. Bardes, A.; Ponce, J.; and LeCun, Y. 2022. VICReg: Variance-invariance-covariance regularization for self-supervised learning. In ICLR. Bishop, C. M. 1994. Mixture density networks. Technical report, Aston University. Eysenbach, B.; Zhang, T.; Levine, S.; and Salakhutdinov, R. 2022. Contrastive learning as goal-conditioned reinforcement learning. In NeurIPS. Ghosh, D.; Gupta, A.; Reddy, A.; Fu, J.; Devin, C.; Eysenbach, B.; and Levine, S. 2021. Learning to reach goals via iterated supervised learning. In ICLR. Guzman-Rivera, A.; Batra, D.; and Kohli, P. 2012. Multiple choice learning: Learning to produce multiple structured outputs. In NeurIPS. Huang, W.; et al. 2026. VJEPA: Variational joint embedding predictive architectures as probabilistic world models. arXiv preprint arXiv:2601.14354. LeCun, Y. 2022. A path towards autonomous machine intelligence. OpenReview.
Lee, S.; Purushwalkam, S.; Cogswell, M.; Ranjan, V.; Crandall, D.; and Batra, D. 2016. Stochastic multiple choice learning for training diverse deep ensembles. In NeurIPS. Lei, H.; Cheng, X.; Qin, Q.; Wang, D.; Fan, K.; Huang, H.; Gu, Q.; Wu, Y.; Jiang, Z.; Chen, Y.; and Ji, L. 2024. M3JEPA: Multimodal alignment via multi-gate MoE based on the joint-embedding predictive architecture. arXiv preprint arXiv:2409.05929. Lloyd, S. 1982. Least squares quantization in PCM. IEEE Transactions on Information Theory, 28(2): 129–137. Park, S.; Frans, K.; Eysenbach, B.; and Levine, S. 2025. OGBench: Benchmarking offline goal-conditioned RL. In ICLR. Park, S.; Ghosh, D.; Eysenbach, B.; and Levine, S. 2023. HIQL: Offline goal-conditioned RL with latent states as actions. In NeurIPS. Pellegrini, S.; Ess, A.; Schindler, K.; and Van Gool, L. 2009. You’ll never walk alone: Modeling social behavior for multitarget tracking. In ICCV. Wang, T.; Torralba, A.; Isola, P.; and Zhang, A. 2023. Optimal goal-reaching reinforcement learning via quasimetric learning. In ICML. Zhou, G.; Pan, H.; LeCun, Y.; and Pinto, L. 2024. DINOWM: World models on pre-trained visual features enable zero-shot planning. arXiv preprint arXiv:2411.04983.