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

arXiv: https://arxiv.org/abs/2605.25313 · HTML: https://mlsanity.com/paper/2605.25313/

## Key facts

- Cited by 3 papers in the map (data through 14 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.

## Placed in

- World models / [Latent predictive world models](https://mlsanity.com/wm/latent-predictive-world-models/)

## Builds on (most cited first, 26 of 26)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Emerging Properties in Self-Supervised Vision Transformers](https://mlsanity.com/paper/2104.14294/) (2021)
- [Bootstrap your own latent: A new approach to self-supervised Learning](https://mlsanity.com/paper/2006.07733/) (2020)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation](https://mlsanity.com/paper/1406.1078/) (2014)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Exploring Simple Siamese Representation Learning](https://mlsanity.com/paper/2011.10566/) (2020)
- [Efficiently Modeling Long Sequences with Structured State Spaces](https://mlsanity.com/paper/2111.00396/) (2021)
- [Barlow Twins: Self-Supervised Learning via Redundancy Reduction](https://mlsanity.com/paper/2103.03230/) (2021)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [Understanding intermediate layers using linear classifier probes](https://mlsanity.com/paper/1610.01644/) (2016)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture](https://mlsanity.com/paper/2301.08243/) (2023)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [Hamiltonian Neural Networks](https://mlsanity.com/paper/1906.01563/) (2019)
- [Revisiting Feature Prediction for Learning Visual Representations from Video](https://mlsanity.com/paper/2404.08471/) (2024)
- [Deep Recurrent Q-Learning for Partially Observable MDPs](https://mlsanity.com/paper/1507.06527/) (2015)
- [Unitary Evolution Recurrent Neural Networks](https://mlsanity.com/paper/1511.06464/) (2015)
- [Embed to Control: A Locally Linear Latent Dynamics Model for Control from Raw Images](https://mlsanity.com/paper/1506.07365/) (2015)
- [Action-Conditional Video Prediction using Deep Networks in Atari Games](https://mlsanity.com/paper/1507.08750/) (2015)
- [Deep Variational Reinforcement Learning for POMDPs](https://mlsanity.com/paper/1806.02426/) (2018)
- [Full-Capacity Unitary Recurrent Neural Networks](https://mlsanity.com/paper/1611.00035/) (2016)
- [QMDP-Net: Deep Learning for Planning under Partial Observability](https://mlsanity.com/paper/1703.06692/) (2017)
- [A quantum generative model for multi-dimensional time series using Hamiltonian learning](https://mlsanity.com/paper/2204.06150/) (2022)

## Built on by (most cited first, 3 of 3)

- [Branch-JEPA: Finite-Support Predictive Distributions for JEPA World Models](https://mlsanity.com/paper/2607.05238/) (2026)
- [Optimizing ML Workload Partitioning between CPUs and CIM Accelerators for Heterogeneous Computing](https://mlsanity.com/paper/2607.05240/) (2026)
- [From Generation to Simulation: How Far Are World Models from Being True Simulators?](https://mlsanity.com/paper/2608.23070/) (2026)

## 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.