1. SwAV是什么为什么说它解决了对比学习的“痛点”如果你最近关注过无监督学习尤其是视觉领域的自监督学习肯定会听过SimCLR、MoCo这些如雷贯耳的名字。它们都属于对比学习Contrastive Learning的范畴核心思想很简单让同一张图片的不同“视角”比如经过裁剪、变色等数据增强后的版本在特征空间里靠得近一些让不同图片的特征离得远一些。这个想法很直观效果也确实惊人但有个大问题——太“烧”显存和算力了。为什么因为要让“不同图片的特征离得远”你就需要大量的“负样本”。想象一下你教一个模型认识猫不仅要给它看很多猫的图片正样本还得给它看海量的狗、车、房子的图片负样本告诉它“这些不是猫”。在对比学习中为了学到一个好的特征通常需要成千上万个负样本同时参与计算这就意味着你的训练批次Batch Size必须非常大。SimCLR当年在ImageNet上刷出新高背后是4096的超大Batch Size和巨量的TPU算力支持这对绝大多数研究者和开发者来说简直是“可望而不可及”的奢侈品。这时候SwAVSwapping Assignments between multiple Views出现了。我第一次读到这篇NeurIPS 2020的论文时感觉就像在密不透风的对比学习框架里打开了一扇新的天窗。它没有走“堆砌负样本”的老路而是巧妙地引入了“聚类”的思想。简单来说SwAV不再让特征直接去和成千上万个其他特征对比而是让特征去学习匹配一个“聚类分配结果”。你可以这样理解我们预先定义好K个“原型”Prototypes你可以把它们想象成K个虚拟的、可学习的类别中心。对于一张图片的某个视角我们不仅提取它的特征还计算它属于这K个原型中每一个的概率分布这个分布叫做“code”。SwAV的核心创新在于“交换预测”用图片A视角一的特征去预测图片A视角二的code同时用视角二的特征去预测视角一的code。模型的目标是让这两个预测尽可能准确。这样做的好处是颠覆性的。首先它彻底摆脱了对海量显式负样本的依赖。模型不再需要记住一个巨大的负样本队列如MoCo也不需要在一个超大的批次里进行两两比较如SimCLR。它只需要关注当前批次的特征和那K个原型计算开销和内存占用直线下降。论文里的实验显示用仅仅256的Batch SizeSwAV就能取得非常接近甚至超过那些需要超大Batch Size的对比学习方法的效果。这对于我们这些“平民玩家”来说意味着在有限的GPU资源上也能跑起来前沿的自监督学习模型意义重大。其次SwAV实现了一种真正的“在线聚类”。早期的基于聚类的方法比如同一作者团队的DeepCluster是离线的需要先把整个数据集的特征都提取出来做一遍聚类得到伪标签再用这些标签去训练网络步骤是割裂的。而SwAV把聚类和特征学习融合到了一个统一的前向-反向传播过程中模型一边学习如何更好地把特征映射到原型上一边也在优化这些原型本身让它们能更好地代表数据分布。这种“在线”特性使得SwAV能够处理海量数据并且更适配深度学习端到端训练的习惯。所以SwAV到底是个啥我更喜欢把它称为“对比学习与在线聚类的优雅融合”。它抓住了对比学习“学习不变性表征”的精髓但用聚类的“分配-预测”机制巧妙地规避了其最大的计算瓶颈。它不是为了创新而创新而是实实在在地解决了一个工程实践上的核心难题。2. 深入核心SwAV的“交换预测”机制是如何工作的理解了SwAV要解决的问题我们再来拆解它的核心引擎——“交换预测”Swapped Prediction。这是整个方法最精妙的部分也是它区别于传统对比学习和传统聚类方法的关键。我们一步步来看。假设我们有一张图片X。按照自监督学习的标准流程我们会对它进行两次随机数据增强比如随机裁剪、颜色抖动等得到两个不同的视图View记作 X_t 和 X_s。把它们分别送入同一个编码器网络比如ResNet得到两个归一化后的特征向量 z_t 和 z_s。这一步和SimCLR等方法是完全一样的。不一样的事情从现在开始。在SwAV中我们维护了一个可学习的参数矩阵C它的尺寸是 [D, K]。这里D是特征的维度K是我们预先设定的原型数量。你可以把C的每一列 c_k 看作一个原型向量代表了特征空间中的一个“锚点”或“概念”。接下来我们要为每个特征计算一个“软分配”向量也就是前面提到的code记作 q_t 和 q_s。这个code是一个K维的概率分布表示特征 z 属于每个原型的概率。那么这个code是怎么来的呢这里用到了一个非常有趣的工具最优传输Optimal Transport。我们不是简单地对特征和原型做softmax因为那样容易导致一个糟糕的“捷径解”——模型可能偷懒把所有特征都分配给同一个原型这样损失函数也能很低但学到的特征毫无意义。为了防止这种崩塌CollapseSwAV借鉴了之前工作SeLa的思想对code的分布施加了两个约束每个特征的概率分布之和为1这很自然。更重要的是要求所有样本对K个原型的分配是均匀的。也就是说在一个批次Batch里平均每个原型会被分配到大约 B/K 次B是批次大小。这个均匀分配约束非常关键它强制模型去发掘数据中多样的模式用尽所有原型从而避免了表征崩塌。求解这个带约束的优化问题就用到了Sinkhorn-Knopp算法。这是一个快速迭代算法能在GPU上高效地求解出最优的分配矩阵Q即所有特征的code。论文中提到将4000个特征分配到3000个原型上只需要大约35毫秒开销几乎可以忽略不计。好了现在我们有了特征 z_t, z_s 和它们对应的code q_t, q_s。SwAV的损失函数登场了。它的形式是一个“交叉预测”的交叉熵损失损失 交叉熵(用z_t预测q_s) 交叉熵(用z_s预测q_t)具体来说“用z_t预测q_s”是什么意思我们先用z_t和所有原型C计算一个相似度分布 p_t softmax(C^T * z_t / τ)τ是温度系数。这个 p_t 可以看作是特征z_t自己“认为”它属于各个原型的概率。然后我们让这个自认为的概率分布 p_t去逼近另一个视角下通过最优传输计算出的“权威”分配结果 q_s。反过来也一样。这个“交换”为什么有效它本质上是在强化“视角不变性”。两个视角来自同一张图片它们的内容本质是相同的因此它们被分配到的原型分布code也应该是一致的。模型通过交换预测学习到的是不管图片经历了怎样的裁剪、变色它都应该被映射到特征空间中相似的区域并且对应到同一组原型概念上。这比直接让两个特征向量靠近对比学习的做法增加了一层“语义抽象”因为原型本身是在训练中不断学习的可以理解为一些基础的视觉概念。我打个比方。传统的对比学习像是教孩子认苹果拿一个红苹果和一个青苹果两个视角告诉孩子“这两个都是苹果”拉近特征。SwAV则像是先定义好“水果”这个原型概念然后拿着红苹果问孩子“你觉得它属于‘水果’这个概念的概率有多大”同时拿着青苹果通过一套更复杂的规则最优传输已经得出了它属于“水果”的概率很高。然后让孩子用红苹果的观察去预测青苹果的结论。孩子为了预测准确就必须学会忽略颜色、大小的变化抓住“苹果”的本质特征。这样一来学习的目标从具体的像素级对比上升到了概念分配的一致性显得更加高效和鲁棒。3. 实战利器Multi-Crop策略——小技巧带来大提升如果说“交换预测”是SwAV的理论核心那么“Multi-Crop”就是它的实战利器一个简单却极其有效的技巧。这个技巧甚至独立于SwAV算法本身被后续很多自监督学习方法所采纳堪称“提点神器”。在讲Multi-Crop之前我们先看看之前的做法。无论是SimCLR、MoCo还是SwAV的基础版本在构造正样本对时通常都是从一张原图中随机裁剪出两个尺寸较大的区域例如把图片缩放到256x256再裁剪出两个224x224的区域作为两个视图。这两个大视图覆盖了原图的大部分区域重叠部分很多因此它们包含的语义信息高度一致非常适合作为正样本对。然而作者们思考了一个问题只依赖两个全局视图够吗显然不够。大视图能捕捉全局的、场景级别的信息但可能会忽略一些细节的、局部的特征。而人类识别物体既看整体也看关键局部。为了让模型学到更鲁棒、更细粒度的特征引入多尺度、多区域的视图是非常自然的想法。最直接的做法就是增加视图的数量比如从一张图裁剪出4个、6个甚至更多的大视图。但这样做的代价是显存和计算量会成倍增加因为每个视图都要经过编码器前向传播。对于已经因为大模型而“捉襟见肘”的GPU内存来说这无疑是雪上加霜。Multi-Crop策略的精妙之处就在于它用了一种“高低搭配”的方式在不显著增加开销的前提下引入了多尺度信息。具体做法是我们仍然裁剪两个标准分辨率的“全局视图”例如224x224它们被称为global views。此外我们再裁剪V个低分辨率的“局部视图”例如96x96它们被称为local views。这些不同尺寸的视图被一起送入同一个编码器网络。关键来了在计算SwAV损失时我们只对两个全局视图的特征计算code即q1和q2。而对于所有视图包括2个全局视图和V个局部视图的特征z我们都用它们去预测这两个全局视图的code。用公式表示总损失变为总损失 Σ_{i∈{1,2}} Σ_{v≠i} 交叉熵损失(z_v, q_i)这里的求和遍历了所有视图v共V2个去预测两个全局视图的code。也就是说局部视图的特征z_local要去预测全局视图的codeq_global。为什么这个设计如此巧妙首先它极大地丰富了正样本对的数量。原本只有一对view1, view2现在有了 (V2) * 2 对预测关系虽然有些是重复的。模型要从一个非常小的局部 patch可能只包含物体的一部分去推断出整个图像的语义分配这迫使编码器必须从局部特征中提取出具有高度辨识度的信息。这大大增强了特征的局部一致性学习。其次它几乎不增加计算负担。计算量最大的部分是什么是计算code最优传输。而Multi-Crop策略只对两个全局视图计算code局部视图只计算特征z。计算特征z的前向传播开销远小于计算最优传输的Sinkhorn迭代。同时由于局部视图的分辨率低如96x96经过下采样后其网络前向计算的实际开销也比全局视图小很多。论文中使用了V4个局部视图总计算量仅增加了不到20%但带来的性能提升却非常显著。在实际代码实现中这通常意味着你的数据加载器需要支持输出一个列表的图像块crops而你的网络需要能处理一个批次内不同分辨率的输入。在PyTorch中你可以将不同分辨率的视图视为独立的批次分别通过编码器然后在损失计算层再按照上述规则组织起来。这个技巧的普适性很强我后来在MoCo v2、BYOL等方法的复现中也尝试加入了Multi-Crop基本上都能稳定带来1-2个百分点的精度提升。4. 从理论到代码手把手解析SwAV的实现细节理解了原理和技巧我们最关心的还是这东西到底怎么实现光说不练假把式下面我就结合论文伪代码和核心代码片段带大家捋一遍SwAV的训练流程并分享几个我踩过的坑和调试经验。首先我们明确一下训练所需的组件一个编码器网络f通常是ResNet-50去掉最后的全连接分类层。输出一个D维的特征向量。一个投影头g一个小型MLP将特征映射到另一个空间用于计算对比损失。在SwAV中它输出的是用于计算code的向量z。原型矩阵C一个可学习的参数尺寸为[D, K]K是原型数量通常取3000。一个队列可选如果批次大小很小如256为了满足最优传输求解的条件特征数不能远小于原型数需要将过去几个批次的特征缓存起来形成一个队列Z与当前批次的特征拼接后再一起计算code。训练循环的核心步骤如下步骤一数据准备与特征提取。对于一个批次中的每张图片x我们应用Multi-Crop策略生成2个全局视图和V个局部视图。将所有视图假设总共B*(2V)个送入编码器f和投影头g得到归一化后的特征向量z。这里要注意对于小分辨率视图我们可以使用编码器中更早的层进行下采样或者直接调整网络开头的卷积步长来适配。# 伪代码示意 images batch # [B, C, H, W] multi_crops multi_crop_augment(images) # 返回一个列表例如 [global1, global2, local1, local2, local3, local4] features [] for crop in multi_crops: z projection_head(encoder(crop)) # 得到归一化后的特征 features.append(z) # features: 一个包含 (2V) 个张量的列表每个张量形状为 [B, D]步骤二计算Code聚类分配。我们只对两个全局视图的特征计算code。假设当前两个全局视图的特征是z1和z2形状都是[B, D]。如果使用了队列我们需要将队列中的特征z_queue形状[queue_size, D]与z1、z2在批次维度拼接得到Z concat([z1, z2, z_queue], dim0)形状为[M, D]其中 M 2*B queue_size。然后我们求解最优传输问题得到分配矩阵Q。Q的尺寸是[K, M]每一列对应一个特征向量的codeK维概率分布。我们只取出对应z1和z2的那部分code记为q1和q2。# 伪代码示意使用Sinkhorn-Knopp算法 def sinkhorn(scores, epsilon0.05, n_iters3): # scores: 原型C和特征Z的相似度矩阵[K, M] Q torch.exp(scores / epsilon).T # 转置为[M, K]并初始化 Q / Q.sum(dim1, keepdimTrue) # 归一化每行每个特征 for _ in range(n_iters): Q / Q.sum(dim0, keepdimTrue) # 归一化每列每个原型实现均匀分配约束 Q / Q.sum(dim1, keepdimTrue) # 再次归一化每行 return Q.T # 转置回[K, M] # 计算相似度 scores torch.matmul(C.T, Z.T) # [K, D] [D, M] - [K, M] # 计算code Q sinkhorn(scores) q1, q2 Q[:, :B], Q[:, B:2*B] # 取出对应z1和z2的code步骤三计算交换预测损失。对于每个全局视图的codeq1,q2我们用所有视图的特征z1,z2,z_local1,z_local2, ...去预测它。计算预测分布p与目标codeq之间的交叉熵损失。def swapped_prediction_loss(z, q_target, C, temperature0.1): # z: 某个视图的特征[B, D] # q_target: 目标code例如q2[K, B] # C: 原型矩阵[D, K] logits torch.matmul(z, C) / temperature # [B, K] log_probs F.log_softmax(logits, dim-1) # [B, K] # 将q_target从[K, B]转置为[B, K]作为目标概率 loss -torch.mean(torch.sum(q_target.T * log_probs, dim-1)) return loss total_loss 0 # 用z1预测q2 total_loss swapped_prediction_loss(features[0], q2, C) # 用z2预测q1 total_loss swapped_prediction_loss(features[1], q1, C) # 用所有局部视图的特征预测q1和q2 for v in range(2, 2V): total_loss swapped_prediction_loss(features[v], q1, C) total_loss swapped_prediction_loss(features[v], q2, C) total_loss / (2 2*V) # 平均一下步骤四反向传播与更新。执行反向传播更新编码器f、投影头g和原型矩阵C的参数。同时如果使用了队列需要用当前批次的两个全局视图特征更新队列通常是先进先出。我踩过的几个坑温度系数 τ这是一个非常敏感的超参数。论文中设置为0.1。如果设置太大softmax分布会太平滑学习信号弱如果设置太小分布会太尖锐训练不稳定。需要仔细调整。原型数量 KK的大小需要与数据集的复杂度匹配。ImageNet上常用3000。如果K太小原型不足以覆盖数据的多样性如果K太大计算开销增加且可能引入噪声。对于更小的数据集如CIFAR可以适当减小K。Sinkhorn迭代次数论文中只用了3次迭代。实际上这个算法收敛很快3次通常足够了。增加迭代次数收益很小但会增加计算时间。特征归一化在计算特征z和原型C时一定要进行L2归一化将它们约束在一个超球面上。这是对比学习相关方法稳定训练的关键。学习率与优化器对于原型矩阵C通常使用一个更大的学习率例如10倍于编码器学习率因为它需要更快地适应变化的特征分布。使用LARS优化器对于大Batch Size训练很有帮助即使SwAV对小Batch友好但在尝试大Batch时LARS也能提升稳定性。5. 效果如何SwAV在ImageNet与下游任务的表现说了这么多SwAV的实际效果到底怎么样毕竟在AI领域再优雅的理论也需要靠硬核的指标来说话。我们直接看它在ImageNet线性评估Linear Evaluation这个自监督学习的“标准考场”上的成绩。在ImageNet数据集上使用标准的ResNet-50作为编码器在无标签数据上预训练后冻结特征提取器只训练一个线性分类器SwAV达到了75.3%的top-1准确率。这个数字是什么概念我们对比一下同期的主流方法SimCLR需要4096的Batch Size达到了74.2%。MoCo v2使用负样本队列达到了71.1%。BYOL不需要负样本达到了74.3%。SwAV不仅超过了需要海量负样本的SimCLR也超过了同为“免对比”的BYOL在当时创造了新的记录。更重要的是SwAV在仅使用256的小Batch Size时就能达到72.7%的top-1准确率。而SimCLR在Batch Size为256时性能会暴跌至60%以下。这充分证明了SwAV在小批量设定下的巨大优势让更多资源有限的研究者能够参与进来。下游任务迁移能力自监督学习的终极目标是学到的特征能够很好地迁移到各种下游任务中证明其通用性。SwAV在这方面表现同样出色。在PASCAL VOC、COCO等目标检测和实例分割任务上用SwAV预训练的模型做初始化其性能全面超越了有监督的ImageNet预训练模型。例如在PASCAL VOC目标检测上SwAV比有监督预训练高出了超过4个点AP。这说明SwAV学到的特征比人工标注的类别标签所监督学到的特征更具有泛化性和语义信息。为什么SwAV迁移效果好我个人认为这得益于其“在线聚类”的机制。模型在学习过程中被迫将图像特征与一组可学习的原型进行匹配。这些原型在训练过程中逐渐演化可以理解为学习到了一些数据驱动的、跨类别的“视觉概念基元”。例如可能有一个原型对应“轮子纹理”另一个对应“动物眼睛”再一个对应“天空的蓝色”。这些基元比具体的“猫”“狗”类别标签更底层、更通用因此在遇到新的、未见过的任务时组合这些基元的能力就体现出了强大的迁移优势。Multi-Crop的贡献论文中也通过消融实验验证了Multi-Crop策略的威力。在SwAV基础上加入Multi-Crop2个全局4个局部视图能在ImageNet线性评估上带来约2-3个百分点的稳定提升。这个提升是免费的计算开销增加很小并且是通用的。我自己的实验也验证了这一点这个策略几乎成了我后续做自监督学习的标配。对小批量数据的友好性这是SwAV最吸引我的一个点。论文中详细展示了随着Batch Size从256增加到4096性能的提升曲线。对于SwAV即使Batch Size只有256性能也已经很高随着Batch Size增大性能增长平稳。而对于SimCLR性能对Batch Size极度依赖小批量时性能很差。这意味着如果你只有一两张消费级显卡例如RTX 3090每张卡可能只能放下128或256的Batch Size你几乎无法有效训练SimCLR但却可以顺利训练SwAV并获得相当不错的结果。这极大地降低了自监督学习的研究和应用门槛。6. 横向对比SwAV与MoCo、SimCLR、BYOL的异同要真正理解SwAV的价值最好把它放在当时自监督学习的“星系图”中和它的“邻居们”做个比较。我们选取三个最具代表性的方法MoCo动量对比、SimCLR端到端对比、BYOL免对比。1. SwAV vs. SimCLR计算效率之战SimCLR负样本依赖者。它在一个超大批次内进行所有样本的两两对比。优点是简单直接性能上限高。缺点是计算复杂度是O(B^2)对Batch SizeB极度敏感需要巨大的显存和算力TPU集群。SwAV原型对比者。它将负样本对比转化为与一组固定数量K的原型进行匹配。计算复杂度主要取决于K和特征维度D与批次大小B是线性关系。因此它对大批次的需求大大降低内存效率极高。核心区别SimCLR进行的是“样本级”的对比而SwAV进行的是“原型级”的对比。后者用一个可学习的“字典”替代了庞大的负样本库。2. SwAV vs. MoCo内存管理之争MoCo队列管理者。它通过一个动量编码器和一个先进先出的负样本队列来维护一个庞大且一致的负样本库。它解耦了批次大小和负样本数量使得在小批次下也能使用大量负样本。SwAV无需队列。它完全不需要维护负样本队列。所有的计算都集中在当前批次和原型矩阵上。内存占用更小实现更简洁。核心区别MoCo通过精巧的工程设计动量更新、队列来模拟大批次对比。SwAV则从算法层面改变了对比的对象从样本到原型从根本上避免了维护大负样本库的需求。3. SwAV vs. BYOL如何避免坍塌CollapseBYOL孪生网络预测者。它使用一个在线网络和一个动量更新的目标网络让在线网络去预测目标网络的输出。它完全抛弃了负样本通过不对称的网络结构如预测头、停止梯度来防止模型学到平凡解所有输出相同。SwAV均匀分配约束者。它也没有显式的负样本但它通过最优传输中的均匀分配约束强制模型将特征均匀地分配到所有原型上从而避免了所有特征坍缩到同一个原型的捷径解。核心区别两者都避免了显式负样本对比。BYOL依靠架构不对称性和动量更新来防止坍塌SwAV则依靠损失函数的设计交换预测和优化约束均匀分配来防止坍塌。BYOL的机制更隐晦而SwAV的机制有更清晰的数学解释最优传输。总结一下如果你追求极致的性能且有海量算力SimCLR及其后续变种如SimCLR v2可能仍是很好的选择。如果你希望在有限内存下获得不错的性能且喜欢简洁的实现SwAV是绝佳的选择。它的代码比MoCo更简单没有动量编码器和队列管理的逻辑。如果你对“负样本”这个概念心存疑虑想探索更纯粹的“自预测”范式那么BYOL和SwAV代表了两种不同的哲学。BYOL更像“自我蒸馏”SwAV更像“自我组织”。从我个人的项目经验来看SwAV在工业界的落地上有独特优势。它的训练稳定对超参数除了温度系数相对不敏感内存占用可控并且学到的特征具有很强的可解释性可以通过观察原型矩阵来理解模型学到了什么。在数据标注成本高昂的领域如医疗影像、遥感图像SwAV这类高效的无监督预训练方法为我们提供了一个强大的特征提取器初始化方案能有效提升下游小样本任务的性能。