SupCon有监督对比学习:原理推导、实现细节与调参经验 对比学习这两年热度一直没降下来从最早的SimCLR、MoCo那一批自监督方案到后来有人往里面加标签信息SupConSupervised Contrastive Learning算是把有监督和对比这两件事捏到一起的一个代表性工作。我第一次读这篇的时候其实有点懵——明明已经有标签了为什么还要费劲做对比学习直接交叉熵训分类头不就完了后来自己动手复现、又在几个小数据集上折腾了一圈才慢慢理解它到底在解决什么问题、为什么在某些场景下确实比纯交叉熵更稳。这篇就把我对SupCon的理解从头捋一遍包括它背后的动机、损失函数怎么推导出来的、实现时那些论文里不会写但实际会踩的坑以及我自己的调参经验。适合已经了解基础对比学习、想搞清楚SupCon细节的读者也适合想把它用到自己分类任务里的人参考。1. 先搞清楚SupCon到底想解决什么1.1 从交叉熵的够用但不够好说起大部分人做分类任务的第一反应就是交叉熵损失加一个softmax分类头这套组合用了十几年确实够用。但如果你仔细观察过训练出来的特征空间会发现一个问题交叉熵只关心每个样本被分到正确类别的概率够不够高它并不直接约束同类样本在特征空间里要聚得多紧、不同类之间要隔得多开。换句话说交叉熵优化的是决策边界而不是特征分布本身。这就导致一个现象模型在训练集上分类很准但特征空间里同类样本可能散得到处都是类与类之间的间隔也可能很窄。遇到分布稍微偏移一点的测试数据或者样本量少的类别泛化就容易出问题。我自己在一个小规模图像分类任务上就遇到过交叉熵训出来的模型在验证集上准确率还行但把特征降维可视化一看几个类的簇互相咬在一起边界模糊得很。SupCon的思路就是既然标签已经告诉我们哪些样本是同一类那为什么不直接利用这个信息去约束特征空间让同类样本在嵌入空间里靠得更近、异类推得更远这其实就是把对比学习的拉近正样本、推远负样本的机制从自监督里的同一样本的不同增强视图扩展到了同一类别的不同样本。1.2 对比学习给分类任务带来的额外约束自监督对比学习比如SimCLR里正样本是同一张图的两个增强视图负样本是batch里其他所有图。它不需要标签靠的是同一张图不管怎么变换语义不变这个假设。SupCon把这个假设换成了同一个类别的不同样本语义相同正样本从同一图的两个视图变成了同一类的任意两个样本。这个改动看起来小但意义挺大。自监督对比学习的正样本对数量受限于增强方式一个batch里每张图只能贡献一对正样本它自己和它的另一个视图。而SupCon里如果一个batch有N个样本、C个类别每个类有若干样本那么正样本对的数量是每个类内样本两两组合数量级一下子就上去了。标签信息让正样本的构造变得极其丰富这是SupCon相比自监督对比学习在分类任务上的天然优势。另外一个关键点是SupCon学到的特征表示对下游任务更友好。因为它在训练时就把类内紧凑、类间分离这个结构编码进了嵌入空间所以即使你把这个backbone拿去做few-shot、做检索、做聚类效果往往也比纯交叉熵训出来的特征好。我在一个检索式的小任务上对比过SupCon预训练的backbone提取的特征做最近邻检索召回率比交叉熵版本高了一截这个差距在样本少的类别上尤其明显。1.3 它和 triplet loss、center loss 这些度量学习方法的区别说到拉近同类、推远异类很多人会想到triplet loss或者center loss这些也是度量学习里常用的手段。SupCon和它们的区别主要在两点一是正负样本的利用效率二是损失的构造方式。triplet loss每次只挑一个正样本对和一个负样本对来算对样本挖掘hard negative mining非常敏感挖得不好训练就不稳定。center loss则是给每个类维护一个中心让样本往中心靠但它只约束了类内距离没有显式地推远类间距离。SupCon则是在一个batch内对每个锚点样本把它和所有同类样本都作为正样本、所有异类样本都作为负样本一次性全部纳入损失计算。这种全对全的构造方式让每个样本的梯度信号更充分也不那么依赖精细的样本挖掘策略。当然代价是计算量。一个batch里如果有N个样本两两之间都要算相似度复杂度是O(N²)。这也是为什么SupCon对batch size比较敏感batch太小正负样本都不够batch太大显存吃不消。这个权衡后面会细说。2. 损失函数是怎么一步步推出来的2.1 从InfoNCE到SupCon的形式要理解SupCon的损失最好先从它借鉴的InfoNCE损失看起。InfoNCE是自监督对比学习里常用的损失对于一个锚点样本i它的正样本是i的另一个视图记作j负样本是batch里其他所有样本。损失的形式是$$ \mathcal{L}{i} -\log \frac{\exp(\text{sim}(z_i, z_j)/\tau)}{\sum{k \neq i} \exp(\text{sim}(z_i, z_k)/\tau)} $$其中sim是余弦相似度τ是温度系数。这个式子的直觉是分子是锚点和正样本的相似度经过温度缩放后的指数分母是锚点和所有其他样本相似度的指数和。最小化这个损失就是让分子相对分母尽量大也就是让正样本的相似度相对负样本更高。SupCon把这个形式扩展了一下。现在锚点i的正样本不再是一个而是它所属类别里的所有其他样本记作集合P(i)。负样本则是batch里所有不属于i这个类别的样本记作集合A(i)\P(i)。损失变成$$ \mathcal{L}{i} -\frac{1}{|P(i)|} \sum{p \in P(i)} \log \frac{\exp(\text{sim}(z_i, z_p)/\tau)}{\sum_{a \in A(i)} \exp(\text{sim}(z_i, z_a)/\tau)} $$注意这里有两个变化一是对每个正样本都算一次log项然后取平均二是分母里包含了所有样本除了i自己包括正样本和负样本。这个形式和InfoNCE的结构是一致的只是正样本从单个变成了多个。2.2 为什么分母要包含正样本第一次看这个公式的时候我有个疑问分母里为什么要把正样本也算进去按直觉分母应该只放负样本才对这样正样本的相似度在分子里越大、负样本在分母里越小损失就越小逻辑更直接。后来想明白了这其实是InfoNCE的标准形式。分母放所有样本除锚点自己是为了让这个损失成为一个在全部候选里挑出正样本的分类问题。你可以把分子看作正确类别的得分分母看作所有类别的得分之和这本质上是一个softmax。如果分母只放负样本那这个损失就不是一个规整的概率分布形式了梯度性质会变差。从梯度角度看分母包含正样本意味着当正样本的相似度已经很高时它在分母里也会贡献较大的值从而抑制损失继续无脑下降起到一种自适应的调节作用。这个细节在实现时如果搞错了比如分母漏掉正样本训练也能跑但效果会打折扣我一开始就犯过这个错后来对着公式逐项核对才发现。2.3 温度系数τ到底在调什么温度系数τ是SupCon里最需要调的超参之一论文里一般取0.07到0.1之间。它的作用是对相似度做缩放τ越小softmax的分布越尖锐损失对难负样本相似度较高的负样本越敏感τ越大分布越平滑模型对难负样本的关注度下降。用一个生活化的类比假设你在一个班里选班长τ小相当于只看得票最高的几个人竞争很激烈稍微落后一点就被放大τ大相当于看整体票数分布比较温和。τ太小容易训练不稳定因为梯度会被少数几个难负样本主导τ太大则对比学习的效果会弱化趋近于普通的softmax分类。我自己的经验是τ从0.1开始试比较稳如果发现训练loss震荡厉害就往上调一点到0.15如果发现特征区分度不够类间间隔小就往下调到0.07左右。这个值跟batch size、数据集的类别数都有关系没有一个放之四海皆准的最优值得结合验证集表现来定。2.4 一个容易忽略的细节归一化与相似度计算SupCon在算相似度之前会对特征做L2归一化然后用点积等价于余弦相似度。这一步很关键如果不做归一化特征的模长会进入相似度计算模型可能会通过放大某些样本的特征模长来作弊让相似度虚高但实际的语义区分度并没有提升。实现上就是在投影头输出之后加一个F.normalize(z, dim1)然后相似度矩阵用z z.T算。这个矩阵是对称的对角线是自身相似度恒为1计算损失时要记得把对角线mask掉否则锚点会把自己当成正样本虽然影响不大但逻辑上不对。另外如果batch里有某个类别只有一个样本那这个样本没有正样本对它的损失项应该跳过不然除以|P(i)|0会出问题。这些边界情况论文里不会写但实际写代码时都得处理。3. 动手实现从零搭一个能跑的SupCon3.1 整体结构backbone 投影头 分类头SupCon的训练结构一般是这样一个backbone比如ResNet提取特征接一个投影头projection head通常是两层MLP把特征映射到对比学习用的嵌入空间损失在这个嵌入空间上算。同时如果你还要做分类可以在backbone特征上再接一个分类头用交叉熵一起训或者干脆只用SupCon预训练、之后冻结backbone再训分类头。论文里的做法是SupCon损失和交叉熵损失联合训练两个损失加权求和。我实测下来联合训练在小数据集上收敛更快但如果你主要目的是拿一个好特征做下游任务也可以只用SupCon损失训然后单独评估特征质量。投影头的维度一般设成128或256太小表达能力不够太大容易过拟合128是个比较安全的起点。这里有个细节投影头只在训练时用推理或者做下游任务时用的是backbone输出的特征不是投影头的输出。这一点和SimCLR一致投影头的存在是为了让对比损失在一个更适合做对比的空间里优化而不影响backbone特征的通用性。3.2 损失函数的核心代码实现下面是我自己写的一个SupCon损失实现基于PyTorch逻辑上对着公式逐项来的import torch import torch.nn as nn import torch.nn.functional as F class SupConLoss(nn.Module): def __init__(self, temperature0.1): super().__init__() self.tau temperature def forward(self, features, labels): # features: [batch_size, dim], 已经过投影头 # labels: [batch_size] device features.device batch_size features.shape[0] # L2归一化 features F.normalize(features, dim1) # 相似度矩阵 [B, B] sim_matrix torch.matmul(features, features.T) / self.tau # 构造mask同类为1异类为0对角线为0 labels labels.contiguous().view(-1, 1) mask torch.eq(labels, labels.T).float().to(device) mask mask - torch.eye(batch_size, devicedevice) # 去掉对角线 # 为了数值稳定减去每行最大值 sim_matrix sim_matrix - sim_matrix.max(dim1, keepdimTrue)[0].detach() # 分母所有样本除自己的exp和 exp_sim torch.exp(sim_matrix) denom exp_sim.sum(dim1, keepdimTrue) - exp_sim.diag().view(-1, 1) # 分子同类样本的exp和 numer (exp_sim * mask).sum(dim1, keepdimTrue) # 每个正样本的log项取平均 pos_count mask.sum(dim1, keepdimTrue).clamp(min1) loss -(torch.log(numer / denom 1e-8) / pos_count).sum() / batch_size return loss这段代码有几个地方值得说明。第一sim_matrix减去每行最大值是为了数值稳定防止exp溢出这是softmax类损失的常规操作但要注意用.detach()不然会影响梯度。第二分母里减去了对角线上的exp_sim.diag()因为锚点自己不应该算进分母。第三pos_count用了clamp(min1)防止某个类别只有一个样本时除零。第四最后除以batch_size是对batch内所有锚点取平均。这个实现是简化版没有处理分布式训练里跨卡的正负样本如果batch分散在多张卡上需要额外的all_gather操作把各卡的特征聚起来再算损失否则每张卡只能看到自己那部分样本正负样本数量会打折扣。3.3 数据增强在SupCon里的角色自监督对比学习里数据增强是核心因为正样本就是靠增强构造的。SupCon里正样本来自同类不同样本增强的作用相对弱化了一些但依然重要。增强的作用变成了增加同类样本的多样性让模型见到的类内变化更丰富学到的特征更鲁棒。我一般用比较标准的增强组合随机裁剪加缩放、随机水平翻转、颜色抖动亮度、对比度、饱和度、随机灰度化。强度不用像SimCLR那么激进因为SupCon本身正样本就多不需要靠强增强来制造难度。如果增强太强反而可能把同类样本的语义破坏掉让正样本对变得不可靠。有个经验如果你的数据集类别之间本身就很像比如细粒度分类增强强度要适当降低尤其是颜色相关的增强因为颜色可能是区分细粒度类别的关键线索。我在一个细粒度任务上就吃过亏颜色抖动开太大模型把颜色信息丢掉了准确率反而下降。3.4 batch size 与显存的权衡SupCon对batch size比较敏感因为batch越大每个类别的样本越多正样本对越丰富负样本也越多对比学习的效果越好。论文里用的batch size都比较大几百到上千但实际中很多人显存有限跑不了那么大的batch。我的建议是如果显存允许尽量把batch size往大了堆至少保证每个类别在batch里有若干个样本比如每类至少4到8个。如果实在堆不上去有两个补救办法一是用梯度累积模拟大batch但要注意梯度累积对对比学习的效果不如真实大batch因为损失是在每个小batch内算的正负样本数量并没有真正增加二是用memory bank或者队列的方式缓存历史特征作为额外的负样本这个思路来自MoCo能缓解batch小的问题但实现复杂度会上去。还有一个折中方案是混合精度训练用fp16能把显存占用降下来不少batch size能翻倍。不过混合精度下要注意损失计算的数值稳定性温度缩放后的相似度矩阵容易出问题建议在算exp之前做好减最大值的操作并且用fp32算损失。4. 训练中那些论文不会告诉你的坑4.1 损失不下降先检查标签对齐我第一次跑SupCon的时候遇到一个诡异的现象loss在前几个epoch降得挺正常然后就卡住不动了准确率也上不去。排查了半天最后发现是数据加载的时候标签和图像没对齐——用了两个不同的DataLoader分别加载图像和标签shuffle的随机种子不一致导致标签错位。这种bug很隐蔽因为loss不会报错只是学不到东西。所以如果你发现SupCon训练异常第一件事是检查标签和特征是不是严格对应的。一个简单的验证方法是故意把batch里所有标签设成一样看loss是不是趋近于0因为全是正样本没有负样本理论上损失应该很小再把标签全设成不同看loss是不是很大。如果这两个极端情况的表现符合预期说明损失实现基本没问题问题可能出在数据管道上。4.2 温度系数和batch size的联动效应前面说了τ要调但τ的最优值其实和batch size有关系。batch size大的时候负样本多分母大τ可以稍微小一点让分布更尖锐batch size小的时候负样本少τ太小会让损失对少数几个负样本过度敏感训练不稳定。我做过一组粗略的对比batch size 256的时候τ0.1比较稳batch size降到64的时候同样的τ就开始震荡调到0.15才稳定下来。所以调参的时候不要把τ和batch size分开看要一起调。一个实用的做法是固定batch size在[0.05, 0.2]区间里扫一遍τ看验证集表现选最优。4.3 投影头维度和特征质量的取舍投影头的输出维度也是个需要试的参数。维度太低比如32嵌入空间的表达能力不够对比损失优化起来受限维度太高比如512以上在小数据集上容易过拟合而且计算相似度矩阵的开销也大。我一般从128开始如果数据集比较大几万张以上可以试256。另外投影头的结构论文里用的是两层MLP加ReLU隐藏层维度和输出维度一致。我试过加BatchNorm效果有轻微提升但不是必须的。如果你的backbone本身比较深投影头可以简单一点一层线性也行因为backbone已经提供了足够的表达能力。4.4 联合训练时两个损失的权重怎么定如果SupCon损失和交叉熵损失联合训练两者的权重需要调。论文里给的是一个固定权重但实际中我发现这个权重跟数据集难度有关。交叉熵权重太大模型会偏向直接分类对比学习的作用被削弱SupCon权重太大分类头的训练信号不足收敛慢。我的经验是让SupCon损失占主导交叉熵作为辅助权重比大概在1:0.5到1:1之间。可以先只用SupCon训一段时间等特征空间初步成型了再加入交叉熵联合训。这种分阶段训练的方式在一些任务上比从头联合训更稳。5. 什么场景下SupCon真的比交叉熵强5.1 样本不均衡和小样本类别SupCon在样本不均衡的场景下表现往往比交叉熵好。原因是交叉熵对每个样本一视同仁少数类样本少梯度贡献就少模型容易偏向多数类。而SupCon里少数类的样本虽然少但它们在对比损失里作为正样本被反复利用每个少数类样本都会和其他少数类样本组成正样本对相当于变相增加了少数类的训练信号。我在一个长尾分布的数据集上对比过交叉熵版本在尾部类别上的准确率明显偏低换成SupCon之后尾部类别的准确率有提升整体准确率也更均衡。当然这不是万能药如果某个类别只有一两个样本正样本对根本构造不出来SupCon也帮不上忙。5.2 需要好特征做下游任务的场景如果你的目标不只是分类而是想拿一个预训练好的backbone去做检索、聚类、few-shot那SupCon的优势更明显。因为它显式优化了特征空间的类内紧凑和类间分离提取出来的特征本身就有很好的判别性不需要下游任务再费劲去学。我做过一个简单的实验用同样的backbone一个用交叉熵预训练一个用SupCon预训练然后在few-shot设置下每类5个样本做分类。SupCon版本的平均准确率比交叉熵版本高了几个百分点而且方差更小说明特征更稳定。这个差距在类别数多、每类样本少的时候更突出。5.3 什么时候不建议用SupConSupCon不是银弹有些场景下用它反而麻烦。一是类别数特别多、每类样本特别少的场景比如每类只有一两个样本正样本对构造不出来SupCon退化成近似交叉熵还多了计算开销。二是训练资源极其有限的场景SupCon对batch size和显存的要求比交叉熵高如果只能跑很小的batch效果可能还不如老老实实用交叉熵。三是任务本身对特征空间结构要求不高的场景比如你只是想要一个能分类的模型不关心特征质量那交叉熵足够没必要上SupCon增加复杂度。工具要匹配需求不要为了用而用。6. 几个我踩过的具体坑和解决方式6.1 分布式训练时正负样本被割裂在多卡训练时如果每张卡各自算SupCon损失那每张卡只能看到自己那部分样本正负样本数量被除以了卡数对比学习效果大打折扣。解决办法是用all_gather把各卡的特征聚到每张卡上再算损失这样每张卡都能看到全局的batch。实现上要注意梯度的问题all_gather本身不传梯度需要配合一些技巧比如用gather后的特征算损失梯度通过原来的特征回传。这块实现起来有点绕如果卡不多、batch够大也可以考虑单卡训练避免这个复杂度。6.2 特征归一化漏做导致loss虚低有一次我忘了在算相似度前做L2归一化结果loss降得特别快看起来训练很顺利但验证集准确率上不去。原因是没归一化时模型可以通过放大特征模长让同类样本的相似度虚高损失是降了但特征的实际区分度没提升。加上归一化之后loss曲线变得没那么漂亮但验证集表现反而更好。这个坑很典型归一化这一步千万别漏。6.3 类别数少时对比学习信号不足在一个只有几个类别的任务上我发现SupCon的效果和交叉熵差不多甚至略差。分析下来是因为类别少负样本的种类有限对比学习推远异类的信号不够丰富。这种情况下SupCon的优势发挥不出来反而多花了计算。所以类别数太少比如少于10类的时候要谨慎评估SupCon是否值得。6.4 学习率和warmup的设置SupCon训练对学习率比较敏感尤其是联合训练的时候。我一般用比纯交叉熵稍小的学习率配合warmup。warmup的作用是让模型在训练初期不要被对比损失的大梯度冲击得太厉害等特征空间初步稳定了再放开学习率。warmup的步数大概占总步数的5%到10%具体看数据集大小。另外优化器我用AdamW比较多weight decay设0.0001到0.001之间。SGD也能用但收敛慢一些需要更仔细地调学习率。7. 一些延伸思考SupCon本质上是在说一件事标签信息不应该只用来算分类损失还可以用来塑造特征空间。这个思路可以往很多方向延伸。比如有人把SupCon用到半监督场景用少量标签构造正样本对大量无标签数据做自监督对比两者结合。也有人把SupCon的思想用到多模态任务里用不同模态的同类样本做正样本对。我自己比较感兴趣的是SupCon和聚类的关系。SupCon学到的特征空间天然适合做聚类因为类内紧凑、类间分离的结构已经编码进去了。我在一个无监督聚类任务上试过用SupCon预训练的特征做KMeans聚类纯度和用交叉熵特征相比有明显提升。这个方向后续还可以继续挖比如用聚类结果反过来指导正样本的构造形成一个迭代的循环。最后分享一个小技巧如果你手头的数据集标签有噪声SupCon对噪声标签的鲁棒性其实比交叉熵好一些。因为对比损失是batch内多个样本共同作用的单个错误标签只会影响它参与的那几个正样本对不会像交叉熵那样直接给一个错误的强梯度。当然噪声太大还是得先清洗数据SupCon只能缓解不能根治。