SHViT: Single-Head Vision Transformer with Memory Efficient Macro Design——具有内存高效宏观设计的单头视觉Transformer

这篇文章提出了一种名为**SHViT（Single-Head Vision Transformer）**的新型视觉Transformer架构，旨在通过减少计算冗余和内存占用，提升模型在资源受限设备上的推理速度和性能。以下是文章的主要内容总结：

1. 背景与动机

• **视觉Transformer（ViT）**在计算机视觉任务中表现出色，但其计算复杂度高，尤其是在多头自注意力（MHSA）机制中，导致在资源受限设备上的推理速度较慢。

• 现有的高效ViT模型通常使用4×4的patch嵌入和4阶段结构，但这些设计在宏观和微观层面存在冗余，导致计算效率低下。

2. 主要贡献

• 宏观设计优化：提出使用更大步长的patchify stem（如16×16），减少早期阶段的空间冗余，从而降低内存访问成本并提升性能。

• 微观设计优化：提出**单头自注意力（SHSA）**模块，替代传统的多头自注意力机制，减少计算冗余，同时通过并行处理全局和局部信息提升准确性。

• SHViT架构：基于上述优化，提出了SHViT模型，该模型在ImageNet分类、目标检测和实例分割等任务上实现了最先进的速度-准确性权衡。

3. 实验结果

• ImageNet分类：SHViT在GPU、CPU和移动设备上均表现出色，比现有的高效模型（如MobileViT、EfficientNet等）更快且更准确。

• 目标检测与实例分割：在COCO数据集上，SHViT作为骨干网络，显著降低了推理延迟，同时保持了高准确性。

• 高分辨率微调：SHViT在高分辨率图像上的表现也优于现有模型，展示了其在不同输入尺寸下的有效性。

4. 创新点

• 单头自注意力（SHSA）：通过减少注意力头的数量，显著降低了计算复杂度和内存占用，同时保持了模型的表达能力。

• 内存高效设计：通过优化patch嵌入和阶段结构，减少了早期阶段的计算冗余，提升了整体模型的推理速度。

5. 结论

• SHViT通过宏观和微观设计的优化，显著提升了视觉Transformer的计算效率，使其在资源受限设备上能够实现快速且准确的推理。

• 实验结果表明，SHViT在多个视觉任务上均优于现有的高效模型，展示了其在速度和准确性之间的卓越平衡。

这篇文章的核心思想是通过减少计算冗余和内存占用，设计出一种高效且快速的视觉Transformer架构，适用于各种资源受限的设备。这里是自己的论文阅读记录，感兴趣的话可以参考一下，如果需要阅读原文的话可以看这里，如下所示：

官方项目地址在这里，如下所示：

摘要

近年来，高效的视觉Transformer（ViT）在资源受限的设备上展示了出色的性能和低延迟。传统上，它们在宏观层面使用4×4的patch嵌入和4阶段结构，而在微观层面则使用多头配置的复杂注意力机制。本文旨在以内存高效的方式解决所有设计层面的计算冗余问题。我们发现，使用更大步长的patchify stem不仅减少了内存访问成本，还通过利用早期阶段减少空间冗余的token表示，实现了竞争性的性能。此外，我们的初步分析表明，早期阶段的注意力层可以用卷积替代，而后期的多个注意力头在计算上是冗余的。为此，我们引入了一种单头注意力模块，该模块不仅避免了头的冗余，还通过并行结合全局和局部信息来提高准确性。基于这些解决方案，我们提出了SHViT（单头视觉Transformer），它在速度和准确性之间取得了最先进的平衡。例如，在ImageNet-1k上，我们的SHViT-S4在GPU、CPU和iPhone12移动设备上分别比MobileViTv2 ×1.0快3.3倍、8.1倍和2.4倍，同时准确性提高了1.3%。在使用Mask-RCNN头的MS COCO上进行目标检测和实例分割时，我们的模型在性能上与FastViT-SA12相当，同时在GPU和移动设备上的骨干网络延迟分别降低了3.8倍和2.0倍。

