Source-linked AI summary
Shuffle Transformer: Rethinking Spatial Shuffle for Vision Transformer
Zilong Huang, Youcheng Ben, Guozhong Luo, Pei Cheng, Gang Yu, Bin Fu
TL;DR
Efficient window-based transformers reduce computation but leave limited cross-window communication, weakening receptive fields for visual representation. Shuffle Transformer adds spatial shuffle and alignment, complemented by depth-wise convolution for neighboring windows, and performs favorably across classification, detection, and segmentation tasks. Its main scope boundary is that spatial shuffle can exhibit a grid issue when image size greatly exceeds window size, motivating the convolutional enhancement.
Problem
Window-based self-attention is computationally efficient but lacks rich cross-window connections, limiting receptive fields and information flow across windows.
Method
Shuffle Transformer integrates spatial shuffle and spatial alignment into window-based self-attention, adding a depth-wise convolutional residual connection to strengthen neighbor-window communication.
Results
Across image classification, object detection, and semantic or instance segmentation, both proposed architectures perform favorably against vision transformers with similar computational complexity.
Takeaways & Limitations
The architecture maintains linear complexity in the number of input tokens while enabling information flow across windows through successive Shuffle Transformer Blocks.
Takeaways & Limitations
Spatial shuffle may produce a grid issue when image size is far greater than the window size, and the reported base-model comparison matches Swin-B, ViT-B, and DeiT-B in model size and computation complexity.
Abstract
from arXiv · showhide
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.
1 Introduction
Window-based self-attention improves efficiency by restricting computation to local windows, but limits receptive fields and cross-window communication. Shuffle Transformer addresses this with spatial shuffle, alignment, and depth-wise convolution while retaining efficient computation.
- 1 Introduction: Global self-attention has quadratic complexity, whereas local windows reduce computation but restrict the receptive field and cross-window communication.These limitations are especially relevant to high-resolution dense prediction tasks.
- 1 Introduction: Spatial shuffle rearranges tokens between windows to provide long-range cross-window connections in window-based self-attention.A spatial alignment operator restores feature alignment with image content after shuffling.
- 1 Introduction: The Shuffle Transformer Block combines spatial shuffle, spatial alignment, and a depth-wise convolutional residual connection.The convolution complements shuffle by enhancing neighbor-window connections and addressing the grid issue.
- 1 Introduction: Successive Shuffle Transformer Blocks allow information to flow across all windows while self-attention remains confined to non-overlapping local windows.This preserves linear computational complexity in the number of input tokens.
- 1 Introduction: Across image classification, object detection, and semantic or instance segmentation, both proposed architectures perform favorably against vision transformers with similar computational complexity.The reported evaluation spans image-level and dense visual tasks.
2 Related Work
Related work develops efficient vision transformers through local attention, hierarchical representations, convolutional integration, and shuffle-based information exchange. These approaches motivate spatial shuffle as a way to restore cross-window communication in efficient window-based models.
- Vision Transformers: Vision Transformers model image tokens with transformer layers, while later variants address training data, token length, and hierarchical representation requirements.PVT produces feature pyramids for dense prediction tasks such as detection and segmentation.
- Integrating Convolution and Vision Transformer: Convolutions are also integrated into vision transformers to enhance local context and combine transformer and convolutional benefits.CPVT uses conditional position encodings, while CvT and CCT introduce convolutions into transformer architectures.
- Window-based Self-Attention: Global self-attention has quadratic complexity in the number of tokens, motivating local or window-based attention for high-resolution vision tasks.Window-based methods include sliding-window and non-overlapping-window categories.
- Window-based Self-Attention: Non-overlapping windows reduce computation but lack cross-window connections, while shifted windows alleviate this limitation only partially for long-range communication.Figure 1 contrasts isolated windows with cross-window interaction after rearranging tokens.
- Channel Shuffle and Spatial Shuffle: ShuffleNet uses channel shuffle to exchange information across convolution groups, inspiring spatial shuffle for exchanging information across attention windows.Spatially shuffled convolution and ISSA are related approaches, but ISSA retains quadratic complexity with input size.
3 Shuffle Transformer
Shuffle Transformer builds efficient cross-window communication by alternating regular and spatially shuffled window attention, then strengthens local cross-window flow with depth-wise convolution. Its hierarchical architecture retains linear token complexity and is evaluated across vision tasks.
- Window-based Multi-head Self-Attention: Window-based self-attention reduces global attention cost to linear growth in image tokens when window size is fixed.For M × M windows, its stated complexity is O(M^2HWC), versus O(H^2W^2C) for global attention.
- Window-based Multi-head Self-Attention: Non-overlapping windows limit receptive fields and weaken cross-window connections, particularly for high-resolution segmentation.Stacked window attention remains confined within each window unless information is exchanged across partitions.
- Spatial Shuffle for Cross-Window Connections: Spatial shuffle rearranges tokens so distant windows are grouped together, creating long-range cross-window connections before spatial alignment restores image-content correspondence.Spatial alignment is the inverse operation that returns shuffled spatial tokens to their original positions.
- Spatial Shuffle for Cross-Window Connections: Merging shuffle with window partitioning and alignment with windows-to-image conversion makes the operations computation-free and implementable by modifying two lines of code.The spatial shuffle is also differentiable for end-to-end network training.
- Neighbor-Window Connection Enhancement: A depth-wise convolution with a residual connection enhances neighbor-window communication and alleviates the grid issue when images are much larger than windows.The convolution is inserted between WMSA and MLP, using a kernel size equal to the window size.
- Shuffle Transformer Block: Shuffle Transformer Blocks alternate WMSA and Shuffle-WMSA, combining spatial shuffle, neighbor-window connection, and MLP modules into the basic block.The architecture uses token embedding, several blocks, and stride-2 token merging layers to produce hierarchical representations.
- Architecture and Variants: The proposed Shuffle Transformer maintains linear token complexity while building rich cross-window connections and reports favorable performance across image classification, segmentation, and detection tasks.Variants include Shuffle-T, Shuffle-S, and Shuffle-B, with default window size M = 7.
4 Experiments
Experiments evaluate Shuffle Transformer on image classification, semantic segmentation, object detection, instance segmentation, and module ablations. Across tasks, the models generally compare favorably with existing backbones, while ablations assess spatial shuffle, neighbor-window connections, shuffle variants, and insertion position.
- Experimental settings: Experiments cover ImageNet-1K classification, ADE20K semantic segmentation, and COCO object detection and instance segmentation.The study uses ImageNet-1K, ADE20K, and COCO benchmarks with task-specific training and evaluation settings.
- ImageNet-1K classification: Shuffle Transformer achieves a better speed-accuracy trade-off than RegNet and stably outperforms similarly complex Transformer counterparts such as Swin on ImageNet classification.All compared models are trained and evaluated at 224 × 224 resolution.
- ADE20K semantic segmentation: Under multi-scale testing on ADE20K, Shuffle-T exceeds Swin-T by 1.4% mIoU, while Shuffle-B reaches 50.5% mIoU and exceeds Swin-B by 0.8% mIoU.Shuffle-S achieves performance comparable to Swin-S, and the variants have comparable inference speed.
- COCO detection and instance segmentation: On COCO, Shuffle-T and Shuffle-S surpass several ConvNet and Transformer baselines for Mask R-CNN, while Shuffle Transformer stably outperforms other networks in AP_b and AP_m with Cascade Mask R-CNN.The reported comparisons use comparable parameter sizes and GFLOPs where stated.
- Ablation studies: Adding spatial shuffle improves ImageNet top-1 accuracy by 1.3% and ADE20K mIoU by 2.1, while adding neighbor-window connection contributes a further 1.0% and 2.8 mIoU.The ablations compare Shuffle-T with regular window-based multi-head self-attention and then add the neighbor-window connection.
- Ablation studies: Long-range spatial shuffle performs best among the tested shuffle designs on both classification and segmentation, while random spatial shuffle achieves comparable performance.The study considers long-range, short-range, and random spatial shuffle variants.
- Ablation studies: Inserting the neighbor-window connection between shuffle WMSA and the MLP block achieves the best performance among the three tested positions.The neighbor-window connection is implemented as a depth-wise convolution with a residual connection.
5 Conclusion
Shuffle Transformer uses efficient local self-attention together with spatial shuffle and depth-wise convolution to connect information across windows. Experiments across diverse visual tasks show favorable performance, supporting its use as a strong vision backbone.
- Window-based multi-head self-attention computes attention within non-overlapping windows for efficient modeling.
- Spatial shuffle introduces cross-window connections, while depth-wise convolution with a residual connection strengthens neighbor-window connections.
- Successive Shuffle Transformer Blocks enable information flow across all windows.
- Extensive experiments show both proposed architectures perform favorably against state-of-the-art vision transformers with similar computational complexity.
- The architectures may serve as stronger backbones for image classification, semantic and instance segmentation, and object detection.