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

1. Problem Statement

Joint-Embedding Predictive Architectures (JEPAs) learn representations by predicting the embedding of one augmented view from another in latent space. A central difficulty in training JEPAs is preventing representation collapse, where the encoder maps all inputs to a constant or trivial solution.

Historical remedies have been largely heuristic (e.g., stop-gradients, contrastive negatives, covariance penalties). LeJEPA [1] advanced this by replacing heuristics with a single explicit target-distribution regularizer: Sketched Isotropic Gaussian Regularization (SIGReg). SIGReg enforces that the embedding distribution be an isotropic Gaussian $N(0, I_D)$ across the full ambient dimension $D$.

However, this target conflicts with the manifold hypothesis, which posits that real-world data embeddings concentrate on a low-dimensional subset of intrinsic dimension $n \ll D$. Forcing embeddings to fill the full $D$-dimensional space isotropically contradicts the prior that motivates representation learning. The paper proposes UR–JEPA, which replaces the isotropic Gaussian target with a uniformly $n$-rectifiable measure. This target enforces that embeddings concentrate on a set that is quantitatively $n$-dimensional at small scales, aligning the regularization with the geometric structure of real data.

2. Methodology

UR–JEPA retains the LeJEPA predictive loss $L_{pred}$ but substitutes the SIGReg regularizer with geometric measure-theoretic objectives derived from Uniform $n$-Rectifiability (UR). Two complementary characterizations of UR are utilized to construct differentiable losses suitable for Stochastic Gradient Descent (SGD).

2.1 Theoretical Foundations

• Uniform $n$-Rectifiability: A measure $\mu$ is uniformly $n$-rectifiable if it is $n$-Ahlfors–David regular (density scales as $r^n$) and possesses "Big Pieces of Lipschitz Images."
• Characterizations:
  1. CGLT Square Function: Based on the Chousionis–Garnett–Le–Tolsa (CGLT) theorem, which characterizes UR via the boundedness of a Carleson-type square function involving scale-normalized ball densities.
  2. Jones $\beta$-numbers: Measures the local deviation of the data from the best-fitting affine $n$-plane.

2.2 The UR–JEPA Loss Families

A. Gaussian-Kernel Smoothed Carleson Loss ($L_{CGLT}$)

This family approximates the CGLT square function using a Gaussian kernel to smooth the empirical density, avoiding the noise issues of raw ball counts.

• Density Estimation: A scale-normalized density $\theta_r(x)$ is computed using a Gaussian kernel with bandwidth $r$.
• Log-Increment Variant (Default): To avoid coupling the loss to the absolute density level (which is unknown), the loss operates on the log-density increment: $\Delta_r \log \theta(x) = \log \theta_{2r}(x) - \log \theta_r(x)$.
  • Loss: $L_{CGLT}^{UR} = \sum_{x \in A} \sum_k (\Delta_{r_k} \log \theta(x))^2$.
  • AD-Regularity Anchor: A separate term $L_{AD}$ penalizes the variance of $\log \theta_r$ across anchors to enforce the $n$-Ahlfors–David regularity condition (constant density level).
• Scale-Derivative Variants: Alternative formulations using the scale derivative $t \partial_t \log \theta_t$ or the raw kernel derivative. The paper finds the log-derivative variant essential for stability, as the raw form suffers from parasitic coupling to local density levels.

B. Jones $\beta$-Number Loss ($L_\beta$)

This family enforces that local neighborhoods lie close to an $n$-dimensional affine plane.

• Local PCA: For each anchor $x$ and scale $r$, a weighted scatter matrix is formed. The loss is the sum of the eigenvalues corresponding to the orthogonal complement of the top-$n$ eigenvectors (the residual variance).
• Anti-Collapse Mechanisms: The raw $\beta$-loss admits a trivial point-mass minimizer (collapse to a single point yields zero residual). The paper introduces two remedies:
  1. Intrinsic Log-Trace Penalty: Adds $-\gamma \log \text{tr}(S_r(x))$ to penalize the collapse of the local neighborhood volume.
  2. Adaptive Eigenvalue Thresholding: Instead of fixing the tangent dimension to $n$, it selects eigenvectors where the variance share exceeds a threshold $\tau$. This allows the local intrinsic dimension to vary per anchor while maintaining the global geometric target.

2.3 Anti-Collapse Properties

Unlike the $\beta$-loss, the CGLT-based losses are anti-collapse by construction. On a point-mass configuration, the smoothed density scales as $r^{-n}$, causing the log-increment $\Delta_r \log \theta$ to be a non-zero constant ($n \log 2$). Thus, the loss remains bounded away from zero, preventing the optimizer from settling in a collapsed state without needing auxiliary penalties.

3. Key Contributions