1. 引言

视觉Transformer（ViT）由于其高模型能力，在各种计算机视觉任务中展示了令人印象深刻的性能[1, 2, 3]。与卷积神经网络（CNN）[4, 5]相比，ViT在建模长距离依赖性和有效扩展大量训练数据和模型参数方面表现出色[6]。尽管有这些优势，普通ViT缺乏归纳偏差，需要更多的训练数据，并且全局注意力模块在图像大小方面具有二次计算复杂度。为了解决这些挑战，之前的研究要么将ViT与CNN结合，要么引入了成本高效的注意力变体。

最近，针对实时约束问题的研究也提出了遵循类似策略的高效模型。这些策略可以分为两类：1）高效架构——宏观设计；2）高效多头自注意力（MHSA）——微观设计。探索架构设计的研究[7, 8, 9, 10, 11, 12]利用卷积处理高分辨率/低层次特征，并采用注意力机制处理低分辨率/高层次特征，展示了在不复杂操作下的优越性能。然而，这些方法主要关注使用哪些模块来聚合token，而不是如何构建token（关于patchify stem和阶段设计）。另一方面，高效的MHSA技术通过实现稀疏注意力[13, 14, 15, 16, 17, 18]或低秩近似[19, 20, 21, 22]来降低注意力成本。这些模块通常与多头机制一起应用。尽管取得了很大进展，但宏观/微观设计中的冗余仍未完全理解或解决。在本文中，我们探索了所有设计层面的冗余，并提出了内存高效的解决方案。

为了识别宏观设计中的计算冗余，我们专注于patch嵌入大小，观察到大多数最近的高效模型使用4×4的patch嵌入。我们进行了如图2所示的实验，以分析传统宏观设计中的空间冗余，并发现了一些有趣的结果。首先，尽管早期阶段的通道较少，但由于token数量庞大（在224×224分辨率下，阶段1：3136个token；阶段2：784个token），早期阶段表现出严重的速度瓶颈。其次，使用16×16的patchify stem处理第一阶段196个token的3阶段设计并不会导致性能显著下降。为了进一步比较，我们设置了一个基本模型（表6（2）），采用上述宏观设计，使用3×3深度卷积进行简单的token混合。与高效模型MobileViT-XS [18]相比，我们的简单模型在ImageNet-1k [23]上实现了1.5%的更高准确性，同时在A100 GPU / Intel CPU上运行速度快5.0倍 / 7.6倍。这些结果表明，早期阶段存在相当大的空间冗余，与专门的注意力方法相比，高效的宏观设计对于模型在严格的延迟限制下实现竞争性性能更为关键。 注意，这一观察并不意味着token mixer是微不足道的。

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

基于这些发现，我们引入了一种基于内存高效设计原则的单头视觉Transformer（SHViT），作为一种在各种设备上运行速度极快的新网络家族。实验表明，我们的SHViT在分类、检测和分割任务中在速度和准确性方面都达到了最先进的性能，如图1所示。例如，我们的SHViT-S4在ImageNet上实现了79.4%的top-1准确性，在Nvidia A100 GPU上的吞吐量为14283 images/s，在Intel Xeon Gold 5218R CPU @ 2.10GHz上的吞吐量为509 images/s，比EfficientNet-B0 [24]在准确性上提高了2.3%，在GPU推理速度上提高了69.4%，在CPU速度上提高了90.6%。此外，SHViT-S4比MobileViTv2×1.0 [21]的准确性高1.3%，在iPhone12移动设备上快2.4倍。在使用Mask-RCNN [26]检测器的MS COCO [25]上进行目标检测和实例分割时，我们的模型在各种设备上的骨干网络延迟较小的情况下，显著优于EfficientViT-M4 [27]，分别提高了6.2 AP^box和4.9 AP^mask。

图2. 宏观设计分析。所有阶段均由MetaFormer块[28]组成。蓝色和红色所示的阶段分别使用深度卷积和注意力层作为token mixer。在下表中，宏观设计数字表示通道数，括号中的数字表示块的数量。

