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层在不考虑不同底层频率的特征的情况下，在所有图像补丁上强制相同的全局注意。这促使我们提出将MSA层分离为两条路径，其中一条路径通过局部自我注意编码高频交互，具有相对高分辨率的特征映射，而另一条路径通过全局注意和下采样的特征映射对低频交互进行编码，从而大大提高了效率。

具体来说，HiLo 采用两个有效的注意力来解开特征图中的高/低频。如图 1 所示，在上路径中，我们将几个头分配给高频注意力 (Hi-Fi)，通过局部窗口自注意力（例如 2 × 2 窗口）来捕获细粒度的高频，这比标准 MSA 更有效。实现低频注意 (Lo-Fi) 的较低路径首先将平均池化应用于每个窗口以获得低频信号。然后，我们为 Lo-Fi 分配剩余的头来模拟输入特征图中每个查询位置与每个窗口中平均池化的低频键和值之间的关系。得益于键和值的长度减少，Lo-Fi 也实现了显着的复杂性降低。最后，我们将细化的高频/低频特征连接起来，并将结果输出转发到后续层。由于 Hi-Fi 和 Lo-Fi 都没有配备耗时的操作，例如扩展窗口和递归，HiLo 的总体框架在 CPU 和 GPU 上都很快。我们通过综合基准表明，HiLo 在性能、FLOPs、吞吐量和内存消耗方面比现有的注意力机制具有优势。

此外，由于不同图像分辨率的插值，我们发现 LITv1 中的固定相对位置编码显着降低了其在密集预测任务上的速度。为了提高效率，我们建议在每个FFN中采用一个带有零填充的3 × 3深度卷积层，从零填充[27]中合并隐式学习的位置信息。此外，3 × 3卷积滤波器同时有助于扩大 LITv1 中早期多层感知器 (MLP) 块的感受野。最后，我们在 ImageNet、COCO 和 ADE20K 上进行了大量实验，以评估 LITv2 的性能。与 SoTA 模型的综合比较表明，我们的架构以更快的吞吐量实现了具有竞争力的性能，这使得 ViT 更适合为现实场景运行低延迟应用程序。

2 Related Work

Vision Transformers.视觉变形金刚是神经网络，它将自我注意机制应用于计算机视觉任务。在 [18] 中，Dosovitskiy 等人。提出了一种用于图像分类的 ViT，它继承了自然语言处理 (NLP) 任务中标准 Transformer [48] 的类似架构。从那时起，后续工作被提出通过结合更多的卷积层来改进ViT[54,61]，引入金字塔特征映射[51,32]，增强局部性[62]，并使用神经结构搜索(NAS)自动搜索性能良好的架构[5,3]。其他一些人也寻求令牌修剪来加速VITS[37]的推理速度，或者将VIT应用于低级视觉任务[47]。与现有的作品相比，本文重点研究了计算机视觉(CV)任务的通用基于vitc的主干，旨在在保持竞争性能的同时，在gpu上实现更好的效率。

Efficient attention mechanisms.有效的注意力机制旨在减少标准 MSA 的二次复杂度。NLP 的现有努力可以大致分为低秩分解 [50]、核化 [28, 39]、内存 [40] 和稀疏机制 [10]。然而，简单地采用这些方法通常在 CV 任务中表现不佳 [32, 63]。在CV中，具有代表性的高效自我注意机制包括空间缩减注意(SRA)[51]、局部窗口注意[32,26]和Twins注意[12]。然而，他们只关注同一层的局部和全局注意力。为了解决这个问题，TNT [21] 引入了额外的全局标记和 MixFormer [6] 将局部窗口注意力与深度卷积层混合。其他一些注意力机制同时考虑两者，例如 Focal[60] 和 QuadTree [44]。然而，由于操作效率低，操作不友好，不能反映在FLOPs中(例如，多尺度窗口分区、递归)，即使与标准MSA相比，它们在gpu上也很慢。为此，所提出的 HiLo 注意力同时在同一 MSA 层捕获丰富的局部全局信息，并且与现有方法相比具有更高的内存效率。

Frequency domain analysis in vision.CV中的频域分析在文献中得到了很好的研究。根据[13，16]，图像中的低频通常捕获全局结构和颜色信息，而高频包含对象的精细细节(例如，尖锐的边缘)。基于这一见解，已经提出了大量的图像超分辨率解决方案[66,19]，泛化[25]，图像重新缩放[56]和神经网络压缩[59,7]。此外，Octave卷积[9]针对卷积层，提出分别在高分辨率/低分辨率特征图上局部应用卷积。与之不同的是，所提出的 HiLo 是一种新颖的注意力机制，它同时捕获与自注意力的局部和全局关系。

