论文研读 | fastervit: fast vision transformers with hierarchical attention

阿航的大佬梦

已于 2024-07-25 09:43:48 修改

阅读量737

点赞数 26

分类专栏：人工智能论文学习录文章标签：人工智能计算机视觉笔记

于 2024-07-25 09:38:55 首次发布

本文链接：https://blog.csdn.net/Verhan_Cao/article/details/140680164

版权

人工智能论文学习录专栏收录该内容

2 篇文章 0 订阅

订阅专栏

（以下是我小组的研究成果，将开源与大家共同分享。论文我附上了文件）

论文下载地址：阿里云盘分享

一、背景

目前ViT模型普遍存在着这样的问题：

ViT的核心工作是注意力机制，但其时间复杂度为O(n^2)，这会影响ViT模型的运行速度。对于高分辨率的图片，带来的巨大计算成本阻碍其在高分辨率图像应用中的使用。
ViT没有CNN所具备的归纳偏置，比如平移等变性和局部性，导致ViT要重新学习那些先验知识，在中小训练集上训练达不到较好的性能。

另外另外些工作发现，多尺度图片的训练方式学习特征表示通常会产生更好的性能，于是Liu等人提出了Swin Transformer。

但NVIDIA研究团队发现Swin Transformer仍依赖原始的窗口易懂机制，在大尺寸图像下的应用仍得不到好的效果。于是NVIDIA研究团队提出了他们的ViT家族，为FasterViT。

二、模型介绍

FasterViT的宏观架构图如上图所示，主要分为4个stage，前面两个是密集卷积阶段，后面两个阶段是他们提出的新的模块——Hierarchical Attention。

FasterViT是如何解决上面这些问题的？

对于ViT本身注意力时间上的局限，NVIDIA研究团队提出了HAT设计，分层注意力 (HAT) 方法将具有二次复杂度的全局自注意力分解为多级注意力，从而降低了计算成本。每个窗口都可以访问专门的载体标记（Carrier Tokens），这些标记参与局部和全局表示学习。在高层面上，全局自注意力机制能够以较低的成本实现高效的跨窗口通信。在这个模块中，整张图片以一个相对较小的计算成本达到了比仅局部注意力好的多的效果。新板块的设计是其一，其二是架构在最前设计了下采样的卷积操作，这缩小了图片的大小，减少后续的计算量，让模型能处理高分辨率的图片提供了另一种的可能性。

同时FasterViT是一个混合CNN-ViT的架构，这种架构就很好的吸收了CNN的优点，使ViT也具备了CNN的归纳偏置。在模型的前半阶段，进行大量的密集卷积操作，以至于首次输入到HAT进行注意力时，已拥有了一个较好的特征，减少模型较好拟合规律所需的数据量。

架构中可分为四类模块：Conv(s2)、Conv Block、Downsample和Hierarchical Attention，下面我一一介绍其工作机制。

（一）Conv(s2)模块

输入图像通过连续的两个3×3×3卷积层投影为D维embedding，每个卷积层的步长为2。embedding会进一步批归一化，每次卷积后都会使用ReLU激活函数。

（二）Conv Block模块

阶段1和阶段2由残差卷积块组成，定义为

其中BN表示批归一化。遵循设计原则，这些卷积是密集的。

（三）Down sample模块

FasterViT的下采样块先对空间特征应用2D层归一化，然后使用内核为3×3×3且步长为2的卷积层，将空间分辨率降低2倍。

（四）HAT模块

1、直观解释

这是作者团队提供的直观解释图，我们先聚焦Hierarchical attention图，这个图展示的是Carrier tokens相互间做注意力机制，要提前说的是每个CT都是对应划分窗口的总结。CT间做注意力后，每个CT已含有了全局上的一个感受。接着进行local window attention，在每个窗口内部与CT做注意力，这样全局的信息通过CT带到了窗口内部的每个tokens。这就是Global attention，分层注意力，也就是HAT核心的逻辑。

2、详细框架

以上就是HAT模块的设计细节，接下来我介绍其中几个关键的过程。

（1）CT初始化（CT init）

CT是怎么产生的？它是窗口内特征的总结。CT的产生经历两个阶段，一是卷积，二是池化。

3×3×L的滤波器组对划分窗口进行卷积，并且使用的是same卷积，不改变卷积前后图片的大小。卷积完后会得到L个通道的图片特征组，最后池化到2（默认）个token组成的CT。

（2）CT的注意力

将组成的CT同样进行tokennization后，成为token序列。然后与主观解释类似，再CTs赋予了二维的位置信息后，CT间进行注意力、交换彼此的信息，再经过MLP层，MLP是由两个神经层组成，中间经历GeLU激活函数。输入层与输出层的节点数相同，不改变其个数，目的是增加其非线性。

（3）CT-Tokens的注意力

