从 MaxViT 到 ACC-ViT：视觉 Transformer 的紧凑型网络实现参数与性能的双重优化！

AIGC 先锋科技

发布于 2024-07-08 14:04:42

1060

发布于 2024-07-08 14:04:42

现代视觉 Transformer 通过在窗口或网格区域内计算注意力来利用像素之间的视觉启发的局部交互，这与原始ViT中使用的全局注意力形成对比。区域注意力将像素交互限制在特定区域内，而稀疏注意力则将它们分散在稀疏的网格中。这些不同的方法在保持层次关系与捕捉全局上下文之间提出了挑战。在本文中，作者从扩张卷积获得灵感，提出了扩张注意力，这是一种结合了区域注意力和稀疏注意力的方法，它动态地整合局部和全局信息，同时保持层次结构。基于此，作者引入了一种多功能混合视觉 Transformer Backbone 网络ACC-ViT，专为标准视觉任务定制。作者紧凑的模型拥有不到2850万个参数，在ImageNet-1K上的准确率约为84%，比最先进MaxViT的准确率高0.42%，同时需要的参数减少了8.4%。

1 Introduction

早期的视觉 Transformer （ViT）探索了利用文本 Transformer 的看似无限可扩展性[9]来处理图像的可行性，但由于缺乏足够的感应偏置，且没有进行任何视觉特定的调整[27]，因此它们无法达到当时最先进的CNN模型[10]的水平。那些早期的ViT忽视了一个特别重要的方面，即局部模式，这些模式通常携带强烈的上下文信息[17]。这导致了Swin Transformer [18]提出的窗口注意力的发展，这是第一个真正有竞争力的视觉 Transformer 模型。

这个概念经过进一步创新，最终分为两个方向：区域注意力和稀疏注意力。

在区域注意力中，考虑不同大小的窗口来计算不同尺度上的局部注意力，而稀疏注意力在分散的稀疏网格中计算简化全局注意力的像素。区域注意力包含了一种层次感[3]，但由于只考虑几个区域以考虑计算复杂性，全局视角受到了严格的限制[34]。

相反，稀疏注意力可以以合理的计算成本计算更好的全局上下文近似，但牺牲了层次结构间的嵌套交互[30]。尽管层次信息极具价值[35]，由于可以获得更丰富的全局上下文，稀疏注意力被用于最先进的模型，而相对较小的窗口注意力可以弥补有限的局部交互。

在图1中，作者提供了区域注意力和稀疏注意力的简化示例。在区域注意力（图1a）中，猫被两个 Patch 识别。蓝色和绿色 Patch 分别检查 Head 和猫的大部分身体。通过结合信息层次，作者可以推理出它是一只猫。然而，由于计算大区域上的注意力成本高昂，区域注意力没有捕捉到腿和尾巴。在稀疏注意力（图1b）的情况下，作者可以观察到整个图像正在被由四种不同颜色标记的网格集分析。尽管看起来作者可以计算整个猫身体上的注意力，但几乎不存在信息层次，而且考虑了许多不相关的像素。例如，计算鸭像素上的注意力可能对分类猫没有太大贡献。

基于以上的直观分析，作者可以得出以下两个观察结果。首先，最好能以合理的计算成本覆盖更多区域。其次，尽管全局信息有益，作者可能通过将感受野限制在一定程度内而受益。结合这两个观察结果，作者可以在检查区域中引入稀疏性，或使网格受到限制。这归结为扩张区域，在atrous卷积[4]中已经分析过。因此，从atrous卷积中汲取灵感，作者提出了Atrous注意力。作者可以考虑多个以不同程度扩张的窗口来计算注意力，然后将它们结合起来。这同时解决了两个问题，保持了目标的层次结构并捕捉了全局上下文。例如，在图1(c)中，可以观察到具有不同扩张 Level 的3个较小 Patch 能够覆盖到猫的全部。

