CSWin Transformer: A General Vision Transformer Backbone with Cross-Shaped Windows

Xiaoyi Dong; Jianmin Bao; Dongdong Chen; Weiming Zhang; Nenghai Yu; Lu Yuan; Dong Chen; Baining Guo

Source-linked AI summary

CSWin Transformer: A General Vision Transformer Backbone with Cross-Shaped Windows

Xiaoyi Dong, Jianmin Bao, Dongdong Chen, Weiming Zhang, Nenghai Yu, Lu Yuan, Dong Chen, Baining Guo

arXiv:2107.00652v3Jul 1, 2021cs.CVcs.LG

TL;DR

Vision Transformers face a trade-off between expensive global self-attention and local attention’s limited interaction range. CSWin Transformer addresses this with parallel horizontal and vertical stripe attention, depth-varying stripe widths, and locally enhanced positional encoding. It achieves state-of-the-art performance across image classification, object detection, and semantic segmentation under constrained computation, with further gains after ImageNet-21K pretraining.

Problem
Global self-attention is expensive, whereas local or windowed attention limits each token’s interaction area and enlarges the receptive field slowly.
Method
CSWin Transformer uses parallel horizontal and vertical stripe self-attention, adjusts stripe width with network depth, and adds locally enhanced positional encoding in a hierarchical architecture.
Results
The architecture achieves state-of-the-art performance across image classification, object detection, and semantic segmentation under constrained computation complexity.
Takeaways & Limitations
Depth-dependent stripe widths enlarge attention areas efficiently, while CSWin’s combined design supports strong general-purpose vision performance under similar computation and model-size settings.
Takeaways & Limitations
The stripe-width computation analysis assumes a square input with H = W.

Abstract

from arXiv · show

We present CSWin Transformer, an efficient and effective Transformer-based backbone for general-purpose vision tasks. A challenging issue in Transformer design is that global self-attention is very expensive to compute whereas local self-attention often limits the field of interactions of each token. To address this issue, we develop the Cross-Shaped Window self-attention mechanism for computing self-attention in the horizontal and vertical stripes in parallel that form a cross-shaped window, with each stripe obtained by splitting the input feature into stripes of equal width. We provide a mathematical analysis of the effect of the stripe width and vary the stripe width for different layers of the Transformer network which achieves strong modeling capability while limiting the computation cost. We also introduce Locally-enhanced Positional Encoding (LePE), which handles the local positional information better than existing encoding schemes. LePE naturally supports arbitrary input resolutions, and is thus especially effective and friendly for downstream tasks. Incorporated with these designs and a hierarchical structure, CSWin Transformer demonstrates competitive performance on common vision tasks. Specifically, it achieves 85.4\% Top-1 accuracy on ImageNet-1K without any extra training data or label, 53.9 box AP and 46.4 mask AP on the COCO detection task, and 52.2 mIOU on the ADE20K semantic segmentation task, surpassing previous state-of-the-art Swin Transformer backbone by +1.2, +2.0, +1.4, and +2.0 respectively under the similar FLOPs setting. By further pretraining on the larger dataset ImageNet-21K, we achieve 87.5% Top-1 accuracy on ImageNet-1K and high segmentation performance on ADE20K with 55.7 mIoU. The code and models are available at https://github.com/microsoft/CSWin-Transformer.

1. Introduction

CSWin Transformer addresses the tension between expensive global attention and limited local attention by combining cross-shaped stripe attention, depth-dependent stripe widths, and locally enhanced positional encoding. Under similar FLOPs and model size, it reports strong results across classification, detection, and segmentation.

