Search papers, labs, and topics across Lattice.
This paper extends the theoretical analysis of Joint Embedding Predictive Architectures (JEPA) to Riemannian manifolds, specifically the hypersphere, to determine the optimal geometry for learned representations in self-supervised learning. They demonstrate that uniform distributions on manifolds, particularly the uniform distribution on the sphere, are optimal for k-NN and kernel ridge regression under a worst-case formulation. Based on this, they introduce SPHERE-JEPA, a modified JEPA framework that enforces hyperspherical uniformity, achieving significant empirical improvements in texture retrieval and ImageNet classification compared to LeJEPA.
Gaussian embeddings might be suboptimal for self-supervised learning: enforcing hyperspherical uniformity in learned representations yields substantial gains in texture retrieval and ImageNet classification.
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).
