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

• URL: http://arxiv.org/abs/2605.25313v1
• Date: Mon, 25 May 2026 00:28:51 GMT
• Title: UWM-JEPA: Predictive World Models That Imagine in Belief Space
• Abstract summary: We introduce the Unitary World Model JEPA, a JEPA world model with a density-matrix latent on a joint system-environment space and a learned unitary predictor.<n>The construction preserves the joint-state spectrum exactly during rollout, so the predictor itself cannot dissipate the represented uncertainty.<n>For JEPA world models to imagine under partial observability, latent geometry and predictor dynamics matter, not frozen context-encoding capacity alone.
• Score: 0.2864713389096699
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• Orthogonal JEPA: Factorized Predictive States for Latent World Models [95.20158869896393]
  We introduce method, a latent world-modeling framework based on predictive factorization.<n>We show that method can allocate redundant capacity to dominant signals while providing weak or conflicting gradients to less dominant predictive structure.<n>Experiments on controlled vision, single-cell transcriptomics, longitudinal health records, continuous control, and molecular dynamics evaluate representation quality, forecasting, planning, and long-horizon stability.
  arXiv  Detail & Related papers  (2026-08-20T13:59:57Z)
• UniJEPA: A Unified Joint-Embedding Predictive Architecture for Task-Agnostic Visual World Modeling [24.745715007832953]
  Joint-Embedding Predictive Architectures (JEPAs) have emerged as a principled framework for self-supervised learning of world models in compact latent spaces.<n>We present UniJEPA, a unified JEPA that jointly learns photometric prediction (image-level transformations) and temporal prediction (video-level next-state dynamics) in one shared latent space.
  arXiv  Detail & Related papers  (2026-08-07T16:55:58Z)
• The SIGReg Objective as Variational Free Energy: A Theoretical Active-Inference Account of JEPA World Models [47.471225029744424]
  Joint-Embedding Predictive Architectures (JEPAs) are the dominant design for latent world models.<n>We show that the choice of anti-collapse regulariser determines whether a JEPA's training objective, a prediction loss plus a weighted embedding regulariser, is a valid Active Inference (AIF) variational free energy.
  arXiv  Detail & Related papers  (2026-07-15T08:59:36Z)
• Predictive Objectives Discard Exogenous Control-Relevant Features: A Controlled Mechanistic Study [0.0]
  Joint-embedding predictive (JEPA) objectives learn representations by predicting future latents.<n>This occurs because the objective optimize temporal predictability rather than control-relevance.<n>We show that as little as 2% of reward-labeled transitions recovers the feature, the effect holds across two environments.
  arXiv  Detail & Related papers  (2026-06-29T10:00:50Z)
• Sensorimotor World Models: Perception for Action via Inverse Dynamics [59.121736411156384]
  We introduce a sensorimotor world model trained end-to-end with inverse dynamics regularization.<n>It prevents representation collapse and induces action-aligned representations.<n>It learns compact, interpretable latent spaces and enables competitive planning performance across simple 2D and 3D control tasks.
  arXiv  Detail & Related papers  (2026-06-18T11:25:16Z)
• Subspace-Decomposed JEPAs: Disentangling Progression and Content in Latent World Models [1.034052616244602]
  Joint-Embedding Predictive Architectures learn compact latent world models by predicting future embeddings.<n>We carve the JEPA latent into two subspaces with disjoint roles.<n>We prove that the two anti-collapse forces act on disjoint coordinates, so they compose additively rather than competing on the same dimensions.
  arXiv  Detail & Related papers  (2026-05-29T10:23:48Z)
• Entity-Centric World Models: Interaction-Aware Masking for Causal Video Prediction [0.0]
  Joint Embedding Predictive Architectures (JEPA) have set new benchmarks in semantic classification.<n>We propose Interaction-Aware JEPA (IA-JEPA), which utilizes a self-supervised motion-centric masking strategy.<n>IA-JEPA achieves 14.26% accuracy on causal reasoning tasks, a significant lead over the 3.22% achieved by standard patch-masked baselines.
  arXiv  Detail & Related papers  (2026-05-14T23:10:04Z)
