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

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

The first family of statistics we consider are moment-based. Taking the standard Gaussian as an instanciation for the moments, we can define the Jarque-Bera ( Jarque and Bera, 1980 ) test that compares the third and fourth moments, i.e., skewness and kurtosis, as JB ⁡ ( 𝒖 ) ​ \triangleq \displaystyle{\rm JB}({\bm{u}})\triangleq N 6 ​ ( skew ^ ​ ( 𝒖 ) 2 + ( kurt ^ ​ ( 𝒖 ) − 3 2 ) 2 ) , \displaystyle\frac{N}{6}\left(\widehat{\rm skew}({\bm{u}})^{2}+\left(\frac{\widehat{\rm kurt}({\bm{u}})-3}{2}\right)^{2}\right), (Jarque-Bera) where skew ^ \widehat{\rm skew} is the skewness computed from the data as 1 n ​ ∑ i = 1 n ( x i − μ ^ ) 3 σ ^ 3 \frac{\frac{1}{n}\sum_{i=1}^{n}\left(x_{i}-\hat{\mu}\right)^{3}}{\hat{\sigma}^{3}} and kurt ^ \widehat{\rm kurt} is the kurtosis 1 n ​ ∑ i = 1 n ( x i − μ ^ ) 4 σ ^ 4 \frac{\frac{1}{n}\sum_{i=1}^{n}\left(x_{i}-\hat{\mu}\right)^{4}}{\hat{\sigma}^{4}} . Typically, the ( Jarque-Bera ) test is used to see if a density follows a Gaussian distribution of any mean and variance–hence it only looks at moments 3 and 4. In our case we aim for a standard Gaussian test and thus add the usual statistics on the first two moments, leading to the extended test EJB ⁡ ( 𝒖 ) ​ \triangleq \displaystyle{\rm EJB}({\bm{u}})\triangleq N ​ μ ^ ​ ( 𝒖 ) 2 σ ^ ​ ( 𝒖 ) 2 + ( N − 1 ) ​ ( σ ^ ​ ( 𝒖 ) 2 − 1 ) 2 2 + JB ⁡ ( 𝒖 ) . \displaystyle\frac{N\hat{\mu}({\bm{u}})^{2}}{\hat{\sigma}({\bm{u}})^{2}}+\frac{(N-1)\left(\hat{\sigma}({\bm{u}})^{2}-1\right)^{2}}{2}+{\rm JB}({\bm{u}}). (Extended Jarque-Bera) The ( Extended Jarque-Bera ) acts as a moment matching problem over the first four moments. Such moment matching methods have proven powerful not only for statistical tests but also as mean to learn parametric and nonparametric models of data.

The Stability and Identifiability Conundrum. We now explain why moment-based tests–albeit powerful–will not be suited for LeJEPA. The k t ​ h k^{th} of a distribution P P is denoted as m k ​ ( P ) m_{k}(P) . The first observation is that well-behaved distributions abiding the Carleman’s condition ∑ k = 1 ∞ m 2 ​ k ( Q ) − 1 / ( 2 k ) = ∞ \sum_{k=1}^{\infty}m_{2k}(Q)^{-1/(2k)}=\infty ( Carleman, 1926 ) , such as the Gaussian, or for distributions with finite interval ( Hausdorff, 1923 ) are uniquely determined by their moments. However, using a finite number of moments creates the following non-identifiability issue which well-known in statistics and often used as a motivation to use all moments ( Lehmann and Romano, 2005 ) .

Hence Section 4.2.1 prescribes us with the guideline to employ as large K K as possible to remove collapsed shortcut solution by making sure our distribution matching is accurate. Yet, doing so leads to unstable gradient-based training due to the gradient norm scaling as O ⁡ ( k ) O(k) , and the variance of Monte Carlo gradient estimates growing as O ⁡ ( k 2 ​ m 2 ​ ( k − 1 ) ) O(k^{2}m_{2(k-1)}) for the k k -th moment since ‖ ∇ θ m k ​ ( P 𝜽 ( 𝒂 ) ) ‖ = ‖ 𝔼 ⁡ [ k ​ ( 𝒂 ⊤ ​ f 𝜽 ​ ( 𝒙 ) ) k − 1 ​ 𝒂 ⊤ ​ J f 𝜽 ​ ( 𝒙 ) ] ‖ \big\|\nabla_{\theta}m_{k}(P_{{\bm{\theta}}}^{({\bm{a}})})\big\|=\|\mathbb{E}\big[k({\bm{a}}^{\top}f_{{\bm{\theta}}}({\bm{x}}))^{k-1}{\bm{a}}^{\top}J_{f_{\bm{\theta}}}({\bm{x}})\big]\| , with J f 𝜽 ​ ( 𝒙 ) ∈ ℝ K × P J_{f_{\bm{\theta}}}({\bm{x}})\in\mathbb{R}^{K\times P} the Jacobian matrix–hereby creating an impractical situation where training stability and identifiability can not be achieved simultaneously.

