PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2610.28074v1cs.AIcs.CV

TL;DR

Dense patch representations leave object organization and temporal correspondence implicit, complicating downstream prediction and planning. Slot-JEPA factorizes them into stable semantic slots and spatial masks, adds transition memory, and uses UOT for mask prediction and goal evaluation. The paper reports transferable slot-reconstructed features across visual-understanding benchmarks and improved planning across planar and robotic-arm tasks.

  • Problem

    Dense patch representations leave object structure and temporal correspondence implicit, while pointwise distances can be insensitive to displacement when predicted and goal supports are largely disjoint.

    Source: §1 · p. 2

  • Method

    Slot-JEPA factorizes dense features into stable semantic slots and time-varying masks, uses TMP to retrieve transitions, and applies UOT to aligned masks for prediction and goal evaluation.

    Source: §1 · p. 2

  • Results

    Across seven visual-understanding benchmarks, slot-reconstructed features remain transferable, and Slot-JEPA improves planning across six planar environments and four REALM-JEPA robotic-arm tasks.

    Source: §1 · p. 2, §6 · p. 11

  • Takeaways & Limitations

    The reported results show the value of combining structured mask prediction, spatially informed goal costs, and transition memory.

    Source: §6 · p. 11

    Slot decomposition does not guarantee instance- or part-level masks, and fixed state and mask cues may miss contact or appearance changes or retrieve redundant examples.

    Source: §5 · p. 11

Abstract

from arXiv · show

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.

1 INTRODUCTION

Dense patch representations leave object structure and temporal correspondence implicit, making downstream prediction, memory, and goal evaluation less direct. Slot-JEPA exposes stable semantic slots and time-varying masks, then uses transition memory and transport-based planning; reported evaluations span visual understanding and goal-conditioned planning.

1 INTRODUCTION

  • Slot-JEPA factorizes dense features into stable semantic slots and time-varying spatial masks, exposing reusable object organization while preserving the original feature field through reconstruction.The decomposition is pretrained once and frozen for downstream tasks.Source: §1 · p. 2
  • TMP stores visual history as compact masks and retrieves relevant, diverse, potentially non-contiguous transitions using statistics of masks associated with corresponding slots.Its two branches separate slot-specific dynamics from inter-slot interaction dynamics.Source: §1 · p. 2
  • UOT compares aligned predicted and goal masks by transporting mass, providing a smooth planning signal over long displacements where pointwise losses can become insensitive.The contribution also uses mask-space action-effect prediction for goal evaluation.Source: §1 · p. 2
  • Across seven visual-understanding benchmarks, slot-reconstructed features retain most semantic information, while Slot-JEPA improves planning across six planar tasks and four REALM-JEPA robotic-arm tasks.The reported planning gains are especially pronounced on more complex manipulation tasks.Source: §1 · p. 2

2 RELATED WORK

The related work connects joint-embedding prediction, object-centric representations, memory-based video segmentation, and optimal transport. Slot-JEPA is positioned as a structured interface and spatial objective that brings these threads to action-conditioned prediction and planning.

2 RELATED WORK

  • The related work frames Slot-JEPA as connecting four lines: joint-embedding prediction, object-centric decomposition and dynamics, memory-based video segmentation, and optimal transport.Together, these motivate its slot-and-mask state, transition memory, and spatial planning objective.Source: §2 · p. 2

2.1 JOINT-EMBEDDING PREDICTIVE WORLD MODELS

  • Representation-learning advances and action-conditioned JEPA world models form distinct strands; Slot-JEPA complements them with a slot-and-mask interface and spatial objective.The interface sits between a pretrained dense encoder and action dynamics rather than replacing joint-embedding pretraining.Source: §2.1 · p. 3

2.2 OBJECT-CENTRIC REPRESENTATION AND DYNAMICS

  • Object-centric methods include feature-reconstruction approaches, while Slot-JEPA keeps a reference-derived semantic slot bank fixed and predicts spatial masks rather than updating each slot's full latent content.Its slots are not guaranteed to correspond to physical instances.Source: §2.2 · p. 3

