[论文解读] Fast Vision Transformers with HiLo Attention

Zizheng Pan, Jianfei Cai+1 more|arXiv (Cornell University)|May 26, 2022

Advanced Neural Network Applications被引用 112

一句话总结

本文介绍了 LITv2，一种使用 HiLo 注意力的快速 Vision Transformer 主干网络，它将高频和低频模式分离，以实现局部高频与全局低频交互，在各任务中实现更快的吞吐量与具有竞争力的准确率。

ABSTRACT

Vision Transformers (ViTs) have triggered the most recent and significant breakthroughs in computer vision. Their efficient designs are mostly guided by the indirect metric of computational complexity, i.e., FLOPs, which however has a clear gap with the direct metric such as throughput. Thus, we propose to use the direct speed evaluation on the target platform as the design principle for efficient ViTs. Particularly, we introduce LITv2, a simple and effective ViT which performs favourably against the existing state-of-the-art methods across a spectrum of different model sizes with faster speed. At the core of LITv2 is a novel self-attention mechanism, which we dub HiLo. HiLo is inspired by the insight that high frequencies in an image capture local fine details and low frequencies focus on global structures, whereas a multi-head self-attention layer neglects the characteristic of different frequencies. Therefore, we propose to disentangle the high/low frequency patterns in an attention layer by separating the heads into two groups, where one group encodes high frequencies via self-attention within each local window, and another group encodes low frequencies by performing global attention between the average-pooled low-frequency keys and values from each window and each query position in the input feature map. Benefiting from the efficient design for both groups, we show that HiLo is superior to the existing attention mechanisms by comprehensively benchmarking FLOPs, speed and memory consumption on GPUs and CPUs. For example, HiLo is 1.4x faster than spatial reduction attention and 1.6x faster than local window attention on CPUs. Powered by HiLo, LITv2 serves as a strong backbone for mainstream vision tasks including image classification, dense detection and segmentation. Code is available at https://github.com/ziplab/LITv2.

研究动机与目标

• 通过在目标硬件上直接进行速度测量来评估 ViT，而不仅仅是 FLOPs，作为动机。
• 提出 HiLo 注意力，以分别建模高频局部细节和低频全局结构。
• 展示基于 HiLo 的 LITv2 在 ImageNet、COCO 与 ADE20K 上实现更快的吞吐量并具备竞争力的准确性。
• 证明相较于最先进的 ViT 以及某些 CNN，内存使用减少、效率更优的优势。

提出的方法

• 引入将注意力头分成高频（Hi-Fi）和低频（Lo-Fi）组的 HiLo 注意力。
• Hi-Fi 在高分辨率特征图上使用局部窗口自注意力（如 2x2 窗口）。
• Lo-Lo 应用对每个窗口进行平均池化以获得低频信号，并使用来自这些池化信号的降采样键和值进行注意力计算。
• 按比例 α 划分注意力头，其中 (1−α)N_h 给 Hi-Fi，αN_h 给 Lo-Fi，然后将输出拼接。
• 在 FFN 中用深度卷积 3x3 替代固定的相对位置编码，以注入位置信息并增大感受野。
• 在 LITv2 中对三种模型尺寸（S/M/B）采用 HiLo，并在 ImageNet-1K 上训练，在下游任务上评估。

实验结果

研究问题

• RQ1 HiLo 注意力在保持与标准自注意力及其他高效注意力方法相近的准确性同时，是否能在 CPU/GPU 上带来加速？
• RQ2 HiLo 的高/低频分离如何影响在高分辨率和密集预测任务上的性能？
• RQ3 结构性修改（如卷积 FFN、移除相对位置编码）是否能进一步提升在视觉任务上的速度和准确性？
• RQ4 在不同工作负载和分辨率下，最优的头部分割比 α 与窗口大小 s 是多少？
• RQ5 与 SoTA ViT 相比，基于 HiLo 的主干在下游检测/分割任务中是否提升吞吐量和内存使用？

主要发现

• 搭载 HiLo 的 LITv2 在 Top-1 准确率上优于 LITv1，同时 FLOPs 更少，吞吐量更快。
• 在 GPU 上测量时，HiLo 相对于 LITv1 在 ImageNet-1K 的推理加速分别为 13%（S）、27%（M）、36%（B）。
• 在可比较的 FLOPs 下，LITv2-S 在 ImageNet-1K 上的吞吐量比 PVT-S 和 Twins-PCPVT-S 更高。
• 在 COCO 目标检测/实例分割上，LITv2 主干比 ResNet 主干具有更高的 AP，并且在 SoTA ViT 中具有竞争力或更优的结果，同时 FPS 和 FLOPs 更有利。
• 在 ADE20K 语义分割上，LITv2 变体在 mIoU 和 FPS方面优于若干基线（例如 LITv2-S/LITv2-M/LITv2-B 在 mIoU 上超过 Swin-Ti/Swin-S/Swin-B，同时时速提升）。
• HiLo 在 ImageNet-1K 及下游任务中，一致地降低 FLOPs 并提升速度，相较于其他高效注意力机制（SRA、W-MSA、T-MSA）。