Organizations: School of Computer Science and Technology, University of Chinese Academy of Sciences · State Key Laboratory of AI Safety, Institute of Computing Technology, Chinese Academy of Sciences · Beijing Academy of Artificial Intelligence · Institute of Information Engineering, Chinese Academy of Sciences · School of Cyber Security, University of Chinese Academy of Sciences
As a prototypical latent world model, Joint-Embedding Predictive Architectures (JEPAs) provide strong semantic representations. However, dense patch representations expose a redundant interface to downstream predictors, which must implicitly infer spatial organization and temporal correspondence, presenting three challenges. First, object information is distributed across patches, forcing each downstream predictor to redundantly recover object structure that could otherwise be reused across tasks before modeling interactions. Second, the lack of explicit entity correspondence across time makes historical transitions expensive, restricting the predictor's flexibility in utilizing context. Third, without explicit transport structure, existing JEPA models typically use pointwise l1/l2 distances, which are insensitive to spatial displacement when predicted and goal supports are largely disjoint. We propose the Slot-based Joint-Embedding Predictive Architecture (Slot-JEPA), which decomposes dense representations into stable semantic slots and time-varying spatial masks. Pretrained once and frozen downstream, this decomposition exposes a reusable object-structured interface. Building on it, our Transition Memory Predictor (TMP) stores visual transitions as compact masks and retrieves useful context using statistics of corresponding slots. Its two branches separate slot-specific dynamics from shared cross-slot dynamics. We further use Unbalanced Optimal Transport (UOT) to supervise mask prediction and evaluate candidate plans by transporting mass between aligned predicted and goal masks, providing a signal better aligned with spatial progress. Across seven visual-understanding benchmarks, the slot-reconstructed features remain transferable. Slot-JEPA also improves goal-conditioned planning across six planar environments and four robotic arm planning tasks on the REALM-JEPA benchmark.
Figures & tables
Figure 1: Overview of Slot-JEPA. The slot decomposition module factorizes patch representations into stable slots S and time-varying masks M . TMP predicts M^t+1 with sampled transition memory. The prediction is supervised against the next-frame mask by Lpred , including UOT loss.
Method
Segmentation (mIoU ↑ )
Depth ( δ1↑ )
VOS ( J&F↑ )
COCO
ADE
VOC
NYUv2
KITTI
SUNRGBD
DAVIS
MoCo v3 ( Chen et al., 2021 )
43.4
24.4
61.7
69.8
77.2
50.6
61.7
iBOT ( Zhou et al., 2022a )
52.2
34.3
73.3
77.7
88.3
69.1
64.3
Mugs ( Zhou et al., 2022b )
51.4
34.6
72.7
74.9
87.7
66.6
62.1
MAE ( He et al., 2022 )
37.9
20.6
52.3
65.3
84.3
59.1
48.9
DINOv2 † ( Oquab et al., 2024 )
53.0
35.5
74.9
78.1
88.0
68.0
64.2
Table 1: Visual-understanding results on all three probes.
Method
Goal State
ID Environments
OOD Environments
TwoRoom
Maze
PushT
TwoRoomR
MazeR
PushObj
DINO-CLS ( Caron et al., 2021 )
✗
6
88
20
34
70
2
✓
84
–
28
48
–
8
ResNet ( He et al., 2016 )
✗
6
98
34
12
68
4
✓
84
–
40
60
–
16
R3M ( Nair et al., 2023 )
✗
38
100
16
8
32
2
Table 2: Planar-planning success rates (%). Goal State indicates whether the planner also receives the low-dimensional goal state.
Figure 4
Method
Reach
Pick
Rotate
Place
Overall
V-JEPA 2.1
13/18 (72.2)
19/30 (63.3)
6/30 (20.0)
0/30 (0.0)
38/108 (35.2)
V-JEPA 2.1 + TMP
16/18 (88.9)
22/30 (73.3)
13/30 (43.3)
12/30 (40.0)
63/108 (58.3)
Slot-JEPA
18/18 (100.0)
27/30 (90.0)
25/30 (83.3)
23/30 (76.7)
93/108 (86.1)
Table 3: Robotic arm planning success rates on REALM-JEPA (count and %).
Method
Obs. loss vs. EE–goal dist.
Pred. loss vs. EE–goal dist.
Pred. loss vs. Obs. loss
Ext. 1
Ext. 2
Wrist
Ext. 1
Ext. 2
Wrist
Ext. 1
Ext. 2
Wrist
V-JEPA 2.1
0.350
0.407
0.533
0.503
0.453
0.297
0.511
0.609
0.549
V-JEPA 2.1 + TMP
0.154
0.103
0.652
0.327
0.319
0.454
0.409
0.421
0.720
Slot-JEPA
0.792
0.891
0.807
0.677
0.788
0.814
0.668
0.825
0.865
Table 4: Spearman correlations among observed/predicted loss and EE–goal distance.
Figure 4: Goal-loss trajectories for Pick and Place. Panels use independent scales.
Figure 5: Slot-mask overlays from Kinetics-400, DROID, and REALM-JEPA (top to bottom). Colors track slots across each sequence.
ID
Predictor
Planning loss
Context
TwoRoom
Maze
PushT
TwoRoomR
MazeR
PushObj
Reach
1
Sequential
UOT + mask ℓ1
Recent transitions
100
100
76
94
98
40
–
2
TMP
Mask ℓ1
Sampled via Eq. ( 6 )
96
92
82
78
90
50
88.9
3
TMP
UOT + mask ℓ1
Recent transitions
94
86
76
84
80
58
55.6
4
TMP
UOT + mask ℓ1
Random transitions
96
88
80
40
72
48
94.4
5
TMP
UOT + mask ℓ1
Sampled via Eq. ( 6 )
100
100
92
90
96
58
100.0
Table 5: Ablation success rates (%) across six planar environments and robot Reach. Bold marks the highest value in each column.
Appendix figures & tables16 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 7: Network architecture of the slot decomposition module.
Figure 8: TMP architecture. Observed transitions form the context blocks, while inverse-encoded pseudo-action and learned future tokens specify the query transition. The reconstructed-feature tower and the three-scale object feature pyramid run in parallel and are fused to predict the future slot mask.
u←(max(Kv,δ)a)θ,v←(max(K⊤u,δ)b)θ.
Appendix
Algorithm 1 Generalized Sinkhorn scaling for one aligned slot-mask pair
Figure 9: Examples of slots identified from their predicted changes. Moving-slot masks follow the controlled entity across probe transitions, while static-slot masks remain aligned with the wall or maze geometry. The labels are produced by the change-score rule rather than ground-truth segmentation.
COCOStuff-27
ADE20K
PASCAL VOC
Method
Backbone
Data
Ep.
mIoU ↑
Acc. ↑
mIoU ↑
Acc. ↑
mIoU ↑
Acc. ↑
DINOv2
ViT-S/14
LVD-142M
–
57.4
77.7
43.5
78.9
79.4
95.3
DINOv2
ViT-B/14
LVD-142M
–
59.1
78.4
46.3
79.8
81.3
95.7
MoCo v2
RN-50
IN1K
800
37.6
64.0
20.0
61.2
53.3
88.4
SwAV
RN-50
IN1K
800
42.2
68.0
23.3
65.8
55.4
89.0
DINO
RN-50
IN1K
800
40.7
66.6
22.4
64.5
53.1
88.4
Appendix
Table C1: Complete linear semantic-segmentation results (%). “Data” and “Ep.” denote representation-pretraining data and epochs; IN1K, COCO, K400, and VM163M denote ImageNet-1K, COCO, Kinetics-400, and VisionMix-163M. † marks reproduced DINOv2.
NYUv2
KITTI
NYU → SUNRGBD
Method
Backbone
δ1↑
AbsRel ↓
RMSE ↓
δ1↑
AbsRel ↓
RMSE ↓
δ1↑
AbsRel ↓
RMSE ↓
DINOv2 (LVD-142M)
ViT-S/14
84.8
0.134
0.437
91.0
0.096
3.119
77.4
0.269
0.473
DINOv2 (LVD-142M)
ViT-B/14
87.8
0.120
0.408
92.0
0.092
2.931
80.7
0.255
0.443
MoCo v2
RN-50
62.9
0.230
0.719
82.4
0.130
4.319
55.5
0.366
0.749
SwAV
RN-50
61.9
0.225
0.721
84.3
0.123
3.985
55.0
0.367
0.753
DINO
RN-50
61.6
0.228
0.726
83.7
0.127
4.068
54.7
0.370
0.755
Appendix
Table C2: Complete monocular-depth and NYUv2-to-SUNRGBD transfer results. † marks reproduced DINOv2.
Method
Backbone
J↑
F↑
J&F↑
DINOv3
ViT-7B/16
–
–
71.1
MoCo v3
ViT-S/16
–
–
61.7
SwAV
RN-50
–
–
41.3
DINO
ViT-B/16
61.5
64.9
63.2
iBOT
ViT-B/16
63.2
65.5
64.3
MAE
ViT-S/16
–
–
48.9
Appendix
Table C3: Complete DAVIS video object segmentation results (%). MoCo v3, SwAV, and MAE report best-checkpoint J&F only.
Figure 10: The three planar ID/OOD environment pairs.
Training environment
Resolution
Epochs
Batch size
Weight decay
TwoRoom / Maze
384
100
128
10−5
PushT
512
100
128
10−5
TwoRoomR / MazeR
384
50
128
10−2
PushObj
512
500
64
10−4
Appendix
Table C4: Training settings for the planar environments.
Figure 11: Representative REALM-JEPA start and goal images for all four Robotic Arm Planning categories. Each task panel shows the same scene and goal stages from the two external cameras and the wrist camera. Goal 1 and Goal 2 denote successive visual goals when a task has multiple stages.
Figure 12: Additional goal-loss landscapes in MazeR and TwoRoomR. Each sampled state is rendered independently; the panels visualize energy geometry rather than executed rollouts. Stars mark goals, and lighter colors indicate lower loss.
Figure 13: Predicted slot masks along successful TwoRoomR, MazeR, PushObj, and Place planning rollouts (top left, top right, bottom left, and bottom right, respectively). Each column pairs the RGB observation obtained after action execution with masks predicted for the same time step. The planar panels show mask heatmaps, whereas the Place panel overlays sharpened predictions on the executed wrist-view RGB.
Figure 14: Slot decomposition on Kinetics-400 pretraining example. The first row contains the RGB frames, and each subsequent row overlays one selected soft slot mask.
Figure 15: Slot decomposition on Kinetics-400 pretraining example. The first row contains the RGB frames, and each subsequent row overlays one selected soft slot mask.
Figure 16: Slot decomposition on DROID downstream-training example. The first row contains the RGB frames, and each subsequent row overlays one selected soft slot mask.
Figure 17: Slot decomposition on DROID downstream-training example. The first row contains the RGB frames, and each subsequent row overlays one selected soft slot mask.
Joint-Embedding Predictive Architectures (JEPAs) have emerged as a principled framework for self-supervised learning of world models in compact latent spaces, yet existing methods are fragmented: some predict masked parts of a single image in latent space (I-JEPA), others learn to predict global photometric transformations (Image World Models), while video-scale JEPAs predict future temporal states and are post-trained for action-conditioned planning (V-JEPA~2, DINO-World, DINO-WM). These objectives are treated as distinct recipes with separate encoders, predictors, and anti-collapse regularizers, hindering a single model from unifying image-level and video-level world modeling. We present UniJEPA, a unified JEPA that jointly learns photometric prediction (image-level transformations) and temporal prediction (video-level next-state dynamics) in one shared latent space. A single end-to-end objective, composed of a next-embedding prediction loss and a Gaussian regularizer, yields a provably anti-collapse encoder-predictor pair trainable from raw pixels without EMA, stop-gradient, or pre-trained encoders. We show that the same latent space supports controllable abstraction: photometric prediction learns invariant structure while temporal prediction learns equivariant dynamics. After action-conditioned post-training on offline trajectories, UniJEPA enables zero-shot planning by treating goal features as prediction targets. On image, video, and control benchmarks, UniJEPA matches or surpasses task-specific JEPAs while requiring a single loss hyperparameter, and plans up to tens of times faster than generative world models at comparable accuracy.
An Lanji, Dawei Liu, Jin Li +3
University of Electronic Science and Technology of China, Chengdu, China.
Joint-Embedding Predictive Architectures (JEPAs) learn world models by predicting in representation space rather than reconstructing pixels, making them a natural backbone for latent model predictive control from offline demonstration logs. JEPA-style training optimizes short-horizon latent prediction, whereas planning requires a multi-step ranking of imagined futures by goal progress. Prior JEPA planners often inherit that ranking from embedding geometry, typically latent Euclidean distance, which arises as a byproduct of representation learning rather than as a progress cost mined from the logs. We propose Temporal-Distance-JEPA, which retains the LeWM encoder--predictor backbone and mines a directed temporal cost from reward-free trajectories: same-trajectory step order supplies positive targets, cross-trajectory pairs act as heuristic negatives, and a rollout-consistency term matches the planner horizon. The mined supervision serves two roles: as the deployed planning cost when progress is topological, and as a representation signal that improves Euclidean planning when contact geometry dominates. Under locked evaluation, deploying the mined cost raises Two-Room success to 100.0% versus LeWM's 97.4%, while shared Euclidean planning on the same temporally trained checkpoint raises OGB-Cube by 14.2 points over LeWM and improves Push-T. Against LeWM and the concurrent RC-aux baseline under locked evaluation, Temporal-Distance-JEPA matches or exceeds both methods on every environment. Ablations show that the directed head, cross-trajectory negatives, and rollout consistency each contribute. Temporal-Distance-JEPA narrows the train--plan gap for JEPA world-model planners by discovering temporal progress structure in offline logs and co-designing cost form with plan-time deployment. Code is available at https://github.com/HKBU-KnowComp/Temporal-Distance-JEPA.
Joint-embedding predictive architectures (JEPAs) learn world models in latent space, enabling efficient planning but leaving their states opaque. We introduce Inside-JEPA, which places relational channels for ordered entity pairs inside the JEPA latent state and shows that an inverse-dynamics loss can ground these channels without relation labels when its head reads them, either alone or as part of the full state. On Something-Something V2, latent prediction alone performs no better than a random relation head, while inverse-dynamics supervision that recovers actions from channel changes improves relational AUROC by .065-.103 across ten settings with annotated boxes, by .081 on RH20T robot episodes with held-out tasks, and by .050-.061 across five relations using off-the-shelf detection and tracking. An inverse head on entity latents decodes actions more accurately but barely improves relational grounding; shuffled labels, two label-free objectives, and action labels on only 5% of the pairs also fail to ground the channels. Given an action, predicted relational changes approach the quality of observed future relations (AUROC .599 vs .621). For one-step relational goals, a supervised readout of the 16 channels ranks the correct action among the top five of 65 in 43% of cases, compared with 24% for prediction-only latents and 17% for a frequency prior.
Hossein Resani, Javen Qinfeng Shi
Australian Institute for Machine Learning, Adelaide University · Responsible AI Research Centre