PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Kai Zhao, Dongliang Nie, Yuchen Lin, Zhehan Luo, Yixiao Gu, Deng-Ping Fan, Dan Zeng

arXiv:2605.09241v1cs.LGcs.AI

TL;DR

JEPA world models must balance collapse prevention against preserving rich latent representations, while ambient-space Gaussian regularization can impose excessive bias on low-dimensional dynamics. Sub-JEPA applies Gaussian regularization across multiple random low-dimensional subspaces and consistently outperforms LeWM across four continuous-control environments.

  • Problem

    JEPA training lacks a simple, well-guided way to balance representation-collapse prevention with preservation of rich latent representations under low-dimensional dynamics.

  • Method

    Sub-JEPA replaces full-space Gaussian regularization with Gaussian constraints over multiple frozen row-orthonormal random projections.

  • Results

    Sub-JEPA consistently outperforms LeWM by clear margins across four continuous-control environments.

  • Takeaways & Limitations

    Subspace regularization offers a simple JEPA world-model baseline that better matches the intrinsic dimensionality of the underlying dynamics.

  • Takeaways & Limitations

    The model is evaluated in a fully offline, reward-free setting using pre-collected observation-action trajectories without rewards or task labels.

Abstract

from arXiv · show

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.

1 Introduction

JEPA world models enable latent-space planning but face a bias-variance tradeoff in stable end-to-end training: weak constraints can cause collapse, while strong constraints limit representation richness. Sub-JEPA moves Gaussian regularization into multiple low-dimensional subspaces, preserving anti-collapse behavior with a bias better matched to intrinsic dynamics, and consistently outperforms LeWM across four continuous-control environments.

  • Background: World models learn compact predictive representations of observations and future states under actions, supporting latent-space planning without modeling every observation pixel.This provides the broader motivation for predictive latent world models.
  • Background: JEPA encodes observations into latent representations and predicts future latent states conditioned on current representations and actions, avoiding full observation reconstruction.The formulation shifts modeling capacity toward task-relevant abstractions.
  • Problem: Stable end-to-end JEPA training faces a bias-variance tradeoff: insufficient constraints cause representation collapse, whereas excessive constraints suppress learned representation richness.Earlier regularization and heuristic recipes were often complex, not end-to-end, and lacked theoretical guidance.
  • Problem: LeWM regularizes latent embeddings toward an isotropic Gaussian distribution, but this ambient-space prior imposes excessive bias when representations lie on low-dimensional manifolds.The mismatch arises because natural control-task representations occupy low-dimensional manifolds within high-dimensional ambient spaces.
  • Contribution: Sub-JEPA applies Gaussian regularization to multiple low-dimensional row-orthonormal projected views, relaxing the global prior while preserving its anti-collapse effect.The method retains LeWM’s simple training recipe and better matches the intrinsic dimensionality of the underlying dynamics.
  • Results: Across four continuous-control environments, Sub-JEPA consistently outperforms LeWM by clear margins, with gains correlated with reductions in effective rank.The reported correlation supports the claim that subspace regularization better respects the intrinsic dimensionality of the dynamics.

2 Related Work

Prior work develops latent world models and JEPA-based predictors to learn environment dynamics or future embeddings without reconstructing pixels. Related anti-collapse methods impose structure through contrastive, asymmetric, decorrelation, whitening, Gaussian, and random-projection-based regularization.

  • Latent World Models: World models learn compressed environment dynamics from observations, enabling planning without repeated real-world interaction.Pixel-space prediction is computationally expensive and captures irrelevant visual details; IRIS and DreamerV3 use autoregressive generative models with image decoders.
  • JEPA and End-to-End Training: JEPAs jointly train encoders and predictors to anticipate future or masked-view embeddings without reconstructing raw pixels.This framework appears in image representation learning, video modeling, and action-conditioned latent dynamics.
  • JEPA and End-to-End Training: Representation collapse is the central JEPA challenge when explicit structural constraints are absent.The supplied passage identifies collapse as a key issue in jointly trained JEPA systems but does not provide its full mechanism.
  • Anti-Collapse Regularization: Anti-collapse methods structure embedding distributions using contrastive separation, teacher-student asymmetry, dimension decorrelation, whitening, or Gaussianity constraints.SimCLR and MoCo require large batches or memory banks; BYOL uses stopgradient, while Barlow Twins and VICReg decorrelate dimensions.
  • Subspace Methods and Random Projections: Random projections support scalable dimensionality reduction and tractable high-dimensional distribution matching through one-dimensional comparisons.The Johnson-Lindenstrauss lemma and sliced Wasserstein distances exemplify this role, while Gaussian regularization sketches embeddings with random directions.

