【论文精读】Swin Transformer: Hierarchical Vision Transformer using Shifted Windows

前言

虽然最近的论文精读都涉及CV领域，与博主研究方向不同，但是这些文章大多都涉及Transformer架构，或者和自监督学习、无监督学习有关，所谓万变不离其宗，虽然在不同领域方法不同，但是很多情况下其思想上有很多的共性，特别值得读者去思考。本篇论文精读来自ICCV2021最佳论文，也奠定了Transformer架构在CV领域的地位，其方法精妙，效果显著，是一篇值得精读的佳作。

---

Paper: https://arxiv.org/pdf/2103.14030.pdf
Code: https://github.com/microsoft/Swin-Transformer
From: ICCV Best Paper

Abstract

本文提出一个新的视觉Transformer，称作Swin Transformer，可以作为CV领域的通用架构。将Transformer应用到CV领域的挑战来源于NLP与CV的差异：

1. 和文本相比，视觉对象的尺度存在巨大变化。
2. 像素的高分辨率大大增加了计算成本和复杂度。

为了解决上述问题，本文提出了层级Transformer，通过移动窗口计算层级表征。移动窗口的机制带来了更高的效率，更好的灵活性和更小的线性复杂度。这些特征让Swin Transformer在多个CV任务上取得了SOTA，比之前最好的方法提升了近3个点。实验结果证明了基于Transformer模型作为CV领域backbone的潜力，以及分层设计和移动窗口的优越性。

1. Introduction

CNNs占据CNN领域模型良久，并且性能优良。但是NLP中对Transformer架构的成功应用促使CV工作者研究其对CV领域的适应，尤其是Transformer的注意力机制可以解决长程依赖问题。因此，本文试图扩展Transformer的适用性，使其可以作为CV领域的backbone，代替CNN架构。
作者发现Transformer从语言领域迁移到视觉领域主要的挑战可以通过两个模式之间的差异来解释。第一是视觉元素在规模上存在很大差异，另一个是图像中像素分辨率比文本token要多得多，这会导致模型计算复杂度大大增多。为了解决上述问题，本文提出了Swin Transformer，它构建了分层特征图，并且计算复杂度与图像大小成线性关系。下图是Swin Transformer与ViT的比较：

Swin Transformer通过从小图像块开始，逐步合并更深Transformer层的相邻图像块来构建分层表示。借助这些分层特征图，Swin Transformer可以解决密集型的任务。线性复杂度体现在红框包围的区域，注意力的计算仅在框内进行，因此随着图像增大，计算也只是线性增大。这些优点和ViT形成鲜明对比，后者缺少局部特征的捕获，并且具有二次的复杂度。
Swin Transformer另一个关键设计是移动窗口，如下图所示：

移动的窗口桥接了前一层的窗口，让同一个红框内的图像块也能获取相邻图像块的信息，从而显著提升了模型的建模能力，此外窗口内所有查询图像块共享相同的key集，有利于硬件访问。
基于上述得到的Swin Transformer在图像分类、目标检测和语义分割任务上取得强大的性能。作者相信跨CV和NLP的统一架构可以让这两个领域同时受益，因为二者的联合建模能够共享领域知识。

2. Related Work

略。

3. Method

3.1. Overall Architecture

Swin Transformer架构总览图如上图所示。它首先通过图像分割模块将RGB图像划分为不重叠的大小为4×4×3的图像块并铺平，接着用一个线性层将图像块token投影到任意维度C。此时整个序列的维度为$\frac{H}{4}\times \frac{W}{4}\times C$。然后在这些图像token上应用Swin Transformer块，Swin Transformer块和线性embedding合称stage1。
为了产生分层的表征，随着网络变深，token的数量通过图像块合并层逐层递减。具体来说，每经过一个stage，图像块的数量减少为原来二等四分之一，输出维度增大为原来的两倍。上图的Swin Transformer一共有4个stage，这些stage共同产生分层表示，具有和典型卷积网络相同的特征图分辨率，可以方便替换现有的CV任务中的骨干网络。
Swin Transformer块如上右图所示，它将Transformer块中的多头自注意力机制MSA替换为了基于移动窗口的模块，其它层保持不变。

3.2. Shifted Window based Self-Attention

标准的Transformer用于图像分类使用的是全局自注意力，但是这种架构计算复杂度是token的二次方，使其不适合需要大量token进行密集预测或高分辨率的图像问题。

Self-attention in non-overlapped windows

