只需两行code即将置换思想与Transformer巧妙衔接

首发：AIWalker
作者：HappyAIWalker

标题&作者团队

本文是腾讯光影实验室俞刚团队在Transformer的的最新力作，他们创造性的将ShuffleNet的置换思想与Transformer进行结合。考虑到近期Swin为代表的非重叠窗口自注意力存在的问题，巧妙了借鉴了ShuffleNet的置换思想，提出了空域置换并于Transformer组合得到了本文的Shuffle Transformer。所提方案在ImagNet分类、ADE20K语义分割以及COCO实例分割等任务上均取得了SOTA性能。原文链接：https://arxiv.org/abs/2106.03...

Abstract

最近，基于窗口的Transformer(即在非重叠局部窗口内计算自注意力)在图像分类、语义分割以及目标检测等任务上证实了其令人满意的性能。然而，鲜少研究致力于跨窗口连接，而它会是提升表达能力的关键因素，

我们对空域置换(Spatial Shuffle)进行了回顾并将其作为一种有效方式构建窗口连接。更进一步，我们提出了一种称之为Shuffle Transformer的架构，它具有高度有效性，同时易于实现(仅需修改两行code！)。与此同时，我们引入深度卷积对空域置换进行补充以增强近邻窗口连接。所提方案在不同视觉任务(包含分类、检测以及语义分割)上均取得了卓越性能。

Method

接下来，我们从标准窗口多头自注意力出发。为构建长距离跨窗口关系，提出了空域置换，近邻窗口连接模块用于增强近邻窗口之间的连接性。然后，我们将空域置换与近邻窗口连接进行集成构建Shuffle Transformer Block以构建丰富的跨窗口连接。最后，我们提供了关于网络架构及其变种的信息。

Window-based Multi-head Self-Attention

为高效建模，有些研究对自注意力添加了局部约束以降低计算复杂度。基于窗口的多头自注意力(WMSA)在局部窗口内计算自注意力，这些窗口采用均匀无重叠方式进行划分。假设输入为，每个窗口包含个词，全局MSA与WMSA的计算量分别是。因此，WMSA更为高效，与HW成线性关系。

相比全局MSA，WMSA需要在计算自注意力之前进行窗口划分，在完成计算之后再进行窗口到图像的变换。然而，实际应用时额外的操作的计算复杂度是不能忽视的。除了窗口划分外，WMSA具有与全局MSA相同的架构。

Spatial Shuffle for Cross-Window Connections

尽管WMSA具有计算友好性，然而，图像被拆分成了无重叠子窗口。如果存在多个WMSA进行了堆叠，那么就会出现单边效应：感受野会受限于窗口而不利于分割类任务。下图a给出了堆叠两个WMSA的示意图，很明显：特定窗口的输出仅与对应窗口的输入相关。该属性会限制窗口间的信息流动，弱化表达能力。有没有感觉这个图跟ShuffleNet中的图非常相似呢？？？

为解决该问题，最直接的一种方案是：第二个WMSA能接收来自不同窗口的输入，见上图b。受启发于ShuffleNet，我们可以铜鼓空域置换操作对上述方案进行优雅而高效实现，见上图c。

不失一般性，假设输入为1D序列，WMSA的窗口尺寸为M，输入包含N个词。我们首先reshape输出空域维度为，然后transpose并平展后送入到下一层。这种类型的操作将来自不同窗口的输入打包，有助于构建长距离跨窗口连接。不同于通道置换，空域置换需要进行空域对齐操作调整空域词到原始的位置以确保特征与图像内容在空间上对齐。空域对齐操作首先将输入空域维度reshape为然后transpose并平展，它是空域置换的逆过程。

考虑到在计算自注意力之前与之后总是有窗口划分和窗口转图像的操作，我们可以将窗口划分与空域置换操作进行合并，窗口转图像与空域对齐进行合并。因此，空域置换与空域对齐不会带来额外的计算，且仅需修改两行code即可实现。更重要的是，空域置换是可微的，这意味着它空域嵌入到网络中进行端到端训练。

