swin-transformer原理介绍

作者：Vincent Liu
链接：https://www.zhihu.com/question/521494294/answer/2492957365
来源：知乎
 

Swin Transformer 的提出可以说是里程碑式的，在学术界引起了巨大的反响，网络上也有许多讲解的教程，这一篇图解Swin Transformer就写得非常棒，另外强烈推荐Zhu老师在b站的精读视频：Swin Transformer论文精读【论文精读】_哔哩哔哩_bilibili。

这篇文章就来记录一下我对 Swin Transformer 的理解。

首先，在 PVT (Pyramid Vision Transformer) 中我已经分析过 ViT 存在的一些缺陷，本文就不再讨论，简单来说就是网络设计的问题和自注意力机制的显存占用问题。下面直接讨论 Swin Transformer 的网络结构。

Swin Transformer 的网络结构

金字塔结构

与 PVT 一样，Swin Transformer 在网络的设计上也实现了金字塔的结构，如下图所示（实际上下面这张图就是从 PVT 论文中摘取的）。

网络结构对比

上图从左到右，我们可以看到

• CNN：金字塔结构，网络越深，feature map 尺寸越小，channel 数越多。
• ViT：柱状结构，单一的 feature map 尺寸。
• Swin Transformer：金字塔结构，网络越深，feature map 尺寸越小，channel 数越多。

这样设计的好处就是能够像 CNN 网络一样处理不同尺度的特征。

Swin Transformer 的总体结构

Swin Transformer 总体结构

从上图我们可以观察到在输入端有一个 Patch Partition 的操作，也就是 Vision Transformer 常规的切图。然后是经过一个线性映射进入第一个 Swin Transformer Block，从而完成 Stage 1 这个模块。Stage 2 的输入特征图大小为 H4×W4×2C\frac{H}{4} \times \frac{W}{4} \times 2C ，之后每个 stage 都会把特征图大小减小一半，通道数增加一倍，也就实现了所谓的金字塔结构。

注意上图 Stage 1 里面的 Swin Transformer Block 下面写了个 x2，这个表示的是这个模块中有两个部分，分别是 Window-based Multi-head Self-attention (W-MSA) 和 Shifted Window-based Multi-head Self-attention (SW-MSA)，这个 SW-MSA 是紧跟着 W-MSA 的，属于串联的关系。

在上图 Stage 3 里面我们可以看到 x6，这代表了 Stage 3 有3组 W-MSA + SW-MSA 串联在一起。

W-MSA & SW-MSA

W-MSA

从 W-MSA 说起，它的设计主要是为了解决 Vision Transformer 的自注意力机制显存占用高的问题。顾名思义，Window-based Multi-head Self-attention 就是把自注意力机制限制在了一个窗口中。

如下图所示，假设输入特征图的大小为 H×W=56×56H \times W = 56 \times 56 ，num_patches 为 8×88 \times 8 ，每个 patch 的大小为 7×77 \times 7 ，在这个设定下，MSA 计算时 Q×KQ \times K 会涉及到 8×8=648 \times 8 = 64 个 patches 的相乘。相对而言，W-MSA 将自注意力的计算限制在了每一个窗口中，也就是下图的绿色窗口中，每个窗口中的 patch 只会和同一个窗口内部的其他 patch 计算注意力，Q×KQ \times K 的计算只会涉及到窗口中 4×4=164 \times 4 = 16 个 patches 的相乘。这样就限制了 Q×KQ \times K 运算时的显存占用。

W-MSA

这里分别对 MSA 和 W-MSA 的自注意力机制的计算复杂度做一个估计。

如下图所示，对于 MSA：

• 将大小为 HWCHWC 的输入特征图通过全连接层转换成大小为 HWCHWC 的 、、Q、K、VQ、K、V ：3HWC23HWC^2
• Q×KQ \times K 计算得到 Attention Matrix： (HW)2C(HW)^2C
• Attention Matrix 乘以 V： (HW)2C(HW)^2C
• 最后再接一个全连接层： HWC2HWC^2

全部加起来就是 4HWC2+2(HW)2C4HWC^2+2(HW)^2C ,相对于输入特征图大小而言计算复杂度就是 O(N2)O(N^2) 。

对于 W-MSA：

• 前面 MSA 的推导中所有的 HWHW 都被限制在了一个尺寸为 M×MM \times M 的窗口中
• 对于一个窗口中的计算复杂度，我们可以用 M2M^2 替代 4HWC2+2(HW)2C4HWC^2+2(HW)^2C 式子里的 HWHW ，得到 4M2C2+2(M)4C4M^2C^2+2(M)^4C