2.3 MEMORY-BASED VOS AND TRANSITION MEMORY

  • TMP transfers memory from perceptual propagation to controlled prediction: it stores observed source–action–target transitions and retrieves potentially non-contiguous examples for a new action outcome.This differs from VOS memory, which stores evidence for locating a target.Source: §2.3 · pp. 3–4

2.4 UNBALANCED OPTIMAL TRANSPORT FOR VISUAL CORRESPONDENCE AND CONTROL

  • UOT charges for spatial displacement while relaxing exact mass matching; Slot-JEPA applies it between same-index masks for training and displacement-aware goal costs.Its ground cost can also incorporate obstacle geometry in planar environments.Source: §2.4 · p. 4

3 METHODOLOGY

Slot-JEPA uses a frozen, reference-derived slot interface to organize visual representations for action-conditioned world modeling. TMP reuses mask-based transition context, while UOT supports mask prediction and goal evaluation.

3.1 OVERVIEW

  • Slot-JEPA inserts a reusable object-structured interface between a visual encoder and task-specific predictors, then freezes that interface while transferring TMP downstream.The encoder and slot decomposition module are frozen after pretraining; TMP is fine-tuned for downstream control tasks.Source: §3.1 · p. 4

3.2 PRETRAINED SLOT DECOMPOSITION MODULE

  • Reference-derived slots summarize content shared across frames, while time-varying masks locate each slot’s content and factorize dense features more compactly.The slots provide a fixed latent basis within a planning context; the masks encode temporal variation relative to that basis.Source: §3.2 · pp. 4–5
  • Cross-attention and slot-wise softmax assign patches to slots, encouraging object-aligned grouping without requiring one slot per physical instance.A slot may cover multiple entities, a part, or background; increasing K permits finer decomposition but increases prediction cost.Source: §3.2 · p. 5

3.3 TRANSITION MEMORY PREDICTOR

  • TMP stores environment transitions as masks and retrieves context using slot-mask area and centroid statistics, mixing distance-weighted and uniform sampling for relevance and diversity.At inference, it samples relevant and diverse transitions, which need not be contiguous in time.Source: §3.3 · pp. 5–6
  • TMP processes context and query in symmetric blocks, with a slot-wise mask branch for spatial changes and a shared feature branch for interactions among slots.The query uses learnable tokens for unknown future inputs, and the combined branches produce a future assignment map.Source: §3.3 · p. 6

3.4 TRANSPORT-BASED OBJECTIVE FOR MASK PREDICTION

  • Because pointwise L1 can become insensitive when predicted and target supports are largely disjoint, mask supervision combines UOT transport cost with mask and feature reconstruction losses.At test time, the first two loss terms become the planning objective using goal masks decoded against the same slots as the current observation.Source: §3.4 · pp. 6–7
  • At test time, aligned current and goal slots support mask-based goal evaluation; CEM optimizes candidate action sequences before an action chunk is executed and replanning begins.The goal uses the same slot set, keeping slot indices aligned across current, predicted, and goal masks.Source: §3.4 · p. 7

4 EXPERIMENTS

Slot-JEPA preserves task-relevant semantics under strong compression and improves goal-conditioned planning across planar and robotic-arm tasks. The experiments also test transfer, memory, and loss components across five subsections.

4.1 EXPERIMENTAL SETUP

  • The experiments assess semantic transfer, planar visual planning, robotic-arm planning, and component contributions; the slot module and visual encoder are frozen for downstream tasks.Semantic evaluation covers segmentation, depth estimation—including NYUv2-to-SUNRGBD transfer without adaptation—and zero-shot VOS.Source: §4.1 · p. 7, §4.2 · p. 7

4.2 PRESERVATION OF SEMANTIC INFORMATION

  • With K = 16, slot representations use about 2.5% as many scalars as the dense patch field while retaining task-relevant semantics across segmentation, depth, and VOS.Slot reconstruction also propagates object-level grouping into patch features, supplying object-aware supervision absent from the original V-JEPA 2.1 objective and helping explain consistent downstream gains.Source: §4.2 · pp. 7–8

