【计算机视觉】Swin Transformer：基于滑动窗口的层次视觉Transformer

Swin Transformer：基于滑动窗口的层次视觉Transformer 《Swin Transformer: Hierarchical Vision Transformer using Shifted Windows》

论文地址：https://arxiv.org/abs/2103.14030

一、简介

​ 在计算机视觉建模中，长期以来都是$\text{CNN}$占据主导地位。从$\text{AlexNet}$和其在$\text{ImageNet}$图像分类任务上的革命性表现开始，$\text{CNN}$的架构更大的规模、更广泛的连接和更复杂的卷积形式进化的越来越强大。随着$\text{CNNs}$作为各种视觉任务的骨干网络，这些架构的改善为整个领域带来了性能的提升。

​ 另一方面，自然语言处理中的网络架构则采取了不同的路线，如今的流行架构为$\text{Transformer}$。$\text{Transformer}$是为序列建模和迁移任务设计的，其显著之处在于能够对数据中的远程依赖进行建模。在语言领域中的巨大成功令研究人员尝试将其应用在计算机视觉中，其近期已经在图像分类和视觉语言联合模型上都显示了非常好的效果。

​ 在本文中，作者尝试扩展$\text{Transformer}$的适用性，使其可以作为计算机视觉的主干网络，就像其在$\text{NLP}$和$\text{CNN}$在计算机视觉中那样。作者观察到，将语言上的优良表现迁移至视觉领域的重要挑战可以通过两个模态间的区别来解释。其中一个区别是规模。不同于单词token，其可以在语言$\text{Transformer}$中作为基础元素被处理，而视觉元素在规模上的差异特别大，特别是在目标检测任务中需要被注意。在现有的基于$\text{Transformer}$的模型中，所有的token规模都是固定的，这并不适合视觉应用。另一个区别是图像中像素的分辨率比文本段落中单词高很多。许多视觉任务，如语义分割，需要在像素级别进行稠密预测。高分辨率图像对于$\text{Transformer}$来说是棘手的，自注意力计算复杂度是图像尺寸的平方。为了解决这些问题，作者提出了一种称为$\text{Swin Transformer}$的通用$\text{Transformer}$，其构造层次$\text{feature map}$并且拥有图像尺寸的线性复杂度。如上图(a)中，$\text{Swin Transformer}$通过从小尺寸$\text{patch}$开始并逐步合并相邻的$\text{patch}$，从而构建一个层次表示。基于这些层次$\text{feature map}$，$\text{Swin Transformer}$能够方便的利用先进的技术进行稠密预测，例如$\text{FPN}$和$\text{U-Net}$。通过将图像分割为非重叠窗口来计算局部自注意力机制，从而实现了线性复杂度。每个窗口中的$\text{patch}$数量是固定的，其计算复杂度是图像尺寸的线性规模。这些优点使$\text{Swin Transformer}$适合作为各种视觉任务的主干网络。

​ $\text{Swin Transformer}$的关键设计元素是连续自注意力层间的滑动窗口划分。滑动窗口链接了前一层的窗口，通过在它们之间提供链接来极大的改善模型的建模能力。$\text{Swin Transformer}$在图像分类、目标检测和语义分割等任务上取得了良好的性能。其在三个任务上显著的优于$\text{ViT/DeiT}$和$\text{ResNe(X)t}$。

​

二、$\text{Swin}$

1. 整体架构

​ 上图是一个$\text{tiny}$版本的$\text{Swin Transformer}$架构。该模型会通过$\text{patch}$划分模块将$\text{RGB}$图像划分为不重叠的$\text{patches}$，类似于$\text{ViT}$。每个$\text{patch}$被当做一个"token"，$\text{patch}$的特征被设置为原始像素$\text{RGB}$值的拼接。在具体实现中，使用$4\times 4$的$\text{patch size}$且每个$\text{patch}$的特征维度为$4\times 4\times 3=48$。使用一个线性embedding层将原始特征投影至一个任意维度，记为$\text{C}$。

​ 将若干个具有改进自注意力的$\text{Transformer}$块($\text{Swin Transformer}$块)应用在这些$\text{patch tokens}$。这些$\text{Transformer}$块会维护$(\frac{H}{4}\times \frac{W}{4})$个token，并且该层和线性嵌入层一起统称为阶段一。

