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

JEPA world models commonly predict the next latent state with one regressor. Under stochastic transitions, squared and cosine regression return the conditional mean and its normalized direction, respectively: a single compromise that may match no valid successor. MoP-JEPA instead uses K hard-assigned heads and a context-only router to produce a finite candidate set in one pass. On held-out OGBench transitions, graph search with single-output predictors succeeds on 0.02--0.09 of queries, whereas MoP-JEPA reaches 0.85. To distinguish useful successors from indiscriminate coverage, we also measure verified-route success (realroute), which checks after graph construction whether the proposal contains a path of real transitions. MoP-JEPA leads this same-protocol metric on all three mazes; an MDN attains high raw coverage but predicts many nonexistent edges.