Swin-transformer论文阅读笔记(Swin Transformer: Hierarchical Vision Transformer using Shifted Windows)

发布时间 2023-10-31 17:39:01作者: 风終會離開

论文标题:Swin Transformer: Hierarchical Vision Transformer using Shifted Windows
论文作者:Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, Baining Guo
论文来源:ICCV 2021,Paper
代码来源:Code


1. 背景介绍

长期以来,计算机视觉建模一直由卷积神经网络(CNN)主导。随着CNN作为各种视觉任务的骨干网络,这些架构上的进步已经导致了性能的提高,并广泛地提升了整个领域。

另一方面,自然语言处理(NLP)中网络架构的发展已经走了一条不同的道路,今天流行的架构是Transformer。Transformer是为序列建模和转导任务而设计的,它以关注数据中的远程依赖关系建模而闻名。它在语言领域的巨大成功促使研究人员研究它对计算机视觉的适应性,最近它在某些任务上显示了有希望的结果,特别是图像分类和联合视觉语言建模。

本文观察到,将其在语言领域的高性能转移到视觉领域的重大挑战可以通过两种模式之间的差异来解释。

  • 作为transformer处理的基本元素不同。在NLP中,token都是固定规模的,这一特性不适合视觉应用。
  • 图像的像素分辨率比文本段落中的单词高得多。许多视觉任务,如语义分割,需要在像素级别进行密集预测,这对于高分辨率图像上的Transformer来说是难以处理的,因为其自关注的计算复杂度是图像大小的二次方。

为了克服这些问题,我们提出了一种通用的Transformer主干,称为Swin Transformer,它构建分层特征映射,并且具有与图像大小相关的线性计算复杂度。
image

如图1(a)所示,Swin Transformer通过从小尺寸的patch(灰色轮廓)开始并逐渐合并更深Transformer层中的相邻补丁来构建分层表示。有了这些分层特征映射,Swin Transformer模型可以方便地利用高级技术进行密集预测,如特征金字塔网络。
线性计算复杂度是通过在分割图像(红色轮廓)的非重叠窗口内局部计算自关注来实现的。
每个窗口的补丁数量是固定的,因此复杂度与图像大小成线性关系。这些优点使得Swin Transformer适合作为各种视觉任务的通用主干,而以前基于Transformer的体系结构产生单一分辨率的特征映射并且具有二次复杂度。

Swin Transformer的一个关键设计元素是它在连续的自关注层之间的窗口分区的移位,如图2所示。

image

在第1层(左),采用规则的窗口划分方案,并在每个窗口内计算自关注。在下一层1 + 1(右)中,窗口分区被移动,产生新的窗口。新窗口中的自关注计算跨越第1层先前窗口的边界,提供它们之间的连接。
实验表明,移位窗口方法比目前常用的滑动窗口方法具有更低的延迟,但在建模能力上是相似的。
image

image

2. 研究现状

CNN及其变体

cnn在整个计算机视觉中都是标准的网络模型。虽然CNN已经存在了几十年,但直到AlexNet的引入,CNN才起飞并成为主流。此后,人们提出了更深入、更有效的卷积神经架构,进一步推动了计算机视觉领域的深度学习浪潮,如VGG、GoogleNet、ResNet、DenseNet、HRNet和EfficientNet。除了这些架构上的进步,在改进单个卷积层方面也有很多工作,比如深度卷积和可变形卷积。
虽然CNN及其变体仍然是计算机视觉应用的主要骨干架构,但本文强调了类似transformer的架构在视觉和语言之间统一建模的强大潜力。

基于自注意的骨干架构

同样受到自注意层和Transformer架构在NLP领域成功的启发,一些作品使用自注意层来取代流行的ResNet中的部分或全部空间卷积层。在这些工作中,自注意是在每个像素的局部窗口内计算的,以加速优化,并且它们比对应的ResNet架构实现了稍微更好的精度/FLOPs权衡。然而,它们昂贵的内存访问导致它们的实际延迟明显大于卷积网络。本文建议在连续层之间移动窗口,而不是使用滑动窗口,这允许在一般硬件中更有效地实现。

自注意/Transformer来补充CNNs

另一项工作是用自注意层或Transformer来增强标准的CNN架构。自注意层可以通过提供编码远程依赖关系或异构交互的能力来补充主干网络或头部网络。最近,Transformer中的编码器-解码器设计已被应用于对象检测和实例分割任务。我们的工作探索了Transformer的基本视觉特征提取,是对这些工作的补充。

基于Transformer的视觉主干

