[論文レビュー] 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を提案。空間シャッフルを用いてウィンドウベースのVision 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.

研究の動機と目的

• ウィンドウベースのビジョン・トランスフォーマーにおける効率的なクロスウィンドウ通信を動機づける。
• 非重複ウィンドウを接続する空間シャッフル機構を提案する。
• 隣接ウィンドウの相互作用を深層畳み込みモジュールで強化する。
• 情報の流れを改善するため、空間シャッフルの有無で交互に配置されたShuffle Transformerブロックを作成する。
• 画像分類、セマンティック/インスタンス分割、物体検出のすべてで強力な性能を示す。

提案手法

• 入力サイズの線形複雑度を達成するためにウィンドウベースの自己注意を採用する。
• 遠方ウィンドウ間の情報を結合する空間シャッフルを導入し、内容配置を元に復元する逆空間整列を適用する。
• WMSAとMLPの間に残差付きの深度-wise畳み込みを組み込み、隣接ウィンドウの相互作用を強化する。
• Shuffle-MHSA、Neighbor-Window Connection、MLPを組み合わせたShuffle Transformer Blockを定義し、レイヤー間で通常のWMSAとShuffle-WMSAの交互配置をとる。
• LayerNormをBatchNormに置換し、2D互換性のある一部レイヤーで1x1畳み込みを使用する。
• Shuffle-T、Shuffle-S、Shuffle-Bの派生モデルを提供し、Swinベースのバックボーンと同等のFLOPsを実現する。

実験結果

リサーチクエスチョン

• RQ1 空間シャッフルは非重複ウィンドウの自己注意において効果的なクロスウィンドウ通信を可能にするか。
• RQ2 隣接ウィンドウ畳み込みを追加することで局所的なクロスウィンドウ相互作用を改善し、グリッド問題を緩和できるか。
• RQ3 Shuffle Transformerブロックは効率と精度の点で既存のウィンドウベースのトランスフォーマーと比べてどうか。
• RQ4 Shuffle Transformerバックボーンを使用した場合、ImageNet、ADE20K、COCOでのパフォーマンス向上はどの程度か。

主な発見

• Shuffle TransformerはSwinと同等の複雑度で最先端に近い結果を達成でき、しばしば精度で上回る。
• ImageNet-1Kで Shuffle-TはTop-1 82.5%、Shuffle-SはTop-1 83.5%（Swin派生と同等のGFLOPs）に達する。
• ADE20Kで Shuffle-Bは50.5% mIoU（マルチスケール）、Shuffle-Tは46.6%、Shuffle-Sは48.4%（Swinベースラインと同等のFLOPsで上回る）。
• COCOのMask R-CNNで Shuffle-TはSwin-Tを上回るAP指標を示す（例：APb 46.8 vs 46.0； APm 42.3 vs 41.6）。
• Cascade Mask R-CNNで Shuffle-T と Shuffle-S は、Swin-T および Swin-S と比較してAPb、APb50、APm などの指標で競争力があるか、あるいは上回る。
• アブレーション研究は、長距離空間シャッフルと隣接ウィンドウ接続の組み合わせが、バニラのウィンドウベース自己注意に対して一貫したゲインをもたらすことを示している。