VISReg: Variance-Invariance-Sketching Regularization for JEPA training

Abstract: 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.

• The paper introduces VISReg, a novel regularization scheme that decouples scale control from shape alignment to prevent embedding collapse in JEPA training.
• It employs Sliced Wasserstein Distance for precise shape control, aligning embedding distributions to an isotropic Gaussian while maintaining robust gradients.
• VISReg demonstrates superior out-of-distribution robustness and transfer learning performance compared to SIGReg, VICReg, and DINO.

VISReg: Variance-Invariance-Sketching Regularization for JEPA Training

Introduction and Motivation

Joint-Embedding Predictive Architectures (JEPA) represent a dominant approach in self-supervised learning (SSL), supplanting contrastive learning with formulations that eliminate the need for negative sampling and heuristic collapse prevention. Despite advances, a fundamental limitation persists: embedding collapse and model instability, which past frameworks like VICReg and SIGReg address only partially. VICReg decomposes the training objective into variance, invariance, and covariance, but covariance constraints offer only second-order regularization, insufficient for fully controlling distributional shape. SIGReg enforces alignment to an isotropic Gaussian using distributional sketching (Epps-Pulley, Cramér–Wold), but lacks scale-shape decoupling and suffers from vanishing gradients when collapse occurs.

The “VISReg: Variance-Invariance-Sketching Regularization for JEPA training” (2606.02572) presents a rigorous solution by combining the distributional rigor of sketching-based regularization with scale control, robust gradients, and scalability. VISReg replaces covariance regularization with a Sliced Wasserstein Distance (SWD)-based sketching loss on normalized embeddings, retains a variance term for scale, and adds a centering constraint, thus decoupling scale and shape regularization and achieving stable, generalizable, foundation representations—even in long-tailed, low-rank, and extreme OOD scenarios.

VISReg Regularization Objective

Decoupled Regularization

VISReg explicitly decouples regularization into three components:

• Scale control: Regularization of per-dimension variance to prevent collapse and ensure consistent information bandwidth.
• Shape control: Sliced-Wasserstein-based sketching regularization that aligns the embedding distribution (along random 1D projections) to an isotropic Gaussian, enforcing full distributional shape.
• Centering: Regularization of the embedding center to enhance robustness and convergence speed.

The overall loss is then:

$$\mathcal{L}_{\mathrm{Reg}} = \lambda_\mathrm{scale}\, \mathcal{L}_\mathrm{scale} + \lambda_\mathrm{shape}\, \mathcal{L}_\mathrm{shape} + \lambda_\mathrm{center}\, \mathcal{L}_\mathrm{center}$$

where hyperparameter weights allow for tuning the trade-off between objectives.

Scale Loss Robustness

The scale constraint penalizes deviation of the per-dimension standard deviation from unity, which, unlike SIGReg or covariance constraints, retains strong gradients under collapse conditions.

Figure 1: Embedding collapse prevention; in collapse, VISReg provides robust gradients for recovery, outperforming SIGReg.

Sliced Wasserstein Distance Shape Loss

VISReg employs random projection-based alignment of embedding slices to Gaussian quantiles using 1D closed-form Wasserstein-2 distances, addressing higher-order moments and not merely covariance or decorrelation. Separation of scale control is enforced using stop-gradient during normalization, ensuring shape optimization does not interfere with variance gradients.

Scalability and Complexity Analysis

VISReg’s regularization is linear in all scaling factors, with an effective complexity $O(NDK)$, where $N$ is batch size, $D$ the projection dimension, and $K$ the number of slices. In practice, $K$ must be proportional to $D$ for optimal distributional regularization ($K = C D$), but distributed training allows $K$ to remain constant per GPU, as slices are independent.

Figure 2: Scaling cost; VISReg is more memory efficient and scalable than other sketching or second-order methods.

VISReg maintains high linear-probe accuracy as both $D$ and $O(NDK)$0 scale.

Figure 3: Linear probe accuracy versus projection dimension; VISReg exhibits consistent gains and robustness as $O(NDK)$1 increases.

Scaling experiments confirm that distributing the projections across multiple GPUs preserves accuracy at scale, further reducing per-GPU memory cost and wall-clock time.

Empirical Performance: OOD Robustness and Downstream Transfer

VISReg is benchmarked extensively on both canonical and distributionally-shifted datasets. Two points are emphasized:

• Robustness to low-quality data: VISReg outperforms SIGReg, VICReg, and DINO on long-tailed (ImageNet-LT) and low-rank (Galaxy10) datasets, consistently preventing collapse and extracting meaningful representations.
• OOD generalization: While DINO attains higher linear probe accuracy on in-domain validation, VISReg exhibits superior transfer, OOD classification, and dense prediction performance—critical for foundation model adaptability.

Fine-tuning analyses reveal that even when DINO outperforms VISReg on the linear probe, VISReg’s transfer learning renders it more effective after finetuning, with improved results across CIFAR, Flowers, Galaxy10, and ImageNet.

VISReg also achieves image generation guidance scores exceeding DINO in IS, gFID, and sample diversity metrics, benefitting generative modeling workflows.

Ablation and Hyperparameter Analysis

VISReg introduces only a small set of tunable hyperparameters (loss weights, learning rate, batch size, projection dimension). Ablations confirm:

• Each component (scale, shape, center) is necessary for convergence and final performance.
• Increasing the shape loss weight enhances performance on low-quality or highly unbalanced distributions.
• Detachment of gradients between scale and shape loss, while a minor improvement in some settings, is favored for stability.

The strong negative Pearson correlation between VISReg loss and online accuracy (–0.996) enables loss curve monitoring as a proxy for downstream probe accuracy.

Figure 4: Pearson correlation between loss and online accuracy during ViT-L/14 training, quantifying the utility of the VISReg loss for early stopping or learning rate scheduling.

Learned Representation Structure

Qualitative analysis of patch representations (PCA visualization) demonstrates that VISReg, compared to DINO, yields spatially more fine-grained and detailed features, contributing to its OOD generalization advantages.

Figure 5: PCA visualization of three video frames; VISReg representations capture more detailed semantic structure than DINO.

Figure 6: PCA visualization of ImageNet1K images; VISReg features show enhanced detail and disentangling compared to DINO.

Implications and Future Directions

VISReg demonstrates that (i) full distributional sketching, (ii) robust scale-regularization, and (iii) carefully decoupled objectives constitute a sufficient collapse-preventing, scalable, and generalizable SSL recipe. This methodology eliminates reliance on heuristic architectural asymmetry (e.g., EMA, stop-gradient, momentum teacher) and achieves resilient performance across severe distributional shifts and scale regimes.

Potential avenues for future work include extensions to multimodal JEPA, sparse or structured sketching to further reduce computational overhead, adaptive loss reweighting, and enhanced projection strategies for specific downstream domains. Empirically, enhanced OOD generalization by VISReg may inform improved robustness benchmarks and model selection in foundation model pipelines.

Conclusion

VISReg introduces a principled and scalable regularization scheme for JEPA-based SSL, combining variance-controlled scale with Sliced Wasserstein-based sketching for shape. The approach consistently features strong gradients in degenerate regimes, exemplary OOD robustness, and efficient scaling properties. VISReg broadens the applicability of self-supervised foundation models, especially in settings where transfer, adaptation, and computational efficiency are paramount (2606.02572).