Swin Transformer原理梳理

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

Transformer 是一种和CNN不同思路的机器视觉方法,开始被应用于自然语言处理,最近几年被引入视觉方向,发展十分迅速,但一直没有超过CNN,直到Swin Transformer的出现,以它作为backbone的网络霸榜了视觉的多个领域。我们来一起学习一下它的结构和原理。

单头注意力(Single-headed attention)

首先要看的是单头注意力模型,它是transformer模型的基础。也是和CNN卷积运算最大的不同。

$$y_{i j}=\sum_{a, b \in \mathcal{N}_{k}(i, j)} \operatorname{softmax}_{a b}\left(q_{i j}^{\top} k_{a b}\right) v_{a b}$$
$ \text { queries } q_{i j}=W_{Q} x_{i j}, \text { keys } k_{a b}=W_{K} x_{a b}, \text { and values } v_{a b}=W_{V} x_{a b} $
是位置ij和其相邻像素的线性变换。如图1所示

1

图1

多个注意头(multi-headed attention)用于学习输入的多个不同表示。将$x_{ij}$沿着深度拓展到N组,使用不同的变换分别计算每个组的单头注意力参数,再将输出表示连接到最终输出中。

早期的实验表明,使用相对位置嵌入可以显着提高精度。这样每个元素有两个距离,行偏移和列偏移。如图2所示

2

图2

这种相对空间的注意力定义为

$$\begin{equation} y_{i j}=\sum_{a, b \in \mathcal{N}_{k}(i, j)} \operatorname{softmax}_{a b}\left(q_{i j}^{\top} k_{a b}+q_{i j}^{\top} r_{a-i, b-j}\right) v_{a b} \end{equation} $$

$r_{a-i, b-j}$是行列偏移量相加。

Vision Transformer (VIT)

然后看一下VIT,它是Swin Transformer借鉴最多的一篇文章。 ViT原理中是将$H×W×C$维的图像,分解成一系列$N×(P^2·C)$维的图像。HW是原图像的尺寸,C代表图像的通道数。$(P,P)$ 是分解后小图像的分辨率;$N = HW/P^2$是 分解的图像个数,即Patches个数,是Transformer的有效输入序列长度。 然后将每个Patch压平(记作tokens),这部分功能就叫patch embeddings。如图3

3

图3

每一个Patch经过嵌入(embeddings)后,在经过batch标准化LayerNormal,然后经过上述的多头注意力算法,经过类似ResNet的前项短路结构,再经过归一化,再经过多层感知机MLP,这部分称作转换编码Transformer Encoder。

Vision Transformer通常只适合于大数据集。它采用的是固定的Patch分解方法。

Swin Transformer

Swin Transformer首先通过ViT类型的方法将输入RGB图像分割成不重叠的patches。每个patch被视为一个“标记”,即tokens,它的特征被设置为原始像素RGB值的串联。文中使用4×4的patch大小,因此每个patch的特征维数为4×4 × 3(通道数) = 48。 注意力模型修改为Swin Transformer。即,将标准多头自注意(MSA)模块替换为基于移位窗口的模块SW-MSA或W-MSA,即不计算全局注意力,而是计算窗口内的注意力。如图4。 这里要注意,原始的代码只接受特定尺寸的图像,自己定义的尺寸很难给出一个合适的window size和patch size。而且无论是window size和patch size都需要是正方形的。 4

图4

Swin Transformer Block是Swin Transformer的基本单元,它由MLP、Layer Normal、W-MSA或SW-MSA以及前项连接结构(shotcut)构成。W-MSA和SW-MSA构成的SWin Block是成对出现的,所以模型结构中的SWin Block总是偶数。W-MSA是常规的多头注意力模块,与ViT中类似,只不过是在窗口内计算注意力。SW-MSA是移动窗口后的多头注意力模块,移位操作的尺寸是按照窗口尺寸的一半,这样可以让原来划分的窗口产生交叉,防止计算注意力时忽视跨窗口的特征。在原始特征上应用上面提到的线性嵌入层,使用4×4的patch大小,因此最初分解的序列patch数为(H/4×W/4),在自然语言处理中一般叫“标记”tokens。根据前面TiV原理,压平变为1D处理后其投射维度记为C。用Swin Transformer block处理这些tokens, 为了产生分层表示,随着网络加深,将部分patchs合并,来减少tokens的数量。第一个合并层 先将每组相邻的2 × 2特征相连。这样就减少了4倍的tokens的数量,相当于2倍下采样。而且此时特征维度变为4C,再应用一个线性层,将特征维度输出变为2C。再用Swin Transformer block进行特征变换,整图分辨率变为(H/8×W/8),同样再经过Swin Transformer 变为(H/16×W/16). 这样就产生了和传统CNN相同的分层表示方法。如图5所示,这和b图中ViT的方法是不同的。