3 Background

Multi-head self-attention.Transformers 建立在多头自注意力之上，它能够捕获不同位置的标记的远程关系。具体来说，设 X ∈ RN ×D 是输入序列到标准 MSA 层，其中 N 是输入序列的长度，D 表示隐藏维度的数量。每个自注意力头使用 X 的线性变换计算查询 Q、键 K 和值 V 矩阵，

Transformer blocks. [18]中描述的标准视觉转换器由一个补丁嵌入层、几个块和预测头组成。设 l 是块的索引。然后每个块包含一个 MSA 层和一个位置前馈网络 (FFN)，可以表示为

FFN 一般由两个线性变换（全连接层）和一个非线性激活函数（通常是ReLU）组成。

其中 LN 表示 LayerNorm [2]，FFN 由两个 FC 层组成，中间具有 GELU [24] 非线性。最近关于 ViT 的工作提出将块划分为几个阶段（通常为 4 个阶段），以生成用于密集预测任务的金字塔特征图。此外，为了降低早期高分辨率特征图的计算成本，将式(4)中的MSA替换为有效的替代方案，如SRA[51]和W-MSA[32]。

Bottlenecks of LITv1.最近的研究表明，模型早期阶段的 MSA 层仍然专注于局部模式 [14]。在相同的观察下，LITv1 [36] 在后期应用标准 MSA 时删除了所有早期的 MSA（即在每个块中排除等式（4））。与PVT[51]和Swin[32]相比，该设计原则在ImageNet上取得了更好的性能。然而，LITv1 的速度仍然有两个主要瓶颈：1）给定一个高分辨率图像，后期的标准 MSA 仍然会导致巨大的计算成本。2)固定的相对位置编码[32]在处理不同的图像分辨率时大大减慢了速度。这是因为为每个不同的图像分辨率插值固定大小的位置编码。在下一节中，我们描述了一种具有零填充位置编码的新型注意力机制，以全面加速 LITv1。

4 Method

4.1 HiLo Attention

我们建议分别在注意力层的特征图中分别处理高频/低频。我们将新的注意力机制命名为 HiLo，如图 1 所示。本质上，低频注意力分支 (Lo-Fi) 是捕获输入（图像/特征）的全局依赖关系，它不需要高分辨率的特征图，但需要全局注意力。另一方面，高频注意力分支 (Hi-Fi) 是捕获精细细节的局部依赖，这需要高分辨率的特征图，但可以通过局部注意力来完成。接下来，我们将详细描述这两个注意力。

High-frequency attention.直观地说，由于高频编码对象的局部细节，因此在特征图上应用全局注意力可能是冗余且计算成本高的。因此，我们建议设计 Hi-Fi 来捕获具有局部窗口自注意力（例如 2 × 2 窗口）的细粒度高频，这节省了显着的计算复杂性。此外，我们在 Hi-Fi 中使用简单的非重叠窗口分区，与窗口移位 [32] 或多尺度窗口分区 [60] 等耗时的操作相比，它更硬件友好。

Low-frequency attention.最近的研究表明，MSA 中的全局注意力有助于捕获低频 [38]。然而，直接将 MSA 应用于高分辨率特征图需要大量的计算成本。由于平均是一个低通滤波器[49]，Lo-Fi首先将平均池化应用于每个窗口，得到输入X中的低频信号。接下来，平均池化的特征图被投影到键 K ∈ R N/s^2×Dh 和值 V ∈ R N/s^2×Dh 中，其中 s 是窗口大小。Lo-Fi 中的查询 Q 仍然来自原始特征图 X。然后我们应用标准注意力来捕获特征图中丰富的低频信息。请注意，由于 K 和 V 的空间缩减，Lo-Fi 同时降低了等式的复杂性。 (1) 和等式。 (2)。

Head splitting. 头部分配的一个简单解决方案是分配与标准 MSA 层相同数量的头部 Hi-Fi 和 Lo-Fi。然而，加倍头会导致更多的计算成本。为了获得更好的效率，HiLo 将 MSA 中相同数量的头分为两组，拆分比为 α，其中 (1 − α)Nh 头将用于 Hi-Fi，其他 αNhhead 用于 Lo-Fi。通过这样做，由于每个注意力的复杂度都低于标准 MSA，HiLo 的整个框架保证了低复杂度并确保 GPU 上的高吞吐量。此外，头部拆分的另一个好处是可学习参数 Wo 可以分解为两个较小的矩阵，这有助于减少模型参数。最后，HiLo 的输出是每个注意力的输出的串联