Motivation: Global self-attention models long-range dependencies but is computationally inefficient, while local attention limits receptive-field growth and requires many blocks to approach global interaction.Efficiently enlarging the receptive field is especially important for high-resolution detection and segmentation inputs.
Cross-shaped attention: CSWin self-attention splits multi-heads into parallel groups that attend to horizontal and vertical stripes, enlarging each token’s attention area without extra computation from the parallel strategy.Each stripe is formed by splitting the input feature into equal-width stripes.
Adaptive stripe widths: Depth-dependent stripe widths use smaller widths in shallow layers and larger widths in deep layers to balance modeling capability and computation cost.The paper provides mathematical analysis of how stripe width affects both capacity and computation.
Positional encoding: Locally-enhanced Positional Encoding adds positional information within each Transformer block and directly operates on attention results for downstream vision tasks.The architecture follows a hierarchical design for general-purpose vision tasks.
Results: 85.4% Top-1 accuracy, 53.9 box AP, 46.4 mask AP, and 51.7 mIOU are reported for CSWin-B across ImageNet-1K, COCO detection, and ADE20K segmentation, surpassing Swin counterparts by +1.2, +2.0, +1.4, and +2.0 respectively.With ImageNet-21K pretraining, the model reaches 87.5% Top-1 accuracy on ImageNet-1K and 55.7 mIOU on ADE20K.

2. Related Work

The related work situates CSWin among vision Transformers, efficient self-attention methods, and positional encoding schemes. It motivates a hierarchical backbone with cross-shaped window attention and locally enhanced positional encoding.

Vision Transformers: Vision Transformers demonstrated competitive results alongside CNNs and motivated architectures spanning multiple vision tasks.The cited progression begins with ViT and includes later Transformer-based designs.
CSWin Transformer: CSWin is introduced as a hierarchical vision Transformer backbone combining cross-shaped window self-attention with locally enhanced positional encoding.The architecture is presented as a general-purpose design for vision tasks.
Efficient Self-attentions: Efficient self-attention is important for high-resolution vision because full self-attention has computation complexity quadratic in image size.Existing approaches reduce cost through local attention windows, halo operations, or shifted windows.
Positional Encoding: Positional encoding restores spatial information omitted by permutation-invariant self-attention, with prior mechanisms including APE, RPE, and CPE.The paper connects its positional design to stronger local inductive bias.

3. Method

CSWin Transformer combines hierarchical processing with cross-shaped stripe attention and locally enhanced positional encoding to balance modeling capacity, computation, and positional information for vision tasks.

Hierarchical architecture: CSWin Transformer uses four hierarchical stages, with convolutional downsampling between adjacent stages to reduce tokens and double channel dimensions.Each stage contains sequential CSWin Transformer blocks, following a feature hierarchy similar to CNN backbones.
Cross-shaped window attention: CSWin attention splits feature maps into equal-width horizontal or vertical stripes, where stripe width balances learning capacity and computation complexity.Horizontal stripes contain sw × W tokens, while vertical stripes are derived analogously.
Cross-shaped window attention: Multi-head grouping assigns separate head groups to horizontal and vertical stripe attention, then concatenates their outputs to enlarge each token’s attention area within one block.The two attention directions are computed in parallel rather than applying the same operation across all heads.
Stripe-width scheduling: Stripe width is chosen smaller in early stages and larger in later stages because feature-map dimensions relative to channel dimensions change across the hierarchy.The default widths for the four stages are 1, 2, 7, and 7 for 224 × 224 inputs.
Locally-enhanced positional encoding: LePE incorporates local positional information within each Transformer block as a learnable per-element bias and supports varying input sizes through a distance threshold.The default Chebyshev-distance threshold is τ = 3, beyond which the positional contribution is set to 0.
Model variants: CSWin Transformer variants differ mainly in base channel dimension and block counts, with Tiny, Small, Base, and Large configurations evaluated using FLOPs calculated at 224 × 224 input size.The MLP expansion ratio is 4, and head counts vary by stage and model variant.

4. Experiments

Experiments evaluate CSWin Transformer across ImageNet classification, COCO detection and segmentation, ADE20K segmentation, inference speed, and ablations. Across these settings, CSWin generally improves performance over prior Transformer and CNN backbones while maintaining comparable computational or inference costs.

