PaperlayerView exact revision on arXiv

Source-linked AI summary

VISReg: Variance-Invariance-Sketching Regularization for JEPA training

Haiyu Wu, Randall Balestriero, Morgan Levine

arXiv:2606.02572v1cs.CV

TL;DR

VICReg’s covariance misses full distributional shape, while SIGReg lacks scale-shape decoupling and loses gradient strength under collapse. VISReg separates scale from shape with variance and Sliced-Wasserstein sketching, achieving the highest accuracy and most stable training particularly on low-rank and long-tailed datasets, alongside superior OOD generalization.

  • Problem

    VICReg captures only second-order statistics, while SIGReg’s sketching lacks scale-shape decoupling and loses gradient strength under collapse.

    Source: §1 · p. 1, §3 · p. 3

  • Method

    VISReg retains variance regularization for scale control and replaces covariance with Sliced Wasserstein sketching to align embeddings with an isotropic Gaussian.

    Source: §1 · p. 2

  • Results

    VISReg achieves the highest accuracy and most stable training among compared methods, particularly on low-rank and long-tailed datasets, and yields superior OOD generalization.

    Source: §1 · p. 2

  • Takeaways & Limitations

    VISReg’s embedding regularization broadens the practical utility of self-supervised foundation models.

    Source: §1 · p. 2

    SIGReg’s gradient vanishes under collapse, and its failure to decouple scale from shape limits flexibility across training regimes.

    Source: §3 · p. 3

Abstract

from arXiv · show

Self-supervised learning methods prevent embedding collapse via modeling heuristics or explicit regularization of the embedding space. Among the latter, VICReg decomposes regularization into variance and covariance objectives, offering flexibility and interpretability. However, covariance captures only second-order statistics -- encouraging decorrelation but failing to enforce the full distributional shape needed for stable training. Sketching-based methods such as SIGReg address this by aligning embeddings to an isotropic Gaussian, but lack flexibility and suffer from vanishing gradients under collapse. We propose Variance-Invariance-Sketching Regularization (VISReg), which replaces covariance with a Sliced-Wasserstein-based sketching objective that enforces full distributional shape, while retaining a variance term for scale control. By decoupling scale and shape, VISReg combines VICReg's flexibility with the distributional rigor of sketching methods, providing robust gradients even under collapse. We show that VISReg scales linearly, outperforms existing regularization on low-quality datasets, and is resilient to long-tailed and low-rank regimes. Pre-trained on ImageNet-1K, VISReg achieves state-of-the-art performance on out-of-distribution datasets. Pre-trained on ImageNet-22K, it matches DINOv2's OOD performance despite the latter using 10x more data (LVD-142M). Project and code: https://haiyuwu.github.io/visreg.

1. Introduction

The introduction presents VISReg as a way to combine flexible scale control with distributional shape regularization, while addressing weaknesses in covariance- and sketching-based methods. Its evaluations report strong results on low-quality data and across in-domain and OOD tasks.

1. Introduction

  • VICReg's covariance term captures only second-order statistics, whereas SIGReg's sketching addresses distributional shape but lacks scale-shape decoupling and has gradients that vanish as embeddings collapse.A distribution can match in mean and covariance yet remain far from Gaussian.Source: §1 · p. 1
  • VISReg retains VICReg's variance term for scale control and replaces covariance with Sliced Wasserstein sketching to align normalized embeddings with an isotropic Gaussian along random 1D projections.Combined with an invariance loss, the method is presented as heuristic-free and as providing a robust gradient signal even under collapse.Source: §1 · p. 2
  • On standard and low-quality datasets, VISReg achieves the highest accuracy and most stable training among VISReg, SIGReg, VICReg, and DINO, especially on low-rank and long-tailed data.DINO struggles to learn meaningful embeddings without careful hyperparameter tuning; the other three methods are reported as robust.Source: §1 · p. 2
  • VISReg achieves the best OOD results after ImageNet-1K pretraining and matches DINOv2's OOD results on ImageNet-22K despite DINOv2 using the 10× larger LVD-142M dataset.It also outperforms DINO after fine-tuning on in-domain and OOD datasets, despite DINO's higher in-domain linear-probe accuracy.Source: §1 · p. 2
  • The paper analyzes hyperparameters for VISReg and related Cramér-Wold-based methods to provide guidance for scaling and training stability.The evaluations span linear classification, transfer learning, dense prediction, and image-generation guidance in in-domain and OOD settings.Source: §1 · p. 2

2. Related Work

Related work covers contrastive, masked-image-modeling, asymmetric joint-embedding, and embedding-regularization approaches. VISReg is positioned as combining VICReg's scale control with sketching-based distributional regularization.

