【图像识别】Swin Transformer

一、引言

论文： Swin Transformer: Hierarchical Vision Transformer using Shifted Windows
作者： Microsoft Research Asia
代码： Swin Transformer
特点： 提出滑动窗口自注意力 (Shifted Window based Self-Attention) 解决Vision Transformer输入高分辨率图像注意力模块计算复杂度高的问题。

二、框架

Swin Transformer的整体框架图如下：

可见，Swin Transformer主要包括Patch Partition+Linear Embedding、Swin Transformer Block、Patch Merging几个组件。每经历一个Stage特征图都会缩小2倍、通道数放大2倍，这与ResNet网络的特征尺度变化过程很类似。

Swin Transformer Block的最小单位为上图右侧两个连续的Block，两个Block的差别仅在于注意力模块是窗口多头自注意力 (W-MSA) 还是滑动窗口多头自注意力 (SW-MSA)。所以每个Stage中包含的Block数量都是2的整数倍。

2.1 Patch Partition+Linear Embedding

Patch Partition是要将图像拆分成多个$4\ast 4\ast 3$大小的Patch（4*4是每个Patch包含的像素点数，3是原始图像通道数），如下图所示：

之后将每个Patch按行展平，于是每个Patch形成一个$1\ast 1\ast 48$的像素点。$48=4\ast 4\ast 3$，原本一个像素点通道数为3，一个Patch有$4\ast 4$个像素点，所以展平后一个Patch变成一个通道数为48的像素点。所有Patch展平后按原位置拼接则形成$\frac{H}{4}\ast \frac{W}{4}\ast 48$的特征图。

通过Linear Embedding（例如核为$1\ast 1$的2D卷积）将通道数从48映射到$C$（原文中$C=96$）。最后通过LayerNorm将所有通道数为$C$的像素点归一化。

上图为一个Patch上16个像素点的Patch Partition+Linear Embedding过程（图中一个小方格代表一个有3个通道的像素点，$c=\frac{C}{3}$）。

事实上，Patch Partition+Linear Embedding就是要将图像尺寸从$H\ast W\ast 3$转换为$\frac{H}{4}\ast \frac{W}{4}\ast C$。 与Vision Transformer取Patch时采取的策略一致，Swin Transformer并没有分两步完成该操作，而是通过一个卷积核大小为$4\ast 4$、步长为4、输出通道数为96的2D卷积层进行特征提取，再通过LayerNorm对所有token进行归一化（一个token对应$\frac{H}{4}\ast \frac{W}{4}\ast C$的特征图中的一个像素点）。

2.2 Swin Transformer Block

下图为Vision Transformer（左）和Swin Transformer（右）中Block的结构图：

其中，Swin Transformer中也有DropPath但是没有画出来，并且他们的MLP模块均为如下结构：

可见，两者结构基本一致，主要区别如下：
(1) Vision Transformer的每一个Block都一样，而Swin Transformer中两个不同的Block是一个基本单位，两个Block顺序连接，不会单独出现。
(2) Vision Transformer使用多头自注意力，而Swin Transformer中两个Block注意力模块不同，前一个使用窗口多头自注意力，后一个使用滑动窗口多头自注意力。
(3) Vision Transformer使用绝对位置编码，而Swin Transformer使用相对位置编码。

所以我们下面主要介绍窗口多头自注意力和滑动窗口多头自注意力，其余部分请参考我关于Vision Transformer的博客。

2.2.1 窗口多头自注意力

引入窗口多头自注意力 (W-MSA) 是为了减少计算复杂度。

如下为多头自注意力 (MSA) 与窗口多头自注意力的窗口划分方法：

多头自注意力不对特征图进行划分，左图中共16个像素点（也叫token），每个像素点作为query时都要与所有像素点进行自注意力计算。而窗口多头自注意力会先将特征图拆分为一个个窗口（右图中窗口高宽为$2\ast 2$），然后在每个窗口内部单独进行自注意力计算。

两者的计算复杂度如下：

其中，$h$为特征图高度，$w$为特征图宽度，$C$为特征图通道数，$M$为每个窗口的大小；自注意力的计算是以特征图或窗口上的token为单位的，所以未被拆分的特征图中应包含$hw$个长度为$C$的token。该计算复杂度不包括除以$\sqrt{d_k}$和SoftMax产生的计算量。

在讲解两者计算复杂度推导过程前，我们需要明确两个矩阵相乘所产生的计算量，如果$A$矩阵高宽为$a\ast b$，$B$矩阵高宽为$b\ast c$，则$A\cdot B$的计算量为$a\ast b\ast c$。

下图实线部分为单头自注意力的计算过程：

其中，黑色字体为当前矩阵含义，红色字体为当前矩阵高宽，绿色和橙色字体为当前运算复杂度。多头自注意力在通道上对token的拆分不会影响计算复杂度，所以除最后将$A\cdot V$通过$W_o$映射为输出时产生的计算复杂度（上图右侧虚线部分），其余与单头自注意力计算复杂度相同。于是，我们可以得到多头自注意力的计算复杂度为$4hw{C}^2+2(hw{)}^{2}C$。

对于窗口多头自注意力，特征图大小为$h\ast w$，窗口大小为$M\ast M$时，则从特征图上可拆分出$\frac{h}{M}\ast \frac{w}{M}$个窗口。每个窗口单独执行多头自注意力产生计算复杂度为$4M^2{C}^2+2M^{4}C$。则$\frac{h}{M}\ast \frac{w}{M}$个窗口的总计算复杂度为$4hw{C}^2+2M^{2}hwC=\frac{h}{M}\ast \frac{w}{M}\ast (4M^2{C}^2+2M^{4}C)$。

