【IJCV】基于概率表征的半监督对比学习框架,FaceChain团队出品

文章链接：https://arxiv.org/abs/2402.18117

代码链接：https://github.com/Haoyu-Xie/PRC2

一、引言

半监督语义分割算法(Semi-Supervised Semantic Segmentation, S4)能够使用较少的有标注数据对模型进行训练，并且获得一个性能良好的图像分割模型，因此得到了研究者们较为广泛的关注。最近的S4研究中，研究者们引入了像素级别的对比学习(Pixel-wise contrastive learning)，进一步提升了模型的能力。然而，最近的许多研究都存在着以下问题：模型在无标注数据上训练时鲁棒性较差，容易受到错误的指导影响。为了提升模型的鲁棒性，我们提出了基于概率表征的对比学习框架(Probabilistic Representation Contrastive Learning framework, PRCL)。这个框架将像素级别的表征用高斯分布进行建模，并且根据表征语义的可靠性来微调它们在对比学习过程中的贡献。这样，模型就获得了容忍错误语义的表征的能力。随后，在概率表征(Probabilistic Representation, PR)的基础上，我们建立了全局分布原型(Global Distribution Prototype, GDP)和虚拟负样本(Virtual Negatives, VNs)来加入到对比学习的过程中，以解决传统像素级别的对比学习由小批量而造成的问题。实验表明，我们提出的PRCL框架在半监督设定下模型的分割能力达到了state-of-the-art的性能。此外，大量的消融实验也表明我们提出的模块的有效性。

二、基于概率表征的对比学习框架

1.）概率表征

近期的基于对比学习的S4算法框架大多都在原来的S4算法框架下(如老师-学生框架)，将像素级别的对比学习视作一个辅助任务，让模型能够在潜在域上获得更规范的输出。具体的做法是将每一个表征分配上对应的语义信息，之后通过优化对比损失的方式将相同类的表征聚集在一起，将不同类的表征分散开来。由于每一个表征需要分配类别信息，在半监督条件设定下，模型在无标注数据上训练时获得的语义信息可能存在错误，这种错误会影响对比学习的效果，最终导致模型的准确度下降。为了解决这个问题，我们在对比学习过程中使用了概率表征。概率表征与传统确定性表征不同的是，它不仅能够反映表征的位置，同时也能够反映表征出现在这个位置上的可能性，即表征分配到的语义信息的可靠性。在实际训练中，我们使用高斯分布来拟合概率表征，高斯分布中的均值μ，即代表概率表征的具体大小以及方向，高斯分布中的方差σ，即代表概率表征的可靠性。均值和方差均由两个不同的MLP预测而得。有了概率表征之后，我们在衡量相似度与对比学习的过程中不仅考虑两个表征之间的距离，还考虑两个表征的可靠性，即两个表征如果距离很近但是可靠性很低，我们也不会认为这两个表征十分相似。由此，通过概率表征中方差的影响，模型在训练中能够容忍一定程度的噪声，极大地提高了鲁棒性。

图1 概率表征（右）与传统确定性表征（左）的差异

2.）全局分布原型

在传统的基于对比学习的S4算法中中，由于批量大小(batch size)的限制，原型总是只集合了较少数量的同类表征的信息。我们认为这种做法会导致原型在训练的过程中产生比较大的漂移，从而无法给表征的聚合提供一个稳定的方向。这是由于每个迭代中的表征有可能存在着错误的语义信息，或者比较大的类间差距。在概率表征的基础之上，我们提出了一种新的原型(Prototype)构建的方法。与传统原型构建方法不同的是，全局分布原型的构建是将训练过程中出现过的所有表征都考虑进原型的构建，即构建出的原型包含了训练中这一类的所有表征的信息。通过这种构建方式，全局分布原型能够让训练过程中的不同的迭代之间进行信息交互，容忍迭代中的瞬时噪声与表征的类间差异，从而让对比学习过程中的原型更加稳定，为对比学习过程提供一个稳定的指导方向。

