LeJEPA: Provable and Scalable Self-Supervised Learning Without the Heuristics (2511.08544v2)

Abstract: Learning manipulable representations of the world and its dynamics is central to AI. Joint-Embedding Predictive Architectures (JEPAs) offer a promising blueprint, but lack of practical guidance and theory has led to ad-hoc R&D. We present a comprehensive theory of JEPAs and instantiate it in {\bf LeJEPA}, a lean, scalable, and theoretically grounded training objective. First, we identify the isotropic Gaussian as the optimal distribution that JEPAs' embeddings should follow to minimize downstream prediction risk. Second, we introduce a novel objective--{\bf Sketched Isotropic Gaussian Regularization} (SIGReg)--to constrain embeddings to reach that ideal distribution. Combining the JEPA predictive loss with SIGReg yields LeJEPA with numerous theoretical and practical benefits: (i) single trade-off hyperparameter, (ii) linear time and memory complexity, (iii) stability across hyper-parameters, architectures (ResNets, ViTs, ConvNets) and domains, (iv) heuristics-free, e.g., no stop-gradient, no teacher-student, no hyper-parameter schedulers, and (v) distributed training-friendly implementation requiring only $\approx$50 lines of code. Our empirical validation covers 10+ datasets, 60+ architectures, all with varying scales and domains. As an example, using imagenet-1k for pretraining and linear evaluation with frozen backbone, LeJEPA reaches 79\% with a ViT-H/14. We hope that the simplicity and theory-friendly ecosystem offered by LeJEPA will reestablish self-supervised pre-training as a core pillar of AI research (\href{https://github.com/rbalestr-lab/lejepa}{GitHub repo}).

• The paper introduces LeJEPA, a joint-embedding predictive model that mathematically ensures optimal isotropic Gaussian embeddings for downstream tasks.
• It leverages SIGReg, a sliced regularization method that efficiently matches 1D marginals to an isotropic Gaussian with minimal hyperparameter tuning.
• Empirical results across 60+ architectures demonstrate LeJEPA’s state-of-the-art performance and robustness in both in-domain and transfer settings.

Provable and Scalable Self-Supervised Learning via LeJEPA

Overview and Context

This paper introduces Latent-Euclidean JEPA (LeJEPA), a principled joint-embedding predictive architecture for self-supervised representation learning, focused on establishing theoretical optimality, practical scalability, and eliminating common heuristics prevalent in the literature. The central contribution is a rigorous characterization of the optimal embedding distribution for foundation models, coupled with a novel, tractable regularization (SIGReg), leading to a lean JEPA formulation. LeJEPA is validated empirically across >60 architectures and multiple domains, demonstrating competitive or superior performance compared to leading baselines, while offering a nearly hyperparameter-free setup and efficient implementation.

Theoretical Foundations: Embedding Distributions for Downstream Optimality

A key question addressed is: Which embedding distribution should a JEPA produce to minimize the expected risk across downstream tasks post-pretraining? The analysis considers both linear and nonlinear downstream probing regimes.

• Linear Probing: The isotropic Gaussian emerges as the unique minimizer of bias and variance in OLS and ridge settings, under fixed total variance constraints. For any anisotropic embedding, there always exists a downstream task that is adversely affected (higher estimator bias/variance), quantitatively demonstrated via analytical and simulation studies.
• Nonlinear Probing: Both k-NN and kernel regression setups exhibit minimized integrated squared bias when embedding distributions are isotropic Gaussian, even under arbitrary smooth nonlinear label functions.

These insights are formalized through lemmas and theorems specifying variance/bias behavior, and are visualized with empirical boundary and $\beta$-distribution plots to clarify the dependence of estimation error on embedding geometry.

Figure 1: Embedding geometry strongly influences estimator variance; isotropic distributions yield lower variance and more stable decision boundaries than anisotropic alternatives.

Sketched Isotropic Gaussian Regularization (SIGReg)

Obtaining the optimal embedding distribution in high dimensions requires stable, differentiable, and computationally lean regularization. Standard multivariate goodness-of-fit metrics scale poorly and/or lead to ill-conditioned gradients. SIGReg addresses this through a sliced/“sketched” approach:

• Directional Distribution Matching: Instead of estimating multivariate distances, SIGReg projects embeddings onto randomly sampled directions and matches each induced 1D marginal to the isotropic Gaussian along these directions, establishing consistency via the Cramér–Wold theorem.
• Statistical Tests: Three families are considered, with preference given to Empirical Characteristic Function (Epps–Pulley statistic). Moment-based tests lack identifiability for finite order and have unstable gradients; CDF-based tests pose differentiability and scalability obstacles on typical hardware. The Epps–Pulley test delivers bounded loss, gradient, and curvature, is amenable to parallelization (supporting DDP), and is trivial to implement.

SIGReg’s implementation costs are linear in both batch size and embedding dimension, validated by runtime statistics for large batches and high-dimensional projections.

Figure 2: SIGReg constrains embedding distributions via 1D projections and characteristic function tests, efficiently matching to the isotropic Gaussian in high dimensions.

