具有 HiLo 注意力机制的快速视觉 Transformer

Vision Transformers (ViTs，视觉 Transformer) 引发了计算机视觉领域最近且最重大的突破。其高效设计大多受间接计算复杂度度量（即 FLOPs）的指导，然而这与吞吐量等直接度量存在明显差距。因此，我们提出以目标平台上的直接速度评估作为高效 ViT 的设计原则。特别地，我们引入了 LITv2，一种简单而有效的 ViT，其在一系列不同模型规模下以更快的速度优于现有的最先进方法。LITv2 的核心是一种新颖的自注意力机制，我们称之为 HiLo。HiLo 的灵感来自于这样的洞察：图像中的高频捕获局部精细细节，低频关注全局结构，而多头自注意力层忽视了不同频率的特性。因此，我们提出通过将头部分成两组来解耦注意力层中的高/低频模式，其中一组通过每个局部窗口内的自注意力编码高频，另一组通过每个窗口和输入特征图中每个查询位置的平均池化低频键与值之间的全局注意力编码低频。受益于两组的高效设计，我们通过对 GPU 和 CPU 上的 FLOPs、速度和内存消耗进行全面基准测试，表明 HiLo 优于现有的注意力机制。例如，在 CPU 上 HiLo 比空间缩减注意力快 1.4 倍，比局部窗口注意力快 1.6 倍。在 HiLo 的驱动下，LITv2 作为主流视觉任务（包括图像分类、密集检测与分割）的强骨干网络。代码见 https://github.com/ziplab/LITv2。