Fast Vision Transformers with HiLo Attention

总结

提出了 HiLo Attention 机制：

• 该机制将自注意力层分为两部分：Hi-Fi（高频注意力） 和 Lo-Fi（低频注意力）。
• Hi-Fi 捕捉局部细节，通过在局部窗口内应用自注意力，减少了计算复杂度。
• Lo-Fi 捕捉全局结构，利用低频特征实现全局注意力，并通过平均池化减少计算开销。
• 通过这种分离方式，HiLo 在保持模型性能的同时，大幅度提高了计算效率，尤其是在高分辨率图像的处理上。
• 将所有注意力层替换为 HiLo Attention，并在每个 FFN 中加入 3×3 深度卷积以提高位置编码效率。

---

Abstract

Vision Transformers (ViTs) 引发了计算机视觉的最新和重大突破。它们的高效设计主要是由计算复杂度的间接度量来指导的，即FLOPs，然而与吞吐量等直接度量有明显的差距。因此，我们建议在目标平台上使用直接速度评估作为高效 ViT 设计原则。特别是，我们引入了 LITv2，这是一种简单有效的 ViT，它在一系列不同模型大小上以更快的速度优于现有的最先进方法。LITv2 的核心是一种新颖的自注意力机制，我们称之为 HiLo。HiLo的灵感来自于这样一种观点，即图像中的高频捕捉局部精细细节，低频聚焦于全局结构，而多头自我注意层忽略了不同频率的特征。因此，我们建议通过将头部分为两组来解开注意力层中的高/低频模式，其中一组通过每个局部窗口内的自注意力对高频进行编码，另一组通过在每个窗口的平均值汇集的低频键和值以及输入特征图中的每个查询位置之间执行全局注意力来编码低频。得益于两组的高效设计，我们表明HiLo通过全面基准测试gpu和cpu上的FLOPs、速度和内存消耗，优于现有的注意机制。例如，HiLo 比空间缩减注意力快 1.4 倍，比 CPU 上的局部窗口注意力快 1.6 倍。由 HiLo 提供支持，LITv2 作为主流视觉任务的强大主干，包括图像分类、密集检测和分割。代码可在 https://github.com/ziplab/LITv2 获得。

1 Introduction

现实世界的应用程序通常需要模型在有限的计算预算(如无人机和自动驾驶)下具有最佳的速度和准确性权衡。这激发了高效视觉转换器(ViT)设计的大量工作，如PVT[51]、Swin[32]和Focal Transformer[60]等。为了衡量计算复杂度，最近ViT设计中广泛采用的度量是浮点操作的数量，即FLOPs。然而，FLOPs 是一个间接度量，不能直接反映目标平台上的真实速度。例如，Focal-Tiny 在 GPU 上比 Swin-Ti 慢得多，尽管它们的 FLOPs 具有可比性。

一般来说，最近vit中间接度量(FLOPs)和直接度量(速度)之间的差异可以归因于两个主要原因。首先，虽然自注意力在低分辨率特征图上是有效的，但由于内存访问成本的密集 [34]，内存和时间的二次复杂度使得高分辨率图像上的速度较慢，其中从片外 DRAM 获取数据可能很耗时。其次，ViTs中的一些高效注意机制具有较低的理论复杂度保证，但由于特殊的操作不是硬件友好的，或者不能并行化，如多尺度窗口分区[60]、递归[44]和扩张窗口[20]。

图 1：HiLo 注意力的框架。Nh 是指该层自注意力头的总数。α 表示高频/低频头的拆分率。最好以彩色观看。

通过这些观察结果，在本文中，我们建议通过直接度量(即吞吐量)来评估ViT，而不仅仅是FLOPs。基于这一原理，我们引入了LITv2，这是一种新颖而准确的视觉转换器，在标准基准上优于大多数最先进的(SoTA) ViT，同时在gpu上实际上更快。LITv2 在 LITv1 [36] 上进行 bulit，这是一个简单的 ViT 基线，它在早期阶段删除所有多头自注意力层 (MSA)，同时在后面的阶段应用标准 MSA。得益于这种设计，LITv1 比 ImageNet 分类中的许多现有工作更快，因为早期 MSA 没有计算成本，而后面的 MSA 只需要处理下采样的低分辨率特征图。然而，标准的MSA在高分辨率图像上仍然存在巨大的计算成本，特别是对于密集预测任务。

为了解决这个问题，我们提出了一种新颖的高效注意力机制，称为 HiLo。HiLo 的动机是自然图像包含丰富的频率，其中高/低频在编码图像模式（即局部精细细节和全局结构）方面发挥着不同的作用。典型的MSA层在不考虑不同底层频率的特征的情况下&