为了高效建模，本文提出计算局部窗口内的自注意力。窗口以不重叠的方式均匀划分图像。假设窗口包含$M\times M$个图像块，全局MSA模块和窗口MSA模块的计算复杂度如下：
$$\begin{array}{l}\Omega (\mathrm{MSA})=4hw{C}^2+2(hw{)}^{2}C,\\ \Omega (\mathrm{W}-\mathrm{MSA})=4hw{C}^2+2M^{2}hwC,\end{array}$$

两个式子的第一项都相同，第二项中，前者是关于的二次方，后者是关于的线性变换。对于大的来说，全局自注意计算量是无法承受的，而基于窗口的自注意是可扩展的。

计算复杂度的计算其实可以很容易推导出来，这里画一张示意图就可以很好理解。以单头多注意力机制为例（多头最后的结果是一样的），其计算流程如下：

Shifted window partitioning in successive blocks

但是基于窗口的自注意力机制缺乏跨窗口的连接，这限制了其建模能力。为了引入跨窗口连接，同时保持高效计算，本文提出了一种移动窗口分区方法， 该方法在连续Swin Transformer块中的两个分区配置之间交替。
如图二所示，第一个模块对图像进行常规的划分，将8×8的特征图均匀划分为大小为4×4的2×2窗口，然后下一个模块采用和前一层不同的窗口配置，将窗口移动$\left(\lfloor \frac{M}{2}\rfloor ,\lfloor \frac{M}{2}\rfloor \right)$个像素。使用移动窗口的分区方法，连续的Swin Transformer块计算如下：
$$\begin{array}{l}{\hat{\mathbf{z}}}^l=W-MSA\left(\mathrm{LN}\left({\mathbf{z}}^{l-1}\right)\right)+{\mathbf{z}}^{l-1},\\ {\mathbf{z}}^l=\mathrm{MLP}\left(\mathrm{LN}\left({\hat{\mathbf{z}}}^l\right)\right)+{\hat{\mathbf{z}}}^l,\\ {\hat{\mathbf{z}}}^{l+1}=SW-MSA\left(\mathrm{LN}\left({\mathbf{z}}^l\right)\right)+{\mathbf{z}}^l,\\ {\mathbf{z}}^{l+1}=\mathrm{MLP}\left(\mathrm{LN}\left({\hat{\mathbf{z}}}^{l+1}\right)\right)+{\hat{\mathbf{z}}}^{l+1}\end{array}$$

其中${\hat{\mathbf{z}}}^l$和${\mathbf{z}}^l$分别是和 MLP 模块的输出，表示使用常规的窗口多头自注意力，表示使用移动窗口的窗口多头自注意力。
移动窗口分割方法引入了相邻不重叠窗口之间的连接，在多个下游任务上有效。

Efficient batch computation for shifted configuration

移动窗口的一个问题是，它会导致更多窗口的生成，并且窗口的大小不一。一个简单的解决方法是将较小的窗口填充到大小，并在计算注意力时mask掉填充的值。但是这种方法带来的计算量是客观的。
本文提出了一个更高效的批量计算方法，通过向左上方循环移位实现，如下图所示：

通过转换后，批处理窗口可能由特征图中不相邻的子窗口组成，因此采用随机掩码机制将自注意力计算限制在每个子窗口内。批处理窗口的数量保持和常规窗口数量相同，因此也是高效的。那么掩码的具体操作如下图所示，即将计算好的自注意力矩阵与右侧的掩码矩阵相乘。

Relative position bias

自注意力的计算遵循常规的计算公式，如下：

只不过作者加入了一个偏置矩阵，该偏置矩阵对模型性能有着显著改进。此外，如果向输入中添加绝对位置编码会稍微降低性能，因此本文采用的预训练中学习到的相对位置编码。

3.3. Architecture Variants

作者构建了四种不同大小规模的模型，如下所示：

其中Swin-B的大小和计算量和ViT-B类似。Swin-T和Swin-S的复杂度分别与ResNet-50和ResNet-101类似。默认情况下，窗口大小设置为7。

4. Experiments

本文对ImageNet-1K、COCO目标检测和ADE20K语义分割进行实验。实验部分首先和其他先进模型进行对比，再进行消融实验。

4.1. Image Classification on ImageNet-1K

作者分别在ImageNet-1K和ImageNet-22K上进行预训练，后者预训练后再ImageNet-1K上进行微调。训练好的模型在ImageNet-1K测试集上执行分类任务。结果分别如下面两张表所示：

常规的ImageNet-1K预训练模型分别与之前最先进的Transformer模型DeiT和最先进的卷积神经网络RegNet、EfficientNet相比，Swin Transformer实现了稍微更好的速度和精度平衡。
ImageNet-22K上预训练的模型显然受到大规模预训练数据的增益，实现了明显更好的速度和精度的平衡。Swin-B获得了86.4%的最高精度，比ViT高出2.4%，以及略低的FLOP（47.0G和55.4G）。较大的Swin-L模型实现了87.3%的最高精度。