3 Method

Sub-JEPA learns latent dynamics from offline, reward-free state-action trajectories by predicting future encoded observations and regularizing embeddings across multiple frozen row-orthonormal random subspaces. This replaces LeWM’s full-space Gaussian constraint with a subspace-based constraint applied independently across projections.

  • Framework: Sub-JEPA extends LeWM by replacing full-space Gaussian regularization with Gaussian regularization over multiple frozen row-orthonormal random projections.The projections are constructed by Gaussian sampling followed by QR-based orthogonalization and are frozen during training.
  • Latent dynamics: The latent world model encodes observations as z_t = f(o_t) and predicts the next latent with P(z_t, a_t), without pixel-level reconstruction.The predictor is trained to match the target latent obtained by encoding the subsequent observation.
  • Problem setting: Training uses only pre-collected observation-action trajectories in a fully offline, reward-free setting without rewards or task labels.The trajectories have length T and may contain RGB observations and continuous control actions.
  • Subspace projections: Each latent z ∈ R^D is mapped into K subspaces of dimension d_s, with d_s = ⌊D/K⌉, using row-orthonormal projection matrices.For subspace k, the projected embedding is z^(k) = P_k z ∈ R^d_s.
  • Subspace Gaussian regularization: Gaussian regularization is applied independently in every subspace by sampling M random unit directions, forming scalar projections, and evaluating the Epps–Pulley normality statistic.The regularization objective averages over the M directions and K subspaces, alongside the latent prediction loss weighted by λ.

4 Experiments · 4.1 Experimental Setup

The experiments evaluate Sub-JEPA on four continuous-control benchmarks using raw RGB observations and compare it with three representative world-model baselines. Sub-JEPA matches LeWM’s training setup while replacing full-space Gaussian regularization with Multi-Subspace Gaussian regularization.

  • 4.1 Experimental Setup: Sub-JEPA is evaluated on Two-Room, Reacher, PushT, and OGB-Cube, spanning navigation, reaching, block-pushing, and visually rich 3D manipulation.All benchmarks use continuous action spaces and raw RGB observations.
  • 4.1 Experimental Setup: Two-Room is a 2D navigation task with low intrinsic dimensionality.It is one of the four continuous-control benchmarks used in evaluation.
  • 4.1 Experimental Setup: Reacher is a two-link planar reaching task, while PushT is a 2D block-pushing manipulation task.Both tasks are evaluated from raw RGB observations with continuous actions.
  • 4.1 Experimental Setup: OGB-Cube is a visually rich 3D manipulation environment included among the evaluation benchmarks.The benchmark uses continuous actions and raw RGB observations like the other tasks.
  • 4.1 Experimental Setup: The baselines are LeWM, PLDM, and DINO-WM.LeWM is the primary baseline; PLDM is pixel-based and uses multiple heuristic objectives, while DINO-WM uses a frozen pretrained DINOv2 visual encoder.
  • 4.1 Experimental Setup: Sub-JEPA keeps LeWM’s encoder, latent predictor, optimizer, training schedule, and loss-weight tuning protocol, replacing full-space Gaussian regularization with Multi-Subspace Gaussian regularization.The embedding dimension is D = 192 for all experiments.
  • 4.1 Experimental Setup: Sub-JEPA uses K = 32 for all environments except PushT, where K = 16 is used for better performance.These values were selected through ablation studies on a held-out validation set.

4.2 Planning Performance