4.3 PLANAR VISUAL PLANNING

  • Using only visual goals, Slot-JEPA succeeds in 100%, 100%, and 92% of ID cases and 90%, 96%, and 58% of OOD cases, outperforming every baseline given goal-state input in four available environments.The authors report smoother, obstacle-aware UOT loss landscapes than ℓ1 losses, including far from the goal.Source: Table 2 · p. 8, §4.3 · p. 8, Figure 2 · p. 8

4.4 ROBOTIC ARM PLANNING

  • 86.1% overall success for Slot-JEPA, versus 35.2% for V-JEPA 2.1 and 58.3% for V-JEPA 2.1 + TMP, across four robotic-arm task families.Success was especially higher on Rotate and Place: 83.3% and 76.7%, versus 20.0% and 0.0% for V-JEPA 2.1.Source: Table 3 · p. 9, §4.4 · p. 10

4.5 ABLATION STUDIES

  • Ablations show UOT improves success in every planar environment and robot Reach; TMP's diverse historical context helps on PushT and PushObj, while query-conditioned relevance matters especially under layout shift.Recent transitions perform comparably to TMP on the simpler TwoRoom and Maze tasks; broader history and relevance balancing benefit more complex settings.Source: Table 5 · p. 10, §4.5 · pp. 10–11

5 LIMITATIONS

Slot decomposition does not guarantee instance- or part-level masks, and TMP retrieval cues may miss relevant changes or return redundant examples.

5 LIMITATIONS

  • Without instance or part annotations, slots may merge objects or split one, limiting reasoning about specific components and contact points.Fixed state and mask cues may miss contact or appearance changes and select redundant transitions; balancing retrieval relevance with diversity remains open.Source: §5 · p. 11

6 CONCLUSION

Slot-JEPA combines object-structured mask prediction, spatially informed goal costs, and transition memory. It remains transferable across seven visual-understanding benchmarks and improves planning across six planar environments and four robotic-arm tasks.

6 CONCLUSION

  • Across seven visual-understanding benchmarks, mixtures of encoder and slot-reconstructed features remain transferable, while planning success improves across six planar environments and four REALM-JEPA robotic-arm tasks.Improvements are especially pronounced on the more complex manipulation tasks.Source: §6 · p. 11
  • Slot-JEPA factorizes dense visual representations into stable object-structured slots and time-varying spatial masks for action-conditioned prediction.The conclusion identifies structured mask prediction, spatially informed goal costs, and transition memory as the combination supporting visual planning.Source: §6 · p. 11

A Additional Method Details

The appendix details the shared slot-decomposition architecture and procedures used in pretraining and downstream tasks. The decomposition extracts slots from context frames, decodes masks, and reconstructs dense features.

A Additional Method Details

  • The shared appendix procedures cover pretraining and downstream tasks, while environment-specific initialization, UOT settings, and planning procedures are reported separately.The appendix organizes the empirical study into representation pretraining, visual understanding, planar planning, robotic-arm planning, and visualizations.Source: A ADDITIONAL METHOD DETAILS · p. 17

A.1 COMPLETE SLOT DECOMPOSITION

  • Slots are initialized from context-frame patch features in non-overlapping 4 × 4 windows, using four evenly spaced valid patch tokens per window.Target frames are not used to construct the slots.Source: §A.1 · p. 17
  • Three local Slot Attention iterations feed pooled local slots, from which farthest-point sampling selects K = 16 candidates for three global iterations that produce the shared slot bank.The complete extraction and mask-decoding network is shown in Fig. 7.Source: §A.1 · p. 17
  • The decoder assigns each patch to slots with a softmax over assignment logits, then uses the mask-weighted slot sum to reconstruct the dense representation.Local and global Slot Attention modules and the decoder are jointly optimized with feature-reconstruction and slot-orthogonality terms; the same slot bank decodes observed, predicted, and goal frames within an environment.Source: §A.1 · p. 17

A.2 PREDICTOR ARCHITECTURE, PRETRAINING, AND MEMORY

TMP predicts future masks through object-wise and shared feature branches, with observed transitions supplying context. The appendix also describes memory retrieval, unbalanced transport, and change-based static-slot identification.

