[Transformer] Next-ViT: Next Generation Vision Transformer

Next-ViT: Next Generation Vision Transformer for Efficient Deployment in Realistic Industrial Scenarios

https://arxiv.org/abs/2207.05501

introduction

由于复杂的注意力机制和模型设计，大多数现有的ViTs在现实的工业部署场景中不能像CNNs那样高效地执行。本文重点关注如何在现实工业场景中有效部署ViTs，提出了Next-ViT，从延迟/准确性权衡的角度来看，它在CNN和ViT中均占主导地位。

在这项工作中，分别开发了Next Convolution Block（NCB）和Next Transformer Block（NTB），以通过部署友好的机制捕获局部和全局信息。然后，Next Hybrid Strategy(NHS) 旨在以高效的混合范式堆叠NCB和NTB，从而提高各种下游任务的性能。

为了提供一个公平的比较，作者将特定硬件上的延迟视为直接的效率反馈。TensorRT 和 CoreML 分别代表了服务器端和移动设备。Next-ViT在ImageNet-1K分类任务上实现了最佳的延迟/准确性权衡。更重要的是，Next-ViT在下游任务上显示出了更显著的延迟/准确性权衡优势。

Method

Next-ViT遵循分层的金字塔体系结构，在每个阶段都有一个patch embedding层和一系列的卷积或Transformer blocks。空间分辨率将逐步降低32×，而通道尺寸将在不同的阶段中扩大。

2.1 Next Convolution Block(NCB)

首先重新审视卷积和 Transformer blocks 的一些经典结构设计，如图 3 所示。ResNet 提出的 BottleNeck 块因其固有的归纳偏差和部署而在视觉神经网络中长期占据主导地位。与 Transformer 模块相比，BottleNeck 模块的有效性并不能令人满意。

ConvNeXt 模块通过模仿 Transformer 模块的设计使 BottleNeck 模块现代化。虽然 ConvNeXt 块部分提高了网络性能，但它在 TensorRT/CoreML 上的推理速度受到低效组件的严重限制，例如 7×7深度卷积、LayerNorm 和 GELU。

Transformer blocks在各种视觉任务中取得了优异的成绩，其内在优势是由 MetaFormer 的范式和基于注意力的token mixer模块共同赋予的。然而，Transformer blocks的推理速度比 TensorRT 和 CoreML 上的 BottleNeck 块要慢得多，因为其注意力机制复杂，这在大多数现实工业场景中是难以承受的。

为了克服上述Block的失败，作者引入了 Next Convolution Block (NCB)，它在保持 BottleNeck 块的部署优势的同时获得了作为 Transformer 块的突出性能。如图 3(f) 所示，NCB 遵循 MetaFormer 的一般架构。设计了一种新颖的多头卷积注意力（MHCA）作为高效token mixer。最后，在 MetaFormer 的范式中使用 MHCA 和 MLP 层构建 NCB。

Multi-Head Convolutional Attention(MHCA)

从h个head的卷积注意力中捕获信息，用一个投影层Wp来促进head之间的信息交互。

输入特征z在通道维度上划分为多头形式。

CA是一种单头卷积注意力，其中Tm和Tn是输入特征z中的相邻token，W代表可训练参数，O代表内积。

CA 能够通过迭代优化可训练参数 W来学习局部感受野中不同 Token 之间的关系。

具体而言，MHCA 的实现是通过组卷积（多头卷积）和逐点卷积进行的。在所有 MHCA 中将 head dim 统一设置为 32，以便在 TensorRT 上使用各种优化进行快速推理。

在 NCB 中采用了高效的 BatchNorm (BN) 和 ReLU 激活函数，而不是传统 Transformer Block 中的 LayerNorm (LN) 和 GELU，进一步加快了推理速度。

消融研究中的实验结果表明，与现有块（例如 BottleNeck 块、ConvNext 块、LSA 块等）相比，NCB 具有优越性。

2.2 Next Transformer Block (NTB)

虽然通过NCB已经有效地学习了局部表示，但全局信息的捕获迫切需要解决。

Transformer Block具有较强的捕获低频信号的能力，从而提供全局信息（例如全局形状和结构）。然而，相关研究观察到，Transformer Block可能会在一定程度上恶化高频信息，如局部纹理信息。不同频率段的信号在人类视觉系统中是不可缺少的，并将以某种特定的方式融合，提取更基本和明显的特征。

