VISReg: Variance-Invariance-Sketching Regularization for JEPA training — part 1/4

Cleaned-text code points [0, 15807). Bibliography omitted here; source.md includes it.

# 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 10 × \times more data (LVD-142M). Project and code: https://haiyuwu.github.io/visreg .

###### Keywords:

## 1 Introduction

Self-supervised learning (SSL) has evolved from contrastive learning ( Chen et al., 2020a ; Chen et al., 2020b ; He et al., 2020 ; Chen et al., 2020c ; Chen et al., 2021 ) to Joint-Embedding Predictive Architectures ( LeCun et al., 2022 ; Assran et al., 2023 ; Caron et al., 2021 ; Zhou et al., 2021 ; Oquab et al., 2024 ; Siméoni et al., 2025 ) , which are more scalable and achieve stronger performance. Despite these advantages, many methods rely on heavy heuristics ( e.g., EMA, frozen layers, teacher-student architectures) to ensure training stability.

To remove such heuristics, VICReg ( Bardes et al., 2022 ) decomposes the training objective into variance, invariance, and covariance optimization. This approach largely reduces the engineering burden while achieving competitive performance. More recently, LeJEPA ( Balestriero & LeCun, 2025 ) proved that sketching the embedding space toward an isotropic Gaussian is an effective principle for ensuring training stability and strong downstream performance, and proposed SIGReg based on the Epps-Pulley test ( Epps & Pulley, 1983 ) and the Cramér-Wold theorem ( Cramér & Wold, 1936 ) to realize this.

However, both methods have clear limitations. VICReg regularizes covariance, which captures only second-order statistics. While this encourages decorrelation, it cannot enforce the full distributional shape of the embedding space—a distribution can match in mean and covariance yet remain far from Gaussian. This makes covariance regularization a comparatively weak proxy for the isotropy that stable, information-rich training requires. On the other hand, SIGReg addresses distributional shape directly through sketching, but it does not decouple scale from shape, limiting flexibility across training regimes. More critically, the gradient of the Epps-Pulley test diminishes as the embedding collapses (Figure 2 ), eventually vanishing entirely—precisely when a strong corrective signal is needed most.

Motivated by these complementary shortcomings, we propose V ariance- I nvariance- S ketching Reg ularization ( VISReg ). VISReg retains the variance term from VICReg to control the scale of the embedding space, but replaces covariance regularization with sketching regularization : we use the Sliced Wasserstein Distance (SWD) ( Bonneel et al., 2015 ) to align the normalized embedding distribution with an isotropic Gaussian prior along random 1D projections, thereby enforcing the full distributional shape. By decoupling scale and shape into separate objectives, VISReg inherits the interpretability and flexibility of VICReg’s decomposed losses while leveraging the distributional rigor of sketching-based methods—and provides a robust gradient signal even under collapse. Combined with a standard invariance loss, VISReg forms a complete, heuristic-free self-supervised learning method.

We compare VISReg with SIGReg, VICReg, and DINO on both standard and low-quality datasets. We find that DINO struggles to learn meaningful embeddings without careful hyperparameter tuning, while VISReg, SIGReg, and VICReg are all robust—but VISReg achieves the highest accuracy and the most stable training, particularly on low-rank and long-tailed datasets. Our hyperparameter analyses further provide clear guidance for methods grounded in the Cramér-Wold theorem.

We evaluate VISReg on linear classification, transfer learning, dense prediction, and image generation guidance, covering both in-domain and out-of-distribution (OOD) settings. We pretrain backbones on ImageNet-1K and evaluate on downstream datasets. First, despite a linear probe accuracy gap relative to the best method on in-domain data, VISReg achieves the best OOD results—one of the most important properties of a useful foundation model. Second, VISReg outperforms DINO ( Caron et al., 2021 ) with the same backbone after fine-tuning on both in-domain and OOD datasets, even though DINO has over 3% higher linear probe accuracy on in-domain data, indicating strong transfer learning capability. Third, a linear segmentation experiment shows VISReg performs on par with DINO for dense prediction, though a gap to the best models ( e.g., MoCoV3 ( Chen et al., 2021 ) , iBOT ( Zhou et al., 2021 ) ) remains. Finally, to test scaling, we pretrain ViT-L/14 on ImageNet-22K ( Ridnik et al., 2021 ) . VISReg achieves results comparable to DINOv2 ( Oquab et al., 2024 ) on OOD datasets, despite the latter being trained on a 10 × \times larger dataset (LVD-142M), demonstrating the strong potential of the VISReg approach.

The contributions of this work are: • We propose VISReg, which replaces the covariance regularization of VICReg with a sketching objective grounded in optimal transport, achieving stronger distributional control, better training stability, and resilience to low-quality datasets.

• We comprehensively analyze the hyperparameter landscape of VISReg and related Cramér-Wold-based methods, providing clear guidance for scaling and training stability within this paradigm.

• We demonstrate that VISReg’s embedding regularization yields superior OOD generalization and strong downstream task performance, broadening the practical utility of self-supervised foundation models.

