Sketched Isotropic Gaussian Regularization (SIGReg)

• Sketched Isotropic Gaussian Regularization (SIGReg) is a method that enforces high-dimensional representations to follow an isotropic Gaussian law using randomized sketching.
• It employs random projections and univariate goodness-of-fit tests, such as the Epps–Pulley statistic, to approximate multivariate distribution matching with computational efficiency.
• SIGReg achieves optimal risk minimization in both regularized regression and self-supervised learning, enabling stable, collapse-free training in large-scale settings.

Sketched Isotropic Gaussian Regularization (SIGReg) is a statistical regularization technique designed to constrain high-dimensional representations—whether regression coefficients in sketched ridge regression or neural embeddings in self-supervised learning—to match the law of an isotropic Gaussian. As a unifying principle, SIGReg serves two complementary domains: (1) efficient solution of regularized linear least squares via randomized sketching and (2) provably optimal embedding regularization in high-dimensional self-supervised learning, where it arises as the unique approach to minimize worst-case downstream risk, both linear and nonlinear. SIGReg achieves its effect by efficiently approximating the match-to-Gaussian constraint through randomized projections (sketches) and 1-D statistics, offering scalability, numerical stability, rigorous theoretical guarantees, and practical ease of deployment in large-scale or distributed settings (Meier et al., 2022, Balestriero et al., 11 Nov 2025).

1. Mathematical Foundation and Formulation

SIGReg generically aims to enforce that a vector-valued variable $z \in \mathbb{R}^K$ (e.g., regression solution, neural embedding) obeys $z \sim \mathcal{N}(0, I_K)$. The regularization term measures the divergence or difference between the empirical law of $\{z_i\}_{i=1}^N$ and the isotropic Gaussian target $Q = \mathcal{N}(0, I_K)$.

To render this approach scalable, full multivariate goodness-of-fit is replaced by testing along $M$ random directions $\{a_m\}_{m=1}^M \subset \mathbb{S}^{K-1}$ (the unit sphere), leveraging the Cramér–Wold theorem, which states that two distributions in $\mathbb{R}^K$ coincide if and only if all one-dimensional projections coincide. For a batch $\{z_i\}_{i=1}^N$, given a univariate goodness-of-fit statistic $T$, the SIGReg objective is

$$\mathrm{SIGReg}_T(\theta; \{z_i\}_{i=1}^N) = \frac{1}{M} \sum_{m=1}^M T(\{a_m^\top z_i\}_{i=1}^N).$$

In self-supervised learning, the effect is instantiated by projecting the embeddings onto random slices, computing for each slice the Epps–Pulley (EP) statistic, which compares the empirical characteristic function to that of the standard normal via a weighted $L^2$ distance:

$$\mathrm{EP} = N \int_{\mathbb R} |\widehat{\phi}_z(t) - \phi_Q(t)|^2 w(t)\,dt,$$

where $$\widehat{\phi}_z(t) = \frac{1}{N}\sum_{i=1}^N e^{it a_m^\top z_i}$$, $\phi_Q(t) = e^{-t^2/2}$, and $w(t)$ is a Gaussian window.

The regularization term enters the training or optimization objective as a weighted sum with a trade-off parameter $\lambda>0$:

$$\mathcal{L}_\text{total} = \mathcal{L}_\mathrm{pred} + \lambda\,\mathrm{SIGReg}.$$

2. Optimality of the Isotropic Gaussian Constraint

The rationale for enforcing an isotropic Gaussian distribution is established both for linear and nonlinear downstream tasks. In linear probing, when using ridge regression, anisotropic embedding covariance (i.e., unequal eigenvalues) increases bias and variance relative to the isotropic case. Formally, OLS with Tikhonov regularization

$$\hat{\beta} = \arg\min_{\beta \in \mathbb{R}^K} \|Y - X\beta\|_2^2 + \lambda\|\beta\|_2^2$$

demonstrates that, whenever $\lambda > 0$, anisotropy strictly increases OLS bias, and variance is minimized only for isotropic covariance.

For nonlinear probing (e.g., $k$NN, kernel smoothing), the leading Integrated Squared Bias (ISB) term depends on the Fisher information of the embedding density $J(p) = \int \|\nabla \log p\|^2 p$. Among all densities with equal trace covariance, $J(p)$ is minimized by the isotropic Gaussian, thus minimizing ISB as well. This establishes the necessity and sufficiency of the isotropic Gaussian law for minimizing worst-case prediction risk across a broad class of downstream tasks.

3. Sketching and Randomized Projection Methodology

Full multivariate distribution matching (e.g., via Maximum Mean Discrepancy or Wasserstein metrics) is computationally expensive, scaling quadratically in $N$ or worse in $K$. SIGReg achieves computational efficiency by sketching: sampling $M$ random unit directions, projecting the high-dimensional data, and applying univariate statistical tests.

Averaging the univariate statistics across $M$ random directions approximates the full multivariate constraint. Theoretical results establish that the average error over random slices decays as $M^{-2\alpha/(K-1)}$ for Sobolev-$\alpha$ smooth densities, so $M = O(K)$ suffices for high-dimensional fidelity, and, with fresh sampling per minibatch (e.g., in SGD), coverage improves rapidly in practice. Pseudorandom generation of directions synchronized across devices (by seeding with global step) ensures consistency in distributed environments.

4. Algorithms and Implementation

SIGReg is highly efficient in large-scale optimization:

