[論文レビュー] Shunted Self-Attention via Multi-Scale Token Aggregation

Sucheng Ren, Daquan Zhou+3 more|arXiv (Cornell University)|Nov 30, 2021

Advanced Neural Network Applications参考文献 34被引用数 10

ひとこと要約

本論文は、自己注意機構を用いて空間的に隣接するトークンを特徴量の類似性と空間的文脈に基づいて選択的に統合することで、1つの注意層内で細粒度および粗粒度の特徴を同時にモデル化できる、新しいマルチスケールトークン集約メカニズム「シャントド自己注意（SSA）」を提案する。SSAは計算コストを低減しながら、画像分類、物体検出、セマンティックセグメンテーションの性能を向上させ、SOTAのFocal Transformerと同等の効率制約下で、ImageNetではモデルサイズとFLOPsを半分に抑えながら84.0%のTop-1精度を達成し、COCOでは1.3 mAP、ADE20Kでは2.9 mIOUの上回り、性能を上回る。

ABSTRACT

Recent Vision Transformer~(ViT) models have demonstrated encouraging results across various computer vision tasks, thanks to their competence in modeling long-range dependencies of image patches or tokens via self-attention. These models, however, usually designate the similar receptive fields of each token feature within each layer. Such a constraint inevitably limits the ability of each self-attention layer in capturing multi-scale features, thereby leading to performance degradation in handling images with multiple objects of different scales. To address this issue, we propose a novel and generic strategy, termed shunted self-attention~(SSA), that allows ViTs to model the attentions at hybrid scales per attention layer. The key idea of SSA is to inject heterogeneous receptive field sizes into tokens: before computing the self-attention matrix, it selectively merges tokens to represent larger object features while keeping certain tokens to preserve fine-grained features. This novel merging scheme enables the self-attention to learn relationships between objects with different sizes and simultaneously reduces the token numbers and the computational cost. Extensive experiments across various tasks demonstrate the superiority of SSA. Specifically, the SSA-based transformer achieves 84.0\% Top-1 accuracy and outperforms the state-of-the-art Focal Transformer on ImageNet with only half of the model size and computation cost, and surpasses Focal Transformer by 1.3 mAP on COCO and 2.9 mIOU on ADE20K under similar parameter and computation cost. Code has been released at https://github.com/OliverRensu/Shunted-Transformer.

研究の動機と目的

• 各注意層におけるトークン間で一様で静的な受容 field を有するため、ビジョントランスフォーマーがマルチスケール特徴を捉える能力に制限があることに対処すること。
• 画像内に異なるスケールのオブジェクトが存在する場合でも、細粒度のディテールを保持しつつ、計算コストとメモリ消費量を低減すること。
• 1つの自己注意層が同時に大きなオブジェクトと小さなオブジェクトに注目できるようにし、下流のビジョンタスクにおける性能を向上させること。
• 既存のViTアーキテクチャを大規模に再設計することなく、汎用的で即座に統合可能なメカニズムを設計すること。

提案手法

• 注意計算を2つのブランチに分割するシャントド自己注意メカニズムを導入：1つは統合済みの粗粒度トークン用、もう1つは未統合の細粒度トークン用。
• 特徴量の類似性と空間的文脈に基づいて、空間的に隣接するトークンを選択的に統合する学習可能なトークン集約関数を適用し、大規模なオブジェクトに適したより大きな受容 field を形成する。
• 粗粒度特徴と細粒度特徴のそれぞれに別々の注意マップを維持することで、同じ層内で異なるスケール間の関係を学習可能にする。
• 隣接するトークンの文脈を統合することで局所的特徴表現を強化する、ディテールに特化した前向き伝搬ネットワークを統合し、特徴の精錬を向上させる。
• 空間解像度を保持し、初期特徴抽出段階での情報損失を低減するマルチスケールパッチ埋め込みヘッドを採用する。
• 統合済みトークンによりシーケンス長とFLOPsを削減する一方で、未統合トークンが高解像度のディテールを維持するハイブリッド計算戦略を適用する。

実験結果

リサーチクエスチョン

• RQ1 階層的特徴マップに依存せずに、1つの自己注意層が複数スケールの特徴を効果的にモデル化できるか？
• RQ2 選択的なトークン統合が、物体検出やセマンティックセグメンテーションなどの密度予測タスクにおいて、計算コストを低減しつつ性能を向上させるか？
• RQ3 分離されたブランチやダウンサンプリング戦略を用いるモデルと比較して、統一された注意メカニズムが優れた性能を発揮できるか？
• RQ4 標準的な畳み込み的または線形削減手法と比較して、提案されたトークン集約関数は、精度と効率の面で優れているか？
• RQ5 前向き伝搬層に局所的ディテール情報を組み込むことで、ViTベースのモデルにおける特徴表現がどの程度向上するか？

主な発見

• シャントドトランスフォーマーは、パラメータ数とFLOPsをSOTAのFocal Transformerと比較して50%削減した状態で、ImageNetで84.0%のTop-1精度を達成した。
• COCO物体検出タスクにおいて、同等のパラメータ数と計算量制約下で、Focal Transformerを1.3 mAP上回った。
• ADE20Kセマンティックセグメンテーションにおいて、同等の効率制約下でFocal Transformerを2.9ポイント上回るmIOUを達成した。
• SegFormerフレームワークを用いた場合、提案されたバックボーンはADE20Kで48.3%のmIoUを達成し、27.5MパラメータのMiT-B2を25.1Mパラメータで1.8ポイント上回った。
• 標準の前向き伝搬層と比較して、ディテールに特化した前向き伝搬層はViTで0.2%、シャントドトランスフォーマーで0.3%のTop-1精度向上を示し、局所的ディテールの保持効果を裏付けた。
• 非オーバーラップパッチ埋め込みと比較して、提案されたパッチ埋め込みヘッドはImageNet精度を1.4%向上させ、オーバーラップパッチ埋め込みと比較しても0.3%向上させたが、FLOPsの増加は最小限に抑えられた。