4.2. Object Detection on COCO

对象检测和实例分割实验在 COCO 2017 上进行，为了证明Swin Transformer可以作为CV任务的backbone，作者将部分方法的卷积backbone换为了Swin Transformer。结果如下表所示：

a表列出了Swin-T和ResNet-50在四种目标检测框架上的结果。Swin Transformer在性能上全方位碾压，但是在模型尺寸、FLOP和延迟上稍大。
表2与R-CNN的不同模型进行比较，Swin Transformer实现了51.9 box AP和45.0 maskAP，显著优于其他模型。
表c是抛开所有的限制，Swin Transformer与之前最先进模型的比较。本文的模型实现了58.7的 box AP，超过了之前最优2.7个点。

4.3. Semantic Segmentation on ADE20K

语义分割任务在ADE20K上进行，下表列出了不同方法、backbone的mIoU、模型大小，FLOPs和FPS。从这些结果可以看出，Swin-L在val集上达到了53.5mIoU，比之前最好的模型高出了3.2。

4.4. Ablation Study

作者在三种下游任务上进行了消融实验，结果如下所示：

采用移动窗口的方法在每个阶段都要优于基于单窗口分区构建的方法，在ImageNet上提升1.1%，在密集预测型任务COCO和ADE20K上分别有2.8/2.2和2.8的提升。这表明了移动窗口在窗口之间建立连接的有效性，并且开销延迟也很小。
上表还显示了不同位置编码方法的比较，相对位置编码在三个任务上都是最优的方法，表明相对位置偏置的有效性。此外，绝对的位置编码在密集任务上对性能有所影响。
虽然ViT放弃了图像中的平移不变性，但是CNN中的长久经验证明其重要性，根据上表也能发现，鼓励平移不变的归纳偏置仍然可取，尤其是针对目标检测和语义分割的密集预测任务。
不同的自注意力方法在下表进行了比较：

移动窗口的循环实现具有更高的硬件效率，下表比较了它们在三个任务上的准确性，表明其在视觉建模上同样精准。此外，与最快的Transformer架构Performer相比，Swin Transformer也要稍快。

5. Conclusion

本文提出Swin Transformer，作为新的视觉Transformer架构，它提出分层特征表示，并且相对于输入的图像大小具有线性的计算复杂度。Swin Transformer在目标检测数据集COCO和语义分割数据集ADE20K上都显著超越之前最佳方法。作者希望Swin Transformer可以作为多模态的统一架构。
当然Swin Transformer还有一个关键要素移动窗口，要让Swin Transformer能够在NLP中受益，需要研究该方法在NLP中的应用。

阅读总结

Swin Transformer作为ViT的延伸工作，将图像中的特性（局部性和平移不变性）融入到Transformer模型中，帮助模型在密集预测任务上取得更好的结果。此外，作者也展露其野心，他们希望Swin Transformer架构不只是替换CV领域CNN架构的backbone，还将成为多模态大一统模型的架构，其想法远大，工作细腻，方法独特，效果显著，不愧为ICCV2021的最佳论文。
在我看来，这篇文章在工作量和创新点上都拉满了。首先工作量从实验部分就可以看出，作者将Swin Transformer和近年来领域最佳或里程碑式的模型进行了全方位对比，包括计算量、模型参数、延迟、精度、领域等，此外消融实验也在多个领域上进行，可谓论据十足。在创新点上，它创造性将图像特性和Transformer结合，通过分层特征学习和移动窗口两个创新性技术，让Transformer也能学习到图像中局部性和平移不变性的归纳偏置，同时又将计算复杂度降到线性，可谓是精妙绝伦。
其实Swin Transformer这样里程碑式的工作并不是空穴来风，也不是遥不可及的。当ViT出来的时候，它的核心在于在大规模数据集上进行预训练的Transformer也能做CV任务，相当于将Transformer引入了CV的大门，但是它并没有其它的创新点，或者说并没有引入图像的特性。因此很容易想到，将图像特性引入到Transformer中，很可能对模型的性能有很大的提升，那么图像的相较于文本的特性是什么，就是后序工作的重点。本文提出的特性是尺度变化和高分辨率，因此作者设计了分层特征提取和窗口自注意力，前者应对尺寸的变化，后者应对高分辨率带来的计算复杂度。当然图像相对于文本的特性远远不止这些，或者部分CV任务也有其独特的特性，因此针对不同的特性，设计不同的方法，才能得到有创新性的工作。