【神经网络架构】Swin Transformer_window-based attention

Swin Transformer: Hierarchical Vision Transformer using Shifted Windows

Swin Transformer:使用Shifted Windows的分层视觉Transformer

2021.08 微软亚洲研究院

摘要

        Swin Transformer：计算机视觉的通用backbone。将Transformer从语言调整到视觉的挑战 来自于这两个领域之间的差异。例如视觉实体规模的巨大变化以及与文本中的文字相比，图像中像素的高分辨率。为了解决这些差异，我们提出了一个分层的Transformer，它的表示是通过Shifted Windows计算的。Shifted Windows方案通过将self-attention计算限制在不重叠（non-overlapping）的局部窗口，同时允许跨窗口连接（cross-window connection），从而提高了计算效率。这种分层结构具有在不同尺度上建模的灵活性，并具有相对于图像大小的线性计算复杂度。Swin Transformer的这些特性使它能够兼容广泛的视觉任务，包括图像分类、以及密集的预测任务：如目标检测、语义分割。 state-of-the-art。分层设计和Shifted Windows方法也证明对所有的mlp架构是有益的。

 总结

        本文介绍了Swin Transformer，这是一种新的视觉Transformer，它可以产生层次特征表示，并与输入图像大小相比具有线性计算复杂度。Swin Transformer在COCO对象检测和ADE20K语义分割方面取得了最先进的性能，大大超过了以往最好的方法。我们希望Swin Transformer在各种视觉问题上的强大表现将鼓励视觉和语言信号的统一建模。作为Swin Transformer的关键元素，基于shifted window的self-attention在视觉问题上显示出了有效和高效的效果，我们也期待着研究它在自然语言处理中的应用。

Introduction

        我们观察到，将语言域的高性能转移到视觉域的重大挑战可以用两种模态的差异来解释：

