[Paper Review] Shuffle Transformer: Rethinking Spatial Shuffle for Vision Transformer

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

Advanced Neural Network ApplicationsComputer Science36 references123 citations

TL;DR

Introduces Shuffle Transformer, which uses spatial shuffle to enable cross-window communication in window-based Vision Transformers, with an added neighbor-window convolution; achieves strong results across classification, segmentation, and detection with linear token complexity.

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.

Motivation & Objective

• Motivate efficient cross-window communication in window-based vision transformers.
• Propose a spatial shuffle mechanism to connect non-overlapping windows.
• Enhance neighbor-window connections with a depth-wise convolution module.
• Create Shuffle Transformer blocks that alternate with and without spatial shuffle for better information flow.
• Demonstrate strong performance across image classification, semantic/instance segmentation, and object detection.

Proposed method

• Adopt window-based self-attention to achieve linear complexity in input size.
• Introduce spatial shuffle to merge information across distant windows and an inverse spatial alignment to restore content alignment.
• Incorporate a depth-wise convolution with a residual connection between WMSA and MLP to strengthen neighbor-window interactions.
• Define a Shuffle Transformer Block combining Shuffle-MHSA, Neighbor-Window Connection, and MLP, with alternation between regular WMSA and Shuffle-WMSA across layers.
• Replace LayerNorm with BatchNorm and use 1x1 convolutions in certain layers for 2D compatibility.
• Provide Shuffle-T, Shuffle-S, and Shuffle-B variants with comparable FLOPs to Swin-based backbones.

Experimental results

Research questions

• RQ1 Can spatial shuffle enable effective cross-window communication in non-overlapping window attention?
• RQ2 Does adding a neighbor-window convolution improve local cross-window interactions and alleviate grid issues?
• RQ3 How do Shuffle Transformer blocks compare to existing window-based transformers in efficiency and accuracy across tasks?
• RQ4 What are the performance gains on ImageNet, ADE20K, and COCO when using Shuffle Transformer backbones?

Key findings

• Shuffle Transformer achieves competitive to state-of-the-art results with similar complexity to Swin, often surpassing it in accuracy.
• On ImageNet-1K, Shuffle-T reaches 82.5% Top-1, and Shuffle-S reaches 83.5% Top-1 (with similar GFLOPs to Swin variants).
• On ADE20K, Shuffle-B achieves 50.5% mIoU (multi-scale), with Shuffle-T 46.6% and Shuffle-S 48.4% (higher than Swin baselines at similar FLOPs).
• On COCO with Mask R-CNN, Shuffle-T outperforms Swin-T in AP metrics (e.g., APb 46.8 vs 46.0; APm 42.3 vs 41.6).
• On COCO with Cascade Mask R-CNN, Shuffle-T and Shuffle-S show competitive or superior APs compared to Swin-T and Swin-S across APb, APb50, APm, and related metrics.
• Ablation studies show long-range spatial shuffle plus neighbor-window connection yields consistent gains over vanilla window-based self-attention.