Swin Transformer:使用移位窗口的分层视觉Transformer

原文:《Swin transformer with multiscale 3D atrous convolution for hyperspectral image classification》

摘要

本文提出了一种新的视觉 Transformer,称为 Swin Transformer,它能够作为计算机视觉的通用骨干。将 Transformer 从语言适应到视觉的挑战源于两个领域之间的差异,例如视觉实体的规模差异较大,以及图像中的像素与文本中的单词相比的高分辨率。为了解决这些差异,我们提出了一个分层 Transformer,其表示是使用移位窗口计算的。移位窗口方案通过将自注意力计算限制在非重叠的本地窗口,同时还允许跨窗口连接,带来了更高的效率。这种层次结构具有在各种尺度上建模的灵活性,并且具有相对于图像大小的线性计算复杂性。Swin Transformer 的这些品质使其能够兼容广泛的视觉任务,包括图像分类(ImageNet-1K 上的 87.3 top-1 准确度)和物体检测等密集预测任务(COCO testdev 上的 58.7 box AP 和 51.1 mask AP) 和语义分割(ADE20K val 上为 53.5 mIoU)。其性能在 COCO 上大幅超越了之前的最先进水平,为 +2.7 box AP 和 +2.6 mask AP,在 ADE20K 上为 +3.2 mIoU,展示了基于 Transformer 的模型作为视觉骨干的潜力。分层设计和移位窗口方法也被证明对全 MLP 架构有益。代码和模型可在https://github.com/microsoft/Swin-Transformer上公开获取。

主要思路

计算机视觉中的建模长期以来一直由卷积神经网络 (CNN) 主导。从 AlexNet 及其在 ImageNet 图像分类挑战上的革命性表现开始,CNN 架构通过更大的规模、更广泛的连接和更复杂的卷积形式发展变得越来越强大。随着 CNN 作为各种视觉任务的骨干网络,这些架构的进步带来了性能的改进,广泛提升了整个领域的水平。 另一方面,自然语言处理 (NLP) 中网络架构的演变采取了不同的路径,今天流行的架构是 Transformer 。Transformer 专为序列建模和转导任务而设计,因其利用注意力来对数据中的远程依赖性进行建模而闻名。它在语言领域的巨大成功促使研究人员研究它对计算机视觉的适应性,最近在某些任务上展示了有希望的结果,特别是图像分类和联合视觉语言建模。在本文中,我们寻求扩展 Transformer 的适用性,使其可以作为计算机视觉的通用骨干,就像它在 NLP 中的作用以及 CNN 在视觉中的作用一样。我们观察到,将其在语言领域的高性能转移到视觉领域的重大挑战可以通过两种模式之间的差异来解释。其中一个差异涉及规模。与作为语言 Transformer 中处理基本元素的单词 token 不同,视觉元素在规模上可能存在很大差异,这是在目标检测等任务中需要注意的问题。在现有的基于 Transformer 的模型中,token 都是固定规模的,这一特性不适合这些视觉应用。另一个区别是图像中像素的分辨率比文本段落中的单词要高得多。存在许多视觉任务,例如语义分割,需要在像素级进行密集预测,这对于高分辨率图像上的 Transformer 来说是很棘手的,因为其自注意力的计算复杂度与图像大小成二次方。为了克服这些问题,我们提出了一种通用的Transformer主干,称为 Swin Transformer,它构建分层特征映射,并且具有与图像大小相关的线性计算复杂度。 如图1(a)所示,Swin Transformer 通过从小尺寸的补丁(用灰色表示)开始并逐渐合并更深 Transformer 层中的相邻补丁来构建分层表示。借助这些分层特征图,Swin Transformer 模型可以方便地利用先进技术进行密集预测,例如特征金字塔网络 (FPN) [42] 或 U-Net [51]。线性计算复杂度是通过在分割图像的非重叠窗口(以红色框出)内局部计算自注意力来实现的。每个窗口中的 patch 数量是固定的,因此复杂度与图像大小成线性关系。这些优点使 Swin Transformer 适合作为各种视觉任务的通用骨干网,与之前基于 Transformer 的架构 [20] 形成鲜明对比,后者生成单一分辨率的特征图并具有二次方复杂度。 Swin Transformer 的一个关键设计元素是连续自注意力层之间窗口分区的移动,如图 2 所示。移动的窗口桥接了前一层的窗口,提供了它们之间的连接,从而显着增强了建模能力(参见表 4)。这种策略在现实世界的延迟方面也很有效:窗口内的所有查询补丁共享相同的 key 集,这有利于硬件中的内存访问。相比之下,早期基于滑动窗口的自注意力方法由于不同 query 像素的 key 集不同,因此在通用硬件上延迟较低。我们的实验表明,所提出的移位窗口方法的延迟比滑动窗口方法低得多,但建模能力相似(参见表 5 和表 6)。事实证明,移位窗口方法对于全 MLP 架构也是有益的。

