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

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

For one-dimensional distributions, the p p -th Wasserstein distance equals the L p L_{p} distance between quantile functions. For discrete empirical samples of size N N : 𝒲 p p ​ ( μ ^ , ν ^ ) = 1 N ​ ∑ i = 1 N ‖ x ( i ) − y ( i ) ‖ p , \mathcal{W}_{p}^{p}(\hat{\mu},\hat{\nu})=\frac{1}{N}\sum_{i=1}^{N}\|x_{(i)}-y_{(i)}\|^{p}, (4) where x ( i ) x_{(i)} denotes the i i -th order statistic.

Leveraging Lemma 3.2 with p = 2 p=2 , the shape loss is: ℒ shape = 1 K ​ ∑ k = 1 K ‖ sort ⁡ ( 𝐙 ~ ​ w k ) − 𝐪 𝒩 ‖ 2 2 , \mathcal{L}_{\mathrm{shape}}=\frac{1}{K}\sum_{k=1}^{K}\left\|\mathrm{sort}(\mathbf{\widetilde{Z}}w_{k})-\mathbf{q}_{\mathcal{N}}\right\|_{2}^{2}, (5) where sort ⁡ ( ⋅ ) \mathrm{sort}(\cdot) sorts the projected values in each direction, and 𝐪 𝒩 ∈ ℝ N \mathbf{q}_{\mathcal{N}}\in\mathbb{R}^{N} represents the fixed quantiles of the standard Gaussian distribution. This is strictly more expressive than covariance regularization: it enforces not just decorrelation but the full marginal distribution along every projected direction.

Additionally, empirical results suggest that regularizing the embedding center increases training robustness, so we include a centering loss: ℒ center = ‖ μ ‖ 2 2 \mathcal{L}_{\mathrm{center}}=\|\mu\|_{2}^{2} (6) where μ \mu is the batch mean.

###### Proposition 3.3 (VISReg Regularization Objective) .

The regularization loss ℒ Reg \mathcal{L}_{\mathrm{Reg}} optimizes variance and distributional shape independently: ℒ Reg = λ scale ​ ℒ scale + λ shape ​ ℒ shape + λ center ​ ℒ center \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}} (7)

The code is shown in Algorithm 1 . Decoupling introduces three hyperparameters; we conduct ablations in Table 12 . The default λ ∗ = 1 \lambda_{*}=1 works well for high-quality datasets, but increasing the shape loss weight improves performance on low-quality datasets.

For the invariance objective, we follow LeJEPA ( Balestriero & LeCun, 2025 ) : ℒ pred = 1 V ​ ∑ i = 1 V ‖ μ g − z i ‖ 2 2 \mathcal{L}_{\mathrm{pred}}=\frac{1}{V}\sum_{i=1}^{V}\|\mu_{g}-z_{i}\|_{2}^{2} (8) where V V is the number of views, μ g \mu_{g} is the mean embedding of global views, and z i z_{i} includes both global and local view embeddings. The full VISReg objective is: ℒ VISReg = ( 1 − λ ) ​ ℒ pred + λ ​ ℒ Reg \mathcal{L}_{\mathrm{VISReg}}=(1-\lambda)\,\mathcal{L}_{\mathrm{pred}}+\lambda\,\mathcal{L}_{\mathrm{Reg}} (9)

Ablation results for each component are in Section B .

### 3.2 VISReg Is Friendly to Scale Up

One practical advantage inherited from the Cramér-Wold framework is favorable scaling behavior. Due to limited computational resources, we analyze scalability through algorithm complexity, simulated scaling cost, and experiments on a small yet challenging dataset.

###### Definition 3.4 .

Let the input feature 𝐙 ∈ ℝ N × D \mathbf{Z}\in\mathbb{R}^{N\times D} , where N N is the mini-batch size and D D is the projection dimension. The number of random slices is K K .

