CPU推理1ms的SHViT Backbone 来啦

近期，高效的视觉 Transformer 在资源受限的设备上展现出优异的低延迟性能。传统上，它们在宏观层面上采用4×4的块嵌入和4阶段结构，同时在微观层面上使用具有多头配置的复杂注意力机制。
本文旨在以节省内存的方式解决所有设计层面的计算冗余问题。作者发现，使用更大步长的块状处理初始部分不仅能降低内存访问成本，还通过利用早期阶段减少空间冗余的 Token 表示来达到具有竞争力的性能。此外，作者初步分析表明，早期阶段的注意力层可以用卷积替代，且后期阶段有许多注意力头在计算上是冗余的。为了处理这个问题，作者引入了一个单头注意力模块，它内在地防止了头冗余，并通过并行组合全局和局部信息同时提高准确性。
基于作者的解决方案，作者提出了SHViT，一个单头视觉 Transformer ，它获得了最佳的速率-精度权衡。例如，在ImageNet-1k上，SHViT-S4在GPU、CPU和iPhone12移动设备上分别比MobileViTv2 ×1.0快3.3倍、8.1倍和2.4倍，同时精度高出1.3%。
对于在MS COCO上进行目标检测和实例分割，使用Mask-RCNN Head ，SHViT在GPU和移动设备上的 Backbone 延迟分别比FastViT-SA12低3.8倍和2.0倍，同时表现出与之相当的性能。

1 Introduction

视觉 Transformer （ViT）已经在各种计算机视觉任务中展示了其卓越的性能，这归功于它们的高模型能力。与卷积神经网络（CNN）相比，ViT在建模远程依赖关系方面表现出色，并且能够随着大量训练数据和模型参数的有效扩展。尽管具有这些优势，但标准的ViT缺乏归纳偏置，这需要更多的训练数据，并且全局注意力模块与图像大小成二次计算复杂度。为了应对这些挑战，先前的研究要么将ViT与CNN结合，要么引入了成本效益高的注意力变体。

近期，针对实时性约束问题的研究也提出了遵循类似策略的高效模型。它们的策略可以分为两组：

高效架构 - 宏观设计
高效的多头自注意力（MHSA）- 微观设计

探索架构设计的研究利用卷积处理高分辨率/低级特征，并采用注意力处理低分辨率/高级特征，展现了无需复杂操作的优越性能。然而，这些方法大多数主要关注于用哪些模块来聚合Token，而不是如何构建Token（关于切块茎和阶段设计）。另一方面，高效的MHSA技术通过实施稀疏注意力或低秩近似来降低注意力的成本。这些模块与通常采用的多头机制相结合使用。尽管取得了所有这些巨大进步，宏观/微观设计中的冗余仍然没有得到充分的理解或解决。在本文中，作者探索了所有设计层次的冗余，并提出了内存高效解决方案。

为了识别宏观设计中的计算冗余，作者专注于 Patch 嵌入大小，观察到最近的高效模型使用了44的 Patch 嵌入。作者进行了如图2所示的实验，以分析传统宏观设计中的空间冗余，并发现了一些有趣的结果。

首先，尽管通道较少，但由于大量的Token（在224224的分辨率下，第一阶段有3136个Token；第二阶段有784个Token），早期阶段表现出了严重的速度瓶颈。
其次，使用一个3阶段设计，在第一阶段通过1616的块状处理196个Token，并不会导致性能显著下降。