本文方法

总体结构

图 3 概述了 Swin Transformer 架构,其中展示了微型版本 (Swin-T)。它首先通过补丁分割模块(如 ViT )将输入 RGB 图像分割成不重叠的补丁。每个补丁都被视为一个“token”,其特征被设置为原始像素 RGB 值的串联。在我们的实现中,我们使用 的 patch 大小,因此每个 patch 的特征维度为 。线性嵌入层应用于此原始值特征,将其投影到任意维度(记为C)。 在这些 patch 标记上应用了几个经过修改的自注意力计算的 Transformer 块(Swin Transformer 块)。Transformer 块维护标记的数量(),并与线性嵌入一起称为“阶段 1”。 为了产生分层表示,随着网络变得更深,通过 patch 合并层来减少标记的数量。第一个 patch 合并层连接每组 相邻 patch 的特征,并在 维连接特征上应用线性层。这将标记数量减少了 倍(分辨率下采样 2 倍),并且输出维度设置为 。随后应用 Swin Transformer 块进行特征转换,分辨率保持为 。patch 合并和特征转换的第一个块被表示为“阶段2”。该过程重复两次,即“阶段 3”和“阶段 4”,输出分辨率分别为 。这些阶段共同产生分层表示,具有与典型卷积网络相同的特征图分辨率,例如VGG和ResNet。因此,所提出的架构可以方便地替换现有方法中用于各种视觉任务的主干网络。 Swin Transformer 模块 Swin Transformer 是通过将 Transformer 模块中的标准多头自注意力(MSA)模块替换为基于移位窗口(第 3.2 节中描述)的模块而构建的,其他层保持不变。如图 3(b) 所示,Swin Transformer 模块由基于移位窗口的 MSA 模块组成,后跟中间带有 GELU 非线性的 2 层 MLP。在每个 MSA 模块和每个 MLP 之前应用 LayerNorm (LN) 层,并在每个模块之后应用残差连接。

基于移位窗口的自注意力

