More than 1 year has passed since last update.

DW-ViT: 動的マルチスケールウィンドウを活用したビジョントランスフォーマー

Last updated at 2024-02-23Posted at 2024-02-22

論文とコード

Pengzhen Ren, Changlin Li, Guangrun Wang, Yun Xiao, Qing Du, Xiaodan Liang, Xiaojun Chang, Beyond Fixation: Dynamic Window Visual Transformer, arXiv:2203.12856, 2022
コード(https://github.com/pzhren/DW-ViT)

挿入している画像は、本論文からの引用となります。

概要

この論文は、視覚トランスフォーマーにおけるウィンドウサイズの重要性を指摘し、新しい手法であるDynamic Window Vision Transformer(DW-ViT)を提案しています。従来のウィンドウベースのViTは、固定されたウィンドウサイズで処理を行うため、性能向上に限界があります。DW-ViTでは、異なるサイズのウィンドウをマルチヘッドアテンションの各ヘッドに割り当てることで、マルチスケールの情報を効率的に取得することができます。さらに、各スケールの情報に動的に重み付けをして融合させることで、マルチスケールの情報を効果的に活用することができます。これらの工夫により、DW-ViTは、画像分類、インスタンス/セマンティックセグメンテーション、物体検出において、従来の手法を上回るパフォーマンスを達成しています。

動的なマルチスケールウィンドウによるViTの性能向上

Swin Transformerなどのウィンドベースのアテンションは優れた性能を発揮しますが、固定された単一スケールのウィンドウを使用することにより、いくつかの疑問が生じます。

ウィンドウサイズは最適なのか？
ウィンドウサイズが大きければ性能も向上するのか？
マルチスケールウィンドウは単一スケールウィンドウよりも有利か？
動的なマルチスケールウィンドウ¹はより良い結果をもたらすか？

これらの疑問を明らかにするために、論文ではウィンドウサイズがモデル性能に与える影響を評価しています。以下の図は、4つの単一スケールウィンドウ(win$\in$[7,14,21,23])を用いたSwin-T、マルチスケールウィンドウ(win=[7,14,21])を持つMSW-Swin、動的なマルチスケールウィンドウ(win=[7,14,21])によるDW-T(DW-ViTのTiny版)の性能を比較しています。

ウィンドウサイズが大きくなるにつれてモデルの性能は向上しますが、ある程度大きくなると、ほとんど向上しなかったり、逆に低下したりします。そのため、単にウィンドウを大きくすれば性能が向上するわけではありません。また、MSW-SwinはSwin-T(win=21)よりも精度が低いため、非最適なマルチウィンドウ設定は性能を低下させる可能性があります。一方で、DW-TはSwin-Tよりも高い性能を示すため、動的なマルチスケールウィンドウは良い結果をもたらします。
DW-ViTは、動的なマルチスケールウィンドウを用いて、マルチスケールウィンドウ設定を最大限に引き出すための新しい手法です。以下の図(a)は、DW-ViTを構成する2つの機構、Multi-Scale Window Multi-head Self-Attention(MSW-MSA)と、Dynamic Multi-Scale Window(DMSW)を簡潔に説明した図です。MSW-MSAでは、異なるウィンドウサイズを用いてそれぞれのアテンションを適用し、情報を出力します。DMSWでは、動的に生成される重み($\alpha_1$と$\alpha_2$)を出力された各情報に掛け合わせ、最後にこれらの情報を融合して新たな特徴マップを生成します。

DW-ViT

以下の図はDW-ViTの全体構成を示しています。

DW-ViTは4つのステージから構成されています。各ステージは、Patch Mergingと複数のDynamic Window Blockで構成されています。入力画像はPatch EmbeddingのConv(stride=kernel=4)によって、出力特徴次元が$\frac{H}{4}\times\frac{W}{4}\times C$となります。Dynamic Window Blockでは特徴次元を維持し、Patch Mergingでは隣接するパッチを連結し全結合を行うことで、解像度を半分に縮小し、チャンネル次元を2倍にします。

Dynamic Window Block

Dynamic Window Blockは、Transformer Blockの基本構成であるマルチヘッドセルフアテンション(MSA)をDynamic Window Module(DWM)に置き換えることで構築されています。DWMの内部では、Swin Transformerと同様なWindow-based MSA(W-MSA)を使用します。そのため、DWMはウィンドウ間での情報交換を行いません。ウィンドウ間の情報交換を実現するために、Swin Transformerと同様に循環シフトが組み込まれたW-MSAが採用されています。このシフトウィンドウ戦略を採用したDWMは、Dynamic Shifted Window(DSW)と呼ばれます。連続するDynamic Window Blockで、DWMとDSWを交互に呼び出すことで、ウィンドウ間での情報交換を効率的に行います。具体的には、2つの連続するDynamic Window Blockは、以下の順序で処理されます。

 Input --> LN --> DWM --> Add --> LN --> MLP --> Add --> Output
        |                  |   |                  |         |
        +------------------+   +------------------+         |
   +--------------------------------------------------------+
   |
 Input --> LN --> DSW --> Add --> LN --> MLP --> Add --> Output
        |                  |   |                  |
        +------------------+   +------------------+

Dynamic Window Module

DWMは、MSW-MSAとDMSWで構成されています。それぞれの詳細な構造は、以下の通りです。

図の左側はMSW-MSAの動作を示しています。$h$はヘッド数、$n_{win}$は異なるサイズのウィンドウ数を表します。図では$h=6$、$n_{win}=3$を例として挙げています。$h$個のヘッドを持つ特徴マップ($h\times H\times W\times \frac{C}{h}$)を、ヘッド方向に$n_{win}$個のグループに分割します。これにより、各グループの特徴マップは($\frac{h}{n_{win}}\times H\times W\times \frac{C}{h}$)となります。その後、各グループに対して異なるウィンドウサイズ$(win_i)$で特徴マップを分割し、各ウィンドウ内に相対位置エンコーディングを含むマルチヘッドアテンション(ヘッド数=$\frac{h}{n_{win}}$)を適用します。この手法は、Swin Transformerで採用されているW-MSAと同様のものです。W-MSAの後、異なるスケールの出力特徴マップが連結され、その結果がDMSWに渡されます。
DMSWの構造は、図の右側に示されています。このプロセスは、FuseとSelectという2つのステップで構成されます。Fuseは、softmax演算の前までの処理を行い、$n_{win}$個の異なるグループの情報を集めて処理するため全体情報を出力します。Selectは、得られた全体情報から各グループの重要度を決定し、その重要度に基づいて情報を合算します。この2つの部分の具体的な処理は以下の通りです。
Fuseの具体的な処理は、まず全結合層$F_{fc}$とGELU活性化関数$F_a$を適用します。その後、グローバル平均プーリング$F_{gp}$を行い、2回目の全結合層とGELUにより$1\times 1\times C'$の特徴マップが生成されます。なお、1回目の全結合層ではチャンネル次元を維持し、2回目の全結合層では$C'=\frac{C}{2 n_{win}}$に縮小されます。
Selectの具体的な処理は、まずチャンネルを$n_{win}$個のグループに分解し、各グループに全結合層とsoftmaxを適用します²。全結合層ではチャンネル次元を$\frac{C}{n_{win}}$に拡大し、softmaxでは$n_{win}$次元に沿って処理することで、各グループの重みを生成します。次に、生成された重みとW-MSAの出力を掛け合わせることで、各グループの情報が重要度に応じて強調されます。その後、各グループの情報を足し合わせて、情報を統合します。最後に、2つの全結合層を経由して³チャンネル次元を$C$に戻した後、残差結合が行われます。
Fuse-and-Selectにおける$F_{gp}$から重要度に基づく合算までの処理は、Squeeze-and-Excitation(SE)の構造と類似しています。SEでは、チャンネル次元の重要度を出力するために、$F_{gp}\rightarrow F_{fc}(縮小)\rightarrow F_{a}\rightarrow F_{fc}(拡大)\rightarrow \text{sigmoid}$の操作を行います。一方、Fuse-and-Selectでは、sigmoidの代わりにsoftmaxを用います。これは、チャンネルの重要度ではなく、$n_{win}$個のグループに対する重要度を計算するためです。

計算複雑度

特徴マップ$h\times w\times C$におけるDWMの計算複雑度は、以下のようになります。$win_i$と$n_{win}$は定数であるため、DWMの計算複雑度は空間サイズ$hw$に線形に依存します。

\displaylines{
  \Omega(\text{MSW-MSA}) = 4 h w C^2 + 2 h w \frac{C}{n_{win}} \sum_i^{n_{win}} win_i^2 \\
  \Omega(\text{DMSW}) = \Biggl(1 + h w \Bigl(1 + \frac{1}{n_{win}}\Bigr)\Biggr) \frac{C^2}{n_{win}} \\
  \begin{align}
    \Omega(\text{DWM}) &= \Omega(\text{MSW-MSA})+ \Omega(\text{DMSW}) \\
                       &= \Biggl(1 + 4 n_{win} + \frac{h w + n_{win}}{h w n_{win}}\Biggr) \frac{h w}{n_{win}} C^2 + 2 h w \frac{C}{n_{win}} \sum_i^{n_{win}} win_i^2
  \end{align}
}

実験

DW-ViTは、ImageNet-1Kの画像分類、ADE20Kのセマンティックセグメンテーション、COCO 2017の物体検出とインスタンスセグメンテーションにおいて、最先端の手法と性能を比較しています。また、アブレーションスタディも実施しています。
各種タスクにおいて、DW-ViTはSwin Transformerと比較して、類似したパラメータ数とFLOPsを持ちながら、高い精度を実現しています。これは、DW-ViTが採用する動的なウィンドウ効果によるものと考えられます。さらに、ImageNet-1Kでは0.7ポイント、ADE20Kでは1.2ポイント、Swin Transformerを上回る性能を達成しています。これは、複雑なタスクにおいて、動的なウィンドウがより効果的である可能性を示唆しています。以下は、画像分類における性能比較の結果を示しています。

アブレーションスタディでは、初めに提示した動的ウィンドウ効果に関する実験を拡張しています。以下の表は、Swin-TとDW-T(DW-ViTのTiny版)のDMSWを3つの状態('1'、'-'、'$\checkmark$')に設定したときの性能を比較しています。DMSW='1'の場合は、動的重み生成を削除し、すべてのグループに同じ重み(1/3)を割り当てます。DMSW='-'は、DMSWを削除することを意味します(MSW-Swin)。DMSW='$\checkmark$'は、通常のDW-Tです。DW-TとDMSW='1'の比較から、動的な重み設定が大きな改善(3.25%)を達成します。ただし、DMSW='1'はMSW-Swinよりも精度が低いため、適切な重みを割り当てないと、DMSWによって逆に精度が低下します。一方、MSW-SwinはSwin-T(win=21)よりも精度が低いため、マルチスケールウィンドウを組み込むだけでは性能が向上しない可能性があります。しかし、DW-Tは、DMSWのすべての状態とすべてのSwin-Tよりも高い精度を達成しています。これは、動的なウィンドウが重要な役割を果たしていることを示唆しています。

結論

この論文では、ウィンドウベースのViTにおけるウィンドウサイズがモデル性能に与える影響について考察し、最適なウィンドウ設定を実現する動的なマルチスケールウィンドウ機構を提案しています。この機構を導入したDW-ViTは、マルチスケール情報のモデリング能力を向上させ、シングルスケールウィンドウを使用するViTよりも優れた性能を発揮します。

ここでの動的なマルチスケールウィンドウとは、ウィンドウのサイズ、位置、形が動的に変化するというのではなく、異なるウィンドウサイズに対する重要度を状況に応じて変化させることを指します。例えば、特定のウィンドウサイズにおいて重要な情報が含まれていると判断された場合、そのウィンドウサイズに対する重みを高くすることを意味します。 ↩
論文では、処理手順として、分解を行った後に全結合層を適用すると述べられています。しかし、実際の実装では、無駄な処理をなくすために、先に全結合層を適用し、その後に分解を行います。 ↩
論文では、2つの全結合層が使用されていると述べられていますが、実際の実装では1つの全結合層しか使用されていません。 ↩

You get articles that match your needs
You can efficiently read back useful information
You can use dark theme

What you can do with signing up