[Transformer]Swin Transformer: Hierarchical Vision Transformer using Shifted Windows

from 微软亚洲研究院

Paper
Code

Abstract

本文提出一种新的视觉Transformer：Swin Transformer,可以作为视觉Transformer的主干网络；将Transformer从NLP迁移到计算机视觉领域有以下挑战：物体尺度变化较大、图像像素信息比文本单词需要更高的分辨率。为了解决以上难题，文本提出一种基于滑窗的层次Transformer，通过将self-attention的计算限制在互
不重叠的局部窗口上，从而提升计算效率；同时也允许跨窗口的连接。
这种层次结构可以灵活的在不同尺度上建模，计算只是线性复杂度O(n),n为图像size。
Swin Transformer可以用于诸多视觉任务，如图像分类，密集预测，目标检测等。层次设计和滑窗方法也被证明适用于各种MLP架构。

Section I Introduction

CNN已经是CV建模的主流模型，CNN正变得越来越复杂，如连接更多、结构更精密，这些改进也进一步提升了CNN的性能。
NLP领域则沿着另一不同的路径演进，目前Transformer是最流行的模型，因为它可以有效的捕获数据的长程依赖关系，在语言领域取得了优异的性能；同时也激发了CVer的研究热情，目前已经将Transformer应用于图像分类、视觉-语言联合建模等任务。


本文拓展了Transformer的适用性，使其可以作为CV任务的一种通用主干网络。将Transformer从NLP迁移到CV主要存在以下两方面的挑战：
一是目标的尺度，在NLP中word token是基础单位，但是CV中的目标尺度存在较大差异，比如目标检测中不同目标尺度差别很大；二是图像中像素的分辨率比文本中单词的分辨率要高得多。



在一些密集预测任务，比如语义分割，需要对每一个像素进行预测，这种分辨率Transformer是难以处理的，因为计算的时间复杂度是序列长度的平方。

为了克服以上挑战，本文的Swin Transformer构造了不同尺度的feature map，并且将计算复杂度降到线性。
Fig 1(a)展示了Swin Transformer的示意图，可以看到Swin Transformer通过在更深层次中合并相邻的patch来获得不同尺度的特征表述；借助这些层次feature map就可以进行密集预测，比如送入FPN或UNet。相比Fig 1(b)中的纯Transformer（ViT）就是基于单一尺度的feature map计算的结果，并且是平方的复杂度。
每个窗口中patch的数目是固定的，这样计算的复杂度就与图像大小成线性关系，这一特点使得Swin Transformer适合作为各种数据任务的通用主干。

Swin Transformer的关键在于窗口在连续的self-attention层之间的移动，移动的窗口有效的连接了当前层和前一层，显著提升了建模能力；移窗方法也有效减少了延迟，因为窗口内所有查询patch共享同一个key set，也方便硬件中的内存访问。
而对比早期其他的移窗法，不同pixel执行query对应不同的keyset，因此对内存访问不友好。

Fig 2展示了Swin Transformer中移动窗口窗的计算，在第l层划分了多个窗口，会在每个窗口内计算self-attention，在下一层，即l+1层可以看到通过移位组合成了新窗口，新窗口的self-attention的计算就跨越了低l层中窗口的边界，从而建立了窗口之间的连接。
而Swin Transformer在图像分类、目标检测、语义分割任务中都取得了优异的性能，超过了ViT,DeiT，ResNeXT等模型。
本文认为设计跨CV和NLP的模型对这两类任务都有益，可以帮助视觉和文本的联合建模，使得两个领域的信息可以进一步共享；而SwinTransformer有助于视觉和文本信号的联合建模。

Section II Related Work

CNN及其变体