• Sub-JEPA: Subspace Gaussian Regularization for Stable End-to-End World Models [26.928611870426383]
  JointEmbedding Predictive Architectures (JEPAs) provide a simple framework for learning world models by predicting future latent representations.<n>The recent LeWorldModel (LeWM) shows that this issue can be alleviated bysimply constraining latent embeddings with an isotropic Gaussian prior.<n>We propose ame, which seeks a favorable operatingpoint on the bias-variance frontier by applying Gaussian constraints inmultiple random subspaces rather than in the originalembedding space.
  arXiv  Detail & Related papers  (2026-05-10T00:51:47Z)
• Probing the Latent World: Emergent Discrete Symbols and Physical Structure in Latent Representations [0.0]
  Video world models trained with Joint Embedding Predictive Architectures (JEPA) acquire rich representations by predicting masked regions in latent space rather than reconstructing pixels.<n>This pathway removes the visual verification of generative models, creating a structural interpretability gap: the encoder has learned physical structure inaccessible in any inspectable form.<n>We propose the AI Mother Tongue (AIM) framework as a passive quantization probe: a lightweight, vocabulary-free probe that converts V-JEPA 2 continuous latent vectors into discrete symbol sequences without task-specific supervision or modifying the encoder.
  arXiv  Detail & Related papers  (2026-03-20T01:30:59Z)
• LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels [49.35636088613484]
  Joint Embedding Predictive Architectures (JEPAs) offer a compelling framework for learning world models in compact latent spaces.<n>We introduce LeModelWorld, the first JEPA that trains stably end-to-end from raw pixels using only two loss terms.<n>With 15M parameters trainable on a single GPU in a few hours, LeWM plans up to 48x faster than foundation-model-based world models.
  arXiv  Detail & Related papers  (2026-03-13T19:48:14Z)
• Causal-JEPA: Learning World Models through Object-Level Latent Interventions [46.562961546550895]
  C-JEPA is a simple and flexible object-centric world model that extends masked joint embedding prediction from image patches to object-centric representations.<n>By applying object-level masking that requires an object's state to be inferred from other objects, C-JEPA induces latent interventions with counterfactual-like effects.
  arXiv  Detail & Related papers  (2026-02-11T21:47:26Z)
• VJEPA: Variational Joint Embedding Predictive Architectures as Probabilistic World Models [0.0]
  We introduce emphVariational JEPA (VJEPA), a textitprobabilistic generalization that learns a predictive distribution over future latent states via a variational objective.<n>VJEPA representations can serve as sufficient information states for optimal control without pixel reconstruction, while providing formal guarantees for collapse avoidance.<n>We propose emphBayesian JEPA (BJEPA), an extension that factorizes the predictive belief into a learned dynamics expert and a modular prior expert.
  arXiv  Detail & Related papers  (2026-01-20T18:04:16Z)
• LeJEPA: Provable and Scalable Self-Supervised Learning Without the Heuristics [53.247652209132376]
  Joint-Embedding Predictive Architectures (JEPAs) offer a promising blueprint, but lack of practical guidance and theory has led to ad-hoc R&D.<n>We present a comprehensive theory of JEPAs and instantiate it in bf LeJEPA, a lean, scalable, and theoretically grounded training objective.
  arXiv  Detail & Related papers  (2025-11-11T18:21:55Z)
• TD-JEPA: Latent-predictive Representations for Zero-Shot Reinforcement Learning [63.73629127832652]
  We introduce TD-JEPA, which leverages TD-based latent-predictive representations into unsupervised RL.<n> TD-JEPA trains explicit state and task encoders, a policy-conditioned multi-step predictor, and a set of parameterized policies directly in latent space.<n> Empirically, TD-JEPA matches or outperforms state-of-the-art baselines on locomotion, navigation, and manipulation tasks across 13 datasets.
  arXiv  Detail & Related papers  (2025-10-01T10:21:18Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.