VISReg: A Robust Regularization Framework for JEPA

• VISReg is a regularization framework that enhances JEPA architectures by explicitly controlling scale, enforcing invariance, and aligning embedding shapes through Sliced-Wasserstein sketching.
• It decomposes the loss into variance, centering, and sketching terms to prevent collapse and ensure full-distributional matching, addressing limitations in VICReg and SIGReg.
• Empirical results demonstrate VISReg’s superior out-of-distribution generalization and improved performance on low-quality and transfer tasks across diverse datasets.

Variance-Invariance-Sketching Regularization (VISReg) is a regularization framework for self-supervised learning, designed specifically to improve the stability and robustness of Joint Embedding Predictive Architectures (JEPA) through a principled decomposition of embedding space regularization. VISReg combines scale control, invariance, and full-distributional matching using a variance term, an invariance (mean matching) term, and a Sliced-Wasserstein (“sketching”) shape regularizer, decoupled via stop-gradient. This novel approach addresses key limitations in prior methods such as VICReg and SIGReg, demonstrating superior out-of-distribution generalization, scalability, and resilience to pathological data regimes (Wu et al., 1 Jun 2026).

1. Motivation and Limitations of Prior Art

VICReg (Variance-Invariance-Covariance Regularization) introduced an explicit decomposition of regularizers into variance, invariance, and covariance objectives, controlling scale, clustering positive views, and decorrelating embedding dimensions, respectively. However, the covariance objective targets only second-order statistics, leaving higher-order mismatches between the embedding distribution and a desired isotropic form possible; distinct distributions can share identical means and covariances yet exhibit differing “shapes.” SIGReg (Sketching-based Regularization in LeJEPA), by contrast, utilizes the Epps–Pulley statistic to match the empirical embedding distribution to an isotropic Gaussian, thus targeting full-distributional alignment. SIGReg, however, (a) fails to separate scale from shape—so misaligned scales undermine sketching, (b) provides no corrective gradient as embeddings collapse (the sketch loss vanishes too rapidly), and (c) does not afford flexible reweighting of its components for adaptation to varied data regimes.

VISReg overcomes these issues through three key modifications:

• Explicit scale control: It retains the VICReg variance term, ensuring gradients do not vanish under collapse.
• Distributional alignment: It replaces covariance regularization with a Sliced-Wasserstein (SW) sketching penalty which matches entire distributional shape by aligning projected one-dimensional marginals with those of an isotropic Gaussian.
• Decoupling via stop-gradient: Scale and shape regularization are disentangled by stopping the gradient path on the per-dimension standard deviation when computing the SW penalty, enabling independent weighting and greater flexibility.

2. Mathematical Formulation

Let $Z \in \mathbb{R}^{N \times D}$ denote the batch of projected embeddings:

• Batch mean: $\mu = (1/N)\sum_{i=1}^N Z_i$
• Centered embeddings: $\hat Z = Z - \mathbf{1}_N \mu^\top$
• Per-dimension std: $\sigma_j = \text{std}_j(\hat Z)$

The loss decomposes as follows:

1. Variance (scale) term:

$$L_\text{var} = \frac{1}{D} \sum_{j=1}^D (1 - \sigma_j)^2$$

This term applies a strong gradient against collapse (as $\sigma_j \to 0$, $\nabla_{\sigma_j} L_\text{var} \rightarrow$ constant).

2. Centering term:

$L_\text{center} = \|\mu\|_2^2$

This maintains batch mean proximity to zero, empirically stabilizing training.

1. Sketching (shape) term (Sliced-Wasserstein):
   • Scale normalized embeddings: $$\tilde Z = \hat Z \cdot \text{diag}(1 / (\sigma \circ \text{sg}(\sigma) + \varepsilon))$$, with stop-gradient on $\sigma$.
   • Random projections: Draw $\mu = (1/N)\sum_{i=1}^N Z_i$0 random unit directions $\mu = (1/N)\sum_{i=1}^N Z_i$1, compute 1D projections $\mu = (1/N)\sum_{i=1}^N Z_i$2.
   • Empirical Wasserstein: For sorted projections $\mu = (1/N)\sum_{i=1}^N Z_i$3 and Gaussian quantiles $\mu = (1/N)\sum_{i=1}^N Z_i$4,

   $\mu = (1/N)\sum_{i=1}^N Z_i$5