总之，我们的贡献如下：

• 我们对大多数现有研究中被忽视的冗余进行了系统分析，并提出了内存高效的设计原则来解决它。

• 我们引入了单头视觉Transformer（SHViT），它在各种设备（如GPU、CPU和iPhone移动设备）上实现了良好的准确性-速度权衡。

• 我们在各种任务上进行了广泛的实验，验证了SHViT的高速度和有效性。

2. 分析与方法

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

2.1 宏观设计中的冗余分析

大多数高效模型[5, 7, 8, 9, 11, 17, 19, 21, 29, 30]采用4×4的patchify stem / 4阶段配置（图2（a））。相比之下，普通ViT模型[1, 31]采用16×16的patchify stem来为后续的MHSA层生成有意义的输入token。我们关注这一差异，并进一步假设，更大步长的patchify stem不仅对MHSA层是必要的，而且对于在严格的延迟机制内进行有效的表示学习也至关重要。

图3. DeiT [31]上的多头冗余分析。为了更好地分析头的冗余，我们将DeiT-T的头数从3增加到6并重新训练模型。我们计算了注意力图，并计算了来自ImageNet的128个测试样本在不同层中每个头之间的平均余弦相似度。每个头的重要性通过其被移除和保留时的得分来确定。放大以获得更好的可视性。

图4. Swin [40]上的多头冗余分析。我们通过将Swin-T的宽度减半来缩小模型。左图：平均余弦相似度。右图：头掩码结果。结果的推导过程与DeiT实验相同。（见图3）

为了验证我们的假设，受[10, 27, 32]启发，我们采用16×16的patchify stem和3阶段设计。我们基于MetaFormer块[28]和上述两种宏观设计构建了两个模型（见图2的详细信息）。具体来说，我们配置了两个模型，使其具有相似数量的通道以实现等效的特征图大小。令人惊讶的是，模型（b）在GPU / CPU上分别快3.0倍 / 2.8倍，尽管其性能比（a）低1.5%。此外，当在256×256分辨率下训练时，（b′）不仅与（a）相当，而且速度显著更快。

如上所述，我们提出的高效宏观设计具有以下优势：

1）可以在早期阶段利用具有大感受野和减少空间冗余的token表示。
2）它可以将特征图大小减少多达16倍，从而显著降低内存访问成本。
3）由于激进的步长设计，当分辨率增加时，吞吐量仅略有下降，从而实现了有效的性能提升（如图2（b′）、8和表2所示）。

2.2 微观设计中的冗余分析

MHSA层在多个子空间（头）中独立计算和应用注意力图，这始终显示出增强的性能[1, 33]。然而，尽管注意力图在计算上要求很高，但最近的研究表明，许多注意力图并不是至关重要的[34, 35, 36, 37, 38, 39]。我们还通过三个实验深入研究了当前小型ViT模型（DeiT-T [31], Swin-T [40]）的多头冗余：注意力图可视化、头相似性分析_和_头消融研究。

对于头相似性分析，我们测量了同一层中每个头与其他头之间的平均余弦相似度。对于头消融研究，我们通过在某些层中取消某些头的输出来评估性能影响，同时在其他层中保持完整的头。并报告最高得分。每个实验的详细信息和进一步结果见补充材料。

首先，在早期阶段（图3（a）），我们观察到表现最好的头往往以类似卷积的方式操作，而移除时对性能影响最小的头通常更全局地处理信息。此外，如图2（b′′）所示，在第一阶段使用注意力层的模型与在第一阶段使用深度卷积层的模型相比，表现出不太有利的速度-准确性权衡。因此，为了效率，我们在初始阶段使用具有空间归纳偏差的卷积作为token mixer。

在后期阶段，我们发现特征和预测层面都存在大量冗余。例如，DeiT-T的后期阶段（图3（b））显示出78.3%的平均头相似度（6个头为64.8%），Swin-T也显示出显著的高值（图4左）。在移除一个头的实验中，我们观察到大多数头可以在不偏离原始准确性的情况下被移除。值得注意的是，在某些Swin-T的情况下（图4右），移除一个头甚至会导致得分略有提高。此外，当仅使用Swin-T中的12或24个头中的一个头时，性能下降平均仅为0.95个百分点。

