Self supervised learning improved with local geometric embedding methods
LocaLeJEPA: Maximum Entropy Representations under Geometric Priors
Machine LearningComputer Vision and Pattern Recognition
Summary
Learning useful features from data without labels is hard because models can find shortcuts that don’t capture meaningful information. Previous methods tried to fix this by pushing the model’s internal signals into simple shapes like round clouds in regular space. The authors propose a new way that respects more complex shapes that data might naturally form, like curved or symmetrical spaces. They show this new method works better on standard image recognition tasks, especially when using special types of neural networks that understand geometry better.
What this means in practice
- •For computer vision engineers: Improve training of image recognition models on large datasets by using geometric-aware self-supervised learning techniques for better feature quality.
- •For machine learning infrastructure teams: Deploy scalable regularization methods that maintain stable and efficient model training with geometric embedding constraints.
Authors
Kartik Tandon, Alessio Borgi, Pavel Snopov, Inés García-Redondo, Davide D'Ascenzo, Sebastiano Cultrera di Montesano, Fabrizio Silvestri, Pietro Liò, Francesca Dominici, Randall Balestriero, Claudio Battiloro
Abstract
Joint-Embedding Predictive Architectures (JEPA) offer an appealing approach to self-supervised learning, but are typically prone to representation collapse. LeJEPA resolves this by introducing a theoretically grounded distribution-matching objective that drives the embedding towards an isotropic Gaussian. However, this approach is fundamentally tied to Euclidean assumptions, and its global target does not reflect domain-specific geometric priors, such as the use of non-Euclidean embedding geometries, e.g., equivariant CNNs or hyperbolic ViTs. We thus introduce LocaLeJEPA, a framework motivated by a manifold-theoretic reformulation of LeJEPA's entropy-maximization property. The accompanying LocaLeReg, a novel locally Euclidean distribution-matching regularizer, enables a stable and scalable implementation of the LocaLeJEPA framework. We then perform extensive experiments in the vision setting on models trained on ImageNet-1K. LocaLeJEPA significantly outperforms LeJEPA when paired with hyperbolic ViT-S/8 and equivariant ResNet-50 backbones, and also with ordinary Euclidean ViT-S/8 and ViT-L/14 backbones. This holds both in terms of linear and k-NN classification performance on ImageNet and on a series of downstream tasks tested. We view the LocaLeJEPA formalism as a principled approach to the JEPA paradigm that better respects intrinsic data geometry while also accommodating generalizations to exotic embedding geometries. Code available at https://github.com/alessioborgi/LocaLeJEPA.