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 github.com/kreasof-ai/sigreg.
Figures & tables
Figure 1: Visualizing Optimization Dynamics. (Left) Stochastic Collapse : In standard SGD without normalization, representations (blue dots) tend to collapse into low-dimensional manifolds (red arrows). (Center) Strong-SIGReg : Forces the distribution towards a perfect isotropic sphere. (Right) Weak-SIGReg (Ours) : Targets only the covariance, preventing collapse (green arrows) while allowing more geometric flexibility (star shape), sufficient for supervised stability.
Figure 2: Weak SIGReg Architecture. A high-dimensional batch Z ( N×C ) is projected via a random sketch matrix S ( C×K ) into a lower-dimensional embedding. The covariance of this sketch is computed and forced towards the Identity matrix I . This avoids computing the generic C×C covariance, reducing memory cost to O(CK) .
Optimizer
SIGReg
Top-1 Acc
Status
AdamW
None
20.73%
Collapse
AdamW
Strong (LeJEPA)
70.20%
Converged
AdamW
Weak (Ours)
72.02%
Converged
Table 1: ViT Optimization Stability (CIFAR-100). The baseline collapses while SIGReg converges.
Model Setup
SIGReg
Top-1 Acc
Expert-Tuned Baseline
None
70.76%
Expert-Tuned Baseline
Strong
72.71%
Expert-Tuned Baseline
Weak
71.65%
Table 2: ViT Performance: Expert Tuning vs. SIGReg.
Augmentation
SIGReg
Top-1 Acc
None
None
26.77%
None
Strong
35.99%
None
Weak
42.17%
Table 3: Vanilla MLP (6-Layers, Pure SGD, No BN). SIGReg significantly improves gradient flow.
Appendix figures & tables4 assets
Supplementary material from the paper’s appendix.
Appendix
Model
Optimizer
Augmentation
SIGReg
Epochs
Top-1 Acc
ResNet18
SGD
No
None
1600
79.03%
ResNet18
SGD
No
Strong
1600
78.86%
ResNet18
SGD
No
Weak
1600
79.42%
ResNet18
SGD
Yes
None
400
82.13%
ResNet18
SGD
Yes
Strong
400
81.18%
ResNet18
SGD
Yes
Weak
400
82.13%
Appendix
Table 4: Complete Experimental Results across all architectures.
Model
Optimizer
Augmentation
SIGReg
Top-1 Acc
ViT (Fixed)
AdamW
Yes
None
70.76%
ViT (Fixed)
AdamW
Yes
Strong
72.71%
ViT (Fixed)
AdamW
Yes
Weak
71.65%
Appendix
Table 5: ViT Performance after Manual Intervention. SIGReg provides additive gains even after heavy manual tuning.
Model Version
Augmentation
SIGReg
Top-1 Acc
Gain over Baseline
Standard ViT
No
None
58.77%
-
Standard ViT
No
Strong
63.16%
+4.39%
Standard ViT
No
Weak
67.52%
+8.75%
Standard ViT
Yes
None
62.44%
-
Standard ViT
Yes
Strong
74.34%
+11.90%
Standard ViT
Yes
Weak
74.56%
+12.12%
Appendix
Table 6: ViT Performance with Muon Optimizer. Muon alone stabilizes the baseline significantly compared to AdamW (Table 1 ), yet SIGReg still provides substantial additive gains.