在传统的信号和图像处理中[12]引入的atrous或扩张卷积的概念，通过deeplab[4]恰当地引入到计算机视觉中，并迅速流行起来。通过忽略一些连续的邻近像素，扩张卷积增加了深层网络中的感受野[20]。此外，Atrous空间金字塔池化可以提取层次特征[5]，这也可以根据条件自适应地使用[22]。然而，在ViT时代，扩张卷积的使用急剧下降，仅限于 Patch 计算[14]。最近，有研究调查了在注意力计算中使用扩张。DiNat[11]涉及以固定速率计算注意力图，而DilateFormer[16]利用了较小 Kernel 的多重扩张速率。

作者为视觉 Transformer 提出了一种新型的注意力机制，并根据该注意力机制设计了一种混合视觉 Transformer 。在视觉 Transformer 时代，作者重新关注几乎被遗忘的atrous（或扩张）卷积，并发现atrous卷积的特性对视觉 Transformer 非常有用。因此，作者基于atrous卷积设计了自己的注意力机制和卷积块。此外，从atrous空间金字塔池化（ASPP）[4]中汲取灵感，作者尝试并行设计，偏离了现有堆叠不同类型 Transformer 层的方式[29]。

2 Methods

Atrous Attention

作者提出了一个结合了稀疏注意力和区域注意力的新型注意力机制：膨胀注意力（图1）。作者从膨胀卷积[4]中获得灵感，膨胀卷积通过从图像中删除某些行和列，有效地增大了感受野，同时使用的参数数量相似。这使得作者能够在合理的计算复杂度下检查全局上下文，同时通过在多个膨胀 Level 上考虑窗口，保留了层次信息。

对于一个输入特征图，其中分别指通道数、高度和宽度，作者计算不同膨胀 Level 上的窗口如下，

这里，作者只考虑膨胀率为2的幂，因为它们可以使用einops操作[23]在GPU上高效计算。与膨胀卷积类似，通过移动或滑动膨胀模式可以捕获所有值。作者在每个以不同膨胀 Level 计算出的窗口上应用了带相对位置嵌入[25]的窗口化多头自注意力，同时也应用在未膨胀的输入上()。

Gating Operation for Adaptive Fusion of Different Branches

合理聚合分层上下文可以显著提高视觉模型的视觉表征能力[35]。因此，作者开发了一种简单、轻量级且自适应的门控操作，用于合并从不同膨胀 Level 提取的特征。

基于输入特征图，通过使用单个线性层结合softmax操作为每个输出分支计算门控权重因子。作者假设，基于输入，门控操作能够理解哪些不同计算特征的部分应该被强调或丢弃。

对于不同的输出分支，计算相应的权重因子，并且可以通过哈达玛积计算融合的特征图。

Shared MLP Layer across Parallel Attentions

在作者的公式中，作者在融合的注意力图上应用了一个共享的MLP层，与传统的 Transformer 层不同，后者在每个注意力操作后使用MLP层。这既降低了计算复杂性，又使模型的学习变得更加容易（如消融研究中所示）。作者推测，拥有一个共享的MLP层减少了模型需要协调的变化程度。

Parallel Atrous Inverted Residual Convolution

倒置残差卷积在视觉 Transformer 中是实际上的卷积选择[29, 33]。然而，为了在整个模型中利用稀疏性和层次性，作者用3个并行的扩张深度可分离卷积替换了中间的深度可分离卷积，其膨胀率分别为1、2和3。这三个卷积操作的输出通过门控机制进行合并。因此，形式上作者提出的扩张倒置残差卷积块可以表示为如下：

这里，和分别指的是挤压和激励块[13]和门控操作。和的下标表示核的大小，上标指的是扩展比或膨胀率，视情况而定。

Acc-ViT

