LeJEPA: Provable and Scalable Self-Supervised Learning Without the Heuristics

Best AI papers explained - A podcast by Enoch H. Kang

Podcast artwork

Categories:

This paper introduces a novel self-supervised learning framework designed to resolve the pervasive issue of representation collapse in existing Joint-Embedding Predictive Architectures (JEPAs). It establishes a theoretical foundation by proving that an isotropic Gaussian distribution is the optimal embedding distribution for minimizing the worst-case risk across various downstream tasks. To enforce this optimal distribution, the paper proposes SIGReg (Sketched Isotropic Gaussian Regularization), a scalable method that uses directional statistical tests, specifically recommending the Epps-Pulley test, to match the empirical feature distribution to the target Gaussian. The core contribution is the resulting LeJEPA loss function, which combines the standard JEPA prediction objective with SIGReg, effectively eliminating the need for complex anti-collapse heuristics like stop-gradients or teacher-student networks, and demonstrating robust, state-of-the-art performance with significantly reduced training complexity.