以前的方法[34, 35, 36, 37, 38, 39]通常首先训练完整网络，然后修剪不必要的头。尽管这些方法有效，但它们以增加计算资源和训练期间的内存占用为代价。为了解决上述问题，我们设计了具有单头的注意力模块，以固有地避免头冗余。这种方法确保了训练和推理过程的简化和高效。

2.3 单头自注意力

在图6中，我们还探索了各种单头设计。最近的研究[7, 12, 27, 29, 32, 41, 42]依次结合卷积和注意力层，将局部细节融入全局上下文中。不幸的是，这种方法只能在给定的token mixer中提取局部细节或全局上下文。此外，[6]指出，一些通道处理局部细节，而另一些通道处理全局建模。这些观察表明，当前的串行方法在处理每层中的所有通道时存在冗余（图6（a）、（b））。相比之下，我们的部分通道方法通过在一个token mixer中并行利用两个互补特征，内存高效地解决了上述问题[12, 43]。

为了有效利用注意力层，层归一化[44]是必不可少的；同时，为了实现多头方法，需要像reshape操作这样的数据移动。因此，如图7所示，MHSA的大部分运行时间被reshape和归一化等内存绑定操作占用[45, 46, 47, 27]。通过最小化内存绑定操作的使用或将其应用于较少的输入通道，SHSA模块可以充分利用GPU/CPU的计算能力。

2.4 单头视觉Transformer

单头视觉Transformer（SHViT）架构的概述如图5所示。给定输入图像，我们首先对其应用四个3×3步长卷积层。与标准ViT模型[31, 1]使用的步长16的16×16卷积token化相比，我们的重叠patchify stem可以提取更好的局部表示[32, 10, 27]。然后，token通过三个阶段的堆叠SHViT块进行分层表示提取。SHViT块由三个主要模块组成（见图5）：用于局部特征聚合或条件位置嵌入的深度卷积（DWConv）层[48, 49]，用于建模全局上下文的单头自注意力（SHSA）层，以及用于通道交互的前馈网络（FFN）。FFN中的扩展比设置为2。DWConv和SHSA的结合以计算和内存高效的方式捕获局部和全局依赖关系。基于第2.2节的发现，我们在第一阶段不使用SHSA层。为了在不丢失信息的情况下减少token，我们使用了一个高效的下采样层，该层由两个阶段1块组成，中间放置了一个倒置残差块[50, 27, 51]（步长2）。最后，使用全局平均池化和全连接层输出预测。

除了上述操作符外，归一化和激活层在决定模型速度方面也起着至关重要的作用。我们仅在SHSA层使用层归一化[44]，同时将批归一化（BN）[52]集成到其余层中，因为BN可以合并到其相邻的卷积或线性层中。我们还使用ReLU [53]激活而不是其他复杂的替代方案[54, 55, 51]，因为它们在各种推理部署平台上要慢得多[56, 7, 27]。

我们构建了四个具有不同深度和宽度设置的SHViT变体。由于大尺寸的patch嵌入和单头设计，我们可以使用比以前的高效模型更多的通道和块。模型规格见表1。

3. 实验

3.1 实现细节

