随着Transformer结构在自然语言处理领域的成功，怎样将transformer应用到视觉任务，设计出更好的视觉主干网络 (vision transformer) 从而为其他下游任务提供更好的特征表示引发了研究者的广泛关注。 然而，当前的vision transformer模型往往需要在更大的数据集进行训练，或者需要更长的训练时长。 这样的训练要求限制了对vision transformer结构的进一步探索和应用。
因此， 京东探索研究院联合悉尼大学 提出了一种新型视觉神经主干网络——ViTAE，在改善模型收敛性的同时，提高模型性能，助力多种下游任务涨点，如语义分割、实例分割、物体检测、视频物体分割，以及姿态估计。 该工作目前已被NeurIPS2021接收。
GitHub ：
https://github.com/Annbless/ViTAE
文章：
https://openreview.net/pdf?id=_RnHyIeu5Y5
一、研究背景
由于Transformer中的自注意力机制具有很强地建模长距离依赖的能力，其在自然语言处理（NLP）领域已经获得了巨大的成功。如何更好的将transformer架构应用到视觉任务中吸引了广泛的关注，现有工作（比如ViT[2], T2T[4]）将整张图像分为若干小块，每个小块记为一个token，将所有token重排列为一个序列后，直接输入transformer网络，提取特征。然而，为了获得较好的图像分类性能，Vision transformer往往需要大量的额外数据以及更长的训练时间，为transformer在视觉任务的实际应用造成困难。这种现象的一个重要的原因是现有工作将图像作为一维序列，忽略了对视觉任务特有的归纳偏置的建模，即对图像局部相关性、物体的尺度不变性的建模，导致模型无法高效地利用数据，影响收敛速度和模型性能。
图1 在不同数量的训练数据、以及不同的训练步数下，T2T、DeiT和ViTAE(本文方法)的分类表现对比。
图2 ViTAE模型结构图
二、 方法原理简述
为了解决这一问题，我们提出从模型设计的角度在transformer中引入归纳偏置，并由此设计了一种新的视觉主干架构（ViTAE），其显着改善了transformer的收敛速度和性能，如图1所示。 具体地，我们在ViTAE中引入两种模块结构设计，即卷积旁路分支用以改善模型对图像局部相关性的建模，和多尺度卷积用以提取不同尺度物体的特征，从而改善ViTAE对物体的尺度不变性的建模。 基于这两种模块结构，我们设计了两种基础模块，我们分别称其为reduction cell (RC)和normal cell (NC)。 如图2上半部分所示，ViTAE由两种基础模块堆叠组成，RC和NC分别由绿色和蓝色表示，其中RC模块将图像降采样，NC模块则负责进一步提取特征。
1、RC模块具体设计
如图2左下所示，给定输入特征，RC使用多尺度卷积核嵌入不同尺度的上下文信息，保证其对不同尺度的物体均有较好的响应，并将特征图降采样到指定大小。然后RC使用多头注意力层（MHSA）建模长距离信息以及一个卷积旁路分支建模局部相关性。最后使用FFN层用以特征转换。由此，RC模块中包含了多尺度卷积以及旁路卷积，分别引入了本文提出的两种归纳偏置。
2、 NC模块具体设计
NC模块具体设计在图2右下给出。由于之前堆叠的RC模块已经将特征图降采样到相对较小的尺度
，因此与RC相比，NC没有使用多尺度卷积核，只使用了旁路卷积保证模型对局部相关性的建模。
图3 ViTAE-Stage变体
3、ViTAE-Stage具体设计
RC模块与NC模块以不同方式堆叠构成了不同的ViTAE变体。其中，为了适应下游分割、检测任务对多尺度特征图的需要，我们以交错方式重排列RC与NC模块，构成了ViTAE-Stage变体，其每个stage的最后一个NC模块的输出被用来执行对下游任务的预测。
三、实验结果
表1 ViTAE和现有方法在ImageNet的分类结果对比
表1展示了不同模型大小的ViTAE变体在ImageNet数据集的分类结果。如图所示，ViTAE在所有模型大小下均达到了最优性能，显着超出ViT[2]、DeiT[3]、T2T[4]、PVT[5]等模型，证明了将归纳偏置引入transformer可以显着改善模型收敛，提高模型性能。
表2 ViTAE-Stage在下游视觉任务与其他方法的性能对比
表2展示了ViTAE-Stage与其他模型在下游视觉任务的性能对比，其中可以看出，ViTAE作为主干网络，在各个下游任务中均带来显着提升。该试验证明了将归纳偏置引入transformer模型对不同任务的泛化性，以及引入归纳偏置对transformer的重要性。
图4 每个注意力层的平均注意力距离对比
为了展示归纳偏置对模型的作用，图4对比了ViTAE和其他方法在不同层的平均注意力距离。 由于引入了对局部相关性的建模，ViTAE中的多头注意力层只需关注长距离依赖关系，因此有较长的注意力距离。 与之相反，其他方法在浅层时则需要更关注于局部，因此其浅层时的注意力距离较短。 该试验揭示了引入的归纳偏置在ViTAE中的作用机制，证明了引入的归纳偏置的有效性。
图5 ViTAE模型的主观结果图
图5展示了Grad-CAM[6]对ViTAE最后一个NC模块的多头注意力层的输出的可视化结果。与其他方法对比，ViTAE可以准确覆盖不同尺度大小的物体(右三列均属于‘鸟’类，但具有不同的物体尺度)，并且对背景信息有较少的关注，证明了将归纳偏执引入transformer中可以帮助模型学习到更有区分性的特征，从而帮助改善模型在不同视觉任务上的性能表现。
四、理性复杂度分析
表3 模型推理速度比较
对于视觉主干网络来说，推理吞吐量是影响其应用的重要因素。为此，我们分析了所提出的ViTAE架构和传统Vision Transformer架构的推理时间消耗。结果如表3所示，其中T2T-ViT[4]是典型的Vision Transformer方法。相比较于T2T-ViT，ViTAE的推理速度和同等模型大小的T2T-ViT相当，但是取得了更好的分类准确度。在分类准确度相似的情况下，ViTAE取得了更快的推理速度，为其在下游任务更广泛的应用提供了可能性。
五、未来工作
除了局部相关性和物体多尺度不变性，我们会继续探索怎样利用其他的归纳偏置进一步改善vision transformer模型。此外，更大规模的模型和更高的准确度是视觉transformer模型的发展方向和追求。目前，我们已经将ViTAE模型扩大到100M和200M的模型尺寸，并取得了具有潜力的结果。除此以外，Vision Transformer的模型依然需要大量有标注的数据，能否通过自监督的方式来利用大量无标注的数据来进一步训练Vision Transformer模型也是一个值得探索的方向。
六、总结
在这篇论文中，我们重新设计了transformer架构，并提出两种基础模块（reduction cell和normal cell）用以引入两种归纳偏置，即图像局部相关性和物体的尺度不变形，并由此提出了新的视觉主干架构ViTAE。我们通过大量的实验证明了ViTAE在模型性能、数据利用效率、模型收敛性以及对不同下游任务的泛化性均具有显着优势，验证了归纳偏置对transformer架构的重要性。
参考文献
[1]. Yufei Xu, Qiming Zhang, Jing Zhang, and Dacheng Tao. ViTAE: Vision Transformer Advanced by Exploring Intrinsic Inductive Bias. Advances in Neural Information Processing Systems., 2021. arXiv:2106.03348 (2021).
[2]. Dosovitskiy, Alexey, et al. “An image is worth 16×16 words: Transformers for image recognition at scale.” ICLR 2020.
[3]. Touvron, Hugo, et al. “Training data-efficient image transformers & distillation through attention.” ICML 2021.
[4]. Yuan, Li, et al. “Tokens-to-token vit: Training vision transformers from scratch on imagenet.” ICCV 2021.
[5]. Wang, Wenhai, et al. “Pyramid vision transformer: A versatile backbone for dense prediction without convolutions.” ICCV 2021.
[6]. Selvaraju, Ramprasaath R., et al. “Grad-cam: Visual explanations from deep networks via gradient-based localization.” ICCV 2017.
本文来自：公众号【 京东探索研究院 】
作者： 京东探索研究院
Illustrastion b y Delesign Graphics
-The End-
Be First to Comment