【源头活水】Shuffle Transformer 高效快速的基础模型

“问渠那得清如许，为有源头活水来”，通过前沿领域知识的学习，从其他研究领域得到启发，对研究问题的本质有更清晰的认识和理解，是自我提高的不竭源泉。为此，我们特别精选论文阅读笔记，开辟“源头活水”专栏，帮助你广泛而深入的阅读科研文献，敬请关注。

来源：知乎—黄子龙

地址：https://zhuanlan.zhihu.com/p/384695881

第一次写知乎文章，主要介绍Tencent GY-Lab关于Transformer的一些探索，希望能对大家有所帮助和启发。之前只是在知乎上看文章，学到了很多，现在觉得自己可以有所反馈和贡献，期待和大家一起讨论~

文章链接: https://arxiv.org/abs/2106.03650

代码：https://github.com/mulinmeng/Shuffle-Transformer

Highlight：

既强又快的Transformer网络，轻松大幅提升分割和检测等任务的性能，CVPR2021人脸分割竞赛冠军方案。

01

关于Vision Transformer

最近Google将Transformer [1]成功运用到了视觉领域，其所提出的Vision Transformer [2]由于在视觉任务中具有建模长距离依赖关系的强大能力，并且引入了较少的归纳偏置，已经取得了与CNNs模型可比拟的性能，获得了越来越多的关注。

但是其使用的Global Self-Attention的具有O(N^2)计算复杂度，N为输入tokens数量。这使得很难将ViT应用于密集预测任务，比如语义分割和物体检测等需要输入高分辨率图像的任务。一种解决方法是使用Window-based Multi-head Self-Attention (WMSA)，将图像划分成固定大小的non-overlap的窗口，然后在每个窗口里（即固定tokens数量）计算self-attention，从而实现线性计算复杂度。但是这种方法也会导致有限的感受野，这对密集预测任务是不利的。为了解决不同窗口间的信息流通，当前主要有3种方案。

02

相关工作之跨窗口连接

1）Swin [3] 在两个Transformer encoder之间通过移动窗口的方式，使得信息可以在相邻窗口间流通。

2）HaloNet [4] 通过放大窗口，使得相邻窗口存在重叠（overlap），从而使得信息可以在附近窗口间流通。

3）加入卷积，卷积操作天然具备局部连接属性，可以在相邻窗口间交换信息。

尽管上述方案有效的建立了相邻窗口间的连接，但是其并不擅长建立长距离的窗口连接，从而更快的扩大感受野。因此我们提出将空间混编（Spatial Shuffle）技术与Window-based Multi-head Self-Attention进行结合，从而快速建立长距离的窗口连接。

在介绍Shuffle Transformer之前，想再多聊一下Swin。由于Swin在下游任务分割、检测等基础任务上的杰出表现，出现了很多基于Swin的改进工作和相关讨论，我们的工作也深受Swin的启发。为什么Swin表现得这么好？Swin采用了逐步下采样和局部连接（local self-attention），这些属性在CNNs中已经具备了。不同之处在于Swin具备 1. 更大的Window size （kernel size）; 2. 动态权重值；3. Transformer特殊的二元结构，空间信息的聚合(Self-Attention) + 信息重新提取(MLP)。现在有些实验证明1和2是有效的，对于3我们正在做一些探索。

最近MSRA王井东老师组最新的工作Demystifying Local Vision Transformer[5]做了一些探索，实验表明Sparse Connectivity, Weight Sharing, and Dynamic Weight是三个关键，大家可以关注一下。

03

Shuffle Transformer

基于空间混编操作的跨窗口连接

先前的一些工作为了高效建模将局部约束应用到自注意力上来降低计算成本。它们提出了窗口自注意力模块，其原理是只计算局部窗口内的自注意力，并且这些局部窗口是以不重叠的方式对图像进行均匀划分。假设每个窗口内包含 MxM个特征，对于一个包含 HxW个特征并且每个特征的维度为C的输入而言，全局MSA模块的计算复杂度和基于窗口的WMSA模块的计算复杂度分别为  和  。因此，当M<<W并且M<<H时，若M固定，窗口自注意力模块在计算复杂度上是随着图像的高 (H) 和宽 (W) 线性增长时，其计算效率相比于传统的全局自注意力模块明显更高。与全局自注意力模块相比，窗口自注意力模块在计算自注意力之前需要做一个窗口分区操作，并且在计算自注意力之后需要一个窗口到图像的转换操作。但是，这些额外操作的计算成本在实际实现中可以忽略不计。除了窗口划分操作外，窗口自注意力模块模块与全局自注意力模块模块共享相同的模型结构。

图 1 空间混编技术示意图

尽管窗口自注意力在计算上比较高效，但是在这种方式下图像被划分为多个不重叠的子窗口。如果将多个窗口自注意力模块叠加在一起，则会产生一个副作用：特征的感受野会受制于窗口范围内，这会对诸如图像分割等应用任务（特别是对高分辨率输入）产生不利影响。图1（a）展示了两个堆叠的基于窗口的WMSA模块情况。很明显可以看出，基于某个窗口的输出仅与该窗口内的输入相关。这一特性阻塞了窗口之间的信息交流并削弱了模型提取得到的特征表示。

为了解决这个问题，一个简单直接的解决方案是允许堆叠的第二个窗口自注意力模块从不同的窗口获取输入信息（如图1（b）所示，窗口自注意力2），在这种方式下不同窗口中的特征将是相关的。受ShuffleNet[6]和IGC[7]启发，这一操作可以通过空间混编的方式有效而优雅地实现（如图1（c）所示）。不失一般性，假设输入是一维序列，对于一个窗口大小为M的窗口自注意力模块，其输入有N个特征。我们首先将输出的空间维度变形为(M,N/M) ，然后转置并将其展平作为下一层的输入。这种操作可以将来自远处窗口的特征放到一起，从而有助于建立远程的跨窗口连接。与传统的通道混编操作不同，完成空间混编操作后需要在空间上（HxW, 高和宽两维度上）将不同特征调整到原始位置，以便在空间上对齐特征和图像内容。空间对齐操作首先将输出的空间维度变形成(N/M,M)，然后再转置并展平，这是前述空间混编操作的逆过程。

