[论文解读] Shuffle Transformer: Rethinking Spatial Shuffle for Vision Transformer

Zilong Huang, Youcheng Ben+4 more|arXiv (Cornell University)|Jun 7, 2021

Advanced Neural Network Applications参考文献 36被引用 123

一句话总结

提出 Shuffle Transformer，利用空间洗牌在基于窗口的视觉 Transformer 中实现跨窗口通信，且增加了邻窗口卷积；在分类、分割和检测任务上以线性令牌复杂度取得强结果。

ABSTRACT

Very recently, Window-based Transformers, which computed self-attention within non-overlapping local windows, demonstrated promising results on image classification, semantic segmentation, and object detection. However, less study has been devoted to the cross-window connection which is the key element to improve the representation ability. In this work, we revisit the spatial shuffle as an efficient way to build connections among windows. As a result, we propose a new vision transformer, named Shuffle Transformer, which is highly efficient and easy to implement by modifying two lines of code. Furthermore, the depth-wise convolution is introduced to complement the spatial shuffle for enhancing neighbor-window connections. The proposed architectures achieve excellent performance on a wide range of visual tasks including image-level classification, object detection, and semantic segmentation. Code will be released for reproduction.

研究动机与目标

• 推动在基于窗口的视觉 Transformer 中实现高效的跨窗口通信。
• 提出一种空间洗牌机制，以连接非重叠的窗口。
• 通过深度可分离卷积模块增强邻窗口连接。
• 创建 Shuffle Transformer 块，在层之间交替使用带空间洗牌与不带空间洗牌的结构，以提升信息流。
• 在图像分类、语义/实例分割和目标检测任务上展示强性能。

提出的方法

• 采用基于窗口的自注意力以实现输入规模的线性复杂度。
• 引入空间洗牌以跨越远处窗口融合信息，并通过反向空间对齐恢复内容对齐。
• 在 WMSA 与 MLP 之间加入带残差连接的深度卷积，以增强邻窗口交互。
• 定义 Shuffle Transformer Block，结合 Shuffle-MHSA、邻窗口连接和 MLP，并在层之间交替使用常规 WMSA 与 Shuffle-WMSA。
• 用 BatchNorm 替代 LayerNorm，并在某些层中使用 1x1 卷积以实现 2D 兼容性。
• 提供 Shuffle-T、Shuffle-S、Shuffle-B 三种变体，其 FLOPs 与基于 Swin 的骨干网络相近。

实验结果

研究问题

• RQ1 空间洗牌是否能够在非重叠窗口注意力中实现有效的跨窗口通信？
• RQ2 添加邻窗口卷积是否能改善局部跨窗口交互并缓解网格问题？
• RQ3 Shuffle Transformer 块在效率和准确性方面与现有基于窗口的变换器相比有何差异？
• RQ4 使用 Shuffle Transformer 骨干在 ImageNet、ADE20K 和 COCO 上的性能提升是多少？

主要发现

• Shuffle Transformer 在与 Swin 相近的复杂度下实现了与最先进方法相当甚至超越的结果的竞争力。
• 在 ImageNet-1K 上，Shuffle-T 达到 82.5% Top-1，Shuffle-S 达到 83.5% Top-1（与 Swin 变体的 GFLOPs 相近）。
• 在 ADE20K 上，Shuffle-B 实现 50.5% mIoU（多尺度），Shuffle-T 46.6%，Shuffle-S 48.4%（在相似 FLOPs 下高于 Swin 基线）。
• 在 COCO 与 Mask R-CNN 的评测中，Shuffle-T 在 AP 指标上优于 Swin-T（如 APb 46.8 vs 46.0； APm 42.3 vs 41.6）。
• 在 COCO 与 Cascade Mask R-CNN 中，Shuffle-T 与 Shuffle-S 在 APb、APb50、APm 等相关指标上表现出与 Swin-T、Swin-S 的竞争力或更优。
• 消融研究表明，长程空间洗牌再加上邻窗口连接相较于原生窗口自注意力带来持续的增益。