提出VISReg正则化方法,提升JEPA训练稳定性与泛化性能。
VISReg: Variance-Invariance-Sketching Regularization for JEPA training

- 用Sliced-Wasserstein距离约束嵌入分布形状,取代传统协方差项。
- 在低质量、长尾数据上表现优于现有方法,且训练更稳定。
- 适合追求高鲁棒性自监督模型的科研与工程人员使用。
自监督学习通过建模启发式或显式正则化嵌入空间防止嵌入坍塌。其中,VICReg将正则化分解为方差与协方差目标,具备灵活性和可解释性。然而,协方差仅捕捉二阶统计量,仅促进解相关,无法强制所需的完整分布形态以实现稳定训练。基于草图的方法如SIGReg通过将嵌入对齐至各向同性高斯分布来弥补此缺陷,但缺乏灵活性且在坍塌时梯度消失。我们提出方差-不变性-草图正则化(VISReg),以基于切片-沃瑟斯坦距离的草图目标替代协方差,强制完整分布形态,同时保留方差项控制尺度。通过解耦尺度与形状,VISReg结合了VICReg的灵活性与草图方法的分布严谨性,在坍塌情况下仍提供稳健梯度。实验表明,VISReg线性可扩展,在低质量数据集上超越现有正则化方法,对长尾和低秩情形具有鲁棒性。在ImageNet-1K上预训练,其在分布外数据集上达到顶尖性能;在ImageNet-22K上预训练,虽仅使用DINOv2的1/10数据(LVD-142M),性能仍与之相当。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。