SPHERE-JEPA: Spherical Prediction with Homogeneous Embeddings

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

• The paper introduces SPHERE-JEPA, a self-supervised learning algorithm that enforces uniform spherical embeddings to optimize k-NN and kernel ridge regression tasks.
• It leverages Sketched Uniform Spherical Regularization using Cramér–Wold projections, ensuring isotropic neighborhoods and minimizing estimator bias.
• Empirical results on multiple benchmarks demonstrate improved linear probing accuracy and enhanced transfer learning performance over Gaussian-based methods.

Spherical Prediction with Homogeneous Embeddings: Theoretical and Empirical Advances in SSL

Overview

The paper "SPHERE-JEPA: Spherical Prediction with Homogeneous Embeddings" (2605.26900) addresses a fundamental open problem in self-supervised learning (SSL): the optimal geometric structure for learned representation spaces when evaluated under nonparametric and kernel-based downstream tasks. Building on the minimax analysis of LeJEPA, which proved isotropic Gaussian embeddings are optimal for Euclidean settings, this work generalizes the analysis to embeddings supported on manifolds, notably the hypersphere, and demonstrates that uniform spherical embeddings are optimal for $k$-NN and a broad class of kernel ridge regression protocols. The paper introduces SPHERE-JEPA, an SSL algorithm enforcing hyperspherical uniformity, and offers strong empirical evidence showing substantial performance improvements over isotropic Gaussian-based methods in both standard and nonparametric settings.

Theoretical Contributions

Minimax Optimality Criteria on Manifolds

The authors rigorously study the minimax risk in downstream protocols commonly used in SSL:

• Linear Ridge Regression: As previously established, isotropic (Gaussian) representations are confirmed optimal in Euclidean space under regularized least squares, as they minimize the top eigenvalue of the covariance.
• $k$-Nearest Neighbor Regression on Manifolds: Extending prior analysis to smooth Riemannian manifolds, the paper shows uniform distributions on the support manifold are minimax optimal under worst-case integrated squared bias (ISB) for $k$-NN regression. Gaussian embeddings, with their density concentration, produce anisotropic, density-biased local neighborhoods, damaging the estimator's properties. In contrast, uniform spherical distributions yield isotropic neighborhoods everywhere, eliminating local estimator bias.
• Kernel Ridge Regression with Dot-Product and Linear Kernels: The theoretical results establish that, for exponential dot-product kernels $K(x, y) = \exp(\kappa x^\top y)$ and the linear kernel, uniform distributions on the hypersphere optimally minimize the population worst-case bias.

This geometric insight reveals a limitation of existing Gaussian priors for SSL representation learning in high dimensions and motivates targeting uniform angular embedding structure instead. The central visualization of this density-induced neighborhood anisotropy is shown below.

Figure 1: $k$-NN neighborhoods under Gaussian embeddings are spatially biased and anisotropic; spherical uniformity yields unbiased, isotropic neighborhoods for all query points.

Cramér–Wold Regularization on the Sphere

SPHERE-JEPA develops Sketched Uniform Spherical Regularization (SUSReg), adapting Cramér–Wold projection-based tests to enforce uniform distribution on the hypersphere, as opposed to the univariate Gaussian projections used in LeJEPA. The core insight is that one-dimensional projections of a uniform distribution on the sphere follow a specific beta-like law, which can be efficiently enforced over mini-batches via the Epps–Pulley test. The method relies on matching all projected directions to the target spherical projection distribution:

$$\rho_d(t) = C_d(1 - t^2)^{(d-3)/2}, \qquad t\in (-1,1)$$

This mechanism is robust to high dimensions, converging to the Gaussian as $d$ grows large. The illustrations compare mixtures of non-uniform spherical distributions (von Mises–Fisher) versus uniform, highlighting the regularizer's effect.

Figure 2: Non-uniform vMF mixtures yield projection histograms mismatched from the $d$-sphere law; uniform distributions yield matching projections in all directions, as required by SUSReg.