为了进一步比较，作者设置了一个基本模型（表6（2）），采用前述宏观设计，仅使用33的深度卷积进行简单的Token混合。与高效模型MobileViT-XS相比，SHViT简单模型在ImageNet-1k上的准确度提高了1.5%，同时在A100 GPU / Intel CPU上的运行速度提高了**5.0**/**7.6**。这些结果表明，在早期阶段存在大量的空间冗余，并且与专门的注意力方法相比，高效的宏观设计对于模型在严格的延迟限制内实现具有竞争力的性能更为关键。需要注意的是，这一观察并不意味着Token混合器是微不足道的。

作者还探讨了微观设计中的冗余，特别是在MHSA层中。大多数高效的MHSA方法主要关注有效的空间Token混合。由于高效的宏观设计，作者能够使用具有更高语义密度的紧凑型Token表示。因此，作者将关注点转向了注意力层中存在的通道（头）冗余，这也是大多数以前的工作忽略的一个重要方面。通过全面的实验，作者发现多头机制中存在明显的冗余，特别是在后期的阶段。
然后，作者提出了一种新颖的单头自注意力（SHSA）作为竞争性替代方案，以减少计算冗余。在SHSA中，单头自注意力仅应用于输入通道的一个子集，而其余通道保持不变。SHSA层不仅消除了由多头机制产生的计算冗余，还通过处理部分通道降低了内存访问成本。此外，这些效率使得可以在不增加计算预算的情况下堆叠更多块并扩大宽度，从而在同一计算预算内提高性能。

基于这些发现，作者引入了一种基于内存高效设计原则的单头视觉 Transformer （SHViT），作为在多种设备上运行速度极快的新一类网络。实验证明，SHViT在速度和准确性方面为分类、检测和分割任务达到了最先进的表现，如图1所示。

例如，SHViT-S4在Nvidia A100 GPU上以14283张图像/秒的吞吐量和在2.10GHz的Intel Xeon Gold 5218R CPU上以509张图像/秒的速度，在ImageNet上达到了79.4%的top-1准确率，比EfficientNet-B0在准确率上高出2.3%，GPU推理速度上快69.4%，CPU速度上快90.6%。

同时，SHViT-S4比MobileViT1.0准确率高1.3%，在iPhone12移动设备上快2.4倍。对于在MS COCO上进行目标检测和实例分割，使用Mask-RCNN检测器，SHViT在各种设备上具有更小的主干延迟，显著优于EfficientViT-M4，AP高出6.2，AP高出4.9。

总之，作者的贡献如下：

作者对大部分现有研究中被忽视的冗余进行了系统分析，并提出了内存高效的设计原则来解决这个问题。
作者引入了单头Vision Transformer（SHViT），它在各种设备（如GPU、CPU和iPhone移动设备）上都能取得良好的准确度与速度的平衡。
作者针对各种任务进行了广泛的实验，并验证了SHViT在高速和有效性方面的高性能。

2 Analysis and Method

在本节中，作者首先通过前所未有的实验对宏观和微观设计中的冗余进行分析，然后讨论了各种减轻这些冗余的解决方案。之后，作者介绍了单头视觉 Transformer （SHViT）并解释其细节。

Analysis of Redundancy in Macro Design

最高效的模型采用了44的分块起始/4阶段配置（图2（a））。相比之下，普通的ViT模型采用1616的分块起始来为后续的MHSA层生成有意义的输入Token。作者关注这种差异，并进一步假设在大步幅的分块起始不仅对于MHSA层是必要的，而且对于在严格延迟限制下的有效表示学习也至关重要。

为了证实作者的假设，作者采用了1616的块状起始（patchify stem）和3阶段设计。作者基于MetaFormer块和上述两种宏观设计构建了两个模型（具体细节见图2）。

具体来说，作者配置了两个模型，使其在等效特征图尺寸下拥有相似数量的通道。令人惊讶的是，尽管模型（b）在GPU/CPU上的速度分别是（a）的3.0/2.8，但其性能比（a）差1.5%。此外，当在256256的分辨率下训练时，（b）不仅与（a）相当，而且显著更快。

如上述观察所示，作者提出的有效宏观设计具有以下优点：

1) 可以在早期阶段利用具有大接受域和减少的空间冗余的token表示。

2) 它可以将特征图的大小缩小多达16倍，从而显著降低内存访问成本。

3) 由于激进的步幅设计，当分辨率提高时，吞吐量只会轻微下降，这导致了有效的性能提升（如图2(b)、8、表2所示）。

Analysis of Redundancy in Micro Design

MHSA层在多个子空间（头）中独立计算并应用注意力图，这已一致显示出性能的增强。然而，尽管注意力图在计算上要求较高，但最近的研究表明，其中许多并不是至关重要。

作者还通过三个实验深入探讨了流行的微小ViT模型（DeiT-T，Swin-T）的多头冗余性：注意力图可视化，头相似性分析，以及头消融研究。

对于 Head 分相似性分析，作者衡量了同一层中每个头与其他头之间的平均余弦相似度。在进行 Head 分消融研究时，作者评估了通过取消给定层中某些头的输出，同时保持其他层的头完整时对性能的影响。并报告了最高分数。每个实验的细节及进一步的成果在补充材料中提供。

首先，在早期阶段（图3（a）），作者观察到表现最佳的头（heads）倾向于以卷积方式操作，而在移除后对性能影响最小的头通常更全局地处理信息。

此外，如图2（b）所示，与在第一阶段使用深度卷积层的模型相比，在第一阶段使用注意力层的模型在速度-准确度权衡上表现较差。因此，为了效率，作者在初始阶段使用带有空间归纳偏置的卷积作为Token混合器（token mixer）。

在后期阶段，作者发现特性和预测层面上都存在很多冗余。例如，DeiT-T（图3（b））的后阶段显示出平均 Head 相似度为78.3%（6个 Head 的相似度为64.8%），Swin-T也表现出显著的高值（图4左侧）。在移除一个 Head 的实验中，作者观察到大部分 Head 可以被移除，而不会与原始准确度偏离太远。

值得注意的是，在某些Swin-T的情况下（图4右侧），移除一个 Head 甚至会导致分数略有提升。此外，当只使用Swin-T中12个或24个 Head 中的一个时，平均性能下降仅为0.95个百分点。

先前的方法通常首先完全训练网络，然后剪除不必要的 Head 以解决 Head 冗余问题。尽管这些方法有效，但它们在训练过程中以增加计算资源和内存占用为代价。为了高效地解决上述问题，作者设计了一个单头注意模块，从本质上避免 Head 冗余。这种方法确保训练和推理过程都得到简化和高效。

Single-Head Self-Attention

基于上述分析，作者提出了一种新的单头自注意力（SHSA）机制，具体细节展示在图5的右下角。它仅仅在一个输入通道的部分（=）应用单头注意力层进行空间特征聚合，而剩下的通道保持不变。作者默认将设置为。正式地说，SHSA层可以描述为：

其中是投影权重，是 Query 和键的维度（默认设置为16），是连接操作。为了一致的内存访问，作者将最初的通道作为整个特征图的代表。此外，SHSA 的最终投影应用于所有通道，而不仅仅是最初的通道，确保了注意力特征有效地传播到剩余的通道。SHSA 可以被解释为沿着块轴顺序地拉伸先前并行计算的冗余头。

在图6中，作者也探索了各种单头设计。近期研究依次结合卷积和注意力层，将局部细节融入到全局语境中。不幸的是，这种方法在给定的Token混合器中只能提取局部细节或全局语境。同时，在[6]中指出，某些通道处理局部细节，而其他通道处理全局建模。这些观察表明，当前的串行方法在处理每一层的所有通道时存在冗余（图6（a），（b））。

相比之下，作者通过在一个Token混合器内并行利用两个互补特性，以及先前的卷积记忆高效地解决了上述问题。

为了有效利用注意力层，层归一化是必不可少的；同时，为了实现多头注意力方法，需要诸如 Reshape 操作之类的数据移动。因此，如图7所示，MHSA的大部分运行时间被内存受限的操作，如 Reshape 和归一化所占。通过最小化内存受限操作的使用或将它们应用于更少的输入通道，SHSA模块可以充分利用GPU/CPU的计算能力。

Single-Head Vision Transformer

单头视觉 Transformer （SHViT）架构的概述如图5所示。给定一个输入图像，作者首先对它应用四个33的带步长的卷积层。与标准ViT模型使用的步长为16的1616卷积Token方法相比，作者的重叠切块干细胞可以提取更好的局部表示。然后，Token通过三个堆叠的SHViT块阶段进行分层表示提取。一个SHViT块由三个主要模块组成（见图5）：深度卷积（DWConv）层用于局部特征聚合或条件位置嵌入，单头自注意力（SHSA）层用于建模全局上下文，以及用于通道交互的前馈网络（FFN）。FFN中的扩展比例设置为2。

DWConv和SHSA的组合以一种计算和内存效率的方式捕捉到局部和全局的依赖关系。基于第2.2节的发现，作者在第一阶段不使用SHSA层。为了在不损失信息的情况下减少Token，作者使用了一个有效的下采样层，该层由两个第一阶段块组成，它们之间放置了一个倒置的残差块（步长-2）。最后，全局平均池化和全连接层被用于输出预测。

除了上述运算符外，归一化和激活层在决定模型速度方面也起着关键作用。作者只为SHSA层采用层归一化，而将批量归一化（BN）整合到其余层中，因为BN可以被合并到其相邻的卷积或线性层中。作者还使用ReLU激活函数，而不是其他更复杂的替代方案，因为它们在各种推理部署平台上的速度要慢得多。

作者构建了四种不同深度和宽度设置的SHViT变体。由于大尺寸的 Patch 嵌入和单头设计，作者可以使用比之前高效模型更多的通道和块。模型规格在表1中提供。

3 Experiments

Implementation Details

作者在ImageNet-1K数据集上进行图像分类，该数据集包含1.28M个训练图像和50K个验证图像，分为1000个类别。所有模型都是从头开始训练，使用AdamW优化器进行300个周期的训练，学习率为，总批处理大小为2048。作者采用余弦学习率调度器，并配合5个周期的线性预热。权重衰减分别设置为0.025/0.032/0.035/0.03，对应SHViT-S1至S4。

为了公平比较，数据增强方法包括Mixup，随机擦除，以及自动增强。对于和分辨率，作者以权重衰减为，学习率为0.004，对模型进行30个周期的微调。此外，还评估了在各种硬件平台上的吞吐量性能。作者在Nvidia A100 GPU上以256的批处理大小测量GPU吞吐量。

对于CPU和CPU，作者在Intel(R) Xeon(R) Gold 5218R CPU @ 2.10GHz处理器上评估运行时间，批处理大小为16（使用单线程）。对于CPU，作者将模型转换为ONNX运行时格式。移动端延迟是在iPhone 12上使用iOS版本16.5进行测量的。作者使用CoreML工具导出模型（批处理大小设为1），并在1000次运行中报告中位数延迟。作者还验证了SHViT作为高效的视觉 Backbone 网络，用于COCO数据集上的目标检测和实例分割，分别使用RetinaNet和Mask R-CNN。所有模型都在mmdetection库中进行1训练计划（12个周期）。

SHViT on ImageNet-1K Classification

如图1、表2和表4所示，作者比较了单头视觉 Transformer （SHViT）与现有最先进模型。

比较结果清楚地表明，SHViT在各种设备上在准确性和吞吐量/延迟之间取得了更好的平衡。

与高效卷积神经网络的比较。 SHViT-S1比MobileNetV3-Small的准确率高出5.4%，同时在A100 GPU和Intel CPU上的速度保持相似。与ShuffleNetV2 2.0相比，SHViT-S2在A100 GPU和Intel CPU上分别实现了2.2和2.5的速度提升，同时性能略有提高。

此外，当SHViT转换为ONNX运行时格式时，速度提高了3.8。与近期提出的FasterNet-T1相比，SHViT-S3不仅准确率高出1.2%，而且在A100 GPU上的运行速度提高了15.1%，在Intel CPU上的运行速度提高了32.4%。值得注意的是，在Top-1准确率达到79.179.4%的情况下，SHViT在A100 GPU/Intel CPU上，以ONNX格式比EfficientNet-B1快2.9/ 3.3，比FastViT-T12快3.4/ 5.5，比MobileOne-S4快2.7/ 1.8。当以节省内存的方式利用最小的注意力模块时，变体 Transformer （VTs）仍然可以展现出类似于高效CNN的快速推理速度。

与高效ViTs和混合模型的比较。 SHViT-S1在A100 GPU和Intel CPU上的ONNX运行时，分别比EfficientViT-M2的吞吐量高出10%和42%，同时显示出较大的性能优势（70.8% 72.8%）。SHViT-S3的准确度与PoolFormer-S12相似，但它使用的FLOPs减少了3倍，在A100 GPU上快3.8倍，在ONNX模型上快6.1倍。

值得注意的是，SHViT-S4在准确度上超过了最近的EdgeViT-XS，高出1.9%，同时在A100 GPU上快2.6倍，在Intel CPU上快1.7倍，在ONNX实现上快4.2倍。

如上述结果所示，当转换成ONNX格式时，SHViT与最近的SOTA模型相比，显示出显著的性能提升。这种增强主要是因为作者的单头设计使用了更少的 Reshape 操作，这些操作通常会在ONNX运行时引起开销。总结来说，上述结果表明，作者提出的内存高效宏观设计对速度-准确度权衡的影响比高效的注意力变体或高度简化的操作（如池化）更为重要。

高分辨率下的微调。作者也把SHViT-S4微调到更高的分辨率。与最先进的EfficientViT-M5相比，SHViT-S4即便在较低分辨率下训练，也取得了具有竞争力的性能。此外，SHViT-S4在A100 GPU上快了77.4%，在Intel CPU上快了55.6%，在ONNX运行时格式上更是快了3.6倍。而且，SHViT-S4在A100 GPU上实现了82.0%的top-1准确率，吞吐量为3957张图像/秒，证明了其在不同输入尺寸上的有效性。

蒸馏结果。作者在表3中报告了使用DeiT 蒸馏方法的作者模型的性能。值得注意的是，SHViT在速度和准确性上都超过了竞争模型。具体来说，SHViT-S3甚至超过了作为ONNX模型的慢5.2倍的FastViT-T8。SHViT-S4在GPU / ONNX运行时的性能优于EfficientFormer-L1，速度分别快2.1倍/1.9倍。

移动端延迟评估。作者在表4中验证了SHViT在移动设备上的有效性。与高效模型EfficientNet-B1 / MobileOne-S4相比，SHViT-S4在iPhone 12设备上实现了类似的准确度，同时运行速度提高了0.2毫秒/0.1毫秒。SHViT-S4在面对为移动延迟优化程度很高的模型时也获得了有竞争力的性能，表明其在不同推理平台上一致的性能表现。

图8中的进一步结果显示，SHViT明显优于近期模型FastViT和EfficientFormer，特别是在更高分辨率上。在低分辨率上，SHViT-S4稍慢一些，但在1024 1024的分辨率下，SHViT分别比FastViT和EfficientFormer的延迟低34.4%和69.7%。这些结果源于宏观和微观设计上内存效率的提高。

SHViT on downstream tasks

在表5中，作者使用两种框架来评估SHViT的迁移能力：

用于目标检测的RetinaNet
用于实例分割的Mask R-CNN

目标检测。 SHViT-S4在移动设备上比MobileNetV3 快2.3倍，并且准确率AP高出+8.9。与MobileFormer相比，SHViT在性能更优的同时，在A100 GPU上快3.2倍，在移动设备上快8.2倍。

实例分割。 SHViT-S4在速度上超过了针对GPU或移动设备优化的模型，如EfficientViT和EfficientNet，同时提供了显著的性能提升。

值得注意的是，SHViT相较于PoolFormer获得了1.7 AP和1.3 AP的提升，但在GPU、CPU和移动设备上的运行速度分别快4.3倍、8.1倍和2.1倍。

如上述结果所示，采用3阶段的大步幅切块式stem不仅降低了计算成本，还生成了有意义的token表征，尤其是在更高分辨率下。此外，与采用类似宏观设计的EfficientViT相比，显著的性能差距证明了作者微观设计选择的有效性。

Ablation Study

在本节中，作者首先验证所提出的单头自注意力（SHSA）层的效果，然后对SHSA层的部分比例值进行简化的消融研究。结果在表6中提供。

SHSA的有效性。为了评估SHSA层是否能够像多头自注意力（MHSA）层那样有效地捕捉全局上下文，作者通过用MHSA层替换SHSA层或者直接移除SHSA层来进行消融研究。如表6（1, 2 与 4）所示，与MHSA层相比，SHSA层显示出更好的速度与准确度权衡。

虽然移除SHSA层会使模型速度更快，但这会导致准确度显著下降。同时，模型（2）与表2中的SOTA模型相比，也能达到高度竞争力的性能，这表明作者在严格的延迟约束下提出的宏观设计提供了一个坚实的基础架构 Baseline 。

寻找合适的SHSA部分比例。默认情况下，作者将所有SHViT模型的部分比例设置为1 / 4.67，以获得最佳的速度-精度权衡（3, 5 对比 4）。与非常小的值相比，适度增加通道以进行 Token 交互可以在低成本下实现有效的性能提升。同时，过大的值并不能提供足够的性能提升来补偿随之而来的成本。

4 Conclusion

在这项工作中，作者调查了许多成熟模型常用的架构设计在空间和通道维度上的冗余。然后，作者提出了带有3级分层表示的16×16 Patch 嵌入和单头自注意力来处理计算上的冗余。此外，作者还展示了基于作者提出的宏观/微观设计的多功能单头视觉 Transformer ，该 Transformer 在各种设备和视觉任务上实现了超快速推理速度和高性能。

虽然作者的宏观设计是有效的，但仍然需要细粒度（高分辨率）特征来进一步提升性能或识别小物体。因此，作者未来的工作将聚焦于设计成本效益高的方法来利用这些特征。将单头设计整合到现有的复杂注意力机制中，将是一种有趣的探索方式。

参考

[1].SHViT: Single-Head Vision Transformer with Memory Efficient Macro Design.

作者：AI视界引擎
来源：AI视界引擎

推荐阅读

本文章著作权归作者所有，任何形式的转载都请注明出处。更多动态滤波，图像质量，超分辨相关请关注我的专栏AIWalker。欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。