Related Work

  • Contrastive methods use negative pairs or prototypes, with sampling bias and hard-negative-mining overhead; VISReg eliminates negative sampling.Examples include SimCLR and MoCo [Chen et al., 2020a;b; He et al., 2020; Chen et al., 2020c], and SwAV.Source: §2 · p. 2
  • Masked-image-modeling methods reconstruct pixels, discrete tokens, or features, but typically learn lower-level spatial statistics and trail joint-embedding methods in linear-probe semantic separability.The passage cites MAE and SimMIM, BEiT, and MaskFeat.Source: §2 · p. 2
  • Asymmetric joint-embedding methods prevent collapse through stop-gradient and predictor networks, momentum teachers, or prototype-based asymmetric updates.Examples include BYOL, SimSiam [Chen & He, 2021], Mean Teacher, DINO [Siméoni et al., 2025], OBoW, and MSN.Source: §2 · pp. 2–3
  • Barlow Twins and VICReg regularize correlations or covariance, but covariance captures only second-order statistics and covariance matrices scale as O(D^2) with embedding dimension D.VICReg decomposes variance, invariance, and covariance objectives to constrain embedding statistics.Source: §2 · p. 3
  • SIGReg provides distributional control with O(D) complexity but weakens under collapse, while KerJEPA incurs O(N^2) batch-size complexity; VISReg combines variance-based scale control with SWD sketching and linear complexity.LpJEPA proposes Rectified Distribution Matching Regularization to enforce embedding sparsity.Source: §2 · p. 3

3. VISReg: Variance-Invariance-Sketching Regularization

VISReg separates scale from shape, replacing covariance regularization with SWD sketching; the analyses report linear complexity and robustness on low-quality datasets.

3. VISReg: Variance-Invariance-Sketching Regularization

  • VISReg retains VICReg’s variance term for scale control, replaces covariance with SWD sketching for full distributional shape, and decouples scale and shape objectives.Covariance captures only second-order statistics, while SIGReg does not separate scale from shape and its gradient diminishes during collapse.Source: §3 · p. 3

3.1. Regularization Loss

  • Normalized embeddings are aligned to an isotropic Gaussian along random 1D projections using SWD, enforcing full marginal distributions rather than only decorrelation.A separate centering loss is included because empirical results suggest it increases training robustness.Source: §3.1 · pp. 3–4

3.3. VISReg Is Robust to Low-Quality Datasets

  • Increasing the shape-loss weight improves performance on low-quality datasets, supporting VISReg’s resilience in long-tailed and low-rank settings.Source: §3.1 · p. 4, §3.3 · p. 6

3.2. VISReg Is Friendly to Scale Up

  • VISReg has O(NDK) complexity versus VICReg’s O(ND^2); distributing slices across GPUs can keep K constant as training scales.K otherwise must exceed D by a factor C>1 for optimal accuracy; with 8 GPUs, accuracy gaps between K=128 and K=1024 shrink to 0.27% for SIGReg, 0.24% for SWD, and 0.22% for VISReg.Source: §3.2 · pp. 5–6
  • SWD-based methods are more efficient in speed and memory than alternatives, and VISReg scales efficiently with batch size.At batch size 50K on one H100, VISReg has slightly better speedup and a 13.7% memory demand over SIGReg.Source: §3.2 · p. 5, Figure 3 · p. 5
  • With D=256, VISReg is the most robust approach across slice counts, consistently achieving the highest linear-probe accuracy, and OT-based methods remain robust at K=1/8D.Source: §3.2 · p. 6

3.3. VISReg Is Robust to Low-Quality Datasets

  • VISReg outperforms all methods at every ImageNet-LT shot level after shape-loss reweighting; on Galaxy10, regularizers prevent collapse with good accuracy, while DINO struggles.Source: Table 1 · p. 6, §3.3 · p. 6, Table 2 · p. 6

4. Experiment

The experiments examine hyperparameters and compare VISReg across classification, transfer, OOD, segmentation, and generation tasks. VISReg’s strongest reported comparative result is on OOD evaluation, while projection-dimension optima vary by task.

4.1. Ablation study

  • For ImageNet1K, 5e-4 to 1e-3 is the reported optimal learning-rate range; 9e-4 is a starting point for large datasets, and λ=0.9 is a starting point there.VISReg is robust to small batches but benefits from large batches for embedding regularization; the authors recommend reducing λ if accuracy saturates quickly.Source: §4.1 · p. 7
  • 256 was selected as the projection dimension: larger dimensions improve in-domain accuracy, while smaller dimensions can favor OOD performance and dense prediction.The authors also report that the smallest tested dimension outperforms the largest on OOD data, while warning that larger dimensions can over-parameterize and memorize training data.Source: §4.1 · p. 7, Table 4 · p. 7