CNN已经成为计算机视觉领域中的主流模型，从Alexnet发展至今，网络结构正逐渐变得更加深入、更高效，如VGG,GoogleNet,ResNet,DenseNet,HRNet,EfficientNet等；还有诸多研究针对卷积操作进行了优化，如深度卷积、可变卷积等。这些CNN及其变体依旧是CV任务的基准框架，而本文强调的是Transformer在CV和NLP之间统一建模的潜力。
基于Self-Attention的基准框架

一些工作也尝试用self-attention layer替代Resnet中部分或全部卷积层；而self-attention的计算是在每个像素的局部窗口内进行的，主要是为了更好的权衡精度和FLOPs，但是局限就是访存代价使得这些工作的实际延迟远超CNN。因此本文不建议使用华创发，而是在连续层之间移动窗口，这样对硬件更友好。


基于Transformer的基准框架


ViT开创性的直接将Transformer用于图像，通过输出互不重叠的patch进行图像分类；与CNN相比更好的权衡了精度与速度；但是ViT需要在大规模数据集上进行训练效果才会比较好，DeiT则提出了一些训练策略，使之可以使用较小的预训练数据集。虽然ViT取得了令人满意的结果但是并不适合作为视觉任务的通用backbone，比如无法适应密集预测。
也有的工作结合上采样、反卷积等将ViT用来做目标检测、语义分割等但性能略有不足。



本文的Swin Transformer虽然作为一种通用框架被提出，但是在图像分类任务中实现了精度-速度的最佳权衡。也有peer work（PVT）通过构建多尺度特征图这种类似的思路进行密集预测，但是计算复杂度仍然是二次的；而本文优化到了线性复杂度，Swin Transformer在COCO目标检测和ADE20K易于分割上均达到了SOTA。

Section III Method

Part 1 Overall Architecture

Fig 3展示了Swin Transformer的示意图，仅展示了Tiny结构。

Set Up:
首先或将输入图像切成互不重叠的patch，每一个patch作为token，patch中所有像素值的级联作为最初的feature；本文将patch size设置为4x4，因此每个patch的feature维度是4x4x3，随后对这一原始特征值进行线性映射到C维度。

Stage 1：

随后将线性投影后的tokens送入Swin Transformer Block进行处理，处理过程中保持token的维度(H/4 X W/4)


Satge 2：


为了获得不同尺度的特征表述，随着网络层次的加深会合并patch，减少token的数量。比如第一个patch会与临近的2x2个patch的特征进行级联，然后在进行线性映射。这会使得token的数量减少4倍，输出维度设置为2C。因此Swin Transformer会将维度变换为H/8 * W/ 8 * 2C



Stage 3 和Stage 4重复这一过程，输出分辨率分别为W/16*H/16, W/32 * H/32.
通过不同的stage获得了不同尺度的特征表述，类似典型CNN网络中不同分辨率的feature map,如VGG和ResNet。
因此本文提出的网络结构可以方便的替代现有各种视觉任务方法中的主干网络。



Swin Transformer block

Swin Transformer block在Transformer的MSA基础上做了一定改动，其他结构保持不变。将SA中替换成了基于移位窗口的MSA（W-MSA->SW-MSA），然后block中还包含两层LN和一个MLP，具体链接方式参见Fig 3（b）.

Part 2 Shifted Window based Self-Attention

原始Transformer以及一些变体中计算的是每一个token与其他全部token的self-attention，因此计算复杂度是序列长度的平方，对于视觉任务中一些密集预测或者高分辨率任务就不是很友好。

Self-Attention in non-overlapped windows

为了高效计算，本文建议在局部窗口内计算self-attention，因此这些窗口以互不重叠的方式均匀的分割图像，每个窗口包含MxM个patch。这样常规的MSA和本文的SW-MSA的计算复杂度对比为：

可以看到一个是(hw)^2,一个是（hw）线性的，并且M是固定的，本文默认设为7，这种设定是可扩展的。

Shifted window partitioning in successive blocks