图2:Hi-Fi和Lo-Fi在不同图像分辨率和相同数量的头部下的FLOPs比较(图a和b)。更大的窗口大小有助于 HiLo 在高分辨率图像上实现更好的效率（图 c）。

Complexity Analysis. 不失一般性，我们假设 Hi-Fi 和 Lo-Fi 有相同数量的头（即 α = 0.5），特征图的宽度和高度相等。如图2-(a)和(b)所示，在小输入图像分辨率和s的小值(如s = 2)下，Hi-Fi和Lo-Fi都相当有效。然而，在更高的分辨率下，Lo-Fi将导致巨大的计算成本，因为它在Eq.(2)中的N方面仍然具有二次复杂度。在这种情况下，略微增加（例如，s = 4）有助于 Lo-Fi 在保持准确性的同时实现更好的效率。将这两个注意力组合在一起，更大的窗口大小也有助于 HiLo 的总体框架，以减少高分辨率图像上的更多 FLOP，如图 2-© 所示。因此，我们提出了一种将 HiLo 应用于现有框架的实用指南：增加窗口大小以获得更好的高分辨率图像效率。我们进一步在第 5.2 节中表明，这一原则有助于 LITv2 在下游任务（例如密集对象检测）上实现更好的速度和准确性权衡。

4.2 Positional Encoding

由于其排列不变的特性，位置编码对于自我注意是必不可少的。在 LITv1 中，后面的 MSA 采用与 Swin [32] 相同的相对位置编码 (RPE) 方案。与使用绝对位置编码[32]相比，该方法在ImageNet上的Top-1精度上显著提高了0.7%的Swin。然而，在密集预测任务上，必须针对不同的图像分辨率对固定的RPE进行插值，这大大减慢了LITv1的训练/推理速度。作为最近的一项研究[27]表明，位置信息可以从cnn中的零填充隐式学习，我们建议在每个FFN中采用一层3×3深度卷积层，用零填充来代替耗时的RPE。值得注意的是，由于消除了早期的MSAs，LITv1的早期块只剩下FFNs，导致1 × 1的小接受域。为此，我们在第5.4节中表明，每个FFN中采用的3 × 3卷积滤波器也通过在早期阶段同时扩大接受域来改善LITv2。

4.3 Model Architecture

LITv2 有三种变体：LITv2-S、LITv2-M 和 LITv2-B，分别对应于 LITv1 中的小、中和基本设置。为了公平比较，我们保持网络宽度和深度与 LITv1 相同。整体修改只是分两步：1）在每个 FFN 中添加一个零填充的深度卷积，并删除所有 MSA 中的所有相对位置编码。2) 用提出的 HiLo 注意力替换所有注意力层。详细的架构配置可以在补充材料中找到。

5 Experiment

在本节中，我们通过实验来验证所提出的LITv2的有效性。按照惯例[51,32,12,60]，我们在三个任务上对LITv2进行了实验，包括ImageNet-1K[43]上的图像分类、COCO[31]上的目标检测和实例分割以及ADE20K上的语义分割[65]。

表 1：ImageNet-1K 上的图像分类结果。默认情况下，FLOPs、吞吐量和内存消耗是根据分辨率224 × 224来衡量的。我们报告了批处理大小为64的吞吐量和训练/测试时间内存消耗。吞吐量在一个NVIDIA RTX 3090 GPU上进行了测试，平均超过30次运行。ResNet 结果来自“ResNet Stikes Back”[53]。“↑ 384”表示模型在 384 × 384 分辨率下进行了微调。 “OOM”表示“内存不足”。

5.1 Image Classification on ImageNet-1K

5.2 Object Detection and Instance Segmentation on COCO

5.3 Semantic Segmentation on ADE20K

5.4 Ablation Stud

在本节中，我们提供了LITv2的消融研究，包括与其他有效注意力变体的比较，α在HiLo中的影响，以及结构修改的影响。默认情况下，吞吐量和内存消耗是在一个批量大小为64的RTX 3090 GPU上测量的，分辨率为224 × 224。