Experimental setup: Experiments cover ImageNet-1K classification, COCO detection and instance segmentation, and ADE20K semantic segmentation, with additional ablation and inference-speed studies.The evaluation uses standard vision benchmarks and compares CSWin with CNN and Transformer backbones.
ImageNet-1K classification: 82.7% Top-1 accuracy with 4.3G FLOPs lets CSWin-T surpass CvT-13, Swin-T, and DeiT-S by 1.1%, 1.4%, and 2.9%, respectively.A similar trend is observed after fine-tuning with 384 × 384 inputs.
COCO object detection: CSWin-T exceeds Swin-T by +4.5 box AP and +3.1 mask AP with the 1× schedule, and by +3.0 box AP and +2.0 mask AP with the 3× schedule.Cascade Mask R-CNN results also show CSWin variants surpassing counterparts under different model configurations.
ADE20K semantic segmentation: +6.7, +4.0, and +3.9 mIOU over Swin are achieved by CSWin-T, CSWin-S, and CSWin-B with Semantic FPN, while Upernet gains are +4.8, +2.8, and +3.0 mIOU.With ImageNet-21K pretraining, CSWin-L reaches 55.7 mIoU and surpasses the previous best model by +2.2 mIoU.
Inference speed: CSWin-S is +1.9% box AP and +1.7% mask AP higher than Swin-S at similar inference speed, while CSWin-T is faster and more accurate than Swin-B.In most cases, CSWin inference speed is less than 10% slower than Swin.

5. Conclusion

CSWin Transformer enlarges each token’s attention area efficiently through grouped horizontal and vertical stripe attention, while increasing stripe width with depth at small additional computation cost. The architecture also incorporates locally-enhanced positional encoding and achieves state-of-the-art performance across vision tasks under constrained computation complexity.

5. Conclusion: CSWin Self-Attention splits multi-heads into parallel groups that perform horizontal and vertical stripe attention, enlarging each token’s attention area within one Transformer block.The grouping design performs the two stripe operations in parallel.
5. Conclusion: Increasing stripe width along network depth further enlarges the attention area with subtle extra computation cost.The paper links this depth-dependent schedule to its mathematical analysis of stripe width.
5. Conclusion: CSWin Transformer incorporates locally-enhanced positional encoding for downstream tasks.
5. Conclusion: CSWin Transformer achieves state-of-the-art performance on various vision tasks under constrained computation complexity.

Experiment Details

The experiments use specified training, fine-tuning, augmentation, stochastic-depth, input-size, and testing settings for ImageNet classification and downstream detection, segmentation, and semantic-segmentation evaluations.

Experiment Details: ImageNet-1K models are trained for 300 epochs with 224 × 224 inputs using AdamW and augmentations including RandAugment, Mixup, CutMix, and Random Erasing.
Experiment Details: Stochastic-depth probabilities are 0.2, 0.4, and 0.5 for CSWin-T, CSWin-S, and CSWin-B during ImageNet training.
Experiment Details: Fine-tuning at 384 × 384 uses 30 epochs, weight decay 1e-8, learning rate 5e-6, and batch size 256.The stochastic-depth ratio remains unchanged from training because a large ratio benefits fine-tuning.
Experiment Details: COCO detection and instance-segmentation experiments use Mask R-CNN and Cascade Mask R-CNN implementations.
Experiment Details: Semantic-segmentation models use 512 × 512 inputs with specified augmentation, stochastic-depth, single-scale, and multi-scale testing settings.Multi-scale testing uses scales from 0.5× to 1.75× of the training scale.

More Experimetns

The paper supplements its limited in-paper comparisons with a broader ImageNet-1K comparison, grouping models by computation complexity. CSWin performs best among concurrent works in that comparison.

More Experimetns: The authors provide a comprehensive comparison with more current methods on ImageNet-1K because the paper compares only a few classical methods.
More Experimetns: CSWin performs best among concurrent works in the broader ImageNet-1K comparison.
More Experimetns: Table 11 compares ImageNet-1K classification models grouped by computation complexity.An asterisk marks EfficientNet models trained with other input sizes.