Transformer论文解读五（Swin Transformer）

本文中要介绍的Swin Transformer全名为Swin Transformer: Hierarchical Vision Transformer using Shifted Windows，是2021年微软研究院发表在ICCV上的一篇文章，在多项视觉任务中展现出优异性能。

1. 改进

相比于Vision Transformer（讲解），Swin Transformer做出了几点改进：

1. 层次化构建方法：使用了类似卷积神经网络中的层次化构建方法（Hierarchical feature maps），比如特征图尺寸中有对图像下采样4倍的，8倍的以及16倍的，这样的backbone有助于在此基础上构建目标检测，实例分割等任务。而Vision Transformer中直接下采样16倍，后面的特征图也是维持这个下采样率不变。
2. 特征图划分：在Swin Transformer中使用了Windows Multi-Head Self-Attention(W-MSA)的概念，比如在下图的4倍下采样和8倍下采样中，将特征图划分成了多个不相交的窗口，并且Multi-Head Self-Attention只在每个窗口内进行。相对于Vision Transformer中直接对整个特征图进行Multi-Head Self-Attention，这样做的目的是减少计算量，尤其是在浅层特征图很大的时候。
3. 窗口变换：特征图划分虽然减少了计算量，但会隔绝不同窗口之间的信息传递，因此所以在论文中作者又提出了Shifted Windows Multi-Head Self-Attention(SW-MSA)的概念，目的在于使信息在相邻窗口中传递。

2. 网络架构

论文中给出了Swin Transformer（Swin-T）网络的架构图。可以看出整个框架的基本流程如下：

1. 图片分块：首先将图片输入到Patch Partition模块中进行分块，即每$4\times 4$相邻的像素为一个Patch，然后在channel方向展平。由于每个patch就有$4\times 4=16$个像素，每个像素有R、G、B三个值，所以展平后的深度是$16\times 3=48$，所以通过Patch Partition后图像shape由 $[H,W,3]$变成了 $[\frac{H}{4},\frac{W}{4},48]$。
2. 线性变换：通过Linear Embeding层对每个像素的channel数据做线性变换，由$48$变成$C$，即图像shape再由 $[\frac{H}{4},\frac{W}{4},48]$变成了$[\frac{H}{4},\frac{W}{4},C]$。在源码中通过一个卷积层实现。
3. 网络模块堆叠：通过四个Stage构建不同大小的特征图，除了Stage1中先通过一个Linear Embeding层外，剩下三个stage都是先通过一个Patch Merging层进行下采样。然后都是重复堆叠Swin Transformer Block。注意，Transformer Block有两种结构，如图(b)中所示，这两种结构的不同之处仅在于一个使用了W-MSA结构，一个使用了SW-MSA结构。而且这两个结构是成对使用的，先使用一个W-MSA结构再使用一个SW-MSA结构。所以堆叠Swin Transformer Block的次数都是偶数（因为成对使用）。
4. 分类头：最后，对于分类网络，在后面接上一个Layer Norm层、全局池化层以及全连接层，得到最终输出。

3. Patch Merging

除了Stage1以外的Stage中，首先要通过一个Patch Merging层进行下采样，如下图所示

假设输入Patch Merging的是一个4x4大小的单通道特征图，Patch Merging的过程如下：

1. 将每个2x2的相邻像素划分为一个patch
2. 将每个patch中相同位置（同一颜色）像素拼在一起，就得到了4个feature map
3. 将四个feature map在深度方向进行concat拼接
4. 通过一个LayerNorm层
5. 通过一个全连接层在feature map的深度方向做线性变化，将feature map的深度由C变成C/2。

可以看出，通过Patch Merging层后，feature map的高和宽会减半，深度会翻倍。

4. W-MSA

引入Windows Multi-head Self-Attention（W-MSA）模块是为了减少计算量。

如下图所示，在Self-Attention计算过程中，普通的Multi-head Self-Attention（MSA）模块将feature map中的每个patch和所有的像素去计算。但Windows Multi-head Self-Attention（W-MSA）模块首先将feature map按照MxM（例子中的M=2）大小划分成一个个Windows，然后单独对每个Windows内部进行Self-Attention。

4. SW-MSA详解

采用W-MSA模块时，只会在每个窗口内进行自注意力计算，所以窗口与窗口之间是无法进行信息传递的。为了解决这个问题，作者引入了Shifted Windows Multi-Head Self-Attention（SW-MSA）模块，即进行偏移的W-MSA。

如下图所示，W-MSA和SW-MSA是成对使用，假设第L层使用W-MSA，那么第L+1层使用的就是SW-MSA。根据左右两幅图对比能够发现窗口（Windows）发生了偏移。偏移后的窗口使得之前的临近窗口之间进行交流，就解决了不同窗口之间无法进行信息交流的问题。

但是，可以发现通过将窗口进行偏移后，原来的4个窗口变成9个窗口了。为了解决这个麻烦，作者又提出而了Efficient batch computation for shifted configuration，一种更加高效的计算方法：

主要思想是将左上角的小窗口移动到右下角，和右下角的小窗口合并在一起，这样又和原来一样是4个4x4的窗口了，所以能够保证计算量是一样的。

为了防止将不同的区域合并在一起导致信息混淆的问题，在实际计算中使用的是masked MSA，通过设置蒙板隔绝其他区域的信息，使得每个像素只与自己所在区域的像素进行计算，从而解决了这个问题。

参考文献