LITv2来袭 | 使用HiLo Attention实现高精度、快速度的变形金刚，下游任务均实时 - 极术社区

近两年来，ViT 在计算机视觉领域的取得了很多重大的突破。它们的高效设计主要受计算复杂度的间接度量（即 FLOPs）指导，但与吞吐量等直接度量存在明显差距。因此，LIT作者建议使用目标平台上的直接速度评估作为高效 ViT 的设计原则。
因此本文提出了 LITv2，一种简单而有效的 ViT，它以更快的速度在各种不同模型大小的频谱上与现有的最先进方法相比表现出色。LITv2 的核心是一种新颖的Self-Attention机制，这里称之为 HiLo。HiLo 的灵感来源于图像中的高频捕捉局部细节，低频聚焦于全局结构，而Multi-Head Self-Attention忽略了不同频率的特征。因此，建议通过将 Head 分解成两组来解耦注意力层中的高/低频模式，其中一组通过每个局部窗口内的Self-Attention对高频进行编码，另一组通过执行注意力来建模来自每个窗口的平均池化的低频key与输入特征图中的每个query位置之间的全局关系。
受益于前面2组的高效设计，通过对 FLOPs、GPU 上的速度和内存消耗进行全面基准测试，证明 HiLo 优于现有的注意力机制。在 HiLo 的支持下，LITv2 是主流视觉任务的强大支柱，包括图像分类、密集检测和分割。

1简介

现实世界的应用程序通常需要模型在有限的计算预算下进行最佳速度和精度权衡，例如无人机和自动驾驶。这激发了朝着高效 ViT 设计的大量工作，例如 PVT、Swin 和 Focal Transformer 等。为了衡量计算复杂度，最近 ViT 设计中广泛采用的度量是浮点运算的数量，即 FLOPs。但是，FLOPs 是一个间接指标，不能直接反映目标平台上的真实速度。例如，Focal-Tiny 在 GPU 上比 Swin-Ti 慢得多，尽管它们的 FLOPs 相当。

一般来说，最近 ViT 中间接指标 (FLOP) 和直接指标 (速度) 之间的差异可归因于2个主要原因。

首先，尽管 self-attention 在低分辨率特征图上很有效，但由于内存访问成本高，内存和时间的二次复杂度使得在高分辨率图像上的速度要慢得多，从片外 DRAM 获取数据可能会加快速度消耗。
其次，ViTs 中的一些高效注意力机制理论上的复杂度保证较低，但实际上在 GPU 上速度很慢，因为特定的操作对硬件不友好或无法并行化，例如多尺度窗口划分、递归和扩张窗口等。

有了这些观察结果，在本文中通过直接度量来评估 ViT，即吞吐量，而不仅仅是 FLOPs。基于此原理引入了 LITv2，这是一种新型高效且准确的Vision Transformer，它在标准基准测试中优于大多数最先进的 ViT，同时在 GPU 上实际上速度更快。

LITv2 是在 LITv1 的基础上构建的，LITv1 是一个简单的 ViT Baseline，它在早期 Stage 移除所有Multi-Head Self-Attention，同时在后期 Stage 应用标准Multi-Head Self-Attention。得益于这种设计，LITv1 比许多现有的 ImageNet 分类工作更快，因为早期的 Multi-Head Self-Attention 没有计算成本，而后来的 Multi-Head Self-Attention 只需要处理下采样的低分辨率特征图。然而，标准的 Multi-Head Self-Attention 在高分辨率图像上仍然存在巨大的计算成本，尤其是对于密集预测任务。

为了解决这个问题，本文提出了一种新的高效注意力机制，称为 HiLo。HiLo 的动机是自然图像包含丰富的频率，其中高/低频在编码图像模式中扮演不同的角色，即分别为局部细节和全局结构。典型的 MSA 层在所有图像块上强制执行相同的全局注意力，而不考虑不同基础频率的特征。这促使提出将 MSA 层分成两条路径，其中一条路径通过局部Self-Attention和相对高分辨率的特征图对高频交互进行编码，而另一条路径通过全局注意力和下采样特征对低频交互进行编码，从而大大提高了效率。

图1

具体来说，HiLo 采用两种有效的注意力来解耦特征图中的高/低频率。如图 1 所示，在上面的路径中，将几个Head分配给高频注意力（Hi-Fi），以通过Local Window Self-Attention（例如，2×2 窗口）捕获细粒度的高频，即比标准 MSA 更有效。在下面的路径中，实现低频注意力（Lo-Fi），首先将平均池化应用于每个窗口以获得低频信号。然后，将剩余的Head配给 Lo-Fi，以对输入特征图中的每个query位置与每个窗口的平均池化低频key之间的关系进行建模。受益于key和value长度的减少，Lo-Fi 复杂度显著降低。

最后，将细化的高/低频特征连接起来，并将结果输出转发到后续层。由于 Hi-Fi 和 Lo-Fi 都没有配备扩张窗口和递归等耗时操作，因此 HiLo 的整体框架在 GPU 上是快速的。通过全面的基准测试表明，HiLo 在性能、FLOPs、吞吐量和内存消耗方面优于现有的注意力机制。

此外，作者发现 LITv1 中的固定相对位置编码由于针对不同图像分辨率的插值而显著降低了其在密集预测任务上的速度。为了提高效率，建议在每个 FFN 中采用一个具有零填充的 3×3 深度卷积层，以结合从零填充中隐式学习的位置信息。此外，3×3 卷积同时有助于扩大 LITv1 中早期多层感知器 (MLP) 块的感受野。