但是基于窗口的SA模块缺少窗口之间的连接，限制了其建模能力，为了加入窗口之间的连接同时又保证计算的效率，本文提出一种移位窗口分配的方法，会在两个连续的Swin Transformer模块中交替进行。
以Fig 2为例，左图是常规的划分方法，8x8个feature map被分到4个窗口中，每个窗口包含4x4个feature map；第二个模块则会合并窗口，将前层窗口移动[m/2,m/2]，移动到右图的位置，因此consecutive Swin Transformer模块的计算结果为：

W-MSA表示常规的基于窗口的多头自注意力，SW-MSA表示移动后的窗口多头自注意力。

通过窗口移动就建立起了前一层与当前层之间互不重叠窗口的联系，也证明这种方法对图像分类、目标检测、语义分割等均有效。
Efficient batch computation for shifted configuration

一个问题是移动窗口后窗口数目变多了，有的窗口尺寸还会小于MxM；一种简单的解决办法就是将其填充会MxM大小，计算注意力的时候把补0的这一部分mask掉不参与计算。
假设原始窗口数为2x2，移动后窗口数目变成3x3，计算量增加了2.25倍。
因此本文提出了一种更高效的batch computation的方法,参见Fig 4，通过循环向左上角移动，此时窗口中的feature map来自不同的子窗口，此时利用mask计算每个子窗口的那部分值；通过这种batch computation的方式，批处理窗口的数目与常规窗口的计算数目相同，保证了效率，Table 5中延迟指标也显示出了这种方法的有效性。


Relative position bias

下面是SA的计算公式：

因为在窗口内计算，每个feature map的相对位置在[-M+1,M-1]范围内，因此本文使用了一个参数化尺寸较小的偏置矩阵Bhat,B的值取自于Bhat。
我们观察到，如果不加入偏置项或者使用绝对位置嵌入，并不会有明显的性能提升，详情参见Table 4。
还有的研究表明输入嵌入绝对位置甚至会导致精度略有下降，因此本文也没有使用绝对位置。
在预训练中学习到的相对位置bias也可以通过插值的方式，用于初始化不同窗口大小的网络模型，比如对网络进行微调时。

Part 3 Architecture Variants

本文搭建的Swin -B参数量与ViT-B/Deit-B相近，此外还有Swin -T,-S,-L,参数量分别是Swin-B的0.25x,0.5x.2x.此外，窗口大小默认设置为M=7，head的维度为d=32,其他网络结构超参数为：
Table 1展示了模型参数量、计算量等信息。

Section IV Experiments

本文在ImageNet，COCO目标检测和ADE20K语义分割数据集上进行了测试。

Part 1 ImageNet Classification

ImageNet-1K数据集包含1.28M训练图像，50K验证图像，共1000类。


Table 1展示了top-1 accuracy。



其他训练细节：Adam Optimizer 训练300epoches



首先在ImageNet–22K预训练后再在ImageNet-1K上微调



Table 1(a)展示了直接在Imagenet-1K上训练的对比结果，与目前SOTA的DeiT结构相比，SwinTransformer结构复杂性没有超过DeiT，但精度获得了进一步的提升；



与卷积网络相比（RegNet,EfficientNet）,Swin Transformer可以更好的权衡精度-速度；此外虽然RegNet和Efficient是完全通过架构搜索得到，而Transformer是根据标准Transformer做的改进，具有进一步优化的潜力。
Table 1(b)则是预训练之后的结果，可以看到预训练比预训练精度进一步提升，而且更好的权衡了精度-速度。





Part 2 COCO Object Detection

目标检测对比的经典框架有：Mask R-CNN,ATSS,RepPoint,Sparse RCNN.对比的卷积网络则是ResNeXt和Trasformer的Deit。





Table2是对比结果，可以看到在Swin T做backbone下Cascaded Mask R-CNN取得了最佳精度，在参数量相近的情况下，精度增益更高。
与DeiT的对比结果显示，Swin-T的精度优于Deit-S;DeiT推理速度较低的主要原因在于计算复杂度是输入图像的二次。