2. Invariance (prediction) term:

   • For $\mu = (1/N)\sum_{i=1}^N Z_i$6 augmentations per image and mean over global crops $\mu = (1/N)\sum_{i=1}^N Z_i$7,

   $\mu = (1/N)\sum_{i=1}^N Z_i$8

3. Composite VISReg regularizer:

$\mu = (1/N)\sum_{i=1}^N Z_i$9

1. Final VISReg objective:

$\hat Z = Z - \mathbf{1}_N \mu^\top$0

Default weights: $\hat Z = Z - \mathbf{1}_N \mu^\top$1 for high-quality data; $\hat Z = Z - \mathbf{1}_N \mu^\top$2 on ImageNet-1K.

3. Training Workflow and Algorithmic Features

• Augmentation strategy: Each batch comprises $\hat Z = Z - \mathbf{1}_N \mu^\top$3 images, each yielding $\hat Z = Z - \mathbf{1}_N \mu^\top$4 augmentations (e.g., multi-crop with $\hat Z = Z - \mathbf{1}_N \mu^\top$5 large + $\hat Z = Z - \mathbf{1}_N \mu^\top$6 small, $\hat Z = Z - \mathbf{1}_N \mu^\top$7). These are encoded and projected into $\hat Z = Z - \mathbf{1}_N \mu^\top$8 via a shared encoder and MLP.

• Loss computation:

1. Invariance: Split $\hat Z = Z - \mathbf{1}_N \mu^\top$9 into global/all, compute $\sigma_j = \text{std}_j(\hat Z)$0, and $\sigma_j = \text{std}_j(\hat Z)$1 via MSE to mean.
2. Regularization: Center embeddings for $\sigma_j = \text{std}_j(\hat Z)$2, compute $\sigma_j = \text{std}_j(\hat Z)$3 (for $\sigma_j = \text{std}_j(\hat Z)$4) and per-dimension std for $\sigma_j = \text{std}_j(\hat Z)$5.
3. Sketching: Normalize $\sigma_j = \text{std}_j(\hat Z)$6 by per-dim std via stop-grad, draw $\sigma_j = \text{std}_j(\hat Z)$7 random directions, project and sort, then compute SWD to Gaussian quantiles.

• Final loss: Weighted sum $\sigma_j = \text{std}_j(\hat Z)$8.

• Optimization: AdamW, weight decay on encoder/projector, mixed precision.

• Distributed training: Slices $\sigma_j = \text{std}_j(\hat Z)$9 are independently re-sampled on each GPU, permitting effective scaling with batch and dimension.

• Implementation details: Detaching $$L_\text{var} = \frac{1}{D} \sum_{j=1}^D (1 - \sigma_j)^2$$0 for $$L_\text{var} = \frac{1}{D} \sum_{j=1}^D (1 - \sigma_j)^2$$1 prevents shape/scale gradients from interfering, supporting faster and more stable convergence.

4. Theoretical Properties and Computational Complexity

• Distributional matching: By the Cramér–Wold theorem, matching all 1D projections’ marginals ensures alignment of the full $$L_\text{var} = \frac{1}{D} \sum_{j=1}^D (1 - \sigma_j)^2$$2-dimensional distribution.

• Wasserstein closed-form: 1D SWD is computable by aligning sorted projections with Gaussian quantiles.

• Gradient behavior:

  • $$L_\text{var} = \frac{1}{D} \sum_{j=1}^D (1 - \sigma_j)^2$$3 provides non-vanishing gradients as $$L_\text{var} = \frac{1}{D} \sum_{j=1}^D (1 - \sigma_j)^2$$4, addressing collapse directly.
  • $$L_\text{var} = \frac{1}{D} \sum_{j=1}^D (1 - \sigma_j)^2$$5 maintains a corrective gradient under collapse, in contrast to the Epps–Pulley test employed by SIGReg, which vanishes.
• Complexity:
  • VISReg: Overall regularizer cost $$L_\text{var} = \frac{1}{D} \sum_{j=1}^D (1 - \sigma_j)^2$$6 (projection) plus $$L_\text{var} = \frac{1}{D} \sum_{j=1}^D (1 - \sigma_j)^2$$7 (sorting per slice). Sorting is subdominant at large $$L_\text{var} = \frac{1}{D} \sum_{j=1}^D (1 - \sigma_j)^2$$8.
  • VICReg: Covariance regularizer scales as $$L_\text{var} = \frac{1}{D} \sum_{j=1}^D (1 - \sigma_j)^2$$9.
  • SIGReg: Epps–Pulley-based sketching involves Fourier evaluation, $\sigma_j \to 0$0, and vanishing gradients.
