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

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.

• The paper introduces a subspace Gaussian regularization method that prevents representational collapse in JEPA-based world models.
• It employs frozen, orthogonal random projections to enforce Gaussianity in low-dimensional subspaces, matching intrinsic task geometry.
• Empirical results demonstrate enhanced planning performance and stable latent trajectories across continuous-control benchmarks compared to global priors.

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

Motivation and Problem Formulation

The paper introduces Sub-JEPA, a refinement of Joint-Embedding Predictive Architectures (JEPAs) for learning compact, predictive world models in continuous-control environments. It addresses the bias-variance tradeoff encountered when training JEPA-based latent dynamics models end-to-end: insufficient regularization leads to representational collapse, while excessive constraint (notably a global isotropic Gaussian prior as in LeWorldModel, or LeWM) suppresses representational richness and mismatches task geometry. Latent representations in world modeling tasks typically reside on low-dimensional manifolds, but models often enforce priors in high-dimensional ambient spaces, inducing unnecessary bias.

Sub-JEPA proposes to alleviate this by regularizing latent embeddings through Gaussian priors, not globally but in multiple random low-dimensional orthogonal subspaces. This design preserves the collapse-prevention effect, relaxes the restrictive global bias, and aligns the structural prior with the intrinsic dimensionality of task dynamics. The architecture retains the end-to-end simplicity of LeWM, modifying only the regularization component; all other JEPA elements (encoder, latent predictor, optimizer, training schedule) are conserved.

Methodological Contributions

Subspace Gaussian Regularization

• Orthogonal Subspace Projection: The method constructs $K$ row-orthonormal random projection matrices, each mapping the latent embedding ($\mathbb{R}^D$) to a low-dimensional subspace ($\mathbb{R}^{d_s}$, with $d_s = D/K$). Projections are orthogonally initialized and frozen during training, ensuring statistical independence and geometric isometry.
• Multi-Subspace Regularization: In each subspace, random unit vectors are sampled to form one-dimensional marginals; Gaussianity is enforced via the Epps-Pulley statistic—averaged across directions and subspaces, yielding the regularization loss. The total loss thus combines latent prediction and weighted multi-subspace Gaussian regularization.
• Frozen Projection Justification: Ablation confirms that frozen orthogonal projections outperform randomly initialized or trainable alternatives. Non-orthogonality introduces redundancy; adaptively learned projections co-adapt with the encoder, undermining anti-collapse regularization.

Empirical Protocol

Experiments span four continuous-control benchmarks: Two-Room (2D navigation), Reacher (planar arm), PushT (block manipulation), and OGB-Cube (visually rich 3D task), all trained from raw RGB without rewards. Sub-JEPA is contrasted primarily against LeWM (global Gaussian prior), PLDM (multi-term regularization), and DINO-WM (pretrained visual backbone).

Experimental Results and Analysis

Planning Performance

• Success Rate Superiority: Sub-JEPA consistently surpasses LeWM across all tasks. Pronounced gains appear in environments with low intrinsic dimensionality (e.g., Two-Room), directly validating the hypothesis that subspace regularization aligns latent geometry with task structure. Notably, DINO-WM only outperforms Sub-JEPA in visually complex settings where pretraining confers an advantage.
• Rank Compression Correlation: Effective rank analysis demonstrates that Sub-JEPA compresses latent representations more aggressively — those reductions strongly correlate with planning success uplift. The mismatch between ambient-space Gaussian priors and low-dimensional task structure is mitigated, allowing the latent space to contract toward task-aligned manifolds.

Design Ablations

• Number of Subspaces ($K$): Increasing $K$ (hence lowering $d_s$ per subspace) improves performance for tasks with low intrinsic dimensionality, up to a threshold (e.g., $K=32$). Excessive partitioning (too small $d_s$) undermines the normality estimation reliability, particularly in manipulation-intensive environments. This demonstrates a bias-variance tradeoff: greater $K$ yields flexibility but risks statistical instability.
• Joint Effect of $\mathbb{R}^D$0 and $\mathbb{R}^D$1: Performance is robust within a broad mid-range of subspace configurations; degradation occurs when $\mathbb{R}^D$2 becomes too small. The anti-collapse effect is preserved only when subspaces can reliably sample latent statistics.

Latent Representation Quality

• Physical Probing: Linear and MLP probes trained to decode agent and block state variables from PushT embeddings show that Sub-JEPA's latent space is at least as, or more, recoverable than LeWM's, especially for translational features. For rotational features, Sub-JEPA slightly underperforms with linear probes but matches with nonlinear ones, indicating fragmentation of angular structure across subspaces.
• Temporal Coherence and Path Straightening: Sub-JEPA produces more temporally coherent and straighter latent trajectories, as measured by mean cosine similarity of latent velocities. This geometric regularity emerges without explicit optimization.
• Long-Horizon Stability: In open-loop rollouts, Sub-JEPA demonstrates substantially lower long-term reconstruction error compared to LeWM, maintaining spatial fidelity and resisting drift.

Theoretical and Practical Implications

Sub-JEPA establishes a principled framework for structure-matched regularization in latent world models, leveraging the manifold hypothesis by relaxing unfounded global constraints. The use of subspace-wise Gaussian priors is theoretically justified by the Cramer-Wold theorem; their practical implementation bypasses the curse of dimensionality in high-dimensional embedding spaces, yielding efficient and robust model training.

The design choice of frozen, orthogonal projections is non-trivial—orthogonality preserves geometric balance, and freezing eliminates co-adaptive weakening of regularization. The approach avoids multi-term heuristic regularization, maintains JEPA's end-to-end nature, and does not rely on external pretraining.

Future Directions

Potential extensions include:

• Adaptive selection of projection dimensionality based on task complexity or observed latent rank statistics.
• Integration of reward signals or online data collection for transfer to reinforcement learning.
• Application to more diverse or visually ambiguous environments, combining subspace Gaussian regularization with other forms of structured latent constraints.
• Theoretical investigation into the optimal tradeoff of $\mathbb{R}^D$3, $\mathbb{R}^D$4, and the regularization strength as a function of task manifold properties.

Conclusion

Sub-JEPA advances JEPA-based world modelling by introducing multi-subspace Gaussian regularization, achieving stable end-to-end training and improved planning success, particularly in low-dimensional task regimes. It empirically and theoretically respects the underlying geometry of environment dynamics, enhances the recoverability and coherence of latent representations, and serves as an effective baseline for future world model research (2605.09241).