VSA: Learning Varied-Size Window Attention in Vision Transformers

この記事は、論文の内容を5分くらいで読めるようにまとめた記事です。そのため、前提となる知識や関連研究に関する説明は大幅に省略しています。
基本的には筆者の備忘録ですが、面白そうと思ったら是非ご自身でも読んでみてください。

概要

TransformerのMultihead Self Attention（MSA）では、入力トークン長の2乗のオーダーの計算コストがかかる。Swin Transformerなどでは、WindowによるAttention範囲の限定を通じてこの問題を軽減しようとしている。しかし、既存のWindowによる手法は、Windowとして手作りの固定位置・固定スケールのものが使用されてきた。本研究では、Windowの位置とスケールを入力データに応じて柔軟に変化させることのできるVaried-Size Window Attention（VSA）を提案している。
これにより、Swin Transformerで提案されたShifted Windowを使わずとも、大域的な特徴量を集約することに成功している。

書誌情報

• Zhang, Qiming, et al. "VSA: Learning Varied-Size Window Attention in Vision Transformers." arXiv preprint arXiv:2204.08446 (2022).
• https://arxiv.org/abs/2204.08446
• 公式実装

ポイント

論文中では、Swin Transformerをベースに、VSAをどのように適用するのかが説明されている。

以下、上図についての簡単な説明であるが、(b)と(c)についてはもう少し細かい内容を後述する。

• (a): VSA Transformer Blockを用いたネットワークの全体像。基本的にはSwin Transformerを踏襲しており、Swin Transformer BlockがVSAのTransformer Blockに置き換わっている。
• (b): VSA Transformer Block内のコアとなるVaried-Size window Regression（VSR）モジュール。Queryについては、Swinのようにあらかじめ用意されたWindowをそのまま使用して計算するが、KeyとValueについては変形されたWindowを用いて作成する。Window内の特徴量をPoolingして得られる特徴量をもとに、Windowの変形に使用するパラメーター$S_w, O_w$を推定する。これらのパラメーターは、Multiheadのheadごとに推定される。変形されたWindowからグリッドサンプリングを行い、改めてトークンとし、そこからKeyとValueが算出される。
• (c): VSA Transformer Blockの詳細。図中のVWAと記載されている箇所が(b)のVSRに相当する。CPEはConditional Positional Encodingの略で、入力に応じて動的に変化するPositional Encodingである。

上図ではモジュールの表記（VWA? VSA?）が混乱しているようです。適宜読み替えてください。

Varied-Size window Regression

(b)について、もう少し詳細に確認してみる。$S_w, O_w$は下式のように算出される。$S_w$は縦横方向それぞれのスケールを表し、$O_w$はWindowの中心位置の並行移動であるオフセットを表す。それぞれ2次元なので、ヘッド数を$N$とすると、$S_w, O_w \in \mathbb{R}^{2\times N}$となる。また、下式のConv層では、1x1Convが使用されている。

$$S_w, O_w=\text { Conv } \circ \text { LeakyReLU } \circ \text { AveragePool }\left(X_w\right)$$

変形したWindowからのトークンのグリッドサンプリングは、ベースとなる座標を$S_w, O_w$によって操作した後、PyTorchのF.grid_sampleのような関数に渡すことで実現できる。詳細は公式実装を参照。
この操作をまとめてReshapeと記載すると、$K, V$は以下のような式で表される。

$$K, V=\text { Reshape } \circ \operatorname{Linear}(X)$$

VSA Transformer Block

次に、(c)についてもう少し詳しく確認する。重要なのは、Conditional Positional Encodingである。比較のために、Swinで用いられているRelative Positional Encodingについても簡単に説明する。

Relative Positional Encoding

Swinでは、Relative Positional Encoding（RPE）というテクニックが用いられており、Window内の位置に応じたバイアスがAttention時に適用される。
下図がその式にあたる。Windowの大きさを$M$とすると、$B \in \mathbf{R}^{M^2 \times M^2}$となり、訓練対象のパラメーター$B$はかなり大きな次元数になる。実際には内積計算を行う$Q, K$内の、トークン同士の相対的な位置関係だけを捉えたいので、訓練対象のパラメーターは$\mathbf{R}^{(2M+1) \times (2M+1)}$のサイズで済むような工夫が施されている。

