本文分享NeurIPS 2021 论文『Revitalizing CNN Attentions via Transformers in Self-Supervised Visual Representation Learning』,由港大&腾讯AI Lab&牛津大学联合提出 CARE,让 CNN 和 Transformer 能在对比学习中“互帮互助”!
详细信息如下:
导言:
自监督视觉表征学习(SSL)的研究提升了编码器主干网络的性能,使其能够区分无标签的训练样本。虽然通过SSL实现的CNN编码器的识别性能与通过监督学习实现的编码器相当,但其网络中的注意力机制还未得到探索,有待进一步改进。
受在识别场景中带视觉注意力机制的Transformer的启发,作者提出了一个CNN注意力恢复(CNN Attention REvitalization,CARE) 框架来训练SSL中由Transformer引导注意力的CNN编码器。本文提出的CARE框架由CNN流(C流,C-stream)和Transformer流(T流,T-stream)组成,其中每个流包含两个分支。C-stream遵循现有的SSL框架,带有两个CNN编码器、两个投影模块和一个预测模块。
T-stream包含两个Transformer、两个投影模块和一个预测器。T-stream连接到CNN编码器,并与剩余的C-stream并行。在训练期间,作者在两个流中同时执行SSL,并使用T-stream输出来监控C-stream。CNN编码器的特征在T-stream中进行调整,以增强视觉注意力。然后使用这些调整的特征来监督C-stream,以便学习attentive的CNN编码器。
在多个标准视觉识别基准(包括图像分类、目标检测和语义分割)上的实验表明,本文提出的CARE框架将CNN编码器主干提高到SOTA的性能。
01
Motivation
有效地学习视觉特征对识别性能有着深远的影响。在处理大规模自然图像时,自监督视觉表征学习中的代理特征训练(pretext feature training),有利于下游识别任务。现有的SSL方法通常利用两个分支来度量来自同一输入图像的不同视图表示之间的相似性。通过最大化一幅图像中相关视图之间的相似性(例如,BYOL、SimSiam和Barlow Twins),或最小化不同图像中视图之间的相似性(例如,MoCo和SimCLR),这些方法被证明是学习自监督视觉表示的有效方法。
Transformer在自然语言处理中被提出,自ViT以来,它在处理大规模视觉数据方面也显示了其优势。Transformer中的编码器-解码器架构探索了无卷积的全局注意力。使用Transformer编码器的成功表明,视觉注意力有利于SSL中的编码器主干网络。另一方面,作者观察到现有的SSL方法并没有将视觉注意纳入CNN编码器中。这激发了作者对SSL中注意力机制的探索。作者希望CNN编码器能够具有Transformer的类似视觉注意力,以降低计算复杂度和内存消耗,从而提高识别性能。
在本文中,作者提出了一个CNN注意力恢复框架(CARE),通过Transformer引导使CNN编码器具备注意力机制。上图(d)显示了CARE的图示,并将其与其他SOTA的SSL框架进行了比较。CARE中有两个流(即C-Stream和T-Stream),其中每个流包含两个分支。C-stream与现有的SSL框架类似,它有两个CNN编码器、两个投影模块和一个预测模块。T-stream由两个Transformer、两个投影模块和一个预测模块组成。T-stream将CNN编码器特征作为输入,并通过Transformer提高特征的注意力。
在训练过程中,作者同时在两个流中执行SSL,并使用T-Stream输出来监督C-Stream。T-stream中的自监督学习确保了Transformer能够产生适用于SSL场景的注意力特征。同时,作者在C-stream上使用注意力来监督,这种监督使C-stream和T-stream都能产生类似的特性。CNN编码器的特征表示通过来自Transformer的视觉注意力得到了改进。
因此,预训练好的CNN编码器能够产生具有注意力的特征,这有利于下游识别场景。在图像分类、目标检测和语义分割基准上的实验表明,本文提出的CARE框架提高了CNN编码器主干网络的性能。
02
方法
本文的框架由C-stream和T-stream组成,如上图所示。
本文的C-stream类似于现有的SSL框架,有两个CNN编码器、两个投影模块和一个预测模块。两个编码器的结构相同,两个投影模块的结构相同。给定一个训练图像X,首先使用一组随机增强来处理它,以创建两个不同的增强视图。然后将这两个视图输入到C-stream,分别获得相应的输出和。然后计算用损失函数来惩罚两个输出的不相似度,该损失项是归一化特征向量的均方误差,如下:
其中为正则化,为点积运算。由于C-stream的输入来自一幅图像,因此在训练过程中C-stream的输出应该是相似的。
T-stream将CNN编码器的输出特征映射作为其输入,与C-stream并行设置。它由两个Transformer、两个投影模块和一个预测模块组成。投影模块和预测模块的结构与C-stream中的结构相同。
两个Transformer共享相同的架构,由连续的注意块组成,其中每个注意块包含两个多层感知机(MLP)层和一个位于中间的多头自注意(MHSA)层,其结构如上图所示。MHSA由多个平行处理输入特征的注意头组成。
在每一个注意头中,如上图右侧所示,输入特征映射分别通过3个不同的MLP层映射到query特征(q)、key特征(k)和 value特征(k)。然后,q和k相乘形成基于内容的注意,而q和位置编码p相乘形成基于位置的注意,如下所示:
其中是q和k的维度。在T-stream中,两个Transformer的输出是2D特征映射,然后将其平均池化并送到投影模块和预测模块。T-Stream的输出可以表示为和,损失函数如下所示:
除了在两个流中引入SSL损失项外,作者还使用T-Stream输出来监督C-Stream。这一项可以写成:
在训练过程中,要求C-stream输出与T-stream输出尽可能相似。相比于知识蒸馏(KD),这个公式主要有三点区别:
网络训练的最终目标函数可以表示为:
其中,λ是控制注意力损失影响的超参数。在计算之后,对C-stream和T-stream的上面分支执行反向传播。如上图中,CNN encoder1, projector1, predictor1通过C-stream中计算的梯度进行更新。同时,Transformer1、projector2和 predictor2通过T-stream中计算的梯度进行更新。然后剩下的部分,基于EMA进行更新。
本文的CARE框架,通过Transformer引导提高了CNN编码器的注意力。在上图中,作者展示了Attention Map的可视化结果。输入图像如(a)所示,而使用C-stream和CARE训练的编码器的Attention Map分别显示在(b)和(C)中。
总的来说,使用CARE训练的编码器的注意力比使用C-stream训练的编码器更集中,这表明T-stream为CNN编码器学习关注目标区域提供了有效的监督。T-stream帮助CNN编码器自适应地选择关注局部区域或全局区域。
例如,当分类需要全局信息时,CARE学习的CNN编码器将更加关注整个对象,如(c)中最后一列所示。另一方面,当局部信息足以进行分类时,CNN编码器将更加关注特定区域(例如,第一列和第二列(c)中的动物头部)。
03
实验
上表展示了ImageNet上自监督学习的实验结果,(a)表展示了以ResNet-50为编码器的实验结果,可以看出,本文方法比其他的自监督方法效果都要更好;(b)表展示了不同Encoder的实验结果,可以看出,本文的方法同样有性能优势。
上表展示了ImageNet上半监督的实验结果,可以看出本文方法的性能优越性。
上表展示了本文方法转移到目标检测和实例分割上的实验结果。
上表展示了用了FPN时,本文方法在目标检测和实例分割上的实验结果。
可以看出,在目标检测和实例分割任务上,本文方法同样比其他自监督方法效果更好。
上表展示了不同λ的实验结果,可以看出λ=100时,实验结果是最好的。
上表展示了不同注意块的实验结果,可以看出注意块为5时,实验结果是最好的。
从上表可以看出,加入可学习的相对位置编码时,实验结果是最好的。
04
总结
Transformer通过注意力探索结构实现了高级的视觉识别。在自监督视觉表征学习中,出现了利用Transformer主干网络提高识别性能的研究,这表明Transformer的视觉注意有益于SSL方法。受此启发,作者研究了如何有效地探索视觉注意力,使SSL中的CNN编码器能够受益。
因此,作者提出了CNN流和Transformer流并行的CARE框架。在SSL过程中,视觉注意力通过Transformer来监督CNN流。虽然在SSL和注意力监督损失项上会消耗更多的计算成本,但学习到的CNN编码器在Transformer引导下能够关注到正确的区域,并且不会在下游任务中消耗更多的计算成本。
在标准视觉识别基准(包括图像分类、对象检测和语义分割)上的实验表明,CARE让CNN编码器主干网络达到了更高的SOTA性能。