从SwAV看无监督学习的未来:为什么说聚类+对比学习=下一代视觉表征?

如果你在过去几年里关注过计算机视觉领域,尤其是自监督学习这片热土,你可能会被各种层出不穷的“对比学习”方法弄得眼花缭乱。从SimCLR到MoCo,再到BYOL,大家似乎都在玩一个“找不同”的游戏:把同一张图片的不同“视角”拉近,把不同图片的视角推远。这个范式效果惊人,但代价也显而易见——它像个“内存怪兽”,对计算资源有着近乎贪婪的需求,尤其是对大批量数据(large batch size)的依赖,让许多研究者和工程师望而却步。

就在大家思考如何为这场“对比盛宴”降温时,一篇名为SwAV的论文在NeurIPS 2020上横空出世。它没有选择在“对比”的道路上继续内卷,而是做了一次优雅的“范式融合”:将经典的聚类思想与前沿的对比学习框架巧妙地编织在一起。结果呢?它不仅在小批量设置下取得了媲美甚至超越当时SOTA的性能,更重要的是,它为我们打开了一扇新窗,让我们重新思考:学习一个好的视觉表征,是否一定要进行显式的、一对一的特征比较?这篇文章,我们就来深入拆解SwAV的智慧,并探讨它如何为无监督学习的未来指明了一条更具可扩展性和实用性的道路。

1. 破局之路:SwAV如何化解对比学习的“内存之困”

要理解SwAV的价值,我们得先回到它试图解决的问题原点。传统的基于实例判别(Instance Discrimination)的对比学习,其核心逻辑可以概括为“拉近正样本,推远负样本”。这里的“正样本”通常来自同一原始图像的不同数据增强视图(例如,随机裁剪、颜色抖动后的版本),而“负样本”则来自批次中所有其他图像。

这种范式存在两个天然的瓶颈:

  1. 计算复杂度高:损失函数(如InfoNCE Loss)需要计算当前样本与批次内所有其他样本的相似度。为了获得足够多的负样本以提供有意义的“推远”信号,批次大小(batch size)必须非常大。这在实践中意味着对GPU显存的巨大压力。
  2. 信息利用效率低:每个负样本都被平等地(或以某种加权方式)“推远”,但并非所有负样本都同等重要。模型可能会浪费大量计算在区分早已远离的、不相关的样本对上。

SwAV的创始人Mathilde Caron等人洞察到了一个关键点:我们真的需要和每一个负样本“硬碰硬”地比较吗?或许,我们可以引入一个“中间人”——聚类中心(prototypes)——来作为特征的抽象锚点。

提示:你可以把prototypes想象成特征空间中的一些“地标”或“概念原型”。模型不再学习区分“这是A图片还是B图片”,而是学习判断“这个特征更靠近地标1还是地标2”。

SwAV的核心创新“交换预测”(Swapped Prediction)机制,正是基于这个想法。给定一张图片的两个增强视图,我们分别提取它们的特征 z_tz_s。同时,我们维护一组可学习的聚类中心 C。对于每个特征,我们不是直接用它去对比其他特征,而是计算它属于每个聚类中心的“软分配”概率,得到一个编码(code)q。这个编码本质上是一个概率分布,描述了特征在各个“概念原型”上的归属情况。

接下来,SwAV的巧妙之处来了:它要求用视图 t 的特征 z_t,去预测视图 s 的编码 q_s;同时,用视图 s 的特征 z_s,去预测视图 t 的编码 q_t。损失函数就是这两个预测任务的交叉熵损失之和。

# 伪代码示意SwAV的核心损失计算
def swapped_prediction_loss(z_t, z_s, prototypes_C, temperature_tau):
    # 计算特征与所有原型之间的相似度
    scores_t = torch.matmul(z_t, prototypes_C.T) / temperature_tau
    scores_s = torch.matmul(z_s, prototypes_C.T) / temperature_tau

    # 通过softmax得到“预测”的概率分布 p
    p_t = F.softmax(scores_t, dim=1)  # 用z_t预测的原型分布
    p_s = F.softmax(scores_s, dim=1)  # 用z_s预测的原型分布

    # 通过Sinkhorn-Knopp算法计算“目标”编码 q (详见下文)
    q_t = compute_swav_code(z_t, prototypes_C)
    q_s = compute_swav_code(z_s, prototypes_C)

    # 交换预测损失:用z_t去匹配q_s,用z_s去匹配q_t
    loss = cross_entropy(p_t, q_s) + cross_entropy(p_s, q_t)
    return loss

