Weak-SIGReg: Covariance Regularization for Stable Deep Learning

TL;DR

Weak-SIGReg introduces a covariance regularization technique that stabilizes training of deep neural networks, especially Vision Transformers, by preventing representation collapse and improving convergence without architectural modifications.

Contribution

The paper adapts Sketched Isotropic Gaussian Regularization into Weak-SIGReg, a computationally efficient covariance regularization method for supervised learning stabilization.

Findings

Recovered ViT training on CIFAR-100 from collapse to high accuracy

Significantly improved convergence of deep MLPs trained with SGD

Demonstrated effectiveness without architectural hacks

Abstract

Modern neural network optimization relies heavily on architectural priorssuch as Batch Normalization and Residual connectionsto stabilize training dynamics. Without these, or in low-data regimes with aggressive augmentation, low-bias architectures like Vision Transformers (ViTs) often suffer from optimization collapse. This work adopts Sketched Isotropic Gaussian Regularization (SIGReg), recently introduced in the LeJEPA self-supervised framework, and repurposes it as a general optimization stabilizer for supervised learning. While the original formulation targets the full characteristic function, a computationally efficient variant is derived, Weak-SIGReg, which targets the covariance matrix via random sketching. Inspired by interacting particle systems, representation collapse is viewed as stochastic drift; SIGReg constrains the representation density towards an isotropic Gaussian, mitigating this drift. Empirically, SIGReg recovers the training of a ViT on CIFAR-100 from a collapsed 20.73\% to 72.02\% accuracy without architectural hacks and significantly improves the convergence of deep vanilla MLPs trained with pure SGD. Code is available at \href{https://github.com/kreasof-ai/sigreg}{github.com/kreasof-ai/sigreg}.