LeJEPA: Provable and Scalable Self-Supervised Learning Without the Heuristics — part 2/10

Cleaned-text code points [15344, 29910). Bibliography omitted here; source.md includes it.

For decades, the two major solutions for AI were supervised learning ( LeCun et al., 2015 ) and learning by reconstruction ( Rumelhart et al., 1986 ) –sometimes combined together, e.g., for semi-supervised learning ( Kingma et al., 2014 ) . In supervised learning, the labels both ensure that semantically similar samples are close to each other in embedding space while preventing complete representation collapse. In particular, it is possible to measure the amount of collapse in supervised learning as a function of the number of classes ( Papyan et al., 2020 ) . The reconstruction objective is similarly well suited to prevent representation collapse as the original input must be recovered from the embeddings, i.e., the embeddings must be as informative about the input as possible–up to some optional denoising tasks that users can setup as part of the training ( Vincent et al., 2010 ) .

Because supervised and reconstruction-based learning have been widely studied for decades, there exists a large body of work to explain and inform practical designs–as well as studying their limitations in producing foundation models ( Balestriero and LeCun, 2024 ; Van Assel et al., 2025 ) . This is not the case for the more recent JEPAs where empirical advances quickly outpace anyone hoping to delve into their inner workings. This dynamic led the community to focus on post-hoc theoretical justification of already found solutions ( Liu et al., 2021 ; Shwartz Ziv and LeCun, 2024 ; Shwartz-Ziv et al., 2022 ; Zhang et al., 2023 ) . In most cases, those studies involve the Mutual Information (MI) ( Shannon, 1948 ; Cover, 1999 ) whose different bounds recover established methods ( Gutmann and Hyvärinen, 2010 ; Ma and Collins, 2018 ; Oord et al., 2018 ; Poole et al., 2019 ; Hjelm et al., 2018 ; McAllester and Stratos, 2020 ) . Because existing studies focus on explaining and interpreting already developed JEPAs, too little principled guidance and innovation has been brought forward. Instead, most of the recent empirical advances take the form of collecting larger dataset, scaling up pre-existing training recipes ( Goyal et al., 2019 ; Chen et al., 2020b ; Oquab et al., 2023 ; Fan et al., 2025 ) , and deriving novel data curation processes ( Vo et al., 2024 ; Kerdreux et al., 2025 ) .

In contrast, our goal in the following Sections 3 , 4 and 5 will be to derive a novel JEPA solution from first principles, i.e., whose design relies on proved necessary conditions for optimality, and with a pretraining recipe that can finally reconcile exploratory research, scalability, and state-of-the-art performances.

## 3 Latent Euclidean: Embeddings Should be Isotropic Gaussian

We address a fundamental question: which distribution should Enc ⁡ ( 𝐱 ) {\rm Enc}({\bm{x}}) follow to minimize empirical risk on any downstream task? We prove that the isotropic Gaussian is the unique optimal distribution for both linear ( Section 3.1 ) and nonlinear probing ( Section 3.2 ), with geometric intuition provided in Section 3.3 . This theoretical result establishes the necessary design principle for our JEPA; Section 4 then provides the practical implementation to achieve it.

### 3.1 Linear Probing

We begin by identifying the optimal distribution for f 𝜽 f_{\bm{\theta}} ’s embeddings by analyzing linear probes–one of the most popular methods for frozen encoder evaluation. Specifically, we ask: which distribution for f 𝛉 ​ ( 𝐱 ) f_{\bm{\theta}}({\bm{x}}) would be most favorable for solving arbitrary downstream tasks, i.e., for any realization of targets 𝐲 {\bm{y}} ?