Sub-JEPA consistently outperforms LeWM across four environments, with the largest improvement on Two-Room, by applying Gaussian regularization in lower-dimensional subspaces. Effective-rank analysis links greater latent compression to larger planning gains, supporting reduced bias toward task-aligned representations.

  • 4.2 Planning Performance: Sub-JEPA consistently outperforms LeWM across all four environments, demonstrating the benefit of relocating Gaussian regularization into subspaces.LeWM and Sub-JEPA are compared across six random seeds using mean ± std; PLDM and DINO-WM rates are taken from LeWM under the same protocol.
  • 4.2 Planning Performance: The largest improvement occurs on Two-Room, where full-space isotropic Gaussian regularization imposes excessive bias on a low-intrinsic-dimension task.Sub-JEPA instead regularizes lower-dimensional orthogonal views, relaxing the global constraint while retaining anti-collapse benefits.
  • 4.2 Planning Performance: On Reacher and PushT, Sub-JEPA outperforms PLDM and surpasses DINO-WM without pretrained visual features, while on OGB-Cube it substantially improves upon LeWM.DINO-WM still leads on visually complex 3D manipulation tasks, whereas Sub-JEPA remains a simple end-to-end alternative without external pretraining.
  • 4.2 Planning Performance: Effective-rank reductions from LeWM to Sub-JEPA directly correlate with planning improvements across all four environments.Two-Room and OGB-Cube show the largest rank compressions and success-rate gains, while PushT and Reacher show smaller rank gaps and more moderate improvements.
  • 4.2 Planning Performance: The results support that subspace regularization contracts latent geometry toward task intrinsic dimensionality, with greater contraction producing larger planning gains.The analysis uses N = 2000 shared evaluation observations per environment and frozen trained encoders to compare latent representations.

4.3 Ablation Studies

The ablations show that Sub-JEPA’s performance depends on balancing subspace flexibility with reliable normality estimates, while frozen orthogonal projections provide the strongest projection design. Moderate subspace counts generally help, but overly narrow subspaces and adaptive projections weaken regularization.

  • Number of subspaces K: On Two-Room, performance peaks around K = 32 before slightly degrading, while OGB-Cube achieves its best result at K = 32 and Reacher benefits slightly from moderate K increases.The effect of K is task-dependent, with larger K improving performance on OGB-Cube and moderately benefiting Reacher.
  • Number of subspaces K: PushT performance collapses at K = 32 because six-dimensional subspaces can produce unstable and uninformative normality estimates for tightly coupled manipulation.Excessive partitioning can make individual subspaces too low-dimensional for stable regularization signals.
  • Subspace design: Increasing K relaxes the global Gaussian prior, but excessively small subspaces weaken regularization effectiveness by making the statistical signal unreliable.Subspace design therefore trades latent flexibility against statistical reliability.
  • Projection strategy: Orthogonal frozen projection consistently achieves the best performance across all environments and provides balanced, non-redundant views for uniform subspace-level constraints.Its geometric isometry makes projection design a key component rather than an implementation detail.
  • Projection strategy: Random frozen and trainable soft-orthogonal projections underperform the frozen orthogonal variant, as redundancy, uneven scaling, and encoder–projection co-adaptation can weaken anti-collapse regularization.Learned projections may align with directions that reduce the effective strength of the regularizer during training.

4.4 Physical State Probing

The section probes whether PushT latent representations encode task-relevant physical information by freezing the encoder and training lightweight probes. Sub-JEPA generally matches or exceeds LeWM, while block-angle linear decoding is a noted exception that nonlinear probing largely resolves.

  • Probing protocol: On PushT, frozen encoders are evaluated by lightweight probes mapping latent embeddings to ground-truth physical variables.The protocol follows LeWM and assesses task-relevant quantities including agent location.
  • Overall results: Sub-JEPA matches or outperforms LeWM across most physical properties and probe types.The results suggest Multi-Subspace Gaussian regularization preserves physically meaningful latent structure while improving nonlinear recoverability.
  • Linear-probe limitation: Sub-JEPA slightly underperforms LeWM for block-angle prediction with a linear probe.The passage hypothesizes that subspace projection fragments angular structure across independent subspaces, reducing linear decodability.
  • Nonlinear recovery: The block-angle performance gap closes with an MLP probe, indicating that nonlinear decoding can compensate for the fragmented angular structure.This pattern illustrates a tradeoff between subspace regularization and linear decodability.
  • Temporal visualization: Sub-JEPA produces more temporally coherent latent trajectories than LeWM in representative Two-Room episodes.The visualization colors points by normalized timestamp and shows less organized temporal evolution for LeWM in several episodes.