The complexity of ℒ Reg \mathcal{L}_{\mathrm{Reg}} is dominated by two operations: 𝒞 Reg = O ⁡ ( N ​ D ​ K ) ⏟ projection + O ⁡ ( K ​ N ​ log ⁡ N ) ⏟ sorting \mathcal{C}_{\mathrm{Reg}}=\underbrace{O(NDK)}_{\text{projection}}+\underbrace{O(KN\log N)}_{\text{sorting}} (10) Since log ⁡ N ≪ D \log N\ll D at scale, the effective complexity is: 𝒞 Reg = O ⁡ ( N ​ D ​ K ) \mathcal{C}_{\mathrm{Reg}}=O(NDK) (11) This is linear in all scaling parameters—compared to VICReg’s O ⁡ ( N ​ D 2 ) O(ND^{2}) from covariance computation. We next analyze the effects of batch size N N , projection dimension D D , and number of slices K K .

Analysis in N N . We simulate the running time and memory demand of popular regularization methods ( Zbontar et al., 2021 ; Bardes et al., 2022 ; Balestriero & LeCun, 2025 ) at scale. Since VISReg is based on SWD, we also include vanilla SWD. Figure 3 shows that SWD-based methods are more efficient in both speed and memory. The 17-knot sampling required by the Epps-Pulley test slows SIGReg down. We conclude that VISReg scales efficiently in batch size.

Lemma 3.1 establishes that we can regulate D D -dimensional space by aligning K K 1D slices, so the relationship between K K and D D is important for scaling. We analyze this by reporting the online linear probe accuracy of ViT-S/8 ( Dosovitskiy et al., 2021 ) on ImageNette 1 1 1 https://github.com/fastai/imagenette ( Deng et al., 2009 ) , comparing SIGReg ( Balestriero & LeCun, 2025 ) (CF-based) with SWD ( Bonneel et al., 2015 ) and VISReg (OT-based). Unless stated otherwise, we use a batch size of 256, a learning rate of 10 − 3 10^{-3} without decay, 4 global views with a cropping ratio (0.08, 1). The λ \lambda weights for SIGReg, SWD, and VISReg are 0.02, 0.6, and 0.6, respectively. Models are trained on a single H100 GPU for 800 epochs; we report the highest accuracy.

Analysis in D D . With a sufficient number of slices ( K = 4096 K=4096 ), we vary the projection dimension D D . Figure 4 reveals three patterns: (1) OT-based methods regularize dimensions more efficiently than the CF-based method; (2) with sufficient K K , VISReg learns better semantically meaningful embeddings; (3) K K must exceed D D by a factor C > 1 C>1 for optimal accuracy, and VISReg requires the smallest C C . The third observation suggests that K K cannot be treated as independent of D D , converting the effective complexity from O ⁡ ( N ​ D ​ K ) O(NDK) to O ⁡ ( N ​ D ⋅ C ​ D ) O(ND\cdot CD) . We address this below.

Analysis in K K . Fixing D = 256 D=256 and varying K K , Figure 5 shows: (1) OT-based methods remain robust even at K = 1 8 ​ D K=\frac{1}{8}D ; (2) VISReg is the most robust approach, consistently achieving the highest linear probe accuracy.

Despite VISReg’s robustness, the correlation between K K and D D remains a concern for complexity. Revisiting Lemma 3.2 and Algorithm 1 , we observe that K K random slices are generated independently per GPU, so one can generate C ​ D M \frac{CD}{M} slices on each of M M GPUs to obtain K = C ​ D K=CD total slices. For example, 128 slices per GPU on 8 GPUs should match 1024 slices on one GPU.

Figure 6 confirms this. With one GPU, the accuracy gap between K = 128 K=128 and K = 1024 K=1024 reaches 13.88% for SIGReg, 2.21% for SWD, and 2.44% for VISReg. With 8 GPUs and the same per-GPU K K , the gap shrinks to 0.27%, 0.24%, and 0.22% respectively. Given nondeterministic training, these results support our claim. Thus, K K can remain constant when scaling, preserving the O ⁡ ( N ​ D ​ K ) O(NDK) complexity.

### 3.3 VISReg Is Robust to Low-Quality Datasets