一共有HWM2\frac{HW}{M^2} 个窗口，所以总的计算复杂度是 HWM2×(4M2C2+2(M)4C)=4HWC2+2M2HWC\frac{HW}{M^2} \times (4M^2C^2+2(M)^4C) = 4HWC^2 + 2M^2HWC 。因为 MM 为窗口的尺寸，是一个较小的固定值，取决于我们对网络的设计，所以可以看成一个常数，在这种情况下相对于输入特征图大小而言总的计算复杂度就是 O(N)O(N) 。

W-MSA 计算复杂度

W-MSA 把 MSA 的计算复杂度从 O(n2)O(n^2) 降低到了 O(n)O(n) ，使得 Vision Transformer 的显存占用不再是个瓶颈。然而 W-MSA 也不是完美的，将注意力限制在窗口中也必然导致全局注意力丧失的问题，为了解决这个问题，作者在 W-MSA 后面加了一个 SW-MSA。

SW-MSA

从下图中可以很直观地看到，SW-MSA 的提出主要是为了补充位于不同窗口中的 patch 之间的注意力。

W-MSA（左）SW-MSA（右）

SW-MSA 最有意思的地方在于其代码的实现上。先说一下 W-MSA 在代码上的实现，其实是很直观的，首先是把特征图平均切分一下，然后在每个窗口内部做自注意力的计算。然而到了 SW-MSA 时，就不是平均切分特征图了，从上图可以看到需要将其切分成好几个尺寸不一的窗口，再在几个大小不一的窗口中计算自注意力，感觉实现起来会很麻烦的样子。Swin Transformer 的作者就想到能不能用同一套代码实现 W-MSA 与 SW-MSA 呢？答案是肯定的，只是要对 SW-MSA 做一些额外的处理。

如下图的中间部分所示，作者将不同的窗口做了平移，例如把 C 部分下移，把 B 部分右移，这样就把下图的左边那个样子转换成了下图的右边那个样子，也就是变成和 W-MSA 的切图一样的4个窗口了，也就可以直接复用 W-MSA 的自注意力计算的代码了。

当然，还需要解决另外一个问题，举个例子，下图的右边部分，3和6被分到了同一个窗口中，那么被标记成3的区域的 patch 在计算注意力的时候就会和被标记成6的区域的 patch 有交集。然而，我们只想要3和3的注意力、以及6和6的注意力，所以需要借助 mask 来遮挡我们不需要的注意力。

Cyclic Shift

还是拿区域3和区域6这个窗口举例，在下图中，假设输入特征图 H=W=14 patchesH=W=14 \text{ patches} ，我们先把区域3和区域6里面的每一个 patch 拿出来，3里面有 4×7=284 \times 7 = 28 个 patches，6里面有 3×7=213 \times 7 = 21 个 patches，一共是 28+21=4928+21=49 个 patches。

计算自注意力的时候，这49个 patches 要和自己的 transpose 相乘，得到一个尺寸为 (28+21)×(28+21)(28+21) \times (28+21) 的 attention matrix。如下图所示，attention matrix 里左上角的部分是 28×2828 \times 28 个区域3和区域3里的 patch 之间计算的注意力，右上角的部分是 28×2128 \times 21 个区域3和区域6里的 patch 之间计算的注意力，左下角的部分是 21×2821 \times 28 个区域6和区域3里的 patch 之间计算的注意力，右下角的部分是 21×2121 \times 21 个区域6和区域6里的 patch 之间计算的注意力。

而我们实际上只想保留 attention matrix 里左上角区域3和区域3里的 patch 之间计算的注意力，以及右下角区域6和区域6里的 patch 之间计算的注意力。所以要给 attention matrix 加上一个 mask matrix，在不要的部分值为-100，在保留的部分值为0，这样后面通过一个 softmax 的时候，值为-100的部分就为趋于0。

对于其他窗口的计算也是类似的，这里就不再赘述，感兴趣可以参考一下这篇文章：图解Swin Transformer。

Masking in SW-MSA

W-MSA + SM-MSA 对比 MobileNet

本质上通过 W-MSA + SW-MSA 替代 MSA 是一个减少计算量的工作，其思想有点像 MobileNet 用深度卷积+逐点卷积代替普通的卷积。

如下图所示，MobileNet 先通过深度卷积核处理特征图，每个深度卷积核只作用于单个通道，再用逐点卷积补充通道之间的信息交流。