A.2 PREDICTOR ARCHITECTURE, PRETRAINING, AND MEMORY

  • TMP fuses three-scale object-wise mask predictions with a shared feature branch that models interactions across slots.Its reconstructed-feature branch uses 12 blocks; the mask branch uses separate predictors at three spatial scales, as illustrated in Fig. 8.Source: §A.2 · p. 17, Figure 8 · p. 19
  • Pretraining uses both forward and backward contexts from each seven-frame clip, with four observed frames predicting the other three and no persistent memory bank.Because natural videos lack action labels, inverse transition encoders compress paired frame reconstructions and masks into pseudo-actions and pseudo-states; these encoders are discarded downstream.Source: §A.2 · pp. 17–18
  • At inference, TMP retrieves five observed transitions using normalized state, per-slot area, and per-slot centroid statistics, sharing the sampled context across candidate actions.When slots are refreshed, retained frames are decoded again so memory masks remain slot-index consistent.Source: §A.2 · p. 18

A.3 UOT AND STATIC-SLOT IDENTIFICATION

  • UOT is computed independently for corresponding slot masks while retaining their original mass, with relaxed marginals allowing mass to disappear or emerge at finite cost.The implementation evaluates spatial transport and relaxed-marginal terms after generalized Sinkhorn scaling.Source: §A.3 · p. 19
  • Static slots are identified from predicted mask changes using separate static and moving thresholds, leaving intermediate changes unclassified.The reported examples show moving masks following controlled entities and static masks aligning with persistent scene geometry; labels come from the change-score rule, not ground-truth segmentation.Source: §A.3 · p. 20
  • For obstacle-aware planning, selected static masks form an obstacle map that adds obstacle occupancy along a path to the ground cost.The obstacle-aware ground-cost variant and thresholds are environment-specific.Source: §A.3 · p. 20

B EXPERIMENTAL DETAILS AND COMPLETE RESULTS

The appendix outlines the study’s experimental stages and gives pretraining data and optimization details. Slot-JEPA is pretrained on Kinetics-400, then its encoder and slot decomposition module are frozen for downstream tasks.

B EXPERIMENTAL DETAILS AND COMPLETE RESULTS

  • The appendix organizes experimental details across representation pretraining, visual understanding, planar planning, robotic-arm planning, and additional visualizations.The visualizations compare representations, planning objectives, and predicted rollouts.Source: B EXPERIMENTAL DETAILS AND COMPLETE… · p. 20

B.1 PRETRAINING

  • Pretraining uses the Kinetics-400 training split of roughly 240,000 clips, starting from a released distilled V-JEPA 2.1 ViT-B/16 checkpoint.The complete dataset has 306,245 clips across 400 human-action classes; videos are decoded at 384 × 384 resolution.Source: §B.1 · pp. 20–21
  • Figure 9 illustrates moving- and static-slot labels derived from predicted changes rather than ground-truth segmentation.Moving masks follow controlled entities, while static masks remain aligned with walls or maze geometry.Source: Figure 9 · p. 21
  • Pretraining jointly optimizes slot decomposition, TMP, and inverse transition encoders while fine-tuning the final six visual-encoder blocks; the encoder and slot module are frozen downstream.Training runs for 10 epochs on eight RTX 4090 GPUs and takes about 30 hours.Source: §B.1 · p. 21

B.2 VISUAL UNDERSTANDING

Slot-JEPA’s frozen slot decomposition retains transferable visual features, while the planar planning setup uses transition memories and UOT-based candidate evaluation across in-distribution and out-of-distribution settings.

B.2 VISUAL UNDERSTANDING

  • Slot-JEPA improves segmentation, NYUv2 depth and transfer, and DAVIS J&F over its V-JEPA 2.1 initialization; segmentation mIoU gains are 3.3%, 3.4%, and 3.0% on COCOStuff-27, ADE20K, and PASCAL VOC.It achieves the strongest results across six segmentation metrics among comparison models excluding gray reference rows; DAVIS J&F rises from 65.9% to 66.6%, the strongest comparison entry. KITTI performance decreases slightly.Source: Table C1 · p. 22, §B.2 · p. 22, Table C2 · p. 23
  • NYUv2 δ1 rises from 80.5% to 83.1%, while AbsRel falls from 0.139 to 0.130 and RMSE from 0.486 to 0.456.All three NYUv2-to-SUNRGBD transfer metrics improve, while KITTI performance decreases slightly.Source: §B.2 · p. 22, Table C2 · p. 23

