SPHERE-JEPA: Spherical Prediction with Homogeneous Embeddings

2026-05-26

AI Summary

This work investigates the optimal geometric structure of representations in self-supervised learning when constrained to Riemannian manifolds, with a focus on the hypersphere, and highlights the limitations of Gaussian embeddings on low-dimensional manifolds. Through a minimax analysis, it establishes the theoretical property that k-nearest neighbors and kernel ridge regression induce uniform distributions on the hypersphere, thereby extending the study of optimal representation geometry—previously limited to Euclidean spaces—to Riemannian manifolds for the first time. Building on these insights, the authors propose SPHERE-JEPA, a novel framework that enforces hyperspherical uniformity via Cramér–Wold projection mechanisms combined with explicit uniformity constraints. Experiments demonstrate that SPHERE-JEPA improves linear probe accuracy by 1.8% on ImageNet-1K, achieves over 6% gains in texture retrieval mAP, and matches or outperforms LeJEPA across multiple benchmarks.

Abstract

A fundamental open question in self-supervised learning (SSL) is the explicit characterization of the optimal geometry of the learned representations. Recently, LeJEPA identified isotropic Gaussian embeddings as optimal for minimizing downstream prediction risk in Euclidean spaces. However, the corresponding problem for distributions supported on lower-dimensional manifolds, such as the hypersphere, remains unexplored. In this work, we demonstrate that extending this minimax analysis to smooth distributions on Riemannian manifolds fundamentally changes the optimal solution. We show that, under a worst-case formulation, both k-nearest neighbors and kernel ridge regression induce hyperspherical uniformity. More precisely, we show that uniform distributions on manifolds are optimal for k-nearest neighbors, and that the uniform distribution on the sphere is optimal for kernel ridge regression with both the exponential dot-product kernel and the linear kernel. This theoretical insight reveals a fundamental limitation of Gaussian embeddings: their non-uniform density induces anisotropic k-NN neighborhoods, severely biasing the estimator. To correct this, we introduce SPHERE-JEPA, a theoretically grounded SSL framework. We adapt LeJEPA's Cram{é}r-Wold projection mechanism to enforce hyperspherical uniformity rather than a Gaussian prior. Empirically, SPHERE-JEPA yields significant improvements, boosting texture retrieval mAP by over 6%, while consistently matching or outperforming LeJEPA on standard benchmarks-including a +1.8% linear probing gain on ImageNet-1K (ViT-B/14).

Problem

• self-supervised learning
• representation geometry
• hyperspherical uniformity
• Riemannian manifolds
• Gaussian embeddings

Innovation

• self-supervised learning
• spherical uniformity
• Riemannian manifolds
• representation geometry
• SPHERE-JEPA