Skip to content

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

Year
Algorithm
LeJEPA
Identifiers

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 approx50 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}).

Similar Papers

5
74%

A Cookbook of Self-Supervised Learning

('Randall Balestriero', 'Mark Ibrahim', 'Vlad Sobal', 'Ari Morcos', 'Shashank Shekhar', 'Tom Goldstein', 'Florian Bordes', 'Adrien Bardes', 'Gregoire Mialon', 'Yuandong Tian', 'Avi Schwarzschild', 'Andrew Gordon Wilson', 'Jonas Geiping', 'Quentin Garrido', 'Pierre Fernandez', 'Amir Bar', 'Hamed Pirsiavash', 'Yann LeCun', 'Micah Goldblum') / 2023