4.2. General comparison

  • 75.7% with ViT-B/16 and 77.0% with ViT-L/14 are VISReg’s ImageNet1K accuracies, exceeding MAE and LeJEPA among methods without heuristics.VISReg also has the best average accuracy on downstream datasets, but an in-domain accuracy gap remains relative to methods using heuristics.Source: §4.2 · p. 8, Table 5 · p. 8
  • VISReg achieves the best average accuracy across OOD datasets without training heuristics, and ImageNet-22K-pretrained VISReg matches DINOv2 with one-tenth the training data.The evaluation covers six OOD datasets from distinct domains.Source: §4.2 · p. 8, Table 6 · p. 9
  • VISReg produces good linear-segmentation results, but its ADE20K mIoU remains notably below the best, including MoCoV3.The comparison uses a ViT-B/16 backbone.Source: Table 8 · p. 9, §4.2 · p. 9
  • Fine-tuned VISReg consistently exceeds supervised learning and DINO on the evaluated datasets, while retaining its advantage on OOD datasets.The transfer evaluation uses CIFAR10, CIFAR100, Flowers, ImageNet1K, and Galaxy10.Source: §4.2 · pp. 8–9, Table 7 · p. 9
  • VISReg achieves better results across all reported image-generation metrics than DINO using SiT-B/2 trained for 100K steps.The evaluation follows the standard 50K-generation setting without CFG.Source: Table 9 · p. 9, §4.2 · p. 9

5. Conclusion

The conclusion presents VISReg as a heuristic-free self-supervised method and reports effectiveness in scaling, training stability, and efficiency. It highlights robustness to low-quality data and stronger performance on OOD data and transfer learning.

5. Conclusion

  • The paper reports that VISReg and similar methods are more robust to low-quality datasets than DINO, a property it considers helpful for real-world applications.It also highlights stronger performance on OOD data and transfer learning.Source: §5 · p. 9
  • The conclusion presents VISReg as a self-supervised method that does not rely on heuristics for training stability.The authors report effectiveness in model scaling, stability, and training efficiency.Source: §5 · p. 9
  • The authors describe extensive experiments evaluating VISReg across aspects they consider important for foundation-model training.They suggest its OOD and transfer-learning performance may enhance foundation-model usefulness.Source: §5 · p. 9

A.1. Training details

The appendix specifies ImageNet-1K and ImageNet-22K pretraining configurations, then describes downstream evaluation protocols for classification and semantic segmentation.

A.1. Training details

  • ImageNet-1K pretraining uses DINO-style multi-crop augmentation with four global and six local views per image.A 3-layer MLP projects concatenated CLS tokens from the backbone's last two layers; optimization uses AdamW and a warmup-then-cosine learning-rate schedule.Source: §A.1 · p. 13
  • VISReg-B and VISReg-L are trained from scratch on ImageNet-1K for 400 epochs, requiring approximately 1,120 and 2,060 GPU-hours, respectively.The models use ViT-B/16 and ViT-L/14 backbones, with effective batch size 512 across 32 GPUs.Source: §A.1 · p. 13
  • ImageNet-22K pretraining uses ViT-L/14 for 100 epochs with 10 views per image and an effective batch size of 1,024.Training uses 16 H100 GPUs and requires approximately 2,720 GPU-hours.Source: §A.1 · p. 13

A.2. Testing details

  • Downstream classification evaluates frozen pretrained features on eight classification benchmarks and six OOD benchmarks using a linear classifier.The protocol concatenates the last four transformer-layer CLS tokens and reports the best test accuracy after a grid search over learning rates.Source: §A.2 · p. 13
  • ImageNet-1K linear evaluation trains ten learning-rate-specific classifier heads in parallel, while ADE20K segmentation uses a linear probe on frozen patch features.The ImageNet-1K protocol reports the best validation accuracy; ADE20K performance is reported as validation mIoU.Source: §A.2 · p. 13

A.3. Datasets

The appendix describes 15 datasets spanning multiple domains, tasks, and difficulty levels, including natural-image classification, medical imaging, remote sensing, astronomy, and semantic segmentation.

