mlsanityData to 7 Oct 2026

UWM-JEPA: Predictive World Models That Imagine in Belief Space

arXiv 2605.25313 ↗First posted 25 May 2026cs.LG cs.AI cs.ROcited by 3

On the map: World models → Latent predictive world models

Abstract (arXiv)

World models for partially observed environments must imagine multiple compatible hidden futures and steer between them under counterfactual actions. Joint Embedding Predictive Architectures (JEPAs) do this in latent space, but a vector-valued latent has no internal structure for carrying the belief over hidden continuations through blind rollout. We introduce the Unitary World Model JEPA (UWM-JEPA), a JEPA world model with a density-matrix latent on a joint system-environment space and a learned unitary predictor. The construction preserves the joint-state spectrum exactly during rollout, so the predictor itself cannot dissipate the represented uncertainty. On a hidden-velocity indicator task requiring five-step forward simulation under a given action sequence with the target observation masked, UWM-JEPA reaches 0.77 accuracy and degrades monotonically as actions are perturbed; a parameter-matched LSTM-JEPA trained under the same counterfactual-target objective and action head collapses to majority-class accuracy (0.53) under every action condition. Under blind rollout, UWM-JEPA loses fewer than ten points of probe R^2 at short horizons while vector-latent baselines lose forty-one and sixty-eight; both nevertheless tie on a held-out context probe, locating the separation in the predictor rather than the encoder. Action sensitivity itself requires training against counterfactual rather than teacher-forced targets, a finding that applies beyond the unitary parameterisation. For JEPA world models to imagine under partial observability, latent geometry and predictor dynamics matter, not frozen context-encoding capacity alone.

Key facts. Cited by 3 papers in the map (data through 7 Oct 2026). Placed in Latent predictive world models, one of 11 directions in World models. Builds on 26 of its 41 references that are in the map. First posted on arXiv 25 May 2026.

Builds on

26 of 41 references.

Closest ancestors · 2

Same direction, newest first

  1. 2024Revisiting Feature Prediction for Learning Visual Representations from Video15 Feb
  2. 2023Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture19 Jan

Other references · 24

  1. 2017Attention Is All You Needcited by 59,932
  2. 2018Representation Learning with Contrastive Predictive Codingcited by 7,662
  3. 2021Emerging Properties in Self-Supervised Vision Transformerscited by 6,893
  4. 2020Bootstrap your own latent: A new approach to self-supervised Learningcited by 4,712
  5. 2023Mamba: Linear-Time Sequence Modeling with Selective State Spacescited by 4,616
  6. 2014Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translationcited by 4,016
  7. 2018Neural Ordinary Differential Equationscited by 3,733
  8. 2020Exploring Simple Siamese Representation Learningcited by 2,853
  9. 2021Efficiently Modeling Long Sequences with Structured State Spacescited by 2,424
  10. 2021Barlow Twins: Self-Supervised Learning via Redundancy Reductioncited by 1,818
  11. 2019Dream to Control: Learning Behaviors by Latent Imaginationcited by 1,598
  12. 2018World Modelscited by 1,444
12 more
  1. 2016Understanding intermediate layers using linear classifier probescited by 1,279
  2. 2023Mastering Diverse Domains through World Modelscited by 1,193
  3. 2020Mastering Atari with Discrete World Modelscited by 990
  4. 2019Hamiltonian Neural Networkscited by 702
  5. 2015Deep Recurrent Q-Learning for Partially Observable MDPscited by 457
  6. 2015Unitary Evolution Recurrent Neural Networkscited by 291
  7. 2015Embed to Control: A Locally Linear Latent Dynamics Model for Control from Raw Imagescited by 266
  8. 2015Action-Conditional Video Prediction using Deep Networks in Atari Gamescited by 171
  9. 2018Deep Variational Reinforcement Learning for POMDPscited by 122
  10. 2016Full-Capacity Unitary Recurrent Neural Networkscited by 90
  11. 2017QMDP-Net: Deep Learning for Planning under Partial Observabilitycited by 51
  12. 2022A quantum generative model for multi-dimensional time series using Hamiltonian learningcited by 3

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2026: 3.

3 papers.

Similar, not cited

Close in wording, no citation either way

  1. 2026Hamiltonian JEPA: Action-Conditioned World Models with an Inherited Control State
  2. 2026Diagnosing JEPA World Models with Action-Conditioned Predictive Consistency
  3. 2026JEPA-x: Cross-Predictive Physics Grounding for Forecastable Latent Dynamics
  4. 2026LRC-JEPA: Disentangling Dynamics and Residual Context for Efficient World Models
  5. 2026UniJEPA: A Unified Joint-Embedding Predictive Architecture for Task-Agnostic Visual World Modeling
  6. 2026No Gaussian Required: Contrastive Inverse Dynamics for JEPA World Models
  7. 2026JEPA-WAM: Learning Vision-Language-Action Policies with Joint-Embedding World Modeling
  8. 2026Toward Physically Grounded JEPA World Models for Goal-Conditioned Robotic Planning