这个设计带来了几个立竿见影的好处:

  • 摆脱大批量依赖:模型不再需要与批次内所有其他样本直接比较,只需与固定数量的K个原型进行比较。因此,即使在小批量(如256)下也能有效训练。
  • 计算效率提升:计算复杂度从 O(B^2) (B为批次大小)降低到了 O(B*K) (K为原型数,通常为几千),且K远小于B。
  • 隐式对比:虽然形式上没有负样本,但通过迫使同一图像的不同视图预测出相同的原型分配,并利用原型间的竞争关系(一个特征主要归属于少数原型),模型实际上在学习一种结构化的、隐式的对比。

2. 算法引擎:在线聚类与Sinkhorn-Knopp算法的精妙共舞

SwAV中另一个关键突破是实现了真正意义上的“在线聚类”。早期的基于聚类的方法,如DeepCluster,采用“聚类-训练”交替进行的离线模式。即先用所有数据聚类得到伪标签,再用这些标签训练网络,如此反复。这种方式效率低下,且难以扩展到海量数据。

SwAV将聚类过程无缝地嵌入到了每个训练步骤中。具体来说,对于当前批次计算出的所有特征 Z 和原型矩阵 C,我们需要为每个特征分配一个编码 Q(即 q 的集合)。SwAV将这个问题优雅地转化为一个最优传输(Optimal Transport)问题

其目标是找到一个分配矩阵 Q,使得所有特征被“运输”到原型上的总成本(定义为负的相似度)最小,同时满足两个约束:

  1. 每个原型被选择的平均次数是均匀的(防止所有特征都坍缩到同一个原型上)。
  2. 分配矩阵本身具有一定的平滑性(通过熵正则化实现)。

数学上,这个优化问题可以表述为:

[ \max_{Q \in \mathcal{Q}} \text{Trace}(Q^T C^T Z) + \epsilon H(Q) ] 其中约束集 (\mathcal{Q} = { Q \in \mathbb{R}_+^{K \times B} \mid Q \mathbf{1}_B = \frac{1}{K} \mathbf{1}_K, Q^T \mathbf{1}_K = \frac{1}{B} \mathbf{1}_B }),H(Q)是矩阵Q的熵,ε是正则化系数。

注意:这里的均匀分配约束(Q1_B = (1/K)1_K)是防止模型坍缩的关键。它强制要求每个原型在每个批次中平均被选中约 B/K 次,从而保证了特征在原型空间的均匀分布,这与对比学习中对“均匀性”(uniformity)的追求异曲同工。

直接求解这个带约束的优化问题并不容易。SwAV借鉴了SeLa论文中的方法,采用Sinkhorn-Knopp算法来高效地求取近似解。该算法通过迭代的行、列归一化操作,能快速收敛到一个满足边际约束的分配矩阵。作者甚至在GPU上进行了高效实现,使得每次计算开销极小。

# Sinkhorn-Knopp算法的简化示意
def sinkhorn_knopp(scores, epsilon=0.05, num_iters=3):
    """
    scores: 相似度矩阵,形状为 [K, B]
    返回:近似最优传输分配矩阵 Q,形状为 [K, B]
    """
    Q = torch.exp(scores / epsilon)  # 指数化
    for _ in range(num_iters):
        # 行归一化,满足对原型的均匀分配约束
        Q /= torch.sum(Q, dim=1, keepdim=True)
        Q /= K  # 使得行和为 1/K
        # 列归一化,满足每个特征分配概率和为1的约束
        Q /= torch.sum(Q, dim=0, keepdim=True)
        Q /= B  # 使得列和为 1/B
    return Q

这种在线最优传输聚类,是SwAV的灵魂。它使得模型能够动态地、在训练过程中同步地学习原型和特征分配,实现了聚类与表征学习的统一优化。

3. 实战加速器:Multi-Crop策略为何是通用“涨点神器”

除了核心算法,SwAV论文中还提出了一个简单却极其有效的技巧——Multi-Crop数据增强。这个策略虽然独立于SwAV算法本身,但因其强大的通用性和有效性,迅速成为了自监督学习领域的标配。

当时的主流方法(如SimCLR)通常使用“两次裁剪”策略:从一张原图中随机裁剪出两个标准尺寸(例如224x224)的区域作为一对正样本。作者发现,增加视图(crop)的数量能显著提升性能,因为这为模型提供了更多样化的局部上下文信息。然而,简单地增加高分辨率裁剪的数量会线性增加计算成本。

