导读
视觉 Transformer（ViT）在各种计算机视觉任务中展现了最先进的性能，但其高计算需求使其在资源有限的边缘设备上不切实际。本文提出了 MicroViT，这是一种轻量级的视觉 Transformer 架构，通过显著降低计算复杂度，同时保持高精度，针对边缘设备进行了优化。MicroViT 的核心是高效单头注意力（ESHA）机制，该机制利用分组卷积来减少特征冗余，并仅处理部分通道，从而减轻了自注意力机制的负担。MicroViT 采用多阶段 MetaFormer 架构进行设计，通过堆叠多个 MicroViT 编码器来提升效率和性能。在 ImageNet-1K 和 COCO 数据集上的全面实验表明，MicroViT 在保持与 MobileViT 系列相当精度的同时，显著提高了 3.6 倍的推理速度，并提高了 40%的效率，减少了能耗，使其适用于移动和边缘设备等资源受限的环境。
1. 引言
近年来，Transformer 在计算机视觉领域受到了广泛关注，并取得了显著成就。该领域的一个显著发展是视觉 Transformer（ViT）[1]的引入，它利用纯 Transformer 进行图像分类任务。继 ViT 之后，提出了多个模型以提高性能，在包括图像分类、目标检测和分割等在内的多种视觉任务中取得了有希望的结果[2]-[6]。
尽管 vanilla Vision Transformer（ViT）[1]表现强劲，但在处理 ImageNet 分类任务时，它需要介于 8500 万到 6.32 亿个参数。其前向传播需要大量的计算资源，导致推理速度缓慢，使其不适合许多特定应用。在内存、处理能力和电池寿命受限的低成本环境中，如移动和边缘设备上应用 Transformer 模型尤其具有挑战性[7]。因此，作者的工作重点在于构建一个轻量级且高效的深度学习模型，以降低计算需求和功耗，同时确保在边缘设备上实现快速推理和高性能。
多项研究尝试通过将视觉 Transformer 与卷积神经网络（CNNs）集成来降低其计算复杂度[8]、[9]。例如，MobileViT 在早期阶段使用 MobileNetV2 的倒置 Bottleneck 卷积块[10]，以降低 ViT 的计算复杂度。众多研究[11]-[13]指出，自注意力（SA）机制，尤其是在 Transformer 的空间特征混合器中，是最耗计算的资源。金字塔视觉 Transformer（PVT）通过应用空间缩减注意力（SRA）来缩短 Token 长度，从而降低了 SA 的二次复杂度（O（n²））。MobileViTv2 提出了可分离线性注意力[12]，以减轻 SA 的负担，而 EdgeNeXt[13]则采用转置 SA 来解决边缘设备实现的复杂度挑战。另一种方法，SHViT，提出多头自注意力（MHSA）可能在头之间出现特征冗余，并引入了仅处理四分之一图像 Token 或特征的单头注意力。然而，这些方法大多数都是独立开发的，没有考虑在边缘计算设备等有限资源环境中的功耗。
本文介绍了 Micro ViT，这是一种针对边缘设备部署优化的新型视觉 Transformer 模型。该模型通过使用高效单头注意力（ESHA）显著降低了计算复杂度和功耗，通过分组卷积和单头 SA 最小化特征冗余，仅处理整体通道的四分之一。MicroViT 建立在 MetaFormer 架构之上，确保在边缘设备上进行快速推理和低功耗使用，同时保持高精度，使其非常适合能量受限的环境。
2. 提出方法
MicroViT 新型视觉模型的主要思想是降低图像特征处理的计算复杂性和冗余。MicroViT 模型集成了高效单头注意力（ESHA）技术，该技术能够以低复杂度 SA 生成低冗余的特征图。
A. 高效单头注意力（ESHA）
ESHA 将局部和全局空间操作结合在一个块中，以实现高效的 Token 信息提取。与原始 SA 不同，原始 SA 使用线性或 PWConv 来形成 Query 、 Key 和 Value ，而 ESHA 使用基于核的分组卷积进行局部信息获取。在 SA 中， Token 经过缩放点积操作以检索全局特征上下文。
B. 总体架构
MicroViT 采用了三阶段的金字塔结构，起始阶段是一个 16×16 的茎干，包含四个 3×3 的卷积，以 16 倍的比例缩小特征。采用 MetaFormer [18]框架，它使用了两个残差块进行空间混合，随后通过残差 FFN （FFN）进行通道混合，具体内容如方程 5 所述。
MicroViT 模型采用一系列可分离卷积和一个残差全连接神经网络（FFN）进行图像块嵌入，第二阶段采用 32 倍降采样率，第三阶段采用 64 倍降采样率，并使用 3×3 图像块嵌入。在早期阶段，深度可分离卷积（DW convolution）作为空间混合器以满足更高的内存需求。最终阶段利用表 1 中概述的高效单头注意力（ESHA）机制。批归一化（Batch Normalization，BN）用于更好地与相邻的卷积层集成并减少 Reshape ，从而提高推理速度。该架构使用全局平均池化，随后通过全连接层进行特征提取和分类。
3. 结果
为了评估 MicroViT，使用了包含 128 万张训练图像和 5 万张验证图像，涵盖 1000 个类别的 ImageNet-1K 数据集[19]。遵循 DeiT 的训练方法[20]，模型在 224×224 的分辨率下训练了 300 个 epoch，初始学习率为 0.004，并使用了多种数据增强。使用了 AdamW 优化器[21]，批大小为 512，在三个 A6000GPU 上运行。
作者评估了模型在不同计算环境中的吞吐量，包括 GPU（RTX-3090）、CPU（Intel i5-13500）以及 Jetson Orin Nano 边缘设备。对于吞吐量，GPU 和 CPU 的批次大小为 256，而边缘设备使用 ONNX Runtime，批次大小为 64。为了提高推理性能，尽可能地将 BN 层与相邻层融合。在 Jetson Orin Nano 上，作者还检查了在 1000 张图像的延迟测试中，保持一致分辨率时的功耗和能耗。
A. ImageNet-1K 分类结果
表 2 展示了在 ImageNet-1K 数据集上，各种 MicroViT 变体与最先进（SOTA）模型的比较。评估重点在于模型的计算效率和准确性，突出了资源消耗与性能之间的权衡。
MicroViT-S1 在性能上优于传统的 CNN 模型，超越了 MobileNetV2-1.0 [10]和 Fasternet-T0 [12]，在 GPU 上的速度提升了 3.6 倍，在 CPU 上提升了 6.7 倍，同时保持了比 MobileNetV2-1.0 高 0.8 的准确率。此外，MicroViT-S2 在效率指标相似的情况下，超越了 EfficientFormerV2-S0 [15]和 EfficientViT-M4 [24]等移动 Transformer ，实现了 0.3%的更高准确率。在整个 MicroViT 模型系列中，CPU 的吞吐量非常稳定，特别是 MicroViT-S1 达到了 552img/s ，这比几个 EfficientViT 变体快了 8 倍，展示了 MicroViT 对高端 GPU 和 CPU 设置的适应性。
表 3 展示了 MicroViT 变体在 Edge 设备上使用 ONNX 与各种 SOTA 模型的表现。MicroViT-S1 的吞吐量达到 773img/s，在 Jetson Orin Nano 上高效管理大量图像。这超过了多个 SOTA 模型，如 MobileViTXS [8]和 EfficientFormer-V2-S0 [15]，使 MicroViT-S1 成为快速图像处理应用的理想选择。此外，它具有 9.1 毫秒的延迟，优于 MobileNetV2-1.0 [10]和 EdgeNeXt-XS [13]，支持实时应用。它消耗 2147 焦耳，实现了高能效，η=3.7。同样，MicroViT-S2 和 MicroViT-S3 在吞吐量和能源使用之间取得平衡，保持准确性，因此对于资源受限的边缘设备而言，它们具有比其他轻量级视觉 Transformer 更高的能效。
下一轮消融实验去除了分组卷积，显著增加了模型复杂度至 20.5M 个参数，同时计算量（GFLOPs）也有所增加。这导致吞吐量和效率（η=1.8）降低，但实现了最高的准确率。然而，这种变体消耗的能量更多，使其在优先考虑准确率而非效率的应用场景中更为理想。
消融研究表明，尽管空间降维可以提高吞吐量，但延迟会增加，从而导致效率下降。与 MobileViT 中的 vanilla attention 等其他注意力模型相比，分组卷积在 ESHA 中成功提高了效率，同时保持了较低的复杂度和能耗。
B. 目标检测结果
作者将 MicroViT-3 与高效的模型[10]、[24]、[27]在 COCO[22]目标检测任务上进行比较，并在表 4 中展示了结果。具体来说，MicroViT-3 在具有可比 Flops 的情况下，AP 值比 MobileNetV2[10]高出 7.7%。与 EfficientViT-M4 相比，作者的 MicroViT-3 使用了 46.8%更少的 Flops，同时 AP 值提高了 3.3%，这证明了其在不同视觉任务中的能力和泛化能力。
4. 结论
本文介绍了 MicroViT，这是一种针对边缘设备优化的新型轻量级视觉 Transformer 架构，考虑了计算能力和能效。通过采用高效单头注意力（ESHA）机制，MicroViT 在保持视觉任务中具有竞争力的准确率的同时，实现了计算复杂度和能耗的显著降低。在 ImageNet-1K 和 COCO 数据集上的大量实验表明，MicroViT 不仅提高了 3.6 倍的吞吐量和推理速度，而且在边缘设备上的效率和性能方面超过了多个 MobileViT 模型，效率提升了 40%。这些结果证实，MicroViT 是部署视觉 Transformer 在资源受限环境中的一种有前景的解决方案。未来的工作将探索对该架构的进一步优化以及在其他边缘计算任务中的更广泛应用。
参考
[1]. MicroViT: A Vision Transformer with Low Complexity Self Attention for Edge Device
END
作者：小书童
来源：集智书童
推荐阅读
- 碾压YOLO系 | Fast-COS横扫BDD100K/TJU双冠，边缘端实时推理，效率起飞
- 基于 Yolact 的检测优化，训练优化不损推理效率，解决边缘模糊问题 ！
- SGLang DP MLA 特性解读
- 推理模型新路线开源！与 DeepSeek 截然不同，抛弃思维链不用人类语言思考
欢迎大家点赞留言，更多 Arm 技术文章动态请关注极术社区嵌入式AI专栏欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。