我们在ImageNet-1K [23]上进行图像分类，该数据集包括128万张训练图像和5万张验证图像，涵盖1000个类别。所有模型均使用AdamW [57]优化器从头开始训练300个epoch，学习率为10^-3，总批量大小为2048。我们使用余弦学习率调度器[58]，并进行5个epoch的线性预热。SHViT-S1到S4的权重衰减分别设置为0.025/0.032/0.035/0.03。为了公平比较，我们遵循[31]中提出的相同数据增强方法，包括Mixup [59]、随机擦除[60]和自动增强[61]。对于384^2和512^2分辨率，我们微调模型30个epoch，权重衰减为10^-8，学习率为0.004。此外，我们评估了各种硬件平台上的吞吐量性能。我们在Nvidia A100上测量GPU吞吐量，批量大小为256。对于CPU和CPU_ONNX，我们在Intel(R) Xeon(R) Gold 5218R CPU @ 2.10GHz处理器上评估运行时，批量大小为16（使用单线程）。对于CPU_ONNX，我们将模型转换为ONNX [62]运行时格式。移动延迟使用iPhone 12（iOS版本16.5）测量。我们使用CoreML工具[63]导出模型（批量大小设置为1），并报告1000次运行的中位数延迟。我们还验证了我们的模型作为COCO [25]上目标检测和实例分割的高效视觉骨干，分别使用RetinaNet [64]和Mask R-CNN [26]。所有模型均在mmdetection库[65]上按照[40]的1×计划（12个epoch）进行训练。

3.2 SHViT在ImageNet-1K分类上的表现

如图1、表2和表4所示，我们将单头视觉Transformer（SHViT）与最先进的模型进行了比较。比较结果清楚地表明，我们的SHViT在各种设备上实现了更好的准确性与吞吐量/延迟的权衡。

与高效CNN的比较。 我们的SHViT-S1在A100 GPU和Intel CPU上保持相似速度的同时，比MobileNetV3-Small [51]的准确性高5.4%。与ShuffleNetV2 ×2.0 [66]相比，SHViT-S2在A100 GPU和Intel CPU上分别实现了2.2倍和2.5倍的速度提升，同时性能略优。此外，当转换为ONNX运行时格式时，我们的模型快3.8倍。与最近的FasterNet-T1 [30]相比，SHViT-S3不仅准确性高1.2%，而且运行速度更快：在A100 GPU上快15.1%，在Intel CPU上快32.4%。值得注意的是，在Top-1准确性为79.1-79.4%时，我们的模型在A100 GPU/Intel CPU上分别比EfficientNet-B1 [24]快2.9倍/3.3倍，比FastViT-T12 [7]快3.4倍/5.5倍，比MobileOne-S4 [56]快2.7倍/1.8倍。当以内存高效的方式利用最少的注意力模块时，ViT仍然可以像高效CNN一样展示快速的推理速度。

与高效ViT和混合模型的比较。 我们的SHViT-S1在A100 GPU和Intel CPU上的吞吐量分别比EfficientViT-M2 [27]高10%和42%，同时性能显著提高（70.8% → 72.8%）。SHViT-S3与PoolFormer-S12 [28]的准确性相似，但使用了3倍少的FLOPs，在A100 GPU上快3.8倍，在ONNX模型上快6.1倍。值得注意的是，我们的SHViT-S4比最近的EdgeViT-XS [17]准确性高1.9%，同时在A100 GPU上快2.6倍，在Intel CPU上快1.7倍，在ONNX实现上快4.2倍。如上所示，当转换为ONNX格式时，我们的模型与最近的SOTA模型相比表现出显著的性能提升。这种提升主要是因为我们的单头设计减少了reshape操作，这些操作在ONNX运行时中常常导致开销。

总结来说，上述结果表明，我们提出的内存高效宏观设计对速度-准确性权衡的影响比高效注意力变体或高度简单的操作（如池化）更为显著。

高分辨率微调。 根据[27]，我们还对SHViT-S4进行了高分辨率微调。与最先进的EfficientViT-M5_{512} [27]相比，我们的SHViT-S4_{384}在较低分辨率下训练时仍表现出竞争性性能。此外，SHViT-S4_{384}在A100 GPU上快77.4%，在Intel CPU上快55.6%，在ONNX运行时格式上快3.6倍。此外，SHViT-S4_{512}在A100 GPU上的吞吐量为3957 images/s，达到了82.0%的top-1准确性，展示了其在各种输入尺寸下的有效性。

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

表4. 移动延迟比较。括号中的结果是使用蒸馏[31]训练的。