标准 Transformer 架构及其对图像分类的适应都进行全局自注意力,其中计算标记与所有其他标记之间的关系。全局计算导致标记数量的二次复杂度,使其不适合许多需要大量 token 来进行密集预测或表示高分辨率图像的视觉问题。 非重叠窗口中的自注意力 为了提高建模效率,我们建议在局部窗口内计算自注意力。所述窗口被布置为以非重叠的方式均匀地划分图像。假设每个窗口包含 个补丁,则全局 MSA 模块和基于 个 patch 图像的窗口的计算复杂度为: 其中前者与补丁数 成二次方,后者在 固定(默认设置为 7)时呈线性。对于大型硬件来说,全局自注意力计算通常无法承受,而基于窗口的自注意力是可扩展的。 连续块中的移动窗口分区 基于窗口的自注意力模块缺乏跨窗口的连接,这限制了其建模能力。为了引入跨窗口连接,同时保持非重叠窗口的高效计算,我们提出了一种移位窗口分区方法,该方法在连续的Swin Transformer块中交替使用两种分区配置。 如图 2 所示,第一个模块使用从左上角像素开始的规则窗口划分策略,将 特征图均匀划分为大小为 个窗口。 然后,下一个模块采用与前一层的窗口配置不同的窗口配置,通过将窗口从规则划分的窗口中移动个像素。 使用移位窗口分区方法,连续的 Swin Transformer 块计算如下: 其中 分别表示块 的 (S)W-MSA 模块和 MLP 模块的输出特征;W-MSA 和 SW-MSA 分别表示使用规则和移位窗口分区配置的基于窗口的多头自注意力。 移位配置的高效批处理计算 移位窗口分区的一个问题是,它将导致更多的窗口,从移位配置中的 ,并且一些窗口将小于 。一个简单的解决方案是将较小的窗口填充到 的大小,并在计算注意力时屏蔽掉填充的值。当规则分区中的窗口数量较少时,例如 ,这种朴素解决方案增加的计算量是相当大的(,增加了2.25倍)。 在这里,我们提出一种更有效的批处理计算方法,即向左上方向循环移位,如图 4 所示。在此移位之后,批处理窗口可能由特征图中不相邻的几个子窗口组成,因此采用掩蔽机制将自注意力计算限制在每个子窗口内。通过循环移位,批处理窗口的数量与规则窗口分区的数量保持一致,因此也是高效的。这种方法的低延迟如表 5 所示。 相对位置偏差 在计算自注意力时,我们遵循[49,1,32,33],在计算相似度时每个头包含一个相对位置偏差

其中 是查询矩阵、键矩阵和值矩阵; 是查询/键维度,是窗口中的 patch 数。由于沿每个轴的相对位置在 范围内,我们参数化一个较小尺寸的偏置矩阵 中的值为取自。 我们观察到与没有此偏差项或使用绝对位置嵌入的同类相比有显着改进,如表 4 所示。如 [20] 中那样向输入进一步添加绝对位置嵌入会稍微降低性能,因此在我们的实现中没有采用。 预训练中学习到的相对位置偏差也可用于初始化模型,以便通过双三次插值对不同的窗口大小进行微调[20, 63]。

结构变量

我们构建了名为 Swin-B 的基本模型,其模型大小和计算复杂度与 ViTB/DeiT-B 类似。我们还介绍了 Swin-T、Swin-S 和 Swin-L,它们分别是模型大小和计算复杂度约为 的版本。请注意,Swin-T 和 Swin-S 的复杂度分别与 ResNet-50 (DeiT-S) 和 ResNet-101 相似。默认情况下,窗口大小设置为 。对于所有实验,每个头的查询维度为 ,每个 MLP 的扩展层为 。 这些模型变体的架构超参数是:

  • Swin-T:, 层数
  • Swin-S:, 层数
  • Swin-B:, 层数
  • Swin-L:, 层数

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

与 Vi Transformer 的区别

  • Swin Transformer使用了类似卷积神经网络中的层次化构建方法(Hierarchical feature maps),特征图尺寸中有对图像下采样4倍的,8倍的以及16倍的,这样的结构有助于更好地捕捉全局特征,有助于在此基础上构建目标检测,实例分割等任务。简单来说即为多尺度的VIT,把多尺度的设计融入到了 Transformer 中。而在之前的 Vision Transformer 中是一开始就直接下采样16倍,后面的特征图也是维持这个下采样率不变。
  • 在Swin Transformer中使用了 Windows Multi-Head Self-Attention(W-MSA) 的概念,比如在下图的4倍下采样和8倍下采样中,将特征图划分成了多个不相交的区域(Window),并且 Multi-Head Self-Attention 只在每个窗口(Window)内进行。相对于Vision Transformer中直接对整个特征图进行 Multi-Head Self-Attention,这样做的目的是能够减少计算量的,尤其是在浅层特征图很大的时候(即此时的下采样倍率比较低地情况下)。这样做虽然减少了计算量但也会隔绝不同窗口之间的信息传递,所以在论文中作者又提出了 Shifted Windows Multi-Head Self-Attention(SW-MSA) 的概念,通过此方法能够让信息在相邻的窗口中进行传递。