#### 4.2.2 Cumulative Density Functions are Impractical

The second family of tests acts upon the CDF. Because those tests require sorting, let’s denote the k th k^{\rm th} order-statistics of N N samples by x k : N x_{k:N} . Two highly standard tests are quadratic Empirical Density Function statistics with different weighting known as Cramér-von Mises ( Cramér, 1928 ; Von Mises, 1981 ) and Anderson Darling ( Anderson and Darling, 1952 ) , and given by T w \displaystyle T_{w} = N ​ ∫ − ∞ ∞ ( F N ​ ( x ) − F ⁡ ( x ) ) 2 ​ w ​ ( x ) ​ 𝑑 F ​ ( x ) \displaystyle=N\int_{-\infty}^{\infty}\left(F_{N}(x)-F(x)\right)^{2}w(x)dF(x) w ⁡ ( x ) \displaystyle w(x) = 1 , \displaystyle=1, (Cramér-von Mises) w ⁡ ( x ) \displaystyle w(x) = [ F ⁡ ( x ) ​ ( 1 − F ⁡ ( x ) ) ] − 1 , \displaystyle=[F(x)(1-F(x))]^{-1}, (Anderson-Darling) where w ⁡ ( x ) w(x) is a weighting function. Adding the U 2 U^{2} statistics on top of Equation Cramér-von Mises recovers the Watson test ( Watson, 1961 ) U 2 = T w − N ​ ( F ¯ − 1 2 ) 2 . U^{2}=T_{w}-N\left(\bar{F}-\frac{1}{2}\right)^{2}. (Watson) We do not consider the Kolmogorov-Smirnov test ( Kolmogorov, 1933 ) as it employs the ℓ ∞ \ell_{\infty} -norm instead of the ℓ 2 \ell_{2} -norm hereby producing sparse gradients. Another common test is the Shapiro-Wilk test ( Shapiro and Wilk, 1965 ) which we found to be unstable in practice–details are provided in Appendix E .

Lack of Scalability and Differentiability. CDF-based tests require sorting that have been highly optimized, e.g., with the 𝒪 ⁡ ( N ​ log ⁡ ( N ) ) \mathcal{O}(N\log(N)) Quicksort algorithm ( Hoare, 1962 ) but that nonetheless breaks the embarrassingly parallel nature of SGD–especially on multi-GPU ( Tanasic et al., 2013 ; Maltenberger et al., 2022 ) due to synchronization requirements. Moreover, these tests involve non-differentiable operations (sorting and order statistics), making them unsuitable for gradient-based optimization without relaxations ( Cuturi et al., 2019 ; Grover et al., 2019 ; Petersen et al., 2022 ) . While there exists intricate sketching solutions ( Dunning and Ertl, 2019 ; Masson et al., 2019 ; Dunning, 2021 ) , each of those solutions introduce numerous additional hyper-parameters–going against our first motivation for LeJEPA.

#### 4.2.3 Characteristic Functions are Stable, Scalable and Identifiable

The third family of tests is concerned with Empirical Characteristic Functions (ECF) which are the Fourier transform of the density function. The Epps–Pulley test ( Epps and Pulley, 1983 ) is one of the most popular test and simply compares in weighted ℓ 2 \ell_{2} -norm the ECF of the data against a target CF E ​ P = N ​ ∫ − ∞ ∞ | ϕ ^ X ​ ( t ) − ϕ ⁡ ( t ) | 2 ​ w ​ ( t ) ​ 𝑑 t . EP=N\int_{-\infty}^{\infty}\left|\hat{\phi}_{X}(t)-\phi(t)\right|^{2}w(t)dt. (Epps–Pulley) The first crucial observation is that the ECF being defined as ϕ ^ X ​ ( t ) = 1 n ​ ∑ j = 1 n e i ​ t ​ X j \hat{\phi}_{X}(t)=\frac{1}{n}\sum_{j=1}^{n}e^{itX_{j}} is naturally differentiable and easily computed in distributed settings via efficient all_reduce operations, as the ECF is a simple average of complex exponentials. The weight function is typically Gaussian, such as w ( t ) = e − t 2 / σ 2 w(t)=e^{-t^{2}/\sigma^{2}} with σ \sigma commonly set to 1 1 .

