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

2026-05-09

AI Summary

This work addresses the tendency of JEPA world models to collapse into trivial solutions during training due to excessive representational variance, a problem exacerbated by the strong inductive bias introduced by conventional isotropic Gaussian priors in high-dimensional spaces. To mitigate this, the authors propose a subspace Gaussian regularization strategy that replaces global constraints in the full embedding space with Gaussian priors applied independently across multiple random subspaces. This approach effectively prevents representational collapse while enhancing both flexibility and training stability. By achieving a more favorable bias–variance trade-off, the method consistently outperforms LeWorldModel across four continuous control tasks, demonstrating superior performance and markedly improved training robustness.

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.

Problem

• JEPA
• world models
• bias-variance tradeoff
• representation collapse
• Gaussian regularization

Innovation

• Subspace Gaussian Regularization
• Joint-Embedding Predictive Architecture
• World Models
• Bias-Variance Tradeoff
• Latent Representation