論文とコード

Sitong Wu, Tianyi Wu, Haoru Tan, Guodong Guo, Pale Transformer: A General Vision Transformer Backbone with Pale-Shaped Attention, arXiv:2112.14000, 2021
非公式コード(https://github.com/qwopqwop200/pale-transformer-pytorch)

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

概要

トランスフォーマーは画像認識や物体検出などの視覚タスクにおいて、高い性能を発揮します。しかし、グローバルセルフアテンションは、計算量が膨大になるという課題があります。この課題を解決するために、アテンションの範囲をローカル領域に限定する手法が提案されています。これらの手法は、計算量を削減する一方で、アテンションの範囲が狭くなるため、画像全体の情報を十分に捉えることができず、性能が低下するという問題があります。この問題に対処するため、論文では、pale形状¹のセルフアテンション(PS-Attention)を提案しています。PS-Attentionは、以前のローカルセルフアテンションと比較して、計算複雑度が似ているにもかかわらず、より広範囲の情報を考慮して処理することができます。このPS-Attentionと階層型アーキテクチャを採用したPale Transformerは、画像分類、物体検出、セマンティック/インスタンスセグメンテーションにおいて、従来の最先端のCSWin Transformerよりも精度を大幅に向上しています。このように、Pale Transformerは、視覚タスクにおいて高い性能を発揮し、汎用的な視覚トランスフォーマーのバックボーンとして有用であることが示唆されています。

既存のSelf-Attention

提案されているローカルセルフアテンションは、さまざまです。

図(b)のウィンドウベースのセルフアテンションは、特徴マップを小さなウィドウに分割し、各ウィンドウ内でセルフアテンションを行います。また、受容野を広げるために、ウィンドウ同士を接続しています。接続には、さまざまな戦略が導入されています。Swin TransformerとShuffle Transformerは、それぞれウィンドウの位置をずらして分割する方法と、各ウィンドウ内の要素を組み合わせて分割する方法を提案しています。これらの方法は、連続するTransformer Block内で、通常のウィンドウと提案されたウィンドウの2つの異なるウィンドウ分割を交互に使用することで、ウィンドウ間の関連を構築しています。MSG Transformerでは、各ウィンドウに配置された、ウィンドウ内の情報を持つメッセンジャートークンを、異なるウィンドウのメッセンジャートークンと混ぜ合わせることで、ウィンドウ間で情報を交換します。
アテンションの範囲をストライプ状にしたセルフアテンションも提案されています(図(c)-(e))。Axial-Attentionは単一行または単一列を扱うのに対し、CSWin-Attentionはそれを拡張し、複数行または複数列を扱うことができます。また、Axial self-attentionは、レイヤーごとに異なるストライプ形状のセルフアテンションを交互に適用するのに対し、CSWin-Attentionでは、レイヤー内で異なるストライプ形状のセルフアテンションを同時に適用します。これにより、CSWin-Attentionは、単一のレイヤーで、縦横のストライプが交差する十字路形状のアテンション範囲を実現しています。

Pale-Shaped Self-Attention

Pale-Shaped Self-Attention(PS-Attention)は、特徴マップを複数のpale形状に分割し、各pale形状にセルフアテンションを適用します。pale形状とは、等間隔に配置された、幅が1のストライプで構成されるグリッドのことです。縦の間隔と横の間隔は必ずしも同一であるとは限りません。1つのpale形状には、$s_r$個の行と$s_c$個の列が含まれます。アテンション範囲をpale形状に設定して効率よく計算するには、pale形状を行方向と列方向に分解し、各領域のセルフアテンションを同じレイヤーで実行することです。これにより、セルフアテンションの計算量が、pale形状の範囲ではなく、$s_r$または$s_c$個のストライプ形状の範囲に削減されます。これはCSWin-Attentionとよく似ています。ただし、CSWin-Attentionではマルチヘッド方向を2つのグループに分け、それぞれのグループに異なるストライプを割り当てていますが、PS-Attentionではチャンネル方向を半分に分けてグループ化します(下図(c))。

これにより、PS-Attentionは、行方向と列方向を組み合わせたpale形状にアテンション範囲を拡大することができます。セルフアテンションの処理が終了すると、チャンネル方向の結合によって、元の形状に戻ります。PS-Attentionは、CSWin-Attentionよりも特徴マップ全体にアテンションの範囲を広く取っているため、より長距離の依存関係を学習することができます。
特徴マップサイズが$h\times w\times c$で、paleサイズが$(s_r, s_c)$の場合、PS-Attentionの計算複雑度は、次の式で表されます。

\displaylines{
\Omega(\mathrm{Pale}) = hwc * (4c + s_ch + s_rw + 27)
}

CSWin-Attentionの計算複雑度は、ストライプ幅を$s_w$で表すと、$hwc * (4c + s_wh + s_ww)$となります。$s_r = s_c = s_w$の場合、PS-Attentionでは計算複雑度が$27hwc$分だけ増加します。これは、PS-Attentionでは、特徴マップをクエリ、キー、バリューに変換するために、separable convolution層を用いるためです。CSWin-Attentionでは、全結合層を用いていますが、PS-Attentionではこれをseparable convolution層に変えたため、計算複雑度が増加しています。

全体的なアーキテクチャ

Pale Transformerは、4つの階層ステージから構成され、多段階の特徴抽出を行っています(上図(a))。各ステージは、Patch Mergingと複数のPale Transformer Blockで構成されています。Patch Mergingは、重複する畳み込み層を用いて、特徴マップを空間的にダウンサンプリングし、チャネル次元を2倍に拡張します。具体的には、最初のステージでは、7x7畳み込み(ストライド4)を使用して、特徴マップを4分の1に縮小します。最後の3ステージでは、3x3畳み込み(ストライド2)を使用して、特徴マップを2分の1に縮小します。Patch Mergingの出力は、Pale Transformer Blockに入力され、トークンの数を変更せずに処理されます。
上図(b)に示すように、Pale Transformer Blockは、3つのコンポーネントで構成されています。位置埋め込みを動的に生成するための条件付き位置エンコーディング(CPE)、文脈情報を捉えるためのPS-Attentionモジュール、特徴量変換のためのMLPモジュールです。CPEは、depth-wise convolutionで実装されています。MLPは、2つの全結合層とその間にGELU活性化関数を挟んで構成されています。PS-AttentionとMLPのそれぞれ前後には、レイヤーノルム層と残差結合を行います。

実験

Pale Transformerのバックボーンとしての有効性と汎化性を示すため、ImageNet-1Kで画像分類、ADE20Kでセマンティックセグメンテーション、COCOで物体検出とインスタンスセグメンテーションの実験を実施しています。また、Pale Transformerの主要コンポーネント設計を検証するため、アブレーションスタディも実施しています。ここでは、その一部の検証結果を紹介します。

ImageNet-1Kによるクラス分類

次表は、Pale Transformerと最先端のCNNおよびトランスフォーマーベースのバックボーンによるImageNet-1Kでの性能を比較したものです。LV-VitとVOLOおよびPaleの上に付いている*は、MixToken拡張とtoken labeling lossを使用してトレーニングを行ったことを意味します。

Pale Transformerは、CNNと比較して、RegNetモデルよりも優れた性能を示します。計算複雑度はおおよそ同じです。また、Pale Transformerは、最先端のトランスフォーマーベースのバックボーンを上回り、似たようなモデルサイズとFLOPs下で、すべての変種においてCSWin Transformerよりも0.7%高い精度を達成しています。さらに、MixToken拡張とtoken labeling lossを使用した場合においても、他のモデルよりも優れた性能を示します。

paleサイズの効果

次表は、paleサイズを変更したときのモデルの性能と計算コストの比較を示しています。

paleサイズを1から7に増やすと、すべてのタスクで性能が継続的に向上します。ただし、9まで増やすと、性能が向上しないタスクが発生し、FLOPsが増加します。そのため、デフォルトではすべてのタスクで7に設定しています。ImageNet-1Kにおいて、paleサイズを5から7に変更すると、FLOPsが減少しますが、これは誤字の可能性があります。

PS-Attentionの性能比較

次表は、PS-Attentionを3つの実装で比較した結果と、既存のストライプベースのセルフアテンションとの比較結果を示しています。

Pale(vanilla)は、pale形状を行方向と列方向に分解せず、pale形状全体を対象としてセルフアテンションを行います。Pale(sequential)は、連続するTransformer Blockで、行方向と列方向のセルフアテンションを交互に計算し、Pale(parallel)は、各Block内で行方向と列方向のセルフアテンションを同時に実行します。Pale(parallel)は、すべてのタスクで最も良い結果を達成し、Pale(vanilla)よりもわずかに優れています。Pale(parallel)がPale(vanilla)よりも性能が高い理由として、次のことが考えられます。Pale(vanilla)では、行方向と列方向のpaleの数を等しくなるように($h/s_r=w/s_c$)、パディング操作を行います。このとき、$s_r=s_c$の場合、非正方形の特徴マップが入力されると、過剰なパディング操作が行われます。これが原因で、Pale(vanilla)の性能が低下したのではないかと考えられています。なお、Pale(sequential)とPale(parallel)では、特徴マップの縦横サイズがpaleサイズで割り切れるように($h\%s_r=0$と$w\%s_c=0$)、パディング操作を行います²。
PS-Attentionと既存のストライプベースのセルフアテンションを比較するため、Pale TransformerのPS-Attentionを、Axial-AttentionとCSWin-Attentionにそれぞれ置き換えて評価を行っています。その結果、表に示すように、PS-Attentionは2つのメカニズムよりも優れた性能を示します。

結論

論文では、新しいセルフアテンションのメカニズムとしてPS-Attentionを提案しています。PS-Attentionは、pale領域内でセルフアテンションを行います。これにより、従来のローカルセルフアテンションよりも、より広範囲の文脈情報を考慮して関係性を捉えることができます。また、効率を向上させるため、pale形状を行方向と列方向に分解して、同じレイヤーで同時に計算するように設計しています。PS-Attentionをベースに開発したPale Transformerは、最先端のパフォーマンスを実現し、さまざまな視覚タスクに適用できる高い汎用性を示しています。

paleは、名詞として使われると「杭、柵、フェンス」といった意味があります。ここでのpale形状とは、メッシュフェンスのような、ワイヤーが等間隔に組まれた格子状の形状のことを指します。 ↩
最初に記載した非公式コードでは、パディング操作は行われません。特徴マップの縦横サイズがpaleサイズで割り切れない場合、assert文によってプログラムが終了します。 ↩