2021-Twins: Revisiting the Design of Spatial Attention

1. Title

• Twins: Revisiting the Design of Spatial Attention in Vision Transformers
• https://github.com/Meituan-AutoML/Twins

2. Summary

本文整体设计思路与之前的Vision Transformer Backbone一致，如何获取多尺度特征（PVT）以及如何降低Self-Attention的计算复杂度（SwinTransformer）。
第一点PVT的方法就是一种比较合适的方法，各个Stage将Embedding进行合并，从而减少Embedding的数目，以实现不同Stage具有不同分辨率的金字塔结构。
第二点SwinTransformer通过在Local Window内部计算Self-Attention的方式降低了计算复杂度，同时提出了一种比较复杂的Shift-Window的方式去捕获各个窗口之间的依赖关系。
作者认为SwinTransformer的Shift-Window方法较为复杂，且现代深度学习框架支持性不够好，因此，本文提出了一种更为简单的方法去实现。
整体思路可以认为是PVT+SwinTransformer的结合：在局部窗口内部计算Self-Attention（SwinTransformer），同时对每个窗口内部的特征进行压缩，然后再使用一个全局Attention机制去捕获各个窗口的关系（PVT）。
个人觉得Twins这种思路更加简单，而且也比较高效，SwinTransformer的实现包含了作者大量精巧的设计（阅读过SwinTransformer的代码应该都会觉得其实现是十分巧妙且美妙的），实际上也是比较复杂的。大道至简，还是Twins的设计更符合认知一点，不过其性能还有待验证。

3. Problem Statement

相较于CNN来说，Transformer由于其能高效地捕获远距离依赖的特性，近期在计算机视觉领域也引领了一波潮流。Transformer主要是依靠Self-Attention去捕获各个token之间的关系，但是这种Global Self-Attention的计算复杂度太高$O(N^2)$，不利于在token数目较多的密集检测任务（分割、检测）中使用。

基于以上考虑，目前主流有两种应对方法：

1. 一种是以SwinTransformer为代表的Locally-Grouped Self-Attention。其在不重叠的窗口内计算Self-Attention，当窗口大小固定时，整体的计算复杂度将下降至$O(N)$，然后再通过其他方法例如SwinTransformer中的Shift-Window的方法去实现窗口间的互动。但这种方法的缺点在于窗口的大小会不一致，不利于现代深度学习框架的优化和加速。
2. 一种是以PVT为代表的Sub-Sampled Version Self-Attention。其在计算Self-Attention前，会先对QKT Token进行下采样，虽然计算复杂度还是$O(N^2)$，但是在实际过程中也已经可以接受。

总结下来，目前的VisionTransformer的关键设计在于Spatial Attention的设计。
因此，本文将重点在Spatial Attention的设计上，期望提出一个高效同时简单的Spatial Attention方法。

4. Method(s)

作者有两个发现：

1. PVT中的Global Sub-Sample Attention是十分高效的，当配合上合适的Positional Encodings（Conditional Positional Encoding）时，其能取得媲美甚至超过目前SOTA的Transformer结构。
2. 更进一步，基于Separable Depthwise Convolution的思想，本文提出了一个Spatially Separable Self-Attention（SSSA）。该模块仅包含矩阵乘法，在现代深度学习框架下能够得到优化和加速。

4.1 Twins-PCPVT

PVT通过逐步融合各个Patch的方式，形成了一种多尺度的结构，使得其更适合用于密集预测任务例如目标检测或者是语义分割，其继承了ViT和DeiT的Learnable Positional Encoding的设计，所有的Layer均直接使用Global Attention机制，并通过Spatial Reduction的方式去降低计算复杂度。
作者通过实验发现，PVT与SwinTransformer的性能差异主要来自于PVT没有采用一个合适的Positional Encoding方式，通过采用Conditional Positional Encoding（CPE）去替换PVT中的PE，PVT即可获得与当前最好的SwinTransformer相近的性能。关于CPE的具体介绍，可以参见我的另一篇博客：Conditional Positional Encodings for Vision Transformers。

4.2 Twins-SVT

通过提出的Spatially Separable Self-Attention（SSSA）去缓解Self-Attention的计算复杂度过高的问题。SSSA由两个部分组成：Locally-Grouped Self-Attention（LSA）和Global Sub-Sampled Attention（GSA）。

4.2.1 Locally-Grouped Self-Attention（LSA）

首先将2D feature map划分为多个Sub-Windows，并仅在Window内部进行Self-Attention计算，计算量会大大减少，由$\left(H^2{W}^{2}d\right)$下降至$\mathcal{O}\left(k_1{k}_{2}HWd\right)$，其中$k_1=\frac{H}{m},k_2=\frac{W}{n}$，当$k_1,k_2$固定时，计算复杂度将仅与$HW$呈线性关系。

4.2.2 Global Sub-Sampled Attention（GSA）

LSA缺乏各个Window之间的信息交互，比较简单的一个方法是，在LSA后面再接一个Global Self-Attention Layer，这种方法在实验中被证明也是有效的，但是其计算复杂度会较高：$\mathcal{O}\left(H^2{W}^{2}d\right)$。

另一个思路是，将每个Window提取一个维度较低的特征作为各个window的表征，然后基于这个表征再去与各个window进行交互，相当于Self-Attention中的Key的作用，这样一来，计算复杂度会下降至：$\mathcal{O}(mnHWd)=\mathcal{O}\left(\frac{H^2{W}^{2}d}{k_1{k}_2}\right)$。
这种方法实际上相当于对feature map进行下采样，因此，被命名为Global Sub-Sampled Attention。
综合使用LSA和GSA，可以取得类似于Separable Convolution（Depth-wise+Point-wise）的效果，整体的计算复杂度为：$\mathcal{O}\left(\frac{H^2{W}^{2}d}{k_1{k}_2}+k_1{k}_{2}HWd\right)$。同时有：$\frac{H^2{W}^{2}d}{k_1{k}_2}+k_1{k}_{2}HWd\ge 2HWd\sqrt{HW}$，当且仅当$k_1\cdot k_2=\sqrt{HW}$。
考虑到分类任务中，$H=W=224$是比较常规的设置，同时，不是一般性使用方形框，则有$k_1=k_2$，第一个stage的feature map大小为56，可得$k_1=k_2=\sqrt{56}=7$。
当然可以针对各个Stage去设定其窗口大小，不过为了简单性，所有的$k$均设置为7。

整个Transformer Block可以被表示为：
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}​同时在每个Stage的第一个Block中会引入CPVT中的的PEG对位置信息进行编码。

4.3 Model Variants

5. Evaluation

5.1 Experiments

通过以上实验结果可以看出，Twins系列在各个任务上均取得了与SwinTransformer相当甚至是超过的水平，不过相比较而言，除了Small，Twins的模型参数比SwinTransformer系列都稍微大一点，而且运行速度似乎也没有明显优势。

5.2 Ablation Studies

6. Conclusion

本文提出了两种Vision Transformer Backbone，同时适用于图片级的分类任务或是其他密集预测任务，并且在分类、分割、检测等多个任务上，均取得了新的SOTA。