论文阅读《Swin Transformer: Hierarchical Vision Transformer using Shifted Windows》

论文链接：https://arxiv.org/abs/2103.14030

代码地址： https://github.com/microsoft/Swin-Transformer

开源15天，star超3k的作品。作者提出了一个基于transformer的backbone，可用于多种视觉任务。和以往ViT，DETR等结构不同的是，Swin Transformer通过shifted windows操作，实现了CNN里面的hierarchical的结构。这类hierarchical的结构更适用于不同的scale，同时其计算复杂性仅与image size线性相关。实验证明，以Swin Transformer为backbone的模型，在分类、检测和分割等多个任务上实现了霸榜。

Motivation

将NLP领域的Transformer迁移到CV的task上，需要考虑这两个模态之间的不同：（1）scale问题：像object detection，目标的尺度不一样，而现有的Transformer里面的token大多对应一个固定的尺度，因此并不适合这类task；（2）图像任务需要更高的resolution：以semantic segmentation为例，对现存的Transformer来说，想要得到更高分辨率的结果，意味着要承担image size的平方的时间复杂度，代价显然非常高。

基于以上两点考量，作者提出了Swin Transformer以解决上述问题：（1）首先该backbone的特征图构建是hierarchical的，因此可以适应多尺度。通过不断的融合临近的patch，即可得到hierarchical的特征图，类似于ResNet，后面可以接FPN或U-Net进行多尺度密集预测；（2）然后他的计算复杂性是和image size线性相关的，相比于现存方法的平方相关优化了很多。用到了non-overlapping window，即一个window包含若干个patch，然后所有的self-attention均在window内部locally计算。由于window内patch数量是固定的，所以时间复杂度与image size成线性关系。

由于上述的window是non-overlapping的，如果只对每个window独立地做self-attention，并不能有效的融合周围的特征，不适合提取到hierarchical的feature map（对比CNN里的kernal在整个feature map上进行滑动）。因此Swin Transformer中的一个核心结构shifted window，用来划分相邻self-attention层的window。

Architecture

Swin Transformer总体结构如上。输入是RGB图像，首先经过一个Patch Partition模块，这里以$4\ast 4$大小为一个patch，所以划分后维度变成$\frac{H}{4}\ast \frac{W}{4}\ast 48$。然后经过一个Linear Embedding层，可以将特征嵌入到任意维度，这里记作$C$。随后经过一个核心的Swin Transformer Block模块，token数不变。以上属于Stage1。

为了产生一个hierarchical的特征表示，token的数量随着网络的加深应该不断减少，因此这里采用Patch Merging模块，来对$2\ast 2$区域内的patch进行融合，这样每一个新patch特征维度变成$4C$。为了减少计算量，融合后紧跟一个降维操作，将维度降到$2C$。随后经过Swin Transformer Block，维度保持不变。以上为Stage2。

后续Stage3-5和之前是一样的，不断融合相邻的patch，直到Stage5输出的特征图大小为$\frac{H}{32}\ast \frac{W}{32}\ast 8C$。这五个Stage联合起来产生了hierarchical的特征表示。和VGG，ResNet这些网络是一样的，都产生了多层级的feature map。因此Swin Transformer这个backbone相比于其他的视觉Transformer，更适合用在CV的task上面。

其中的核心模块Swin Transformer Block如Figure3(b)所示，关键之处使用W-MSA和SW-MSA取代了传统Transformer中的MSA（Multi-head self attention module），其他基本不变。MLP，LayerNorm和shortcut结构都用在了这个Block中，不过不是重点。重点是SW-MSA用到了shifted window操作，使得hierarchical feature和线性时间复杂度成为可能。

Shifted Window based Self-Attention

关于window与patch的关系，上图表示的很清楚。首先分析一下以前的global self-attention，和本文采用在每个window内部做self-attention的时间复杂度。两种时间复杂度如下式：

两个公式的由来是这样的：
（1）$Q=x\ast W^Q$，需要$hw{C}^2$。K和V的复杂度与Q一样。随后是$Q{K}^T$需要$(hw{)}^{2}C$，其结果与V相乘同样需要$(hw{)}^{2}C$。得到结果$Z$以后乘$W^Z$需要$hw{C}^2$。因此总计$4hw{C}^2+2(hw{)}^{2}C$。这个结果显然是和image size平方相关的。
（2）如采用window内部self attention，KQVZ的复杂度都不变。$Q{K}^T$需要计算$\frac{h}{M}\ast \frac{w}{M}$个window，每个window的复杂度为$(M^2{)}^{2}C$，因此总计$M^{2}hwC$。后续结果与V相乘的计算量也为$M^{2}hwC$。所以总体的时间复杂度为$4hw{C}^2+2M^{2}hwC$。由于M是固定的，因此该结果与image size线性相关。

通过上述分析，Swin Trabsformer已经解决了现存方法时间复杂度高的缺点。但如果采用只计算window内部的self-attention，会导致不同window之间缺少connection，换句话说，global信息将会损失。为了解决这个问题，作者提出了Shifted window，用以划分相邻layer特征图上的window。

一种很直观的方法就是按照上述Figure2右侧图的划分方式，这样在之前层的那些non-overlapping的window就可以建立联系。但是Figure2的Shifted window方式也有弊端，window数目从$2\ast 2$变成了$3\ast 3$，计算复杂度直接上涨2.25倍，并不是特别乐观。因此作者提出了一种shifted方式，如下图：

如上图，依然按照Figure2的方式划分网格，但是采用了cyclic shift（左二）：将浅色区域的A B C平移到深色区域的对应位置。在做了这个shift操作以后，在原来的feature map上一些不相邻的window会变成相邻的，且位于同一个新的window中，因此需要采用一个mask机制，来抑制这两个本不相邻的部分在一个window中做self attention。这一部分最好结合代码，可参考该文章。

最后是在做self attention的时候，加上一个position bias，如下：

Experiment

分类任务惜败于EfficientNet，可能是因为多尺度和分类关系不是那么密切？

检测精度还是很高的。

语义分割最大的网络精度也很高。

总结

作者提出了一个Transformer backbone，不仅可以产生hierarchical的特征表示，还可以使时间复杂度降至和image size线性相关。核心部分就是window的引入降低了复杂度，patch的融合以及shifted window的引入可以提取多尺度的feature。Swin Transformer在分类、检测、分割等多个任务中达到了很高的精度，同时拥有不俗的速度。虽然不认同是对CNN的降维打击，但这种融合了CNN层级思想的Transfomer架构还是非常值得学习和follow的。