FasterViT｜英伟达提出分层注意力，达成SOTA性能-吞吐量均衡

论文地址：https://arxiv.org/abs/2306.06189

代码地址：https://github.com/NVlabs/FasterViT

出发点

本文旨在面向主流硬件设备(如GPU)研发具有高吞吐量的骨干架构。当前主流硬件设备包含多个CUDA与Tensor核计算单元，它需要进行频繁的数据迁移进行计算，可能会受到数据移动带宽影响。因此，以下两种类型算子需要进行精心平衡以最大化吞吐量：

• 受计算量约束的算子称之为math-limited，笔者将其称之为计算约束算子;
• 受访存迁移约束的算子称之为memory-limited，笔者将其称之为访存约束算子。

在分层架构中, 特征分辨率会随推理进行而逐渐缩小, 如从输入的 到最后 一个stage的 。

• 网络的起始阶段特征分辨率比较大(如 )，而通道数较少，此时会受访存迁移约束。在这个阶段，计算密集型的算子更为适合，如稠密卷积，而非depth-wise卷积；类似的，非矩阵操作形式的算子(如非线性、池化、BN)属于访存约束算子，应当尽可能少的使用。
• 在网络的后半部分更倾向于计算约束算子。以网络第三阶段为例，特征分辨率为 ，这为更强表达能力的算子(如LN、SE、自注意力等对吞吐量影响较小)提供了操作空间。

基于上述考量，本文提出了一种新型的混合架构FasterViT，它取得了更优的性能-吞吐量均衡，达成了新的Pareto前沿，可参考下图。

本文方案

上图给出了本文所提FasterViT架构示意图，从中可以看到:

• 在网络的前半部分，特征分辨率比较大，推理效率主要受访存约束，故它仅采用了卷积类操作以充分利用稠密卷积；
• 在网络的后半部分，特征分辨率比较小，推理效率会受计算约束，故它采用了attention等操作提升表达能力。

Stem 该部分通过两个连续 卷积(stride为 2 )将输入图像 转换为重叠 块, 输出维度为 的嵌入向量( embedding )。此外, 每个卷积后接BN与ReLU激活 函数。

DownBlock FasterViT延续了分层结构，相邻stage之间通过下采样模块进行空域分辨 率收缩。下采样模块由LN与 stride 的 卷积构成。

ConvBlock Stage1与Stage2由残差模块构成，定义如下：

HierarchicalAttention 上图给出了本文所提新型窗口注意力示意图，通过如下操作，本文构建了分层注意力(Hierarchical Attention)的概念。

• 它以SwinT所引入的局部窗口为出发点，
• 引入了一种CT(Carrier Tokens)进行完整局部窗口信息聚合，第一个注意力模块作用于CTs以进行全局信息建模；
• 将局部窗口Token与CTs进行拼接并执行自注意力处理以促进局部与全局信息交互。

假设给定输入特征 (为简单起见, 这里假设 ), 我们首先将输入特 征拆分为 局部窗口, 表示如下:

本文方案的关键在于Carrier Tokens(CTs)的构建, 它有助于以更低的消耗保持注意力机 制。首先, 我们通过池化操作初始化CTs, 描述如下:

注1 : 这里的 表示高效位置编码。

注2: 上述CT初始化步骤仅在每个stage执行一次, 每个Stage有自己独一的CTs。

在每个HAT模块内部, CTs执行如下注意力处理:

接下来，为进行长短程空域信息建模，我们在之间进行信息交互，处理流程如下：

最后，上述所得Tokens将被拆分为拆分以用于后续分层注意力模块中：

在每个stage，上述步骤将迭代执行多次。为促进长短程交互，在每个stage的尾部参考EdgeViT进行了全局信息传播。每个stage的输出计算如下：

Merge

此外，对于窗口Tokens与CTs，参考SwinV2注入了绝对位置偏置，而位置编码方式的灵活性促使其可以用于任务输入分辨率。

上图对比了不同全局-局部自注意力模块之间的区别，所提HAT将自注意力拆分为局部与亚全局形式且可压缩为2个稠密注意力。

本文实验

ImageNet分类

上表对比了不同Backbone在ImageNet分类任务上的性能，可以看到：

• 相比卷积架构，在同等吞吐量下，FasterViT具有更高的精度；
• 相比Transformer架构，FasterViT具有更快的推理速度；
• 相比其他混合架构，FasterViT具有更高的吞吐量，同时具有更优的Top1性能。

COCO检测与实例分割

上表给出了COCO检测与实例分割任务上不同方案性能对比，从中可以看到：

• 相比其他方案，FasterViT具有更优的精度-吞吐量均衡。
• 相比ConvNeXt-B，FasterViT指标高出0.2boxAP、0.3MaskAP，同时吞吐量高出15%；
• 相比Swin-B，FasterViT指标高1.0boxAP、1.0maskAP，同时吞吐量高出30%。

ADE20K分割

上表给出了ADE20K语义分割任务上不同方案性能对比，可以看到：

• FasterViT具有更优的性能-吞吐量均衡；
• 相比Swin-B，FasterViT指标高出1.0mIoU@ss, 0.7mIoU@ms，同时吞吐量高出16.94%；
• 相比ConvNeXt-B，FasterViT指标高出0.4mIoU@ms，同时吞吐量高出7.01%。

全文到此结束，更多消融实验与分析请移步原文。

公众号后台回复“极市直播”获取100+期极市技术直播回放+PPT