Slot-JEPA: Object-Structured Prediction Interfaces for Latent World Models

Peisong Wen; Qianqian Xu; Minye Lei; Siran Dai; Yang Liu; Qingming Huang

cs.AIOct 20, 2026

Slot-JEPA: Object-Structured Prediction Interfaces for Latent World Models

Authors: Peisong Wen, Qianqian Xu, Minye Lei, Siran Dai, Yang Liu, Qingming Huang

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 ·

Abstract

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 $\bm{S}$ and time-varying masks $\bm{M}$ . TMP predicts $\hat{\bm{M}}_{t+1}$ with sampled transition memory. The prediction is supervised against the next-frame mask by $\mathcal{L}_{\mathrm{pred}}$ , including UOT loss.

Method	Segmentation (mIoU $\uparrow$ )			Depth ( $\delta_{1}\uparrow$ )			VOS ( $\mathcal{J}\&\mathcal{F}\uparrow$ )
	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 $\dagger$ ( 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.

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.

Panel of Figure 5 — 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 $\ell_{1}$	Recent transitions	100	100	76	94	98	40	–
2	TMP	Mask $\ell_{1}$	Sampled via Eq. ( 6 )	96	92	82	78	90	50	88.9
3	TMP	UOT + mask $\ell_{1}$	Recent transitions	94	86	76	84	80	58	55.6
4	TMP	UOT + mask $\ell_{1}$	Random transitions	96	88	80	40	72	48	94.4
5	TMP	UOT + mask $\ell_{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

Panel of Figure 7 — Figure 7: Network architecture of the slot decomposition module.

Panel of Figure 8 — 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\leftarrow\left(\frac{a}{\max(Kv,\delta)}\right)^{\theta},\qquad v\leftarrow\left(\frac{b}{\max(K^{\top}u,\delta)}\right)^{\theta}.

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 $\uparrow$	Acc. $\uparrow$	mIoU $\uparrow$	Acc. $\uparrow$	mIoU $\uparrow$	Acc. $\uparrow$
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.

\dagger

marks reproduced DINOv2.

		NYUv2			KITTI			NYU $\rightarrow$ SUNRGBD
Method	Backbone	$\delta_{1}\uparrow$	AbsRel $\downarrow$	RMSE $\downarrow$	$\delta_{1}\uparrow$	AbsRel $\downarrow$	RMSE $\downarrow$	$\delta_{1}\uparrow$	AbsRel $\downarrow$	RMSE $\downarrow$
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.

\dagger

marks reproduced DINOv2.

Method	Backbone	$\mathcal{J}\uparrow$	$\mathcal{F}\uparrow$	$\mathcal{J}\&\mathcal{F}\uparrow$
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

\mathcal{J}\&\mathcal{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.

Goal-Conditioned Planning

See all 18 papers

UniJEPA: A Unified Joint-Embedding Predictive Architecture for Task-Agnostic Visual World Modeling
Aug 7, 2026An Lanji, Dawei Liu, Jin Li +3World Model Learning Joint-Embedding Predictive Architecture
University of Electronic Science and Technology of China, Chengdu, China.
Temporal-Distance JEPA: Plan-Aware Representation Learning for Latent World Model Predictive Control
Jul 28, 2026Jiaxin Bai, Jiaxuan XiongWorld Model-Based Planning Joint-Embedding Predictive Architecture
Hong Kong Baptist University Hong Kong, China
Inside-JEPA: From Prediction to Relational Grounding in Latent World Models
Oct 22, 2026Hossein Resani, Javen Qinfeng ShiJoint-Embedding Predictive Architecture Latent World Models
Australian Institute for Machine Learning, Adelaide University · Responsible AI Research Centre

Aug 7, 2026J/K move · Enter open · S save

UniJEPA: A Unified Joint-Embedding Predictive Architecture for Task-Agnostic Visual World Modeling

An Lanji, Dawei Liu, Jin Li, Haoran Xu+2

University of Electronic Science and Technology of China, Chengdu, China.

World Model Learning Joint-Embedding Predictive Architecture Model-Based Planning Predictive Representation Learning Self-Supervised Visual Representation Learning Video World Models

Open paper AskPDF arXiv

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.

Abstract

Goal-Conditioned Planning

Explore similar work

UniJEPA: A Unified Joint-Embedding Predictive Architecture for Task-Agnostic Visual World Modeling

Temporal-Distance JEPA: Plan-Aware Representation Learning for Latent World Model Predictive Control

Inside-JEPA: From Prediction to Relational Grounding in Latent World Models