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

AI summary

Overview

Research area: Artificial intelligence — latent world models, joint-embedding predictive architectures (JEPAs), object-centric representation learning, and visual model-predictive control (planning).

Technical level: Advanced. The paper assumes familiarity with self-supervised visual encoders (DINO/DINOv2, V-JEPA), Slot Attention, attention-based predictors, and optimal transport.

Scope: Slot-JEPA proposes an object-structured "interface" that converts dense patch features from a frozen V-JEPA 2.1 encoder into a small set of stable semantic slots plus time-varying spatial masks, and uses that interface for prediction, memory-based context retrieval, and transport-based goal evaluation in visual planning.

What This Paper Is About

JEPA-style world models represent observations as dense grids of patch features, so any downstream predictor must implicitly rediscover which patches belong to which object and how objects correspond across time. The authors argue this makes prediction, memory reuse, and goal comparison unnecessarily hard, and they replace the dense grid with a reusable object-structured interface: a fixed bank of slots plus per-frame masks that say where each slot's content appears. The goal is a world model that predicts action effects in mask space, reuses compact non-contiguous history, and scores candidate plans with a displacement-aware transport cost rather than pointwise ℓ1/ℓ2 distance.

Key Contributions

1. Slot-JEPA, a decomposition module that factorizes dense patch representations into a shared set of K stable semantic slots and time-varying spatial masks, reconstructing the original feature field. It is pretrained once (on Kinetics-400) and frozen for downstream tasks, exposing object organization as a reusable interface rather than a task-specific byproduct.

2. Transition Memory Predictor (TMP), which stores observed transitions as compact masks in a memory bank, retrieves relevant and diverse (potentially non-contiguous) context via slot-area and centroid statistics with distance-weighted plus uniform sampling, and splits prediction into a slot-wise mask branch and a shared cross-slot feature branch.

3. An Unbalanced Optimal Transport (UOT) objective for both supervising mask prediction and evaluating candidate plans, transporting mass between masks with the same slot index. This replaces pointwise ℓ1/ℓ2 comparisons that become insensitive to displacement when predicted and goal supports are largely disjoint.

4. An action-free pretraining scheme: an inverse transition encoder compresses each observed change into a low-dimensional pseudo-action bottleneck (discarded after pretraining), enabling pretraining on natural video without action labels. The authors also release code.

Main Findings

• Slots preserve semantics under heavy compression. With K = 16 slots, the factorized representation holds only about 2.5% as many scalars as the dense patch field, yet slot-reconstructed features improve over V-JEPA 2.1 on most probes: COCO segmentation mIoU 56.1 vs 52.8, ADE20K 38.9 vs 35.5, PASCAL VOC 75.1 vs 72.1, NYUv2 depth δ1 83.1 vs 80.5, SUNRGBD 72.3 vs 70.0, and DAVIS VOS 66.6 vs 65.9. KITTI depth is the one reported regression (87.8 vs 89.2).

• Slot reconstruction makes patch features more object-aware. Because the slot-reconstruction objective fine-tunes the visual encoder and pushes object-level grouping back into patch features, the authors attribute the downstream gains to object-aware supervision absent from the original V-JEPA 2.1 objective.

• Strong planar planning with visual goals only. On three in-distribution tasks Slot-JEPA reaches 100% (TwoRoom), 100% (Maze) and 92% (PushT); on out-of-distribution variants 90% (TwoRoomR), 96% (MazeR) and 58% (PushObj), over 50 cases per environment. DINO-WM reaches 50/100/74 and 82/86/28; V-JEPA 2.1 reaches 50/94/58 and 50/86/46; LeWM reaches 80/98/88 and 52/88/30. Slot-JEPA also beats every baseline that additionally receives the low-dimensional goal state in all four environments where such results are reported.

• Large gains on robotic arm planning (REALM-JEPA, 108 cases). Slot-JEPA succeeds in 93/108 cases overall (86.1%), versus 38/108 (35.2%) for V-JEPA 2.1 and 63/108 (58.3%) for V-JEPA 2.1 + TMP. Per task: Reach 18/18 (100.0%), Pick 27/30 (90.0%), Rotate 25/30 (83.3%), Place 23/30 (76.7%). V-JEPA 2.1 obtains 20.0% on Rotate and 0.0% on Place. Notably, Slot-JEPA uses only the visual goal, while the patch-based baselines are given ground-truth target-position guidance until the end effector is within 30 cm of the target.

• The planning loss tracks physical progress much better. Spearman correlation of observed goal loss with 7-DoF end-effector–goal distance is 0.792/0.891/0.807 across the two external and one wrist views for Slot-JEPA, versus 0.350/0.407/0.533 for V-JEPA 2.1. Predicted-loss correlations are 0.677/0.788/0.814 versus 0.503/0.453/0.297. The authors note these observed-loss correlations mean predictor error alone cannot explain the baselines' failures.

• Better predictor, not just better features, drives the robotic gains. Swapping V-JEPA 2.1's sequential predictor for TMP while keeping patch representations and the goal loss raises success from 35.2% to 58.3%, consistent with TMP using the full trajectory history rather than temporally contiguous frames.