Multi-Crop的聪明之处在于混合分辨率裁剪。它不再只生成两个大尺寸视图,而是生成:

  • 2个“全局”视图(global views):标准分辨率(如224x224)的大裁剪,覆盖图像的主要部分。
  • V个“局部”视图(local views):更低分辨率(如96x96)的小裁剪,专注于图像的局部细节。

在计算损失时,只对2个全局视图计算编码(code),然后将所有视图(2个全局+V个局部)的特征与这两个编码进行交换预测。也就是说,局部视图的特征被用来预测全局视图的编码,反之亦然。

视图类型分辨率示例数量是否计算编码作用
全局视图224x2242提供场景级语义信息,作为预测的“锚点”
局部视图96x96V (通常为4或6)提供物体局部细节信息,增加数据多样性

这种设计实现了“四两拨千斤”的效果:

  1. 低成本增加多样性:小分辨率裁剪的计算和内存开销远小于大裁剪。增加4个96x96的视图,其总计算量可能还不及一个224x224的视图,但却极大地丰富了模型看到的局部模式。
  2. 促进局部一致性学习:模型被迫学会将一个小局部区域的特征,与一个更大、更全局的上下文所对应的语义原型(编码)对齐。这增强了模型对物体局部特征的识别和不变性理解,对于细粒度分类等下游任务尤其有益。
  3. 通用性强:该策略几乎可以无缝应用到任何基于图像视图对比或预测的方法中,如SimCLR、MoCo等,并能稳定带来性能提升。

在实际代码实现中,Multi-Crop策略使得数据加载和损失计算部分需要做一些调整,但结构清晰:

# 伪代码示意Multi-Crop下的损失计算
def multi_crop_swav_loss(all_features, global_codes, temperature_tau):
    """
    all_features: 列表,包含所有视图的特征 [z_global1, z_global2, z_local1, z_local2, ...]
    global_codes: 仅两个全局视图对应的编码 [q_global1, q_global2]
    """
    loss = 0
    # 遍历所有视图的特征
    for idx, z in enumerate(all_features):
        # 计算该特征与所有原型的相似度,得到预测p
        p = compute_prediction(z, prototypes_C, temperature_tau)
        # 该特征需要与“另一个”全局视图的编码计算损失
        if idx == 0:  # 第一个全局视图的特征,去匹配第二个全局视图的编码
            target_code = global_codes[1]
        elif idx == 1:  # 第二个全局视图的特征,去匹配第一个全局视图的编码
            target_code = global_codes[0]
        else:  # 局部视图的特征,随机与一个全局视图的编码匹配(或与两个都匹配取平均)
            # 论文中通常与两个全局编码都计算损失
            loss_to_global1 = cross_entropy(p, global_codes[0])
            loss_to_global2 = cross_entropy(p, global_codes[1])
            loss += (loss_to_global1 + loss_to_global2) / 2
            continue
        loss += cross_entropy(p, target_code)
    return loss

4. 理论再思考:Alignment与Uniformity在SwAV中的新诠释

在对比学习理论中,“对齐与均匀”(Alignment and Uniformity)是一个被广泛接受的框架,用于解释为何对比学习有效。对齐指的是正样本对在特征空间中的距离应该小;均匀指的是所有样本的特征应该尽可能均匀地分布在单位超球面上,以保留最大信息。

那么,SwAV没有显式的负样本,它是如何实现“均匀性”的呢?这正是其设计最精妙的地方。SwAV通过在线最优传输的均匀分配约束,间接地强制执行了均匀性。

  • 隐式的对齐:交换预测机制强制同一图像的不同视图预测出相同的原型分配,这等价于要求这些视图的特征在原型空间中靠近(即与相同的原型有高相似度),从而实现了特征级别的对齐。
  • 结构化的均匀:约束每个原型被平均选择,意味着学习到的原型向量本身必须在特征空间中均匀散开,以“覆盖”所有不同的数据模式。同时,特征为了被分配给这些均匀散布的原型,它们自身在特征空间中的分布也会被诱导得更加均匀。这是一种通过约束中间变量(原型分配)来间接约束特征分布的高明手段。

我们可以通过一个表格来对比三种范式的差异:

特性传统对比学习 (如SimCLR)离线聚类 (如DeepCluster)SwAV (在线聚类对比)
学习方式在线离线交替在线
对比目标特征 vs 特征特征 vs 聚类中心特征 vs 聚类分配(编码)
均匀性保证靠大批量负样本显式推开靠聚类算法本身(如K-means)靠最优传输的均匀分配约束
计算复杂度O(B^2)O(迭代数 * (数据量 * K))O(B*K)
内存效率低(需大批量或内存队列)中(需存储所有特征用于聚类)高(仅需存储原型)
扩展性受限于GPU内存受限于聚类算法效率优秀,可扩展到极大数据库

这种范式转换带来了更深层的启示:学习一个好的表征,未必需要直接进行样本间的两两比较。通过引入一个结构化的、可学习的中间表示层(如原型),并施加合适的约束,我们可以在保持甚至提升性能的同时,大幅降低算法的复杂度和资源消耗。这为将自监督学习应用到更大规模数据集、更复杂模态(如视频、3D点云)提供了新的思路。

5. 超越ImageNet:SwAV启示下的未来方向与工程实践

SwAV在ImageNet上75.3%的线性评估准确率证明了其强大实力,但其影响远不止于此。它在下游任务(如目标检测、语义分割)上的迁移性能甚至超越了有监督的ImageNet预训练模型,这凸显了其学习到的特征具有更强的泛化性和语义性。

从SwAV出发,我们可以窥见无监督学习未来几个值得探索的方向:

1. 原型的学习与解释性 SwAV中的原型是可学习的参数。这些原型最终学到了什么?它们是否对应着可解释的视觉概念(如“带轮子的物体”、“纹理区域”)?后续工作可以探索如何可视化或解释这些原型,将无监督学习从“黑箱”推向“可解释”。例如,可以尝试将原型反向映射回图像空间,观察其激活的图像区域。

2. 跨模态与层次化原型 当前SwAV应用在单模态(图像)上。一个自然的扩展是将其应用于视频、音频-视觉等多模态数据。可以设计跨模态的交换预测任务,并学习一组共享的原型,作为连接不同模态的“语义枢纽”。此外,是否可以学习层次化的原型结构(粗粒度到细粒度),以更好地捕获数据的层次化语义?

3. 面向超大规模数据的极致优化 SwAV的在线和高效特性使其非常适合超大规模数据训练。未来的工程优化可以集中在:

  • 异步与分布式原型更新:在分布式训练中,如何高效同步和更新全局原型集合。
  • 原型生命周期管理:数据分布可能随时间变化,是否需要动态地增加、合并或淘汰原型?
  • 与Transformer架构的结合:Vision Transformer已成为主流,SwAV如何适配Transformer的特征特性进行优化?例如,原型是否可以作用于patch级别的特征?

4. 实际部署中的技巧与陷阱 如果你打算在自家数据集上尝试SwAV或类似思想,这里有几个从实践中来的心得:

  • 原型数量K的选择:这是一个关键超参数。太小会导致特征区分度不足,太大会增加计算量且可能过拟合。通常可以从数据集类别数的5-10倍开始尝试。ImageNet上常用3000个原型。
  • 温度系数τ:控制预测分布的平滑程度。τ小则分布更“尖锐”,τ大则更平滑。需要与学习率等参数一起仔细调优。
  • 小批量下的特征累积:当批量很小时,为了满足最优传输算法的数值稳定性,需要累积多个批次的特征后再计算编码。这个“累积步数”需要平衡:步数太少,特征集合小,分配不稳定;步数太多,编码更新延迟大,影响学习效率。
  • 初始化的艺术:原型的初始化很重要。一种常见且有效的方法是使用随机正交矩阵进行初始化,这能保证初始原型在超球面上均匀分布。

我在一些中型规模的工业数据集上复现SwAV时发现,直接套用ImageNet的超参有时效果并不理想。特别是当数据域与自然图像差异较大时(例如医学影像、卫星图),需要重新调整数据增强策略(Multi-Crop的尺寸、数量),甚至考虑设计领域特定的增强方法。此外,损失函数的权重平衡有时也需要微调,尤其是在结合了其他辅助任务的情况下。

SwAV更像是一个思想框架,它告诉我们“聚类”与“对比”并非对立,而是可以协同进化的统一体。它降低了自监督学习的入门门槛,让更多人能够在有限资源下探索无监督表征的奥秘。其核心的“交换预测”和“在线约束分配”思想,也正在被后续的许多工作(如DINO、Barlow Twins等)以不同形式吸收和演化。无监督学习的未来,注定是属于这些更高效、更优雅、更接近人类学习本质的范式。

Logo

北京人形旗下天工造物具身智能开源社区,聚焦具身天工与慧思开物两大平台

更多推荐