Sub-JEPA: Latent Subspace Innovations

• Sub-JEPA is a suite of techniques that decomposes the JEPA latent space into modular subspaces to prevent collapse and enhance interpretability.
• It applies targeted regularization methods, including Gaussian subspace loss, progression–content splits, and operator-aligned modules, to improve data efficiency and planning success.
• Sub-JEPA enhances self-supervised learning through theoretically grounded constraints that yield robust, task-structured, and physically-informed latent representations.

Sub-JEPA refers to a suite of architectural and algorithmic innovations within the Joint-Embedding Predictive Architecture (JEPA) paradigm, leveraging latent subspace decompositions or operator-aligned latent modules to address issues of collapse prevention, data efficiency, and interpretability in self-supervised world modeling. The name encompasses three core lines: Subspace Gaussian Regularization for world models (Zhao et al., 10 May 2026), Subspace-Decomposed JEPAs disentangling progression and content (Thil et al., 29 May 2026), and Sub-JEPA modular banks in Physics-Informed JEPA for operator-split PDE surrogates (Yee et al., 1 Apr 2026). Though each originated independently, all rely on the key idea: aligning regularization or predictive structure to latent subspaces—either learned, random, or physics-derived—to yield more robust, interpretable, and data-efficient representations.

1. Sub-JEPA in World Modeling: Subspace Gaussian Regularization

Sub-JEPA, in the context of world models and control, addresses the bias–variance trade-off inherent in JEPA latent collapse prevention. Standard JEPA training, which predicts future latents, is susceptible to trivial solutions (collapse), for example, by mapping all inputs to a constant vector. Prior efforts (LeWM) penalize deviation from an isotropic Gaussian prior in the full latent space, but this injects excessive bias, especially when the true world dynamics occupy a low-dimensional manifold within a high-dimensional embedding.

Sub-JEPA resolves this by partitioning the latent embedding ($z \in \mathbb{R}^D$) into $K$ random, orthonormal subspaces of dimension $d_s \ll D$, enforcing the Gaussian prior only within each subspace. Specifically, for each subspace, random matrices $G_k$ are sampled and orthonormalized to form projection operators $P_k$, which remain frozen throughout training. The batch of latents is projected to each $d_s$-dimensional subspace, and the Epps–Pulley normality statistic is applied along many ($M$) random 1D directions per subspace, averaging the resulting tests to yield the subspace Gaussian loss. The total training objective becomes

$$\min_{\theta, \phi} \; \mathcal{L}_{\mathrm{pred}} + \lambda \mathcal{L}_{\mathrm{reg}},$$

where $\mathcal{L}_{\mathrm{pred}}$ is the latent prediction loss and $\mathcal{L}_{\mathrm{reg}}$ is the mean Epps–Pulley statistic over all $K$0 subspaces and directions (Zhao et al., 10 May 2026).

This subspace approach weakens the global bias, permitting the encoder to concentrate information on dynamic-relevant directions, while still preventing collapse. Empirical benchmarks demonstrate that Sub-JEPA outperforms LeWM by clear margins in planning success across continuous control environments (mean improvements: +10.7% Two-Room, +9% OGB-Cube) and does so with more compact, smooth latent representations.

2. Orthogonal Latent Decomposition: Progression–Content Split

The Subspace-Decomposed JEPA (SD-JEPA) framework introduces an explicit split in the JEPA latent space into orthogonal progression and content subspaces (Thil et al., 29 May 2026). The latent encoding $K$1 is written as

$K$2

with $K$3 and $K$4 being orthogonal injection matrices, and $K$5 intended to represent explicit “progression” (i.e., task phase), while $K$6 encodes the remaining content.

Training combines three terms:

• Next-embedding MSE on the full $K$7.
• SIGReg (“Sketched Isotropic Gaussian Regularizer,” from LeWM) applied to $K$8 only, to prevent collapse in content.
• Cosine-margin triplet loss on $K$9, encouraging monotonic progression along episode time.

Straightening losses may be added but are not essential for the improvement. Crucially, these losses act on orthogonal coordinates (formally proven as “disjoint gradient supports”), so anti-collapse pressures are additive rather than competing.

SD-JEPA yields improved planning on control benchmarks (e.g., Push-T +1.3 pp, Reacher +2 pp vs LeWM) and, for $d_s \ll D$0, produces an interpretable 1D angular coordinate $d_s \ll D$1. This “scene-aware compass” advances with task progress and regresses under backtracking (e.g., Two-Room domain), offering phase-aware surprise signals superior to naive prediction errors (|$d_s \ll D$2| AUROC = 0.414 vs 0.238 for $d_s \ll D$3–MSE at $d_s \ll D$41 step, 97.5% per-episode win rate) (Thil et al., 29 May 2026).

3. Sub-JEPA in Physics-Informed Surrogate Modeling

In Physics-Informed JEPA (PI-JEPA), “Sub-JEPA” refers to a bank of lightweight, chained predictor modules aligned with the Lie–Trotter operator-splitting decomposition prevalent in multiphysics PDE solvers (Yee et al., 1 Apr 2026). The architecture includes:

• Context and EMA target encoders with Fourier+attention backbones.
• $d_s \ll D$5 predictor transformers $d_s \ll D$6, one per physical operator (e.g., pressure, transport, reaction).
• Auxiliary decoders $d_s \ll D$7, used to map latent predictions back to pixel space for PDE-residual computation.

PI-JEPA is trained end-to-end, with a masked latent prediction loss in latent space (never pixel space) and per-sub-predictor PDE residual regularizers, using only “free” unlabeled geostatistical parameter fields (e.g., permeability, porosity). Only during fine-tuning is supervised data (full PDE solves) used, and sample efficiency improves dramatically: PI-JEPA achieves $d_s \ll D$8 lower error than FNO and $d_s \ll D$9 lower than DeepONet at $G_k$0 supervised samples; the benefit is preserved up to $G_k$1, where PI-JEPA remains 24% more accurate than a scratch-trained counterpart (Yee et al., 1 Apr 2026).

By chaining the sub-predictors according to the split-step structure of the underlying physics (e.g., first predict pressure, then saturation, then reaction), each module learns a specialized dynamics, and the latent subspace predictors collectively model the full PDE evolution.

4. Empirical Findings and Practical Considerations

Experiments with Sub-JEPA variants report the following:

• World models: Subspace Gaussian regularization yields higher planning success and more compact latent representations than full-space regularization. Empirical performance peaks with $G_k$2 subspaces of dimension $G_k$3–$G_k$4 ($G_k$5), with degradation observed for excessively small ($G_k$6) or large ($G_k$7) subspaces.
• Projection strategies: Orthogonal frozen subspace projections outperform both random and learned projections, likely because learning projections allows co-adaptation that undermines collapse prevention.
• SD-JEPA: The progression–content split is provably non-interfering; sweep experiments confirm that the split, and not just the presence of the triplet or SIGReg losses, drives gains.
• Planning and surprise localization: The phase-like progression variable $G_k$8 from SD-JEPA captures semantic transitions (e.g., gripper contact) more reliably than scalar latent prediction error, showing unique interpretability advantages.
• In PI-JEPA, the modular Sub-JEPA architecture reduces sample complexity from $G_k$9 to $P_k$0 in large-scale reservoir grids, corresponding to order-of-magnitude gains in simulation budget.

5. Theoretical And Algorithmic Principles

Sub-JEPA methods are characterized by:

• Leveraging the mismatch between ambient latent dimension and the true dynamical manifold to relax anti-collapse regularization selectively (subspaces in (Zhao et al., 10 May 2026), content subspace in (Thil et al., 29 May 2026), physics-aligned modules in (Yee et al., 1 Apr 2026)).
• Formally proving that orthogonalized anti-collapse losses avoid double-counting and span independent coordinates ((Thil et al., 29 May 2026), Propositions 3.1, 3.2).
• Exploiting operator-split structure for modular surrogate pretraining, converting unlimited unlabeled instances into self-supervised signals for each sub-operator (Yee et al., 1 Apr 2026).

A general pattern is moving from single, uniform global latent constraints toward localized, modular, or task-structured constraints that reflect both inductive and physical structure of the modeled system.

6. Applications and Broader Impact

Sub-JEPA architectures find direct application in:

• Data-efficient reinforcement learning and planning from raw observations, with improvements in planning success and interpretability.
• Physics-informed surrogate modeling for multiphysics PDEs, where operator-aligned latent predictors enable label-free pretraining on cheap parameter ensembles and fine-tuning on limited simulations.
• Structured representation learning, particularly where interpretable axes of progression or phase are valuable for analysis, anomaly detection, or directing exploration.

These approaches are especially advantageous whenever:

• Labeled trajectories are scarce or expensive to generate, but large quantities of structured, unlabeled instances are available (e.g., in geostatistics, high-dimensional simulation, or robotics).
• The true data-generating dynamics are much lower-dimensional than high-capacity encoders would suggest.
• Interpretability of long-horizon latent trajectories or phase progression is critical for downstream usage.

7. Summary Table: Key Sub-JEPA Variants

Sub-JEPA Variant	Core Mechanism	Representative Results
Subspace Gaussian (Zhao et al., 10 May 2026)	K random orthonormal subspaces, Gaussian prior in each	+10%–11% planning vs LeWM; robust across K, $P_k$1
SD-JEPA (Thil et al., 29 May 2026)	Explicit progression–content split; orthogonal losses	+3 pp planning, interpretable phase; 97.5% AUROC surprise
PI-JEPA (physics) (Yee et al., 1 Apr 2026)	Operator-aligned bank of latent predictors, PDE regularizers	$P_k$2 lower surrogate error, $P_k$3 data savings

Taken together, Sub-JEPA methodologies represent a systematic exploitation of latent subspace structure, whether statistical, semantic, or physics-imposed, to enhance stability, sample efficiency, and interpretability across both world modeling and surrogate modeling domains.