official implementation：GitHub - bytedance/Next-ViT
third-party implementation：https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/nextvit.py
背景
近年来，视觉Transformer（ViTs）在工业界和学术界受到了越来越多的关注，并在图像分类、目标检测、语义分割等各种计算机视觉任务中取得了显著成功。然而，从实际部署的角度来看，卷积神经网络（CNNs）仍然主导着视觉任务，因为ViTs通常比经典的CNN（如ResNet）要慢得多。Transformer模型的推理速度受限于多头自注意力（MHSA）机制的二次复杂性、不可折叠的LayerNorm和GELU层、复杂的模型设计导致频繁的内存访问和复制等因素。
出发点
本文的出发点是设计一种视觉神经网络，能够像CNN一样快速推理，并且像ViTs一样强大。现有的很多工作尝试通过设计更高效的空间注意力机制来缓解MHSA的二次计算复杂性，或者结合高效的卷积block和强大的Transformer block来设计CNN-Transformer混合架构，以在准确性和延迟之间取得更好的平衡。然而，这些现有的混合架构在下游任务（如分割和检测）中容易导致性能饱和，并且无法同时具备高效性和性能。
创新点
- Next Convolution Block (NCB)：提出了一种新的多头卷积注意力（MHCA）机制，以一种对部署友好的方式捕获视觉数据中的短距离依赖信息。
- Next Transformer Block (NTB)：NTB不仅擅长捕获长距离依赖信息，还能作为一种轻量级的高低频信号混合器，增强建模能力。
- Next Hybrid Strategy (NHS)：设计了一种新的混合策略，在每个阶段高效地堆叠NCB和NTB，大大减少Transformer block的比例，同时保持视觉Transformer网络在各种下游任务中的高精度。
- Next-ViT：提出了一种适用于实际工业部署场景的下一代视觉Transformer，通过在TensorRT和CoreML上的广泛实验，展示了其在各种视觉任务中的优势，显著优于现有的CNNs、ViTs和CNN-Transformer混合架构。
方法介绍
Next Convolution Block（NCB）
作者首先回顾了一些经典的conv block和Transformer block的结构设计，如图3所示。ResNet提出的BottleNeck block在视觉神经网络中长期占据主导地位，但与Transformer block相比性能不够充分。ConvNeXt block模仿Transformer block的设计对BottleNeck block进行改进，在一定程度上提高了性能，但由于低效的组件在TensorRT/CoreML上的推理速度严重受限，例如7x7深度卷积、LayerNorm和GELU。Transformer block在各种视觉任务中表现良好，但推理速度比BottleNeck block慢得多，这在大多数工业场景中是难以忍受的。
为了克服上述问题，作者提出了Next Convolution Block（NCB），它在保持BottleNeck block部署优势的同时获得了Transformer block的性能。具体如图3(f)所示，NCB遵循了MetaFormer（具体介绍见MetaFormer（CVPR 2022，Sea）- CSDN博客）的一般结构，作者设计了一个新的Multi-Head Convolutional Attention（MHCA）作为token mixer，NCB可以表示如下
其中 \(z^{l-1}\) 表示来自 \(l-1\) block的输入，\(\tilde{z}^l\) 和 \(z^l\) 分别表示MHCA的输出和 \(l\) NCB的输出。
Mutil-Head Convolutional Attention（MHCA）
MHCA可以表示如下
这里MHCA从 \(h\) 个并行子空间中捕获信息，\(z=[z_1,z_2,...,z_h]\) 表示沿通道维度将输入特征 \(z\) 划分为多头形式。为了促进多头之间的信息交互，作者还为MHCA配备了一个投影层 \(W^P\)。CA是单头卷积注意力，可以表示如下
其中 \(T_m\) 和 \(T_n\) 是输入特征 \(z\) 中相邻的token，\(O\) 是一个具有可训练参数 \(W\) 和输入token \(T_{\{m,n\}}\) 的内积操作。CA能够通过迭代优化 \(W\) 来学习局部感受野中不同token之间的相关性。具体来说MHCA包含一个group convolution（multi-head convolution）和point-wise convolution，如图3(f)所示。MHCA中的head dim统一设置为32，此外NCB中采用BN和ReLU而不是Transformer block中的LN和GELU以进一步加快推理速度。
Next Transformer Block（NTB）
有研究观察到Transformer block在一定程度上会破坏高频信号，如局部纹理信息。不同频率段的信号在人类视觉系统中是必不可少的，并以某种特定的方式融合来提取基本和明显的特征。
受此观察启发，作者设计了Next Transforme Block（NTB）以一种轻量机制来捕获多频信号。此外，NTB作为一种有效多频信号混合器，进一步提高了整体建模能力，如图2所示。
NTB首先用Efficient Multi-Head Self Attention（E-MHSA）来捕获低频信号，表示如下
其中 \(z=[z_1,z_2,...,z_h]\) 表示沿通道维度将输入特征 \(z\) 划分为多头形式，SA是spatial reduction self-attention受PVT v2（具体介绍见PVT v2 原理与代码解析 - CSDN博客）中的Linear SRA启发，表示如下
其中Attention表示标准的注意力计算即 \(softmax(\frac{QK^T}{d_k})V\)，\(d_k\) 表示scaling factor。\(P_s\) 是步长为 \(s\) 的avg-pool用来在attention计算之间对空间维度进行降采样以降低计算成本。因为E-MHSA的时间消耗很大程度上也受到通道数量的影响，NTB在E-MHSA之前用point-wise卷积进行通道降维，以进一步加速推理。引入了一个收缩比 \(r\) 用于减少通道数，E-MHSA中也使用BN来实现高效的部署。
此外，NTB还配置了一个MHCA与E-MHSA协作来捕获多频信号。两者的输出concat起来以混合高低频信号。最后用一个MLP层来提取更基本和明显的特征。NTB的整体过程可以表示如下
其中 \(\tilde{z}^l,\hat{z}^{l},z^l\) 分别表示E-MHSA、MHCA和NTB的输出。Proj表示point-wise卷积，同样NTB也统一采用BN和ReLU。
Next Hybrid Strategy（NHS）
如图4(b)(c)所示，现有的卷积和Transformer混合的策略几乎都是采用浅层使用卷积block深层使用Transformer block的方式，但作者观察到这种混合策略在下游任务（如检测和分割）中很容易达到性能饱和。原因是分类任务只使用最后一个阶段的输出来进行预测，而下游任务如检测和分割通常依赖每个阶段的输出特征来获得更好的结果。而之前的混合策略只在最后几个阶段使用Transformer block，这样浅层阶段就无法捕获全局信息，例如物体的全局形状和结构，这对分割和检测任务至关重要。
为了克服上述问题，作者提出了一种新的混合策略Next Hybrid Strategy（NHS），具体采用 \((N+1)*L\) 的方式来混合NCB和NTB。如图4(d)所示，采用（NCB x N + NTB x 1）的方式混合两个block，在每个阶段的最后采用一个NTB可以使模型在浅层也能学习到全局表示。
如表1所示，作者比较了不同混合策略的性能和延迟，其中C表示卷积block，T表示Transformer block，\(H_N\) 表示本文提出的（NCB x N + NTB x 1）的混合方式，可以看到本文提出的混合策略在下游任务上显著优于之前的方法。例如 \(C\ H_N\ H_N\ H_N\) 在检测任务中超过了 \(C\ C\ C\ T\) 0.8 mAP，分割任务中超过了0.8% mIoU。此外，\(H_N \ H_N\ H_N\ H_N\) 的结果表明在第一个阶段使用Transformer block会破坏模型的精度-延迟的平衡。
此外，当通过增加block数量来构建更大的模型时，如表2上半部分所示，只增加NCB的数量性能很快达到饱和。相反，以（NCB x N + NTB x 1）x L的方式增加L的数量可以获得更好的性能。因此本文采用后者的方式来扩展模型。
实验结果
作者设计了S/B/L三种不同大小的模型，具体配置如下。
在ImageNet上和其它模型的对比如表4所示
在ADE20K上采用Semantic FPN分割模型的效果如表5所示
采用Mask R-CNN在COCO数据集上的目标检测和实例分割性能对比如表6所示
989
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