Other tests, e.g., based on the Entropy ( Székely and Rizzo, 2005 ) are not considered here as they require numerous additional design choices for the univariate Entropy estimation ( Silverman, 2018 ; Beirlant et al., 1997 ) , e.g., using kernels ( Joe, 1989 ) , or M-estimators ( Miller, 2003 ) .

Epps-Pulley has bounded loss, gradient and curvature. We now consider the remaining two families of tests: moment-based and CF-based. First, recall that moments are polynomial in the data and with extreme growth rate for higher moment–assuming they even exist. Even for well-behaved distributions, raising values to a power of k k can quickly lead to exploding gradients. This comes in sharp contrast with the ECF which is always bounded and with bounded gradients for any input distribution for the projected samples z i = 𝒂 ⊤ ​ f θ ​ ( 𝒙 n ) z_{i}={\bm{a}}^{\top}f_{\theta}({\bm{x}}_{n}) , n = 1 , … , N n=1,\ldots,N .

SIGReg with Epps-Pulley statistic with DDP support and 𝒪 ⁡ ( N ) \mathcal{O}(N) time and memory complexity. x is a (N, K) tensor, num_slices is | 𝔸 | |{\mathbb{A}}| in Section 4.2 , ‘global_step‘ is used for sync. sampling across GPUs and can be omited for single-GPU training. An optimized implementation with caching is also provided in our official codebase, computation times provided in Table 6 .

By the chain rule, Section 4.2.3 directly gives ‖ ∇ θ E ​ P ​ ( 𝐚 ) ‖ ≤ 4 ​ σ 2 N ​ ∑ i = 1 N ‖ 𝐚 ⊤ ​ ∇ θ f θ ​ ( 𝐱 i ) ‖ \left\|\nabla_{\theta}EP(\mathbf{a})\right\|\leq\frac{4\sigma^{2}}{N}\sum_{i=1}^{N}\|\mathbf{a}^{\top}\nabla_{\theta}f_{\theta}(\mathbf{x}_{i})\| , providing stable gradients. The limitations of moment-based and CDF-based tests coupled with Section 4.2.3 justifies our choice of the ( Epps–Pulley ): (i) DDP-friendly and scalable, (ii) uniformly bounded gradients and curvature regardless of input distribution, and (iii) hyper-parameter free implementation. Lastly, we highlight that our implementation has a linear memory and computational complexity of 𝒪 ⁡ ( N ) \mathcal{O}(N) , with N N the minibatch size . The implementation of SIGReg using that statistical test is provided in , along with computation times of the forward-backward pass in Table 6 .

As a last step before introducing LeJEPA, we ought to study the requirements on the number of directions ( | 𝔸 | |{\mathbb{A}}| ) for ( 4.2 ) to be effective in high-dimension.

### 4.3 How SIGReg Beats the Curse of Dimensionality

This last section seeks to characterize how many slices in 𝔸 {\mathbb{A}} one must sample for ( 5 ) to be an effective statistical test. That design is crucial if we hope for LeJEPA to successfully converge towards isotropic Gaussian embeddings.

#### Smoothness Beats the Curse of Dimensionality

Our first argument arguing for a favorable scaling of | 𝔸 | |{\mathbb{A}}| with the embedding dimension K K relies on the smoothness of P 𝜽 P_{{\bm{\theta}}} as measured by its Sobolev regularity α \alpha ( Adams and Fournier, 2003 ) . We formalize below a bound on the directional test from Equation 3 over all possible directions 𝒂 {\bm{a}} when the test statistic is minimized over | 𝔸 | = M |{\mathbb{A}}|=M directions. While we provide bounds on the expected discrepancy over random directions 𝒂 {\bm{a}} when the EP test is satisfied (equals zero) on a finite set of directions, the provided proof includes the case of moment-based and CDF-based tests as well.

As | 𝔸 | → ∞ |{\mathbb{A}}|\to\infty , the bound decays as | 𝔸 | − 2 α / ( K − 1 ) |{\mathbb{A}}|^{-2\alpha/(K-1)} , showing that | 𝔸 | = O ⁡ ( K ) |{\mathbb{A}}|=O(K) directions suffice for ϵ \epsilon -approximation when α \alpha is large. Some examples of embedding densities with varying α \alpha are provided in Figure 4 . The following statement characterizes how the M M directions actually constrain the entire space as a function of α \alpha . The constant C ⁡ ( K , α ) = 2 2 ​ α ​ π ( K − 1 ) / 2 ​ Γ ​ ( α + K − 1 2 ) ( K − 1 ) ​ Γ ​ ( α ) ​ Γ ​ ( K − 1 2 ) C(K,\alpha)=\frac{2^{2\alpha}\pi^{(K-1)/2}\Gamma\left(\alpha+\frac{K-1}{2}\right)}{(K-1)\Gamma(\alpha)\Gamma\left(\frac{K-1}{2}\right)} is visualized in Figure 15 (left) depicting how α \alpha and | 𝔸 | |{\mathbb{A}}| interact. In words, we obtain that thanks to the natural smoothness of DN–either stemming from the architecture or the implicit and explicit regularizers used during training–applying SIGReg on | 𝔸 | |{\mathbb{A}}| directions can be sufficient to tightly constrain the entire space. We note that considering the worst case over 𝒂 {\bm{a}} or using low-discrepancy sequences for 𝒂 {\bm{a}} does not impact the asymptotic bounds, details provided in Appendix D .

