论文链接:https://arxiv.org/pdf/2207.05501.pdf
代码地址:https
Next-ViT: Next Generation Vision Transformer for Efficient Deployment in Realistic Industrial Scenar
由于复杂的注意力机制和模型设计,大多数现有的视觉Transformer(ViT)在现实的工业部署场景(如TensorRT和CoreML)中不能像卷积神经网络(CNN)那样高效地执行。视觉神经网络能否设计得像神经网络一样快速地进行推理,并且表现得像ViT一样强大?本文提出了一种用于在现实工业场景中高效部署的下一代Transformer,即Next-ViT,从延迟/准确性权衡的角度来看,它同时主导了CNN和ViT。下一代卷积块(NCB)和下一代变换块(NTB)分别用于通过部署友好机制捕获局部和全局信息。然后,下一个混合策略(NHS)被设计为将NCB和NTB堆叠在一个有效的混合模式中,这提高了各种下游任务的性能。
一、引言
本文介绍了下一代卷积块(NCB),它擅长于通过一种新的部署友好型多头卷积注意力(MHCA)捕获视觉数据中的短期依赖信息。其次,构建了下一代Transformer块(NTB),NTB不仅捕获长期依赖信息,而且还是一个轻量级的高低频信号混频器,以增强建模能力。最后,设计了下一代混合策略(NHS),在每个阶段以新的混合范式堆叠NCB和NTB,这大大降低了Transformer块的比例,并在各种下游任务中保持视觉Transformer网络的高精度。
基于上述提出的方法,本文提出了适用于现实工业部署场景的下一代视觉Transformer(简称next ViT)。在本文中,为了进行公平的比较,我们提供了一种将特定硬件上的延迟视为直接效率反馈的观点。下图中几个现有模型的准确性和延迟权衡图。如下图(a)(d)所示,Next ViT在ImageNet-1K分类任务上实现了最佳延迟/准确性权衡。更重要的是,Next ViT在下游任务上显示出更显著的延迟/准确性权衡优势。如下图(b)(c)所示,在TensorRT上,Next-ViT在COCO检测上的表现优于ResNet 5.5 mAP(从40.4到45.9),在ADE20K分割上的表现为7.7%mIoU(从38.8%到46.5%)。接下来,ViT实现了与CSWin相当的性能,同时推理速度提高了3.6倍。如下图(e)(f)所示,在CoreML上,NextViT在COCO检测上超过EfficientFormer 4.6mAP(从42.6到47.2),在ADE20K分割上超过3.5%mIoU(从45.1%到48.6%)。
本文主要贡献总结如下:
1、本文开发了强大的卷积块和Transformer块,即NCB和NTB,具有部署友好的机制。下一代ViT堆栈NCB和NTB以构建先进的CNN Transformer混合架构。
2、本文从新的视角设计了一种创新的CNN Transformer混合策略,以提高性能和效率。
二、实现细节
本文展示了Next ViT,如下图所示。Next ViT遵循分层金字塔结构,在每个阶段使用了补丁嵌入层和一系列卷积或Transformer块。空间分辨率将逐步降低32倍,而通道尺寸将在不同阶段扩展。局部和全局信息的融合也在NTB中进行,这进一步提高了建模能力。最后,系统地研究了卷积和Transformer块的集成方式。为了克服现有方法的固有缺陷,引入了下一代混合策略,该策略堆叠了创新的NCB和NTB,以构建先进的CNN Transformer混合架构。
一、Next Convolution Block (NCB)
为了展示所提出的NCB的优势,首先回顾了卷积和Transformer块的一些经典结构设计,如下图所示。ResNet提出的BottleNeck块由于其固有的感应偏差和大多数硬件平台中的部署友好特性,在视觉神经网络中长期占据主导地位。不幸的是,与Transformer模块相比,BottleNeck模块的有效性不足。ConvNeXt模块通过模仿Transformer的设计,使瓶颈模块进行现代化。虽然ConvNeXt块提高了网络性能,但其在TensorRT/CoreML上的推理速度受到严重限制。Transformer块在各种视觉任务中取得了优异的结果,其内在优势由MetaFormer和基于注意力的混合器模块共同赋予。然而,Transformer块的推理速度比BottleNeck块慢得多,因为它的注意力机制复杂,这在大多数现实的工业场景中是无法承受的。
为了克服上述缺点,本文引入了下一代卷积块(NCB),它保持了瓶颈块的部署优势,同时获得了作为Transformer的突出性能。如上图(f)所示,NCB遵循MetaFormer的一般架构,经验证,该架构对Transformer模块至关重要。同时,高效的基于注意力的令牌混合器也同样重要。本文设计了一种新型的多头卷积注意力模块(MHCA),作为一种具有部署友好卷积操作的高效令牌混合器。最后,在MetaFormer的范式中使用MHCA和MLP层构建NCB:
其中
z
l
−
1
z^{l−1}
zl−1表示来自l− 1块的输入。
二、Multi-Head Convolutional Attention (MHCA)
为了将现有的基于注意力的令牌混合器从高延迟困境中解放出来,本文设计了一种具有高效卷积运算的新型注意力机制,即卷积注意力(CA),以实现快速推理速度。同时,受MHSA中有效的多头设计的启发,使用多头范式构建卷积注意力,该范式共同关注不同位置不同表示子空间的信息,以实现有效的局部表示学习。建议的多头卷积注意(MHCA)的定义可概括如下:
MHCA从h个并行表示子空间捕获信息。
z
=
[
z
1
,
z
2
,
…
,
z
h
]
z=[z_1,z_2,…,z_h]
z=[z1,z2,…,zh]表示将输入特征z划分为通道维度上的多头形式。为了促进多个头部之间的信息交互,还为MHCA配备了投影层(WP)。CA是单头卷积注意,可定义为:
其中
T
m
T_m
Tm和
T
n
T_n
Tn是输入特征z中的相邻标记。O是具有可训练参数W和输入标记
T
{
m
,
n
}
T_{\{m,n\}}
T{m,n}的内积运算。CA能够通过迭代优化可训练参数W来学习局部感受野中不同标记之间的权重。MHCA的实现是通过分组卷积(多头卷积)和逐点卷积来实现的,如上图(f)所示。在所有MHCA中统一将head dim设置为32,以便在TensorRT上进行快速推理。此外,在NCB中采用了高效的BatchNorm(BN)和ReLU激活函数,而不是传统Transformer块中的LayerNorm(LN)和GELU,这进一步加快了推理速度。
三、Next Transformer Block (NTB)
尽管已经通过NCB有效地提取局部特征,但全局信息的获取仍亟待解决。Transformer具有很强的捕获低频信号的能力,这些低频信号提供全局信息(例如全局形状和结构)。然而,相关研究观察到,Transformer块可能会在一定程度上恶化高频信息,如局部纹理信息。不同频率段中的信号在人类视觉系统中是不可或缺的,并将以某种特定的方式进行融合,以提取更重要和独特的特征。 基于这些观察结果,本文开发了下一代Transformer(NTB),以在轻量级机制中捕获多频率信号。此外,NTB是一种有效的多频信号混频器,可进一步增强整体建模能力。如上文中结构图所示,NTB首先使用高效的多头自注意力(E-MHSA)捕获低频信号 :
其中
z
=
[
z
1
,
z
2
,
…
,
z
h
]
z=[z_1,z_2,…,z_h]
z=[z1,z2,…,zh]$表示在通道维度中将输入特征z划分为多头形式。SA是一种空间自注意算子,其灵感来自线性SRA,表现为:
p
s
p_s
ps是具有步长s的平均池化操作,用于在注意力操作之前对空间维度进行下采样以降低计算成本。 具体而言,E-MHSA模块的时间消耗也受到其通道数量的极大影响。因此,NTB在E-MHSA模块之前使用逐点卷积执行通道降维,以进一步加速推理。引入收缩比r以减少通道。还利用E-MHSA模块中的批处理规范化实现了极其高效的部署。
此外,NTB配备有MHCA模块,该模块与E-MHSA模块协作以捕获多频率信号。之后,来自E-MHSA和MHCA的输出特征被级联以混合高低频信息。最后借用MLP层来提取更重要和更独特的特征。简言之,NTB的操作如下:
其中,
z
l
~
\tilde{z^l}
zl~、
z
l
^
\hat{z^l}
zl^和
z
l
z^l
zl分别表示E-MHSA、MHCA和NTB的输出。
P
r
o
j
Proj
Proj表示通道投影的逐点卷积层。此外,NTB统一采用BN和ReLU作为有效规范和激活层,而不是LN和GELU。与传统的Transformer块相比,NTB能够在轻量级机制中捕获和混合多频率信息,这大大提高了模型性能。
四、Next Hybrid Strategy (NHS)
最近的一些工作为将CNN和Transformer结合起来以实现高效部署付出了巨大努力。如下图(b)(c)所示。它们在浅层阶段单调地采用卷积块,而在最后一个或两个阶段仅堆叠Transformer块。这些传统的混合策略很容易在下游任务(例如分割和检测)上达到性能饱和。原因是,分类任务仅使用最后阶段的输出进行预测,而下游任务(例如分割和检测)通常依赖于每个阶段的特征以获得更好的结果。然而,传统的混合策略只是在最后几个阶段堆叠Transformer块。因此,浅层阶段无法捕获全局信息,例如对象的全局形状和结构,这对于分割和检测任务至关重要。
为了克服现有混合策略的缺陷,本文从新的视角提出了下一代混合策略(NHS),该策略创造性地将卷积块(NCB)和Transformer块(NTB)用
(
N
+
1
)
∗
L
(N+1)∗ L
(N+1)∗L混合范式。NHS在控制Transformer块的比例以实现高效部署的情况下,显著提高了下游任务中的模型性能。首先,为了赋予浅层阶段捕获全局信息的能力,提出了一种新的(NCB×N+NTB×1)模式的混合策略。如上图(d)所示,在每一级中堆叠N个NCB和一个NTB。具体来说,Transformer块(NTB)放置在每个阶段的末尾,这使模型能够学习浅层中的全局表示。本文进行了一系列实验来验证所提出的混合策略的优越性。差分混合策略的性能如下表所示。C表示在一个阶段中均匀堆叠卷积块(NCB),T表示用Transformer块(NTB)构建一个阶段。
H
N
H_N
HN表示在相应阶段以(NCB×N+NTB×1)模式堆叠NCB和NTB。下表中的所有型号都配备了四级。例如,C C C C表示在所有四个阶段中始终使用卷积块。为了公平比较,在相似的TensorRT延迟下构建了所有模型。如下表所示,与下游任务中的现有方法相比,所提出的混合策略显著提高了模型性能。
C
H
N
H
N
H
N
C H_N H_N H_N
CHNHNHN实现了最佳的整体性能,在检测方面超过
C
C
C
T
C C C T
CCCT 0.8 mAP,在分割方面超过0.8%mIoU。
通过增加第三阶段的块数进一步验证了
C
H
N
H
N
H
N
C H_N H_N H_N
CHNHNHN在大模型上的一般有效性。下表中前三行的实验结果表明,大型模型的性能难以提升,并逐渐达到饱和。这种现象表明,通过放大(NCB×N+NTB×1)模式的N来扩展模型大小,即简单地添加更多卷积块不是最佳选择。这也意味着(NCB×N+NTB×1)模式中的N值可能严重影响模型性能。因此,通过大量实验探索N值对模型性能的影响。如下表(中间)所示,在第三阶段构建了N的不同配置的模型。为了建立具有相似延迟的公平比较模型,当N值很小时,堆叠L组(NCB×N+NTB×1)模式。结果展示(NCB×N+NTB×1)×L模式的堆叠NCB和NTB比(NCB×N+NTB×2)模式获得更好的模型性能。这表明以适当的方式((NCB×N+NTB×1))重复组合低频信号提取器和高频信号提取器可导致更高质量的表示学习。如下表所示,第三阶段中N=4的模型实现了性能和延迟之间的最佳权衡。在第三阶段,通过放大(NCB×4+NTB×1)×L模式的L来进一步构建更大的模型。如下表(底部)所示,Base(L=4)和Large(L=6)的性能与小模型相比,模型得到了显著提升,这验证了所提出的(NCB×N+NTB×1)×L模式的总体有效性。在本文的其余部分中,使用N=4作为基本配置。
将NCB和NTB与上述Next Hybrid Strategy进行叠加,以构建Next ViT,其正式定义为:
三、实验
ImageNet-1K分类的不同最先进方法的比较。
ADE20K语义分割任务中不同主干的比较。
比较基于Mask R-CNN的对象检测和实例分割任务的不同主干。
四、消融实验
为了验证所提出的NCB的有效性,用著名的块替换了NextViT中的NCB,例如ResNet中的瓶颈块、ConvNeXt块、Twins中的LSA块等。
本文探讨了Next Transformer Block的收缩率r对Next ViT整体性能的影响。如上表所示,减小收缩比r,即e-MHSA模块中的通道数量,将减少模型延迟。此外,r=0.75和r=0.5的模型比纯Transformer(r=1)的模型性能更好。这意味着适当的融合多个频率信号处理将提高表征学习的模型能力。
特别是,r=0.75的模型实现了最佳的延迟/准确性权衡。它在分类、检测和分割方面优于基线模型(r=1.0),分别为0.4%、0.5APb和1.0%mIoU,同时更轻。上述结果表明了所提出的NTB块的有效性。
本文进一步研究了Next ViT中不同规范化层和激活函数的影响。如上表所示,LN和GELU都带来了微不足道的性能改善,但TensorRT上的推断延迟明显更高。另一方面,BN和ReLU在总体任务上实现了最佳的延迟/准确性权衡。因此,Next ViT中统一使用BN和ReLU,以便在现实的工业场景中高效部署。