$$\text { Attention }(Q, K, V)=\operatorname{SoftMax}\left(Q K^{T} / \sqrt{d}+B\right) V$$

Conditional Positional Encoding

VSAでは、RPEとは異なる方法として、Conditional Positional Encoding^[1]（CPE）を用いている。
CPEは単純なDepth-wiseの2次元Conv層によって実装される。公式実装を見てみると、カーネルサイズとしてWindowの大きさ（を奇数に補正したもの）を使い、biasをTrueに設定していることがわかる。そのため、biasによって無条件の位置情報を表現する一方で、Conv操作によって条件付きの位置情報を算出している、と考えることができる。これにより、入力された特徴マップをもとに、最適化されたPositional Embeddingが得られると期待される。
下図に示すように、VSAではVSA Transformer BlockごとにCPEが何度も適用される。

CPEとRPEは併用できることに注意が必要です。論文ではあまり強調されていない気がしましたが、VSAではWindowの変形パラメーターを計算する必要があるため、位置情報がどのようにVSRに与えられるかが非常に重要になりそうです。上図でわかる通り、Norm層を挟んでいるとはいえ、VSAのほぼ直前にCPEが存在し、位置情報がVSAに対して直接的に与えられているといえます。それに比べると、RPEの位置情報がVSAに入力されるまでにはいくつかのステップがあり、RPE単体では間接的な効果しかなさそうです。

実験

実験として、ImageNetの画像分類、COCOの物体検出、インスタンスセグメンテーション、CityScapeのセマンティックセグメンテーションを行っている。しかしここでは、提案手法の有効性を示す実験に絞って簡単に触れる。

アブレーションスタディ

VSR + CPEの組み合わせの有効性を確認するために、ImageNetによる画像分類のTop1精度で比較したのが、以下の表になる。

Shiftは、SwinのShifted Window Multihead Self Attention（SW-MSA）を表す。Shifted WindowはWindow位置が固定だと、そのWindow内でしか特徴量をやり取りできず、大域的な特徴を掴めない、という問題を解決するために導入されたモジュールである。VSRによって、Window位置が入力データによって動的に変化することで、Shifted Windowを使用する必要がなくなった、ということが表から読み取れる。
また、CPEがないとVSRの効果は限定的であるということもわかる。

画像解像度とWindowスケールの関係性

VSAの利点は、入力された特徴量に応じてWindowのスケールと位置を変化させることができる点にある。これは、入力画像の解像度が変化したときに、わかりやすく効果を発揮する。

下図は、入力解像度を変えたときに、3つのVSAレイヤーで算出されたスケールがどのような分布になっているのかを表したもので、右の方ほど深い位置のレイヤーである。上段が頻度を表すヒストグラムで、下段が各ビン内の各入力解像度の割合である。注目すべきは一番右の列で、入力解像度が大きいほど、算出されたスケールが大きい方に分布していることが分かる。解像度が大きいと当然そこに写っている物体も大きくなるために、ある程度の特徴が抽出されている段階の層のVSAでは、Windowのスケールを入力に応じて変化させることで物体の全体像を把握しようとしている、というふうに解釈できる。

Windowの可視化

最後に、VSRが出力するWindowがどのようなものなのかを確認しておく。下図の(a)はImageNetに関する、(b)はMS COCOでの例である。(a)(b)からわかる通り、Swinが固定Windowなのに対し、VSAではヘッドごとに異なるWindowを作成していることが見て取れる。
(c)は、ImageNetからサンプルした20のクラスの画像から得られる特徴量を、SwinとVSAそれぞれで抽出し、t-SNEにかけたものである。VSAの方がSwinよりもクラス内の分散が小さくまとまっているように見えることが確認できている。これは、モデルがVSAによって、物体の大きさによらない識別的な特徴量を得られることを示していると言える。

脚注

1. Chu, X., Tian, Z., Zhang, B., Wang, X., Wei, X., Xia, H., Shen, C.: Conditional positional encodings for vision transformers. arXiv preprint arXiv:2102.10882 (2021)