CT注意力完成后，会与窗口tokennization好后的token序列拼接在一起，然后类似的对序列赋予它所处的位置信息。下一步是做每个窗口局部的自注意力，把全局的感受带给局部。最后再经历次MLP，同样的不改变其个数。

四、模型运行机制

假设输入特征图的尺寸为224*224*3。

然后进入Conv(s2)模块，经过3*3*3的卷积核达到下采样的效果，特征图的尺寸变为112*112*D，我们设置卷积核组数为6，那最后卷积后的结果为112*112*6。接着会进行批量归一化操作核ReLU函数激活。第二的Conv(s2)模块与第一块相同。两个卷积过程完成后，特征图的尺寸变为56*56*12。

进入FasterViT模型的主体（4个Stage），前面两个Stage是密集卷积操作，先经历N1个Conv Block后经历N2个Conv Block，为了避免梯度消失问题里面的卷积间采用残差连接的设计。此Conv Bolck模块并不会改变特征图的尺寸大小，但是每次大量密集卷积完成后会进行一个下采样的操作。在此过程，特征图的尺寸从56*56*12变为28*28*24再变为14*14*48。接着进入Stage3，也是作者提出的新模块——HAT。

上图是HAT的细节图，特征图此时的尺寸为14*14*48。

首先进行的是对窗口的划分，我们设置k=7，特征图的分割与图中所画一致，被分为四个local window。

CTs的init过程是先经历3*3*3的卷积核组对local window进行卷积，然后经过特征池化的操作。我们设置总结对应窗口的token为2，也就是每个CT是两个token。

Local window和对应CTs都要进行tokenization过程编程序列，CT间会进行自注意力机制。

上公式展示了所有的细节，在注意力前CTs首先要进行归一化操作，并且将子注意力的结果乘以一个学习效率，将结果累加到原CTs上。

全局注意力后会进入MLP，MLP没有改变输出的个数（与输入相同），只有两层但其中的GeLU激活函数，增加了其表示的非线性。

将序列进行拼接。

进行局部的注意力机制，将全局的感受带到每一个token。

然后是分割拼接的序列。这个ViT模块进行了N3次。最终还原成特征图的形式。

位置信息是有用的，在中间多次还引入了token的位置信息，这有助于注意力学习。因为位置编码由MLP插值，所以这种方法对图像大小变化是具有灵活性的，经过训练的模型可以应用于任何输入分辨率。

从stage3输出的特征图的尺寸没有发生变化，仍是14*14*45，接着进行下采样的操作尺寸变为7*7*96。在经历类似的stage4阶段，最后输出最后高度处理后的特征表示。

最后经过一个简单的分类器，就可以完成分类任务。

五、实验验证

（一）图像分类

NVIDIA研究团队开发的FasterViT模型在ImageNet-1K数据集上表现出色，与各种混合架构、卷积神经网络以及基于Transformer的模型相比，在相同的吞吐量条件下实现了更高的分类精度。特别地，与ConvNeXt-T模型相比，FasterViT的Top-1准确率提高了2.2个百分点。在平衡准确性和吞吐量方面，FasterViT的表现优于基于Transformer的网络，如Swin Transformer系列，在推理速度上有明显优势。此外，与最近提出的高效模型如EfficientFormer和MaxViT相比，FasterViT不仅在平均吞吐量上更具竞争力，而且在ImageNet分类Top-1指标上的性能也更胜一筹。进一步地，通过使用TensorRT等优化技术，FasterViT的延迟-准确性Pareto前沿仍然保持领先。

（二）目标检测

在MS COCO数据集上对目标检测和实例分割任务进行性能评估时，FasterViT模型展现了在准确性和吞吐量之间的出色平衡。例如，FasterViT-4在使用Cascade Mask R-CNN网络进行检测时，相较于ConvNeXt-B和Swin-B，在盒子AP指标上分别提高了0.2和1.0，在掩膜AP指标上分别提高了0.3和1.0，同时在吞吐量方面分别快了15%和30%。这一趋势在其他模型变体中也有体现。

此外，当使用FasterViT-4作为ImageNet-21K预训练的骨干网络，并与最先进的DINO模型结合进行额外的目标检测实验时，FasterViT-4达到了58.7的盒子AP，这进一步证明了FasterViT作为高性能骨干网络的有效性，尤其是在与更复杂的先进模型结合使用时。

（三）语义分割

在ADE20K数据集上进行语义分割实验时，使用UPerNet网络，FasterViT模型在性能和吞吐量的平衡上表现出色。具体来说，FasterViT-4在mIoU指标上优于Swin-B，在单尺度和多尺度推理方面分别提高了1.0和0.7，并且吞吐量提高了16.94%。与ConvNeXt-B相比，在多尺度推理的情况下，FasterViT-4的吞吐量提高了7.01%，同时mIoU提高了0.4。这些结果表明FasterViT在保持高精度的同时，还能够在推理速度上展现出明显的优势。