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

TL;DR

Sub-JEPA introduces a subspace Gaussian regularization technique for JEPA-based world models, improving training stability and representation flexibility by applying Gaussian constraints in multiple random subspaces.

Contribution

It proposes a novel subspace Gaussian regularization method that balances bias and variance, enhancing the stability and performance of JEPA-based world models.

Findings

Outperforms LeWorldModel across four continuous-control environments.

Provides a simple yet effective baseline for future JEPA research.

Balances training stability and representation flexibility effectively.

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.