Low-quality datasets pose challenges from many angles. We evaluate on ImageNet-LT ( Liu et al., 2019 ) (long-tailed) and Galaxy10 ( Leung, 2025 ) (low-rank). Training settings follow the previous section except l ​ r = 10 − 4 lr=10^{-4} , K = 4096 K=4096 , D = 256 D=256 , with images resized to 128px. We include DINO ( Caron et al., 2021 ) and VICReg as baselines. Models are trained from scratch for 400 epochs to simulate real-world scenarios where suitable pretrained models do not exist—a common challenge in domains like AI for Science.

ImageNet-LT is a long-tailed variant of ImageNet-1K containing 115K images from 1K classes, categorized into many-shot, medium-shot, and few-shot. Table 1 shows that VISReg outperforms all methods at all levels after adjusting the shape loss weight (details in Table 12 ), whereas DINO fails to learn meaningful embeddings.

Galaxy10 comprises 17,736 galaxy images from 10 classes. We treat it as low-rank because: (1) it has 10 classes with limited training data, below the capacity of ViT-S/8; and (2) most images contain a large ratio of black pixels, limiting useful content. Table 2 shows that all four regularization methods prevent collapse and achieve good accuracy, but DINO struggles to learn meaningful embeddings.

Summary of Analyses. First, VISReg has complexity O ⁡ ( N ​ D ​ K ) O(NDK) , linear in all scaling factors—an improvement over VICReg’s O ⁡ ( N ​ D 2 ) O(ND^{2}) . Second, we observe that K K is correlated with D D by a factor C > 1 C>1 , but prove that distributing slices across M M GPUs resolves this, keeping K K constant at scale. Third, VISReg outperforms existing methods in training efficiency, effectiveness, and robustness. Fourth, VISReg is more resilient to low-quality datasets through loss reweighting—demonstrating the importance of decoupling scale and shape over the monolithic covariance or sketching approaches. All these results confirm that VISReg is a practical and principled regularization method for real-world self-supervised learning.

## 4 Experiment

This section covers the ablation study of hyperparameter settings, the effect of projection dimension on downstream tasks, and comparisons between VISReg and existing methods in linear probe, transfer learning, domain shifting, dense instance prediction, and image generation guidance.

### 4.1 Ablation study

Unlike previous works ( Caron et al., 2021 ; Chen et al., 2021 ; Assran et al., 2023 ) relying on heuristics for training stability, VISReg only has four hyper-parameters to tune. Unless stated otherwise, the training set is ImageNet1K, the backbone is ViT-B/16, the number of slices is 4096 per GPU, the augmentation settings follow LeJEPA ( Balestriero & LeCun, 2025 ) with 2 global views and 6 local views, and the training epoch is 100. We report the online linear probe accuracy on ImageNet1K to analyze the effect of hyper-parameters, as shown in Table 3 .

Effect of λ \lambda . Different from SIGReg, scaling the regularization loss with batch size to maintain the batch size invariance, VISReg is naturally batch invariant, so a large λ \lambda value is needed to ensure the contribution of VISReg in the gradient. For small datasets, e.g., ImageNette and Galaxy10, 0.6 is a good start. For large datasets, e.g., ImageNet1K, 0.9 is a good start.

Effect of learning rate. Similar to the other methods, 5e-4 to 1e-3 is the optimal range for the training on ImageNet1K. When training on a large dataset, 9e-4 is a good start.

Effect of batch size. Grounded in the same theorem as LeJEPA, VISReg is also robust to a small batch size. Different from LeJEPA, the VISReg algorithm in VISReg benefits from a large batch size in regularizing the embedding space. Hence, we recommend reducing λ \lambda when observing a fast accuracy saturation with a large batch size.

Effect of projection dimension. Similar to previous works, we use a 3-layer MLP as the projection layer to apply regularization. Different from previous works, the final projection dimension not only decides the information bandwidth but also the difficulty of the regularization process, i.e., the lower dimension the easier. To investigate the trade-off, we increase the training epochs to 300 and test the model performance under four projection dimension settings.