Denote as 𝒁 ∈ ℝ N × K {\bm{Z}}\in\mathbb{R}^{N\times K} the matrix of N N embeddings, each K K -dimensional, from f 𝜽 ​ ( 𝒙 n ) f_{{\bm{\theta}}}({\bm{x}}_{n}) . The unknown corresponding labels are denoted as 𝒚 ∈ ℝ N {\bm{y}}\in\mathbb{R}^{N} . Without loss of generality, we consider univariate targets; the following analysis extends to multivariate targets. The linear probe minimizes the following least square problem ( Bishop and Nasrabadi, 2006 ) β ^ = arg ⁡ min β ∈ ℝ K ​ ‖ 𝒚 − 𝒁 ​ β ‖ 2 2 + λ ​ ‖ β ‖ 2 2 , \hat{\beta}=\underset{\beta\in\mathbb{R}^{K}}{\arg\min}\|{\bm{y}}-{\bm{Z}}\beta\|_{2}^{2}+\lambda\|\beta\|_{2}^{2}, (OLS) where β ^ \hat{\beta} is the optimal probe parameters, and λ ≥ 0 \lambda\geq 0 is an hyperparameter controlling the Tikhonov regularizer strength ( Bishop, 1995 ; Golub et al., 1999 ) . Despite not knowing 𝒚 {\bm{y}} , it is possible to describe the bias and variance of the estimator β ^ \hat{\beta} as a function of the distribution of 𝒁 {\bm{Z}} . Consider two embeddings with identical column spans 𝒁 aniso , 𝒁 iso {\bm{Z}}_{\rm aniso},{\bm{Z}}_{\rm iso} . 𝒁 aniso {\bm{Z}}_{\rm aniso} ’s covariance matrix eigenvalues are given by { λ k } k = 1 K \{\lambda_{k}\}_{k=1}^{K} with at least two distinct values, while 𝒁 iso {\bm{Z}}_{\rm iso} ’s covariance matrix eigenvalues are all equal to 1 K ​ ∑ k = 1 K λ k \frac{1}{K}\sum_{k=1}^{K}\lambda_{k} . Hence, the two candidate embeddings 𝒁 aniso , 𝒁 iso {\bm{Z}}_{\rm aniso},{\bm{Z}}_{\rm iso} capture the same intrinsic features and have same energy, but different geometries.

From the above Sections 3.1 and 3.1 we obtain that the distribution of features must be isotropic. We now move to nonlinear probing where the standard Gaussian will emerge as the unique optimum.

### 3.2 Nonlinear Probing

To allow for more flexible evaluation of the pretrained encoder f 𝜽 f_{{\bm{\theta}}} , it has become increasingly common to work with a nonlinear probe. We analyze two widely-used nonlinear methods: radius-based k-NN ( Taunk et al., 2019 ; Sun and Huang, 2010 ; Zhang et al., 2017 ; Abu Alfeilat et al., 2019 ) for its simplicity and kernel methods ( Nadaraya, 1964 ; Watson, 1964 ) for their theoretical tractability.

As in Section 3.1 , we ask ourselves which distribution of embeddings would be preferable for a foundation model. We first define our prediction function. The training data consists of the N N embeddings along with their training labels { ( 𝒛 n , 𝒚 n ) } n = 1 N \{({\bm{z}}_{n},{\bm{y}}_{n})\}_{n=1}^{N} . The prediction, using radius-based k-NN for a query vector 𝒒 {\bm{q}} is formed as 𝒚 ^ ​ ( 𝒒 ) := 1 | 𝒩 r 0 ​ ( 𝒒 ) | ​ ∑ n ∈ 𝒩 r 0 ​ ( 𝒒 ) 𝒚 n , \displaystyle\widehat{{\bm{y}}}({\bm{q}}):=\frac{1}{|\mathcal{N}_{r_{0}}({\bm{q}})|}\sum_{n\in\mathcal{N}_{r_{0}}({\bm{q}})}{\bm{y}}_{n}, (kNN) where 𝒩 r 0 ​ ( 𝒒 ) = { n : ‖ 𝒛 n − 𝒒 ‖ ≤ r 0 } \mathcal{N}_{r_{0}}({\bm{q}})=\{n:\|{\bm{z}}_{n}-{\bm{q}}\|\leq r_{0}\} . The specific choice of radius r 0 r_{0} controls how many neighbors predictions are averaged to form the query’s prediction. The kernel’s prediction at a query 𝒒 ∈ ℝ K {\bm{q}}\in\mathbb{R}^{K} is given by 𝒚 ^ ​ ( 𝒒 ) ​ \triangleq ​ ∑ n = 1 N K h ​ ( 𝒒 − 𝒛 n ) ​ 𝒚 n ∑ n = 1 N K h ​ ( 𝒒 − 𝒛 n ) . \displaystyle\widehat{\bm{y}}({\bm{q}})\triangleq\frac{\sum_{n=1}^{N}K_{h}({\bm{q}}-{\bm{z}}_{n}){\bm{y}}_{n}}{\sum_{n=1}^{N}K_{h}({\bm{q}}-{\bm{z}}_{n})}. (Kernel)

