[论文解读] EfficientViT: Memory Efficient Vision Transformer with Cascaded Group Attention

Xinyu Liu, Houwen Peng+4 more|arXiv (Cornell University)|May 11, 2023

Advanced Memory and Neural Computing被引用 36

一句话总结

EfficientViT 引入了具有三明治布局和级联分组注意力的节省内存的视觉 Transformer，以减少内存访问和计算，在速度和准确性方面超越以往的高效模型。

ABSTRACT

Vision transformers have shown great success due to their high model capabilities. However, their remarkable performance is accompanied by heavy computation costs, which makes them unsuitable for real-time applications. In this paper, we propose a family of high-speed vision transformers named EfficientViT. We find that the speed of existing transformer models is commonly bounded by memory inefficient operations, especially the tensor reshaping and element-wise functions in MHSA. Therefore, we design a new building block with a sandwich layout, i.e., using a single memory-bound MHSA between efficient FFN layers, which improves memory efficiency while enhancing channel communication. Moreover, we discover that the attention maps share high similarities across heads, leading to computational redundancy. To address this, we present a cascaded group attention module feeding attention heads with different splits of the full feature, which not only saves computation cost but also improves attention diversity. Comprehensive experiments demonstrate EfficientViT outperforms existing efficient models, striking a good trade-off between speed and accuracy. For instance, our EfficientViT-M5 surpasses MobileNetV3-Large by 1.9% in accuracy, while getting 40.4% and 45.2% higher throughput on Nvidia V100 GPU and Intel Xeon CPU, respectively. Compared to the recent efficient model MobileViT-XXS, EfficientViT-M2 achieves 1.8% superior accuracy, while running 5.8x/3.7x faster on the GPU/CPU, and 7.4x faster when converted to ONNX format. Code and models are available at https://github.com/microsoft/Cream/tree/main/EfficientViT.

研究动机与目标

• 分析影响 Vision Transformer (ViT) 推理速度的因素，聚焦于内存访问、计算冗余和参数使用。
• 设计一个节省内存的 ViT 构建块，以在不牺牲准确性的前提下提升吞吐量。
• 开发一种新的注意力机制，减少冗余并提升头部特征多样性。
• 重新分配参数以提高关键模块的效率，同时裁剪不那么重要的组件。
• 证明 EfficientViT 在下游视觉任务上的可迁移性。

提出的方法

• 引入三明治布局块，其中一个受内存限制的 MHSA 被 FFN 层夹在中间，以减少内存瓶颈运算。
• 提出 Cascaded Group Attention (CGA)，通过对注意力头输入不同特征分割并在各头之间级联输出来增加多样性并降低计算。
• 应用 Taylor 结构化裁剪来识别重要通道并引导参数重新分配，在扩展关键模块宽度的同时收缩次要部分。
• 使用重叠的补丁嵌入以及包含 BN 全部阶段的三级分层架构，以提高硬件上的速度与实用性。
• 在 ImageNet-1K 上对 GPU/CPU/ONNX 的吞吐量和准确性进行评估，并测试对下游任务的可迁移性。

Figure 1 : Speed and accuracy comparisons between EfficientViT (Ours) and other efficient CNN and ViT models tested on an Nvidia V100 GPU with ImageNet-1K dataset [ 17 ] .

实验结果

研究问题

• RQ1 如何对 ViT 进行重新设计，以在推理阶段最小化内存瓶颈运算？
• RQ2 通过将注意力头输入分割的特征通道（CGA）是否能够在不牺牲准确性的前提下减少冗余并增加注意力多样性？
• RQ3 哪些参数重新分配策略能够在轻量化 ViT 中实现更好的速度-准确度权衡？
• RQ4 所提出的 EfficientViT 模块是否能推广到下游视觉任务并在真实硬件上部署？

主要发现

• EfficientViT-M5 在 Nvidia V100 上实现 77.1% 的 Top-1 准确率，吞吐量为 10,621 图像/秒，较 MobileNetV3-Large 在准确率上高出 1.9%，在吞吐量方面 GPU 提升 40.4%，CPU 提升 45.2%。
• EfficientViT-M2 达到 70.8% Top-1 精度，领先 MobileViT-XXS 1.8%，提供 5.8x GPU 和 3.7x CPU 的加速，ONNX 性能快 7.4x。
• EfficientViT-M4 在 ImageNet 上优于若干高效模型，具有更高吞吐量和具有竞争力的准确性（例如，与 EdgeViT-XXS 相比，GPU 提速 4.4x，CPU 提速 3.0x）。
• 用三明治布局替代内存瓶颈的 MHSA，降低内存时间消耗，同时增加基于 FFN 的通道通信。
• CGA 通过将注意力头分布到特征分割并级联输出来降低注意力计算，提高效率和准确性（消融实验显示优于 MHSA 的提升）。
• 参数重新分配，增加 V 通道宽度、减小 Q/K 维度以及 FFN 展开，带来可观的准确性和吞吐量提升。

Figure 2 : Runtime profiling on two standard vision transformers Swin-T and DeiT-T. Red text denotes memory-bound operations, i.e. , the time taken by the operation is mainly determined by memory accesses, while time spent in computation is much smaller.