深度卷积+逐点卷积

相对位置编码

上文中对 W-MSA + SM-MSA 的讨论其实还缺了一部分，我们看自注意力的计算公式 Attention(Q,K,V)=Softmax(QKTd+B)V\text{Attention}(Q, K, V) = \text{Softmax}(\frac{QK^T}{\sqrt{d}} + B)V ，前文的讨论只涉及了 Q、K、V，这里面还有个 B，就是相对位置编码。

如下图所示，假设 W-MSA 的窗口尺寸为 2×22 \times 2 ，我们自注意力的计算就是在这个窗口内完成的。

W-MSA 的窗口

Swin Transformer 的作者在设计相对位置编码的时候也考虑到了节约计算资源，首先，通过如下代码生成一个可学习的 relative positional embedding table。

relative_position_bias_table = nn.Parameter(torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads))

也就是对于每一个 attention head，都生成一个 (2⋅window size−1)×(2⋅window size−1)(2 \cdot \text{window size} - 1) \times (2 \cdot \text{window size} - 1) 大小可学习的相对位置编码矩阵，在我们的例子中， window size=2\text{window size}=2 ，所以这个相对位置编码矩阵的大小为 3×33 \times 3 。这个相对位置编码矩阵包含了我们尺寸为 2×22 \times 2 的窗口中所有 patch 之间可能存在的相对位置，例如上下一格的关系，左右一格的关系，等等。

相对位置编码矩阵

接着，我们还需要一个 index 来检索相对位置编码矩阵中的值。举个例子，假设我们有一个相对位置编码 index table，如下图中左边那个小图所示，里面的元素代表了不同 patch 之间相对位置对应的 index 值。

例如， q0q_0 和 k2k_2 之间的关系，查看下图右边那个小图，就是上下的关系（0号 patch 在2号 patch 上边），其在下图左边那个小图中的 index 是1。同样的，q1q_1 和 k3k_3 之间的关系，查看下图右边那个小图，也是上下的关系（1号 patch 在3号 patch 上边），其在下图左边那个小图中的 index 也是1。这就是说， q0q_0 和 k2k_2 之间的相对位置编码 index 与 q1q_1 和 k3k_3 之间的相对位置编码 index 相同，都是1，我们就回到上图中的大小为 3×33 \times 3 的相对位置编码矩阵中寻找1号元素的值作为他们的相对位置编码。

相对位置编码矩阵检索

这个大小为 3×33 \times 3 的相对位置编码矩阵好理解，就是一个可学习的矩阵，那么这个相对位置编码 index table 如何获取呢？

相对位置编码 index table

首先，还是假设 W-MSA 的窗口尺寸为 2×22 \times 2 ，我们可以分别标出这个窗口中每个 patch 的x轴坐标和y轴坐标，如下图所示。

W-MSA 的窗口中patch的坐标

计算这个窗口中每个 patch 间所有相对位置的一个好方法就是用所有 patch 的x、y坐标做差值。利用广播机制，我们分别用x轴坐标向量减去x轴坐标向量的转置以及y轴坐标向量减去y轴坐标向量的转置。

坐标向量相减

得到下图的相对位置编码 index 矩阵。

坐标向量相减结果

然而，我们想要的相对位置编码 index 矩阵是可以从里面获取相对位置编码的 index 值，显然上图中的矩阵里的元素并不是合理的 index 数值。所以还要进一步处理。

首先，index 最好是从0开始，不应该有负值，于是我们做如下操作：

去负值

其次，x和y方向的相对位置编码应该是不一样的，于是我们对x坐标单独做如下操作：

x坐标处理

最后再相加一下，就得到了最后的相对位置编码 index 矩阵：

相加

到目前为止，Swin Transformer Block 里的主要内容就讨论完了，再看一下 Swin Transformer 的结构图，我们会发现 Swin Transformer Block 出来接了一个 Patch Merging 模块，上文中讨论的 W-MSA 和 SW-MSA 都是为了解决 Vision Transformer 显存占用的问题，而这个 Patch Merging 解决的是 Vision Transformer 另外一个主要问题，也就是通过改变特征图尺寸使得 Vision Transformer 可以处理不同尺度的特征。

Patch Merging

Patch Merging

Patch Merging 和 Pooling 非常相似。我们先对大小为 HWCHWC 特征图进行窗口大小为 2×22 \times 2 的提取，如下图所示。

Patch Merging 切分