We search over all distributions of Z subject to a fixed total variance constraint, e.g., Tr ⁡ ( Cov ⁡ ( 𝒁 ) ) = κ 1 \Tr(\mathrm{Cov}({\bm{Z}}))=\kappa_{1} or ‖ Cov ⁡ ( 𝒁 ) ‖ F = κ 2 \|\mathrm{Cov}({\bm{Z}})\|_{F}=\kappa_{2} . The specific value of κ \kappa does not affect the optimal distribution shape. Following the same type of derivations as done in the linear regime–with the exception of some additional regularity conditions–we are able to precisely identify the isotropic Gaussian as the unique optimum to minimize bias as formalized below.

Numerous additional details and discussions on the regularity assumptions we employed are provided in Appendix A . Together, these results establish the isotropic Gaussian distribution as the optimal design to minimize the worst-case risk of a foundation model across downstream tasks.

### 3.3 Geometric and Practical Insights

We now empirically validate that the isotropic Gaussian is optimal when no information about downstream tasks is available. We focus on linear probing ( Section 3.1 ), where all considered distributions have the same total variance.

When employing a linear probe, an anisotropic distribution increases both bias (with Tikhonov regularization) and variance. Examining bias first ( Section 3.1 ), we present in Figure 18 visualizations for both continuous regression and discrete classification tasks. We observe that the cosine similarity between estimated and ground-truth parameters equals 1 only for isotropic distributions, degrading for anisotropic cases regardless of sample size or regularization strength. Regarding variance ( Section 3.1 ), we show in Figure 3 that learned parameters vary significantly more across training sets when the covariance is anisotropic (right) compared to isotropic (left)—even when using logistic regression instead of OLS. Figure 17 further illustrates this effect, showing the distribution of learned β \beta parameters across different training samples for both cases. The anisotropic distribution clearly produces higher-variance estimators.

These theoretical and empirical results establish our design principle for LeJEPA: embeddings f 𝛉 ​ ( 𝐱 ) f_{{\bm{\theta}}}({\bm{x}}) should follow an isotropic Gaussian distribution to minimize worst-case risk across downstream tasks encountered post-training . Section 4 introduces a novel regularizer to achieve this distribution.

## 4 SIGReg: Reliable Isotropic Gaussian Regularization in High-Dimension

Having established the isotropic Gaussian as the optimal embedding distribution ( Section 3 ), we now introduce Sketched Isotropic Gaussian Regularization (SIGReg)–a distribution matching objective that is simultaneously (i) differentiable , (ii) scalable , (iii) provable , and (iv) interpretable . SIGReg builds on three key innovations. First, we formulate distribution matching as a statistical test under the null hypothesis P 𝜽 = Q P_{{\bm{\theta}}}=Q ( Section 4.1 ). Second, we identify a test that guarantees bounded gradients and curvature while maintaining linear complexity and efficient multi-GPU scaling ( Section 4.2 ). Third, SIGReg bypasses the curse of dimensionality, eliminating collapsed shortcut solutions entirely ( Section 4.3 ).

### 4.1 Hypothesis Testing as a Judge

Asking for f 𝜽 ​ ( 𝒙 ) f_{{\bm{\theta}}}({\bm{x}}) ’s distribution P 𝜽 P_{{\bm{\theta}}} to match a target distribution Q Q is typically done by creating various measures of distance or divergence, and estimating them in high-dimension. We propose a different starting point grounded in statistics. Consider the hypothesis testing framework ( Fisher, 1928 ; Neyman and Pearson, 1933 ) given by H 0 : P 𝜽 = Q vs. H 1 : P 𝜽 ≠ Q , \displaystyle H_{0}:P_{{\bm{\theta}}}=Q\quad\text{vs.}\quad H_{1}:P_{{\bm{\theta}}}\neq Q, (2) with H 0 H_{0} being referred to as the null hypothesis . That is, we are asking in Equation 2 if there is enough empirical evidence to reject the null. To answer that question, one (i) employs a test-statistic , i.e., a single scalar value summarizing the evidence from the empirical samples, (ii) determines a critical value τ α \tau_{\alpha} for the test-statistic based on the probability α \alpha of Type I error, i.e., of mistakenly rejecting a true null hypothesis, (iii) compares the test-statistic to the critical value τ α \tau_{\alpha} ; if the test-statistic exceeds τ α \tau_{\alpha} , reject the null hypothesis. If the null is not rejected, we can only claim that there is not sufficient empirical evidence against P 𝛉 = Q P_{{\bm{\theta}}}=Q .

