Expanding SPHERE-JEPA: A Family of Statistical Regularizers for the Hypersphere

2026-06-16

AI Summary

This work addresses the instability in existing self-supervised learning methods that rely on slice-based regularization via random one-dimensional projections, which introduces high gradient variance. The authors propose a full-dimensional statistical regularization objective directly defined on the unit hypersphere, eliminating the need for stochastic projection approximations. They establish, for the first time, an analytical equivalence between slice-based regularization and Maximum Mean Discrepancy (MMD), and introduce deterministic regularizers based on MMD, Kernelized Stein Discrepancy (KSD), and KL divergence. Leveraging spectral theory, they construct rotation-invariant kernels—Heat and Bandlimited—to enable unbiased and stable distribution matching over the hypersphere. Experiments on ImageNet and Galaxy10 demonstrate faster convergence, improved training stability, and superior performance. Notably, different statistical criteria induce distinct representation geometries, with KL-based regularization achieving the best results in texture retrieval tasks.

Abstract

In Self-Supervised Learning (SSL), preventing representation collapse by explicitly enforcing a uniform distribution on the unit hypersphere has proven to be effective. However, current frameworks typically rely on sliced statistical regularizers such as SIGReg (used in LeJEPA) and SUSReg (used in SPHERE-JEPA), which approximate this continuous objective via Monte Carlo sampling along random 1D directions. This stochasticity injects projection variance into the training gradients, destabilizing optimization, and hindering convergence. In this work, we first show that analytically integrating out these random projections natively yields a deterministic Maximum Mean Discrepancy (MMD), bypassing the variance of sliced methods. Motivated by this equivalence, we formulate full-dimensional objectives for MMD, Kernel Stein Discrepancy (KSD), and Kullback-Leibler (KL) divergence directly on the sphere to enforce a uniform distribution. To prevent spatial bias, we equip these tests with rotationally invariant kernels constructed via spectral theory, systematically evaluating two canonical families: smooth exponential decay (Heat) and strict frequency cutoff (Bandlimited) filters. Empirically, removing projection-induced noise results in more stable optimization, faster convergence, and consistent improvements over stochastic sliced regularizers on ImageNet and Galaxy10. Furthermore, we reveal that the choice of the statistical test shapes the geometry of the learned latent space: MMD and KSD favor locally clustered organization suitable for object-centric domains, whereas the continuous KDE-based KL divergence promotes fine-grained instance separation, yielding the strongest results on unclustered procedural texture retrieval.

Problem

• Self-Supervised Learning
• Representation Collapse
• Statistical Regularization
• Hypersphere
• Projection Variance

Innovation

• Maximum Mean Discrepancy
• rotationally invariant kernels
• statistical regularizers
• hypersphere uniformity
• self-supervised learning