腾讯提出Shuffle Transformer:重新思考视觉Transformer的空间Shuffle

点击下方卡片,关注“CVer”公众号

AI/CV重磅干货,第一时间送达

转载自:AIWalker

Shuffle Transformer: Rethinking Spatial Shuffle for
Vision Transformer
标题&作者团队

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

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框架下,所提方案具有更佳的 指标。

  • 以上结果验证了所提方案在目标检测与实例分割方面的有效性。

全文到此结束,更多消融实验与分析建议查看原文。

CVPR和Transformer资料下载

后台回复:CVPR2021,即可下载CVPR 2021论文和代码开源的论文合集
后台回复:Transformer综述,即可下载最新的两篇Transformer综述PDF
CVer-Transformer交流群成立
扫码添加CVer助手,可申请加入CVer-Transformer 微信交流群,方向已涵盖:目标检测、图像分割、目标跟踪、人脸检测&识别、OCR、姿态估计、超分辨率、SLAM、医疗影像、Re-ID、GAN、NAS、深度估计、自动驾驶、强化学习、车道线检测、模型剪枝&压缩、去噪、去雾、去雨、风格迁移、遥感图像、行为识别、视频理解、图像融合、图像检索、论文投稿&交流、PyTorch和TensorFlow等群。
一定要备注:研究方向+地点+学校/公司+昵称(如Transformer+上海+上交+卡卡),根据格式备注,可更快被通过且邀请进群
▲长按加小助手微信,进交流群▲点击上方卡片,关注CVer公众号
整理不易,请给CVer点赞和在看
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值