文章摘要
自监督学习是JEPA世界模型核心技术,但面临表征坍塌难题,此前方法各有缺陷。近日,自监督学习新方法VISReg获图灵奖得主认可。它将正则项解耦为「尺度」与「形状」,代码轻量,计算复杂度和扩展性有优势。实验显示,其在多维度表现领先,为表征坍塌问题提供新正则化解法。

自监督学习是JEPA世界模型的核心技术路径,其无需人工标注即可从海量数据中学习通用表征,但长期面临表征坍塌的核心难题:模型倾向于将不同输入映射到高度相似的向量空间,看似完成训练,实则未学到具备判别力的特征。

此前的自监督正则化方法主要分为两类:一类依赖EMA、教师-学生网络等启发式技巧抑制坍塌,但这类方法训练过程脆弱、调参难度高,且可解释性与扩展性不足;另一类通过正则项直接约束表征分布,比如LeCun团队提出的VICReg,将学习目标拆分为方差、不变性与协方差三项,通过协方差约束特征维度间的相关性,但协方差仅能刻画二阶统计量,无法区分均值、方差一致但分布形状迥异的表征。

后续提出的SIGReg基于Cramér–Wold定理,通过sketching技术将嵌入分布对齐到标准高斯分布,完整约束了分布形状,但仍存在两个关键缺陷:一是当表征开始坍塌时,SIGReg的梯度会随之衰减,坍塌越严重修正信号越弱,模型难以自行恢复;二是未将表征的幅度大小(尺度)与分布形态(形状)两个独立属性分离,二者在优化过程中相互干扰,在长尾、低质量、低秩数据上的适配性较差。

近日,自监督学习新方法VISReg(Variance-Invariance-Sketching Regularization)获得图灵奖得主Yann LeCun的高度认可,他评价该方法是VICReg到SIGReg的自然延续,精准解决了此前方法的核心痛点。

核心技术设计

VISReg的核心思路是将防止表征坍塌的正则项解耦为「尺度」与「形状」两个独立目标,通过停止梯度操作彻底切断二者的优化干扰,同时无需依赖任何启发式训练技巧。整个正则目标由三部分组成:

  1. 尺度正则(Scale Regularization)

    该部分约束每一维特征的方差,防止表征幅值坍缩。其关键特性在于:当模型出现坍缩时,该项的梯度会趋近于一个常数,从而保证模型能够稳定地从坍缩状态中恢复,弥补了SIGReg的梯度消失缺陷。

  2. 形状正则(Shape Regularization)

    该部分先通过带停止梯度的归一化操作消除尺度影响,再单独约束分布形状。归一化时对标准差施加停止梯度,确保形状损失的优化不会反过来改变表征尺度,真正实现了「尺度」与「形状」两个目标的解耦。之后通过切片Wasserstein距离将分布的几何形状对齐到各向同性高斯分布,依据Cramér–Wold定理,只要将高维表征沿足够多的随机一维方向切片后逐一对齐到高斯分布,就等价于在高维空间对齐整个分布,从而完整刻画了分布形状而非仅二阶统计量。

  3. 中心化损失

    该部分是一个将batch均值拉向原点的损失,进一步优化表征的分布特性。

预测损失沿用JEPA/LeJEPA的不变性目标,让全局视角与局部视角的嵌入向全局视角的均值对齐,最后通过单一超参λ在预测损失与正则损失之间平衡,得到完整的训练目标。

轻量的代码实现

VISReg的正则目标实现非常轻量,核心逻辑仅需约15行PyTorch代码:

def visreg(z, K=64):
    # 1. 中心化损失
    mu = z.mean(dim=0)
    L_center = mu.pow(2).mean()
    # 2. 尺度损失
    z_cent = z - mu
    std = z_cent.std(dim=0, unbiased=False)
    L_scale = (1.0 - std).pow(2).mean()
    # 3. 形状损失:切片 Wasserstein 距离
    z_norm = z_cent / (std.detach())
    W = torch.randn(D, K)
    W /= W.norm(p=2, dim=0)
    p_sorted = torch.sort(z_norm @ W, dim=0).values
    u = torch.arange(1, N+1) / (N+1)
    target = Normal(0, 1).icdf(u)
    L_shape = (p_sorted - target).pow(2).mean()
    return L_scale + L_shape + L_center

计算复杂度与扩展性

在计算复杂度与扩展性方面,VISReg同样具备优势。其正则部分的计算复杂度为O(N*D*K)(N为batch大小、D为特征维度、K为切片数),所有扩展因子均为线性;而VICReg的协方差项复杂度为O(D²),随特征维度平方增长。在同等batch规模下,VISReg在单块H100 GPU上的运行速度与显存占用均优于SIGReg。