Figure 3: In high dimensions ($d=256$), the spherical beta law closely matches the normal, supporting Gaussian approximation for SUSReg.

Methodological Framework

SPHERE-JEPA maintains the alignment-based invariance loss of JEPA/LeJEPA but constrains the learned embeddings to the hypersphere and regularizes for global uniformity using SUSReg. Specifically:

• Encoder outputs are explicitly normalized to the unit sphere.
• Global invariant prediction loss aligns each crop/view to the mean prototype of its global views, using an MSE between normalized vectors, which approximates squared geodesic distance.
• Uniform spherical distribution regularization is enforced via Cramér–Wold projections and the Epps–Pulley statistical test, matching the theoretical density for random one-dimensional projections.

The objective is a convex combination of the alignment loss and SUSReg. This construction ensures that both view invariance and optimally uniform spherical geometry are achieved.

Empirical Evaluation

Experiments are conducted on ImageNet-1K/100, Galaxy10, and a controlled procedural texture retrieval benchmark. Architectures include both Vision Transformer (ViT-B/14) and ResNet backbones. Key observations include:

• Linear Probing and $k$-NN: On ImageNet-1K, SPHERE-JEPA produces a +1.8% linear probing accuracy gain compared to LeJEPA under identical ViT-B/14 protocol (see Table 1 in the text). $k$0-NN accuracies remain unaffected or slightly improved, particularly in nonparametric retrieval.
• Transfer Learning: SPHERE-JEPA improves average linear probing performance across all tested transfer datasets (e.g., CIFAR, DTD, Flowers), showing enhanced generalizability of learned features.
• Nonparametric Texture Retrieval: The model achieves over 6% mean average precision (mAP) improvement in a nearest-neighbor texture retrieval task, with gains consistent both at the projection head and backbone layers. The retrieval task probes whether representations are suitable for local geometric reasoning, validating the hypothesized superiority for $k$1-NN under spherical uniformity.

Texture datasets and views for retrieval are illustrated for reference.

Figure 4: Samples from cloud, disk, flake, and wood procedural texture families used in retrieval benchmarks.

View augmentations and their geometric variability are shown:

Figure 5: Left: spatial source regions for random affine transforms. Right: augmented views following geometric and photometric transformations.

Implications and Future Directions

The findings fundamentally advance the understanding of optimal representation geometry in SSL:

• For nonparametric and kernel-based downstream protocols, Gaussian priors are suboptimal in high-dimensional embedding spaces. Uniform spherical geometry not only resolves density-induced estimator bias but is also optimal for dot-product kernel methods central to vision and general representation tasks.
• Projection-based regularization leveraging Cramér–Wold characterization offers an efficient and theoretically sound route for enforcing high-dimensional distributional constraints, with minimal computational overhead.
• Implications for practical SSL: Optimized spherical representations could enhance performance in compositionality, retrieval, and any setting where local geometry is critical. This is likely to become increasingly important as non-linear, kernelized, or nonparametric protocols become more prominent in evaluation and downstream use.

Open questions and future work include comprehensive hyperparameter robustness analysis, extensive benchmarking across multiple random seeds, scaling to larger architectures, and evaluation in dense prediction (segmentation/detection) scenarios. Further explorations on how spherical uniformity interacts with other priors and objectives (e.g., clustering, metric learning), and how it might integrate or compete with contrastive methods in complex multimodal domains, are promising research directions.

Conclusion

SPHERE-JEPA establishes that uniform spherical representations align with the minimax criteria for key SSL evaluation protocols, overcoming the geometric limitations of Gaussian-based priors. Via Sketched Uniform Spherical Regularization, the method enforces this geometry efficiently and demonstrates substantial empirical gains in both linear and nonparametric regimes. This contributes significant theoretical clarification and practical advancement in the design of SSL algorithms, with meaningful consequences for future representation learning methods focusing on optimal data geometry.