Day 12: Twin Transformer by 美团

这是美团和澳大利亚阿德莱德大学联合发表的新文章，也是和 Transformer 相关的，以下是一些要点。

• Swin Transformer 的 Shifted Windows 虽然有效，但是由于尺寸不同，因此在用现有的深度学习模型来实现的时候会有一定的困难，因为大部分框架都希望输入的尺寸是相同的。
• 另一种方案是 PVT 中提出来的。在普通的自注意力模块中，每个 query 都需要和所有的输入 token 计算注意力，但 PVT 中，每个 query 只需要和一些取样出来的输入 token 计算注意力。虽然理论上它的计算复杂度仍然是二次方的，但已经是在实际中可以接受的了。
• 上面提到的视觉 Transformer 的核心是如何设计空间注意力（spatial attention），作者发现在 PVT 中提出的全局采样注意力（global sub-sampled attention）非常有效，而且当与位置编码一起使用的时候，它的效果可以与SOTA的视觉 Transformer 的表现持平，甚至更高。
• 根据这些，作者提出了第一个模型架构：Twin-PCPVT。他们提出了一个仔细设计且简单有效的空间注意力机制，用来替换 PVT 中低效的全局采样注意力。这种注意力机制的灵感来源于可分离深度卷积（separable depthwise convolutions），因此作者将它命名为：空间可分离自注意力（spatially separable self-attention SSSA）
• SSSA 由两种注意力操作组成，一种是局部分组自注意力（locally-grouped self-attention LSA），另一种是全局下采样注意力（global sub-sampled attention GSA）。其中 LSA 捕捉比较精细的短距离的信息，而 GSA 捕捉长距离和全局的信息。
• 两种注意力操作都有效且实现起来很简单，这又带来了第二种结构：Twin-SVT

Related Work

这一章节的前两部分没有太多可讲，因为灵感主要是来源于后面两个部分，分组卷积和可分离卷积，以及位置编码。

分组卷积及可分离卷积

分组卷积最早在 AlexNet 中提出，二者在加速网络训练上都非常有效。深度卷积（depthwise convolution）则是分组卷积的一个极端例子，也即所分的组数与通道数相等，再后接一个通道卷积（point-wise convolution），将不同通道之间的信息聚合起来。

位置编码

本文用的位置编码是 PVT 中提出的条件位置编码（conditional positional emcodings），是根据输入进行动态调整的一种编码，它的表现要比绝对位置编码和相对位置编码好。

方法细节

作者提出两种简单且强大的视觉 Transformer 上的空间设计。第一种是基于 PVT 和 CPVT 的，只使用了全局注意力，因此称之为 Twins-PCPVT。第二种是基于上面提出的 SSSA ，交替使用局部注意力和全局注意力，称之为 Twin-SVT。

Twins-PCPVT

PVT 引入了金字塔型多层级设计，来更好地处理密集预测任务，比如目标检测和语义分割。它采用了 ViT 和 DeiT 里的绝对位置编码，而且所有层都使用了全局注意力机制，并依靠空间上的压缩来减少处理整个序列时的计算量。对于 Swin Transformer 的表现强于 PVT，作者是很惊讶的。

后来作者发现，PVT 的表现不佳仅仅是因为绝对位置编码的问题，因为它会使模型在处理不同尺寸的输入时遇到些麻烦。同时，这种编码方式也打破了平移不变性。然而 Swin 使用的是相对位置编码，因此规避了这个问题。

我们用在 CPVT 中提出的条件位置编码来替换掉原始的绝对位置编码。条件位置编码（CPE）是基于输入的，能天然地规避掉绝对位置编码遇到的问题。用来产生条件位置编码的位置编码生成器（position encoding generator），接在每个 stage 的第一个编码模块后面。作者使用最简单的位置编码生成器（比如没有 batch norm 的二维通道卷积）。对于图片分类任务，在 CPVT 后，全都移除了类 token，并在 stage 的最后使用全局池化。

作者也尝试过在 Swin 中使用条件位置编码来取代相对位置编码，但实验结果并不理想。

Twins-SVT

视觉 Transformer 在密集预测任务中，一直受到计算量太大这个问题的严重影响，当输入为$H\times W$时，维度为$d$的自注意力的计算复杂度为$\hat{\mathcal{O}}\left(H^2{W}^{2}d\right)$。在此，我们提出一个空间可分离的自注意力，来减轻这个问题。SSSA 是由局部分组自注意力（LSA）和全局下采样注意力（GSA）组成。

局部分组自注意力（LSA）

