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

Slot-JEPA. Object-Structured Prediction Interfaces for Latent World Models. By Peisong Wen and five others, from the University of Chinese Academy of Sciences, the Institute of Computing Technology, Chinese Academy of Sciences, and two other institutions. Published on the twentieth of October. Its code repository is worked out: about three hundred source files, with a readme and a licence.

Joint-Embedding Predictive Architectures, or JEPAs, are a prototypical latent world model, and they give strong semantic representations. The trouble the authors point to is not the quality of those representations but the interface they offer. Dense patch features hand every downstream predictor a redundant view of the scene, and each predictor has to work out spatial organisation and temporal correspondence for itself, implicitly. That raises three challenges.

First, information about an object is spread across many patches. Every predictor must recover the object structure again before it can begin modelling interactions, even though that structure could have been recovered once and reused across tasks. Second, nothing explicitly tracks which entity is which from one moment to the next. That makes historical transitions expensive to use, and restricts how flexibly a predictor can draw on context. Third, without any explicit transport structure, existing JEPA models typically measure error with pointwise l1 or l2 distances. Those distances are insensitive to spatial displacement when the predicted region and the goal region are largely disjoint: a prediction that is close and one that is far look about equally wrong.

The Slot-based Joint-Embedding Predictive Architecture, Slot-JEPA, decomposes dense representations into two parts: stable semantic slots, and spatial masks that vary over time. The slots say what is there; the masks say where it is at each moment. The decomposition is pretrained once and then frozen downstream, so it becomes a reusable interface organised around objects, which any later predictor can build on.

The first piece built on it is the Transition Memory Predictor, TMP. It stores visual transitions as compact masks, and it retrieves useful context using statistics of the slots that correspond across time. It has two branches, which separate dynamics specific to a single slot from dynamics shared across slots.

The second piece is Unbalanced Optimal Transport, UOT. It supervises mask prediction, and it is also used to evaluate candidate plans, by transporting mass between aligned predicted masks and goal masks. The result is a signal that is better aligned with spatial progress than a pointwise distance, which is exactly what the third challenge called for.

Two kinds of evidence are reported. On representation, the slot-reconstructed features remain transferable across seven visual-understanding benchmarks, so the decomposition does not cost the semantic quality JEPAs are valued for. On control, Slot-JEPA improves goal-conditioned planning across six planar environments and four robotic arm planning tasks on the REALM-JEPA benchmark.

Each of the three complaints meets a matching repair. Objects are recovered once rather than by every predictor, slots give the memory a way to index the same entity across time, and transport replaces a distance that could not tell near misses from far ones.