​ 为了能够产生一个层次表示，随着网络变的更深，token数量会通过一个$\text{path merging}$层进行缩减。第一个$\text{patch merging}$层会拼接$2\times 2$个相邻的$\text{patch}$特征，并在合并后的$4C$维度特征上应用一个线性层。这会导致token的数量减少到$2\times 2=4$倍，并将输出维度设置为$2C$。将$\text{Swin Transformer}$块应用在特征上进行特征转换，分辨率保持为$\frac{H}{8}\times \frac{W}{8}$。$\text{patch merging}$层和特征传播$\text{Swin Transformer}$块被称为阶段二。该过程被重复两次，分别称为阶段3和阶段4，输出分辨率分别为$\frac{H}{16}\times \frac{W}{16}$和$\frac{H}{32}\times \frac{W}{32}$。这些阶段会产生一个层次表示，具有与$\text{VGG}$和$\text{ResNet}$这样典型卷积神经网络相同的特征映射分辨率。因此，$\text{Swin Transformer}$可以非常方便的替换现有视觉任务中的backbone。

• Swin Transformer块

  $\text{Swin Transformer}$通过滑动窗口模块来替换$\text{Transformer}$中的标准多头自注意力机制($\text{MSA}$)，其他部分保持相同。图上图(b)所示，$\text{Swin Transformer}$由一个基于$\text{MSA}$的滑动窗口组成，并在其后加一个具有激活函数$\text{GELU}$的双层$\text{MLP}$。$\text{LayerNorm(LN)}$层应用在每个$\text{MSA}$层和$\text{MLP}$层之前，并且在每个模块中都会应用残差链接。

2. 基于自注意力的滑动窗口

​ 将标准$\text{Transformer}$应用在图像任务中，需要执行全局自注意力机制，每个token与其他所有的token间的关系都需要被计算。这样全局的计算将导致token数量平方的复杂度，这导致其不适合具有大量token的密集预测或者一个高分辨率的视觉任务。

2.1 不重叠窗口的自注意力

​ 为了能够高效建模，提出在局部窗口内计算自注意力。该窗口通过不重叠的方法均匀划分图像。假设每个窗口包含$M\times M$个$\text{patch}$，在一个具有$h\times w$个$\text{patch}$块图像上的全局$\text{MSA}$的计算复杂度和基于窗口的计算复杂度为
$$\Omega (\text{MSA})=4hw{C}^2+2(hw{)}^{2}C\text{\hspace{1em}(1)}$$

$$\Omega (\text{W-MSA})=4hw{C}^2+2M^{2}hwC\text{\hspace{1em}(2)}$$

前者时间复杂度是$\text{patch}$数量$hw$的平方；当$M$固定时，后者的时间复杂度为线性的。当$hw$太大时，全局自注意力机制的负担太大，而基于窗口的自注意力则相对灵活。

2.2 连续块中的滑动窗口划分

​ 基于窗口的自注意力机制缺乏窗口间的连接，限制了其建模的能力。为了能够在保持无重叠窗口高效计算的同时引入跨窗口连接，本文提出了滑动窗口划分方法。该方法在连续$\text{Swin Transformer}$块之间交替使用两种划分配置。

​ 如上图说明，第一个模块使用常规的窗口划分策略。划分从左上角开始，$8\times 8$的$\text{feature map}$被均等的划分为$2\times 2$个窗口，每个窗口的尺寸为$4\times 4(M=4)$。下一个模块则通过将前一层的常规窗口移动$(\lfloor \frac{M}{2}\rfloor ,\lfloor \frac{M}{2}\rfloor )$来获得。

