LocaLeJEPA: Maximum Entropy Representations under Geometric Priors

Joint-Embedding Predictive Architectures (JEPA) offer a compelling approach to self-supervised learning, but are typically prone to representation collapse. LeJEPA resolves this by introducing a provably robust distribution-matching objective that drives the embedding towards an isotropic Gaussian. However, this approach is fundamentally tied to Euclidean assumptions. Therefore, it is incompatible with the incorporation of 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 robust, practical implementation of the LocaLeJEPA framework. We then perform extensive experiments on models trained on ImageNet-1k. LocaLeJEPA outperforms LeJEPA not only when paired with hyperbolic ViT-S/8 and equivariant ResNet-50 backbones, but even with an ordinary Euclidean ViT-S/8, both in terms of linear and kNN classification performance on ImageNet and on a series of downstream tasks tested. Overall, 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.