1. Geometric Regularization for JEPAs: The paper operationalizes Uniform $n$-Rectifiability as a training loss, providing a principled alternative to the isotropic Gaussian target of LeJEPA. This aligns the regularization with the manifold hypothesis.
2. New Loss Families: It introduces two concrete, SGD-friendly loss families:
   • $L_{CGLT}$: Based on the CGLT theorem, utilizing Gaussian-kernel smoothed densities and log-increments.
   • $L_\beta$: Based on Jones $\beta$-numbers via local PCA, augmented with log-trace penalties to prevent collapse.
3. Theoretical-Practical Bridge: It demonstrates how deep learning can implement complex geometric measure theory concepts (Carleson measures, Jones $\beta$-numbers) through differentiable approximations.
4. Methodological Insight: The paper establishes that the log-transform of the density (or scale-derivative) is methodologically essential. The raw scale-derivative form underperforms significantly due to coupling with local density levels, whereas the log-form removes this confound.

4. Empirical Results

Experiments were conducted on ImageNet-10 (Inet10), Galaxy10 SDSS, ImageNet-100 (Inet100), and EuroSAT.

4.1 Performance vs. LeJEPA

• Inet10: UR–JEPA($L_{CGLT}$) achieves 0.9141 ± 0.0014 top-1 accuracy, outperforming the matched-recipe LeJEPA($L_{SIGReg}$) (0.9058 ± 0.0019) by +0.83 percentage points (pp). Crucially, UR–JEPA exhibits ~30% lower seed standard deviation.
• Galaxy10 SDSS: At a matched recipe, UR–JEPA($L_{CGLT}$) (0.8142) and LeJEPA (0.8105) are statistically tied, with UR–JEPA showing lower variance. Both significantly outperform published LeJEPA baselines (attributed to recipe improvements), but the regularizer itself provides a marginal, non-significant gain over the matched LeJEPA baseline.
• Inet100 & EuroSAT: UR–JEPA and LeJEPA converge to the same peak-accuracy band. On EuroSAT, UR–JEPA maintains its signature of ~2x tighter seed variance compared to LeJEPA.
• Scaling: Unlike LeJEPA, which degrades when increasing projector dimension $D$ from 16 to 32, UR–JEPA improves, suggesting its local geometric statistics are less sensitive to the sparsity of the ambient space.

4.2 Geometric Distinction

Despite similar accuracy, the two methods produce structurally distinct representations:

• LeJEPA: Produces a near-isotropic spectrum (flat eigenvalue distribution), consistent with its Gaussian target.
• UR–JEPA: Produces a sharply anisotropic spectrum. The covariance eigenvalues drop by 4 to 5 orders of magnitude at index $\sim 20\text{--}25$ (out of $D=32$). This indicates an effectively low-rank projection, consistent with the intrinsic dimension of natural images, while preserving near-Gaussian per-dimension marginals (a consequence of the Diaconis–Freedman phenomenon).
• Robustness: The "cliff" signature in UR–JEPA is reproducible across seeds and datasets, whereas LeJEPA remains isotropic.

4.3 Ablation Studies

• Intrinsic Dimension ($n$): Performance is robust across a wide plateau ($n \in [6, 10]$) but collapses catastrophically at $n=4$.
• Log-Transform: The log-derivative variant of the CGLT loss significantly outperforms the raw form (by ~1.5 pp), confirming the necessity of the log-transform for stability.
• $\beta$-Loss: The bare $\beta$-loss collapses on Inet10. However, with the log-trace penalty ($L_{\beta, \gamma}$), it becomes competitive on Galaxy10, though it underperforms the CGLT variants.

5. Significance and Claims

The paper claims that UR–JEPA offers a principled geometric alternative to distributional matching in self-supervised learning.

• Structural Recovery: UR–JEPA recovers the local tangent dimension as a tunable parameter and enforces scale-free shape constraints (Carleson-type conditions) that LeJEPA cannot. It produces representations that are effectively low-rank globally while maintaining local $n$-rectifiability.
• Stability: The method demonstrates lower seed variance and better scaling with projector dimension compared to LeJEPA, suggesting a more stable optimization landscape.
• Modesty: The authors acknowledge that while UR–JEPA matches or slightly exceeds LeJEPA in accuracy, the primary gain is structural (low-rank, geometrically consistent representations) and statistical (reduced variance). They do not claim a universal superiority in accuracy across all tasks, noting that on Galaxy10 and EuroSAT, the methods are statistically tied at convergence.
• Theoretical Gap: The paper admits it does not provide a finite-sample convergence guarantee for minimizing the Carleson loss toward a UR target, leaving this as an open theoretical question.

In summary, UR–JEPA successfully translates the geometric measure theory concept of uniform rectifiability into a practical, differentiable regularizer that aligns self-supervised learning with the manifold hypothesis, yielding representations that are geometrically distinct and statistically more stable than those produced by isotropic Gaussian targets.