​ 使用滑动窗口划分方法，连续$\text{Swin Transformer}$则计算方式为
$$\begin{gathered} {\hat{\text{z}}}^l=\text{W-MSA}(\text{LN}({\text{z}}^{l-1}))+{\text{z}}^{l-1} \\ {\text{z}}^l=\text{MLP}(\text{LN}({\hat{\text{z}}}^l))+{\hat{\text{z}}}^l \\ {\hat{\text{z}}}^{l+1}=\text{SW-MSA}(\text{LN}({\text{z}}^l))+{\text{z}}^l \\ {\text{z}}^{l+1}=\text{MLP}(\text{LN}({\hat{\text{z}}}^{l+1}))+{\hat{\text{z}}}^{l+1}\text{\hspace{1em}(3)} \end{gathered}$$
其中，${\hat{\text{z}}}^l$和${\text{z}}^l$表示在块$\text{l}$中$\text{(S)W-MSA}$模块的输出特征和$\text{MLP}$的输出特征。$\text{W-MAS}$和$\text{SW-MSA}$表示使用常规和滑动窗口划分配置的多头注意力机制。

2.3 滑动窗口的高效计算

​ 滑动窗口划分会导致更多的窗口，即窗口数量从$\lceil \frac{h}{M}\rceil \times \lceil \frac{w}{M}\rceil$增加到$(\lceil \frac{h}{M}\rceil +1)\times (\lceil \frac{w}{M}+1\rceil )$。此外，还会导致一些窗口小于$M\times M$。一个简单的方法是将小窗口填充至$M\times M$尺寸，并在计算注意力时将填充值遮蔽掉。若常规窗口划分的数量较小时(例如$2\times 2$)，这种简单的方法增加的计算量并不小($2\times 2\to 3\times 3$，增大$\text{2.25}$倍)。本文提出了一个更加有效的计算方法，该方法如上图所示，通过向左上角方向循环移动。经过移动后，一个batch内的窗口可能由若干个子窗口组成，这些子窗口在$\text{feature map}$中并不相邻。因此，对于每个子窗口，需要使用一个mask机制来限制自注意力计算。通过循环移动后，batch内窗口数量和常规常量划分一致，且同样高效。

2.4 相对位置编码

​ 计算自注意力机制，在每个头计算相似度时引入相对位置偏移$B\in {\mathbb{R}}^{M^2\times M^2}$
$$\text{Attention}(Q,K,V)=\text{Softmax}(Q{K}^T/\sqrt{d}+B)V\text{\hspace{1em}(4)}$$
其中，$Q,K,V\in {\mathbb{R}}^{M^2\times d}$是$\text{query}$,$\text{key}$和$\text{value}$矩阵；$\text{d}$是$\text{query/key}$维度；$M^2$是一个窗口中的$\text{patch}$数量。因为沿每一个轴的相对位置位于范围$[-M+1,M-1]$，因此参数化一个较小尺寸的偏差矩阵$\hat{B}\in {\mathbb{R}}^{(2M-1)\times (2M-1)}$，且$B$中的值从$\hat{B}$中获取。

​ 作者发现相较于并引入偏差或者使用绝对位置嵌入，这种方法都有较大的改善。在输入中添加绝对位置嵌入会轻微降低表现，因此在实现中没有使用绝对位置嵌入。

3. 架构变体

​ 本文构造的base模型称为$\text{Swin-B}$，其模型尺寸和计算复杂度类似于$\text{ViT-B/DeiT-B}$。此外，作者还引入了$\text{Swin-T}$、$\text{Swin-S}$和$\text{Swin-L}$，其分别是base模型尺寸和计算复杂度的$0.25\times$、$0.5\times$和$2\times$版本。$\text{Swin-T}$和$\text{Swin-S}$的复杂度类似于$\text{ResNet-50(DeiT-S)}$和$\text{ResNet-101}$。窗口的尺寸默认为$M=7$。每个注意力头的$\text{query}$维度为$d=32$，每个$\text{MLP}$的扩展层是$\alpha =4$。

​ 这些模型变体的结构超参为：

• $\text{Swin-T}$：$C=96$，层数量={2,2,6,2}
• $\text{Swin-S}$：$C=96$，层数量={2,2,18,2}
• $\text{Swin-B}$：$C=128$，层数量={2,2,18,2}
• $\text{Swin-L}$：$C=192$，层数量={2,2,18,2}

​ 其中，$C$是第一阶段隐藏层的维度。

三、总结

• 通过两层之间窗口的滑动来建立层次链接；
• 模型中，相对位置编码的效果更好；
• 在视觉领域：通过patch、window以及window间的链接来逐步建模；