• Complexity: The core cost is $O(NKM)$ for matrix multiplication plus $O(NMT)$ for characteristic function computation (per batch of size $N$, embedding dimension $K$, $M$ slices, $T$ quadrature points).
• Distributed Training: Designed to be compatible with PyTorch DDP; the only cross-GPU synchronization is all-reduce on the complex-valued averages for CF computation (shape $M\times T$).
• No Custom Kernels: Relies on GEMM, elementwise complex exponential, and trapezoidal integration; no $\mathcal{O}(N^2)$ or $\mathcal{O}(K^2)$ bottlenecks appear.
• Pseudocode:

def SIGReg(z, global_step, M=512): # z: (N, K) embeddings dev = z.device # 1) sample M random directions synchronized across GPUs g = torch.Generator(device=dev) g.manual_seed(global_step) A = torch.randn(K, M, generator=g, device=dev) A = A / A.norm(dim=0, keepdim=True) # unit-norm columns # 2) project embeddings: z_proj shape (N, M) z_proj = z @ A # (N×K)×(K×M)->(N×M) # 3) compute Epps–Pulley on each of the M slices t = torch.linspace(-5, 5, 17, device=dev) # quadrature nodes w = torch.exp(-0.5*t**2) # gaussian window # empirical CF per slice: (M×N×1) -> (M×T) zt = z_proj.unsqueeze(2) * t ecf = (zt.mul(1j).exp()).mean(dim=0) # gather across GPUs with all_reduce phi0 = torch.exp(-0.5*t**2) # target CF err = (ecf - phi0).abs().square() * w # (M×T) EP = N * torch.trapz(err, t, dim=1) # integrate per slice return EP.mean() # average over M slices

Typical hyperparameters are $\lambda \in [0.01, 0.1]$, $M = 16$–$1024$ slices (modest values suffice for stability), bandwidth $\sigma \in [0.5,2]$, integration domain $[-5,5]$ with $\sim 17$ quadrature nodes, and minibatch size $N\geq 128$.

In the context of linear regression, sketching is also used to accelerate Tikhonov-regularized least squares with strong preconditioning, using a single random projection for all regularization parameters and exploiting the statistical dimension $\mathrm{sd}_\lambda(A)$ when feasible. Two specific variants are:

• SIGReg–Chol: Cholesky-based, robust for arbitrary sketches of size $s \gtrsim n$.
• SIGReg–LR: Low-rank, exploiting $\mathrm{sd}_\lambda(A) \ll n$ for cost-efficient preconditioning.

5. Theoretical Guarantees

SIGReg enjoys mathematically rigorous guarantees in both risk minimization and statistical convergence:

• Consistency: SIGReg is a valid level-$\alpha$ test, attaining power $1$ as $N \to \infty$ (Theorem 4.2, (Balestriero et al., 11 Nov 2025)).
• Gradient & Hessian Bounds: The Epps–Pulley slice statistic retains bounded derivatives:

$$|\partial_{z_i}\mathrm{EP}| \leq \frac{4\sigma^2}{N}, \quad |\partial^2_{z_i}\mathrm{EP}| \leq \frac{C\sigma^3}{N},$$

ensuring gradient stability for all embedding magnitudes.

• Bias in Minibatch Estimators: Bias is $\mathcal{O}(1/N)$ in both loss and gradient, vanishing as batch size increases.
• Approximation Accuracy: For embeddings with Sobolev-$\alpha$ smooth laws, global test error decays as $M^{-2\alpha/(K-1)}$.
• Convergence Rate: In regression, preconditioned LSQR converges to $\epsilon$-accuracy in $O(\log(1/\epsilon))$ iterations, provided the preconditioner achieves $\kappa = O(1)$.

6. Empirical Performance and Practical Implications

• Self-Supervised Learning: Empirical validation in LeJEPA covers more than 10 datasets and 60 architectures. For ImageNet-1k (using a ViT-H/14 backbone), LeJEPA with SIGReg achieves 79% top-1 accuracy in linear evaluation mode.
• Stability and Model Selection: The combined objective's value $$(1-\lambda)\,\mathcal{L}_\mathrm{pred} + \lambda\,\mathrm{SIGReg}$$ correlates strongly ($\rho_s \geq 0.85$) with downstream probe accuracy, and after a simple $\lambda$ rescaling the correlation approaches $99\%$, enabling label-free model selection.
• Collapse-Free Training: SIGReg eliminates the need for heuristic collapse-prevention methods; in practice, no stop-gradient, teacher-student, negative-sampling, or whitening is required, and even very large models remain stable.
• Efficiency: For $N=512$, $K\sim 1024$, $M=512$, $T=17$, the forward and backward SIGReg pass is approximately $0.5$ ms on a V100 GPU.
• Emergent Structure: Embeddings regularized with SIGReg demonstrate robust statistical properties and interpretable structure—e.g., unsupervised foreground-background separation and temporally consistent video segmentation—even with minimal regularization.

7. Broader Connections and Significance

SIGReg arises in two fundamental but previously disparate regimes: as an algorithmic device for randomized algorithms in Tikhonov (ridge) regression (Meier et al., 2022), and as an embedding regularizer for self-supervised representation learning (Balestriero et al., 11 Nov 2025). In both cases, the method removes the need for tuning collapse heuristics or repeated expensive matrix factorizations, yielding efficiency, stability, and predictability. The use of random projection-based sketching, with precise error and risk control, provides a template for scalable distribution matching in high dimensions. The approach offers a principled alternative to both traditional multivariate regularization and patchwork heuristic prevention of feature collapse in modern machine learning systems.