## 2 Related Work

Contrastive Learning and Sampling Strategies. Early successes in self-supervised learning relied heavily on contrastive objectives, which maximize the similarity between positive pairs while pushing apart negative samples ( Chen et al., 2020a ; He et al., 2020 ; Misra & Maaten, 2020 ) . SimCLR ( Chen et al., 2020a ; Chen et al., 2020b ) demonstrated the importance of strong data augmentation and large batch sizes. To decouple batch size dependency, MoCo ( He et al., 2020 ; Chen et al., 2020c ) introduced a momentum queue to maintain a dynamic dictionary of negative samples. SwAV ( Caron et al., 2020 ) reformulated contrastive learning as an online clustering problem via the Sinkhorn-Knopp algorithm. However, these methods rely on negative pairs or prototypes, introducing sampling bias ( Chuang et al., 2020 ) and computational overhead for hard-negative mining ( Robinson et al., 2021 ) . Like other non-contrastive methods, VISReg eliminates the need for negative sampling entirely.

Masked Image Modeling (MIM). Inspired by BERT ( Devlin et al., 2019 ) in NLP, MIM approaches learn by reconstructing masked inputs. MAE ( He et al., 2022 ) and SimMIM ( Xie et al., 2022 ) operate on pixel-level reconstruction, demonstrating high scalability for fine-tuning tasks. BEiT ( Bao et al., 2022 ) proposes predicting discrete visual tokens. MaskFeat ( Wei et al., 2022 ) reconstructs HOG features to focus on structural information. Despite excelling in transfer learning, MIM methods typically learn lower-level spatial statistics and lag behind joint-embedding methods in linear probing due to weaker semantic linear separability ( Park et al., 2023 ; Baevski et al., 2022 ) .

Asymmetric Joint-Embedding Architectures. To avoid collapse without negatives, several methods introduce architectural asymmetry. BYOL ( Grill et al., 2020 ) and SimSiam ( Chen & He, 2021 ) rely on stop-gradient operations and predictor networks to break symmetry. Mean Teacher ( Tarvainen & Valpola, 2017 ) and DINO ( Caron et al., 2021 ; Oquab et al., 2024 ; Siméoni et al., 2025 ) utilize a momentum-updated teacher to stabilize training, with DINO further employing centering and sharpening. OBoW ( Gidaris et al., 2021 ) and MSN ( Assran et al., 2022 ) leverage prototype-based learning with asymmetric updates. Although effective, these methods rely on implicit regularization heuristics, making their non-collapse dynamics theoretically opaque ( Li et al., 2022 ) .

Geometric and Information-Theoretic Regularization. VISReg is most closely related to methods that explicitly regularize the statistical properties of embeddings. Barlow Twins ( Zbontar et al., 2021 ) minimizes redundancy in the cross-correlation matrix between twin networks. W-MSE ( Ermolov et al., 2021 ) projects embeddings onto the unit sphere and performs whitening. VICReg ( Bardes et al., 2022 ) explicitly constrains the variance, invariance, and covariance of embeddings to maximize information content. However, covariance regularization captures only second-order statistics—it encourages decorrelation but cannot enforce the full distributional shape of the embedding space. Moreover, methods like Barlow Twins and VICReg require computing covariance matrices, scaling quadratically as 𝒪 ⁡ ( D 2 ) \mathcal{O}(D^{2}) with embedding dimension D D .

LeJEPA ( Balestriero & LeCun, 2025 ) proved that regularizing the embedding space toward an isotropic Gaussian distribution can maintain stable heuristic-free training, and introduced SIGReg—grounded in the Epps-Pulley test ( Epps & Pulley, 1983 ) —to achieve this. While SIGReg provides stronger distributional control than covariance regularization and scales linearly as 𝒪 ⁡ ( D ) \mathcal{O}(D) , its regulation signal diminishes when the embedding collapses. KerJEPA ( Zimmermann et al., 2025 ) leverages MMD to estimate the regulation of infinite projections, but incurs O ⁡ ( N 2 ) O(N^{2}) complexity in batch size N N . A contemporary work, LpJEPA ( Kuang et al., 2026 ) , proposes Rectified Distribution Matching Regularization (RDMReg) to enforce embedding sparsity.

Our VISReg bridges VICReg and SIGReg: it retains VICReg’s variance term for scale control but replaces the covariance term with a sketching objective based on SWD ( Bonneel et al., 2015 ) , achieving full distributional shape regularization with robust gradients, decoupled scale-shape optimization, and linear complexity—making it well suited for scaling.

## 3 VISReg: Variance-Invariance-Sketching Regularization

VICReg ( Bardes et al., 2022 ) decomposes embedding regularization into variance and covariance terms, providing interpretability and flexibility. However, covariance regularization captures only second-order statistics: it encourages decorrelation among embedding dimensions but cannot enforce the full distributional shape of the embedding space. A distribution can match in mean and covariance yet remain far from the isotropic Gaussian that LeJEPA ( Balestriero & LeCun, 2025 ) proved to be optimal for stable, heuristic-free self-supervised training.