使用作者提出的注意力和卷积块，作者设计了ACC-ViT，一种混合的、分层的视觉 Transformer 架构。按照惯例，该模型由输入 Level 的卷积干细胞、4个ACC-ViT块和分类顶部[29]组成。干细胞对输入图像进行下采样，使得自注意力计算变得可行。干细胞之后是4个ACC-ViT块，这些块通过堆叠扩张卷积和注意力层构建，并且在每个块之后对图像进行下采样。由于图像尺寸变小，作者在4个块中分别使用3、2、1和0（无扩张） Level 的扩张注意力。最后，分类顶部基于全局平均池化和全连接层，这是当代架构中使用的[18, 19, 29]。作者提出的架构图示在图2中展示。

按照传统做法，作者设计了几种不同复杂度 Level 的ACC-ViT变体，配置详情在附录中给出。

3 Results

作者展示了在三个标准基准测试上，作者的方法与现有最先进方法的比较结果。第一部分描述了在任务A上的结果，作者评估了作者的方法在准确性和运行时间方面的性能。第二部分展示了在任务B上的结果，重点关注与 Baseline 方法的比较。最后，最后一部分提供了一个消融研究，以证明作者方法中各组件的有效性。

ImageNet-1k Image Classification

作者在标准的图像分类任务上训练了不同变种的ACC-ViT，使用了ImageNet-1K数据集[8]。这些模型是采用Torchvision[26]推荐的超参数值以及近期研究[29, 33, 35]中采用的方法进行训练的。在传统的有监督学习范式下，这些模型在分辨率的图像上进行训练，没有使用任何额外的数据。图3展示了ACC-ViT达到的top-1验证准确度，并与流行的最先进模型进行了比较。

在类似参数范围和FLOPs下，ACC-ViT的表现超越了最先进的MaxViT和MOAT模型。在小型和微型模型中，使用相似的FLOPs，ACC-ViT比MaxViT分别实现了高0.35%和高0.15%的准确度，尽管参数分别减少了9.15%和9.7%。对于所有其他模型，ACC-ViT的表现甚至更为令人印象深刻。ACC-ViT的小型和微型版本比其他模型的小型和基础版本更准确。

在医学图像数据集上的迁移学习实验

通用视觉 Backbone 模型最有用的应用之一是迁移学习，ViT模型已经展示了显著的有效性[38]。为了评估ACC-ViT的迁移学习能力，作者选择了具有多样性和相关挑战的医学图像数据集[21]。选择了3个不同模态和尺寸类别的医学图像数据集，以严格评估在不同配置下视觉表示的可迁移性。来自HAM10000[28]（皮肤黑色素瘤，10,015张图像）、EyePACS[7]（糖尿病视网膜病变，5,000张图像）和BUSI[2]（乳腺超声，1,578张图像）的数据集图像被随机以分层方式分为60:20:20的训练-验证-测试分割。在ImageNet-1K上预训练的Swin、ConvNext、MaxViT和ACC-ViT的微型版本被微调了100个周期，并在测试数据上计算了宏观精确度、召回率、F1分数以及准确度（表1）。可以观察到，在大多数指标上，ACC-ViT超过了其他模型，这在大型数据集上更为明显。对于小型数据集BUSI，ConvNeXt证明是最佳模型，ACC-ViT成为第二佳。这是可以预期的，因为医学成像研究表明，在小型数据集上，卷积网络的表现优于 Transformer [15]。在所有模型中，MaxViT似乎表现最差，特别是在稀有类别上，这可能意味着小型数据集不足以充分调整模型的参数。

Model Interpretation

为了解释和分析模型学习的内容，作者在Swin、MaxViT和ACC-ViT上应用了Grad-CAM [24]（图4）。类激活图显示，MaxViT倾向于关注图像的不相关部分，这可能是因为网格注意力。相反，Swin似乎经常专注于一个小局部区域。ACC-ViT明显能够检查和区分整个金鱼和橡皮擦（如图4a和b所示）。此外，在分类图4c中的火烈鸟示例时，ACC-ViT关注了整个火烈鸟群，而Swin和MaxViT分别关注了子集和不相关的像素。有趣的是，在将图4d分类为锤头鲨时，ACC-ViT只关注锤头鱼，而其他 Transformer 同时强调了两种鱼。