与我们的工作最相关的是Vision Transformer(ViT)及其后续研究。ViT的开创性工作是直接将Transformer架构应用于非重叠的中等大小图像块上进行图像分类。与卷积网络相比,它在图像分类方面实现了令人印象深刻的速度和精度权衡。虽然ViT需要大规模的训练数据集(即JFT-300M)才能表现良好,但DeiT引入了几种训练策略,使ViT在使用较小的ImageNet-1K数据集时也能有效。ViT在图像分类方面的结果令人鼓舞,但由于其低分辨率特征映射和复杂度随图像大小的二次增长,其结构不适合作为密集视觉任务或高分辨率输入图像的通用骨干网络。
有一些工作将ViT模型应用于直接上采样或反卷积的密集视觉任务,如对象检测和语义分割,但性能相对较低。与本文的工作同时进行的是修改ViT架构,以更好地进行图像分类。根据经验,本文发现Swin Transformer架构在这些图像分类方法中实现了最佳的速度精度权衡,即使本文的工作侧重于通用性能而不是专门用于分类。另一项同时进行的工作探索了在Transformer上构建多分辨率特征图的类似思路。它的复杂性仍然是图像大小的二次方,而本文的复杂性是线性的,也在局部运行,这已被证明有利于建模视觉信号的高相关性。本文的方法既高效又有效,在COCO目标检测和ADE20K语义分割上都达到了最先进的精度。

3. 方法

3.1 总体架构

image

图3给出了Swin Transformer体系结构的概述,其中演示了小型版本(Swin-T)。

它首先通过像ViT一样的patch分割模块将输入的RGB图像分割成不重叠的patch。每个patch被视为一个“token”,其特征被设置为原始像素RGB值的串联。本文中,使用4 × 4的patch大小,因此每个patch的特征维数为4 × 4 × 3 = 48。在此原始值特征上应用线性嵌入层,将其投影到任意维度(记为C)。
在这些patch tokens上应用了几个具有改进的自注意力计算的Transformer块(Swin Transformer块)。Transformer块保持令牌的数量(\(\frac{H}{4} \times \frac{W}{4}\))不变,并与线性嵌入一起被称为“Stage 1”。

为了产生分层表示,随着网络的深入,通过patch合并层来减少token的数量。第一个patch合并层将每组2 × 2相邻patch的特征进行拼接,并在拼接后的4C维特征上应用线性层。这将token的数量减少2 × 2 = 4的倍数(2倍分辨率降采样),并将输出维度设置为2C。之后使用Swin Transformer块进行特征变换,分辨率保持在\(\frac{H}{8} \times \frac{W}{8}\)。这第一个块的patch合并和特征转换被称为“Stage 2”。该过程重复两次,作为“Stage 3”和“Stage 4”,输出分辨率分别为\(\frac{H}{16} \times \frac{W}{16}\)\(\frac{H}{32} \times \frac{W}{32}\)

这些阶段共同产生一个分层表示,具有与典型卷积网络相同的特征映射分辨率,例如VGG和ResNet。
因此,该结构可以方便地取代现有方法中的骨干网络,用于各种视觉任务。

Swin Transformer block

Swin Transformer是通过将Transformer块中的标准多头自注意(MSA)模块替换为基于移位窗口的模块(见3.2节)而构建的,其他层保持不变。如图3(b)所示,Swin Transformer模块由一个基于移位窗口的MSA模块组成,其次是一个2层MLP,中间是GELU非线性。在每个MSA模块和每个MLP模块之前应用一个LN (LayerNorm)层,在每个模块之后应用一个残差连接。

3.2 基于移位窗口的自注意

标准Transformer架构及其对图像分类的适应都进行全局自注意力,其中计算一个token与所有其他token之间的关系。全局计算导致token数量的二次复杂度,使得它不适合许多需要大量token进行密集预测或表示高分辨率图像的视觉问题。

非重叠窗口中的自注意

为了提高建模效率,建议在局部窗口内计算自注意。所述窗口被布置成以非重叠方式均匀地分割图像。假设每个窗口包含\(M × M\) patchs,则全局MSA模块和基于\(h × w\) patchs图像的窗口的计算复杂度为:

\[\begin{align} & \Omega(\mathrm{MSA})=4 h w C^{2}+2(h w)^{2} C, \tag{1} \\ & \Omega(\mathrm{W}-\mathrm{MSA})=4 h w C^{2}+2 M^{2} h w C, \tag{2} \end{align} \]

其中,前者与patch数 \(hw\) 成二次关系,后者在 \(M\) 固定时为线性关系(默认为7)。对于大型硬件来说,全局自注意计算通常是负担不起的,而基于窗口的自注意是可扩展的。

在连续块中移位的窗口分区

基于窗口的自注意模块缺乏跨窗口的连接,这限制了它的建模能力。为了引入跨窗口连接,同时保持非重叠窗口的高效计算,本文提出了一种移位窗口分区方法,该方法在连续的Swin Transformer块中交替使用两种分区配置。

image

如图2所示,第一个模块采用从左上角像素开始的常规窗口划分策略,将 \(8 × 8\) 特征图均匀划分为大小为 \(4 × 4 (M = 4)\)\(2 × 2\) 个窗口。然后,下一个模块采用与前一层不同的窗口配置,将窗口从规则分区的窗口移动\(\left( \lfloor{\frac{M}{2}} \rfloor,\, \lfloor{\frac{M}{2}} \rfloor\right)\)的像素。

