Title: Slide-Transformer: Hierarchical Vision Transformer with Local Self-Attention
导读
众所周知,ViT 取得良好性能的一个关键组件便是自注意力机制,它可以从全局上下文中实现自适应特征提取。当然,反复提到的一个问题还是计算复杂度问题。先前的方法主要从两个角度去解决问题:
- 基于稀疏的全局注意力;
- 采用窗口注意力;
然而,以上这些方法要么不利于局部特征学习,要么还需要人工去调参。说到这,我们需要进一步思考局部注意力的意义是什么?作者认为,它能够将每个查询的感受野限制在其自己的相邻像素内,具有卷积和自注意力的双重好处,即局部归纳偏差和动态特征选择。
遗憾的是,目前的局部注意力模块要么使用效率低下的Im2Col
函数,要么依赖于难以推广到没有CUDA
支持的设备的特定CUDA
内核。
为此,本文提出了一种新颖的局部注意力模块Slide Attention
,它利用常见的卷积操作实现高效性、灵活性和通用性。
具体来说,我们首先从新的row-based
的角度重新解释了column-based
的Im2Col
函数,并使用深度卷积作为有效替代。在此基础上,本文提出了一个基于重参数化技术的deformed shifting module
,进一步将固定的键/值位置嵌入到局部区域的可变形特征上。如此一来,该模块便能够以高效灵活的方式实现了局部注意力范式。
广泛的实验表明,本文所提的滑动注意力模块适用于各种先进的vision transformer
模型,并兼容各种硬件设备,在综合基准测试上取得了持续改进的表现。
动机
Figure 1. Comparison of our model and other attention patterns.
上图给出了几种不同自注意力机制的实现方式。与 PVT 中的全局注意力和 Swin-Transformer 中的窗口注意力相比,本文提出了一种新颖的 Slide Attention 模块,它不仅可以像局部注意力一样施加局部归纳偏置,而且极具高效性和灵活性,轻松与各种 Vision Transformer 模型和硬件设备进行有效集成。
如上所述,针对先前局部注意模块中采用的低效 Im2Col 函数,本文将以崭新的视角看待这个问题。具体来说,原始 Im2Col 从基于列的视图生成键和值矩阵,其中每一列代表以输入的特定位置为中心的局部区域。或者,我们从基于行的视图重新制定键和值矩阵,并显示每一行对应于向不同方向移动的输入特征。
这种新的见解让我们有机会更进一步应用精心设计的Depthwise Convolutions
来代替原本的移位操作。如此一来,我们便可以将Im2Col
函数替换为标准的卷积运算,这样便能够很容易部署到不同的硬件设备上。此外,为了进一步增强灵活性,我们引入了一种新颖的可变形移位模块。如图 1(c) 所示,该模块将固定键和值位置relaxes
到局部区域内的可变形特征。通过使用重参数化技术,既有效地增加了模型容量,又同时保持了推理效率。
Figure 2. Performance and inference speed comparison on local attention implementations.
为了对上述方法进行全面比较,本文实际分析了这两种局部自注意力实现的性能和运行时间,并将其与 SwinTransformer 中的窗口注意力进行了比较。如图 2 所示,基于 Im2Col 的局部注意力在效率和性能方面都不太有利。而基于 CUDA 的方法虽然可以通过窗口注意力等矢量化操作保持相当的推理速度,但仅显示出比其他注意力模式(例如 Swin-Transformer 中的窗口注意力)略微的提升。考虑到在不同硬件上适配CUDA
内核的难度,目前为止,我们仍然缺乏兼具高效性和高泛化性的局部注意力模块。
针对这种现象,本文工作首先将Im2Col
函数重新解释为特征移位操作,然后用更有效的深度卷积进行替代方案。可以看到,当进一步配备可变形的移位模块时,所提方法在类似的推理时间下实现了比基线的显着改进。上述 FPS 是基于 NVIDIA RTX3090 GPU 设备上测试的。
方法
Figure 3. Different implementation on the local attention module.
如上所述,局部注意力受到效率问题的困扰,导致实际应用中很难被部署上去。在本节中,作者首先展示了低效的 Im2Col 函数可以从另一个角度重新解释并证明其等价于一组简单的特征变换。在此基础上,本文进一步采用高效的深度卷积代替特征移位操作。最后,通过配备一种新颖的可变形移位模块,本文最终提出了一种局部注意模块,称为 Slide Attention,极具高效性和灵活性。
New Perspective on Im2Col
首先,我们可以将 Im2Col 函数表示为如下形式:
其次,从基于列的视图来看,如图 3(1) 所示,键/值矩阵包含 列,其中每一列对应于以特定查询为中心的局部窗口。具体来说,如果我们仔细检查输出的每一列,上面的等式可以重新表述为:
上述公式表示以 (i, j) 为中心的局部窗口,即对所有查询的窗口进行采样并按顺序放置。
紧接着,我们从基于行的角度来看,如图 3(2) 所示,键/值矩阵包含 行,其中每一行对应于将输入移向某个方向。具体来说,我们可以关注输出的每一行并将上述等式重新表述为:
这相当于将原始特征图向某个特定的方向 进行平移。
通过这种方式,我们提供了一种新的替代方法来理解 Im2Col 函数,方法是将基于列的视图替换为新颖的基于行的视图。以 为例,如果我们首先将原始特征图向 9 个不同方向移动(如图 3(2.b)),然后将这些特征展平为行,最后将它们连接在列中(如图 3(2.c)),得到的键/值矩阵被证明等价于 局部窗口,可以恢复与原始 Im2Col 函数完全等价的输出(如图 3(1.c))。
Shift as Depthwise Convolution
如上所述,我们提供了一种新的视角去理解 Im2Col 函数,但简单地将特征向不同方向移动仍然涉及低效的切片操作,这对提升局部注意力的效率几乎没有帮助。然而,与 Im2Col 中所有查询的采样窗口不同,我们可以采用更高效的方式实现,即深度卷积,如图 3(3) 所示。
具体地,对于窗口大小为 k 的局部注意力,我们可以将 Im2Col 函数重新实现为深度卷积,从而减轻主要开销。此外,这些深度卷积可以进一步归结为单组卷积,这不仅可以避免低效的切片操作,还可以转换为卷积操作从而轻松部署到各种硬件设备上。
Deformed Shifting Module
通过将原来的 Im2Col 函数切换为 depthwise convolutions,局部注意力的效率得到了极大的提升。尽管如此,精心设计的内核权重仍然将键和值限制在固定的相邻位置,这可能不是捕获不同特征的最佳解决方案。
因此,本文提出了一种新颖的可变形移位模块,以进一步增强局部注意力的灵活性。具体来说,我们在 shiftwise 卷积中利用设计范例,并引入并行卷积路径,其中内核参数在训练过程中随机初始化和学习。与将特征向不同方向移动的固定核相比,可学习内核可以解释为所有局部特征的线性组合。这类似于可变形卷积网络 中的不规则感受野。
如上图 4 所示,额外的卷积路径从几个角度改进了局部注意力模块:
- The key and value pairs in the local attention are extracted by a more flexible module, that greatly improves model capacity and can capture diverse features.
局部注意力机制是指在自注意力机制中,只关注输入序列中特定的局部区域,从而减少计算量和存储需求。这里提到的灵活模块指的是用来提取局部区域内特征的模块,其灵活性使得模型能够更好地适应不同的任务和数据集,并能够捕捉到更多的特征。
- The learnable convolution kernel shows a resemblance with the deformable technique in DCN. Similar to the bilinear interpolation of four neighboring pixels in DCN, our deformed shifting module can be viewed as a linear combination of features within the local window. This finally contributes to augmenting the spatial sampling locations and model geometric transformation of inputs.
本文提出的可学习卷积核其实同 DCNs 中的可变性卷积类似,可以对输入进行更灵活的采样,从而增强了模型对几何变换的适应性。具体来说,可变形平移模块中采用了类似于 DCN 中的双线性插值的方法,将局部窗口内的特征进行了线性组合,以获得更多的采样位置和几何变换的信息。
- We use the re-parameterization technique to transform the two parallel paths into a single convolution. In this way, we can improve the model capacity while maintaining inference efficiency.
为了进一步提高效率,本文引入了重参数化技术,以将原本需要两个并行路径来处理的信息,通过一次卷积操作实现,从而提高了模型的容量,同时还保持了模型的推断效率。
实验
重点看下上图,图中展示了所提方法在不同硬件下的实际推理时间,包括边缘端(iPhone 12)和服务器端(RTX 3090)。可以看出,与基线相比,嵌入滑动注意力模块后可以在不同设备上的运行时和模型性能之间表现出明显更好的权衡,并且在 Vision Transformer 模型上实现了高达 2.3 倍的加速。顺便一提,对于其他局部注意力模块,由于基于 CUDA 的方法无法在这些设备上实现,作者这里仅将所提方法与基于 Im2Col 的方法进行比较。
总结
在这篇论文中,作者重新审视了局部注意力机制,并通过提出一种新颖的 Slide Attention 模块来解决其效率问题。该模块仅使用常规卷积操作,将低效的 Im2Col 函数替换为深度卷积,并配备了一个可变形平移模块,从而实现了高效、灵活和通用的局部注意力。广泛的实验表明,所提模块可以广泛应用于各种视觉 Transformer 模型和不同的硬件设备上,同时在计算效率和模型性能之间实现了更好的权衡。
整体来说,本文所提出的滑动注意力中一味增加窗口大小的只会带来边际改进。因此,探索更大的窗口尺寸带来更好的性能-效率权衡值得我们进一步研究,这将会是未来研究的一个重要方向。