接着就是把不同编号的 patch 分别提取出来，得到4个尺寸为 H2×W2×C\frac{H}{2} \times \frac{W}{2} \times C 的特征图，然后再将这个4个特征图在 channel 方向拼接起来，得到尺寸为 H2×W2×4C\frac{H}{2} \times \frac{W}{2} \times 4C 的特征图。

Patch Merging 切分拼接

注意下图 Stage 2，特征图经过 Patch Merging + Swin Transformer Block 后尺寸除以2，通道数乘以2，而上图中尺寸除以2，通道数却数乘以4。

Patch Merging

所以我们需要将其拉直，经过一个 Linear 层降维，再 Reshape 成二维，最后尺寸从 H2×W2×4C\frac{H}{2} \times \frac{W}{2} \times 4C 变成了 H2×W2×2C\frac{H}{2} \times \frac{W}{2} \times 2C 。

Patch Merging Linear + Reshape

对比 PVT

对 Swin Transformer 结构设计做一个小结：

• 通过 Patch Merging 实现了类似于 Pooling 的操作逐层减小特征图尺寸
• 在 Swin Transformer Block 中通过 W-MSA + SW-MSA 解决了显存占用问题

我们再对比一下 PVT，发现 PVT 也是用不同的方法解决了同样的问题：

• 在每一个 Block 开头进行切图 + Linear，逐层减小特征图尺寸
• 通过 Spatial Reduction Attention (SRA) 减少 Attention 计算的显存占用

PVT 结构

在减少显存占用这方面，Swin Transformer 其实更胜一筹，PVT 提出的 SRA 将 K、V 的尺寸从 HWCHWC 降低到 HWR2C\frac{HW}{R^2}C ，从而将计算复杂度从 O(n2)O(n^2) 降低到 O(n2R2)O(\frac{n^2}{R^2}) 。（本来在估计计算复杂度的时候 R2R^2 是应该忽略的，不过在这里缩小系数 R2R^2 相对于输入特征图尺寸 HWHW 是一个较大的值）

然而，显然在特征图尺寸 HWHW 较大时，PVR 的计算复杂度 O(n2R2)O(\frac{n^2}{R^2}) 还是会比 Swin Transformer 的 O(n)O(n) 高很多的。

PVT 的 SRA

实验对比

ImageNet-1K

如下图所示，在分类任务上，Swin Transformer 无论是 FLOPs 还是准确率都是优于其他两个 Vision Transformer 的算法的（ViT 与 DeiT）。

然而，对比起卷积神经网络，似乎没有明显优势，感觉比 EfficientNet 还要弱一些。估计 Swin Transformer 的优势不在分类任务。作者在论文里说的是“Compared with the state-of-the-art ConvNets, i.e. RegNet and EffificientNet, the Swin Transformer achieves a slightly better speed-accuracy trade-off”。

ImageNet-1K

ADE20K

在 ADE20K 分割任务是，Swin Transformer 几乎是横扫了。这也得益于其金字塔结构的设计，使其能够胜任分割这种像素级的任务。

ADE20K

另外，论文里还有关于相对位置编码和 SW-MSA 的消融实验，感兴趣可以去读一下论文，这里就不在赘述。

CNN 和 Swin Transformer 的异同

Swin Transformer 这篇论文的提出是里程碑式的，使得 Vision Transformer 在处理计算机视觉任务时不再有明显短板，也引发了许多关于CNN 和 Swin Transformer 的异同的思考。

相似点

• 操作都在 window 里（卷积核通常为 3×33 \times 3 到 7×77 \times 7 大小的窗口 vs W-MSA）
• 都是金字塔结构
• 都是 pretraining + finetuning 的 pipeline
• transfomer 的切图甚至都可以用 Conv2d 写（kernel_size=stride=patch_size）

nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)

不同点

CNN 有很强的 inductive bias（归纳偏置），也就是针对不同的特定的问题，我们认为模型应该有某些特点，从而做出一系列针对模型设计的人为限制：

• 处理图片时，每个位置的信息与周围的信息相关，因此设计出 Conv
• 处理 NLP 任务时，输出的结果与单词的顺序有关，因此设计出 RNN
• CNN 空间平移不变性的特点
• CNN 网络底层处理简单的特征，逐层进行特征聚合处理的特征越来越复杂

Transformer 相对于 CNN 的 inductive bias 就少很多：

• Patch 内信息建模 --> MLP
• Patch 间信息建模 --> MSA

由于 Transformer 的 inductive bias 少，所以在处理图片任务时就需要更多的数据去训练，但是可能上限会更高。