• Scalability: Empirical accuracy is stable for $\sigma_j \to 0$1, but cross-GPU distribution enables constant $\sigma_j \to 0$2 in aggregate as $\sigma_j \to 0$3 increases.

5. Empirical Performance and Comparative Results

• Out-of-distribution (OOD) linear probe: VISReg-B/16 achieves 70.2% (DTD, Galaxy10, AID, ChestXRay, RetinaMNIST, OrganAMNIST), outperforming DINO, MoCoV3, and iBOT (69.5–69.6%) without reliance on teacher–student or special heuristics.
• Long-tailed regimes (ImageNet-LT): With $\sigma_j \to 0$4, VISReg attains 35.14% overall accuracy (versus VICReg 33.08%, SIGReg ~32%).
• Low-rank structure (Galaxy10): 80.76% (compared to SWD 80.60%, SIGReg 80.50%, VICReg 79.93%, DINO ~73.5%).
• In-domain (ImageNet-1K): VISReg-B/16 at 400 epochs reaches 75.7% (LeJEPA 75.6%, MAE 75.1%; DINO uses teacher–student heuristics for 78.2%).
• Large-scale pretraining (ImageNet-22K): VISReg-L/14 (100 ep) matches DINOv2-LVD142M on OOD accuracy (72.94% vs. 72.93%) with only 10% as much data.
• Transfer learning: VISReg surpasses DINO by 0.2–0.4% on CIFAR10/100, Flowers, IN-1K, and Galaxy10.
• Dense prediction tasks: ADE20K linear segmentation: 30.16 mIoU (compared to MoCoV3 31.69, DINO 29.40).
• Generative guidance: On iREPA with SiT-B/2, VISReg exceeds DINO across IS, gFID, precision, and recall.

6. Ablations and Robustness Analyses

• Regularization vs. prediction interpolation ($\sigma_j \to 0$5): Optimal performance for $\sigma_j \to 0$6 ($\sigma_j \to 0$769.19% on ViT-B/16 at 100 epochs); performance degrades at higher lambda.
• Learning rate sensitivity: Peak observed near $\sigma_j \to 0$8.
• Batch size: Best accuracy at 512.
• Projection dimension $\sigma_j \to 0$9: While $\nabla_{\sigma_j} L_\text{var} \rightarrow$0 is marginally beneficial for ImageNet-1K, lower dimensions benefit OOD and segmentation tasks; $\nabla_{\sigma_j} L_\text{var} \rightarrow$1 robustly performs well.
• Number of slices $\nabla_{\sigma_j} L_\text{var} \rightarrow$2: Robustness observed even for $\nabla_{\sigma_j} L_\text{var} \rightarrow$3; accuracy plateaus above $\nabla_{\sigma_j} L_\text{var} \rightarrow$4 with multi-GPU distribution supporting scalability.
• Component ablation: Removing $\nabla_{\sigma_j} L_\text{var} \rightarrow$5 or $\nabla_{\sigma_j} L_\text{var} \rightarrow$6 dramatically reduces performance (20–33% vs. 91.8% baseline), while removing $\nabla_{\sigma_j} L_\text{var} \rightarrow$7 slows but does not destabilize convergence ($\nabla_{\sigma_j} L_\text{var} \rightarrow$8 impact).
• Detach analysis: Detaching scale from shape gradients yields consistent improvements (0.3–0.6%) across classification, OOD generalization, and segmentation tasks.
• Loss-weighting ratios: Assigning up to $\nabla_{\sigma_j} L_\text{var} \rightarrow$9 in favor of the sketch (shape) term increases accuracy on low-quality datasets (e.g., ImageNet-LT +3.2%, Galaxy10 +1.3%), whereas balanced weights suit high-quality settings like Imagenette.

---

VISReg provides a theoretically principled, computationally efficient, and empirically validated regularizer for JEPA training. By unifying robust scale control with Sliced-Wasserstein-based shape matching, VISReg achieves state-of-the-art out-of-distribution generalization and resilience to distributional pathologies, while allowing flexible adaptive weighting suited to a range of data qualities and tasks (Wu et al., 1 Jun 2026).