3.）虚拟负样本

另外，在传统基于对比学习的S4算法中，负样本的选取有两种方式：在当前迭代中的表征中选取或建立memory bank来存放并取用负样本。在这里，我们在全局分布原型的基础上，提出了虚拟负样本，用于高效的补充当前迭代中表征类别数可能少于数据集的类别数的问题。我们使用采样的方式从全局分布原型中获取虚拟表征，并且将其同真实表征一起用于构建对比学习中的负分布。与只使用当前迭代中的真实表征的方法相比，我们的方法构建出的负分布能够更加全面，更加接近真实负分布。与建立memory bank的策略相比，由于我们的方法只需要存储全局分布原型，不需要存储很大数量的真实表征，因此消耗更少的存储空间，并且能够加快训练速度，获得更好的效果。

图2 两个由有限的batch size带来的问题(L1, L2)以及对应的解决办法(S1, S2)

此外，PRCL整体框架如下图所示：

图3 PRCL的整体框架

三、实验结果

1.）算法性能的对比

我们将PRCL与最近的SOTA S4算法在Pascal VOC 2012数据集和Cityscapes数据集上的性能进行定量比较。所有算法所采用的骨干网络均为ResNet-101。从结果中可以看到，相比于其他SOTA算法，无论在哪个数据集上，我们提出的PRCL都有极强的竞争力，尤其是在标签率较少的时候。这是由于在标签率较低时，模型预训练较少，产生的伪标签质量较差。我们提出的算法能够一定程度上增强模型在低质量伪标签下的表现，因此在低标签率下的表现更为优异。

2.）概率标准的有效性

接下来我们展示了概率表征中方差的热度图与产生的伪标签情况的对比。热度图中

图4 概率表征方差热度图

颜色越红的地方代表这里的表征的语义信息越有可能不正确。从图中可以直观的看出概率表征中的方差有能力找到错误的伪标签与容易出现模棱两可的边缘部分，从而，这些部分的表征在对比学习中的贡献将会减少。下表展现了使用概率表征之后，模型的准确度将在很大的标签率范围内超过基线模型。

3.）全局分布原型的有效性

下面展示的是有无更新策略的原型在训练中的漂移情况。从图中可以看出，如果不

图5 不同更新方式的原型的漂移情况

使用任何原型更新策略，即原型的计算总是通过集合当前迭代中所有的同类表征的信息这一方法，那么在相邻迭代中，原型的位置会发生比较大的漂移，导致在对比学习的过程中不能够给表征聚合提供一个稳定的方向。相反，有了更新策略之后，原型的漂移减小，表征聚合的方向更加稳定。另外，由于概率表征能够在建立原型的时候对错误的语义信息有容忍性，因此原型在更新的过程中漂移更小。下表展现了原型更新策略相比于基线模型的有效性。

4.）虚拟负样本的有效性

下图简单展示了虚拟负样本的生成过程。虚拟负样本是通过全局分布原型采样获得

图6 虚拟负样本的示意图

的。它围绕在全局分布原型周围，相较于真实的表征来说，虚拟负样本能够包含该类表征的全局信息。最重要的是，使用生成虚拟负样本这一策略在对比学习构建负分布时进行补充，能够极大的减少需要的显存，因为我们只需要存储与类别数相同数量的全局分布原型，而不是存储具体的表征。下表展现了虚拟负样本在对比学习中的表现，并且定量地反映了节省的存储空间。

四、结论

        在这一工作中，我们从现有的基于对比学习的S4算法鲁棒性欠佳这一问题出发，提出了概率表征来增强对比学习的鲁棒性，并且建立了PRCL框架。另外，我们在概率表征的基础上，提出了全局分布原型与虚拟负样本，两者都能够解决模型在单个迭代的训练中面临的问题，从而提升模型的最终表现。最后，我们希望我们的工作能够对之后的S4工作有所启发。