移动延迟评估。 我们还在表4中验证了我们的模型在移动设备上的有效性。与高效模型EfficientNet-B1 [24] / MobileOne-S4 [56]相比，我们的SHViT-S4在iPhone 12设备上运行速度快0.2 ms / 0.1 ms，同时准确性相似。SHViT-S4还在针对移动延迟高度优化的模型中表现出竞争性性能，表明其在各种推理平台上的一致性性能。图8中的进一步结果显示，我们的模型显著优于最近的模型FastViT [7]和EfficientFormer [8]，尤其是在高分辨率下。在低分辨率下，SHViT-S4稍慢，但在1024 × 1024分辨率下，我们的模型比FastViT和EfficientFormer分别降低了34.4%和69.7%的延迟。这些结果源于宏观和微观设计中增加的内存效率。

3.3 SHViT在下游任务上的表现

在表5中，我们使用两个框架评估了SHViT的迁移能力：1）用于目标检测的RetinaNet [64]，2）用于实例分割的Mask R-CNN [26]。

目标检测。 SHViT-S4在移动设备上比MobileNetV3 [51]快2.3倍，并且AP高出8.9。与MobileFormer [19]相比，我们的模型在A100 GPU和移动设备上分别快3.2倍和8.2倍，同时表现出更好的性能。

实例分割。 在使用Mask R-CNN的COCO实例分割任务中，SHViT-S4在A100 GPU和移动设备上分别比EfficientNet-B0 [24]快1.9倍和1.2倍，同时AP^{mask}高出4.1。与PoolFormer-S12 [28]相比，我们的模型在A100 GPU和移动设备上分别快4.3倍和2.1倍，同时AP^{mask}高出1.3。

4. 相关工作

近年来，许多研究致力于设计高效的卷积神经网络（CNN）架构，以在资源受限的设备上实现快速推理。这些工作主要集中在减少计算复杂度和内存占用上，同时保持高准确性。例如，MobileNets [71]使用深度可分离卷积来减少计算量，ShuffleNets [66, 72]使用通道混洗操作来增强特征交互，GhostNet [67]使用廉价的线性变换来生成更多特征，EfficientNet [24]使用复合缩放法则来平衡模型的深度、宽度和分辨率，许多工作[7, 56, 73]使用结构重参数化来加速推理。

即使在视觉Transformer（ViT）[1]领域，也有许多努力致力于高效设计，以加速在各种设备上的推理速度。一种有前途的方法是设计一种新的ViT架构，将CNN的局部先验集成进来。这种方法大多只在后期阶段引入注意力机制，从而在不增加大量计算开销的情况下高效提取全局信息[7, 8, 9, 10, 12]。相比之下，其他方法在单个token mixer [14, 43, 74]或逐块[19]中并行使用注意力和卷积，以结合丰富的特征集。另一种方法侧重于减少MHSA的计算复杂度[17, 18, 21, 69]。例如，MobileViTv2 [21]引入了具有线性复杂度的可分离自注意力机制（相对于token数量）。EdgeViT [17]将MHSA应用于子采样特征，以在成本效益高的方式下执行近似全空间交互。与上述方法不同，我们优先考虑以最小的空间冗余组织token，而不是高效地混合token。

此外，最近的工作[27, 34, 35, 36, 37, 38, 39, 75, 76, 77]表明，许多头以相似的方式工作，并且可以在不影响性能的情况下被修剪。EfficientViT [27]提出将注意力头与完整通道的不同分割一起输入，以提高注意力多样性。此外，[76]提出了多头相似性的正则化损失，而[78]探索了不同层之间的头相似性。与减少多头冗余不同，我们设计了单头配置的模块，这不仅固有地防止了多头冗余，还节省了计算成本。

5. 结论

在本工作中，我们研究了在许多现有模型中常见的空间和通道维度上的冗余。然后，我们提出了16×16的patch嵌入和3尺度分层表示以及单头自注意力来解决计算冗余。我们进一步提出了基于我们提出的宏观/微观设计的多功能SHViT，它在各种设备和视觉任务上实现了超快的推理速度和高性能。