两者计算复杂度差别在后半部分，以原文中的参数$h=56,w=56,C=96,M=7$为例，多头自注意力后半部分计算复杂度为$2(hw{)}^{2}C=2(56\ast 56{)}^2\ast 96=1888223232$，窗口多头自注意力后半部分的计算复杂度为$2M^{2}hwC=2\ast 7^2\ast 56\ast 56\ast 96=29503488$。两者相差64倍。

此外，Swin Transformer使用相对位置编码来避免注意力运算未考虑token顺序的问题。与绝对位置编码直接加在输入特征图上不同，相对位置编码以偏执形式存在参与注意力权重的计算，其公式如下：
$$Attention(Q,K,V)=SoftMax(\frac{Q{K}^T}{\sqrt{d_k}}+B)V$$

详情请参考我关于绝对位置编码和相对位置编码的博客。

2.2.2 滑动窗口多头自注意力

采用窗口多头自注意力时，每个窗口都是独立的，会导致窗口间的相关信息无法被捕捉。例如一个目标处于多个窗口上，窗口独立进行注意力计算就导致该目标各个部分的相关信息无法被考虑。Swin Transformer引入滑动窗口来解决这一问题。

2.2.2.1 窗口滑动

对于$M\ast M$的窗口，滑动方向为右下，向右移动$\lfloor \frac{M}{2}\rfloor$个像素点，向下移动$\lfloor \frac{M}{2}\rfloor$个像素点。

下图是窗口高宽为$4\ast 4$时窗口滑动前后的窗口分割情况（因为是向右下滑动，所以我在图中增加了左上方的窗口）：

可见，经滑动之后的窗口中可能包含滑动前多个窗口的内容。例如，滑动后的中心窗口就包含滑动前的4个窗口中的内容。该操作从一定程度上缓解了相关信息完全丢失的问题。

2.2.2.2 窗口移位与掩码

不过，滑动后各个窗口大小不一，只有在padding后才能组成tensor进行并行运算。以上图为例，原本只有4个$4\ast 4$的窗口，padding后会产生9个$4\ast 4$的窗口，这无疑增加了很多计算负担。Swin Transformer通过移位和掩码解决了该问题。

移位流程如下：

可见，$A,B,C$均为左上方新增窗口向右下滑动时产生的窗口，因此$A,C$的高度为$\lfloor \frac{M}{2}\rfloor$，$A,B$的宽度为$\lfloor \frac{M}{2}\rfloor$，$C$的宽度为$w-\lfloor \frac{M}{2}\rfloor$，$B$的高度为$h-\lfloor \frac{M}{2}\rfloor$。

这样，所有新的窗口大小都是统一的$M\ast M$。但是这引入了另一个问题，同一窗口内包括原本不应相互计算注意力的内容。 例如上图中的右下角窗口，同时包含$A,B,C,D$四个部分的内容，但$A$原本在左上角，与$B$的下侧、$C$的右侧、$D$的右下侧一般是没有联系的。就像一张仅包含天空和草地的图片，$A,C$是天空，$B,D$是草地，$A$不必与$B,D$进行注意力计算。

于是Swin Transformer又引入了掩码。首先给移位后的各个窗口一个编号，如下图：

然后给每个窗口每个像素点一个掩码矩阵，如下图（以右下角窗口中两个像素点为例）：

掩码矩阵中仅有两种值，0和-100。与查询像素点属于同一编号的像素点对应掩码为0，否则为-100。窗口注意力会计算每个像素点与所有像素点间的相似度，所以窗口中有几个像素点就有几个掩码矩阵，每个掩码矩阵都与窗口大小一致。

这个掩码矩阵像相对位置偏执一样被施加在注意力计算相似度的公式中，如下：
$$Attention(Q,K,V)=SoftMax(\frac{Q{K}^T}{\sqrt{d_k}}+B+Mask)V$$

我们不考虑相对位置偏执$B$，相似度$\frac{Q{K}^T}{\sqrt{d_k}}$+掩码$Mask$经SoftMax产生效果如下：

可见，对于橙色框中像素点来说，同编号的像素点与其的相似度没有被改变，其它像素点相似度被减去100。因相似度矩阵中的值都比较小，减去100相当于变成一个很小的数，经SoftMax后其值接近于0。

在各个窗口上计算完多头自注意力后，按照之前移位流程逆向复原，滑动窗口多头自注意力就完成了。

综上，移位操作使并行运算成为可能，掩码操作使注意力计算不跨窗口。

2.3 Patch Merging

Swin Transformer Block不改变特征图尺寸，Patch Merging起到使特征图缩小2倍、通道数放大2倍的作用。

以大小为$4\ast 4$、单通道的特征图为例，Patch Merging的过程如下图：

无论特征图是多大，Patch Merging时Patch的高宽均为$2\ast 2$。截取各Patch同相对位置的像素点，之后将截取结果拼接到一起形成新的特征图。假设原特征图为$[H,W,C]$，则中间特征图为$[\frac{H}{2},\frac{W}{2},4C]$。 之后对中间特征图每个像素点所有通道分别执行归一化，并经过一次全连接线性映射将每个像素点的通道数减半。最后特征图为$[\frac{H}{2},\frac{W}{2},2C]$。

致谢：

本博客仅做记录使用，无任何商业用途，参考内容如下：
Swin-Transformer中MSA和W-MSA模块计算复杂度推导
Swin Transformer论文精读【论文精读】
Swin-Transformer网络结构详解