4.5 Latent Trajectories

Sub-JEPA learns latent representations better matched to low-dimensional task geometry, producing temporally coherent trajectories across episodes, unlike LeWM’s less regular temporal structure.

  • Latent Trajectories: Sub-JEPA consistently produces temporally coherent latent trajectories across episodes, with temporally adjacent observations forming well-organized paths.The trajectories are visualized by projecting consecutive observations’ [CLS] embeddings to 2D with UMAP and coloring them by normalized temporal index.
  • Latent Trajectories: LeWM exhibits less regular temporal structure in several episodes, suggesting its full-space Gaussian prior distorts geometry when task dynamics are intrinsically low-dimensional.The comparison is made on Two-Room, where low intrinsic dimensionality makes mismatch with the full-space Gaussian prior most pronounced.

4.6 Temporal Latent Path Straightening

Sub-JEPA produces straighter latent trajectories than LeWM on PushT and OGB-Cube, without explicit optimization. This indicates reduced geometric distortion from subspace-wise regularization and may explain its planning gains.

  • 4.6 Temporal Latent Path Straightening: Temporal path straightening measures how linearly dynamics evolve in latent space, with higher mean consecutive-velocity cosine similarity indicating smoother trajectories.For embeddings z1:T, temporal velocity is defined as vt = zt+1 − zt.
  • 4.6 Temporal Latent Path Straightening: Sub-JEPA consistently produces straighter latent trajectories than LeWM on both PushT and OGB-Cube without explicit optimization.The result suggests that subspace-wise regularization reduces geometric distortion relative to full-space SIGReg.

4.7 Open-loop Rollout Visualization

Open-loop rollouts on Two-Room show that Sub-JEPA preserves scene geometry and coherent long-term state propagation better than LeWM. Its subspace-wise Gaussian regularization improves recursive latent stability over long horizons.

  • 4.7 Open-loop Rollout Visualization: Sub-JEPA maintains scene geometry during long-horizon open-loop rollouts, whereas LeWM develops increasing spatial drift and structural distortion.Both models are conditioned on frames at t = {0, 5, 10} and predict future states from actions alone for 20 steps up to t = 95.
  • 4.7 Open-loop Rollout Visualization: Subspace-wise regularization supports more coherent long-term state propagation by replacing one global Gaussian constraint with structured subspace constraints.This mechanism is presented as improving recursive latent stability.
  • 4.7 Open-loop Rollout Visualization: Sub-JEPA produces more accurate long-horizon predictions and reduced cumulative drift than LeWM in the Two-Room rollout comparison.The comparison uses recursively predicted observations conditioned on three context frames and driven only by actions.
  • 4.7 Open-loop Rollout Visualization: Sub-JEPA achieves consistently higher temporal latent-path straightness than LeWM on both environments, suggesting smoother latent dynamics.Straightness is measured by mean cosine similarity between consecutive latent velocity vectors, with higher values indicating straighter trajectories.

Conclusion

Sub-JEPA extends LeWorldModel by moving Gaussian regularization from the full embedding space into random orthogonal subspaces. This addresses the excessive bias of ambient-space regularization when task dynamics lie on low-dimensional manifolds, while experiments show consistent outperformance over LeWorldModel.

  • Method: Sub-JEPA relocates Gaussian regularization from the ambient embedding space into random orthogonal subspaces.The method is presented as a simple extension of LeWorldModel.
  • Motivation: LeWorldModel’s full-space isotropic Gaussian prior can impose excessive bias when task dynamics occupy a low-dimensional manifold.The conclusion identifies this mismatch between the high-dimensional latent space and the underlying task structure as the motivation for Sub-JEPA.
  • Empirical finding: Experiments demonstrate that Sub-JEPA consistently outperforms LeWorldModel across four tasks.
Loading 2605.09241v1…