1. 差异之一涉及规模。与语言transformer中作为基本处理元素的word tokens不同，视觉元素可以在规模上有很大的差异，这是在目标检测等任务中需要注意的问题。在现有的基于transformer的模型中，tokens都是固定scale的，这一属性不适合这些视觉应用程序。
2. 图像的像素分辨率比文字段落中的words高得多。目前存在许多视觉任务，如语义分割，需要在像素级进行密集预测，而这对于Transformer在高分辨率图像上是难以实现的，因为其self-attention 的计算复杂度是图像大小的二次方。

        为了克服这些问题，我们提出了一个通用的Transformer backbone，称为Swin Transformer，它构建分层feature maps，计算复杂度与图像大小成线性关系。

        如图1(a)所示，Swin Transformer构建了一个分层表示，从较小的块（patches）(用灰色表示)开始，逐渐将相邻的块合并到更深的Transformer层中。通过这些分层feature maps，Swin Transformer模型可以方便地利用高级技术进行密集预测，如特征金字塔网络(FPN)或U-Net。

        线性计算复杂度是通过在划分图像(用红色表示)的非重叠窗口中局部计算 self-attention 来实现的。由于每个窗口中的patches数量是固定的，因此复杂度与图像大小成线性关系。这些优点使Swin Transformer适合作为各种视觉任务的通用backbone，而不像以前的基于Transformer的架构，后者生成单一分辨率的feature maps，复杂度为二次型。

         如图2所示，Swin Transformer的一个关键设计元素是在连续的self-attention层之间移动窗口分区。移位的窗口连接了前一层的窗口，提供了它们之间的连接，显著增强了建模能力(见表4)。这种策略对于real-world latency也是有效的:  窗口中的所有query patches共享相同的key set，这有利于硬件中的内存访问。相比之下，早期基于滑动窗口（sliding window）的self-attention方法在一般硬件上的延迟较低，因为不同的key sets对应不同的query pixels。

 注：虽然在普通硬件上实现基于滑动窗口的卷积层有很多有效的方法，但由于它在feature map上共享kernel weights，因此在实践中，基于滑动窗口的自注意层很难实现高效的内存访问。

        实验结果表明，提出的shifted window方法比sliding window方法有更低的延迟，但在建模能力上是相似的。shifted window方法也被证明对所有的mlp架构是有益的。

        所提出的Swin Transformer在图像分类、目标检测和语义分割等识别任务上都取得了较好的性能。在三个任务的延迟相似的情况下，它显著优于ViT / DeiT和ResNe(X)t模型。它在COCO测试开发集上的58.7 box AP和51.1 mask AP 超过了先前最先进的结果+2.7 box AP((Copy-paste，不需要外部数据)和+2.6 mask AP(DetectoRS)。在ADE20K语义分割上，在val集上得到 53.5 mIoU，比之前的最先进的(SETR)提高了3.2 mIoU。在ImageNet-1K图像分类中，准确率最高达87.3%。

        我们相信，一个跨越计算机视觉和自然语言处理的统一架构可以使这两个领域受益，因为它将促进视觉和文本信号的联合建模，并且两个领域的建模知识可以更深入地共享。我们希望Swin Transformer在各种视觉问题上的出色表现能够在社区中推动这种信念的深入，并鼓励视觉和语言信号的统一建模。

方法

1. 整体架构

        图3展示了Swin Transformer架构的概述，它展示了小型版本(Swin-T)。它首先通过一个patch splitting module(像ViT)将一个输入的RGB图像分割成非重叠的patches。每个patch都被视为一个“token”，其特征被设置为原始像素RGB值的串联（concatenation）。在我们的实现中，我们使用的patch大小为4*4，因此每个patch的特征维数为4*4*3 = 48。在这个原始值特征上用一个linear embedding层，将其投影到任意维度(记作C)。

        在这些patch tokens上应用了几个带有修改的self-attention计算的Transformer块(Swin Transformer块)。Transformer块维护tokens的数量(H/4*W/4)，并与linear embedding一起称为“Stage 1”。

        为了生成层次表示，随着网络的加深，通过patch merging层减少tokens的数量。第一个patch merge层将2 * 2 neighboring patches的每组的features进行拼接，并对拼接后的4C维feature应用线性层。这减少了tokens的数量，其倍数为 2*2 = 4 (2× downsampling of resolution)，输出维度设置为2C。之后使用Swin Transformer块进行特征变换，分辨率保持在H/8 * W/8。这个patch merging和特征转换的第一个块记为“Stage 2”。该过程重复两次，分别为“Stage 3”和“Stage 4”，输出分辨率分别为H/16 * W/16和H/32 * W/32。这些阶段共同产生的层次表示具有与典型卷积网络(如VGG和ResNet相同的feature map resolutions。因此，该体系结构可以方便地替代现有的各种视觉任务的backbone。

Swin Transformer block

        Swin Transformer是通过将Transformer块中的 the standard multi-head self attention (MSA) module替换为基于shifted windows的模块来构建的，其他层保持不变。如图3(b)所示，Swin Transformer块包括一个基于shifted window的MSA模块，然后是一个带有GELU非线性的2层MLP。在每个MSA模块和每个MLP之前应用一个LayerNorm (LN)层，在每个模块之后应用一个residual connection。

2. 基于 Self-Attention的Shifted Window

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

非重叠窗口中的Self-Attention

        为了高效建模，我们提出在local windows内计算self-attention。通过排列窗口，以不重叠的方式均匀地分割图像。假设每个窗口包含 M*M 个patches，在 一张 h * w 个patches的图像上， 一个global MSA模块 和  基于window的MSA模块 的计算复杂度为：

 注：在确定计算复杂度时省略了SoftMax计算。

 其中前者是patch数量（hw）的二次，后者在M固定时(默认为7)是线性的。Global self-attention计算对于大的hw来说通常是昂贵的，而基于window的self-attention是scalable。

在连续blocks中Shifted window划分

        window-based self-attention模块缺乏跨窗口的连接，这限制了它的建模能力。为了引入跨窗口连接，同时保持非重叠窗口的高效计算，我们提出了一种移位窗口划分(shifted window partitioning)方法，它在consecutive Swin Transformer blocks中的两个分块配置之间交替。

        如图2所示, 第一个模块使用一个普通窗口划分策略从左上的像素开始, 8*8的feature map是均匀地划分为 2*2 个窗口，每个窗口大小 4*4  (M = 4)。然后，下一个模块采用了一个窗口配置，来自于上一层窗口的移动，通过用$(\left \lfloor \frac{M}{2} \right \rfloor, \left \lfloor \frac{M}{2} \right \rfloor)$像素从有规则分区的Windows中替换Windows。

        采用移位窗口划分方法，consecutive Swin Transformer blocks计算为：

其中$\hat{z}^{l}$和$z^{l}$表示block L的 (S)W-MSA 模块和MLP模块的输出特征。W-MSA和SW-MSA分别表示基于窗口的使用规则的multi-head self-attention和移位的窗口分区配置。

        移位窗口划分方法引入了上一层相邻无重叠窗口之间的连接，发现该方法在图像分类、目标检测和语义分割方面是有效的，如表4所示。

针对shifted configuration的高效批量计算

        shifted window partitioning的一个问题是在shifted configuration中从 $\small \left \lceil \frac{h}{M} \right \rceil \times \left \lceil \frac{w}{M} \right \rceil$到 $\small \left (\left \lceil \frac{h}{M} \right \rceil + 1 \right ) \times \left (\left \lceil \frac{w}{M} \right \rceil + 1 \right )$, 它会导致更多的windows, 并且一些windows会比 M * M。

注：为了使窗口大小 (M, M) 能被feature map的大小(h, w) 整除，如果需要，在feature map上使用右下padding。

         一个简单的解决方案是将较小的窗口填充到 M * M 的大小，并在计算注意力时屏蔽填充的值。当regular partitioning中的窗口数较小时， 如2 * 2，使用这种简单的方案增加的计算量是相当大的(2 * 2 → 3 * 3, which is 2.25 times greater)。在这里，我们提出了一种更有效的批处理计算方法，即向左上角方向循环移动，如图4所示。在此shift之后，一个批处理窗口可能由几个在feature map中不相邻的子窗口组成，因此采用掩蔽机制将self-attention计算限制在每个子窗口内。通过循环移位，批处理窗口的数量与常规窗口划分的数量保持一致，因此也是高效的。这种方法的低延迟如表5所示。

相对位置偏差

        在计算self-attention时，我们遵循[49,1,32,33]的方法，在计算相似度时将相对位置偏差$\large B \in \mathbb{R}^{M^{2}\times M^{2}}$包含在每个头上。

 这里， $Q,K,V \in \mathbb{R}^{M^{2}\times d}$是 query, key 和 value 矩阵；d is the query/key dimension，and M2 is the number of patches in a window. 由于每个轴上的相对位置在[−M + 1, M −1]范围内，我们参数化了一个较小尺寸的偏置矩阵$\large \hat{B} \in \mathbb{R}^{(2M-1) \times (2M-1)}$， B中的值是从$\large \hat{B}$中取的。

        我们观察到，与没有这个偏差项或使用绝对位置嵌入的对照相比，有显著的改善，如表4所示。像[20]中那样，进一步向输入中添加绝对位置嵌入会略微降低性能，因此在我们的实现中没有采用它。

        在预训练中学习到的相对位置偏差也可以通过双三次插值来初始化不同窗口大小的微调模型[20,63]。

3. 架构变种

        我们构建了称为Swin-B的基本模型，使其模型大小和计算复杂度与ViTB/DeiT-B类似。我们还引入了Swin-T、Swin-S和Swin-L，它们分别是模型大小和计算复杂度分别为0.25x、0.5x和2x的版本。请注意，Swin-T和Swin-S的复杂度分别与ResNet-50 (DeiT-S)和ResNet-101的复杂度相似。窗口大小默认设置为M = 7。所有实验中，每个头部的query维数为 d = 32，每个MLP的扩展层数为α = 4。这些模型变量的结构超参数是:

• Swin-T: C = 96, layer numbers = {2, 2, 6, 2}
• Swin-S: C = 96, layer numbers ={2, 2, 18, 2}
• Swin-B: C = 128, layer numbers ={2, 2, 18, 2}
• Swin-L: C = 192, layer numbers ={2, 2, 18, 2}

这里，C为第一阶段隐藏层的通道数。在ImageNet图像分类上， 模型尺寸、理论计算复杂度(FLOPs)和模型变种的吞吐量如表1所示。

实验

略