The slicing approach is theoretically supported: under modest smoothness (Sobolev regularity $\alpha$), the expected error upper bound decays rapidly (as $M^{-2\alpha/(K-1)}$), where $M$ is the number of directions (scaling much better than quadratic tests).

Figure 3: Example distributions on the sphere demonstrate how increasing smoothness parameters $\alpha$ enable more global coverage with fewer directions in high dimensions.

LeJEPA Architecture and Implementation

LeJEPA combines the JEPA prediction loss (view agreement in embedding space) with the SIGReg regularizer, targeting the isotropic Gaussian. Its loss function is:

$$\mathcal{L}_{LeJEPA} = (1 - \lambda)\,\text{PredictionLoss} + \lambda\,\text{SIGReg}$$

where:

• Prediction Loss ensures embedding agreement between transformed/corrupted views.
• SIGReg enforces the isotropic Gaussian embedding distribution.

Implementation is succinct, requiring no explicit negative samples, stop-gradient, teacher–student, predictor, nor schedule heuristics. The system exhibits strong architectural invariance: ResNets, ViTs, ConvNets, Swin, and MaxViT variants produce high-quality representations with identical training procedures.

Figure 4: JEPA achieves non-degenerate representations by enforcing predictability across views, removing shortcut collapses by construction.

Empirical Validation

• Hyperparameter Stability: Performance of LeJEPA is robust to $\lambda$, the number of views, batch size, projector embedding dimension, and register tokens. No catastrophic collapse observed, and minor tunings marginally improve results.
• Architecture Agnosticism: LeJEPA offers non-trivial representation learning across 8 architecture families and 50+ models. SOTA or near-SOTA accuracy is achievable with frozen linear probes, supporting out-of-the-box deployments.

  Figure 5: Hundreds of models across families demonstrate SOTA downstream linear probe accuracy with LeJEPA, highlighting universal applicability.

• Training Loss as Proxy: SIGReg loss aligns almost monotonic with downstream performance, yielding $0.99$ Spearman correlation after minimal transformation; unsupervised cross-validation and model selection become feasible.

  Figure 6: High Spearman correlation between LeJEPA’s training loss and downstream accuracy across hyperparameters–a robust, label-free validation method.

• Domain Transfer and In-Domain SOTA: On specialized datasets (Galaxy10, Food101, Flowers102), LeJEPA in-domain pretraining outperforms large-scale transfer learning using DINOv2/v3/IJEPA. This contradicts the widespread belief that SSL pretraining requires enormous datasets and expensive hyperparameter search.

  Figure 7: In-domain pretraining on Galaxy10 via LeJEPA overtakes even large foundation model transfer, enabling small-data domain adaptation.

• Semantic and Perceptual Structuring: PCA visualization of last-layer features (ViT-Large, ImageNet-1K) reveals emergent object-centric segmentations, coherent perceptual grouping, and explicit separation of foreground/background without segmentation supervision.

Figure 8: Unsupervised features from LeJEPA naturally encode semantic grouping; foreground objects cluster in feature space and are visually separated from background via projection onto top PCs.

• Video Object Segmentation: Self-supervised representations yield temporally coherent object attention maps across frames, enabling object tracking and segmentation from raw video without explicit supervision.

Figure 9: Thresholding LeJEPA attention maps provides unsupervised video object segmentation with strong temporal consistency.

Performance Results

Strong numerical results are emphasized, such as ViT-H/14 reaching 79% ImageNet-1K accuracy with linear evaluation and frozen backbone. Small architectures (ResNet-18, ConvNeXt Nano) reliably reach SOTA on in-domain datasets even with $<1000$ training samples, dispelling the necessity of data scaling for representation quality.

Practical Implications and Future Directions

LeJEPA provides a methodologically rigorous and hyperparameter-lean solution to SSL representation learning, directly linking theoretical optimality (embedding distribution) and empirical usability. The tractability of SIGReg regularization makes it suitable for scaling to billion-parameter models and multi-node distributed setups. Eliminating brittle heuristic choices and aligning training and representation quality metrics paves the way for rapid, reproducible development and cross-domain deployment.

Future work could investigate extending sliced distribution matching approaches to non-Euclidean and multi-modal embedding spaces, or integrating LeJEPA with fine-tuning protocols in settings where both label efficiency and transfer robustness are critical. Also relevant is further exploration of unsupervised segmentation and object discovery, building on the spontaneous semantic grouping observed in LeJEPA features.

Conclusion

LeJEPA establishes new standards in self-supervised learning for representation quality, stability, and theoretical soundness. By mathematically characterizing and attaining the optimal embedding distribution (isotropic Gaussian), leveraging 1D distributional tests for high-dimensional scaling, and eliminating reliance on prevailing heuristics, LeJEPA advances both the reliability and applicability of JEPA-based architectures. Empirical results corroborate these advances across architectures, scales, and domains. This work underlines the utility of theory-guided algorithm design and will inform future developments in scalable, domain-agnostic representation learning.