The investigation includes three aspects: in-domain, OOD, and segmentation. Following the settings in DINOv2 ( Oquab et al., 2024 ) , we run offline linear probe on ImageNet1K and linear segmentation on ADE20K ( Zhou et al., 2017 ) . The full-shot linear probe performance of the other datasets ( Maji et al., 2013 ; Krause et al., 2013 ; Krizhevsky et al., 2009 ; Nilsback & Zisserman, 2008 ; Parkhi et al., 2012 ; Cimpoi et al., 2014 ; Wang et al., 2017 ) is reported for pattern observation.

Starting at the in-domain results, one observation is that a larger projection dimension results in a higher accuracy. With the embedding size 768 of ViT-B/16, projection dimension 512 gives the highest overall accuracy and 64 gives the lowest average accuracy. This indicates that projection dimension can be the bottleneck for in-domain classification. Focusing on the OOD datasets, the observation is that a smaller projection dimension limits the OOD performance, but a larger dimension might lead to over-parameterization and training set memorization. Interestingly, the smallest the projection dimension outperforms the largest one. Lastly, we observe that a smaller projection dimension leads to a better performance on dense instance prediction. We choose 256 as the optimal setting in the training.

### 4.2 General comparison

Linear probe, transfer learning, and domain shifting are three key aspects of evaluating the efficacy of a SSL foundation model. We compare VISReg with seven existing methods that are widely used in the real-world applications. In addition, we add the segmentation and generation tasks to evaluate VISReg on dense instance prediction and semantic meaning guidance for generation.

Datasets. The base training set is ImageNet1K ( Deng et al., 2009 ) . There are 15 datasets used to cover the general comparison experiment. 8 of them are in-domain datasets: FGVC-aircraft ( Maji et al., 2013 ) , Stanford cars ( Krause et al., 2013 ) , Cifar10 & Cifar100 ( Krizhevsky et al., 2009 ) , Oxford 102 flowers ( Nilsback & Zisserman, 2008 ) , Food 101 ( Bossard et al., 2014 ) , Oxford-IIIT Pet ( Parkhi et al., 2012 ) , and ImageNet1K. 6 of them are OOD datasets: Describable Textures Dataset (DTD) ( Cimpoi et al., 2014 ) , Galaxy10 ( Leung, 2025 ) , ChestXRay ( Wang et al., 2017 ) , Aerial Image Dataset (AID) ( Xia et al., 2017 ) , RetinaMNIST ( Liu et al., 2022 ) , and OrganAMNIST ( Bilic et al., 2023 ) . The last dataset is ADE20K ( Zhou et al., 2017 ) for dense instance prediction. The details are in A.3 .

Training settings. We choose two commonly used backbones, ViT-B/16 and ViT-L/14, to run the experiments. ViT-B/16 uses the best hyperparameters in Table 3 and Table 4 . ViT-L/14 is trained with {learning rate=8e-4, λ \lambda =0.7, batch size=512, projection dim=384}. Both backbones are trained for 400 epochs and 4 global + 6 local views are used. The other settings follow LeJEPA ( Balestriero & LeCun, 2025 ) . We directly use the timm package to create the model and load the pre-trained weights.

Downstream task evaluation settings. The linear probe, transfer learning, and linear segmentation experiments use the same settings as DINOv2 ( Oquab et al., 2024 ) . The only difference is that the training epoch of linear probe on downstream datasets is 10. We only compare with the models that are pre-trained on ImageNet1K.

VISReg has a competitive in-domain performance. Table 5 groups the methods based on the heuristics utilization. Within the w/o heuristics group, VISReg has a stronger in-domain performance than MAE ( He et al., 2022 ) and LeJEPA ( Balestriero & LeCun, 2025 ) , achieving 75.7% accuracy with ViT-B/16 and 77.0% accuracy with ViT-L/14 on ImageNet1K. Moreover, VISReg achieves the best average accuracy on downstream datasets. Comparing to the w/ heuristics methods, there is still an accuracy gap. Despite the accuracy gap on in-domain datasets, VISReg indicates a stronger performance on DTD, the only OOD dataset in the table. Note that the ViT-B/16 of VISReg even outperforms the ViT-L and ViT-H of the other methods. This intriguing observation motivates the extended experiments on more OOD datasets.