Neighbor-Window Connection Enhancement

将空域置换引入到WMSA中可以构建跨窗口链接，尤其是长距离跨窗口链接。然而，当处理高分辨率图像时存在一个潜在问题：当图像尺寸远大于窗口尺寸时会产生网络问题。

幸运的是，有这样几种方法通过增强近邻窗口链接缓解该问题。

提升窗口尺寸
配合移位窗口
引入卷积。

考虑到高效性，我们在WMSA与MLP之间插入了深度卷积。深度卷积的kernel尺寸与窗口尺寸相同，该操作可以提升近邻窗口的信息流动并缓解网络问题。

Shuffle Transformer Block

Shuffle Transformer Block(STB)包含置换MSA(Shuffle MHSA)、近邻窗口连接模块(NWC)以及MLP模块。为引入跨窗口连接，同时保持非重叠窗口的高效计算，我们提出了一种策略：在连续STB中对WMSA与Shuffle-WMSA进行交替执行，见下图。

从上图可以看到：第一个Transformer模块采用常规窗口划分策略，而第二个Transformer模块采用带空域置换的WMSA。从外，在每个模块内部添加NWC以增强近邻窗口的连接。因此，所提STB可以构建丰富的跨窗口连接并增强表达能力。最后，连续的STB计算流程如下：

注：为更好的处理2D输入，我们采用BatchNorm替换了STB中的LayerNorm；同时，线性性层替换为卷积。

Architecture and Variants

上图给出了所提Shuffle Transformer的架构示意图，注：此为tiny版本。Shuffle Transformer包含一个词嵌入层、多个STB以及词合并层。在我们的实现中，我们采用两个堆叠卷积作为词嵌入层。为生成分层表达，我们采用了stride=2，kernel=2的卷积作为词合并层以减少词数量。

为公平对比，我们参考了Swin的配置。我们所构建的基线模型Shuffle-B具有与Swin-B/ViTB/DeiT-B相当的模型大小以及计算复杂度。我们同时还引入了Shuffle-T与Shuffle-S，它们分别与Swin-T、Swin-S相当。窗口尺寸默认为7，每个头的query维度为32，每个MLP的扩展层为。本文所提模型的超参分别如下：

Shuffle-T：C=96，层数=；
Shuffle-S：C=96，层数=；
Shuffle-B：C=128，层数=。

Experiments

为展示所提方案的有效性，我们在ImageNet分类、ADE20K语义分割以及COCO实例分割等任务上进行了充分的实验对比。

Classification on ImageNet

上表对比了不同方案在ImageNet上的性能，从中可以看到：

相比其他ConvNet(比如RegNet)，所提方案取得了更佳的速度-精度均衡；
相比其他Transformer方案(比如Swin)，所提方案在相近计算复杂度下具有更高的精度。

Semantic Segmentation on ADE20K

上表对比了不同方案在ADE20K语义分割方面的性能，从中可以看到：

相比Swin，所提Shuffle具有更佳IoU指标，同时推理速度更快；
在多尺度测试下，相比Swin-T，Shuffle-T取得了1.4%mIoU指标提升；
相比Swin-B，Shuffle-B取得了SOTA指标50.5%mIoU，比Swin-B高0.8%mIoU；
相比Swin-S，Shuffle-T取得了0.7%mIoU指标提升。

Instance Segmentation on COCO

上表对比了不同方案在COCO实例分割方面的性能，从中可以看到：

在Mask R-CNN框架下，Shuffle-T与Shuffle-S大幅超过了其他ConvNet与Transformer方案(包含DeiT、PVT-Medium)；
相比Swin，所提Shuffle在所有指标上取得同等指标；
在Cascade Mask R-CNN框架下，所提方案具有更佳的与指标。
以上结果验证了所提方案在目标检测与实例分割方面的有效性。