更重要的是,K个随机切片可以分摊到多块GPU上,在M块GPU上每块生成K/M个切片即可等价于单卡生成全部K个切片。实验显示,当单卡切片数不足时,使用8卡每卡128个切片(合计1024),与单卡1024切片的精度差距仅为0.22%,大幅缩小了精度损失,使得大规模训练时可以保持常数K,几乎不增加单卡负担。

全面的实验验证

研究团队在15个数据集(8个域内 + 6个分布外 + ADE20K稠密预测)上,将VISReg与MoCoV3、DINO、iBOT、I-JEPA、MAE、data2vec等7种主流自监督方法进行了对比,场景涵盖天文、医疗、遥感、纹理、花卉等,其优势体现在多个维度:

  1. 域内线性探测

    在不使用任何启发式技巧的前提下,VISReg在域内线性探测任务中表现领先:ViT-B/16的域内精度达到75.7%,高于MAE的75.1%;ViT-L/14进一步提升至77.0%,高于LeJEPA的75.6%。与使用启发式技巧的iBOT、DINO相比,VISReg在常规数据集上仅略低,但在纹理数据集DTD上反超全部方法,证明其跨域泛化能力源于方法本身而非人工技巧堆叠。

  2. 分布外泛化:全面最优

    分布外泛化是比域内精度更严格的测试,依赖启发式的方法常在ImageNet域内被充分调优,但难以迁移到差异较大的新分布。研究团队在医疗、天文、遥感、纹理等6个与ImageNet训练域完全无关的OOD数据集上进行评测,结果显示VISReg在所有方法、所有骨干规模上均取得了最佳平均OOD精度,甚至超过部分使用启发式技巧的更大骨干模型。ViT-B/16的VISReg平均OOD精度为70.19%,ViT-L/14为70.63%,明显高于MAE的67.85%,并优于MoCoV3、DINO、I-JEPA等方法。

  3. 数据效率:以1/10数据比肩DINOv2

    将VISReg(ViT-L/14)在ImageNet-22K(约1400万张图像)上预训练后,其6个OOD数据集的平均精度达到72.94%,与在10倍规模的LVD-142M(1.42亿张图像)上训练的DINOv2(72.93%)基本持平,也就是说VISReg仅用约1/10的数据量就达到了同等水平。作为对照,同为ViT-L/14但仅用ImageNet-1K预训练的VISReg平均精度为70.63%,证明其学到的表征具备极强的通用性。

  4. 迁移微调:全面超过DINO

    尽管VISReg在部分域内数据集上的线性探测精度略低于DINO,但经过微调后,它在CIFAR-10、CIFAR-100、Flowers、ImageNet-1K、Galaxy10全部五个数据集上均超过DINO与有监督预训练,表明其表征分布更均匀、冗余更低、可迁移性更强。

  5. 稠密预测与生成引导

    VISReg的优势不止于分类任务。在ADE20K线性语义分割任务(ViT-B/16)中,其mIoU为30.16,高于DINO的29.40与MAE的23.60,仅次于MoCoV3的31.69,在不使用任何启发式技巧的前提下取得了极具竞争力的结果。在生成引导任务中,由VISReg特征引导的生成在四项指标中的三项优于DINO,证明VISReg学到的表征作为生成引导信号同样更优。

  6. 低质量数据鲁棒性

    在长尾分布(ImageNet-LT)与低秩(Galaxy10)等低质量数据集上,VISReg能够稳定防止表征坍塌并学到有意义的表征,而DINO在缺乏精细调参时直接失败。在ImageNet-LT上,ViT-S/8从头训练400 epoch后,DINO的Overall精度仅为5.13%,几近完全失败,而VISReg(增大形状损失权重)取得了全面最优;在Galaxy10低秩任务中,SIGReg、SWD、VISReg均能成功避免训练坍缩并取得良好精度,而DINO难以学到有意义的特征。

结论

VISReg的研究表明,将表征正则解耦为「尺度」与「形状」两个独立组件,可以得到比现有方法更稳定、更高效、泛化性更强的自监督学习方法。在不使用任何训练启发式技巧的前提下,它在图像识别、分割与生成引导等多个维度上取得了领先或接近领先的结果,并以约1/10的数据达到了DINOv2的OOD水平,为JEPA世界模型长期存在的表征坍塌问题提供了一种全新的正则化解法。

以上内容不代表本平台立场,仅供读者参考