A.3. Datasets

  • The dataset collection comprises 15 datasets spanning multiple domains, tasks, and difficulty levels.The appendix provides dataset-specific descriptions in the following entries.Source: §A.3 · p. 14
  • ImageNet-1K spans 1,000 natural-object classes; CIFAR-10 and CIFAR-100 provide 10 and 100 object classes at 32 × 32 resolution, while Stanford Cars covers 196 car models.Source: §A.3 · p. 14
  • The collection includes astronomy, remote-sensing, and medical-image classification through Galaxy10, AID, NIH Chest X-ray, RetinaMNIST, and OrganAMNIST.These datasets span galaxy morphology, aerial scenes, chest radiographs, retinal images, and CT slices.Source: §A.3 · p. 14
  • Flowers-102, DTD, FGVC-Aircraft, Food-101, and Oxford-IIIT Pets test fine-grained or attribute-rich recognition across plants, textures, aircraft, food, and pet breeds.Their class counts range from 37 pet breeds to 101 food categories and 102 flower species.Source: §A.3 · p. 14
  • ADE20K provides semantic-segmentation data with pixel-level annotations for 150 classes across indoor and outdoor scenes.It contains 20,210 training images and 2,000 validation images.Source: §A.3 · p. 14

B. Additional ablations

The additional ablations test VISReg's scale and shape components, their decoupling, and the balance between their loss weights across datasets of differing quality.

B. Additional ablations

  • The ablations use a ViT-S/8 backbone and include ImageNet-LT, Galaxy10, and Imagenette to represent long-tailed, low-rank, and normal datasets.ImageNet-LT and Galaxy10 train for 400 epochs; Imagenette trains for 800.Source: B. Additional ablations · p. 14
  • Scale and shape losses are necessary for convergence on ImageNette, while center loss helps learning become faster and better.Source: Table 10 · p. 15
  • Fully decoupling scale and shape losses helps learning in all three tasks.Source: Table 11 · p. 15
  • A higher shape-regularization weight is helpful on low-quality datasets but not on high-quality datasets.The weight-ratio ablation compares results by dataset.Source: Table 12 · p. 15

B.1. Effect of decoupled components in training.

Ablations show that scale and shape losses strongly affect accuracy, while center loss has a smaller effect on final accuracy but speeds convergence. Detachment between scale and shape losses brings minor improvements across all three datasets.

B.1. Effect of decoupled components in training.

  • Without scale loss, accuracy decreases 71.02%, and without shape loss, it decreases 58.4%.The ablation identifies both components as important to the learning process.Source: §B.1 · p. 15
  • Center loss changes final accuracy by 0.41% but increases convergence speed.The authors conclude that all three components are necessary.Source: §B.1 · p. 15
  • Detachment between scale and shape losses yields minor performance improvements across all three datasets.The authors use detachment in all experiments.Source: §B.1 · p. 15

B.2. Effect of decoupled components in different training set scenarios.

The component-weight ablation holds total loss weight constant to isolate the effects of changing the scale, shape, and center ratios. Shape emphasis improves accuracy on ImageNet-LT and Galaxy10, while the default setting works best on Imagenette.

B.2. Effect of decoupled components in different training set scenarios.

  • With total weight fixed at 3, shape 4:1 beats the equal baseline by +3.2% on ImageNet-LT and +1.3% on Galaxy10.The analysis varies the ratio among scale, shape, and center components.Source: §B.2 · p. 15
  • Shape is the most impactful DSSO objective, and shifting weight toward it monotonically improves accuracy on ImageNet-LT and Galaxy10.The authors interpret this as evidence that higher shape regularization helps learning on low-quality datasets.Source: §B.2 · p. 15
  • Scale 4:1 produces the largest reported drops: -4.6% on ImageNet-LT and -2.7% on Galaxy10.On Imagenette, the default setting performs best; other imbalanced ratios largely reduce learning effectiveness.Source: §B.2 · p. 15

C.1. VISReg loss indicates the performance

VISReg loss tracks online probe accuracy closely in ViT-L/14 training on ImageNet1K. PCA visualizations compare learned patch features with DINOv1 and report finer-grained details for VISReg.

C.1. VISReg loss indicates the performance

  • -0.996 Pearson correlation between loss and online probe accuracy indicates that the loss curve reflects the model's learning curve.The measurement uses ViT-L training on ImageNet1K; Figure 7 reports the correlation for a 100-epoch run.Source: §C.1 · p. 15, Figure 7 · p. 16

C.2. Further comparison with DINOv1 on image and video.

  • PCA visualizations of image and video features show more granular details for VISReg than DINOv1.Patch features are reduced to three PCA components and aligned between models for consistent color semantics.Source: §C.2 · p. 16, Figure 8 · p. 16, Figure 9 · p. 17
Loading 2606.02572v1…