As it stands, Equation 2 remains impractical in large dimension as existing tests have at least quadratic complexity with the number of samples considered (more details in Appendix F ). We thus propose to derive a sketching strategy by decomposing Equation 2 into simpler univariate tests. Denoting the push-forward distributions P 𝜽 ( 𝒂 ) ​ \triangleq ​ ( 𝒂 ⊤ ) # ​ P 𝜽 P_{\bm{\theta}}^{({\bm{a}})}\triangleq({\bm{a}}^{\top})_{\#}P_{\bm{\theta}} and Q ( 𝒂 ) ​ \triangleq ​ ( 𝒂 ⊤ ) # ​ Q Q^{({\bm{a}})}\triangleq({\bm{a}}^{\top})_{\#}Q , we can define the following directional univariate test H 0 ​ ( 𝒂 ) : P 𝜽 ( 𝒂 ) = Q ( 𝒂 ) ​ vs. ​ H 1 ​ ( 𝒂 ) : P 𝜽 ( 𝒂 ) ≠ Q ( 𝒂 ) , \displaystyle H_{0}({\bm{a}}):P_{\bm{\theta}}^{({\bm{a}})}=Q^{({\bm{a}})}\;\;\text{vs.}\;\;H_{1}({\bm{a}}):P_{\bm{\theta}}^{({\bm{a}})}\not=Q^{({\bm{a}})}, (3) for a given directional unit-norm vector 𝒂 ∈ 𝒮 K − 1 {\bm{a}}\in\mathcal{S}^{K-1} . The corresponding directional test-statistic of Equation 3 is computed as T ⁡ ( { 𝒂 ⊤ ​ f 𝜽 ​ ( 𝒙 n ) } n = 1 N ) T(\{{\bm{a}}^{\top}f_{{\bm{\theta}}}({\bm{x}}_{n})\}_{n=1}^{N}) . Examples of tests T T will be provided in the later Section 4.2 . Repeating that process over a set of M M directions 𝔸 ​ \triangleq ​ { 𝒂 1 , … , 𝒂 M } {\mathbb{A}}\triangleq\{{\bm{a}}_{1},\dots,{\bm{a}}_{M}\} and aggregating the individual values lead to the following global test-statistic T 𝔸 ​ ( { f 𝜽 ​ ( 𝒙 n ) } n = 1 N ) ​ \triangleq ​ max 𝒂 ∈ 𝔸 ⁡ T ⁡ ( { 𝒂 ⊤ ​ f 𝜽 ​ ( 𝒙 n ) } n = 1 N ) . \displaystyle T_{{\mathbb{A}}}(\{f_{{\bm{\theta}}}({\bm{x}}_{n})\}_{n=1}^{N})\triangleq\max_{{\bm{a}}\in{\mathbb{A}}}T(\{{\bm{a}}^{\top}f_{{\bm{\theta}}}({\bm{x}}_{n})\}_{n=1}^{N}). (4) We now provide a formal statement asserting the consistency of Equation 4 to test the original multivariate null hypothesis from Equation 2 . Our result leverages the well-known union-intersection principle ( Roy, 1953 ) , and a slightly modified Cramér-Wold theorem. We denote by = d \stackrel{{\scriptstyle d}}{{=}} equality in distribution.

The assumptions required in the proof of Section 4.1 hold for classical consistent univariate tests T T such as the ones presented in the following Section 4.2 .

### 4.2 SIGReg: Sketching the Epps-Pulley Test is Stable and Scalable

Our proposed regularizer–coined Sketched Isotropic Gaussian Regularization (SIGReg)–follows directly from Section 4.1 using any statistical test T T targeted towards the isotropic Gaussian, illustrated in Figures 2 and 5 , and formalized below.

We replace the maximum over 𝒂 ∈ 𝔸 {\bm{a}}\in{\mathbb{A}} in Section 4.1 by an average in ( 5 ) to avoid sparse gradient over the directions in 𝔸 {\mathbb{A}} . We now delve on the choice of T T for which we compare well-known candidate tests in the field of statistics that are categorized into (i) moment based ( Section 4.2.1 ), (ii) CDF based ( Section 4.2.2 ), and (iii) CF based ( Section 4.2.3 ) statistics–ultimately justifying our choice of the Epps-Pulley statistic.

#### 4.2.1 Moments are Unstable and Insufficient