[論文レビュー] Wave-ViT: Unifying Wavelet and Transformers for Visual Representation Learning

Ting Yao, Yingwei Pan+3 more|arXiv (Cornell University)|Jul 11, 2022

Advanced Neural Network Applications被引用数 8

ひとこと要約

Wave-ViTは、ウェーブレット変換とビジョントランスフォーマーを統合する新しいウェーブレットブロックを提案し、自己注意機構における逆引き可能で情報損失のないダウンサンプリングを可能にした。平均プーリングのような不可逆的ダウンサンプリングに代えて、離散ウェーブレット変換（DWT）を用いてキー／値をダウンサンプリングし、逆離散ウェーブレット変換（IDWT）で特徴マップを再構成・強化することで、SOTAモデルと同等のFLOPsで画像認識、オブジェクト検出、セグメンテーションの分野で優れた性能を達成した。

ABSTRACT

Multi-scale Vision Transformer (ViT) has emerged as a powerful backbone for computer vision tasks, while the self-attention computation in Transformer scales quadratically w.r.t. the input patch number. Thus, existing solutions commonly employ down-sampling operations (e.g., average pooling) over keys/values to dramatically reduce the computational cost. In this work, we argue that such over-aggressive down-sampling design is not invertible and inevitably causes information dropping especially for high-frequency components in objects (e.g., texture details). Motivated by the wavelet theory, we construct a new Wavelet Vision Transformer ( extbf{Wave-ViT}) that formulates the invertible down-sampling with wavelet transforms and self-attention learning in a unified way. This proposal enables self-attention learning with lossless down-sampling over keys/values, facilitating the pursuing of a better efficiency-vs-accuracy trade-off. Furthermore, inverse wavelet transforms are leveraged to strengthen self-attention outputs by aggregating local contexts with enlarged receptive field. We validate the superiority of Wave-ViT through extensive experiments over multiple vision tasks (e.g., image recognition, object detection and instance segmentation). Its performances surpass state-of-the-art ViT backbones with comparable FLOPs. Source code is available at \url{https://github.com/YehLi/ImageNetModel}.

研究の動機と目的

• マルチスケールビジョントランスフォーマーにおける、平均プーリングなどの不可逆的ダウンサンプリング操作によって引き起こされる情報損失を解消すること。
• プーリングの代わりに逆引き可能なウェーブレットベースのダウンサンプリングを採用することで、自己注意計算中に高周波成分（例：テクスチャ）を保持すること。
• 逆ウェーブレット変換によるダウンサンプリングされた特徴の再構成を通じて、自己注意出力における局所的文脈モデリングを強化すること。
• 最小限の計算オーバーヘッドで高い精度を維持するスケーラブルで効率的なトランスフォーマーブロックを設計すること。
• 提案されたウェーブレットブロックの有効性を、画像認識および密度予測タスクを含む複数のビジョンタスクにおいて検証すること。

提案手法

• 離散ウェーブレット変換（DWT）を用いてキーと値を4つのサブバンド（LL, LH, HL, HH）に分解し、逆引き可能なダウンサンプリングを実現するウェーブレットブロックを導入する。
• 4つのDWTサブバンドを1つのダウンサンプリング済み特徴マップにスタックし、周波数サブバンド間の局所的不変性を強制するために3×3畳み込みを適用する。
• ダウンサンプリング済みのキー／値と元のクエリを用いてマルチヘッド自己注意を実行し、計算効率を維持する。
• ダウンサンプリング済みのキー／値に逆離散ウェーブレット変換（IDWT）を適用し、拡大された有効受容 field を持つ高分解像特徴マップを再構成する。
• 注目された特徴マップとIDWTによる再構成マップを統合し、最終出力を得ることで局所的文脈集約を強化する。
• マルチステージ・マルチスケールトランスフォーマー構造にウェーブレットブロックを統合し、視覚的表現学習のための新規バックボーンであるWave-ViTを構築する。

実験結果

リサーチクエスチョン

• RQ1 逆引き可能なウェーブレット変換が、ビジョントランスフォーマーにおける不可逆的プーリング操作に効果的に置き換えられ、高周波成分の画像詳細を保持できるか？
• RQ2 逆ウェーブレット変換の統合により、顕著な計算コスト増加なしに自己注意出力における局所的文脈モデリングが向上するか？
• RQ3 標準的なプーリングベースの自己注意ブロックと比較して、提案されたウェーブレットブロックは、ビジョンタスク全体における精度と効率性において優れているか？
• RQ4 Wave-ViTは、既存モデルと同等のFLOPsで、画像認識、オブジェクト検出、インスタンスセグメンテーションにおいてSOTA性能を達成できるか？
• RQ5 ウェーブレットベースのダウンサンプリングが、サリエンシーマップ可視化によって示されるように、特徴表現のロバストネスをどの程度向上させるか？

主な発見

• Wave-ViT-SはADE20K検証セットで51.5%のmIoUを達成し、最高性能を示した競合モデル（Swin-S）を2.0ポイント上回った。
• Wave-ViT-BはImageNet-1Kで82.7%のTop-1精度を達成し、不可逆的プーリングを用いたベースライン（82.0%）を0.7ポイント上回った。
• アブレーションスタディの結果、平均プーリングをDWTに置き換えるだけでTop-1精度が82.0%から82.5%に向上し、逆引き可能なダウンサンプリングの利点を示した。
• DWTベースのブロックにIDWTを追加することで、計算量の増加は無視できるほど小さく、Top-1精度は82.7%まで上昇し、拡大された有効受容 field の有効性を裏付けた。
• スコアCAM可視化により、Wave-ViT-SはPVTv2-B2よりもよりロバストで意味的に焦点を合わせた表現を学習していることが確認された。
• Wave-ViTはタスク間で優れた一般化性能を示し、同等のFLOPsで画像認識、オブジェクト検出、インスタンスセグメンテーションの分野でSOTA性能を達成した。