Comparing HiLo with other attention mechanisms.基于LITv2-S，我们比较了HiLo与ImageNet-1K上其他高效注意机制的性能，包括PVT中的空间缩小注意(SRA) [51]， Swin中的基于移位窗口的注意(W-MSA)[32]和Twins中的局部和全局交替注意(T-MSA)[12]。在我们的实现中，我们直接用每个比较的方法替换HiLo。结果如表4所示。一般来说，HiLo在实现更好的性能和更快的速度的同时减少了更多的flop。此外，在图3中，我们提供了基于不同图像分辨率的更多注意力机制的综合基准，包括Focal[60]、QuadTree[44]和Performer[11]。由于并行性较弱，它们甚至比在gpu上使用标准msa还要慢。与它们相比，HiLo在FLOPs、吞吐量和内存消耗方面取得了具有竞争力的成绩。此外，我们基于ADE20K和语义FPN进行了实验，结果表明HiLo在下游密集预测任务上取得了比其他注意机制更高的性能增益。

表5:基于LITv1-S的架构修改效果。“卷积神经网络”意味着我们在每个FFN中添加一层3 × 3深度卷积层。“RPE”为相对位置编码[32]。

图4:α对LITv2-S的影响。

图5:LITv2-B中Hi-Fi和Lo-Fi 8个输出通道的频率幅值(14 × 14)。该量级在100个样本上取平均值。颜色越浅，星等越大。离中心越近的像素意味着频率越低。可视化代码可以在补充材料中找到。

Effect of α.如图4所示，由于在分辨率为224 × 224、窗口大小为2的情况下，low - fi的复杂度低于Hi-Fi，因此当我们为low - fi分配更多的磁头时，更大的α有助于减少更多的FLOPs。此外，我们发现HiLo在α = 0时表现不佳，在这种情况下只剩下Hi-Fi, HiLo只关注高频。我们推测低频在自我注意中起重要作用。对于其他α值，我们发现性能差异在0.2%左右，其中α = 0.9达到最佳性能。然而，值得注意的是，尽管纯Lo-Fi分支(α = 1.0)在ImageNet-1K上可以取得有竞争力的结果，但高频信号在捕获精细物体细节方面起着重要作用，这对于语义分割等密集预测任务尤为重要。例如，当α = 0.9时，基于LITv2-S的语义FPN比使用α = 1.0时(43.7%)获得了更多的性能增益(+0.6%)。

Effect of architecture modifications.基于 LITv2-S，我们探索了架构修改的效果。如表 5 所示，受益于早期阶段扩大的感受野，深度卷积的采用提高了 ImageNet 和 COCO 的性能。接下来，通过删除相对位置编码，我们显着提高了密集预测任务的 FPS，在两个数据集上的性能略有下降。此外，由于深度卷积通过零填充[27]对位置信息进行编码，因此与之前的工作[32]相比，RPE的消除不会导致显著的性能下降。最后，从 HiLo 中受益，我们在 ImageNet 和 COCO 上实现了更高的模型效率提升。

Spectrum analysis of HiLo. 在图 5 中，我们通过将快速傅里叶变换 (FFT) 分别应用于 Hi-Fi 和 Lo-Fi 注意力的输出特征图来可视化频率分量 [42] 的大小。可视化表明 Hi-Fi 捕获了更多的高频，Lo-Fi 主要关注低频。这与我们的目标是在单个注意力层解开特征图中的高频和低频的目标密切相关。

表 6：LITv2-S 与其他最近 ViT 在不同 GPU 上的速度和性能比较。所有吞吐量结果均在 30 次运行中取平均值，总批量大小为 64，图像分辨率为 224 × 224 在一个 GPU 卡上。我们还报告了 ImageNet-1K 上的 Top-1 准确度。

图 6：基于单个注意力层和 14×14 特征图在 CPU 和 GPU 上与更多注意力机制的吞吐量比较。

6 Conclusion and Future Work

在本文中，我们介绍了 LITv2，这是一种在 GPU 上快速速度的新型高效视觉转换器主干，并在 ImageNet 和下游任务上的表现优于大多数 SoTA 模型。我们还提出了 HiLo attention，即 LITv2 的核心，它有助于提高效率，尤其是在高分辨率图像上。在竞争性能下，HiLo 在 FLOP、吞吐量和内存消耗方面比现有的注意力机制有很大的优势。未来的工作可能包括结合卷积词干[57]和重叠补丁嵌入[52]以获得更好的性能，或者在语音识别和视频处理等更多任务上扩展HiLo。

Limitations and societal impact.HiLo 采用头部分割比率将不同数量的头部分配给 Hi-Fi 和 Lo-Fi。在我们的实验中，这个比率是由 ImageNet 上的网格搜索确定的（即 α = 0.9）。然而，不同的任务可能对高频和低频有不同的重要性。因此，α 的最佳值是特定于任务的，需要手动设置。