最后，在 ImageNet、COCO 和 ADE20K 上进行了广泛的实验，以评估 LITv2 的性能。与 SoTA 模型的全面比较表明，LITv2 架构以更快的吞吐量实现了具有竞争力的性能，使 ViT 在现实场景中运行低延迟应用程序更加可行。

2本文方法

2.1 HiLo Attetion

自然图像包含丰富的频率，其中高频捕获对象的局部细节（例如，线条和形状），而低频编码全局结构（例如，纹理和颜色）。

然而，典型 MSA 层中的全局自注意力没有考虑不同底层频率的特征。为此，建议在注意力层的特征图中分别处理高/低频。将新的注意力机制命名为 HiLo，如图 1 所示。如图所示，HiLo 包含高频注意力 (Hi-Fi) 和低频注意力 (Lo-Fi)，用于对特征图中不同频率的关系进行建模。

1、High-frequency attention (Hi-Fi)

直观地说，由于高频对对象的局部细节进行编码，因此在特征图上应用全局注意力可能是冗余且计算成本高的。因此，建议设计 Hi-Fi 来捕获具有局部窗口自注意力（例如2 × 2窗口）的细粒度高频，这可以节省大量的计算复杂度。

此外，在 Hi-Fi 中采用了简单的非重叠窗口，与窗口移位或多尺度窗口划分等耗时的操作相比，它对硬件更加友好。

2、Low-frequency attention (Lo-Fi)

最近的研究表明，MSA 中的全局注意力有助于捕获低频。然而，直接将 MSA 应用于高分辨率特征图需要巨大的计算成本。

1653967043(1).png

请注意，由于 K 和 V 的空间缩减，Lo-Fi 同时降低了矩阵运算和Self-Attention的复杂度。

3、Head splitting

1653967148(1).png

4、Complexity Analysis

设 N 和 D 是 HiLo 注意力层中的Token数和隐藏维度数。将 s 表示为窗口大小。为简单起见，假设 Hi-Fi 和 Lo-Fi 具有相同的Head数和特征图具有相等的宽度和高度。那么，每个注意力的计算成本来自三个部分：

Q、K、V 矩阵的投影
注意力计算和V的加权和
最终线性投影

对于 Hi-Fi，每个部分的计算成本为

1653967177(1).png

图2

如图2-(a)和(b)所示，在较小的输入图像分辨率和较小的 s 值（例如，s = 2）下，Hi-Fi 和 Lo-Fi 的效率相当。然而，在分辨率更高的情况下，Lo-Fi 将导致巨大的计算成本。在这种情况下，稍微增加 s（例如，s = 4）有助于 Lo-Fi 在保持准确性的同时实现更好的效率。

将这2个注意力结合在一起，更大的窗口大小也有助于 HiLo 的整体框架减少高分辨率图像上的更多 FLOPs，如图 2-(c) 所示。因此，建议在现有框架中采用 HiLo 的实用指南：

增加窗口大小以便在高分辨率图像上获得更好的效率。

2.2 Positional Encoding

位置编码由于其排列不变的特性而对Self-Attention至关重要。在 LITv1 中，后来的 MSA 采用与 Swin 相同的相对位置编码 (RPE) 方案。这个与使用绝对位置编码相比，该方法在 ImageNet 上的 Top-1 准确度上显着提高了 0.7%。

然而，在密集预测任务中，必须针对不同的图像分辨率对固定的 RPE 进行插值，这大大减慢了 LITv1 的训练/推理速度。正如最近的一项研究表明，位置信息可以从 CNN 中的零填充中隐式学习，建议在每个 FFN 中采用具有零填充的一层 3×3 深度卷积层来代替耗时的 RPE。值得注意的是，由于消除了早期Stage中的 MSA，LITv1 中的早期Stage Block只剩下 FFN，这导致了 1×1 的微小感受野。为此，作者还在每个 FFN 中采用的 3×3 卷积滤波器通过同时扩大早期Stage的感受野来改进 LITv2。

2.3 模型架构

LITv2 具有三个变体：LITv2-S、LITv2-M 和 LITv2-B，分别对应 LITv1 中的Small、medium和base设置。为了公平比较，保持网络宽度和深度与 LITv1 相同。整体修改简单地分为2个步骤：

在每个 FFN 中添加一层带有零填充的深度卷积，并删除所有 MSA 中的所有相对位置编码
用建议的 HiLo 注意力替换所有注意力层。

具体架构信息和配分分别参考下图3和表6

图3

表6

3实验

3.1 消融实验

1、HiLo与其他注意力机制的比较

基于 LITv2-S，将 HiLo 的性能与 ImageNet-1K 上的其他有效注意力机制进行了比较，包括 PVT 中的空间缩减注意力 (SRA)、Swin 中的基于移位窗口的注意力 (W-MSA) 以及Twinsz中的交替的局部和全局注意力 (T-MSA)。

表 4

在实现中，直接用每个比较方法替换 HiLo。结果如表 4 所示。总的来说，HiLo 减少了更多的 FLOPs，同时实现了更好的性能和更快的速度。

图3

此外，在图 3 中，提供了基于不同图像分辨率的更多注意力机制的综合基准，包括 Focal、QuadTree 和 Performer。由于并行性较弱，它们甚至比在 GPU 上使用标准 MSA 还要慢。与它们相比，HiLo 在 FLOPs、吞吐量和内存消耗方面取得了具有竞争力的结果。