受到可分离卷积的启发，我们首先将二维的特征图平均地分成几个窗口，让自注意力只发生在每个窗口内。具体点说，一个特征图被分成了$m\times n$ 个窗口。为了更好地范化，我们假设$H\%m=0\text{ and }W\%n=0$，每个组有 $\frac{HW}{mn}$ 个元素，计算量为 $\mathcal{O}\left(\frac{H^2{W}^2}{m^2{n}^2}d\right)$ ，总计算量则为：$\mathcal{O}\left(\frac{H^2{W}^2}{mn}d\right)$ 。若设$k_1,k_2$分别为 $k_1=\frac{H}{m},k_2=\frac{W}{n}$，则计算量可以表示为 $\mathcal{O}\left(k_1{k}_{2}HWd\right)$，可以看出当 $k_1\ll H\text{ and }{k}_2\ll W$ 时，模型会更有效率，且当 $k_1和k_2$ 为固定值时，与$HW$呈线性增长。

尽管上面的设计在计算时比较友好，但可以发现不同的 window 之间缺乏信息交流。在 Swin 中是提出了一个滑动窗口的设计，作者则提出下面 GSA 的设计。

局部下采样注意力（GSA）

作者没有采用简单粗暴地加上全局注意力层的方式，因为会带来计算量上的问题。

作者用一个代表值来总结每个 window 中的重要信息，并用这个代表值来与其它 window 进行信息交流（作用就像自注意力中的 key），将计算量极大地减少为 $\mathcal{O}(mnHWd)=\mathcal{O}\left(\frac{H^2{W}^{2}d}{k_1{k}_2}\right)$。这其实等价于在注意力操作中，将下采样得到的特征图当作 key。

对于下采样过程，我们尝试了很多种方法，最后的实验发现还是普通的步幅卷积效果最好。因此，这个可分享的 Transformer 模块可以写成如下形式：

z ^ i j l =  LSA  (  LayerNorm  ( z i j l − 1 ) ) + z i j l − 1 z i j l =  FFN  (  LayerNorm  ( z ^ i j l ) ) + z ^ i j l z ^ l + 1 =  GSA  (  LayerNorm  ( z l ) ) + z l z l + 1 =  FFN  (  LayerNorm  ( z ^ l + 1 ) ) + z ^ l + 1 , i ∈ { 1 , 2 , … . , m } , j ∈ { 1 , 2 , … , n } \begin{array}{l}\hat{\mathbf{z}}_{i j}^{l}=\text { LSA }\left(\text { LayerNorm }\left(\mathbf{z}_{i j}^{l-1}\right)\right)+\mathbf{z}_{i j}^{l-1} \\\mathbf{z}_{i j}^{l}=\text { FFN }\left(\text { LayerNorm }\left(\hat{\mathbf{z}}_{i j}^{l}\right)\right)+\hat{\mathbf{z}}_{i j}^{l} \\\hat{\mathbf{z}}^{l+1}=\text { GSA }\left(\text { LayerNorm }\left(\mathbf{z}^{l}\right)\right)+\mathbf{z}^{l} \\\mathbf{z}^{l+1}=\text { FFN }\left(\text { LayerNorm }\left(\hat{\mathbf{z}}^{l+1}\right)\right)+\hat{\mathbf{z}}^{l+1}, \\i \in\{1,2, \ldots ., m\}, j \in\{1,2, \ldots, n\}\end{array} z^ijl​= LSA ( LayerNorm (zijl−1​))+zijl−1​zijl​= FFN ( LayerNorm (z^ijl​))+z^ijl​z^l+1= GSA ( LayerNorm (zl))+zlzl+1= FFN ( LayerNorm (z^l+1))+z^l+1,i∈{1,2,….,m},j∈{1,2,…,n}​

此外，我们用 CPVT 中提出的位置编码生成器来编码位置信息，且在程序进行中处理输入不同长度的输入（而不是提前处理好）。这个结构被插在每个 stage 的第一个模块后面。

与 Swin 的比较

Swin 使用了滑动窗口来让不同的窗口之间进行信息交流，但是这种方法比较复杂，而且没有办法根据不同的设备来进行优化。我们的方法更简单而且更友好，因为它是建立在目前常用的操作上的。此外，我们的局部-全局设计能够更好的利用全局的信息，这也是在视觉任务中非常重要的。

最后可能有人注意到，我们的两个网络结构有着细微的区别，比如深度隐藏层维度、头的个数、扩张系数等等。这是因为我们希望与那两个目前比较有名的模型之间，有一个公平的比较。