B.3 PLANAR PLANNING

  • The planar evaluation pairs six environments into in-distribution and out-of-distribution settings; PushObj tests held-out L and plus shapes after training on I, Z, small-T, and square shapes.Object poses vary in both splits, so shape identity defines the PushObj distinction.Source: §B.3 · p. 25
  • For planar out-of-distribution cases, random-action transitions from the new layout or shape initialize memory before goal-directed evaluation, and observed planning transitions are appended afterward.A fixed slot bank is extracted from a small subset of training trajectories; in-distribution memory instead starts with randomly sampled training transitions.Source: §B.3 · p. 25
  • CEM ranks candidate action sequences using TMP rollouts and UOT plus mask-ℓ1 goal loss, with obstacle-aware transport costs for TwoRoom, TwoRoomR, Maze, and MazeR.PushT and PushObj use the standard Euclidean ground cost; the sampled memory contexts stay fixed across candidates at each MPC step.Source: §B.3 · p. 25

B.4 ROBOTIC ARM PLANNING

REALM-JEPA evaluates visual-goal planning across four robotic-arm task families, with synchronized goals and per-view memories. Slot-JEPA gives the strongest reported agreement statistics across comparisons and camera views.

B.4 ROBOTIC ARM PLANNING

  • REALM-JEPA contains 108 visual-goal planning cases across Reach, Pick, Rotate, and Place, with synchronized goals from three cameras.Reach uses one visual goal; Pick, Rotate, and Place use two successive visual goals, and reference actions and intermediate states are unavailable to the planner.Source: §B.4 · p. 26
  • Patch-based V-JEPA 2.1 baselines receive ground-truth target-position guidance for Place and Rotate when positional distance exceeds 30 cm, but Slot-JEPA does not.The baseline proposal bias is disabled at 30 cm or closer, where planning uses visual goal loss alone.Source: §B.4 · pp. 26–27
  • The three views share predictor parameters but maintain independent slot sets, mask histories, and transition memories; fixed cameras guide CEM before wrist-view refinement.Candidate robot actions are transformed into each camera coordinate system, without cross-view attention or a joint TMP memory bank.Source: §B.4 · p. 27
  • Slot-JEPA gives the strongest Kendall’s τ, Spearman’s ρ, standardized slope β, and R2 for all three loss-distance comparisons and every camera view.The statistics are macro-averaged across tasks after within-task standardization and distance-stratified sampling.Source: §B.4 · p. 27, Table C5 · p. 28

B.5 ADDITIONAL VISUALIZATIONS

The visualizations show interpretable slot masks and goal-loss landscapes: slots track distinct scene regions over time, and obstacle-aware UOT assigns lower loss to feasible progress than to visually nearby obstructed states.

B.5 ADDITIONAL VISUALIZATIONS

  • Obstacle-aware UOT loss forms low-loss regions along feasible progress toward the goal and assigns less favorable energy to visually nearby obstructed states.Figure 12 visualizes energy geometry over independently rendered sampled states, not executed rollouts; stars mark goals and lighter colors indicate lower loss.Source: §B.5 · p. 28, Figure 12 · p. 29
  • Predicted slot masks track task-relevant regions across environments and time, making expected object or region movement inspectable alongside executed RGB observations.Figure 13 pairs each executed observation with masks predicted for that time step; the Place panel overlays sharpened predictions on wrist-view RGB.Source: §B.5 · p. 28, Figure 13 · p. 30
  • Across Kinetics-400 and DROID examples, different slots emphasize distinct objects or spatial regions and remain recognizable over time.Figures 14–17 overlay selected soft slot masks on the input RGB frames.Source: §B.5 · p. 28, Figure 14 · p. 31, Figure 15 · p. 32
  • Together, the visualizations present a spatially structured state that TMP uses for coherent future-mask prediction and UOT uses for goal-directed costs.Source: §B.5 · p. 29

Related papers

Loading 2610.28074v1…