基于这些观察结果，作者开发了Next Transformer Block(NTB)来捕获轻量级机制中的多频信号。此外，NTB 作为一种有效的多频信号混频器，进一步提高了整体建模能力。

如图2所示，NTB首先采用高效的多头自注意力(E-MHSA)捕获低频信号，

其中，输入特征z在通道维数上划分为多头形式。

SA是一种空间缩减自注意力算子，其灵感来自于线性SRA，其表现为：

Ps是一种具有Stride的Avg-pool操作，用于在注意力操作前对空间维度进行降采样，以降低计算成本。

作者观察到E-MHSA模块的时间消耗也很大程度地受到其通道数量的影响。因此，NTB在E-MHSA模块之前使用点卷积进行通道降维，以进一步加速推理。

此外，NTB还配备了一个MHCA模块，与E-MHSA模块合作，以捕获多频信号。

NTB的实现可以表述如下：

Pro 表示用于通道投影的逐点 卷积层。此外, NTB 统一采用 BN 和 ReLU 作为有效范数和激活层, 而不是 LN 和 GELU。与传统的 Transformer 模块相比, NTB 能够在轻量级机制中捕获和混合多频信息, 极大地提升模型性能。

2.3 Next Hybrid Strategy (NHS)

最近的一些工作结合CNN和 Transformer 实现了有效部署，如图4(b)(c)所示。

上述方法几乎都是在早期阶段采用卷积块，在最后一两个阶段采用Transformer Block。

但这些传统的混合策略在下游任务（如分割和检测）上可能达到性能饱和。分类任务只使用最后一阶段的输出来进行预测，而下游任务（如分割和检测）通常依赖于每个阶段的特征来获得更好的结果。然而，传统的混合策略只是在最后几个阶段堆叠Transformer Block。因此，浅层阶段无法捕获全局信息，例如物体的全局形状和结构，这对分割和检测任务至关重要。

为了克服现有混合策略的失败，作者提出了一种新的混合策略 (NHS)，它创造性地将卷积块 (NCB) 和Transformer Block (NTB) 堆叠在一起。NHS在控制 Transformer Block的比例的情况下，显著提升下游任务中的模型性能，以实现高效部署。

首先，为了赋予浅层捕获全局信息的能力，作者提出了一种新颖的（NCB×N+NTB×1）模式混合策略，在每个阶段依次堆叠N个NCB和一个NTB，如图所示4(d)。

作者进行了一系列实验来验证所提出的混合策略的优越性。不同混合策略的性能如表1所示。Hn表示堆叠N个NCB和一个NTB。

表2前3行的实验结果表明，大模型的性能难以提升并逐渐达到饱和。这种现象表明，通过扩大 (NCB × N + NTB × 1) 模式的 N 来扩大模型大小，即简单地添加更多的卷积块并不是最佳选择。

因此，作者通过广泛的实验探索 N 的值对模型性能的影响。如表 2（中）所示，在第3阶段构建了具有不同 N 配置的模型。为了建立具有相似延迟的模型以进行公平比较，在 N 的值较小时堆叠 L 组 (NCB × N + NTB × 1) 模式。令人惊讶的是，作者发现 (NCB × N + NTB × 1) × L 模式比 (NCB × N + NTB × 1) 模式获得了更好的模型性能。

如表 2 所示，第3阶段 N = 4 的模型实现了性能和延迟之间的最佳折衷。

通过在第3阶段扩大 (NCB × 4 + NTB × 1) × L 模式的 L 来进一步构建更大的模型，如表2（下）所示。

2.4 Next-ViT Architectures

为了与其他 SOTA 网络进行公平比较，我们提出了三个典型的变体，即 Next-ViT-S/B/L。架构规范如表 3 所示，其中 C 表示阶段的输出通道，S 表示卷积或 Avg 池中的步幅。此外，NTB中E-MHSA和MHCA的信道比例r设置为0.75，E-MHSA中不同阶段的空间缩减率s为[8,4,2,1]。每个 MLP 层的扩展比设置为 2，E-MHSA 和 MHCA 中的 head dim 设置为 32。对于归一化和激活函数，NCB 和 NTB 块都使用 BatchNorm 和 ReLU。

实验

消融实验

r代表MHSA与MHCA的通道数之比。

（a）resnet关注高频信息，Swin关注低频信息，Next-ViT可以同时关注高低频信息。

（b）Next-ViT可以捕获更丰富的上下文信息以及更精确的全局信息。