使用移位窗口划分方法,连续的Swin Transformer块计算为:

\[\begin{aligned} &\hat{\mathbf{z}}^l=\text{W-MSA}\left(\mathrm{LN}\left(\mathbf{z}^{l-1}\right)\right)+\mathbf{z}^{l-1}, \\ &\begin{aligned}\mathbf{z}^l=\mathsf{MLP}\left(\mathrm{LN}\left(\hat{\mathbf{z}}^l\right)\right)+\hat{\mathbf{z}}^l,\end{aligned} \\ &\hat{\mathbf{z}}^{l+1}=\text{SW-MSA}\left(\mathrm{LN}\left(\mathbf{z}^l\right)\right)+\mathbf{z}^l, \\ &\mathbf{z}^{l+1}=\text{MLP}\left(\text{LN}\left(\hat{\mathbf{z}}^{l+1}\right)\right)+\hat{\mathbf{z}}^{l+1}, \end{aligned} \tag{3} \]

式中 \(\hat{\mathbf{z}}^l\)\(\mathbf{z}^l\) 分别表示块 \(l\) 的 (S)W-MSA 模块和 MLP 模块的输出特征。W-MSA 和 SW-MSA 分别使用常规和移位的窗口分区配置表示基于窗口的多头自注意。

移位窗口分割方法引入了前一层相邻非重叠窗口之间的联系,在图像分类、目标检测和语义分割方面效果显著,如表4所示。
image

移位构型的高效批处理计算

移位窗口分区的一个问题是,它将导致更多的窗口,从 \(\lceil\frac{h}{M}\rceil\times\lceil\frac{w}{M}\rceil\)\((\lceil\frac{h}{M}\rceil + 1)\times(\lceil\frac{w}{M}\rceil + 1)\),并且一些窗口小于 \(M × M\)

一个简单的解决方案是将较小的窗口填充到 \(M × M\) 的大小,并在计算注意力时屏蔽填充值。当常规分区的窗口数量很少时,例如:\(2 × 2\),这个朴素的解决方案增加的计算量是相当可观的(\(2 × 2 → 3 × 3\),扩大了\(2.25\)倍)。

在这里,本文提出一种更有效的批处理计算方法,即向左上方向循环移位,如图4所示。

image

在此转换之后,批处理窗口可能由特征映射中不相邻的几个子窗口组成,因此采用屏蔽机制将自关注计算限制在每个子窗口内。

通过循环移位,批处理窗口的数量与常规窗口分区的数量保持一致,因此也是高效的。这种方法的低延迟如表5所示。
image

相对位置偏差

在计算自注意时,本文遵循已有的方法,在计算相似度时为每个头部加入一个相对位置偏差:

\[\operatorname{Attention}(Q,K,V)=\operatorname{SoftMax}(QK^T/\sqrt d+B)V,\tag{4} \]

其中,\(Q,K,V\in\mathbb{R}^{M^2\times d}\) 是查询、键和值矩阵;\(d\) 是查询/键维数,\(M^2\) 是一个窗口的patch数。由于沿各轴的相对位置在 \([−M + 1,M−1]\),本文参数化了一个较小尺寸的偏置矩阵 \(\hat{B}\in\mathbb{R}^{(2M-1)\times(2M-1)}\),且 \(B\) 的取值来自 \(\hat{B}\)

观察到,此步骤与没有此偏差项或使用绝对位置嵌入的同类相比,有了显著的改进,如表4所示。
image

3.3 模型变体

本文建立了基础模型,称为 Swin-B,具有类似于 ViT-B/DeiT-B 的模型大小和计算复杂度。本文还建立了Swin-T, Swin-S 和 Swin-L,它们分别是约为模型大小和计算复杂度的0.25倍、0.5倍和2倍的版本。需要注意的是,Swin-T 和 Swin-S 的复杂度分别与 ResNet-50 (DeiT-S) 和 ResNet-101 相似。窗口大小默认设置为 \(M = 7\)。对于所有实验,每个头部的查询维数为 \(d = 32\),每个 MLP 的扩展层为 \(\alpha = 4\)。这些模型变体的体系结构超参数为:

image

其中 \(C\) 为第一阶段隐藏层的通道数。表1列出了用于ImageNet图像分类的模型大小、理论计算复杂度(FLOPs)和模型变体的吞吐量。
image

实验

此处仅给出了语义分割相关的实验。

image

结论

Swin Transformer是一种新的视觉Transformer,它产生层次特征表示,并且对输入图像大小具有线性计算复杂度。Swin Transformer在COCO对象检测和ADE20K语义分割方面实现了最先进的性能,大大超过了以前的最佳方法。论文作者希望Swin Transformer在各种视觉问题上的强大表现将鼓励视觉和语言信号的统一建模。

作为Swin Transformer的关键元素,基于移位窗口的自注意在视觉问题上是有效和高效的,论文作者期待着研究它在自然语言处理中的应用。