#### SGD Beats the Curse of Dimensionality

Our second argument leverages the iterative nature of DN training. Although we may use only | 𝔸 | |{\mathbb{A}}| to be a few hundreds, the cumulative number of sampled directions grows linearly with training time. This resampling effect (illustrated in Figure 7 , bottom) enables rapid convergence. Even small | 𝔸 | |{\mathbb{A}}| achieves tight distributional matching compared to keeping the set 𝔸 {\mathbb{A}} fixed throughout minibatches (recall Section 4.3 ). Our experiments show that even with | 𝔸 | |{\mathbb{A}}| as low as 16 16 can easily outperform a fixed set with | 𝔸 | |{\mathbb{A}}| of order of thousands thanks to the compounding effect of resampling at each minibatch.

#### Empirical Validation on Synthetic Data

We conclude this section with a controlled experiment applying ( 5 ) with gradient-based training to produce isotropic embeddings. In this setup, we directly consider embeddings 𝒁 {\bm{Z}} which we will differentiate and optimized to minimize ( 5 ). By directly optimizing the embeddings we are able to observe the impact of the loss without any possible constraint and regularization that would come from the architecture. We sample N N i.i.d. samples 𝒙 n {\bm{x}}_{n} in a D D -dimensional space. This sampling is based on an isotropic Gaussian distribution–but the first two dimensions are again set to the adversarial “X” shape. That is, among the D D dimensions, only two must be transformed as all the other ones already obey the isotropic Gaussian target distribution. We then make the samples 𝒙 n {\bm{x}}_{n} differentiable and optimize then to minimize the value of the different statistical tests compute on M M random M M random directions. Those directions are resampled after each gradient step–which follows the procedure we will employ in LeJEPA. We present the results in Figure 6 demonstrating that even in challenging case, i.e., D = 512 D=512 and M = 16 M=16 , SIGReg is able to detect the two degenerate dimensions and unfold them back to how they should look like under the target distribution.

## 5 LeJEPA: Stable and Scalable Implementation

Having established that isotropic Gaussians are the optimal embedding distribution for foundation models ( Section 3 ) and introduced SIGReg to achieve this distribution ( Section 4.2 ), we now present the complete LeJEPA framework. We first evaluate candidate statistical tests ( Sections 4.2.1 and 4.2.2 ) and identify characteristic function-based tests as optimal for gradient-based training ( Section 4.2.3 ). The full LeJEPA implementation follows in Section 5.1 .

### 5.1 LeJEPA: SIGReg + Prediction Loss

LeJEPA implementation–works out-of-the-box on any dataset, with DDP, with any backbone, e.g., torchvision or timm. For non-ViT architectures (e.g., ResNet), set global_views = all_views. We use bs for the minibatch size, SIGReg is from .

We now discuss the implementation of LeJEPA starting with SIGReg and followed by the prediction and total losses.

The SIGReg Loss. We chose ( Epps–Pulley ) for its provable boundedness ( Section 4.2.3 ) and its scalability. Its implementation follows exactly the equation except for the integrate which is estimated using a quadrature approximation. We find that the simple trapezoidal quadrature rule is sufficient even with as few knots as 17 17 , as ablated in Figure 20 . In particular, we leverage the symmetry of the integrand to double the number of knots for free, see the official code. On the other hand, the use of minibatches introduces a bias vanishing at rate 𝒪 ⁡ ( 1 / N ) \mathcal{O}(1/N) , as formalized below.

Hence, the gradients we obtain from using ( Epps–Pulley ) are biased by an explicit 𝒪 ⁡ ( 1 / N ) \mathcal{O}(1/N) term. We found this bias to be minimal and not a concern even for minibatches as small as 16. Unbiased alternatives include using U-statistic debiasing of | ϕ θ | 2 |\phi_{\theta}|^{2} or sample splitting, which we do not explore in this study. Our final implementation of the SIGReg term with Epps-Pulley statistic is provided in .