Part 3 ADE20K Semantic Segmentation

ADE20K是广泛被使用的语义分割数据集，包含150类共25K张图像，其中training:testing:validation = 20K:3K：2K
Table 3则是分割结果，对比了mIoU，参数量，FLOPs等信息.
可以看到在计算成本接近的情况下，Swin-S比Deit-S精度更高；簇ResNet系列也高。
Swin-L预训练之后的结果比模型更大的SETR性能更高。

Part 4 Ablation Study

本节通过消融实验验证Swin Transformer关键部位的作用。


Shifted windows



Table 4显示的是在三个数据集上进行消融实验的结果，可以看到借助移动窗口，在每个任务上几乎都有1%的精度提升，实验结果充分证明了移位窗口在建立跨层窗口连接的有效性，从Table 5也可以看出，引入移动
窗口带来的延迟开销也不大。




Relative position bias




Table 4还展示了使用不同位置嵌入的对比结果引入相对位置嵌入，可以看到引入绝对位置信息有一定的精度提升，但引入相对位置信息更为有效；此外还要注意到虽然引入绝对信息对于图像分类任务有一定的精度提升，但是对语义分割和目标检测，其实是引入了反作用。
可以看到 虽然ViT、Deit网络结构抛弃了保持变换不变性的部分，但是本文发现一些平移不变性仍然适合于通用的视觉任务建模，尤其是对于目标检测或语义分割等这种密集预测任务。









Different self-attention methods




本文还比较了不同注意力方法的速度，可以看到cyclic attention的计算方式更适合硬件加速，尤其是对于一些具有较深层次的网络结构。
从Table5的结果看出，基于shifted windows的注意力计算方法在不同阶段均对总体效果有加速，Table 6展示与滑窗法相比在三种任务上的精度，展示出了两者在视觉任务建模上的相似性，而且Swin-T+shifted window在精度上提升了2.3%.

Section V Conclusion

本文提出一种新的Swin Transformer结构，可以产生不同尺度的特征表述，并且注意力的计算是线性复杂度。Swin Transformer在目标检测和语义分割任务中均达到了SOTA。本文希望Swin Transformer成为视觉任务和翻译任务中的一种通用的网络框架。

shifted window作为Swin Transformer的核心部件，实验结果显示基于移动窗口计算注意力是十分有效且高效的，本文会进一步研究其在NLP领域中的应用。

Appendix

Part 1 Detailed Architecture

Table 1展示了Swin Transformer不同规模的具体结构，输入图像为224*224，Concat表示将一个patch内临近的nxn个feature进行级联-最终效果是将feature map进行n倍的下采样；dim 96表示线性层的输出维度是96;win sz = 7x7表示MHSA的计算窗口大小为7x7.


附录还包含了进行图像分类、目标检测、语义分割的具体的实验设定，有需要的可以参考Part 2部分。






Part 3 More Experiments

本文还进行了其他实验。

1.Image Classcification with different input size

Table 8展示了不同输入尺寸的图像的性能，可以看到输入图像分辨率越高 top-1精度越高，但是响应的推理速度也会变慢。

2.Different Optimizer for ResNeXt on COCO

Table 9展示的则是使用不同优化器（AdamW,SGD）对目标检测任务的性能影响。虽然SGD是Mask R-CNN默认的Optimizer，但本文观察到替换为AdamW可以进一步提升精度，尤其在backbone较小时。

3.Swin MLP Mixer

本文还将提出的分层设计和移动窗口都应用到MLP-Mixer结构上，称之为Swin-Mixer.Table 10也展示了Swin-Mixer与原始MLP-Mixer等性能的对比，可以看到Swin Mixer可以更好的权衡精度-性能。

以上实验结果均表明：本文提出的分层设计和移动窗口方法均是值得进一步推广的。