[Paper Review] CSWin Transformer: A General Vision Transformer Backbone with Cross-Shaped Windows

Xiaoyi Dong, Jianmin Bao+6 more|arXiv (Cornell University)|Jul 1, 2021

CCD and CMOS Imaging SensorsEngineering71 references101 citations

TL;DR

CSWin Transformer introduces cross-shaped window self-attention with parallel horizontal and vertical stripes and locally-enhanced positional encoding to create a scalable, general-purpose vision backbone with strong performance on classification, detection, and segmentation.

ABSTRACT

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.

Motivation & Objective

• Develop a general-purpose Vision Transformer backbone with improved efficiency and long-range modeling.
• Propose a cross-shaped window self-attention mechanism that computes attention in horizontal and vertical stripes in parallel.
• Introduce Locally-Enhanced Positional Encoding (LePE) for better local positional bias and input-resolution flexibility.
• Build a hierarchical CSWin Transformer architecture and validate on ImageNet-1K, COCO, and ADE20K.
• Demonstrate improvements over Swin Transformer under similar FLOPs and scale up with larger pretraining data.

Proposed method

• Define Cross-Shaped Window (CSWin) self-attention that splits multi-heads into two parallel groups for horizontal and vertical stripe attention.
• Compute attention within stripes of equal width sw, with stripe widths varying by depth to balance capacity and cost.
• Split heads into two groups so horizontal and vertical attentions run in parallel within a block, without extra computation.
• Introduce Locally-Enhanced Positional Encoding (LePE) added as a parallel module to self-attention that operates on projected values and supports arbitrary input resolutions.
• Construct a four-stage hierarchical CSWin Transformer with overlapping embedding, and four variants CSWin-T/S/B/L with specified channel counts and head configurations.
• Evaluate on ImageNet-1K classification, COCO object detection/instance segmentation, and ADE20K semantic segmentation, with and without ImageNet-21K pretraining.

Experimental results

Research questions

• RQ1 Can CSWin self-attention achieve larger effective receptive fields efficiently through cross-shaped windows?
• RQ2 Does parallel horizontal/vertical stripe attention improve modeling power without increasing computation significantly?
• RQ3 Does LePE provide robust positional encoding across varying input resolutions and downstream tasks?
• RQ4 How does CSWin Transformer perform as a general backbone across classification, detection, and segmentation compared to Swin and other vision Transformers?

Key findings

• CSWin-T achieves 82.7% Top-1 on ImageNet-1K with 4.3G FLOPs, outperforming Swin-T and DeiT-S under similar budgets.
• CSWin-S and CSWin-B surpass Swin-S and Swin-B in ImageNet-1K when using comparable FLOPs and model sizes.
• On COCO object detection, CSWin-B achieves 53.9 box AP and 46.4 mask AP, surpassing Swin-B by notable margins.
• On ADE20K segmentation, CSWin-B attains 53.9 mIoU (or 52.2 mIoU with certain setups) and CSWin-L reaches 55.7 mIoU with ImageNet-21K pretraining, outperforming previous SOTA backbones.
• With ImageNet-21K pretraining, CSWin-B/L reach 87.0/87.5 Top-1 accuracy on ImageNet-1K respectively, showing benefits from larger data.