• Ablation specifics are not reported in the provided content. The paper states that Table 5 compares five configurations (sequential V-JEPA 2.1-style predictor with the mask interface, no UOT at planning, recent-transitions-only memory, random memory transitions, and full Slot-JEPA with Eq. 6 sampling), but the table is truncated and no ablation numbers appear.

Methodology in Plain English

The system sits between a general visual encoder and a task-specific dynamics predictor.

First, a decomposition module looks at a set of reference frames and compresses their patch features into a small fixed bank of "slots" using Slot Attention — competitive attention where slots vie for patches, so different slots end up covering different content. With those slots frozen, any frame is then described by masks: for each slot, a soft map over image patches saying where that content is. Multiplying masks by slots must reconstruct the original patch features, so the decomposition loses as little semantics as possible; an orthogonality penalty keeps slots from duplicating each other. Slots summarize what is in the scene and stay stable over time; masks carry what moves.

Second, the Transition Memory Predictor stores experience as (mask, state, action, next mask, next state) entries instead of patches, making history cheap. At test time it scores stored transitions by comparing summaries of the current state — optional low-dimensional robot state plus each slot's area and 2-D centroid — and samples a mix of close and random examples, since random draws give diversity. The retrieved transitions act as in-context examples for an attention predictor: one branch works slot by slot on spatial mask fields, and another looks at the whole feature grid plus masks to model interactions between slots. Learnable placeholder tokens stand in for the unknown future. The two outputs are combined into predicted next-frame masks.

Third, because actions are not annotated in natural video, pretraining squeezes each observed transition through a narrow pseudo-action bottleneck, so the predictor must use that compressed code rather than peeking at the answer. After pretraining, this bottleneck and the state encoder are thrown away and replaced by task-specific action and state encoders.

Fourth, masks are trained with a mix of UOT loss (cost of moving mass from predicted to target masks on a spatial grid, with a penalty for creating or removing mass, solved by Sinkhorn iterations), plain mask ℓ1, and feature reconstruction ℓ1. At planning time the same UOT and mask terms become the goal cost. Candidate action sequences are rolled out, optimized with the cross-entropy method, executed in chunks, and replanned. All masks — current, predicted, goal — are decoded against the same frozen slots, so slot index k means the same thing everywhere, and obstacles can be folded into the transport ground cost in planar environments.

Why This Matters

Research impact. The paper reframes a representational question ("how good are the features?") as an interface question ("what structure does the predictor get to see?"). It gives a concrete, frozen, task-agnostic recipe for exposing object structure from a pretrained dense encoder, and it diagnoses a specific failure mode of pointwise latent losses — insensitivity to displacement when supports are disjoint — with an objective that fixes it. The observation that simply swapping in a memory-based predictor nearly doubles robotic success (35.2% to 58.3%) is a useful signal that predictor design, not just encoder quality, is a bottleneck for JEPA world models.

Real-world applications:

• Robot manipulation from image goals, where a task is specified by showing the desired end state rather than hand-coded reward or coordinates.
• Warehouse or mobile-robot navigation and pushing tasks where obstacles and object shapes differ from training and a smooth distance-to-goal signal matters.
• Video object segmentation and tracking, since the slot-mask decomposition is evaluated zero-shot on DAVIS-2017 and supports label propagation.
• Perception backbones for dense prediction (segmentation, monocular depth) that need object-aware features without retraining a large encoder.

Industry relevance. Because the encoder and decomposition are frozen downstream, teams can ship one perception stack and fine-tune only a small predictor per task. Storing history as masks rather than patch grids cuts memory cost for long-horizon context, which matters for on-robot deployment, and a goal cost that correlates with real physical progress is directly useful for model-predictive control stacks.

Future Directions

• Slot count and granularity. K is a fixed hyperparameter; the paper notes slots are object-centric components, not guaranteed physical instances, and that a slot may cover multiple entities, a part, or background. Adaptive slot allocation and how to choose K per scene remain open.
• Richer action-conditioned benchmarks beyond planar and tabletop tasks. The evaluation covers six planar environments and four REALM-JEPA arm tasks with one 7-DoF arm; whether the interface holds for deformable objects, multi-agent scenes, or longer horizons is untested.
• Learning the decomposition end-to-end with dynamics instead of freezing it, and understanding when freezing helps versus limits adaptation, especially given the KITTI depth regression.
• Scale and pretraining data. The interface is pretrained on Kinetics-400; how it behaves with much larger video corpora, and whether the pseudo-action bottleneck remains adequate for complex dynamics, is not established. The paper's ablation table is also truncated in the provided text, so the component-level contribution of UOT versus memory sampling versus the two-branch design cannot be quantified from the content available.

Target Audience

Researchers and graduate students working on world models, self-supervised video representation, object-centric learning, or model-based control. It is most useful to readers who already know JEPA-style architectures and Slot Attention and want to understand how to make a pretrained dense encoder usable for planning. Practitioners building image-goal robot policies or memory-based predictors will find the TMP and UOT components directly actionable, especially the finding that a displacement-aware goal cost outperforms pointwise ℓ1 even when the baselines are handed extra ground-truth guidance.

Authors’ 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.