[論文レビュー] VSA: Learning Varied-Size Window Attention in Vision Transformers

Qiming Zhang, Yufei Xu+2 more|arXiv (Cornell University)|Apr 18, 2022

Advanced Neural Network Applications被引用数 4

ひとこと要約

本論文は、ビジョントランスフォーマーにおける固定サイズウィンドウを代替する学習可能でデータ駆動のウィンドウ構成に適した可変サイズウィンドウ注意（VSA）を提案する。各注意ヘッドごとに最適なウィンドウサイズと位置を予測するウィンドウ回帰モジュールを用いることで、長距離依存関係のモデリングを向上させ、多様なオブジェクトスケールにわたる特徴表現を改善し、計算コストをほとんど増加させずに性能を向上させる。Swin-TではImageNetで1.1%の精度向上を達成し、より大きな入力画像を用いた場合、最大で1.9%の向上を示す。

ABSTRACT

Attention within windows has been widely explored in vision transformers to balance the performance, computation complexity, and memory footprint. However, current models adopt a hand-crafted fixed-size window design, which restricts their capacity of modeling long-term dependencies and adapting to objects of different sizes. To address this drawback, we propose extbf{V}aried- extbf{S}ize Window extbf{A}ttention (VSA) to learn adaptive window configurations from data. Specifically, based on the tokens within each default window, VSA employs a window regression module to predict the size and location of the target window, i.e., the attention area where the key and value tokens are sampled. By adopting VSA independently for each attention head, it can model long-term dependencies, capture rich context from diverse windows, and promote information exchange among overlapped windows. VSA is an easy-to-implement module that can replace the window attention in state-of-the-art representative models with minor modifications and negligible extra computational cost while improving their performance by a large margin, e.g., 1.1\% for Swin-T on ImageNet classification. In addition, the performance gain increases when using larger images for training and test. Experimental results on more downstream tasks, including object detection, instance segmentation, and semantic segmentation, further demonstrate the superiority of VSA over the vanilla window attention in dealing with objects of different sizes. The code will be released https://github.com/ViTAE-Transformer/ViTAE-VSA.

研究の動機と目的

• 固定サイズウィンドウ注意の制限を解消し、長距離依存関係のモデリングとさまざまなサイズのオブジェクトへの適応性を向上させること。
• 入力データに応じて各注意ヘッドごとに最適なウィンドウサイズと位置を動的に学習可能にする。
• 計算コストやメモリ使用量を著しく増加させることなく、視覚タスクの性能を向上させること。
• シフトされたウィンドウのような複雑な機構に依存するのを減らし、重複する適応的ウィンドウを通じた情報交換を可能にすること。
• 分類、検出、セグメンテーションを含むさまざまなタスクに一般化できることを示すこと、特に高解像度入力において有効であることを確認すること。

提案手法

• VSAは、デフォルトの固定サイズウィンドウ内のトークンに基づいて、ターゲット注意ウィンドウのサイズと空間的位置を予測するウィンドウ回帰モジュールを導入する。
• その後、キーとバリューのトークンは予測されたターゲットウィンドウから抽出され、クエリはデフォルトのウィンドウから得られる。これにより、動的な受容野が実現される。
• 各注意ヘッドは独立してVSAを適用するため、ヘッド間で多様なウィンドウ構成が可能となり、マルチスケールのコンテキストを捉えることができる。
• モジュールはプラグアンドプレイであり、Swin-Tのようなモデルに標準的なウィンドウ注意を置き換える際、アーキテクチャの変更を最小限に抑え、FLOPsの増加もほとんどない。
• 重複するウィンドウを通じてウィンドウ間の特徴交換を促進することで、シフトされたウィンドウ機構の必要性が低下する。
• ウィンドウ回帰はモデル全体とエンドツーエンドで学習されるため、オブジェクトスケールや空間配置に応じたデータ駆動の適応が可能になる。

実験結果

リサーチクエスチョン

• RQ1 学習可能なウィンドウ機構は、固定サイズウィンドウの制限を超えてビジョントランスフォーマーにおける長距離モデリングを改善できるか？
• RQ2 適応的ウィンドウサイズと位置の予測は、さまざまなスケールのオブジェクトに対するより優れた特徴表現をもたらすか？
• RQ3 VSAは計算コストを著しく増加させることなく、画像分類および密度予測タスクの性能を向上させられるか？
• RQ4 固定ウィンドウが大規模なオブジェクトを捉えにくくなる高解像度入力において、VSAはどのように性能を発揮するか？
• RQ5 重複する適応的ウィンドウによる効果的なウィンドウ間情報交換を可能にすることで、シフトされたウィンドウ機構の必要性を低減または排除できるか？

主な発見

• Swin-TではImageNetで1.1%のトップ-1精度向上を達成し、より大きな画像（例：480×480）を用いた学習・推論では、精度向上が1.9%にまで増加する。
• オブジェクト検出、インスタンスセグメンテーション、セマンティックセグメンテーションといった下流タスクにおいて、VSAは標準的なウィンドウ注意を常に上回り、特に大規模なオブジェクトに対して顕著な性能向上を示す。
• t-SNE可視化では、VSAを搭載したモデルがより明確な特徴クラスタリングを示しており、より判別力があり意味的に意味のある特徴表現であることが確認された。
• VSAは重複するウィンドウを通じて効果的な情報交換を可能にし、Swinにおけるシフトウィンドウ機構をVSAバージョンでは不要にした。
• PyTorchの非最適なサンプリング操作による12–17%の速度低下があるものの、GPUメモリ使用量はわずか2%増加にとどまり、計算コストの増加は非常に小さい。
• 可視化結果から、VSAは固定サイズウィンドウよりも文脈に適応した多様なウィンドウを生成し、特に顕著な領域ではサイズと位置のばらつきが大きく、オブジェクトをよりよくカバーしていることが確認された。