5

图5

Swin Transformer 的每个stage是怎么样的工作过程才能产生上面所说的维数变化呢?其实就是简单的下采样,将相邻图片融合,然后在用前面说的两个Transformer串联实现的。具体过程为:第一个模块使用常规的窗口划分策略,将8 × 8 feature map均匀划分为2×2个大小为4×4的窗口,如图6所示,此时每个窗口的边长M=4。下一个模块使用不同的模块(交替),将规则的窗口用(M/2, M/2)置换(每个大的规则的部分让出M/2的长和高,如图6)。

6

图6

这时候有一个问题,窗口数比原来多了,长宽方向都多了一个窗口。比如原来是$w/M\times h/M$,就会变成$(w/M+1) \times (h/M+1)$。而且有一部分窗口肯定是小于M×M的,这时候无法进行后续的注意力计算。如果按简单的方法可以直接把不够M的宽度或高度直接用空像素填充成M。但是这种方法比较浪费计算资源。这里文章提出了一种提更高效的填充方法,即向左上角循环移动,如图7,将大窗向左上角移动(小窗补到右下方)后,重新构成了规则窗口。这时候一个批处理窗口可能由几个在特征映射中不相邻的子窗口组成,因此在计算注意力的时候采用掩码将不同的特征屏蔽掉,只计算他们原来的子窗的注意力。这时候对于这一个阶段来说,相邻窗口合并再计算注意力以后就成了1/2下采样,而且经过这种窗口变换,实现了跨窗口的连接。

7

图7

要注意Swin Transformer所采用的自关注Self-attention方法,复杂度相比之前ViT的得到了降低。因为它是在每个Patch中计算的(ViT是将所有的Patch压缩,然后一起放入Transformer encoder中,相当于整幅图一起计算的),这样在高分辨率图像中也同样能够不断分解处理,而ViT处理大图像就有一些困难。

在计算self-attention时,采用的公式和上面的相对空间的注意力定义差不多。

$$\begin{equation} \operatorname{Attention}(Q, K, V)=\operatorname{SoftMax}\left(Q K^{T} / \sqrt{d}+B\right) V \end{equation}$$

Q, K, V 分别是query, key 和 value的矩阵;d是 query/key的维度。

连续的Swin Transformer blocks 计算方法用数学表示应该是

$$\begin{equation} \begin{array}{l} \hat{\mathbf{z}}^{l}=\operatorname{W-MSA}\left(\operatorname{LN}\left(\mathbf{z}^{l-1}\right)\right)+\mathbf{z}^{l-1}, \\ \mathbf{z}^{l}=\operatorname{MLP}\left(\operatorname{LN}\left(\hat{\mathbf{z}}^{l}\right)\right)+\hat{\mathbf{z}}^{l}, \\ \hat{\mathbf{z}}^{l+1}=\operatorname{SW-MSA}\left(\operatorname{LN}\left(\mathbf{z}^{l}\right)\right)+\mathbf{z}^{l}, \\ \mathbf{z}^{l+1}=\operatorname{MLP}\left(\operatorname{LN}\left(\hat{\mathbf{z}}^{l+1}\right)\right)+\hat{\mathbf{z}}^{l+1} \end{array} \end{equation}$$
式中,$\hat{\mathbf{z}}^{l}$是W-MSA/SW-MSA的输出,$\mathbf{z}^{l}$是多(双)层感知机(MLP)的输出。

最后,Swin Transformer的几种backbone参数

  • Swin-T: C = 96, layer numbers = {2,2,6,2}
  • Swin-S: C = 96, layer numbers ={2,2,18,2}
  • Swin-B: C = 128, layer numbers ={2,2,18,2}
  • Swin-L: C = 192, layer numbers ={2,2,18,2}

C是第一层隐含层的通道数channel,layer numbers是图4中每层的加载个数。

相关阅读:Swin Transformer 代码解析

See Also