Sub-JEPA: Subspace Gaussian Regularization for Stable End-to-End World Models

Hugging Face Daily Papers Papers
world-models jepa regularization end-to-end-learning reinforcement-learning

Summary

The authors introduce Sub-JEPA, a method using Subspace Gaussian Regularization to improve the stability of end-to-end world models like LeWM, showing consistent performance gains on continuous-control benchmarks.

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.
Original Article Export to Word Export to PDF
View Cached Full Text

Cached at: 05/12/26, 07:33 AM

Paper page - Sub-JEPA: Subspace Gaussian Regularization for Stable End-to-End World Models

Source: https://huggingface.co/papers/2605.09241 We’re releasing Sub-JEPA 🌐

LeWM (from LeCun’s group) is the first end-to-end trainable JEPA world model — it uses isotropic Gaussian regularization to prevent representation collapse. Clean and effective.

Our take: latent representations sit on low-dimensional manifolds, so enforcing a full-space Gaussian is too strong a bias.

We propose Subspace Gaussian Regularization: instead of constraining the full embedding space, we project latents into multiple orthogonal subspaces and apply Gaussian constraints there. Simple change, better inductive bias.

Results on 4 continuous-control benchmarks consistently outperform LeWM, with gains correlated to reductions in effective rank — the lower the task’s intrinsic dimensionality, the larger the gain.

Similar Articles

LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels

Papers with Code Trending

LeWorldModel introduces a stable, end-to-end Joint-Embedding Predictive Architecture that trains directly from pixels with minimal hyperparameters and provable anti-collapse guarantees. It achieves significant speedups in planning compared to foundation models while maintaining competitive performance on robotic manipulation tasks.

Learning Visual Feature-Based World Models via Residual Latent Action

Hugging Face Daily Papers

This paper introduces RLA-WM, a visual feature-based world model that leverages residual latent actions and flow matching to efficiently predict future visual states. The method outperforms existing video-diffusion and feature-based approaches while enabling novel robot learning techniques from offline, actionless demonstration videos.