考虑到在自注意力计算之前和之后分别有窗口分区操作和窗口到图像的转换操作，我们可以将窗口分区操作与空间混编操作合并，以及将窗口到图像的转换操作与空间对齐操作合并。这样带来的好处是，对于空间混编操作和空间对齐操作并不会增加额外的计算量，并且很容易通过修改两行代码来实现。此外，空间混编操作也是可微的，这意味着它可以嵌入到现有网络结构中进行端到端的训练。

基于深度分离卷积的相邻窗口的连接增强

将空间混编操作引入到窗口自注意力模块可以建立跨窗口的连接，特别是增加了跨远程窗口的信息交流。但是这一方法在处理高分辨率图像输入时仍然存在潜在风险，当输入图像尺寸远大于窗口尺寸时，可能会出现“网格问题”。

幸运的是，我们可以通过几种增强邻接窗口连接的办法来解决这一问题：1）增大窗口大小；2）配合移位窗口机制；3) 引入卷积操作。考虑到效率问题，我们在窗口自注意力模块和前馈网络之间插入了一个带残差连接的深度分离卷积层，并且它的卷积核尺寸大小与窗口大小相同。这一算子可以增强相邻窗口（Neighbor-Window Connection，NWC）之间的信息流动，从而缓解“网格问题”。

基于空间混编的Transformer模块

如图2所示，两个基于空间混编的Transformer模块首尾相连，每一个模块依次包含了归一化层，（空间混编）窗口自注意力模块，深度分离卷积，归一化层和前馈网络。其与标准的Transformer模块有两点不同。1）使用了（空间混编）窗口自注意力。2)引入了深度可分离卷积。

图 2 两个连续的基于空间混编的Transformer模块

04

实验结果

我们在图像分类数据集ImageNet，语义分割数据集ADE20K和实例分割数据集COCO上验证了所提出的方法。在3个数据集上我们均取得了具备竞争力的结果，结果如下表所示，依次为ImageNet的分类TOP 1 accuracy，ADE20K分割的mIoU和COCO数据集上物体检测和实例分割mAP。可以看到在分类，分割任务中，Shuffle Transformer结果明显好于Swin，在检测任务中，结果与Swin相比拟。另外，值得注意的是Shuffle Transformer的推理速度要优于Swin。

05

初步应用之Short-video Face Parsing

我们将Shuffle Transformer应用到短视频人脸分割竞赛（The 3rd Person in Context (PIC) Workshop and Challenge at CVPR 2021 Short-video Face Parsing Track）中取得了第一名的成绩。

FAA-Shuffle Transformer架构

在该项竞赛中我们使用Shuffle Transformer作为encoder，配合改进的Alignseg[8][code]的decoder结构获得了非常不错的分割结果。同时我们对比了当前被广泛应用的HRNet[9]，实验结果表明我们所提出的网络结构取得了更高的结果，特别是在眼影、其他皮肤区域等难以清晰区域边界的标签上获得了更好的效果，具体示例如下所示。

关于人脸分割模型更加详细信息请前往 Arxiv(https://arxiv.org/abs/2106.08650) 。后续我们将Shuffle Transformer应用到更多的下游视觉任务中。

06

总结

我们提出了一种新的Transformer结构，将Spatial Shuffle和Window-based Self-Attention结合，有效的建立起了跨窗口连接，增强了模型的表达能力。实验表明本文所提出的方法在多种视觉任务中取得了很好的性能。

Vaswani A, Shazeer N, Parmar N, et al. Attention is all you need. NIPS, 2017.
Dosovitskiy A, Beyer L, Kolesnikov A, et al. An image is worth 16x16 words: Transformers for image recognition at scale. ICLR, 2021.
Liu Z, Lin Y, Cao Y, et al. Swin transformer: Hierarchical vision transformer using shifted windows[J]. arXiv preprint arXiv:2103.14030, 2021.
Vaswani A, Ramachandran P, Srinivas A, et al. Scaling local self-attention for parameter efficient visual backbones. CVPR 2021.
Han Q, Fan Z, Dai Q, et al. Demystifying Local Vision Transformer: Sparse Connectivity, Weight Sharing, and Dynamic Weight[J]. arXiv preprint arXiv:2106.04263, 2021.
Zhang X, Zhou X, Lin M, et al. Shufflenet: An extremely efficient convolutional neural network for mobile devices[C]. CVPR, 2018: 6848-6856.
Ting Z, Guo-Jun Q, Bin X, et al. Interleaved group convolutions for deep neural networks[C]. ICCV, 2017.
Huang Z, Wei Y, Wang X, et al. Alignseg: Feature-aligned segmentation networks[J]. T-PAMI, 2021.
Wang J, Sun K, Cheng T, et al. Deep high-resolution representation learning for visual recognition[J]. T-PAMI, 2020.
Huang Z, Ben Y, Luo G, et al. Shuffle Transformer: Rethinking Spatial Shuffle for Vision Transformer[J]. arXiv preprint arXiv:2106.03650, 2021.

本文目的在于学术交流，并不代表本公众号赞同其观点或对其内容真实性负责，版权归原作者所有，如有侵权请告知删除。

“源头活水”历史文章

更多源头活水专栏文章，

请点击文章底部“阅读原文”查看

分享、在看，给个三连击呗！