SIGReg ( Balestriero & LeCun, 2025 ) addresses this by directly sketching the embedding distribution toward an isotropic Gaussian via the Epps-Pulley test and the Cramér-Wold theorem. However, we identify two limitations: (1) its gradient diminishes as the embedding collapses (Figure 2 ), vanishing precisely when correction is needed most; and (2) it does not decouple scale from shape, limiting flexibility across training regimes.

VISReg resolves these limitations by replacing VICReg’s covariance term with a sketching objective while retaining the variance term for scale control. By decoupling regularization into distinct scale and shape objectives, VISReg provides robust gradients against collapse, distributional rigor beyond second-order statistics, and the flexibility to reweight objectives for different data regimes.

### 3.1 Regularization Loss

We decouple the regularization into scale and shape components, each operating independently. For simplicity, the number of augmentations V V is omitted from the derivation.

Scale Regularization. We regulate the scale of the embedding space using a variance constraint, following the same intuition as VICReg. Directly minimizing the KL divergence ( Kullback & Leibler, 1951 ) to an isotropic Gaussian prior incurs O ⁡ ( D 3 ) O(D^{3}) complexity. We relax this by factorizing into marginal distributions. Given the centered embedding 𝐙 ^ ∈ ℝ N × D \mathbf{\hat{Z}}\in\mathbb{R}^{N\times D} , the scale loss is: ℒ scale = 1 D ​ ∑ j = 1 D ( 1 − σ j ​ ( 𝐙 ^ ) ) 2 \mathcal{L}_{\mathrm{scale}}=\frac{1}{D}\sum_{j=1}^{D}(1-\sigma_{j}(\mathbf{\hat{Z}}))^{2} (1) where σ j ​ ( ⋅ ) \sigma_{j}(\cdot) denotes the standard deviation of the j j -th dimension. This formulation provides a gradient that approaches a constant during collapse, ensuring a reliable corrective signal.

Shape Regularization. Where VICReg uses covariance to encourage decorrelation—capturing only second-order statistics—we instead sketch the embedding distribution toward an isotropic Gaussian, enforcing full distributional shape. To isolate the geometric structure from magnitude, we normalize 𝐙 ^ \mathbf{\hat{Z}} : 𝐙 ~ = 𝐙 ^ s ​ g ​ ( σ ) + ϵ \mathbf{\widetilde{Z}}=\frac{\mathbf{\hat{Z}}}{sg(\sigma)+\epsilon} (2) The stop-gradient s ​ g ​ ( ⋅ ) sg(\cdot) decouples shape optimization from scale, ensuring gradients from the shape loss do not interfere with variance regulation. Unlike prior uses of stop-gradient as a collapse-prevention heuristic ( Grill et al., 2020 ; Chen & He, 2021 ) , here it serves a principled role in objective decomposition.

To efficiently align the high-dimensional distribution of 𝐙 ~ \mathbf{\widetilde{Z}} with the isotropic Gaussian prior, we leverage the Sliced Wasserstein Distance , grounded in the Cramér-Wold theorem ( Cramér & Wold, 1936 ) :

###### Lemma 3.1 (Cramér-Wold Theorem) .

Let μ \mu and ν \nu be two probability measures on ℝ d \mathbb{R}^{d} . The Radon transform ( Radon, 2005 ) ℛ \mathcal{R} , defined as ℛ ​ μ ​ ( θ , t ) := ∫ ℝ d δ ⁡ ( t − ⟨ x , θ ⟩ ) ​ 𝑑 μ ​ ( x ) \mathcal{R}\mu(\theta,t):=\int_{\mathbb{R}^{d}}\delta(t-\langle x,\theta\rangle)\,d\mu(x) along all directions θ ∈ 𝕊 d − 1 \theta\in\mathbb{S}^{d-1} , is injective. Thus: μ = ν ⇔ ℛ ​ μ ​ ( θ , ⋅ ) = ℛ ​ ν ​ ( θ , ⋅ ) , ∀ θ ∈ 𝕊 d − 1 . \mu=\nu\iff\mathcal{R}\mu(\theta,\cdot)=\mathcal{R}\nu(\theta,\cdot),\quad\forall\theta\in\mathbb{S}^{d-1}. (3)

This allows us to regularize the high-dimensional shape by aligning 1D random projections P k = 𝐙 ~ ​ w k P_{k}=\mathbf{\widetilde{Z}}w_{k} , where w k ∈ ℝ D w_{k}\in\mathbb{R}^{D} . Unlike SIGReg, which operates in the frequency domain via the Epps-Pulley test, we adopt the 2-Wasserstein distance ( 𝒲 2 \mathcal{W}_{2} ) , which admits an efficient closed-form solution in 1D ( Peyré et al., 2019 ; Bonneel et al., 2015 ; Deshpande et al., 2018 ) :

###### Lemma 3.2 (1D Wasserstein Closed-Form) .