VISReg: Variance-Invariance-Sketching Regularization for JEPA training

📄 arXiv: 2606.02572v1 📥 PDF

作者: Haiyu Wu, Randall Balestriero, Morgan Levine

分类: cs.CV

发布日期: 2026-06-01

🔗 代码/项目: PROJECT_PAGE

💡 一句话要点

VISReg：方差-不变性-素描正则化方法，提升JEPA训练的稳定性和泛化性

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 自监督学习 正则化 方差不变性 Sliced-Wasserstein距离 分布对齐

📋 核心要点

  1. 现有自监督学习方法在嵌入空间正则化方面存在不足，如VICReg仅考虑二阶统计量，SIGReg缺乏灵活性且易梯度消失。
  2. VISReg通过引入基于Sliced-Wasserstein的素描目标，在保留方差控制的同时，强制执行完整的分布形状，解耦尺度和形状。
  3. 实验表明，VISReg在低质量数据集、长尾分布和低秩情况下表现出色，并在ImageNet预训练后，OOD性能达到SOTA。

📝 摘要（中文）

自监督学习方法通常通过建模启发式规则或显式正则化嵌入空间来防止嵌入崩溃。其中，VICReg将正则化分解为方差和协方差目标，具有灵活性和可解释性。然而，协方差仅捕获二阶统计量，鼓励去相关性，但未能强制执行稳定训练所需的完整分布形状。基于素描的方法（如SIGReg）通过将嵌入与各向同性高斯对齐来解决此问题，但缺乏灵活性，并且在崩溃下会遭受梯度消失。我们提出了方差-不变性-素描正则化（VISReg），它用基于Sliced-Wasserstein的素描目标代替协方差，该目标强制执行完整的分布形状，同时保留方差项以进行尺度控制。通过解耦尺度和形状，VISReg结合了VICReg的灵活性和素描方法的分布严谨性，即使在崩溃下也能提供鲁棒的梯度。我们表明，VISReg线性缩放，在低质量数据集上优于现有正则化方法，并且对长尾和低秩状态具有弹性。在ImageNet-1K上预训练后，VISReg在分布外数据集上实现了最先进的性能。在ImageNet-22K上预训练后，它与DINOv2的OOD性能相匹配，尽管后者使用了10倍的数据（LVD-142M）。项目和代码：https://haiyuwu.github.io/visreg。

🔬 方法详解

问题定义：现有的自监督学习方法，如VICReg，虽然通过方差和协方差正则化来避免嵌入崩溃，但协方差仅能捕获二阶统计信息，无法保证嵌入的完整分布形状。而SIGReg虽然能对齐到高斯分布，但缺乏灵活性，且在嵌入崩溃时容易出现梯度消失问题。因此，需要一种既能保证分布形状，又能保持灵活性的正则化方法。

核心思路：VISReg的核心思路是将VICReg中的协方差正则化替换为基于Sliced-Wasserstein距离的素描（Sketching）目标。Sliced-Wasserstein距离能够比较两个分布的差异，从而强制嵌入具有期望的分布形状（例如，各向同性高斯分布）。同时，保留方差项来控制嵌入的尺度，从而解耦形状和尺度，提高训练的稳定性和鲁棒性。

技术框架：VISReg的整体框架与VICReg类似，都是通过最小化一个包含方差、不变性和素描（或协方差）的损失函数来实现自监督学习。主要流程包括：1）对输入图像进行数据增强，得到两个不同的视图；2）通过编码器网络提取两个视图的嵌入表示；3）计算方差损失，保证嵌入具有足够的方差；4）计算不变性损失，使两个视图的嵌入尽可能一致；5）计算基于Sliced-Wasserstein距离的素描损失，使嵌入的分布形状接近目标分布。

关键创新：VISReg的关键创新在于使用Sliced-Wasserstein距离来度量嵌入分布与目标分布之间的差异，并将其作为正则化项。与传统的协方差正则化相比，Sliced-Wasserstein距离能够捕获更高阶的统计信息，从而更好地保证嵌入的分布形状。此外，通过解耦尺度和形状，VISReg能够更灵活地控制嵌入的各个方面，提高训练的稳定性。

关键设计：VISReg的关键设计包括：1）使用Sliced-Wasserstein距离作为素描损失，具体实现中，通过随机投影将高维嵌入投影到多个一维空间，然后在每个一维空间上计算Wasserstein距离；2）平衡方差损失、不变性损失和素描损失的权重，以获得最佳的性能；3）选择合适的编码器网络结构，例如ResNet或Vision Transformer。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

VISReg在ImageNet-1K上预训练后，在多个OOD数据集上取得了SOTA性能。在ImageNet-22K上预训练后，VISReg的OOD性能与使用10倍数据的DINOv2相当，证明了其高效性和鲁棒性。此外，实验还表明，VISReg在低质量数据集、长尾分布和低秩情况下表现优于现有的正则化方法。

🎯 应用场景

VISReg可应用于各种自监督学习任务，尤其是在数据质量较低或数据分布不平衡的情况下。例如，在医学图像分析、遥感图像处理等领域，由于数据获取成本高昂，数据量通常较小，且数据分布可能存在偏差。VISReg能够提高模型在这些场景下的泛化能力，从而降低对大量标注数据的依赖。

📄 摘要（原文）

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.