cs.LGMay 31, 2026

UR-JEPA: Uniform Rectifiability as a Regularizer for Joint-Embedding Predictive Architectures

Authors: Triet M. Le

Abstract

A central difficulty in training Joint-Embedding Predictive Architectures (JEPAs) is preventing representation collapse. LeJEPA addresses this by enforcing an isotropic Gaussian target on the embeddings via Sketched Isotropic Gaussian Regularization (SIGReg). This target is in tension with the manifold hypothesis, which expects embeddings to concentrate on a low-dimensional subset of the ambient space. We propose \emph{UR-JEPA}, which targets a uniformly nn-rectifiable measure of local tangent dimension nn at small scales, realized through a Gaussian-kernel smoothed Carleson-type square function LCGLT\mathcal{L}^{\text{CGLT}}, with a complementary Jones ββ-number formulation. On Inet10, UR-JEPA(LCGLT\mathcal{L}^{\text{CGLT}}) attains 0.9141±0.00140.9141 \pm 0.0014 for a +0.83+0.83,pp gain over LeJEPA(LSIGReg\mathcal{L}^{\text{SIGReg}}) with 30%\sim 30\% lower seed standard deviation; on matched-recipe Galaxy10~SDSS, a single-seed ImageNet-100100 run, and a 33-seed EuroSAT remote-sensing run, the two methods lie in the same peak-accuracy band at convergence, with UR-JEPA retaining its lower-seed-variance signature. On EuroSAT the in-domain pair is competitive at 96.096.0 to 96.1%96.1\% with large remote-sensing foundation-model transfer at a 25×25\times smaller backbone. The distinction is geometric: direct visualization of the projector output distribution shows that on all four datasets UR--JEPA(LCGLT\mathcal{L}^{\text{CGLT}}) produces a global PCA spectrum with a 44 to 55 order-of-magnitude drop at index 20\sim 20 to 2525 out of D=32D = 32, while LeJEPA's spectrum is near-flat (top-to-bottom ratio at most 3.63.6). Per-dimension marginals are simultaneously near-Gaussian for both methods (mean Shapiro-Wilk W[0.992,0.996]W \in [0.992, 0.996]) as a Diaconis-Freedman consequence. At matched accuracy the two regularizers therefore yield structurally distinct projected representations.

Explore similar work

May 10, 2026cs.LG

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

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.
Kai Zhao, Dongliang Nie, Yuchen Lin +4
May 29, 2026cs.LG

Subspace-Decomposed JEPAs: Disentangling Progression and Content in Latent World Models

Joint-Embedding Predictive Architectures (JEPAs) learn compact latent world models by predicting future embeddings, but no single coordinate of the latent is designated to encode task progression. We carve the JEPA latent into two orthogonal subspaces with disjoint roles: a low-dimensional progression subspace shaped by a cosine-margin triplet loss, and a high-dimensional content subspace regularised by the existing SIGReg objective of LeWM. We prove that the two anti-collapse forces act on disjoint coordinates, so they compose additively rather than competing on the same dimensions. Our method, SD-JEPA improves over the LeWM baseline on the majority of its control benchmarks at matched compute, and outperforms the strongest non-LeWM JEPA baseline on Push-T; a subspace-ablation falsifier confirms the split is the load-bearing ingredient. Beyond planning, the resulting 1-D angular progression coordinate functions as a scene-aware compass on the latent. It advances with task progress, regresses when the agent backtracks, and under controlled perturbations both spikes and relocalises to a semantically appropriate new task-phase sector, separating the moment of surprise from its meaning in a way that prediction-error scalars cannot.
Lucas Thil, Jesse Read, Rim Kaddah +1
Jul 15, 2026cs.LG

The SIGReg Objective as Variational Free Energy: A Theoretical Active-Inference Account of JEPA World Models

Joint-Embedding Predictive Architectures (JEPAs) are the dominant design for latent world models, yet they are usually justified by empirical performance rather than a normative principle. We show that the choice of anti-collapse regulariser determines whether a JEPA's training objective, a prediction loss plus a weighted embedding regulariser, is a valid Active Inference (AIF) variational free energy. We organise four non-contrastive regularisers (VICReg, LogDet, PairDist, and SIGReg) into an entropy-estimator hierarchy indexed by a prior-miscalibration gap, and show that the gap's sign, whether the estimator bounds the latent entropy from above or below, decides whether the AIF surprise bound survives: VICReg and LogDet are unsafe upper bounds, PairDist a safe lower bound, and SIGReg eliminates the gap. We then prove a correspondence theorem: under the standard constant-noise encoder model and successful SIGReg enforcement (isotropic-Gaussian embeddings), the gap vanishes, the objective becomes an exact information bottleneck, the surprise bound is preserved, and the latent goal cost becomes an exact proxy for AIF pragmatic value, whereas VICReg leaves an irreducible second-order anisotropy term. We extend the correspondence to multi-step expected free energy, ensemble epistemic value, and a learned-policy regime, and we identify the one AIF term no current JEPA world model computes: the state-epistemic value, a future-state coverage signal. The predictions differ in kind, not degree, and are stated here as theoretical consequences left for empirical test in separate work; full proofs are in Appendix A, and the algebraic core of every result is machine-verified in Lean 4 (Appendix D).
Fabio Arnez, Alexandra Gomez-Villa