LeJEPA: Provable and Scalable Self-Supervised Learning Without the Heuristics — part 5/10
Cleaned-text code points [60506, 74849). Bibliography omitted here; source.md includes it.
A hallmark of successful self-supervised learning is the emergence of semantically meaningful attention patterns without explicit supervision ( Caron et al., 2021 ) . To assess whether LeJEPA learns such structure, we visualize the attention maps of the learned representations. Following DINO ( Caron et al., 2021 ) , we apply PCA to the embeddings and visualize the first principal components, which reveal clear correspondence to object boundaries and salient regions ( Figure 14 ). Furthermore, we explore whether these attention patterns can enable unsupervised video segmentation—a challenging task requiring temporal consistency and object understanding. By thresholding the self-attention maps of the [CLS] token, we obtain binary masks that track objects across frames without any segmentation labels during training. As shown in Figure 13 , LeJEPA’s attention naturally segments foreground objects from background with remarkable temporal coherence , suggesting that the learned representations capture both spatial semantics and temporal structure. This emergent capability demonstrates that LeJEPA’s stability-focused objective does not sacrifice the semantic richness of learned features.
## 7 Conclusion
We have established a principled theoretical framework for JEPA-based self-supervised learning that fundamentally resolves its core pathologies. Our contributions span theory and practice: we proved that isotropic Gaussian embeddings uniquely minimize worst-case downstream risk, introduced SIGReg as a tractable and provably correct method to enforce this distribution, and demonstrated that this approach eliminates representational collapse by design–and not through ad-hoc combinations of teacher-student networks, stop-gradients, or asymmetric architectures.
We validate LeJEPA across domains and over 60 60 architectures including gigantic versions with 1.8B parameters. In spite of its simplicify , LeJEPA matches state-of-the-art performance while requiring fewer than 50 lines of core implementation. Critically, our approach provides what SSL has long needed: a mathematically rigorous foundation that directly informs practical algorithm design.
## Acknowledgments
We would like to thank Mike Rabbat and Lucas Maes for providing valuable feedbacks on the manuscript.
## Appendix A Additional Details on Nonlinear Probing
### A.1 kNN Probing
To allow for more flexible evaluation of the pretrained encoder f 𝜽 f_{{\bm{\theta}}} , it is standard to work with a k k -NN prober [ Taunk et al., 2019 ] , both for regression and classification. We rely on the radial k k -NN variation that leverages a sample-dependent k k –improving performance for non uniform distributions of samples [ Sun and Huang, 2010 , Zhang et al., 2017 , Abu Alfeilat et al., 2019 ] .
We denote the underlying embedding density as p z ∈ C 3 p_{z}\in C^{3} with derivatives of order up to 3 3 bounded, and finite Fisher information and covariance. This regularity condition is fulfilled by current encoders. The unknown labels come from the target function η : ℝ K → ℝ \eta:\mathbb{R}^{K}\to\mathbb{R} , assumed C 2 C^{2} . We handle classification tasks by setting η ( 𝒛 ) = ℙ ( Y = 1 ∣ 𝒛 ) \eta({\bm{z}})=\mathbb{P}(Y=1\mid{\bm{z}}) . The training consists of the N N embeddings along with their training labels { ( 𝒛 n , η ( 𝒛 n ) ) } n = 1 N \{({\bm{z}}_{n},\eta({\bm{z}}_{n}))\}_{n=1}^{N} , where we will denote 𝒚 n \triangleq η ( 𝒛 n ) {\bm{y}}_{n}\triangleq\eta({\bm{z}}_{n}) . The prediction for a query vector 𝒒 {\bm{q}} is formed as 𝒚 ^ ( 𝒒 ) := 1 𝒚 ( 𝒒 ) ∑ n : ‖ 𝒛 n − 𝒒 ‖ ≤ r 0 𝒚 n , \displaystyle\widehat{{\bm{y}}}({\bm{q}}):=\frac{1}{{\bm{y}}({\bm{q}})}\sum_{n:\left\lVert{\bm{z}}_{n}-{\bm{q}}\right\rVert\leq r_{0}}{\bm{y}}_{n}, (kNN) with 𝒚 ( 𝒒 ) \triangleq # { n : ‖ 𝒛 n − 𝒒 ‖ ≤ r 0 } {\bm{y}}({\bm{q}})\triangleq\#\{n:\left\lVert{\bm{z}}_{n}-{\bm{q}}\right\rVert\leq r_{0}\} counting the number of samples within a r r -radius ball around 𝒒 {\bm{q}} . The radius r r controls how many neighbors predictions are averaged to form the query’s prediction. As per the linear probing’s Section 3.1 , we can characterize the bias of the estimator Equation kNN at a particular query point, as formalized below.
To obtain the integrated bias, i.e., over the distribution of query points, we consider the following two properties. First, the distribution of query points follow the training distribution, i.e., 𝒒 ∼ p z {\bm{q}}\sim p_{z} , second, target function η \eta has gradient which is mean-zero and isotropic with 𝔼 [ ∇ η ( 𝒛 ) ∇ η ( 𝒛 ) ⊤ ] = τ g 2 I d \mathbb{E}\big[\nabla\eta({\bm{z}})\nabla\eta({\bm{z}})^{\top}\big]=\tau_{g}^{2}I_{d} with τ g 2 ∈ ( 0 , ∞ ) \tau_{g}^{2}\in(0,\infty) uniformly in 𝒛 {\bm{z}} . We also have any finite scalar-constraint on the covariance of the embeddings such as Tr ( Σ ) = c \Tr(\Sigma)=c or ‖ Σ ‖ F = c \|\Sigma\|_{F}=c for a finite constant c c .
As a result, we now have a unique minimizer for the optimal embedding density for both the linear and k-NN probes.
### A.2 Kernel Probing
As an alternative to ( kNN ), it is also common to leverage kernel methods, which we consider in this section.
Consider a kernel K : ℝ K → ℝ K:\mathbb{R}^{K}\to\mathbb{R} with the following standard properties ∫ ℝ d K ( u ) 𝑑 u = 1 , \displaystyle\int_{\mathbb{R}^{d}}K(u)du=1, (normalized) ∫ ℝ d u K ( u ) 𝑑 u = 0 , \displaystyle\int_{\mathbb{R}^{d}}uK(u)du=0, (symmetric) ∫ ℝ d u u ⊤ K ( u ) 𝑑 u = μ 2 ( K ) I d , \displaystyle\int_{\mathbb{R}^{d}}uu^{\top}K(u)du=\mu_{2}(K)I_{d}, (isotropic) R ( K ) \triangleq ∫ ℝ d K ( u ) 2 𝑑 u < ∞ , \displaystyle R(K)\triangleq\int_{\mathbb{R}^{d}}K(u)^{2}du<\infty, (finite roughness) for some μ 2 ( K ) ∈ ( 0 , ∞ ) \mu_{2}(K)\in(0,\infty) , some bandwidth h > 0 h>0 and denoting K h ( t ) \triangleq h − d K ( t / h ) K_{h}(t)\triangleq h^{-d}K(t/h) , we remind the reader that the Nadaraya-Watson estimator, introduced in Nadaraya [1964] , Watson [1964] , at a query 𝒒 ∈ ℝ d {\bm{q}}\in\mathbb{R}^{d} is 𝒚 ^ ( 𝒒 ) \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{x}}_{n}){\bm{y}}_{n}}{\sum_{n=1}^{N}K_{h}({\bm{q}}-{\bm{x}}_{n})}. (NW) Similarly to ( kNN ), we will see that the performance of ( NW ) depends crucially on the distribution of the training points. We have access to our dataset of inputs from p z p_{z} and for each sample 𝒛 n {\bm{z}}_{n} the corresponding target is given from η ( 𝒛 n ) = 𝔼 [ Y n ∣ 𝒛 n ] \eta({\bm{z}}_{n})=\mathbb{E}[Y_{n}\mid{\bm{z}}_{n}] . We also denote the corresponding conditional variance of the target function at that point as v ( x ) = Var ( Y i ∣ X i = x ) v(x)=\mathrm{Var}(Y_{i}\mid X_{i}=x) . We follow the regularity conditions of the k-NN probing derivations and additionally assume that p p has sufficiently light tails so that for each coordinate j j , lim ‖ x ‖ → ∞ p ( x ) = 0 \lim_{\|x\|\to\infty}p(x)=0 and lim ‖ x ‖ → ∞ x j p ( x ) = 0 \lim_{\|x\|\to\infty}x_{j}p(x)=0 . We first derive the pointwise bias and variance for 𝒚 ^ ( 𝒒 ) \widehat{\bm{y}}({\bm{q}}) .
We now show that, under a fixed mean and total-covariance constraint on p z p_{z} , the isotropic Gaussian distribution uniquely minimizes the bias and variance of the kernel regression estimator at any test point. We restrict the smoothness class of the target function using ℳ ( L , B ) \triangleq { m ∈ C 2 ( ℝ d ) : ∥ ∇ 𝒚 ( 𝒒 ) ∥ ≤ L , | Δ 𝒚 ( 𝒒 ) | ≤ B , ∀ 𝒒 ∈ ℝ d } , \mathcal{M}(L,B)\triangleq\Big\{m\in C^{2}(\mathbb{R}^{d}):\|\nabla{\bm{y}}({\bm{q}})\|\leq L,\\ |\Delta{\bm{y}}({\bm{q}})|\leq B,\forall{\bm{q}}\in\mathbb{R}^{d}\Big\}, allowing us to formalize below the worst case integrated bias and the optimal density for z z .
## Appendix B Proofs
### B.1 Proof of Section 3.1
###### Proof.
Our proof follows standard derivations when it comes to studying the bias of an estimator. Let’s consider the ridge regression problem (Tikhonov regularized least squares estimator) with close form estimator 𝜷 ^ = ( 𝐗 T 𝐗 + λ wd 𝐈 ) − 1 𝐗 T 𝐘 . \hat{\bm{\beta}}=(\mathbf{X}^{T}\mathbf{X}+\lambda_{\rm wd}\mathbf{I})^{-1}\mathbf{X}^{T}\mathbf{Y}. (11) The labels are formed from the ground truth parameter β true \beta_{\rm true} with centered error, as per 𝐘 = 𝐗 𝜷 true + 𝜺 \mathbf{Y}=\mathbf{X}\bm{\beta}_{\text{true}}+\bm{\varepsilon} where 𝔼 [ 𝜺 ] = 𝟎 \mathbb{E}[\bm{\varepsilon}]=\mathbf{0} . We can now look at the bias of our estimator given by Bias ( 𝜷 ^ ) \displaystyle\text{Bias}(\hat{\bm{\beta}}) = 𝔼 [ 𝜷 ^ ] − 𝜷 true \displaystyle=\mathbb{E}[\hat{\bm{\beta}}]-\bm{\beta}_{\text{true}} = ( 𝐗 T 𝐗 + λ wd 𝐈 ) − 1 𝐗 T 𝐗 𝜷 true − 𝜷 true \displaystyle=(\mathbf{X}^{T}\mathbf{X}+\lambda_{\rm wd}\mathbf{I})^{-1}\mathbf{X}^{T}\mathbf{X}\bm{\beta}_{\text{true}}-\bm{\beta}_{\text{true}} = − λ wd ( 𝐗 T 𝐗 + λ wd 𝐈 ) − 1 𝜷 true \displaystyle=-\lambda_{\rm wd}(\mathbf{X}^{T}\mathbf{X}+\lambda_{\rm wd}\mathbf{I})^{-1}\bm{\beta}_{\text{true}} = − λ wd 𝐐 ( 𝚲 + λ 𝐈 ) − 1 𝐐 T 𝜷 true \displaystyle=-\lambda_{\rm wd}\mathbf{Q}(\bm{\Lambda}+\lambda\mathbf{I})^{-1}\mathbf{Q}^{T}\bm{\beta}_{\text{true}} We will now compare that bias when 𝑿 {\bm{X}} has isotropic and anisotropic covariance with same total variance: λ 1 + λ 2 + ⋯ + λ p p = λ ¯ . \frac{\lambda_{1}+\lambda_{2}+\cdots+\lambda_{p}}{p}=\bar{\lambda}. (12) For any anisotropic covariance matrix of 𝑿 {\bm{X}} , denote by 𝒒 1 {\bm{q}}_{1} the eigenvector with smallest eigenvalue, and let’s denote by κ > 0 \kappa>0 a positive constant. We now define 𝜷 true = κ ⋅ 𝐪 p , \bm{\beta}_{\text{true}}=\kappa\cdot\mathbf{q}_{p}, (13) leading to ‖ Bias ( 𝜷 ^ ) ‖ isotropic = λ wd λ ¯ + λ wd ‖ 𝜷 true ‖ , \displaystyle\|\text{Bias}(\hat{\bm{\beta}})\|_{\text{isotropic}}=\frac{\lambda_{\rm wd}}{\bar{\lambda}+\lambda_{\rm wd}}\|\bm{\beta}_{\text{true}}\|, ‖ Bias ( 𝜷 ^ ) ‖ non-isotropic = λ wd λ p + λ wd ‖ 𝜷 true ‖ \displaystyle\|\text{Bias}(\hat{\bm{\beta}})\|_{\text{non-isotropic}}=\frac{\lambda_{\rm wd}}{\lambda_{p}+\lambda_{\rm wd}}\|\bm{\beta}_{\text{true}}\| Since λ p < λ ¯ \lambda_{p}<\bar{\lambda} (strict inequality when not isotropic): λ wd λ p + λ wd > λ wd λ ¯ + λ wd \frac{\lambda_{\rm wd}}{\lambda_{p}+\lambda_{\rm wd}}>\frac{\lambda_{\rm wd}}{\bar{\lambda}+\lambda_{\rm wd}} we obtain that ‖ Bias ( 𝜷 ^ ) ‖ non-isotropic > ‖ Bias ( 𝜷 ^ ) ‖ isotropic \|\text{Bias}(\hat{\bm{\beta}})\|_{\text{non-isotropic}}>\|\text{Bias}(\hat{\bm{\beta}})\|_{\text{isotropic}}
As a result, whenever the covariance matrix of 𝑿 {\bm{X}} is anisotropic, there will be downstream tasks for which the estimator bias is increased compared to having isotropic covariance matrix. Anisotropic covariance structure thus amplifies regularization bias when the true parameter vector aligns unfavorably with the data’s covariance structure. ∎
### B.2 Proof of Section 3.1
###### Proof.
We use the same formula as in Section B.1 with λ wd = 0 \lambda_{\rm wd}=0 . We first see that the estimator is unbiased. We will now leverage that result to compute the covariance matrix of the estimator Var ( 𝜷 ^ | 𝐗 ) \displaystyle\text{Var}(\hat{\bm{\beta}}|\mathbf{X}) = 𝔼 [ ( 𝜷 ^ − 𝜷 ) ( 𝜷 ^ − 𝜷 ) T | 𝐗 ] \displaystyle=\mathbb{E}[(\hat{\bm{\beta}}-\bm{\beta})(\hat{\bm{\beta}}-\bm{\beta})^{T}|\mathbf{X}] = 𝔼 [ ( 𝐗 T 𝐗 ) − 1 𝐗 T 𝜺 𝜺 T 𝐗 ( 𝐗 T 𝐗 ) − 1 | 𝐗 ] \displaystyle=\mathbb{E}[(\mathbf{X}^{T}\mathbf{X})^{-1}\mathbf{X}^{T}\bm{\varepsilon}\bm{\varepsilon}^{T}\mathbf{X}(\mathbf{X}^{T}\mathbf{X})^{-1}|\mathbf{X}] = ( 𝐗 T 𝐗 ) − 1 𝐗 T 𝔼 [ 𝜺 𝜺 T | 𝐗 ] 𝐗 ( 𝐗 T 𝐗 ) − 1 \displaystyle=(\mathbf{X}^{T}\mathbf{X})^{-1}\mathbf{X}^{T}\mathbb{E}[\bm{\varepsilon}\bm{\varepsilon}^{T}|\mathbf{X}]\mathbf{X}(\mathbf{X}^{T}\mathbf{X})^{-1} = ( 𝐗 T 𝐗 ) − 1 𝐗 T ( σ 2 𝐈 n ) 𝐗 ( 𝐗 T 𝐗 ) − 1 \displaystyle=(\mathbf{X}^{T}\mathbf{X})^{-1}\mathbf{X}^{T}(\sigma^{2}\mathbf{I}_{n})\mathbf{X}(\mathbf{X}^{T}\mathbf{X})^{-1} = σ 2 ( 𝐗 T 𝐗 ) − 1 \displaystyle=\sigma^{2}(\mathbf{X}^{T}\mathbf{X})^{-1} leading to the total variance tr ( Var ( 𝜷 ^ ) ) = σ 2 tr ( 𝐆 − 1 ) = σ 2 ∑ j = 1 p 1 λ j \text{tr}(\text{Var}(\hat{\bm{\beta}}))=\sigma^{2}\text{tr}(\mathbf{G}^{-1})=\sigma^{2}\sum_{j=1}^{p}\frac{1}{\lambda_{j}} where we used the eigendecomposition: 𝐆 = 𝐐 𝚲 𝐐 T \mathbf{G}=\mathbf{Q}\bm{\Lambda}\mathbf{Q}^{T}
The function f ( x ) = 1 x f(x)=\frac{1}{x} is strictly convex on ( 0 , ∞ ) (0,\infty) allowing us to leverage Jensen’s Inequality: 1 K ∑ k = 1 K 1 λ k > 1 1 K ∑ j = 1 K λ k \displaystyle\frac{1}{K}\sum_{k=1}^{K}\frac{1}{\lambda_{k}}>\frac{1}{\frac{1}{K}\sum_{j=1}^{K}\lambda_{k}} ⇔ 1 K ∑ k = 1 K 1 λ k > 1 K ∑ k = 1 K 1 1 K ∑ j = 1 K λ k \displaystyle\iff\frac{1}{K}\sum_{k=1}^{K}\frac{1}{\lambda_{k}}>\frac{1}{K}\sum_{k=1}^{K}\frac{1}{\frac{1}{K}\sum_{j=1}^{K}\lambda_{k}} ⇔ ∑ k = 1 K 1 λ k > ∑ k = 1 K 1 1 K ∑ j = 1 K λ k \displaystyle\iff\sum_{k=1}^{K}\frac{1}{\lambda_{k}}>\sum_{k=1}^{K}\frac{1}{\frac{1}{K}\sum_{j=1}^{K}\lambda_{k}} ⇔ tr ( Var ( 𝜷 ^ ) ) aniso > tr ( Var ( 𝜷 ^ ) ) iso \displaystyle\iff\text{tr}(\text{Var}(\hat{\bm{\beta}}))_{\text{aniso}}>\text{tr}(\text{Var}(\hat{\bm{\beta}}))_{\text{iso}} The inequality is strict whenever the eigenvalues { λ j } j = 1 p \{\lambda_{j}\}_{j=1}^{p} are not all equal. ∎
### B.3 Proof of Section A.1
###### Proof.
Under PPP, conditional expectations of η ^ ( x ) \widehat{\eta}(x) coincide with the normalized ball average 𝔼 [ η ^ ( x ) ] = ∫ B ( 0 , r 0 ) η ( x + z ) p ( x + z ) 𝑑 z ∫ B ( 0 , r 0 ) p ( x + z ) 𝑑 z to second order in r 0 , \mathbb{E}\big[\widehat{\eta}(x)\big]\;=\;\frac{\int_{\mathrm{B}(0,r_{0})}\eta(x+z)p(x+z)dz}{\int_{\mathrm{B}(0,r_{0})}p(x+z)dz}\quad\text{to second order in }r_{0}, which is the key surrogate used below.
Ball integrals. For computations we use (by symmetry) for any r > 0 r>0 : ∫ B ( 0 , r ) z 𝑑 z = 0 , ∫ B ( 0 , r ) z z ⊤ 𝑑 z = Vol d + 2 d + 2 I d , ∫ B ( 0 , r ) ‖ z ‖ 2 𝑑 z = d Vol d + 2 d + 2 . \int_{\mathrm{B}(0,r)}zdz=0,\qquad\int_{\mathrm{B}(0,r)}zz^{\top}dz=\frac{{\rm Vol}^{d+2}}{d+2}I_{d},\qquad\int_{\mathrm{B}(0,r)}\left\lVert z\right\rVert^{2}dz=\frac{d{\rm Vol}^{d+2}}{d+2}.