Sub-JEPA: Subspace Gaussian Regularization for Stable End-to-End World Models
- URL: http://arxiv.org/abs/2605.09241v1
- Date: Sun, 10 May 2026 00:51:47 GMT
- Title: Sub-JEPA: Subspace Gaussian Regularization for Stable End-to-End World Models
- Abstract summary: 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.
- Score: 26.928611870426383
- License: http://creativecommons.org/licenses/by/4.0/
- Abstract: Joint-Embedding Predictive Architectures (JEPAs) provide a simpleframework for learning world models by predicting future latent representations.However, JEPA training is subject to a bias-variance tradeoff.Without sufficient structural constraints, excessive representationalvariance causes the model to collapse to trivial solutions.The recent LeWorldModel (LeWM) shows that this issue can be alleviated bysimply constraining latent embeddings with an isotropic Gaussian prior.However, latent representations inherently lie on low-dimensional manifoldswithin a high-dimensional ambient space, and enforcing an isotropic Gaussianprior directly in this ambient space introduces an overly strong bias.In this work, 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.This design relaxes the global constraint while preserving itsanti-collapse effect, leading to a better balance between trainingstability and representation flexibility.Extensive experiments across fourcontinuous-control environments demonstrate that consistentlyoutperforms LeWM with very clear margins.Our method is simple yet effective, and serves as a strong baseline for future JEPA-based world model research.fdefinedeeemodeThe code is available at https://github.com/intcomp/Sub-JEPA.
Related papers
- Flow-JEPA: Flow Matching for Robust Latent Dynamics in JEPA World Models [18.035166689862205]
Joint-Embedding Predictive Architectures (JEPAs) have shown strong potential for learning compact predictive representations.<n>We propose Flow-JEPA, a conditional flow matching dynamics model that jointly generates a sequence of future latent states conditioned on the current observation and actions.
arXiv Detail & Related papers (2026-08-29T03:40:06Z) - 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) - Learning Spatio-Temporal Foundation Models from Pure Synthetic Data [85.23980441738811]
NeoST is a scalable synthetic-training corpus to real-world bias, a latent-space reasoning architecture that generates and iteratively refines multiple future trajectories, and latent-space objectives that emphasize structural dynamics and enable inference-time correction under distribution shifts.<n>Extensive experiments across diverse real-world benchmarks show that NeoST consistently outperforms existing STFMs in diverse real-worldtemporal systems.
arXiv Detail & Related papers (2026-06-27T02:56:46Z) - UWM-JEPA: Predictive World Models That Imagine in Belief Space [0.2864713389096699]
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.
arXiv Detail & Related papers (2026-05-25T00:28:51Z) - Language Generation as Optimal Control: Closed-Loop Diffusion in Latent Control Space [99.30250506439678]
This work reformulates language generation as an optimal control problem.<n>It provides a unified theoretical perspective to analyze autoregressive and diffusion models.
arXiv Detail & Related papers (2026-05-14T08:13:43Z) - ThinkJEPA: Empowering Latent World Models with Large Vision-Language Reasoning Model [53.15040805435013]
Vision-language models (VLMs) provide strong semantic grounding and general knowledge by reasoning over uniformly sampled frames.<n>We propose a VLM-guided JEPA-style latent world modeling framework that combines dense-frame dynamics modeling with long-horizon semantic guidance.
arXiv Detail & Related papers (2026-03-23T17:59:42Z) - 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) - 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) - Latent Iterative Refinement Flow: A Geometric-Constrained Approach for Few-Shot Generation [5.062604189239418]
We introduce Latent Iterative Refinement Flow (LIRF), a novel approach to few-shot generation.<n>LIRF establishes a stable latent space using an autoencoder trained with our novel textbfmanifold-preservation loss.<n>Within this cycle, candidate samples are refined by a geometric textbfcorrection operator, a provably contractive mapping.
arXiv Detail & Related papers (2025-09-24T08:57:21Z) - Multi-Agent Path Finding in Continuous Spaces with Projected Diffusion Models [57.45019514036948]
Multi-Agent Path Finding (MAPF) is a fundamental problem in robotics.<n>This work proposes a novel approach that integrates constrained optimization with diffusion models for MAPF in continuous spaces.
arXiv Detail & Related papers (2024-12-23T21:27:19Z) - Denoising with a Joint-Embedding Predictive Architecture [21.42513407755273]
We introduce Denoising with a Joint-Embedding Predictive Architecture (D-JEPA)<n>By recognizing JEPA as a form of masked image modeling, we reinterpret it as a generalized next-token prediction strategy.<n>We also incorporate diffusion loss to model the per-token probability distribution, enabling data generation in a continuous space.
arXiv Detail & Related papers (2024-10-02T05:57:10Z) - To Regularize or Not To Regularize? The Bias Variance Trade-off in
Regularized AEs [10.611727286504994]
We study the effect of the latent prior on the generation deterministic quality of AE models.
We show that our model, called FlexAE, is the new state-of-the-art for the AE based generative models.
arXiv Detail & Related papers (2020-06-10T14:00:14Z)
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.